kinematics

kinematics

Modules:

Name	Description
bimanual_yam_kinematics	Kinematics solver for the bimanual YAM robot.
floating_rum_kinematics
franka_kinematics
i2rt_yam_kinematics	Kinematics solver for the i2rt YAM robot.
mujoco_kinematics	This module provides forward and inverse kinematics functionality for robots in MuJoCo.
rby1_kinematics
stretch_kinematics

bimanual_yam_kinematics

Kinematics solver for the bimanual YAM robot.

Classes:

Name	Description
BimanualYamKinematics	Kinematics solver for the bimanual YAM robot (two 6-DOF arms).

Functions:

Name	Description
main

BimanualYamKinematics

BimanualYamKinematics(model: MjModel, data: MjData | None = None, namespace: str = '', robot_view_factory: RobotViewFactory = BimanualYamRobotView)

Bases: MlSpacesKinematics

Kinematics solver for the bimanual YAM robot (two 6-DOF arms).

Methods:

Name	Description
fk	Compute forward kinematics for all move groups.
ik	Solve inverse kinematics to reach a target pose.
twist_to_joint_vel	Convert an end-effector twist to joint velocities.

Source code in molmo_spaces/kinematics/bimanual_yam_kinematics.py

def __init__(
    self,
    model: MjModel,
    data: MjData | None = None,
    namespace: str = "",
    robot_view_factory: RobotViewFactory = BimanualYamRobotView,
) -> None:
    if data is None:
        data = MjData(model)
    robot_view = robot_view_factory(data, namespace)
    super().__init__(data, robot_view)

fk

fk(move_group_qpos: dict[str, ndarray], base_pose: ndarray, rel_to_base: bool = False) -> dict[str, ndarray]

Compute forward kinematics for all move groups.

This method computes the pose of each move group's leaf frame given the joint positions. The poses can be returned either relative to the world frame or relative to the robot base.

Parameters:

Name	Type	Description	Default
move_group_qpos	dict[str, ndarray]	Dictionary mapping move group IDs to their joint positions	required
base_pose	ndarray	4x4 pose matrix of the robot base.	required
rel_to_base	bool	If True, return poses relative to robot base. If False, return poses in world frame.	False

Returns:

Type	Description
dict[str, ndarray]	Dictionary mapping move group IDs to their 4x4 pose matrices

Source code in molmo_spaces/kinematics/mujoco_kinematics.py

def fk(
    self,
    move_group_qpos: dict[str, np.ndarray],
    base_pose: np.ndarray,
    rel_to_base: bool = False,
) -> dict[str, np.ndarray]:
    """Compute forward kinematics for all move groups.

    This method computes the pose of each move group's leaf frame given the joint positions.
    The poses can be returned either relative to the world frame or relative to the robot base.

    Args:
        move_group_qpos: Dictionary mapping move group IDs to their joint positions
        base_pose: 4x4 pose matrix of the robot base.
        rel_to_base: If True, return poses relative to robot base. If False, return poses in world frame.

    Returns:
        Dictionary mapping move group IDs to their 4x4 pose matrices
    """
    self._robot_view.base.pose = base_pose
    self._robot_view.set_qpos_dict(move_group_qpos)
    mujoco.mj_kinematics(self._mj_model, self._mj_data)
    ret = {}
    for move_group_id in self._robot_view.move_group_ids():
        move_group = self._robot_view.get_move_group(move_group_id)
        if rel_to_base:
            ret[move_group_id] = move_group.leaf_frame_to_robot
        else:
            ret[move_group_id] = self._robot_view.base.pose @ move_group.leaf_frame_to_robot
    return ret

ik

ik(move_group_id: str, pose: ndarray, unlocked_move_group_ids: list[str], q0: dict[str, ndarray], base_pose: ndarray, rel_to_base: bool = False, eps: float = 0.0001, max_iter: int = 1000, damping: float = 1e-12, dt: float = 1.0)

Solve inverse kinematics to reach a target pose.

This method iteratively solves for joint positions that would place the end-effector at the target pose. It uses a damped least squares approach with velocity-based updates.

Parameters:

Name	Type	Description	Default
move_group_id	str	ID of the move group whose end-effector pose is specified	required
pose	ndarray	4x4 target pose matrix	required
unlocked_move_group_ids	list[str]	List of move group IDs whose joints can move	required
q0	dict[str, ndarray]	Dictionary mapping move group IDs to their initial joint positions	required
base_pose	ndarray	4x4 pose matrix of the robot base.	required
rel_to_base	bool	If True, target pose is relative to robot base. If False, pose is in world frame.	False
eps	float	Convergence threshold for the error norm	0.0001
max_iter	int	Maximum number of iterations	1000
damping	float	Damping factor for the damped least squares solution	1e-12
dt	float	Time step for velocity integration	1.0

Returns:

Type	Description
	Dictionary mapping move group IDs to their joint positions if successful, None if failed

Raises:

Type	Description
ValueError	If q0 keys don't match move group IDs

Source code in molmo_spaces/kinematics/mujoco_kinematics.py

def ik(
    self,
    move_group_id: str,
    pose: np.ndarray,
    unlocked_move_group_ids: list[str],
    q0: dict[str, np.ndarray],
    base_pose: np.ndarray,
    rel_to_base: bool = False,
    eps: float = 1e-4,
    max_iter: int = 1000,
    damping: float = 1e-12,
    dt: float = 1.0,
):
    """Solve inverse kinematics to reach a target pose.

    This method iteratively solves for joint positions that would place the end-effector
    at the target pose. It uses a damped least squares approach with velocity-based
    updates.

    Args:
        move_group_id: ID of the move group whose end-effector pose is specified
        pose: 4x4 target pose matrix
        unlocked_move_group_ids: List of move group IDs whose joints can move
        q0: Dictionary mapping move group IDs to their initial joint positions
        base_pose: 4x4 pose matrix of the robot base.
        rel_to_base: If True, target pose is relative to robot base. If False, pose is in world frame.
        eps: Convergence threshold for the error norm
        max_iter: Maximum number of iterations
        damping: Damping factor for the damped least squares solution
        dt: Time step for velocity integration

    Returns:
        Dictionary mapping move group IDs to their joint positions if successful, None if failed

    Raises:
        ValueError: If q0 keys don't match move group IDs
    """
    if set(q0.keys()) != set(self._robot_view.move_group_ids()):
        raise ValueError(
            f"q0 keys must match move group ids: {set(q0.keys())} != {set(self._robot_view.move_group_ids())}"
        )
    self._robot_view.base.pose = base_pose
    self._robot_view.set_qpos_dict(q0)
    if rel_to_base:
        pose = self._robot_view.base.pose @ pose

    move_group = self._robot_view.get_move_group(move_group_id)
    for i in range(max_iter):
        mujoco.mj_fwdPosition(self._mj_model, self._mj_data)
        mujoco.mj_sensorPos(self._mj_model, self._mj_data)

        ee_pose = relative_to_global_transform(
            move_group.leaf_frame_to_robot, self._robot_view.base.pose
        )

        err_trf = inverse_homogeneous_matrix(ee_pose) @ pose
        twist_lin, twist_ang = transform_to_twist(err_trf)

        err = np.concatenate([ee_pose[:3, :3] @ twist_lin, ee_pose[:3, :3] @ twist_ang])
        if np.linalg.norm(err) < eps:
            succ = True
            break
        elif i == max_iter - 1:
            succ = False
            break

        J: np.ndarray = self._robot_view.get_jacobian(move_group_id, unlocked_move_group_ids)
        if (JJT_det := np.linalg.det(J @ J.T)) < 1e-20:
            print(f"WARN: IK Jacobian is rank deficient! det(JJ^T)={JJT_det:.0e}")

        H = J @ J.T + damping * np.eye(J.shape[0])
        q_dot = J.T @ np.linalg.solve(H, err)
        dq = q_dot * dt

        j = 0
        for mg_id in unlocked_move_group_ids:
            mg = self._robot_view.get_move_group(mg_id)
            mg.joint_pos = mg.integrate_joint_vel(mg.joint_pos, dq[j : j + mg.vel_dim])
            j += mg.vel_dim
        self._constrain_state()

    if succ:
        return self._robot_view.get_qpos_dict()
    return None

twist_to_joint_vel

twist_to_joint_vel(move_group_id: str, twist: ndarray, unlocked_move_group_ids: list[str], q0: dict[str, ndarray], base_pose: ndarray, twist_frame: Literal['world', 'base', 'leaf'] = 'world', damping: float = 1e-12)

Convert an end-effector twist to joint velocities.

This method computes the joint velocities that would result in the specified end-effector twist. It uses damped least squares to handle singularities.

Parameters:

Name	Type	Description	Default
move_group_id	str	ID of the move group whose end-effector twist is specified	required
twist	ndarray	6D twist vector (linear and angular velocities)	required
unlocked_move_group_ids	list[str]	List of move group IDs whose joints can move	required
q0	dict[str, ndarray]	Dictionary mapping move group IDs to their current joint positions	required
base_pose	ndarray	4x4 pose matrix of the robot base.	required
twist_frame	Literal['world', 'base', 'leaf']	Frame of the twist, either in the world frame, the robot base frame, or the leaf frame of the move group.	'world'
damping	float	Damping factor for the damped least squares solution	1e-12

Returns:

Type	Description
	Array of joint velocities that would achieve the desired twist

Raises:

Type	Description
ValueError	If q0 keys don't match move group IDs

Source code in molmo_spaces/kinematics/mujoco_kinematics.py

def twist_to_joint_vel(
    self,
    move_group_id: str,
    twist: np.ndarray,
    unlocked_move_group_ids: list[str],
    q0: dict[str, np.ndarray],
    base_pose: np.ndarray,
    twist_frame: Literal["world", "base", "leaf"] = "world",
    damping: float = 1e-12,
):
    """Convert an end-effector twist to joint velocities.

    This method computes the joint velocities that would result in the specified
    end-effector twist. It uses damped least squares to handle singularities.

    Args:
        move_group_id: ID of the move group whose end-effector twist is specified
        twist: 6D twist vector (linear and angular velocities)
        unlocked_move_group_ids: List of move group IDs whose joints can move
        q0: Dictionary mapping move group IDs to their current joint positions
        base_pose: 4x4 pose matrix of the robot base.
        twist_frame: Frame of the twist, either in the world frame, the robot base frame, or the leaf frame of the move group.
        damping: Damping factor for the damped least squares solution

    Returns:
        Array of joint velocities that would achieve the desired twist

    Raises:
        ValueError: If q0 keys don't match move group IDs
    """
    if set(q0.keys()) != set(self._robot_view.move_group_ids()):
        raise ValueError(
            f"q0 keys must match move group ids: {set(q0.keys())} != {set(self._robot_view.move_group_ids())}"
        )
    self._robot_view.base.pose = base_pose
    self._robot_view.set_qpos_dict(q0)

    # convert twist to world frame
    if twist_frame == "base":
        twist[:3] = self._robot_view.base.pose[:3, :3] @ twist[:3]
        twist[3:] = self._robot_view.base.pose[:3, :3] @ twist[3:]
    elif twist_frame == "leaf":
        move_group = self._robot_view.get_move_group(move_group_id)
        twist[:3] = (
            self._robot_view.base.pose @ move_group.leaf_frame_to_robot[:3, :3] @ twist[:3]
        )
        twist[3:] = (
            self._robot_view.base.pose @ move_group.leaf_frame_to_robot[:3, :3] @ twist[3:]
        )
    else:
        assert twist_frame == "world"

    J = self._robot_view.get_jacobian(move_group_id, unlocked_move_group_ids)
    if (JJT_det := np.linalg.det(J @ J.T)) < 1e-20:
        print(f"WARN: IK Jacobian is rank deficient! det(JJ^T)={JJT_det:.0e}")
    H = J @ J.T + damping * np.eye(J.shape[0])
    q_dot = J.T @ np.linalg.solve(H, twist)
    return q_dot

main

main() -> None

Source code in molmo_spaces/kinematics/bimanual_yam_kinematics.py

def main() -> None:
    import mujoco
    import numpy as np
    from mujoco.viewer import launch_passive

    np.set_printoptions(linewidth=np.inf)
    import time

    # Load bimanual YAM robot with a mocap base
    model_xml = """
    <mujoco>
        <asset>
            <model name="bimanual_yam" file="assets/robots/i2rt_yam/bimanual_yam.xml"/>
        </asset>
        <worldbody>
            <body name="base" mocap="true">
                <attach model="bimanual_yam" prefix="" />
            </body>
        </worldbody>
    </mujoco>
    """

    model = MjModel.from_xml_string(model_xml)
    robot_view_factory = BimanualYamRobotView

    data = MjData(model)
    mujoco.mj_forward(model, data)

    ns = ""
    robot_view = robot_view_factory(data, ns)
    left_arm_group = robot_view.get_move_group("left_arm")
    right_arm_group = robot_view.get_move_group("right_arm")
    kinematics = BimanualYamKinematics(
        model, namespace=ns, robot_view_factory=robot_view_factory
    )

    # Set both arms to middle of joint limits
    left_qp = np.mean(left_arm_group.joint_pos_limits, axis=1)
    right_qp = np.mean(right_arm_group.joint_pos_limits, axis=1)
    left_arm_group.joint_pos = left_qp
    right_arm_group.joint_pos = right_qp
    print(f"Initial left arm joint positions: {left_qp}")
    print(f"Initial right arm joint positions: {right_qp}")
    mujoco.mj_forward(model, data)

    # Get current EE poses and create target poses for IK test
    # Left arm: move up/down
    left_pose0 = robot_view.base.pose @ left_arm_group.leaf_frame_to_robot
    left_pose1 = left_pose0.copy()
    left_pose0[2, 3] += 0.05  # Move up 5cm
    left_pose1[2, 3] -= 0.05  # Move down 5cm

    # Right arm: move forward/backward
    right_pose0 = robot_view.base.pose @ right_arm_group.leaf_frame_to_robot
    right_pose1 = right_pose0.copy()
    right_pose0[0, 3] += 0.05  # Move forward 5cm
    right_pose1[0, 3] -= 0.05  # Move backward 5cm

    with launch_passive(model, data) as viewer:
        viewer.sync()
        i = 0
        while viewer.is_running():
            # Alternate between two target poses for each arm
            left_target = left_pose1 if i % 2 == 0 else left_pose0
            right_target = right_pose1 if i % 2 == 0 else right_pose0

            # Run IK for left arm
            qpos_dict = robot_view.get_qpos_dict()
            left_ret = kinematics.ik(
                "left_arm", left_target, ["left_arm"], qpos_dict, robot_view.base.pose
            )
            if left_ret is not None:
                robot_view.set_qpos_dict(left_ret)
                qpos_dict = robot_view.get_qpos_dict()

            # Run IK for right arm
            right_ret = kinematics.ik(
                "right_arm", right_target, ["right_arm"], qpos_dict, robot_view.base.pose
            )
            if right_ret is not None:
                robot_view.set_qpos_dict(right_ret)

            left_status = "Success" if left_ret is not None else "Failed"
            right_status = "Success" if right_ret is not None else "Failed"
            print(f"IK iteration {i}: left={left_status}, right={right_status}")

            i += 1
            mujoco.mj_forward(model, data)
            viewer.sync()
            time.sleep(2)

floating_rum_kinematics

Classes:

Name	Description
FloatingRUMKinematics

Functions:

Name	Description
main

FloatingRUMKinematics

FloatingRUMKinematics(model: MjModel, data: MjData | None = None, namespace: str = '', robot_view_factory: RobotViewFactory = FloatingRUMRobotView)

Bases: MlSpacesKinematics

Methods:

Name	Description
fk	Compute forward kinematics for all move groups.
ik
twist_to_joint_vel	Convert an end-effector twist to joint velocities.

Source code in molmo_spaces/kinematics/floating_rum_kinematics.py

def __init__(
    self,
    model: MjModel,
    data: MjData | None = None,
    namespace: str = "",
    robot_view_factory: RobotViewFactory = FloatingRUMRobotView,
):
    if data is None:
        data = MjData(model)
    robot_view = robot_view_factory(data, namespace)
    super().__init__(data, robot_view)

fk

fk(move_group_qpos: dict[str, ndarray], base_pose: ndarray, rel_to_base: bool = False) -> dict[str, ndarray]

Compute forward kinematics for all move groups.

This method computes the pose of each move group's leaf frame given the joint positions. The poses can be returned either relative to the world frame or relative to the robot base.

Parameters:

Name	Type	Description	Default
move_group_qpos	dict[str, ndarray]	Dictionary mapping move group IDs to their joint positions	required
base_pose	ndarray	4x4 pose matrix of the robot base.	required
rel_to_base	bool	If True, return poses relative to robot base. If False, return poses in world frame.	False

Returns:

Type	Description
dict[str, ndarray]	Dictionary mapping move group IDs to their 4x4 pose matrices

Source code in molmo_spaces/kinematics/mujoco_kinematics.py

def fk(
    self,
    move_group_qpos: dict[str, np.ndarray],
    base_pose: np.ndarray,
    rel_to_base: bool = False,
) -> dict[str, np.ndarray]:
    """Compute forward kinematics for all move groups.

    This method computes the pose of each move group's leaf frame given the joint positions.
    The poses can be returned either relative to the world frame or relative to the robot base.

    Args:
        move_group_qpos: Dictionary mapping move group IDs to their joint positions
        base_pose: 4x4 pose matrix of the robot base.
        rel_to_base: If True, return poses relative to robot base. If False, return poses in world frame.

    Returns:
        Dictionary mapping move group IDs to their 4x4 pose matrices
    """
    self._robot_view.base.pose = base_pose
    self._robot_view.set_qpos_dict(move_group_qpos)
    mujoco.mj_kinematics(self._mj_model, self._mj_data)
    ret = {}
    for move_group_id in self._robot_view.move_group_ids():
        move_group = self._robot_view.get_move_group(move_group_id)
        if rel_to_base:
            ret[move_group_id] = move_group.leaf_frame_to_robot
        else:
            ret[move_group_id] = self._robot_view.base.pose @ move_group.leaf_frame_to_robot
    return ret

ik

ik(move_group_id: str, pose: ndarray, unlocked_move_group_ids: list[str], q0: dict[str, ndarray], base_pose: ndarray, rel_to_base: bool = False, eps: float = 0.0001, max_iter: int = 1000, damping: float = 1e-12, dt: float = 1.0)

Source code in molmo_spaces/kinematics/floating_rum_kinematics.py

def ik(
    self,
    move_group_id: str,
    pose: np.ndarray,
    unlocked_move_group_ids: list[str],
    q0: dict[str, np.ndarray],
    base_pose: np.ndarray,
    rel_to_base: bool = False,
    eps: float = 1e-4,
    max_iter: int = 1000,
    damping: float = 1e-12,
    dt: float = 1.0,
):
    if move_group_id == "gripper":
        ee_to_base = self._robot_view.get_move_group("gripper").leaf_frame_to_robot
        pose = pose @ np.linalg.inv(ee_to_base)
    else:
        assert move_group_id == "base"
    if "base" not in unlocked_move_group_ids:
        return None

    if rel_to_base:
        pose = base_pose @ pose
    pos = pose[:3, 3]
    quat = R.from_matrix(pose[:3, :3]).as_quat(scalar_first=True)
    ret = deepcopy(q0)
    ret["base"] = np.concatenate([pos, quat])
    return ret

twist_to_joint_vel

twist_to_joint_vel(move_group_id: str, twist: ndarray, unlocked_move_group_ids: list[str], q0: dict[str, ndarray], base_pose: ndarray, twist_frame: Literal['world', 'base', 'leaf'] = 'world', damping: float = 1e-12)

Convert an end-effector twist to joint velocities.

This method computes the joint velocities that would result in the specified end-effector twist. It uses damped least squares to handle singularities.

Parameters:

Name	Type	Description	Default
move_group_id	str	ID of the move group whose end-effector twist is specified	required
twist	ndarray	6D twist vector (linear and angular velocities)	required
unlocked_move_group_ids	list[str]	List of move group IDs whose joints can move	required
q0	dict[str, ndarray]	Dictionary mapping move group IDs to their current joint positions	required
base_pose	ndarray	4x4 pose matrix of the robot base.	required
twist_frame	Literal['world', 'base', 'leaf']	Frame of the twist, either in the world frame, the robot base frame, or the leaf frame of the move group.	'world'
damping	float	Damping factor for the damped least squares solution	1e-12

Returns:

Type	Description
	Array of joint velocities that would achieve the desired twist

Raises:

Type	Description
ValueError	If q0 keys don't match move group IDs

Source code in molmo_spaces/kinematics/mujoco_kinematics.py

def twist_to_joint_vel(
    self,
    move_group_id: str,
    twist: np.ndarray,
    unlocked_move_group_ids: list[str],
    q0: dict[str, np.ndarray],
    base_pose: np.ndarray,
    twist_frame: Literal["world", "base", "leaf"] = "world",
    damping: float = 1e-12,
):
    """Convert an end-effector twist to joint velocities.

    This method computes the joint velocities that would result in the specified
    end-effector twist. It uses damped least squares to handle singularities.

    Args:
        move_group_id: ID of the move group whose end-effector twist is specified
        twist: 6D twist vector (linear and angular velocities)
        unlocked_move_group_ids: List of move group IDs whose joints can move
        q0: Dictionary mapping move group IDs to their current joint positions
        base_pose: 4x4 pose matrix of the robot base.
        twist_frame: Frame of the twist, either in the world frame, the robot base frame, or the leaf frame of the move group.
        damping: Damping factor for the damped least squares solution

    Returns:
        Array of joint velocities that would achieve the desired twist

    Raises:
        ValueError: If q0 keys don't match move group IDs
    """
    if set(q0.keys()) != set(self._robot_view.move_group_ids()):
        raise ValueError(
            f"q0 keys must match move group ids: {set(q0.keys())} != {set(self._robot_view.move_group_ids())}"
        )
    self._robot_view.base.pose = base_pose
    self._robot_view.set_qpos_dict(q0)

    # convert twist to world frame
    if twist_frame == "base":
        twist[:3] = self._robot_view.base.pose[:3, :3] @ twist[:3]
        twist[3:] = self._robot_view.base.pose[:3, :3] @ twist[3:]
    elif twist_frame == "leaf":
        move_group = self._robot_view.get_move_group(move_group_id)
        twist[:3] = (
            self._robot_view.base.pose @ move_group.leaf_frame_to_robot[:3, :3] @ twist[:3]
        )
        twist[3:] = (
            self._robot_view.base.pose @ move_group.leaf_frame_to_robot[:3, :3] @ twist[3:]
        )
    else:
        assert twist_frame == "world"

    J = self._robot_view.get_jacobian(move_group_id, unlocked_move_group_ids)
    if (JJT_det := np.linalg.det(J @ J.T)) < 1e-20:
        print(f"WARN: IK Jacobian is rank deficient! det(JJ^T)={JJT_det:.0e}")
    H = J @ J.T + damping * np.eye(J.shape[0])
    q_dot = J.T @ np.linalg.solve(H, twist)
    return q_dot

main

main()

Source code in molmo_spaces/kinematics/floating_rum_kinematics.py

def main():
    import time

    import mujoco
    import numpy as np
    from mujoco.viewer import launch_passive

    np.set_printoptions(linewidth=np.inf)

    xml_path = str(get_robot_paths().get("floating_rum")) + "/model.xml"
    model = MjModel.from_xml_path(xml_path)
    robot_view_factory = FloatingRUMRobotView

    data = MjData(model)
    mujoco.mj_forward(model, data)

    ns = ""
    robot_view = robot_view_factory(data, ns)
    kinematics = FloatingRUMKinematics(
        model, namespace=ns, robot_view_factory=robot_view_factory
    )

    pose0 = np.array(
        [
            [0.7361, -0.6769, 0.0, 0.5325],
            [0.6769, 0.7361, -0.0, -0.0325],
            [0.0, 0.0, 1.0, 0.0734],
            [0.0, 0.0, 0.0, 1.0],
        ]
    )
    pose1 = np.array(
        [
            [0.23935042, 0.84872045, 0.47157711, 1.05149006],
            [0.92900645, -0.34137376, 0.14286698, 0.4176159],
            [0.28223818, 0.40390291, -0.87017472, 0.88686037],
            [0.0, 0.0, 0.0, 1.0],
        ]
    )

    robot_view.base.pose = pose0
    mujoco.mj_forward(model, data)

    groups = [
        "base",
    ]

    with launch_passive(model, data) as viewer:
        _show_poses(viewer, np.array([pose0, pose1]))
        viewer.sync()
        i = 0
        while viewer.is_running():
            if i % 2 == 0:
                ret = kinematics.ik(
                    "gripper",
                    pose1,
                    groups,
                    robot_view.get_qpos_dict(),
                    robot_view.base.pose,
                    dt=0.1,
                )
            else:
                ret = kinematics.ik(
                    "gripper", pose0, groups, robot_view.get_qpos_dict(), robot_view.base.pose
                )
            print(i % 2, ret)
            i += 1
            if ret is not None:
                robot_view.set_qpos_dict(ret)
            mujoco.mj_forward(model, data)
            viewer.sync()
            time.sleep(3)

franka_kinematics

Classes:

Name	Description
FrankaKinematics

Functions:

Name	Description
main

FrankaKinematics

FrankaKinematics(model: MjModel, data: MjData | None = None, namespace: str = '', robot_view_factory: RobotViewFactory = FrankaFR3RobotView)

Bases: MlSpacesKinematics

Methods:

Name	Description
fk	Compute forward kinematics for all move groups.
ik	Solve inverse kinematics to reach a target pose.
twist_to_joint_vel	Convert an end-effector twist to joint velocities.

Source code in molmo_spaces/kinematics/franka_kinematics.py

def __init__(
    self,
    model: MjModel,
    data: MjData | None = None,
    namespace: str = "",
    robot_view_factory: RobotViewFactory = FrankaFR3RobotView,
) -> None:
    if data is None:
        data = MjData(model)
    robot_view = robot_view_factory(data, namespace)
    super().__init__(data, robot_view)

fk

fk(move_group_qpos: dict[str, ndarray], base_pose: ndarray, rel_to_base: bool = False) -> dict[str, ndarray]

Compute forward kinematics for all move groups.

This method computes the pose of each move group's leaf frame given the joint positions. The poses can be returned either relative to the world frame or relative to the robot base.

Parameters:

Name	Type	Description	Default
move_group_qpos	dict[str, ndarray]	Dictionary mapping move group IDs to their joint positions	required
base_pose	ndarray	4x4 pose matrix of the robot base.	required
rel_to_base	bool	If True, return poses relative to robot base. If False, return poses in world frame.	False

Returns:

Type	Description
dict[str, ndarray]	Dictionary mapping move group IDs to their 4x4 pose matrices

Source code in molmo_spaces/kinematics/mujoco_kinematics.py

def fk(
    self,
    move_group_qpos: dict[str, np.ndarray],
    base_pose: np.ndarray,
    rel_to_base: bool = False,
) -> dict[str, np.ndarray]:
    """Compute forward kinematics for all move groups.

    This method computes the pose of each move group's leaf frame given the joint positions.
    The poses can be returned either relative to the world frame or relative to the robot base.

    Args:
        move_group_qpos: Dictionary mapping move group IDs to their joint positions
        base_pose: 4x4 pose matrix of the robot base.
        rel_to_base: If True, return poses relative to robot base. If False, return poses in world frame.

    Returns:
        Dictionary mapping move group IDs to their 4x4 pose matrices
    """
    self._robot_view.base.pose = base_pose
    self._robot_view.set_qpos_dict(move_group_qpos)
    mujoco.mj_kinematics(self._mj_model, self._mj_data)
    ret = {}
    for move_group_id in self._robot_view.move_group_ids():
        move_group = self._robot_view.get_move_group(move_group_id)
        if rel_to_base:
            ret[move_group_id] = move_group.leaf_frame_to_robot
        else:
            ret[move_group_id] = self._robot_view.base.pose @ move_group.leaf_frame_to_robot
    return ret

ik

ik(move_group_id: str, pose: ndarray, unlocked_move_group_ids: list[str], q0: dict[str, ndarray], base_pose: ndarray, rel_to_base: bool = False, eps: float = 0.0001, max_iter: int = 1000, damping: float = 1e-12, dt: float = 1.0)

Solve inverse kinematics to reach a target pose.

This method iteratively solves for joint positions that would place the end-effector at the target pose. It uses a damped least squares approach with velocity-based updates.

Parameters:

Name	Type	Description	Default
move_group_id	str	ID of the move group whose end-effector pose is specified	required
pose	ndarray	4x4 target pose matrix	required
unlocked_move_group_ids	list[str]	List of move group IDs whose joints can move	required
q0	dict[str, ndarray]	Dictionary mapping move group IDs to their initial joint positions	required
base_pose	ndarray	4x4 pose matrix of the robot base.	required
rel_to_base	bool	If True, target pose is relative to robot base. If False, pose is in world frame.	False
eps	float	Convergence threshold for the error norm	0.0001
max_iter	int	Maximum number of iterations	1000
damping	float	Damping factor for the damped least squares solution	1e-12
dt	float	Time step for velocity integration	1.0

Returns:

Type	Description
	Dictionary mapping move group IDs to their joint positions if successful, None if failed

Raises:

Type	Description
ValueError	If q0 keys don't match move group IDs

Source code in molmo_spaces/kinematics/mujoco_kinematics.py

def ik(
    self,
    move_group_id: str,
    pose: np.ndarray,
    unlocked_move_group_ids: list[str],
    q0: dict[str, np.ndarray],
    base_pose: np.ndarray,
    rel_to_base: bool = False,
    eps: float = 1e-4,
    max_iter: int = 1000,
    damping: float = 1e-12,
    dt: float = 1.0,
):
    """Solve inverse kinematics to reach a target pose.

    This method iteratively solves for joint positions that would place the end-effector
    at the target pose. It uses a damped least squares approach with velocity-based
    updates.

    Args:
        move_group_id: ID of the move group whose end-effector pose is specified
        pose: 4x4 target pose matrix
        unlocked_move_group_ids: List of move group IDs whose joints can move
        q0: Dictionary mapping move group IDs to their initial joint positions
        base_pose: 4x4 pose matrix of the robot base.
        rel_to_base: If True, target pose is relative to robot base. If False, pose is in world frame.
        eps: Convergence threshold for the error norm
        max_iter: Maximum number of iterations
        damping: Damping factor for the damped least squares solution
        dt: Time step for velocity integration

    Returns:
        Dictionary mapping move group IDs to their joint positions if successful, None if failed

    Raises:
        ValueError: If q0 keys don't match move group IDs
    """
    if set(q0.keys()) != set(self._robot_view.move_group_ids()):
        raise ValueError(
            f"q0 keys must match move group ids: {set(q0.keys())} != {set(self._robot_view.move_group_ids())}"
        )
    self._robot_view.base.pose = base_pose
    self._robot_view.set_qpos_dict(q0)
    if rel_to_base:
        pose = self._robot_view.base.pose @ pose

    move_group = self._robot_view.get_move_group(move_group_id)
    for i in range(max_iter):
        mujoco.mj_fwdPosition(self._mj_model, self._mj_data)
        mujoco.mj_sensorPos(self._mj_model, self._mj_data)

        ee_pose = relative_to_global_transform(
            move_group.leaf_frame_to_robot, self._robot_view.base.pose
        )

        err_trf = inverse_homogeneous_matrix(ee_pose) @ pose
        twist_lin, twist_ang = transform_to_twist(err_trf)

        err = np.concatenate([ee_pose[:3, :3] @ twist_lin, ee_pose[:3, :3] @ twist_ang])
        if np.linalg.norm(err) < eps:
            succ = True
            break
        elif i == max_iter - 1:
            succ = False
            break

        J: np.ndarray = self._robot_view.get_jacobian(move_group_id, unlocked_move_group_ids)
        if (JJT_det := np.linalg.det(J @ J.T)) < 1e-20:
            print(f"WARN: IK Jacobian is rank deficient! det(JJ^T)={JJT_det:.0e}")

        H = J @ J.T + damping * np.eye(J.shape[0])
        q_dot = J.T @ np.linalg.solve(H, err)
        dq = q_dot * dt

        j = 0
        for mg_id in unlocked_move_group_ids:
            mg = self._robot_view.get_move_group(mg_id)
            mg.joint_pos = mg.integrate_joint_vel(mg.joint_pos, dq[j : j + mg.vel_dim])
            j += mg.vel_dim
        self._constrain_state()

    if succ:
        return self._robot_view.get_qpos_dict()
    return None

twist_to_joint_vel

twist_to_joint_vel(move_group_id: str, twist: ndarray, unlocked_move_group_ids: list[str], q0: dict[str, ndarray], base_pose: ndarray, twist_frame: Literal['world', 'base', 'leaf'] = 'world', damping: float = 1e-12)

Convert an end-effector twist to joint velocities.

This method computes the joint velocities that would result in the specified end-effector twist. It uses damped least squares to handle singularities.

Parameters:

Name	Type	Description	Default
move_group_id	str	ID of the move group whose end-effector twist is specified	required
twist	ndarray	6D twist vector (linear and angular velocities)	required
unlocked_move_group_ids	list[str]	List of move group IDs whose joints can move	required
q0	dict[str, ndarray]	Dictionary mapping move group IDs to their current joint positions	required
base_pose	ndarray	4x4 pose matrix of the robot base.	required
twist_frame	Literal['world', 'base', 'leaf']	Frame of the twist, either in the world frame, the robot base frame, or the leaf frame of the move group.	'world'
damping	float	Damping factor for the damped least squares solution	1e-12

Returns:

Type	Description
	Array of joint velocities that would achieve the desired twist

Raises:

Type	Description
ValueError	If q0 keys don't match move group IDs

Source code in molmo_spaces/kinematics/mujoco_kinematics.py

def twist_to_joint_vel(
    self,
    move_group_id: str,
    twist: np.ndarray,
    unlocked_move_group_ids: list[str],
    q0: dict[str, np.ndarray],
    base_pose: np.ndarray,
    twist_frame: Literal["world", "base", "leaf"] = "world",
    damping: float = 1e-12,
):
    """Convert an end-effector twist to joint velocities.

    This method computes the joint velocities that would result in the specified
    end-effector twist. It uses damped least squares to handle singularities.

    Args:
        move_group_id: ID of the move group whose end-effector twist is specified
        twist: 6D twist vector (linear and angular velocities)
        unlocked_move_group_ids: List of move group IDs whose joints can move
        q0: Dictionary mapping move group IDs to their current joint positions
        base_pose: 4x4 pose matrix of the robot base.
        twist_frame: Frame of the twist, either in the world frame, the robot base frame, or the leaf frame of the move group.
        damping: Damping factor for the damped least squares solution

    Returns:
        Array of joint velocities that would achieve the desired twist

    Raises:
        ValueError: If q0 keys don't match move group IDs
    """
    if set(q0.keys()) != set(self._robot_view.move_group_ids()):
        raise ValueError(
            f"q0 keys must match move group ids: {set(q0.keys())} != {set(self._robot_view.move_group_ids())}"
        )
    self._robot_view.base.pose = base_pose
    self._robot_view.set_qpos_dict(q0)

    # convert twist to world frame
    if twist_frame == "base":
        twist[:3] = self._robot_view.base.pose[:3, :3] @ twist[:3]
        twist[3:] = self._robot_view.base.pose[:3, :3] @ twist[3:]
    elif twist_frame == "leaf":
        move_group = self._robot_view.get_move_group(move_group_id)
        twist[:3] = (
            self._robot_view.base.pose @ move_group.leaf_frame_to_robot[:3, :3] @ twist[:3]
        )
        twist[3:] = (
            self._robot_view.base.pose @ move_group.leaf_frame_to_robot[:3, :3] @ twist[3:]
        )
    else:
        assert twist_frame == "world"

    J = self._robot_view.get_jacobian(move_group_id, unlocked_move_group_ids)
    if (JJT_det := np.linalg.det(J @ J.T)) < 1e-20:
        print(f"WARN: IK Jacobian is rank deficient! det(JJ^T)={JJT_det:.0e}")
    H = J @ J.T + damping * np.eye(J.shape[0])
    q_dot = J.T @ np.linalg.solve(H, twist)
    return q_dot

main

main() -> None

Source code in molmo_spaces/kinematics/franka_kinematics.py

def main() -> None:
    import mujoco
    import numpy as np
    from mujoco.viewer import launch_passive

    np.set_printoptions(linewidth=np.inf)
    import time

    model_xml = """
    <mujoco>
        <asset>
            <model name="franka" file="assets/robots/franka_fr3/fr3.xml"/>
        </asset>
        <worldbody>
            <body name="base" mocap="true">
                <attach model="franka" body="fr3_link0" prefix="" />
            </body>
        </worldbody>
    </mujoco>
    """

    model = MjModel.from_xml_string(model_xml)
    robot_view_factory = FrankaFR3RobotView

    data = MjData(model)
    mujoco.mj_forward(model, data)

    ns = ""
    robot_view = robot_view_factory(data, ns)
    arm_group = robot_view.get_move_group("arm")
    kinematics = FrankaKinematics(model, namespace=ns, robot_view_factory=robot_view_factory)

    qp = np.mean(arm_group.joint_pos_limits, axis=1)
    arm_group.joint_pos = qp
    print(qp)
    mujoco.mj_forward(model, data)

    pose0 = robot_view.base.pose @ arm_group.leaf_frame_to_robot
    pose1 = pose0.copy()
    pose0[2, 3] += 0.1
    pose1[2, 3] -= 0.1

    groups = [
        "arm",
    ]

    with launch_passive(model, data) as viewer:
        viewer.sync()
        i = 0
        while viewer.is_running():
            if i % 2 == 0:
                ret = kinematics.ik(
                    "arm", pose1, groups, robot_view.get_qpos_dict(), robot_view.base.pose
                )
            else:
                ret = kinematics.ik(
                    "arm", pose0, groups, robot_view.get_qpos_dict(), robot_view.base.pose
                )
            print(i % 2, ret)
            i += 1
            if ret is not None:
                robot_view.set_qpos_dict(ret)
            mujoco.mj_forward(model, data)
            viewer.sync()
            time.sleep(3)

i2rt_yam_kinematics

Kinematics solver for the i2rt YAM robot.

Classes:

Name	Description
I2rtYamKinematics	Kinematics solver for the i2rt YAM 6-DOF arm.

Functions:

Name	Description
main

I2rtYamKinematics

I2rtYamKinematics(model: MjModel, data: MjData | None = None, namespace: str = '', robot_view_factory: RobotViewFactory = I2rtYamRobotView)

Bases: MlSpacesKinematics

Kinematics solver for the i2rt YAM 6-DOF arm.

Methods:

Name	Description
fk	Compute forward kinematics for all move groups.
ik	Solve inverse kinematics to reach a target pose.
twist_to_joint_vel	Convert an end-effector twist to joint velocities.

Source code in molmo_spaces/kinematics/i2rt_yam_kinematics.py

def __init__(
    self,
    model: MjModel,
    data: MjData | None = None,
    namespace: str = "",
    robot_view_factory: RobotViewFactory = I2rtYamRobotView,
) -> None:
    if data is None:
        data = MjData(model)
    robot_view = robot_view_factory(data, namespace)
    super().__init__(data, robot_view)

fk

fk(move_group_qpos: dict[str, ndarray], base_pose: ndarray, rel_to_base: bool = False) -> dict[str, ndarray]

Compute forward kinematics for all move groups.

This method computes the pose of each move group's leaf frame given the joint positions. The poses can be returned either relative to the world frame or relative to the robot base.

Parameters:

Name	Type	Description	Default
move_group_qpos	dict[str, ndarray]	Dictionary mapping move group IDs to their joint positions	required
base_pose	ndarray	4x4 pose matrix of the robot base.	required
rel_to_base	bool	If True, return poses relative to robot base. If False, return poses in world frame.	False

Returns:

Type	Description
dict[str, ndarray]	Dictionary mapping move group IDs to their 4x4 pose matrices

Source code in molmo_spaces/kinematics/mujoco_kinematics.py

def fk(
    self,
    move_group_qpos: dict[str, np.ndarray],
    base_pose: np.ndarray,
    rel_to_base: bool = False,
) -> dict[str, np.ndarray]:
    """Compute forward kinematics for all move groups.

    This method computes the pose of each move group's leaf frame given the joint positions.
    The poses can be returned either relative to the world frame or relative to the robot base.

    Args:
        move_group_qpos: Dictionary mapping move group IDs to their joint positions
        base_pose: 4x4 pose matrix of the robot base.
        rel_to_base: If True, return poses relative to robot base. If False, return poses in world frame.

    Returns:
        Dictionary mapping move group IDs to their 4x4 pose matrices
    """
    self._robot_view.base.pose = base_pose
    self._robot_view.set_qpos_dict(move_group_qpos)
    mujoco.mj_kinematics(self._mj_model, self._mj_data)
    ret = {}
    for move_group_id in self._robot_view.move_group_ids():
        move_group = self._robot_view.get_move_group(move_group_id)
        if rel_to_base:
            ret[move_group_id] = move_group.leaf_frame_to_robot
        else:
            ret[move_group_id] = self._robot_view.base.pose @ move_group.leaf_frame_to_robot
    return ret

ik

ik(move_group_id: str, pose: ndarray, unlocked_move_group_ids: list[str], q0: dict[str, ndarray], base_pose: ndarray, rel_to_base: bool = False, eps: float = 0.0001, max_iter: int = 1000, damping: float = 1e-12, dt: float = 1.0)

Solve inverse kinematics to reach a target pose.

This method iteratively solves for joint positions that would place the end-effector at the target pose. It uses a damped least squares approach with velocity-based updates.

Parameters:

Name	Type	Description	Default
move_group_id	str	ID of the move group whose end-effector pose is specified	required
pose	ndarray	4x4 target pose matrix	required
unlocked_move_group_ids	list[str]	List of move group IDs whose joints can move	required
q0	dict[str, ndarray]	Dictionary mapping move group IDs to their initial joint positions	required
base_pose	ndarray	4x4 pose matrix of the robot base.	required
rel_to_base	bool	If True, target pose is relative to robot base. If False, pose is in world frame.	False
eps	float	Convergence threshold for the error norm	0.0001
max_iter	int	Maximum number of iterations	1000
damping	float	Damping factor for the damped least squares solution	1e-12
dt	float	Time step for velocity integration	1.0

Returns:

Type	Description
	Dictionary mapping move group IDs to their joint positions if successful, None if failed

Raises:

Type	Description
ValueError	If q0 keys don't match move group IDs

Source code in molmo_spaces/kinematics/mujoco_kinematics.py

def ik(
    self,
    move_group_id: str,
    pose: np.ndarray,
    unlocked_move_group_ids: list[str],
    q0: dict[str, np.ndarray],
    base_pose: np.ndarray,
    rel_to_base: bool = False,
    eps: float = 1e-4,
    max_iter: int = 1000,
    damping: float = 1e-12,
    dt: float = 1.0,
):
    """Solve inverse kinematics to reach a target pose.

    This method iteratively solves for joint positions that would place the end-effector
    at the target pose. It uses a damped least squares approach with velocity-based
    updates.

    Args:
        move_group_id: ID of the move group whose end-effector pose is specified
        pose: 4x4 target pose matrix
        unlocked_move_group_ids: List of move group IDs whose joints can move
        q0: Dictionary mapping move group IDs to their initial joint positions
        base_pose: 4x4 pose matrix of the robot base.
        rel_to_base: If True, target pose is relative to robot base. If False, pose is in world frame.
        eps: Convergence threshold for the error norm
        max_iter: Maximum number of iterations
        damping: Damping factor for the damped least squares solution
        dt: Time step for velocity integration

    Returns:
        Dictionary mapping move group IDs to their joint positions if successful, None if failed

    Raises:
        ValueError: If q0 keys don't match move group IDs
    """
    if set(q0.keys()) != set(self._robot_view.move_group_ids()):
        raise ValueError(
            f"q0 keys must match move group ids: {set(q0.keys())} != {set(self._robot_view.move_group_ids())}"
        )
    self._robot_view.base.pose = base_pose
    self._robot_view.set_qpos_dict(q0)
    if rel_to_base:
        pose = self._robot_view.base.pose @ pose

    move_group = self._robot_view.get_move_group(move_group_id)
    for i in range(max_iter):
        mujoco.mj_fwdPosition(self._mj_model, self._mj_data)
        mujoco.mj_sensorPos(self._mj_model, self._mj_data)

        ee_pose = relative_to_global_transform(
            move_group.leaf_frame_to_robot, self._robot_view.base.pose
        )

        err_trf = inverse_homogeneous_matrix(ee_pose) @ pose
        twist_lin, twist_ang = transform_to_twist(err_trf)

        err = np.concatenate([ee_pose[:3, :3] @ twist_lin, ee_pose[:3, :3] @ twist_ang])
        if np.linalg.norm(err) < eps:
            succ = True
            break
        elif i == max_iter - 1:
            succ = False
            break

        J: np.ndarray = self._robot_view.get_jacobian(move_group_id, unlocked_move_group_ids)
        if (JJT_det := np.linalg.det(J @ J.T)) < 1e-20:
            print(f"WARN: IK Jacobian is rank deficient! det(JJ^T)={JJT_det:.0e}")

        H = J @ J.T + damping * np.eye(J.shape[0])
        q_dot = J.T @ np.linalg.solve(H, err)
        dq = q_dot * dt

        j = 0
        for mg_id in unlocked_move_group_ids:
            mg = self._robot_view.get_move_group(mg_id)
            mg.joint_pos = mg.integrate_joint_vel(mg.joint_pos, dq[j : j + mg.vel_dim])
            j += mg.vel_dim
        self._constrain_state()

    if succ:
        return self._robot_view.get_qpos_dict()
    return None

twist_to_joint_vel

twist_to_joint_vel(move_group_id: str, twist: ndarray, unlocked_move_group_ids: list[str], q0: dict[str, ndarray], base_pose: ndarray, twist_frame: Literal['world', 'base', 'leaf'] = 'world', damping: float = 1e-12)

Convert an end-effector twist to joint velocities.

This method computes the joint velocities that would result in the specified end-effector twist. It uses damped least squares to handle singularities.

Parameters:

Name	Type	Description	Default
move_group_id	str	ID of the move group whose end-effector twist is specified	required
twist	ndarray	6D twist vector (linear and angular velocities)	required
unlocked_move_group_ids	list[str]	List of move group IDs whose joints can move	required
q0	dict[str, ndarray]	Dictionary mapping move group IDs to their current joint positions	required
base_pose	ndarray	4x4 pose matrix of the robot base.	required
twist_frame	Literal['world', 'base', 'leaf']	Frame of the twist, either in the world frame, the robot base frame, or the leaf frame of the move group.	'world'
damping	float	Damping factor for the damped least squares solution	1e-12

Returns:

Type	Description
	Array of joint velocities that would achieve the desired twist

Raises:

Type	Description
ValueError	If q0 keys don't match move group IDs

Source code in molmo_spaces/kinematics/mujoco_kinematics.py

def twist_to_joint_vel(
    self,
    move_group_id: str,
    twist: np.ndarray,
    unlocked_move_group_ids: list[str],
    q0: dict[str, np.ndarray],
    base_pose: np.ndarray,
    twist_frame: Literal["world", "base", "leaf"] = "world",
    damping: float = 1e-12,
):
    """Convert an end-effector twist to joint velocities.

    This method computes the joint velocities that would result in the specified
    end-effector twist. It uses damped least squares to handle singularities.

    Args:
        move_group_id: ID of the move group whose end-effector twist is specified
        twist: 6D twist vector (linear and angular velocities)
        unlocked_move_group_ids: List of move group IDs whose joints can move
        q0: Dictionary mapping move group IDs to their current joint positions
        base_pose: 4x4 pose matrix of the robot base.
        twist_frame: Frame of the twist, either in the world frame, the robot base frame, or the leaf frame of the move group.
        damping: Damping factor for the damped least squares solution

    Returns:
        Array of joint velocities that would achieve the desired twist

    Raises:
        ValueError: If q0 keys don't match move group IDs
    """
    if set(q0.keys()) != set(self._robot_view.move_group_ids()):
        raise ValueError(
            f"q0 keys must match move group ids: {set(q0.keys())} != {set(self._robot_view.move_group_ids())}"
        )
    self._robot_view.base.pose = base_pose
    self._robot_view.set_qpos_dict(q0)

    # convert twist to world frame
    if twist_frame == "base":
        twist[:3] = self._robot_view.base.pose[:3, :3] @ twist[:3]
        twist[3:] = self._robot_view.base.pose[:3, :3] @ twist[3:]
    elif twist_frame == "leaf":
        move_group = self._robot_view.get_move_group(move_group_id)
        twist[:3] = (
            self._robot_view.base.pose @ move_group.leaf_frame_to_robot[:3, :3] @ twist[:3]
        )
        twist[3:] = (
            self._robot_view.base.pose @ move_group.leaf_frame_to_robot[:3, :3] @ twist[3:]
        )
    else:
        assert twist_frame == "world"

    J = self._robot_view.get_jacobian(move_group_id, unlocked_move_group_ids)
    if (JJT_det := np.linalg.det(J @ J.T)) < 1e-20:
        print(f"WARN: IK Jacobian is rank deficient! det(JJ^T)={JJT_det:.0e}")
    H = J @ J.T + damping * np.eye(J.shape[0])
    q_dot = J.T @ np.linalg.solve(H, twist)
    return q_dot

main

main() -> None

Source code in molmo_spaces/kinematics/i2rt_yam_kinematics.py

def main() -> None:
    import mujoco
    import numpy as np
    from mujoco.viewer import launch_passive

    np.set_printoptions(linewidth=np.inf)
    import time

    # Load YAM robot with a mocap base (similar to how add_robot_to_scene works)
    model_xml = """
    <mujoco>
        <asset>
            <model name="yam" file="assets/robots/i2rt_yam/yam.xml"/>
        </asset>
        <worldbody>
            <body name="base" mocap="true">
                <attach model="yam" body="arm" prefix="" />
            </body>
        </worldbody>
    </mujoco>
    """

    model = MjModel.from_xml_string(model_xml)
    robot_view_factory = I2rtYamRobotView

    data = MjData(model)
    mujoco.mj_forward(model, data)

    ns = ""
    robot_view = robot_view_factory(data, ns)
    arm_group = robot_view.get_move_group("arm")
    kinematics = I2rtYamKinematics(model, namespace=ns, robot_view_factory=robot_view_factory)

    # Set arm to middle of joint limits
    qp = np.mean(arm_group.joint_pos_limits, axis=1)
    arm_group.joint_pos = qp
    print(f"Initial joint positions: {qp}")
    mujoco.mj_forward(model, data)

    # Get current EE pose and create two target poses (up and down)
    pose0 = robot_view.base.pose @ arm_group.leaf_frame_to_robot
    pose1 = pose0.copy()
    pose0[2, 3] += 0.05  # Move up 5cm
    pose1[2, 3] -= 0.05  # Move down 5cm

    groups = ["arm"]

    with launch_passive(model, data) as viewer:
        viewer.sync()
        i = 0
        while viewer.is_running():
            # Alternate between two target poses
            target_pose = pose1 if i % 2 == 0 else pose0
            ret = kinematics.ik(
                "arm", target_pose, groups, robot_view.get_qpos_dict(), robot_view.base.pose
            )
            print(f"IK iteration {i}: {'Success' if ret is not None else 'Failed'}")
            i += 1
            if ret is not None:
                robot_view.set_qpos_dict(ret)
            mujoco.mj_forward(model, data)
            viewer.sync()
            time.sleep(2)

mujoco_kinematics

This module provides forward and inverse kinematics functionality for robots in MuJoCo. It implements both forward kinematics (FK) and inverse kinematics (IK) solvers, as well as methods for converting between joint velocities and end-effector twists.

The main class, MujocoKinematics, provides a general-purpose interface for computing kinematic quantities for any robot that can be represented in MuJoCo. It works with the RobotView abstraction to handle different types of robots and their move groups.

Classes:

Name	Description
MlSpacesKinematics	A general-purpose kinematics solver for robots in MuJoCo.

MlSpacesKinematics

MlSpacesKinematics(data: MjData, robot_view: RobotView)

A general-purpose kinematics solver for robots in MuJoCo.

This class provides methods for computing forward kinematics (FK), inverse kinematics (IK), and converting between joint velocities and end-effector twists. It works with any robot that can be represented in MuJoCo through the RobotView abstraction.

The class handles both world-frame and robot-base-frame relative poses and twists, making it flexible for different use cases.

Initialize the kinematics solver. This constructor will directly use the data and robot_view objects for internal computations. Subclasses can copy the passed-in data before invoking this super constructor to maintain a private copy of the data, so as to not conflict with client code.

Parameters:

Name	Type	Description	Default
data	MjData	The simulation state. This object is directly used internally for kinematics computations.	required
robot_view	RobotView	A RobotView instance bound to data representing the robot to compute kinematics for	required

Methods:

Name	Description
fk	Compute forward kinematics for all move groups.
ik	Solve inverse kinematics to reach a target pose.
twist_to_joint_vel	Convert an end-effector twist to joint velocities.

Source code in molmo_spaces/kinematics/mujoco_kinematics.py

def __init__(self, data: MjData, robot_view: RobotView) -> None:
    """Initialize the kinematics solver.
    This constructor will directly use the data and robot_view objects for internal computations.
    Subclasses can copy the passed-in data before invoking this super constructor to maintain
    a private copy of the data, so as to not conflict with client code.

    Args:
        data: The simulation state. This object is directly used internally for kinematics computations.
        robot_view: A RobotView instance bound to data representing the robot to compute kinematics for
    """
    self._mj_model = data.model
    self._mj_data = data
    self._robot_view = robot_view
    mujoco.mj_forward(self._mj_model, self._mj_data)

fk

fk(move_group_qpos: dict[str, ndarray], base_pose: ndarray, rel_to_base: bool = False) -> dict[str, ndarray]

Compute forward kinematics for all move groups.

This method computes the pose of each move group's leaf frame given the joint positions. The poses can be returned either relative to the world frame or relative to the robot base.

Parameters:

Name	Type	Description	Default
move_group_qpos	dict[str, ndarray]	Dictionary mapping move group IDs to their joint positions	required
base_pose	ndarray	4x4 pose matrix of the robot base.	required
rel_to_base	bool	If True, return poses relative to robot base. If False, return poses in world frame.	False

Returns:

Type	Description
dict[str, ndarray]	Dictionary mapping move group IDs to their 4x4 pose matrices

Source code in molmo_spaces/kinematics/mujoco_kinematics.py

def fk(
    self,
    move_group_qpos: dict[str, np.ndarray],
    base_pose: np.ndarray,
    rel_to_base: bool = False,
) -> dict[str, np.ndarray]:
    """Compute forward kinematics for all move groups.

    This method computes the pose of each move group's leaf frame given the joint positions.
    The poses can be returned either relative to the world frame or relative to the robot base.

    Args:
        move_group_qpos: Dictionary mapping move group IDs to their joint positions
        base_pose: 4x4 pose matrix of the robot base.
        rel_to_base: If True, return poses relative to robot base. If False, return poses in world frame.

    Returns:
        Dictionary mapping move group IDs to their 4x4 pose matrices
    """
    self._robot_view.base.pose = base_pose
    self._robot_view.set_qpos_dict(move_group_qpos)
    mujoco.mj_kinematics(self._mj_model, self._mj_data)
    ret = {}
    for move_group_id in self._robot_view.move_group_ids():
        move_group = self._robot_view.get_move_group(move_group_id)
        if rel_to_base:
            ret[move_group_id] = move_group.leaf_frame_to_robot
        else:
            ret[move_group_id] = self._robot_view.base.pose @ move_group.leaf_frame_to_robot
    return ret

ik

ik(move_group_id: str, pose: ndarray, unlocked_move_group_ids: list[str], q0: dict[str, ndarray], base_pose: ndarray, rel_to_base: bool = False, eps: float = 0.0001, max_iter: int = 1000, damping: float = 1e-12, dt: float = 1.0)

Solve inverse kinematics to reach a target pose.

This method iteratively solves for joint positions that would place the end-effector at the target pose. It uses a damped least squares approach with velocity-based updates.

Parameters:

Name	Type	Description	Default
move_group_id	str	ID of the move group whose end-effector pose is specified	required
pose	ndarray	4x4 target pose matrix	required
unlocked_move_group_ids	list[str]	List of move group IDs whose joints can move	required
q0	dict[str, ndarray]	Dictionary mapping move group IDs to their initial joint positions	required
base_pose	ndarray	4x4 pose matrix of the robot base.	required
rel_to_base	bool	If True, target pose is relative to robot base. If False, pose is in world frame.	False
eps	float	Convergence threshold for the error norm	0.0001
max_iter	int	Maximum number of iterations	1000
damping	float	Damping factor for the damped least squares solution	1e-12
dt	float	Time step for velocity integration	1.0

Returns:

Type	Description
	Dictionary mapping move group IDs to their joint positions if successful, None if failed

Raises:

Type	Description
ValueError	If q0 keys don't match move group IDs

Source code in molmo_spaces/kinematics/mujoco_kinematics.py

def ik(
    self,
    move_group_id: str,
    pose: np.ndarray,
    unlocked_move_group_ids: list[str],
    q0: dict[str, np.ndarray],
    base_pose: np.ndarray,
    rel_to_base: bool = False,
    eps: float = 1e-4,
    max_iter: int = 1000,
    damping: float = 1e-12,
    dt: float = 1.0,
):
    """Solve inverse kinematics to reach a target pose.

    This method iteratively solves for joint positions that would place the end-effector
    at the target pose. It uses a damped least squares approach with velocity-based
    updates.

    Args:
        move_group_id: ID of the move group whose end-effector pose is specified
        pose: 4x4 target pose matrix
        unlocked_move_group_ids: List of move group IDs whose joints can move
        q0: Dictionary mapping move group IDs to their initial joint positions
        base_pose: 4x4 pose matrix of the robot base.
        rel_to_base: If True, target pose is relative to robot base. If False, pose is in world frame.
        eps: Convergence threshold for the error norm
        max_iter: Maximum number of iterations
        damping: Damping factor for the damped least squares solution
        dt: Time step for velocity integration

    Returns:
        Dictionary mapping move group IDs to their joint positions if successful, None if failed

    Raises:
        ValueError: If q0 keys don't match move group IDs
    """
    if set(q0.keys()) != set(self._robot_view.move_group_ids()):
        raise ValueError(
            f"q0 keys must match move group ids: {set(q0.keys())} != {set(self._robot_view.move_group_ids())}"
        )
    self._robot_view.base.pose = base_pose
    self._robot_view.set_qpos_dict(q0)
    if rel_to_base:
        pose = self._robot_view.base.pose @ pose

    move_group = self._robot_view.get_move_group(move_group_id)
    for i in range(max_iter):
        mujoco.mj_fwdPosition(self._mj_model, self._mj_data)
        mujoco.mj_sensorPos(self._mj_model, self._mj_data)

        ee_pose = relative_to_global_transform(
            move_group.leaf_frame_to_robot, self._robot_view.base.pose
        )

        err_trf = inverse_homogeneous_matrix(ee_pose) @ pose
        twist_lin, twist_ang = transform_to_twist(err_trf)

        err = np.concatenate([ee_pose[:3, :3] @ twist_lin, ee_pose[:3, :3] @ twist_ang])
        if np.linalg.norm(err) < eps:
            succ = True
            break
        elif i == max_iter - 1:
            succ = False
            break

        J: np.ndarray = self._robot_view.get_jacobian(move_group_id, unlocked_move_group_ids)
        if (JJT_det := np.linalg.det(J @ J.T)) < 1e-20:
            print(f"WARN: IK Jacobian is rank deficient! det(JJ^T)={JJT_det:.0e}")

        H = J @ J.T + damping * np.eye(J.shape[0])
        q_dot = J.T @ np.linalg.solve(H, err)
        dq = q_dot * dt

        j = 0
        for mg_id in unlocked_move_group_ids:
            mg = self._robot_view.get_move_group(mg_id)
            mg.joint_pos = mg.integrate_joint_vel(mg.joint_pos, dq[j : j + mg.vel_dim])
            j += mg.vel_dim
        self._constrain_state()

    if succ:
        return self._robot_view.get_qpos_dict()
    return None

twist_to_joint_vel

twist_to_joint_vel(move_group_id: str, twist: ndarray, unlocked_move_group_ids: list[str], q0: dict[str, ndarray], base_pose: ndarray, twist_frame: Literal['world', 'base', 'leaf'] = 'world', damping: float = 1e-12)

Convert an end-effector twist to joint velocities.

This method computes the joint velocities that would result in the specified end-effector twist. It uses damped least squares to handle singularities.

Parameters:

Name	Type	Description	Default
move_group_id	str	ID of the move group whose end-effector twist is specified	required
twist	ndarray	6D twist vector (linear and angular velocities)	required
unlocked_move_group_ids	list[str]	List of move group IDs whose joints can move	required
q0	dict[str, ndarray]	Dictionary mapping move group IDs to their current joint positions	required
base_pose	ndarray	4x4 pose matrix of the robot base.	required
twist_frame	Literal['world', 'base', 'leaf']	Frame of the twist, either in the world frame, the robot base frame, or the leaf frame of the move group.	'world'
damping	float	Damping factor for the damped least squares solution	1e-12

Returns:

Type	Description
	Array of joint velocities that would achieve the desired twist

Raises:

Type	Description
ValueError	If q0 keys don't match move group IDs

Source code in molmo_spaces/kinematics/mujoco_kinematics.py

def twist_to_joint_vel(
    self,
    move_group_id: str,
    twist: np.ndarray,
    unlocked_move_group_ids: list[str],
    q0: dict[str, np.ndarray],
    base_pose: np.ndarray,
    twist_frame: Literal["world", "base", "leaf"] = "world",
    damping: float = 1e-12,
):
    """Convert an end-effector twist to joint velocities.

    This method computes the joint velocities that would result in the specified
    end-effector twist. It uses damped least squares to handle singularities.

    Args:
        move_group_id: ID of the move group whose end-effector twist is specified
        twist: 6D twist vector (linear and angular velocities)
        unlocked_move_group_ids: List of move group IDs whose joints can move
        q0: Dictionary mapping move group IDs to their current joint positions
        base_pose: 4x4 pose matrix of the robot base.
        twist_frame: Frame of the twist, either in the world frame, the robot base frame, or the leaf frame of the move group.
        damping: Damping factor for the damped least squares solution

    Returns:
        Array of joint velocities that would achieve the desired twist

    Raises:
        ValueError: If q0 keys don't match move group IDs
    """
    if set(q0.keys()) != set(self._robot_view.move_group_ids()):
        raise ValueError(
            f"q0 keys must match move group ids: {set(q0.keys())} != {set(self._robot_view.move_group_ids())}"
        )
    self._robot_view.base.pose = base_pose
    self._robot_view.set_qpos_dict(q0)

    # convert twist to world frame
    if twist_frame == "base":
        twist[:3] = self._robot_view.base.pose[:3, :3] @ twist[:3]
        twist[3:] = self._robot_view.base.pose[:3, :3] @ twist[3:]
    elif twist_frame == "leaf":
        move_group = self._robot_view.get_move_group(move_group_id)
        twist[:3] = (
            self._robot_view.base.pose @ move_group.leaf_frame_to_robot[:3, :3] @ twist[:3]
        )
        twist[3:] = (
            self._robot_view.base.pose @ move_group.leaf_frame_to_robot[:3, :3] @ twist[3:]
        )
    else:
        assert twist_frame == "world"

    J = self._robot_view.get_jacobian(move_group_id, unlocked_move_group_ids)
    if (JJT_det := np.linalg.det(J @ J.T)) < 1e-20:
        print(f"WARN: IK Jacobian is rank deficient! det(JJ^T)={JJT_det:.0e}")
    H = J @ J.T + damping * np.eye(J.shape[0])
    q_dot = J.T @ np.linalg.solve(H, twist)
    return q_dot

rby1_kinematics

Classes:

Name	Description
RBY1Kinematics

Functions:

Name	Description
main

RBY1Kinematics

RBY1Kinematics(model: MjModel, data: MjData | None = None, namespace: str = '', holo_base: bool = False)

Bases: MlSpacesKinematics

Methods:

Name	Description
fk	Compute forward kinematics for all move groups.
ik	Solve inverse kinematics to reach a target pose.
twist_to_joint_vel	Convert an end-effector twist to joint velocities.

Source code in molmo_spaces/kinematics/rby1_kinematics.py

def __init__(
    self,
    model: MjModel,
    data: MjData | None = None,
    namespace: str = "",
    holo_base: bool = False,
) -> None:
    if data is None:
        data = MjData(model)
    robot_view = RBY1RobotView(data, namespace=namespace, holo_base=holo_base)
    super().__init__(data, robot_view)

fk

fk(move_group_qpos: dict[str, ndarray], base_pose: ndarray, rel_to_base: bool = False) -> dict[str, ndarray]

Compute forward kinematics for all move groups.

This method computes the pose of each move group's leaf frame given the joint positions. The poses can be returned either relative to the world frame or relative to the robot base.

Parameters:

Name	Type	Description	Default
move_group_qpos	dict[str, ndarray]	Dictionary mapping move group IDs to their joint positions	required
base_pose	ndarray	4x4 pose matrix of the robot base.	required
rel_to_base	bool	If True, return poses relative to robot base. If False, return poses in world frame.	False

Returns:

Type	Description
dict[str, ndarray]	Dictionary mapping move group IDs to their 4x4 pose matrices

Source code in molmo_spaces/kinematics/mujoco_kinematics.py

def fk(
    self,
    move_group_qpos: dict[str, np.ndarray],
    base_pose: np.ndarray,
    rel_to_base: bool = False,
) -> dict[str, np.ndarray]:
    """Compute forward kinematics for all move groups.

    This method computes the pose of each move group's leaf frame given the joint positions.
    The poses can be returned either relative to the world frame or relative to the robot base.

    Args:
        move_group_qpos: Dictionary mapping move group IDs to their joint positions
        base_pose: 4x4 pose matrix of the robot base.
        rel_to_base: If True, return poses relative to robot base. If False, return poses in world frame.

    Returns:
        Dictionary mapping move group IDs to their 4x4 pose matrices
    """
    self._robot_view.base.pose = base_pose
    self._robot_view.set_qpos_dict(move_group_qpos)
    mujoco.mj_kinematics(self._mj_model, self._mj_data)
    ret = {}
    for move_group_id in self._robot_view.move_group_ids():
        move_group = self._robot_view.get_move_group(move_group_id)
        if rel_to_base:
            ret[move_group_id] = move_group.leaf_frame_to_robot
        else:
            ret[move_group_id] = self._robot_view.base.pose @ move_group.leaf_frame_to_robot
    return ret

ik

ik(move_group_id: str, pose: ndarray, unlocked_move_group_ids: list[str], q0: dict[str, ndarray], base_pose: ndarray, rel_to_base: bool = False, eps: float = 0.0001, max_iter: int = 1000, damping: float = 1e-12, dt: float = 1.0)

Solve inverse kinematics to reach a target pose.

This method iteratively solves for joint positions that would place the end-effector at the target pose. It uses a damped least squares approach with velocity-based updates.

Parameters:

Name	Type	Description	Default
move_group_id	str	ID of the move group whose end-effector pose is specified	required
pose	ndarray	4x4 target pose matrix	required
unlocked_move_group_ids	list[str]	List of move group IDs whose joints can move	required
q0	dict[str, ndarray]	Dictionary mapping move group IDs to their initial joint positions	required
base_pose	ndarray	4x4 pose matrix of the robot base.	required
rel_to_base	bool	If True, target pose is relative to robot base. If False, pose is in world frame.	False
eps	float	Convergence threshold for the error norm	0.0001
max_iter	int	Maximum number of iterations	1000
damping	float	Damping factor for the damped least squares solution	1e-12
dt	float	Time step for velocity integration	1.0

Returns:

Type	Description
	Dictionary mapping move group IDs to their joint positions if successful, None if failed

Raises:

Type	Description
ValueError	If q0 keys don't match move group IDs

Source code in molmo_spaces/kinematics/mujoco_kinematics.py

def ik(
    self,
    move_group_id: str,
    pose: np.ndarray,
    unlocked_move_group_ids: list[str],
    q0: dict[str, np.ndarray],
    base_pose: np.ndarray,
    rel_to_base: bool = False,
    eps: float = 1e-4,
    max_iter: int = 1000,
    damping: float = 1e-12,
    dt: float = 1.0,
):
    """Solve inverse kinematics to reach a target pose.

    This method iteratively solves for joint positions that would place the end-effector
    at the target pose. It uses a damped least squares approach with velocity-based
    updates.

    Args:
        move_group_id: ID of the move group whose end-effector pose is specified
        pose: 4x4 target pose matrix
        unlocked_move_group_ids: List of move group IDs whose joints can move
        q0: Dictionary mapping move group IDs to their initial joint positions
        base_pose: 4x4 pose matrix of the robot base.
        rel_to_base: If True, target pose is relative to robot base. If False, pose is in world frame.
        eps: Convergence threshold for the error norm
        max_iter: Maximum number of iterations
        damping: Damping factor for the damped least squares solution
        dt: Time step for velocity integration

    Returns:
        Dictionary mapping move group IDs to their joint positions if successful, None if failed

    Raises:
        ValueError: If q0 keys don't match move group IDs
    """
    if set(q0.keys()) != set(self._robot_view.move_group_ids()):
        raise ValueError(
            f"q0 keys must match move group ids: {set(q0.keys())} != {set(self._robot_view.move_group_ids())}"
        )
    self._robot_view.base.pose = base_pose
    self._robot_view.set_qpos_dict(q0)
    if rel_to_base:
        pose = self._robot_view.base.pose @ pose

    move_group = self._robot_view.get_move_group(move_group_id)
    for i in range(max_iter):
        mujoco.mj_fwdPosition(self._mj_model, self._mj_data)
        mujoco.mj_sensorPos(self._mj_model, self._mj_data)

        ee_pose = relative_to_global_transform(
            move_group.leaf_frame_to_robot, self._robot_view.base.pose
        )

        err_trf = inverse_homogeneous_matrix(ee_pose) @ pose
        twist_lin, twist_ang = transform_to_twist(err_trf)

        err = np.concatenate([ee_pose[:3, :3] @ twist_lin, ee_pose[:3, :3] @ twist_ang])
        if np.linalg.norm(err) < eps:
            succ = True
            break
        elif i == max_iter - 1:
            succ = False
            break

        J: np.ndarray = self._robot_view.get_jacobian(move_group_id, unlocked_move_group_ids)
        if (JJT_det := np.linalg.det(J @ J.T)) < 1e-20:
            print(f"WARN: IK Jacobian is rank deficient! det(JJ^T)={JJT_det:.0e}")

        H = J @ J.T + damping * np.eye(J.shape[0])
        q_dot = J.T @ np.linalg.solve(H, err)
        dq = q_dot * dt

        j = 0
        for mg_id in unlocked_move_group_ids:
            mg = self._robot_view.get_move_group(mg_id)
            mg.joint_pos = mg.integrate_joint_vel(mg.joint_pos, dq[j : j + mg.vel_dim])
            j += mg.vel_dim
        self._constrain_state()

    if succ:
        return self._robot_view.get_qpos_dict()
    return None

twist_to_joint_vel

twist_to_joint_vel(move_group_id: str, twist: ndarray, unlocked_move_group_ids: list[str], q0: dict[str, ndarray], base_pose: ndarray, twist_frame: Literal['world', 'base', 'leaf'] = 'world', damping: float = 1e-12)

Convert an end-effector twist to joint velocities.

This method computes the joint velocities that would result in the specified end-effector twist. It uses damped least squares to handle singularities.

Parameters:

Name	Type	Description	Default
move_group_id	str	ID of the move group whose end-effector twist is specified	required
twist	ndarray	6D twist vector (linear and angular velocities)	required
unlocked_move_group_ids	list[str]	List of move group IDs whose joints can move	required
q0	dict[str, ndarray]	Dictionary mapping move group IDs to their current joint positions	required
base_pose	ndarray	4x4 pose matrix of the robot base.	required
twist_frame	Literal['world', 'base', 'leaf']	Frame of the twist, either in the world frame, the robot base frame, or the leaf frame of the move group.	'world'
damping	float	Damping factor for the damped least squares solution	1e-12

Returns:

Type	Description
	Array of joint velocities that would achieve the desired twist

Raises:

Type	Description
ValueError	If q0 keys don't match move group IDs

Source code in molmo_spaces/kinematics/mujoco_kinematics.py

def twist_to_joint_vel(
    self,
    move_group_id: str,
    twist: np.ndarray,
    unlocked_move_group_ids: list[str],
    q0: dict[str, np.ndarray],
    base_pose: np.ndarray,
    twist_frame: Literal["world", "base", "leaf"] = "world",
    damping: float = 1e-12,
):
    """Convert an end-effector twist to joint velocities.

    This method computes the joint velocities that would result in the specified
    end-effector twist. It uses damped least squares to handle singularities.

    Args:
        move_group_id: ID of the move group whose end-effector twist is specified
        twist: 6D twist vector (linear and angular velocities)
        unlocked_move_group_ids: List of move group IDs whose joints can move
        q0: Dictionary mapping move group IDs to their current joint positions
        base_pose: 4x4 pose matrix of the robot base.
        twist_frame: Frame of the twist, either in the world frame, the robot base frame, or the leaf frame of the move group.
        damping: Damping factor for the damped least squares solution

    Returns:
        Array of joint velocities that would achieve the desired twist

    Raises:
        ValueError: If q0 keys don't match move group IDs
    """
    if set(q0.keys()) != set(self._robot_view.move_group_ids()):
        raise ValueError(
            f"q0 keys must match move group ids: {set(q0.keys())} != {set(self._robot_view.move_group_ids())}"
        )
    self._robot_view.base.pose = base_pose
    self._robot_view.set_qpos_dict(q0)

    # convert twist to world frame
    if twist_frame == "base":
        twist[:3] = self._robot_view.base.pose[:3, :3] @ twist[:3]
        twist[3:] = self._robot_view.base.pose[:3, :3] @ twist[3:]
    elif twist_frame == "leaf":
        move_group = self._robot_view.get_move_group(move_group_id)
        twist[:3] = (
            self._robot_view.base.pose @ move_group.leaf_frame_to_robot[:3, :3] @ twist[:3]
        )
        twist[3:] = (
            self._robot_view.base.pose @ move_group.leaf_frame_to_robot[:3, :3] @ twist[3:]
        )
    else:
        assert twist_frame == "world"

    J = self._robot_view.get_jacobian(move_group_id, unlocked_move_group_ids)
    if (JJT_det := np.linalg.det(J @ J.T)) < 1e-20:
        print(f"WARN: IK Jacobian is rank deficient! det(JJ^T)={JJT_det:.0e}")
    H = J @ J.T + damping * np.eye(J.shape[0])
    q_dot = J.T @ np.linalg.solve(H, twist)
    return q_dot

main

main() -> None

Source code in molmo_spaces/kinematics/rby1_kinematics.py

def main() -> None:
    import mujoco
    import numpy as np
    from mujoco.viewer import launch_passive

    np.set_printoptions(linewidth=np.inf)
    import os
    import time

    cwd = os.getcwd()
    os.chdir("assets/robots/rby1")
    model = MjModel.from_xml_path("rby1.xml")
    os.chdir(cwd)

    data = MjData(model)
    mujoco.mj_forward(model, data)

    namespace = ""
    robot_view = RBY1RobotView(data, namespace=namespace)
    right_arm = robot_view.get_move_group("right_arm")
    kinematics = RBY1Kinematics(model, namespace=namespace)

    qp = np.random.uniform(*right_arm.joint_pos_limits.T)
    right_arm.joint_pos = qp
    print(qp)
    mujoco.mj_forward(model, data)

    pose0 = robot_view.base.pose @ right_arm.leaf_frame_to_robot
    pose1 = pose0.copy()
    pose1[2, 3] += 0.2

    groups = [
        "right_arm",
        "torso",
    ]

    with launch_passive(model, data) as viewer:
        viewer.sync()
        i = 0
        while viewer.is_running():
            if i % 2 == 0:
                ret = kinematics.ik(
                    "right_arm", pose1, groups, robot_view.get_qpos_dict(), robot_view.base.pose
                )
            else:
                ret = kinematics.ik(
                    "right_arm", pose0, groups, robot_view.get_qpos_dict(), robot_view.base.pose
                )
            print(i % 2, ret)
            i += 1
            if ret is not None:
                robot_view.set_qpos_dict(ret)
            mujoco.mj_forward(model, data)
            viewer.sync()
            time.sleep(3)

stretch_kinematics