Skip to content

kinematics

kinematics

Modules:

Name Description
floating_rum_kinematics
mujoco_kinematics

General-purpose forward and inverse kinematics solver for robots in MuJoCo.
test_robot_ik

Interactive script to test a robot with the parallel and non-parallel IK solvers.

floating_rum_kinematics

Classes:

Name Description
FloatingRUMKinematics

Kinematics solver that provides a specialized implementation for the Floating RUM robot.

Functions:

Name Description
main

FloatingRUMKinematics 

FloatingRUMKinematics(robot_config: BaseRobotConfig)

Bases: MlSpacesKinematics

Kinematics solver that provides a specialized implementation for the Floating RUM robot.

Create a kinematics solver for a robot.

Parameters:

Name Type Description Default
robot_config BaseRobotConfig

The robot configuration.

 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 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def __init__(self, robot_config: "BaseRobotConfig") -> None:
    """
    Create a kinematics solver for a robot.

    Args:
        robot_config: The robot configuration.
    """
    spec = mujoco.MjSpec()
    robot_config.robot_cls.add_robot_to_scene(
        robot_config,
        spec,
        prefix=robot_config.robot_namespace,
        pos=[0.0, 0.0, 0.0],
        quat=[1.0, 0.0, 0.0, 0.0],
        strip_meshes=True,
    )

    self._mj_model = spec.compile()
    self._mj_data = mujoco.MjData(self._mj_model)
    self._robot_view = robot_config.robot_view_factory(
        self._mj_data, robot_config.robot_namespace
    )
    mujoco.mj_forward(self._mj_model, self._mj_data)
_mj_data instance-attribute 
_mj_data = MjData(_mj_model)
_mj_model instance-attribute 
_mj_model = compile()
_robot_view instance-attribute 
_robot_view = robot_view_factory(_mj_data, robot_namespace)
_constrain_state 
_constrain_state() -> None

Constrain the current state to be within the joint limits.

This method clips all joint positions to their respective limits. It should be called after any operation that might move joints outside their valid ranges.

Source code in molmo_spaces/kinematics/mujoco_kinematics.py 
61
62
63
64
65
66
67
68
69
70
71
def _constrain_state(self) -> None:
    """Constrain the current state to be within the joint limits.

    This method clips all joint positions to their respective limits. It should be
    called after any operation that might move joints outside their valid ranges.
    """
    for mg_id in self._robot_view.move_group_ids():
        mg = self._robot_view.get_move_group(mg_id)
        mg.joint_pos = np.clip(
            mg.joint_pos, mg.joint_pos_limits[:, 0], mg.joint_pos_limits[:, 1]
        )
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 
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
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] | None, 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] | None

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/floating_rum_kinematics.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
def ik(
    self,
    move_group_id: str,
    pose: np.ndarray,
    unlocked_move_group_ids: list[str] | None,
    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 unlocked_move_group_ids is None:
        unlocked_move_group_ids = self._robot_view.move_group_ids()

    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 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
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:
        log.warning(
            f"[MlSpacesKinematics][{self._robot_view.name}] 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

_show_poses 

_show_poses(viewer: Handle, poses: ndarray, color=(1, 0, 0, 1)) -> None
Source code in molmo_spaces/kinematics/floating_rum_kinematics.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
def _show_poses(viewer: Handle, poses: np.ndarray, color=(1, 0, 0, 1)) -> None:
    assert poses.ndim == 3 and poses.shape[1:] == (4, 4)
    ngeom = viewer.user_scn.ngeom
    # Define relative parts of the gripper

    gripper_parts = [
        ("sphere", mujoco.mjtGeom.mjGEOM_CYLINDER, [0.003, 0.015, 0], [0.00, 0.00, 0.05]),
        (
            "cylinder_left",
            mujoco.mjtGeom.mjGEOM_CYLINDER,
            [0.003, 0.02, 0],
            [0, -0.046, 0.0835],
        ),
        (
            "cylinder_right",
            mujoco.mjtGeom.mjGEOM_CYLINDER,
            [0.003, 0.02, 0],
            [0, 0.046, 0.0835],
        ),
        ("connecting_bar", mujoco.mjtGeom.mjGEOM_BOX, [0.002, 0.044, 0.002], [0, 0, 0.065]),
    ]

    i = 0
    for T in poses:
        for _part_name, geom_type, size, offset in gripper_parts:
            T_part = pos_quat_to_pose_mat(offset, [1, 0, 0, 0])
            A = T.copy() @ T_part
            mujoco.mjv_initGeom(
                viewer.user_scn.geoms[ngeom + i],
                type=geom_type,
                size=np.array(size),
                pos=A[:3, 3],
                mat=T[:3, :3].flatten(),
                rgba=color,
            )
            i += 1
    viewer.user_scn.ngeom = ngeom + i

main 

main()
Source code in molmo_spaces/kinematics/floating_rum_kinematics.py 
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
def main():
    np.set_printoptions(linewidth=np.inf)

    robot_config = FloatingRUMRobotConfig()
    xml_path = robot_config.get_robot_xml_path()
    model = MjModel.from_xml_path(str(xml_path))

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

    ns = ""
    robot_view = FloatingRUMRobotView(data, ns)
    kinematics = FloatingRUMKinematics(robot_config)

    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)

mujoco_kinematics

General-purpose forward and inverse kinematics solver for robots in MuJoCo. This solver is not parallelizable and runs on the CPU. While fairly fast, this is not suitable for large batches.

Classes:

Name Description
MlSpacesKinematics

A general-purpose kinematics solver for robots in MuJoCo.

Attributes:

Name Type Description
log

log module-attribute 

log = getLogger(__name__)

MlSpacesKinematics 

MlSpacesKinematics(robot_config: BaseRobotConfig)

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.

Create a kinematics solver for a robot.

Parameters:

Name Type Description Default
robot_config BaseRobotConfig

The robot configuration.

 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 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def __init__(self, robot_config: "BaseRobotConfig") -> None:
    """
    Create a kinematics solver for a robot.

    Args:
        robot_config: The robot configuration.
    """
    spec = mujoco.MjSpec()
    robot_config.robot_cls.add_robot_to_scene(
        robot_config,
        spec,
        prefix=robot_config.robot_namespace,
        pos=[0.0, 0.0, 0.0],
        quat=[1.0, 0.0, 0.0, 0.0],
        strip_meshes=True,
    )

    self._mj_model = spec.compile()
    self._mj_data = mujoco.MjData(self._mj_model)
    self._robot_view = robot_config.robot_view_factory(
        self._mj_data, robot_config.robot_namespace
    )
    mujoco.mj_forward(self._mj_model, self._mj_data)
_mj_data instance-attribute 
_mj_data = MjData(_mj_model)
_mj_model instance-attribute 
_mj_model = compile()
_robot_view instance-attribute 
_robot_view = robot_view_factory(_mj_data, robot_namespace)
_constrain_state 
_constrain_state() -> None

Constrain the current state to be within the joint limits.

This method clips all joint positions to their respective limits. It should be called after any operation that might move joints outside their valid ranges.

Source code in molmo_spaces/kinematics/mujoco_kinematics.py 
61
62
63
64
65
66
67
68
69
70
71
def _constrain_state(self) -> None:
    """Constrain the current state to be within the joint limits.

    This method clips all joint positions to their respective limits. It should be
    called after any operation that might move joints outside their valid ranges.
    """
    for mg_id in self._robot_view.move_group_ids():
        mg = self._robot_view.get_move_group(mg_id)
        mg.joint_pos = np.clip(
            mg.joint_pos, mg.joint_pos_limits[:, 0], mg.joint_pos_limits[:, 1]
        )
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 
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
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] | None, 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] | None

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 
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
def ik(
    self,
    move_group_id: str,
    pose: np.ndarray,
    unlocked_move_group_ids: list[str] | None,
    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())}"
        )

    if unlocked_move_group_ids is None:
        unlocked_move_group_ids = 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:
            log.warning(
                f"[MlSpacesKinematics][{self._robot_view.name}] 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 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
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:
        log.warning(
            f"[MlSpacesKinematics][{self._robot_view.name}] 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

test_robot_ik

Interactive script to test a robot with the parallel and non-parallel IK solvers.

On mac, run with mjpython.

Functions:

Name Description
main

main 

main() -> None
Source code in molmo_spaces/kinematics/test_robot_ik.py 
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "config_class", help="The class name of the robot config to use, e.g. 'FrankaRobotConfig'"
    )
    parser.add_argument(
        "--config_module",
        default="molmo_spaces.configs.robot_configs",
        help="The module name the robot config is in, defaults to 'molmo_spaces.configs.robot_configs'",
    )
    parser.add_argument(
        "--move-group",
        help="The move group to test IK for, defaults to the first gripper move group",
    )
    parser.add_argument(
        "--unlocked-move-groups",
        nargs="+",
        help="The move groups to unlock for the IK solver, defaults to all move groups",
    )
    parser.add_argument(
        "--parallel",
        action="store_true",
        help="Use the parallel IK solver instead of the non-parallel one",
    )
    args = parser.parse_args()

    config_module = importlib.import_module(args.config_module)
    config_class = getattr(config_module, args.config_class)
    robot_config: "BaseRobotConfig" = config_class()

    if args.parallel:
        kinematics = SimpleWarpKinematics(robot_config)
    else:
        kinematics = MlSpacesKinematics(robot_config)

    spec = mujoco.MjSpec()
    robot_config.robot_cls.add_robot_to_scene(
        robot_config,
        spec,
        prefix="",
        pos=[0.0, 0.0, 0.0],
        quat=[1.0, 0.0, 0.0, 0.0],
    )
    model = spec.compile()
    data = mujoco.MjData(model)
    robot_view = robot_config.robot_view_factory(data, "")

    robot_view.set_qpos_dict(robot_config.init_qpos)
    mujoco.mj_forward(model, data)

    if args.move_group is None:
        move_group_id = robot_view.get_gripper_movegroup_ids()[0]
        move_group = robot_view.get_move_group(move_group_id)
    else:
        move_group_id = args.move_group
        move_group = robot_view.get_move_group(move_group_id)

    pose0 = move_group.leaf_frame_to_world.copy()
    pose1 = pose0.copy()
    pose0[2, 3] += 0.05  # Move up 5cm
    pose1[2, 3] -= 0.05  # Move down 5cm

    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(
                move_group_id,
                target_pose,
                args.unlocked_move_groups,
                robot_config.init_qpos,
                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)

parallel kinematics

parallel

Modules:

Name Description
dummy_parallel_kinematics
parallel_kinematics
warp_kinematics

Provides a general-purpose, robot-agnostic, vectorized (and optionally GPU accelerated) kinematics solver.

dummy_parallel_kinematics

Classes:

Name Description
DummyParallelKinematics

Dummy parallel kinematics wrapper that sequentially calls :class:MlSpacesKinematics internally.

Attributes:

Name Type Description
logger

logger module-attribute 

logger = getLogger(__name__)

DummyParallelKinematics 

DummyParallelKinematics(robot_config: BaseRobotConfig, kinematics: MlSpacesKinematics)

Bases: ParallelKinematics

Dummy parallel kinematics wrapper that sequentially calls :class:MlSpacesKinematics internally.

Methods:

Name Description
fk

Compute forward kinematics for all move groups.
ik

Finds joint positions that would place the end-effector at the target pose.
warmup_ik

Incur startup costs up-front to speed up subsequent calls to IK.
Source code in molmo_spaces/kinematics/parallel/dummy_parallel_kinematics.py 
20
21
22
23
24
25
26
def __init__(
    self,
    robot_config: "BaseRobotConfig",
    kinematics: MlSpacesKinematics,
):
    super().__init__(robot_config)
    self._kinematics = kinematics
fk 
fk(qpos_dicts: list[dict[str, ndarray]] | dict[str, ndarray], base_poses: ndarray, rel_to_base: bool = False) -> list[dict[str, ndarray]] | dict[str, ndarray]

Compute forward kinematics for all move groups.

Parameters:

Name Type Description Default
qpos_dicts list[dict[str, ndarray]] | dict[str, ndarray]

The joint positions.

 required
base_pose

The base pose(s) of the robots. Shape: (batch_size, 4, 4) or (4, 4)

 required
rel_to_base bool

Whether the returned pose(s) should be relative to the base frame.

 False

Returns:

Type Description
list[dict[str, ndarray]] | dict[str, ndarray]

A list of qpos dictionaries for each robot in the batch, or a single qpos dictionary if unbatched.
Source code in molmo_spaces/kinematics/parallel/dummy_parallel_kinematics.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def fk(
    self,
    qpos_dicts: list[dict[str, np.ndarray]] | dict[str, np.ndarray],
    base_poses: np.ndarray,
    rel_to_base: bool = False,
) -> list[dict[str, np.ndarray]] | dict[str, np.ndarray]:
    is_batch, batch_size, qpos_dicts, base_poses = self._batchify(qpos_dicts, base_poses)

    ret = []
    for qpos_dict, base_pose in zip(qpos_dicts, base_poses):
        # Ensure base_pose is writable (create copy if read-only)
        base_pose = np.array(base_pose, copy=True)
        ret.append(self._kinematics.fk(qpos_dict, base_pose, rel_to_base))
    return ret if is_batch else ret[0]
ik 
ik(move_group_id: str, poses: ndarray, unlocked_move_group_ids: list[str] | None, q0_dicts: list[dict[str, ndarray]] | dict[str, ndarray], base_poses: ndarray, rel_to_base: bool = False, dt: float = 1.0, max_iter: int = 100, success_eps: float = 0.0001, **kwargs)

Finds joint positions that would place the end-effector at the target pose.

Parameters:

Name Type Description Default
pose

The target pose(s) to reach. Shape: (batch_size, 4, 4) or (4, 4)

 required
q0_dicts list[dict[str, ndarray]] | dict[str, ndarray]

The initial joint positions.

 required
base_pose

The base pose(s) of the robots. Shape: (batch_size, 4, 4) or (4, 4)

 required
rel_to_base bool

Whether the pose(s) are relative to the base frame.

 False
**kwargs Any

Additional keyword arguments for the IK solver, defined by the concrete implementation.

 {}

Returns:

Type Description
list[dict[str, ndarray] | None] | dict[str, ndarray] | None

A list of qpos dictionaries for each robot in the batch, or a single qpos dictionary if unbatched. If the solver fails to converge for a given robot, the corresponding qpos dictionary is None.
Source code in molmo_spaces/kinematics/parallel/dummy_parallel_kinematics.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
def ik(
    self,
    move_group_id: str,
    poses: np.ndarray,
    unlocked_move_group_ids: list[str] | None,
    q0_dicts: list[dict[str, np.ndarray]] | dict[str, np.ndarray],
    base_poses: np.ndarray,
    rel_to_base: bool = False,
    dt: float = 1.0,
    max_iter: int = 100,
    success_eps: float = 1e-4,
    **kwargs,
):
    is_batch, batch_size, q0_dicts, base_poses, poses = self._batchify(
        q0_dicts, base_poses, poses
    )

    ret = []
    for q0_dict, base_pose, pose in zip(q0_dicts, base_poses, poses):
        # Ensure arrays are writable (create copies if read-only)
        # This fixes "ValueError: buffer source array is read-only" when setting base.pose
        # Use copy=True to ensure we have a writable array
        base_pose = np.array(base_pose, copy=True)
        pose = np.array(pose, copy=True)

        ret.append(
            self._kinematics.ik(
                move_group_id,
                pose,
                unlocked_move_group_ids,
                q0_dict,
                base_pose,
                rel_to_base,
                eps=success_eps,
                max_iter=max_iter,
                dt=dt,
            )
        )
    success = np.array([jp_dict is not None for jp_dict in ret])
    if not np.all(success):
        logger.debug(f"[DummyParallelKinematics] IK failed for indices {np.where(~success)[0]}")
    return ret if is_batch else ret[0]
warmup_ik 
warmup_ik(batch_size: int)

Incur startup costs up-front to speed up subsequent calls to IK. The warmed up batch size may need to match that of the subsequent calls to IK.

Parameters:

Name Type Description Default
batch_size int

The batch size to warmup.

 required
Source code in molmo_spaces/kinematics/parallel/dummy_parallel_kinematics.py 
28
29
def warmup_ik(self, batch_size: int):
    pass

parallel_kinematics

Classes:

Name Description
ParallelKinematics

ParallelKinematics 

ParallelKinematics(robot_config: BaseRobotConfig)

Bases: ABC

Methods:

Name Description
fk

Compute forward kinematics for all move groups.
ik

Finds joint positions that would place the end-effector at the target pose.
warmup_ik

Incur startup costs up-front to speed up subsequent calls to IK.
Source code in molmo_spaces/kinematics/parallel/parallel_kinematics.py 
11
12
def __init__(self, robot_config: "BaseRobotConfig"):
    self._robot_config = robot_config
fk abstractmethod 
fk(qpos_dicts: list[dict[str, ndarray]] | dict[str, ndarray], base_poses: ndarray, rel_to_base: bool = False) -> list[dict[str, ndarray]] | dict[str, ndarray]

Compute forward kinematics for all move groups.

Parameters:

Name Type Description Default
qpos_dicts list[dict[str, ndarray]] | dict[str, ndarray]

The joint positions.

 required
base_pose

The base pose(s) of the robots. Shape: (batch_size, 4, 4) or (4, 4)

 required
rel_to_base bool

Whether the returned pose(s) should be relative to the base frame.

 False

Returns:

Type Description
list[dict[str, ndarray]] | dict[str, ndarray]

A list of qpos dictionaries for each robot in the batch, or a single qpos dictionary if unbatched.
Source code in molmo_spaces/kinematics/parallel/parallel_kinematics.py 
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
@abstractmethod
def fk(
    self,
    qpos_dicts: list[dict[str, np.ndarray]] | dict[str, np.ndarray],
    base_poses: np.ndarray,
    rel_to_base: bool = False,
) -> list[dict[str, np.ndarray]] | dict[str, np.ndarray]:
    """
    Compute forward kinematics for all move groups.

    Args:
        qpos_dicts: The joint positions.
        base_pose: The base pose(s) of the robots. Shape: (batch_size, 4, 4) or (4, 4)
        rel_to_base: Whether the returned pose(s) should be relative to the base frame.

    Returns:
        A list of qpos dictionaries for each robot in the batch, or a single qpos dictionary if unbatched.
    """
    raise NotImplementedError
ik abstractmethod 
ik(move_group_id: str, poses: ndarray, unlocked_move_group_ids: list[str] | None, q0_dicts: list[dict[str, ndarray]] | dict[str, ndarray], base_poses: ndarray, rel_to_base: bool = False, **kwargs: Any) -> list[dict[str, ndarray] | None] | dict[str, ndarray] | None

Finds joint positions that would place the end-effector at the target pose.

Parameters:

Name Type Description Default
pose

The target pose(s) to reach. Shape: (batch_size, 4, 4) or (4, 4)

 required
q0_dicts list[dict[str, ndarray]] | dict[str, ndarray]

The initial joint positions.

 required
base_pose

The base pose(s) of the robots. Shape: (batch_size, 4, 4) or (4, 4)

 required
rel_to_base bool

Whether the pose(s) are relative to the base frame.

 False
**kwargs Any

Additional keyword arguments for the IK solver, defined by the concrete implementation.

 {}

Returns:

Type Description
list[dict[str, ndarray] | None] | dict[str, ndarray] | None

A list of qpos dictionaries for each robot in the batch, or a single qpos dictionary if unbatched. If the solver fails to converge for a given robot, the corresponding qpos dictionary is None.
Source code in molmo_spaces/kinematics/parallel/parallel_kinematics.py 
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
@abstractmethod
def ik(
    self,
    move_group_id: str,
    poses: np.ndarray,
    unlocked_move_group_ids: list[str] | None,
    q0_dicts: list[dict[str, np.ndarray]] | dict[str, np.ndarray],
    base_poses: np.ndarray,
    rel_to_base: bool = False,
    **kwargs: Any,
) -> list[dict[str, np.ndarray] | None] | dict[str, np.ndarray] | None:
    """
    Finds joint positions that would place the end-effector at the target pose.

    Args:
        pose: The target pose(s) to reach. Shape: (batch_size, 4, 4) or (4, 4)
        q0_dicts: The initial joint positions.
        base_pose: The base pose(s) of the robots. Shape: (batch_size, 4, 4) or (4, 4)
        rel_to_base: Whether the pose(s) are relative to the base frame.
        **kwargs: Additional keyword arguments for the IK solver, defined by the concrete implementation.

    Returns:
        A list of qpos dictionaries for each robot in the batch, or a single qpos dictionary if unbatched.
            If the solver fails to converge for a given robot, the corresponding qpos dictionary is None.
    """
    raise NotImplementedError
warmup_ik abstractmethod 
warmup_ik(batch_size: int)

Incur startup costs up-front to speed up subsequent calls to IK. The warmed up batch size may need to match that of the subsequent calls to IK.

Parameters:

Name Type Description Default
batch_size int

The batch size to warmup.

 required
Source code in molmo_spaces/kinematics/parallel/parallel_kinematics.py 
14
15
16
17
18
19
20
21
22
23
@abstractmethod
def warmup_ik(self, batch_size: int):
    """
    Incur startup costs up-front to speed up subsequent calls to IK.
    The warmed up batch size may need to match that of the subsequent calls to IK.

    Args:
        batch_size: The batch size to warmup.
    """
    pass

warp_kinematics

Provides a general-purpose, robot-agnostic, vectorized (and optionally GPU accelerated) kinematics solver.

Classes:

Name Description
IKArgs

Preallocated buffers for the IK solver arguments
IKBuffers

Preallocated buffers for the IK solver state
SimpleWarpKinematics

A warp-based general-purpose parallel inverse kinematics solver for robots.
SolverData

Functions:

Name Description
cholesky_solve6

Solve Hx=b via Cholesky decomposition for 6x6 symmetric positive-definite matrix
get_err

Calculate the body-frame position and rotation error between the current and target poses
get_jac_frame_info

Get the position and body ID of the frame
lm_step

Single step of the Levenberg-Marquardt solver.
mask_jacobian

Apply a column-wise mask to a (3, nv) Jacobian matrix

Attributes:

Name Type Description
logger
mat66f
vec6f

logger module-attribute 

logger = getLogger(__name__)

mat66f module-attribute 

mat66f = matrix(shape=(6, 6), dtype=float32)

vec6f module-attribute 

vec6f = vector(length=6, dtype=float32)

IKArgs dataclass 

IKArgs(poses: array(dtype=mat44f), leaf_frame_id: array(dtype=int), leaf_frame_type: array(dtype=int), damping: array(dtype=float32), dt: array(dtype=float32), jacobian_mask: array2d(dtype=int))

Preallocated buffers for the IK solver arguments

Attributes:

Name Type Description
damping array(dtype=float32)
dt array(dtype=float32)
jacobian_mask array2d(dtype=int)
leaf_frame_id array(dtype=int)
leaf_frame_type array(dtype=int)
poses array(dtype=mat44f)
damping instance-attribute 
damping: array(dtype=float32)
dt instance-attribute 
dt: array(dtype=float32)
jacobian_mask instance-attribute 
jacobian_mask: array2d(dtype=int)
leaf_frame_id instance-attribute 
leaf_frame_id: array(dtype=int)
leaf_frame_type instance-attribute 
leaf_frame_type: array(dtype=int)
poses instance-attribute 
poses: array(dtype=mat44f)

IKBuffers dataclass 

IKBuffers(pos_err: array(dtype=vec3f), rot_err: array(dtype=vec3f), jacp: array3d(dtype=float32), jacr: array3d(dtype=float32), q_dot: array2d(dtype=float32), frame_pos: array(dtype=vec3f), frame_mat: array(dtype=mat33f), frame_bodyid: array(dtype=int))

Preallocated buffers for the IK solver state

Attributes:

Name Type Description
frame_bodyid array(dtype=int)
frame_mat array(dtype=mat33f)
frame_pos array(dtype=vec3f)
jacp array3d(dtype=float32)
jacr array3d(dtype=float32)
pos_err array(dtype=vec3f)
q_dot array2d(dtype=float32)
rot_err array(dtype=vec3f)
frame_bodyid instance-attribute 
frame_bodyid: array(dtype=int)
frame_mat instance-attribute 
frame_mat: array(dtype=mat33f)
frame_pos instance-attribute 
frame_pos: array(dtype=vec3f)
jacp instance-attribute 
jacp: array3d(dtype=float32)
jacr instance-attribute 
jacr: array3d(dtype=float32)
pos_err instance-attribute 
pos_err: array(dtype=vec3f)
q_dot instance-attribute 
q_dot: array2d(dtype=float32)
rot_err instance-attribute 
rot_err: array(dtype=vec3f)

SimpleWarpKinematics 

SimpleWarpKinematics(robot_config: BaseRobotConfig, device: str = 'cpu')

Bases: ParallelKinematics

A warp-based general-purpose parallel inverse kinematics solver for robots. This solver only supports optimizing SimplyActuatedMoveGroups to reach a target pose for a given MJCFFrameMixin move group. Most robots satisfy this assumption, but more complicated robots may need custom kinematics implementations.

Parameters:

Name Type Description Default
robot_config BaseRobotConfig

The robot configuration.

 required
device str

The warp device to use for the solver.

 'cpu'

Methods:

Name Description
fk

Compute forward kinematics for all simple move groups.
ik

Solve inverse kinematics to reach a target pose.
warmup_ik

Incur startup costs up-front to speed up subsequent calls to IK.
Source code in molmo_spaces/kinematics/parallel/warp_kinematics.py 
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
def __init__(self, robot_config: "BaseRobotConfig", device: str = "cpu"):
    """
    Args:
        robot_config: The robot configuration.
        device: The warp device to use for the solver.
    """
    super().__init__(robot_config)
    self._device = device

    spec = MjSpec()
    robot_config.robot_cls.add_robot_to_scene(
        robot_config,
        spec,
        prefix=robot_config.robot_namespace,
        pos=[0.0, 0.0, 0.0],
        quat=[1.0, 0.0, 0.0, 0.0],
        strip_meshes=True,
    )
    self._mj_model: MjModel = spec.compile()
    with wp.ScopedDevice(self._device):
        self._mjw_model = mjw.put_model(self._mj_model)

    mj_data = MjData(self._mj_model)
    self._robot_view = robot_config.robot_view_factory(mj_data, robot_config.robot_namespace)

    self._actuated_move_groups: OrderedDict[str, SimplyActuatedMoveGroup] = OrderedDict()
    self._frame_move_groups: dict[str, MJCFFrameMixin] = {}
    for mg_id in self._robot_view.move_group_ids():
        mg = self._robot_view.get_move_group(mg_id)
        assert (
            mg.n_joints == 0
            or isinstance(mg, SimplyActuatedMoveGroup)
            or isinstance(mg, GripperGroup)
        )
        if isinstance(mg, SimplyActuatedMoveGroup):
            self._actuated_move_groups[mg_id] = mg
        if isinstance(mg, MJCFFrameMixin):
            self._frame_move_groups[mg_id] = mg

    if self._mj_model.nq != self._mj_model.nv:
        raise ValueError(
            "Number of position variables (nq) must equal number of velocity variables (nv) for warp-based IK solver"
        )
fk 
fk(qpos_dicts: list[dict[str, ndarray]] | dict[str, ndarray], base_poses: ndarray, rel_to_base: bool = False) -> list[dict[str, ndarray]] | dict[str, ndarray]

Compute forward kinematics for all simple move groups. Non MJCFFrameMixin move groups (e.g. bases) are not included in the output.

Parameters:

Name Type Description Default
qpos_dicts list[dict[str, ndarray]] | dict[str, ndarray]

The joint positions.

 required
base_poses ndarray

The base pose(s) of the robots. Shape: (batch_size, 4, 4) or (4, 4)

 required
rel_to_base bool

Whether the returned pose(s) should be relative to the base frame.

 False

Returns:

Type Description
list[dict[str, ndarray]] | dict[str, ndarray]

A list of qpos dictionaries for each robot in the batch, or a single qpos dictionary if unbatched.
Source code in molmo_spaces/kinematics/parallel/warp_kinematics.py 
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
def fk(
    self,
    qpos_dicts: list[dict[str, np.ndarray]] | dict[str, np.ndarray],
    base_poses: np.ndarray,
    rel_to_base: bool = False,
) -> list[dict[str, np.ndarray]] | dict[str, np.ndarray]:
    """
    Compute forward kinematics for all simple move groups.
    Non MJCFFrameMixin move groups (e.g. bases) are not included in the output.

    Args:
        qpos_dicts: The joint positions.
        base_poses: The base pose(s) of the robots. Shape: (batch_size, 4, 4) or (4, 4)
        rel_to_base: Whether the returned pose(s) should be relative to the base frame.

    Returns:
        A list of qpos dictionaries for each robot in the batch, or a single qpos dictionary if unbatched.
    """
    is_batch, batch_size, qpos_dicts, base_poses = self._batchify(qpos_dicts, base_poses)
    solver_data = self._get_data(batch_size)
    data = solver_data.data

    qpos_arr = self._dicts_to_qpos_arr(qpos_dicts)
    with wp.ScopedDevice(self._device):
        wp.copy(data.qpos, wp.from_numpy(qpos_arr))
        mjw.fwd_position(self._mjw_model, data)

    dol = {}
    for mg_id, mg in self._frame_move_groups.items():
        trf = np.repeat(np.expand_dims(np.eye(4), axis=0), batch_size, axis=0)
        xpos, xmat = (
            (data.xpos, data.xmat)
            if mg.leaf_frame_type == "body"
            else (data.site_xpos, data.site_xmat)
        )
        trf[:, :3, 3] = xpos[:, mg.leaf_frame_id].numpy()
        trf[:, :3, :3] = xmat[:, mg.leaf_frame_id].numpy()
        if isinstance(self._robot_view.base, SimplyActuatedMoveGroup):
            # if the base is actuated, trf is in world frame so we need to convert to base frame if necessary
            if rel_to_base:
                trf = np.linalg.solve(base_poses, trf)
        elif not rel_to_base:
            # if the base is unactuated, trf is in the base frame so we need to convert to world frame if necessary
            trf = base_poses @ trf
        dol[mg_id] = trf

    ret = []
    for i in range(batch_size):
        d = {}
        for mg_id, trf in dol.items():
            d[mg_id] = trf[i]
        ret.append(d)
    return ret if is_batch else ret[0]
ik 
ik(move_group_id: str, poses: ndarray, unlocked_move_group_ids: list[str] | None, q0_dicts: list[dict[str, ndarray]] | dict[str, ndarray], base_poses: ndarray, rel_to_base: bool = False, converge_eps: float = 0.001, success_eps: float = 0.0005, max_iter: int = 50, damping: float = 1e-12, dt: float = 1.0)

Solve inverse kinematics to reach a target pose.

Parameters:

Name Type Description Default
move_group_id str

The ID of the move group to solve for.

 required
poses ndarray

The target poses. Shape: (batch_size, 4, 4) or (4, 4)

 required
unlocked_move_group_ids list[str] | None

The IDs of the move groups that are not locked. If None, all move groups are unlocked.

 required
q0_dicts list[dict[str, ndarray]] | dict[str, ndarray]

The initial joint positions.

 required
base_poses ndarray

The base poses. Shape: (batch_size, 4, 4) or (4, 4)

 required
rel_to_base bool

Whether the target poses are relative to the base frame.

 False
converge_eps float

The convergence threshold in joint space.

 0.001
success_eps float

The success threshold in twist space.

 0.0005
max_iter int

The maximum number of iterations.

 50
damping float

The damping factor for the Levenberg-Marquardt solver.

 1e-12
dt float

The time step for velocity integration.

 1.0

Returns:

Type Description

A list of qpos dictionaries for each robot in the batch, or a single qpos dictionary if unbatched. If the solver fails to converge for a given robot, the corresponding qpos dictionary is None.
Source code in molmo_spaces/kinematics/parallel/warp_kinematics.py 
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
def ik(
    self,
    move_group_id: str,
    poses: np.ndarray,
    unlocked_move_group_ids: list[str] | None,
    q0_dicts: list[dict[str, np.ndarray]] | dict[str, np.ndarray],
    base_poses: np.ndarray,
    rel_to_base: bool = False,
    converge_eps: float = 1e-3,
    success_eps: float = 5e-4,
    max_iter: int = 50,
    damping: float = 1e-12,
    dt: float = 1.0,
):
    """
    Solve inverse kinematics to reach a target pose.

    Args:
        move_group_id: The ID of the move group to solve for.
        poses: The target poses. Shape: (batch_size, 4, 4) or (4, 4)
        unlocked_move_group_ids: The IDs of the move groups that are not locked. If None, all move groups are unlocked.
        q0_dicts: The initial joint positions.
        base_poses: The base poses. Shape: (batch_size, 4, 4) or (4, 4)
        rel_to_base: Whether the target poses are relative to the base frame.
        converge_eps: The convergence threshold in joint space.
        success_eps: The success threshold in twist space.
        max_iter: The maximum number of iterations.
        damping: The damping factor for the Levenberg-Marquardt solver.
        dt: The time step for velocity integration.

    Returns:
        A list of qpos dictionaries for each robot in the batch, or a single qpos dictionary if unbatched.
            If the solver fails to converge for a given robot, the corresponding qpos dictionary is None.
    """
    if move_group_id not in self._frame_move_groups:
        raise ValueError(f"Move group {move_group_id} is not a MJCFFrameMixin")

    if unlocked_move_group_ids is None:
        unlocked_move_group_ids = list(self._actuated_move_groups.keys())
    else:
        for mg_id in unlocked_move_group_ids:
            if mg_id not in self._actuated_move_groups:
                raise ValueError(f"Move group {mg_id} is not a simply actuated move group!")

    is_batch, batch_size, q0_dicts, base_poses, poses = self._batchify(
        q0_dicts, base_poses, poses
    )
    solver_data = self._get_data(batch_size)
    data = solver_data.data
    ik_args = solver_data.ik_args

    with wp.ScopedDevice(self._device):
        ee_mg = self._frame_move_groups[move_group_id]
        assert isinstance(ee_mg, MJCFFrameMixin)
        leaf_frame_id = ee_mg.leaf_frame_id
        leaf_frame_type = (
            mujoco.mjtObj.mjOBJ_BODY
            if ee_mg.leaf_frame_type == "body"
            else mujoco.mjtObj.mjOBJ_SITE
        )

        if isinstance(self._robot_view.base, SimplyActuatedMoveGroup):
            # if the base is actuated, the solving happens in world frame so ensure targets are in world frame
            if rel_to_base:
                poses = base_poses @ poses
        elif not rel_to_base:
            # if the base is unactuated, the solving happens in base frame so ensure targets are in base frame
            poses = np.linalg.solve(base_poses, poses)

        wp.copy(ik_args.poses, wp.from_numpy(poses.astype(np.float32), dtype=wp.mat44f))
        ik_args.leaf_frame_id.fill_(leaf_frame_id)
        ik_args.leaf_frame_type.fill_(leaf_frame_type.value)
        ik_args.damping.fill_(damping)
        ik_args.dt.fill_(dt)
        wp.copy(
            ik_args.jacobian_mask,
            wp.from_numpy(self._create_jacobian_mask(batch_size, unlocked_move_group_ids)),
        )

        q0_arr = self._dicts_to_qpos_arr(q0_dicts)
        wp.copy(data.qpos, wp.from_numpy(q0_arr))

        for i in range(max_iter):
            if self._device.startswith("cuda"):
                if solver_data.ik_capture is None:
                    with wp.ScopedCapture(device=self._device) as capture:
                        self._ik_solve_step(solver_data, batch_size)
                    solver_data.ik_capture = capture
                wp.capture_launch(solver_data.ik_capture.graph)
            else:
                self._ik_solve_step(solver_data, batch_size)
            wp.synchronize()

            q_dot = solver_data.ik_buffers.q_dot.numpy()
            if np.all(np.linalg.norm(q_dot, axis=-1) < converge_eps):
                logger.debug(
                    f"[SimpleWarpKinematics] Batch of size {batch_size} converged in {i} iterations"
                )
                break
        else:
            logger.debug(
                f"[SimpleWarpKinematics] Batch of size {batch_size} failed to converge in {max_iter} iterations"
            )

        mjw.kinematics(self._mjw_model, data)
        err_norm = self._get_err_norm(solver_data, batch_size)
        success = np.array(err_norm < success_eps)

    ret_q_arr = data.qpos.numpy()
    ret_q_dicts = self._qpos_arr_to_dicts(ret_q_arr)
    for q0_dict, ret_q_dict in zip(q0_dicts, ret_q_dicts):
        for k, v in q0_dict.items():
            if k not in ret_q_dict:
                ret_q_dict[k] = v

    ret = [d if s else None for d, s in zip(ret_q_dicts, success)]
    return ret if is_batch else ret[0]
warmup_ik 
warmup_ik(batch_size: int)

Incur startup costs up-front to speed up subsequent calls to IK. The warmed up batch size may need to match that of the subsequent calls to IK.

Parameters:

Name Type Description Default
batch_size int

The batch size to warmup.

 required
Source code in molmo_spaces/kinematics/parallel/warp_kinematics.py 
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
def warmup_ik(self, batch_size: int):
    self._get_data(batch_size)

    mj_data = MjData(self._mj_model)
    robot_view = self._robot_config.robot_view_factory(
        mj_data, self._robot_config.robot_namespace
    )
    for mg_id, qpos in self._robot_config.init_qpos.items():
        robot_view.get_move_group(mg_id).joint_pos = qpos
    mujoco.mj_forward(self._mj_model, mj_data)

    mg_id = next(iter(self._frame_move_groups.keys()))
    pose = np.broadcast_to(
        robot_view.get_move_group(mg_id).leaf_frame_to_robot[None], (batch_size, 4, 4)
    )

    self.ik(
        mg_id,
        pose,
        None,
        robot_view.get_qpos_dict(),
        np.eye(4),
        rel_to_base=True,
        max_iter=1,
    )

SolverData dataclass 

SolverData(data: Data, ik_buffers: IKBuffers, ik_args: IKArgs, ik_capture: ScopedCapture | None = None)

Attributes:

Name Type Description
data Data
ik_args IKArgs
ik_buffers IKBuffers
ik_capture ScopedCapture | None
data instance-attribute 
data: Data
ik_args instance-attribute 
ik_args: IKArgs
ik_buffers instance-attribute 
ik_buffers: IKBuffers
ik_capture class-attribute instance-attribute 
ik_capture: ScopedCapture | None = None

cholesky_solve6 

cholesky_solve6(H: mat66f, b: vec6f) -> vec6f

Solve Hx=b via Cholesky decomposition for 6x6 symmetric positive-definite matrix

Source code in molmo_spaces/kinematics/parallel/warp_kinematics.py 
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
@wp.func
def cholesky_solve6(H: mat66f, b: vec6f) -> vec6f:
    """Solve Hx=b via Cholesky decomposition for 6x6 symmetric positive-definite matrix"""
    # Cholesky decomposition: H = L @ L^T
    L = mat66f()
    for i in range(6):
        for j in range(i + 1):
            s = float(0.0)
            for k in range(j):
                s += L[i, k] * L[j, k]
            if i == j:
                L[i, j] = wp.sqrt(H[i, i] - s)
            else:
                L[i, j] = (H[i, j] - s) / L[j, j]

    # Forward substitution: L @ y = b
    y = vec6f()
    for i in range(6):
        s = float(0.0)
        for k in range(i):
            s += L[i, k] * y[k]
        y[i] = (b[i] - s) / L[i, i]

    # Backward substitution: L^T @ x = y
    x = vec6f()
    for i in range(5, -1, -1):
        s = float(0.0)
        for k in range(i + 1, 6):
            s += L[k, i] * x[k]
        x[i] = (y[i] - s) / L[i, i]

    return x

get_err 

get_err(xpos: array(dtype=vec3f), xmat: array(dtype=mat33f), poses: array(dtype=mat44f), pos_err: array(dtype=vec3f), rot_err: array(dtype=vec3f))

Calculate the body-frame position and rotation error between the current and target poses

Source code in molmo_spaces/kinematics/parallel/warp_kinematics.py 
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
@wp.kernel
def get_err(
    # inputs
    xpos: wp.array(dtype=wp.vec3f),
    xmat: wp.array(dtype=wp.mat33f),
    poses: wp.array(dtype=wp.mat44f),
    # outputs
    pos_err: wp.array(dtype=wp.vec3f),
    rot_err: wp.array(dtype=wp.vec3f),
):
    """Calculate the body-frame position and rotation error between the current and target poses"""
    i = wp.tid()
    frame_pos = xpos[i]
    frame_rotmat = xmat[i]
    target_pose = poses[i]

    # fmt: off
    site_pose = wp.mat44(
        frame_rotmat[0, 0], frame_rotmat[0, 1], frame_rotmat[0, 2], frame_pos[0],
        frame_rotmat[1, 0], frame_rotmat[1, 1], frame_rotmat[1, 2], frame_pos[1],
        frame_rotmat[2, 0], frame_rotmat[2, 1], frame_rotmat[2, 2], frame_pos[2],
        0.0,       0.0,       0.0,       1.0,
    )
    # fmt: on
    err_trf = wp.inverse(site_pose) @ target_pose

    # fmt: off
    rotmat = wp.mat33(
        err_trf[0, 0], err_trf[0, 1], err_trf[0, 2],
        err_trf[1, 0], err_trf[1, 1], err_trf[1, 2],
        err_trf[2, 0], err_trf[2, 1], err_trf[2, 2],
    )
    # fmt: on
    q = wp.quat_from_matrix(rotmat)
    axis, angle = wp.quat_to_axis_angle(q)
    t = wp.vec3f(err_trf[0, 3], err_trf[1, 3], err_trf[2, 3])

    if wp.abs(angle) < 1e-6:
        pos_err[i] = frame_rotmat @ t
        rot_err[i] = wp.vec3f()
    else:
        w = axis * angle
        V = (
            wp.identity(3, dtype=wp.float32)
            + (1.0 - wp.cos(angle)) / angle**2.0 * wp.skew(w)
            + (angle - wp.sin(angle)) / angle**3.0 * wp.skew(w) @ wp.skew(w)
        )
        t = wp.inverse(V) @ t
        pos_err[i] = frame_rotmat @ t
        rot_err[i] = frame_rotmat @ w

get_jac_frame_info 

get_jac_frame_info(frame_id_: array(dtype=int), frame_type_: array(dtype=int), xpos: array2d(dtype=vec3f), xmat: array2d(dtype=mat33f), site_xpos: array2d(dtype=vec3f), site_xmat: array2d(dtype=mat33f), site_bodyid: array(dtype=int), frame_pos: array(dtype=vec3f), frame_mat: array(dtype=mat33f), frame_bodyid: array(dtype=int))

Get the position and body ID of the frame

Source code in molmo_spaces/kinematics/parallel/warp_kinematics.py 
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
@wp.kernel
def get_jac_frame_info(
    # inputs
    frame_id_: wp.array(dtype=int),  # single-element array
    frame_type_: wp.array(dtype=int),  # single-element array
    xpos: wp.array2d(dtype=wp.vec3f),
    xmat: wp.array2d(dtype=wp.mat33f),
    site_xpos: wp.array2d(dtype=wp.vec3f),
    site_xmat: wp.array2d(dtype=wp.mat33f),
    site_bodyid: wp.array(dtype=int),
    # outputs
    frame_pos: wp.array(dtype=wp.vec3f),
    frame_mat: wp.array(dtype=wp.mat33f),
    frame_bodyid: wp.array(dtype=int),
):
    """Get the position and body ID of the frame"""
    i = wp.tid()
    frame_id = frame_id_[0]
    frame_type = frame_type_[0]
    if frame_type == wp.static(mujoco.mjtObj.mjOBJ_SITE.value):
        frame_pos[i] = site_xpos[i, frame_id]
        frame_mat[i] = site_xmat[i, frame_id]
        frame_bodyid[i] = site_bodyid[frame_id]
    elif frame_type == wp.static(mujoco.mjtObj.mjOBJ_BODY.value):
        frame_pos[i] = xpos[i, frame_id]
        frame_mat[i] = xmat[i, frame_id]
        frame_bodyid[i] = frame_id

lm_step 

lm_step(jacp: array3d(dtype=float32), jacr: array3d(dtype=float32), pos_err: array(dtype=vec3f), rot_err: array(dtype=vec3f), damping: array(dtype=float32), dt: array(dtype=float32), nv: int, qpos: array2d(dtype=float32), q_dot: array2d(dtype=float32))

Single step of the Levenberg-Marquardt solver.

Source code in molmo_spaces/kinematics/parallel/warp_kinematics.py 
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
@wp.kernel
def lm_step(
    # inputs
    jacp: wp.array3d(dtype=wp.float32),
    jacr: wp.array3d(dtype=wp.float32),
    pos_err: wp.array(dtype=wp.vec3f),
    rot_err: wp.array(dtype=wp.vec3f),
    damping: wp.array(dtype=wp.float32),  # single-element array
    dt: wp.array(dtype=wp.float32),  # single-element array
    nv: int,
    # outputs
    qpos: wp.array2d(dtype=wp.float32),
    q_dot: wp.array2d(dtype=wp.float32),
):
    """Single step of the Levenberg-Marquardt solver."""
    i = wp.tid()

    err = vec6f(
        pos_err[i][0],
        pos_err[i][1],
        pos_err[i][2],
        rot_err[i][0],
        rot_err[i][1],
        rot_err[i][2],
    )

    # H = J @ J^T + damping * I, where J = [jacp; jacr] is (6, nv)
    H = mat66f()
    for a in range(6):
        for b in range(6):
            val = float(0.0)
            for k in range(nv):
                Ja = jacp[i, a, k] if a < 3 else jacr[i, a - 3, k]
                Jb = jacp[i, b, k] if b < 3 else jacr[i, b - 3, k]
                val += Ja * Jb
            if a == b:
                val += damping[0]
            H[a, b] = val

    # x = H^{-1} @ err
    x = cholesky_solve6(H, err)

    # q_dot = J^T @ x, dq = q_dot * dt
    for k in range(nv):
        val = float(0.0)
        for a in range(3):
            val += jacp[i, a, k] * x[a]
            val += jacr[i, a, k] * x[a + 3]
        q_dot[i, k] = val
        qpos[i, k] += val * dt[0]

mask_jacobian 

mask_jacobian(J: array3d(dtype=float32), mask: array2d(dtype=int), nv: int)

Apply a column-wise mask to a (3, nv) Jacobian matrix

Source code in molmo_spaces/kinematics/parallel/warp_kinematics.py 
145
146
147
148
149
150
151
152
153
154
155
@wp.kernel
def mask_jacobian(
    J: wp.array3d(dtype=wp.float32),
    mask: wp.array2d(dtype=int),
    nv: int,
):
    """Apply a column-wise mask to a (3, nv) Jacobian matrix"""
    i = wp.tid()
    for j in range(3):
        for k in range(nv):
            J[i, j, k] = J[i, j, k] * float(mask[i, k])