while viewer.is_running():
    step_start = time.time()

    # Initial joint positions
    q = data.qpos[dof_ids].copy()

    for _ in range(max_iter):
        # THESE two lines of code, causes bad cube grasping
        # ============================================== #
        data.qpos[dof_ids] = q
        mujoco.mj_forward(model, data)
        # ============================================== #
        
        # Spatial velocity (aka twist).
        dx = data.mocap_pos[mocap_id] - data.site(site_id).xpos
        twist[:3] = Kpos * dx / integration_dt

        twist_norm = np.linalg.norm(twist)
        # Stop iterations
        if twist_norm < tolerance_err:
            break

        # Jacobian.
        mujoco.mj_jacSite(model, data, jac[:3], jac[3:], site_id)

        # Solve system of equations: J @ dq = error.
        dq = jac[:, dof_ids].T @ np.linalg.solve(jac[:, dof_ids] @ jac[:, dof_ids].T + diag, twist)

        # Nullspace control biasing joint velocities towards the home configuration.
        dq += (eye[dof_ids,dof_ids] - np.linalg.pinv(jac[:, dof_ids], rcond=1e-4) @ jac[:, dof_ids]) @ (Kn * (q0 - data.qpos)[dof_ids])

        # Clamp maximum joint velocity.
        dq_abs_max = np.abs(dq).max()
        if dq_abs_max > max_angvel:
            dq *= max_angvel / dq_abs_max

        # Integrate joint velocities to obtain joint positions.
        q = data.qpos[dof_ids].copy()  # Note the copy here is important.
        q += dq * integration_dt
        # mujoco.mj_integratePos(model, q, dq, integration_dt)
        np.clip(q[dof_ids], *model.jnt_range[dof_ids].T, out=q[dof_ids])

    # Set the control signal and step the simulation.
    data.ctrl[actuator_ids] = q[dof_ids]
    mujoco.mj_step(model, data)

    viewer.sync()
    time_until_next_step = dt - (time.time() - step_start)
    if time_until_next_step > 0:
        time.sleep(time_until_next_step)