Skip to content

API reference

Generated from the source docstrings. The SO101Env lerobot config class is not included here because importing it requires lerobot (and therefore torch); its fields are documented in Using it with lerobot.

Environments

lerobot_env_so101.SO101PickCubeGymEnv 

SO101PickCubeGymEnv(
    seed: int = 0,
    control_dt: float = 0.1,
    physics_dt: float = 0.002,
    render_spec: GymRenderingSpec = GymRenderingSpec(),
    render_mode: Literal[
        "rgb_array", "human"
    ] = "rgb_array",
    image_obs: bool = False,
    reward_type: str = "sparse",
    random_block_position: bool = False,
    action_scale: float = 1.0,
    top_down_ik: bool = False,
    jaw_pads: bool = False,
)

Bases: SO101GymEnv

Environment for an SO-101 robot picking up a cube.

Create the pick-cube environment.

Parameters:

Name Type Description Default
seed int

Seed for the environment's own RNG, which drives block placement when random_block_position is set.

 0
control_dt float

Seconds of simulated time per step() call. The default of 0.1 s corresponds to 10 Hz.

 0.1
physics_dt float

MuJoCo integration timestep. control_dt // physics_dt substeps are run per step().

 0.002
render_spec GymRenderingSpec

Frame height, width, and camera id used by render().

 GymRenderingSpec()
render_mode Literal['rgb_array', 'human']

"rgb_array" returns frames from render(); "human" opens a MuJoCo viewer.

 'rgb_array'
image_obs bool

Add a pixels.front camera view to the observation. Note that this replaces environment_state rather than adding to it, on the assumption that a pixel-based policy should infer the block position from the image.

 False
reward_type str

"sparse" or "dense". Sparse mode requires both a lift and the gripper being near the block.

 'sparse'
random_block_position bool

Randomize the block position on each reset instead of using a fixed pose.

 False
action_scale float

Metres per unit of position action. Applies to dx, dy, dz only — grasp is an absolute target and is unaffected. Must be positive.

 1.0
top_down_ik bool

Keep the gripper pointing straight down and move the grasp point between the jaws instead of gripperframe. See SO101GymEnv.

 False
jaw_pads bool

Enable the fingertip collision pads. See SO101GymEnv.

 False
Source code in packages/lerobot_env_so101/lerobot_env_so101/pick_env.py 
 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
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
def __init__(
    self,
    seed: int = 0,
    control_dt: float = 0.1,
    physics_dt: float = 0.002,
    render_spec: GymRenderingSpec = GymRenderingSpec(),  # noqa: B008
    render_mode: Literal["rgb_array", "human"] = "rgb_array",
    image_obs: bool = False,
    reward_type: str = "sparse",
    random_block_position: bool = False,
    action_scale: float = 1.0,
    top_down_ik: bool = False,
    jaw_pads: bool = False,
):
    """Create the pick-cube environment.

    Args:
        seed: Seed for the environment's own RNG, which drives block
            placement when ``random_block_position`` is set.
        control_dt: Seconds of simulated time per ``step()`` call. The
            default of 0.1 s corresponds to 10 Hz.
        physics_dt: MuJoCo integration timestep. ``control_dt //
            physics_dt`` substeps are run per ``step()``.
        render_spec: Frame height, width, and camera id used by
            ``render()``.
        render_mode: ``"rgb_array"`` returns frames from ``render()``;
            ``"human"`` opens a MuJoCo viewer.
        image_obs: Add a ``pixels.front`` camera view to the observation.
            Note that this *replaces* ``environment_state`` rather than
            adding to it, on the assumption that a pixel-based policy
            should infer the block position from the image.
        reward_type: ``"sparse"`` or ``"dense"``. Sparse mode requires both
            a lift and the gripper being near the block.
        random_block_position: Randomize the block position on each reset
            instead of using a fixed pose.
        action_scale: Metres per unit of position action. Applies to
            ``dx, dy, dz`` only — ``grasp`` is an absolute target and is
            unaffected. Must be positive.
        top_down_ik: Keep the gripper pointing straight down and move the
            grasp point between the jaws instead of ``gripperframe``. See
            ``SO101GymEnv``.
        jaw_pads: Enable the fingertip collision pads. See ``SO101GymEnv``.
    """
    self.reward_type = reward_type

    super().__init__(
        seed=seed,
        control_dt=control_dt,
        physics_dt=physics_dt,
        render_spec=render_spec,
        render_mode=render_mode,
        image_obs=image_obs,
        home_position=HOME_POSITION,
        cartesian_bounds=CARTESIAN_BOUNDS,
        action_scale=action_scale,
        top_down_ik=top_down_ik,
        jaw_pads=jaw_pads,
    )

    self._block_z = self._model.geom("block").size[2]
    self._random_block_position = random_block_position

    agent_dim = self.get_robot_state().shape[0]
    agent_box = spaces.Box(-np.inf, np.inf, (agent_dim,), dtype=np.float32)
    env_box = spaces.Box(-np.inf, np.inf, (3,), dtype=np.float32)

    if self.image_obs:
        self.observation_space = spaces.Dict(
            {
                "pixels": spaces.Dict(
                    {
                        "front": spaces.Box(
                            0,
                            255,
                            (self._render_specs.height, self._render_specs.width, 3),
                            dtype=np.uint8,
                        ),
                    }
                ),
                "agent_pos": agent_box,
            }
        )
    else:
        self.observation_space = spaces.Dict(
            {
                "agent_pos": agent_box,
                "environment_state": env_box,
            }
        )

reset 

reset(
    seed=None, **kwargs
) -> Tuple[Dict[str, np.ndarray], Dict[str, Any]]

Reset the environment.

Source code in packages/lerobot_env_so101/lerobot_env_so101/pick_env.py 
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def reset(self, seed=None, **kwargs) -> Tuple[Dict[str, np.ndarray], Dict[str, Any]]:
    """Reset the environment."""
    super().reset(seed=seed)

    mujoco.mj_resetData(self._model, self._data)
    self.reset_robot()

    if self._random_block_position:
        block_xy = self.np_random.uniform(*_SAMPLING_BOUNDS)
    else:
        block_xy = _FIXED_BLOCK_XY
    self._data.jnt("block").qpos[:3] = (*block_xy, self._block_z)
    mujoco.mj_forward(self._model, self._data)

    self._z_init = self._data.sensor("block_pos").data[2]
    self._z_success = self._z_init + _LIFT_SUCCESS_HEIGHT

    obs = self._compute_observation()
    return obs, {}

step 

step(
    action: ndarray,
) -> Tuple[
    Dict[str, np.ndarray], float, bool, bool, Dict[str, Any]
]

Take a step in the environment.

Source code in packages/lerobot_env_so101/lerobot_env_so101/pick_env.py 
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
def step(self, action: np.ndarray) -> Tuple[Dict[str, np.ndarray], float, bool, bool, Dict[str, Any]]:
    """Take a step in the environment."""
    self.apply_action(action)

    obs = self._compute_observation()
    rew = self._compute_reward()
    success = self._is_success()

    block_pos = self._data.sensor("block_pos").data
    exceeded_bounds = np.any(
        block_pos[:2] < (_SAMPLING_BOUNDS[0] - _BLOCK_OUT_OF_BOUNDS_MARGIN)
    ) or np.any(block_pos[:2] > (_SAMPLING_BOUNDS[1] + _BLOCK_OUT_OF_BOUNDS_MARGIN))

    terminated = bool(success or exceeded_bounds)

    return obs, rew, terminated, False, {"succeed": success}

lerobot_env_so101.SO101PickPlaceGymEnv 

SO101PickPlaceGymEnv(
    seed: int = 0,
    control_dt: float = 0.1,
    physics_dt: float = 0.002,
    render_spec: GymRenderingSpec = GymRenderingSpec(),
    render_mode: Literal[
        "rgb_array", "human"
    ] = "rgb_array",
    image_obs: bool = False,
    random_block_position: bool = False,
    action_scale: float = 1.0,
    top_down_ik: bool = True,
    jaw_pads: bool = True,
)

Bases: SO101PickCubeGymEnv

Environment for an SO-101 robot moving a cube to a goal position.

environment_state is the block position followed by the goal position. Success needs the block to have been lifted off the table earlier in the episode, so pushing it to the goal does not count. The IK keeps the gripper pointing straight down by default, so the action moves the grasp point between the jaws (ik_point_pos) rather than the gripperframe site reported in agent_pos.

Create the pick-and-place environment.

Parameters:

Name Type Description Default
seed int

Seed for the environment's own RNG.

 0
control_dt float

Seconds of simulated time per step() call.

 0.1
physics_dt float

MuJoCo integration timestep.

 0.002
render_spec GymRenderingSpec

Frame height, width, and camera id used by render().

 GymRenderingSpec()
render_mode Literal['rgb_array', 'human']

"rgb_array" or "human".

 'rgb_array'
image_obs bool

Add a pixels.front camera view to the observation in place of environment_state.

 False
random_block_position bool

Randomize the block and goal positions on each reset instead of using fixed ones.

 False
action_scale float

Metres per unit of position action. Must be positive.

 1.0
top_down_ik bool

Keep the gripper pointing straight down. Without it the jaws cannot straddle the block.

 True
jaw_pads bool

Enable the fingertip collision pads. Without them the block is pinched at two points and swings out of the jaws.

 True
Source code in packages/lerobot_env_so101/lerobot_env_so101/pick_place_env.py 
 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
101
102
103
104
105
106
107
108
109
110
111
112
113
def __init__(
    self,
    seed: int = 0,
    control_dt: float = 0.1,
    physics_dt: float = 0.002,
    render_spec: GymRenderingSpec = GymRenderingSpec(),  # noqa: B008
    render_mode: Literal["rgb_array", "human"] = "rgb_array",
    image_obs: bool = False,
    random_block_position: bool = False,
    action_scale: float = 1.0,
    top_down_ik: bool = True,
    jaw_pads: bool = True,
):
    """Create the pick-and-place environment.

    Args:
        seed: Seed for the environment's own RNG.
        control_dt: Seconds of simulated time per ``step()`` call.
        physics_dt: MuJoCo integration timestep.
        render_spec: Frame height, width, and camera id used by
            ``render()``.
        render_mode: ``"rgb_array"`` or ``"human"``.
        image_obs: Add a ``pixels.front`` camera view to the observation in
            place of ``environment_state``.
        random_block_position: Randomize the block and goal positions on
            each reset instead of using fixed ones.
        action_scale: Metres per unit of position action. Must be positive.
        top_down_ik: Keep the gripper pointing straight down. Without it the
            jaws cannot straddle the block.
        jaw_pads: Enable the fingertip collision pads. Without them the
            block is pinched at two points and swings out of the jaws.
    """
    super().__init__(
        seed=seed,
        control_dt=control_dt,
        physics_dt=physics_dt,
        render_spec=render_spec,
        render_mode=render_mode,
        image_obs=image_obs,
        reward_type="sparse",
        random_block_position=random_block_position,
        action_scale=action_scale,
        top_down_ik=top_down_ik,
        jaw_pads=jaw_pads,
    )
    self._goal_pos = np.asarray([*_FIXED_GOAL_XY, self._block_z])
    self._block_dof_adr = self._model.joint("block").dofadr[0]
    self._gripper_qpos_adr = self._model.joint("gripper").qposadr[0]
    self._was_carried = False

    if not self.image_obs:
        self.observation_space = spaces.Dict(
            {
                "agent_pos": self.observation_space["agent_pos"],
                "environment_state": spaces.Box(-np.inf, np.inf, (6,), dtype=np.float32),
            }
        )

goal_pos property 

goal_pos: ndarray

Where the block centre should rest, in world coordinates.

reset 

reset(
    seed=None, **kwargs
) -> Tuple[Dict[str, np.ndarray], Dict[str, Any]]

Reset the arm, the block, and the goal.

Source code in packages/lerobot_env_so101/lerobot_env_so101/pick_place_env.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
def reset(self, seed=None, **kwargs) -> Tuple[Dict[str, np.ndarray], Dict[str, Any]]:
    """Reset the arm, the block, and the goal."""
    # Skips SO101PickCubeGymEnv.reset, which samples from the pick task's box.
    super(SO101PickCubeGymEnv, self).reset(seed=seed)

    mujoco.mj_resetData(self._model, self._data)
    self.reset_robot()

    if self._random_block_position:
        block_xy = self.np_random.uniform(*_SAMPLING_BOUNDS)
        goal_xy = self.np_random.uniform(*_SAMPLING_BOUNDS)
        while np.linalg.norm(goal_xy - block_xy) < _MIN_GOAL_DISTANCE:
            goal_xy = self.np_random.uniform(*_SAMPLING_BOUNDS)
    else:
        block_xy, goal_xy = _FIXED_BLOCK_XY, _FIXED_GOAL_XY
    self._data.jnt("block").qpos[:3] = (*block_xy, self._block_z)
    self._goal_pos = np.asarray([*goal_xy, self._block_z])
    mujoco.mj_forward(self._model, self._data)

    self._z_init = self._data.sensor("block_pos").data[2]
    self._was_carried = False

    return self._compute_observation(), {}

step 

step(
    action: ndarray,
) -> Tuple[
    Dict[str, np.ndarray], float, bool, bool, Dict[str, Any]
]

Take a step in the environment.

Source code in packages/lerobot_env_so101/lerobot_env_so101/pick_place_env.py 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
def step(self, action: np.ndarray) -> Tuple[Dict[str, np.ndarray], float, bool, bool, Dict[str, Any]]:
    """Take a step in the environment."""
    self.apply_action(action)

    block_pos = self._data.sensor("block_pos").data
    self._was_carried = self._was_carried or bool(block_pos[2] - self._z_init > _MIN_CARRY_RISE)

    obs = self._compute_observation()
    success = self._is_success()

    block_xy = block_pos[:2]
    exceeded_bounds = np.any(block_xy < _WORKSPACE_XY[0]) or np.any(block_xy > _WORKSPACE_XY[1])
    terminated = bool(success or exceeded_bounds)

    return obs, float(success), terminated, False, {"succeed": success}

lerobot_env_so101.SO101GymEnv 

SO101GymEnv(
    xml_path: Path | None = None,
    seed: int = 0,
    control_dt: float = 0.02,
    physics_dt: float = 0.002,
    render_spec: GymRenderingSpec = GymRenderingSpec(),
    render_mode: Literal[
        "rgb_array", "human"
    ] = "rgb_array",
    image_obs: bool = False,
    home_position: ndarray = HOME_POSITION,
    cartesian_bounds: ndarray = CARTESIAN_BOUNDS,
    action_scale: float = 1.0,
    top_down_ik: bool = False,
    jaw_pads: bool = False,
)

Bases: MujocoGymEnv

SO-101 pick/place base environment: 4-DOF Cartesian position IK control.

Base class only: it implements robot control but not step()/reset(). Instantiate a task subclass such as SO101PickCubeGymEnv.

Native action space is [dx, dy, dz, grasp]. SO-101's 5-DOF arm cannot independently track a 6-DOF Cartesian pose, so the IK controller solves for position only and there are no orientation deltas to command. Use lerobot_env_so101.wrappers.SevenDofToFourDofAdapter for teleop/HIL code written against the upstream 7-dim convention.

Action deltas are unscaled metres, inherited from gym-hil's teleoperation convention: an action of 1.0 saturates the workspace in a single step. See "Known limitations" in the README before training a policy against this.

The grasp component is an absolute target in [0, 1] (0=closed, 1=open), not a delta: it is written directly to the gripper actuator's ctrl each step via lerobot_env_so101.gripper.normalized_to_ctrl.

Create the SO-101 base environment.

Parameters:

Name Type Description Default
xml_path Path | None

MJCF scene to load. Defaults to the bundled pick_scene.xml.

 None
seed int

Seed for the environment's RNG.

 0
control_dt float

Seconds of simulated time per control step.

 0.02
physics_dt float

MuJoCo integration timestep.

 0.002
render_spec GymRenderingSpec

Frame height, width, and camera id used by render().

 GymRenderingSpec()
render_mode Literal['rgb_array', 'human']

"rgb_array" or "human".

 'rgb_array'
image_obs bool

Whether the task subclass should include a camera view in its observation.

 False
home_position ndarray

Arm joint configuration applied by reset_robot().

 HOME_POSITION
cartesian_bounds ndarray

(2, 3) array of min/max end-effector target positions. The IK target is clipped to this box every step.

 CARTESIAN_BOUNDS
action_scale float

Metres per unit of position action. Must be positive.

 1.0
top_down_ik bool

Solve the IK for the grasp point between the jaws with the gripper pointing straight down, instead of for the gripperframe position alone. The action then moves the grasp point.

 False
jaw_pads bool

Enable collision on the box pads at the two fingertips. MuJoCo collides the jaw meshes as convex hulls, which touch a block at one point each; the pads give it flat faces to be held by.

 False

Raises:

Type Description
ValueError

If action_scale is not positive.
Source code in packages/lerobot_env_so101/lerobot_env_so101/mujoco_env.py 
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
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
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
def __init__(
    self,
    xml_path: Path | None = None,
    seed: int = 0,
    control_dt: float = 0.02,
    physics_dt: float = 0.002,
    render_spec: GymRenderingSpec = GymRenderingSpec(),  # noqa: B008
    render_mode: Literal["rgb_array", "human"] = "rgb_array",
    image_obs: bool = False,
    home_position: np.ndarray = HOME_POSITION,
    cartesian_bounds: np.ndarray = CARTESIAN_BOUNDS,
    action_scale: float = 1.0,
    top_down_ik: bool = False,
    jaw_pads: bool = False,
):
    """Create the SO-101 base environment.

    Args:
        xml_path: MJCF scene to load. Defaults to the bundled
            ``pick_scene.xml``.
        seed: Seed for the environment's RNG.
        control_dt: Seconds of simulated time per control step.
        physics_dt: MuJoCo integration timestep.
        render_spec: Frame height, width, and camera id used by
            ``render()``.
        render_mode: ``"rgb_array"`` or ``"human"``.
        image_obs: Whether the task subclass should include a camera view
            in its observation.
        home_position: Arm joint configuration applied by
            ``reset_robot()``.
        cartesian_bounds: ``(2, 3)`` array of min/max end-effector target
            positions. The IK target is clipped to this box every step.
        action_scale: Metres per unit of position action. Must be
            positive.
        top_down_ik: Solve the IK for the grasp point between the jaws
            with the gripper pointing straight down, instead of for the
            ``gripperframe`` position alone. The action then moves the
            grasp point.
        jaw_pads: Enable collision on the box pads at the two fingertips.
            MuJoCo collides the jaw meshes as convex hulls, which touch a
            block at one point each; the pads give it flat faces to be
            held by.

    Raises:
        ValueError: If ``action_scale`` is not positive.
    """
    if xml_path is None:
        xml_path = ASSETS_DIR / "pick_scene.xml"

    if action_scale <= 0.0:
        raise ValueError(f"action_scale must be positive, got {action_scale}")

    self._home_position = home_position
    self._cartesian_bounds = cartesian_bounds
    self._action_scale = action_scale
    self._top_down_ik = top_down_ik

    super().__init__(
        xml_path=xml_path,
        seed=seed,
        control_dt=control_dt,
        physics_dt=physics_dt,
        render_spec=render_spec,
    )

    self.metadata = {
        **self.metadata,
        "render_fps": int(np.round(1.0 / self.control_dt)),
    }
    self.render_mode = render_mode
    self.image_obs = image_obs

    self._arm_dof_ids = np.asarray([self._model.joint(name).id for name in _ARM_JOINT_NAMES])
    self._arm_ctrl_ids = np.asarray([self._model.actuator(name).id for name in _ARM_JOINT_NAMES])
    self._gripper_ctrl_id = self._model.actuator("gripper").id
    self._ee_site_id = self._model.site("gripperframe").id
    if jaw_pads:
        pad_ids = [self._model.geom(name).id for name in _JAW_PAD_GEOM_NAMES]
        self._model.geom_contype[pad_ids] = 1
        self._model.geom_conaffinity[pad_ids] = 1
    self._roll_joint_id = self._model.joint(_ROLL_JOINT_NAME).id
    self._roll_index = _ARM_JOINT_NAMES.index(_ROLL_JOINT_NAME)
    self._ik_point_offset = (
        grasp_point_offset(self._model, self._ee_site_id) if top_down_ik else None
    )
    self._camera_id = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_CAMERA, "front")

    self._target_ee_pos: np.ndarray | None = None
    self._last_action_was_zero = False

    base_obs_space = {
        "agent_pos": spaces.Dict(
            {
                "tcp_pose": spaces.Box(-np.inf, np.inf, shape=(7,), dtype=np.float32),
                "tcp_vel": spaces.Box(-np.inf, np.inf, shape=(6,), dtype=np.float32),
                "gripper_pose": spaces.Box(-1, 1, shape=(1,), dtype=np.float32),
            }
        )
    }
    if self.image_obs:
        self.observation_space = spaces.Dict(
            {
                **base_obs_space,
                "pixels": spaces.Dict(
                    {
                        "front": spaces.Box(
                            low=0,
                            high=255,
                            shape=(self._render_specs.height, self._render_specs.width, 3),
                            dtype=np.uint8,
                        ),
                    }
                ),
            }
        )
    else:
        self.observation_space = spaces.Dict(base_obs_space)

    self.action_space = spaces.Box(
        low=np.asarray([-1.0, -1.0, -1.0, 0.0], dtype=np.float32),
        high=np.asarray([1.0, 1.0, 1.0, 1.0], dtype=np.float32),
        dtype=np.float32,
    )

ik_point_pos property 

ik_point_pos: ndarray

World position of the point the action moves.

The gripperframe site by default, the grasp point between the jaws with top_down_ik.

reset_robot 

reset_robot()

Reset the arm to the home position and resynchronize the IK target.

The IK target (_target_ee_pos) is an accumulated Python attribute rather than a MuJoCo body, so it must be re-read from the end-effector sensor here; otherwise it would still hold the target from before the reset.

Source code in packages/lerobot_env_so101/lerobot_env_so101/mujoco_env.py 
335
336
337
338
339
340
341
342
343
344
345
346
347
348
def reset_robot(self):
    """Reset the arm to the home position and resynchronize the IK target.

    The IK target (``_target_ee_pos``) is an accumulated Python attribute
    rather than a MuJoCo body, so it must be re-read from the end-effector
    sensor here; otherwise it would still hold the target from before the
    reset.
    """
    self._data.qpos[self._arm_dof_ids] = self._home_position
    self._data.ctrl[self._arm_ctrl_ids] = self._home_position
    mujoco.mj_forward(self._model, self._data)

    self._target_ee_pos = self.ik_point_pos
    self._last_action_was_zero = False

apply_action 

apply_action(action: ndarray) -> None

Apply a native 4-dim action [dx, dy, dz, grasp] via position IK.

Position deltas are multiplied by action_scale (metres per unit action); the grasp command is an absolute target in [0, 1] (0=closed, 1=open), written directly to the gripper ctrl each step, and is not scaled.

A near-zero position delta holds the current pose rather than the accumulated target, so repeated zero actions do not drift toward an overshot target.

Parameters:

Name Type Description Default
action ndarray

Array of [dx, dy, dz, grasp]. Position deltas are clipped to the environment's Cartesian bounds after scaling.

 required
Source code in packages/lerobot_env_so101/lerobot_env_so101/mujoco_env.py 
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
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
def apply_action(self, action: np.ndarray) -> None:
    """Apply a native 4-dim action ``[dx, dy, dz, grasp]`` via position IK.

    Position deltas are multiplied by ``action_scale`` (metres per unit
    action); the grasp command is an absolute target in [0, 1]
    (0=closed, 1=open), written directly to the gripper ctrl each step,
    and is not scaled.

    A near-zero position delta holds the current pose rather than the
    accumulated target, so repeated zero actions do not drift toward an
    overshot target.

    Args:
        action: Array of ``[dx, dy, dz, grasp]``. Position deltas are
            clipped to the environment's Cartesian bounds after scaling.
    """
    x, y, z, grasp_command = action

    # Zero-action detection uses the raw input so the hold-position
    # behaviour does not depend on the configured scale.
    raw_delta = np.asarray([x, y, z])
    action_is_zero = np.linalg.norm(raw_delta) < _ZERO_ACTION_TOLERANCE
    delta = raw_delta * self._action_scale

    # On transition from motion to hold, lock the target to the current position
    # instead of the last accumulated (and possibly overshot) target.
    if action_is_zero and not self._last_action_was_zero:
        self._target_ee_pos = self.ik_point_pos
    elif not action_is_zero:
        self._target_ee_pos = self._target_ee_pos + delta
    self._last_action_was_zero = action_is_zero

    self._target_ee_pos = np.clip(
        self._target_ee_pos, self._cartesian_bounds[0], self._cartesian_bounds[1]
    )

    gripper_range = self._model.actuator("gripper").ctrlrange
    self._data.ctrl[self._gripper_ctrl_id] = normalized_to_ctrl(grasp_command, gripper_range)

    # The arm actuators are <position>: ctrl is a target angle held across
    # the substeps, so the IK is solved once per control step.
    target_q = solve_ik(
        model=self._model,
        data=self._data,
        site_id=self._ee_site_id,
        dof_ids=self._arm_dof_ids,
        target_pos=self._target_ee_pos,
        ik_method="levenberg_marquardt",
        ik_damping=_IK_DAMPING,
        ik_iterations=_IK_ITERATIONS,
        seed_q=self._top_down_seed() if self._top_down_ik else None,
        point_offset=self._ik_point_offset,
        approach_axis=_APPROACH_AXIS if self._top_down_ik else None,
    )
    arm_ctrlrange = self._model.actuator_ctrlrange[self._arm_ctrl_ids]
    self._data.ctrl[self._arm_ctrl_ids] = np.clip(target_q, arm_ctrlrange[:, 0], arm_ctrlrange[:, 1])

    for _ in range(self._n_substeps):
        mujoco.mj_step(self._model, self._data)

get_robot_state 

get_robot_state()

Get the current state of the robot.

Source code in packages/lerobot_env_so101/lerobot_env_so101/mujoco_env.py 
438
439
440
441
442
443
444
445
def get_robot_state(self):
    """Get the current state of the robot."""
    ee_pos = self._data.sensor("so101/ee_pos").data
    qpos = self.data.qpos[self._arm_dof_ids].astype(np.float32)
    qvel = self.data.qvel[self._arm_dof_ids].astype(np.float32)
    gripper_pose = self.get_gripper_pose()

    return np.concatenate([qpos, qvel, gripper_pose, ee_pos])

get_gripper_pose 

get_gripper_pose()

Gripper ctrl target normalized to [-1, 1]: -1 = fully closed, +1 = fully open.

Source code in packages/lerobot_env_so101/lerobot_env_so101/mujoco_env.py 
447
448
449
450
451
452
def get_gripper_pose(self):
    """Gripper ctrl target normalized to [-1, 1]: -1 = fully closed, +1 = fully open."""
    gripper_range = self._model.actuator("gripper").ctrlrange
    gripper_pos = self._data.ctrl[self._gripper_ctrl_id]
    normalized = 2 * (gripper_pos - gripper_range[0]) / (gripper_range[1] - gripper_range[0]) - 1
    return np.array([normalized], dtype=np.float32)

render 

render() -> np.ndarray

Render the front camera view.

The renderer is built on first use: constructing it eagerly would open an OpenGL context on every instantiation, which fails on a headless host even for state-only training that never renders.

Source code in packages/lerobot_env_so101/lerobot_env_so101/mujoco_env.py 
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
def render(self) -> np.ndarray:
    """Render the front camera view.

    The renderer is built on first use: constructing it eagerly would open
    an OpenGL context on every instantiation, which fails on a headless
    host even for state-only training that never renders.
    """
    if self._viewer is None:
        self._viewer = mujoco.Renderer(
            model=self._model,
            height=self._render_specs.height,
            width=self._render_specs.width,
        )
    self._viewer.update_scene(self.data, camera=self._camera_id)
    return self._viewer.render()

lerobot_env_so101.MujocoGymEnv 

MujocoGymEnv(
    xml_path: Path,
    seed: int = 0,
    control_dt: float = 0.02,
    physics_dt: float = 0.002,
    render_spec: GymRenderingSpec = GymRenderingSpec(),
)

Bases: Env

A MuJoCo model wrapped in the gymnasium interface.

Owns the model, data, and renderer, and derives the number of physics substeps per control step. Robot-specific control and the task loop belong to subclasses.

Load an MJCF model and configure the simulation timestep.

Parameters:

Name Type Description Default
xml_path Path

Path to the MJCF scene.

 required
seed int

Seed for random_state.

 0
control_dt float

Seconds of simulated time per control step.

 0.02
physics_dt float

MuJoCo integration timestep. The number of substeps per control step is control_dt // physics_dt.

 0.002
render_spec GymRenderingSpec

Frame height, width, and camera id used by render().

 GymRenderingSpec()
Source code in packages/lerobot_env_so101/lerobot_env_so101/mujoco_env.py 
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
def __init__(
    self,
    xml_path: Path,
    seed: int = 0,
    control_dt: float = 0.02,
    physics_dt: float = 0.002,
    render_spec: GymRenderingSpec = GymRenderingSpec(),  # noqa: B008
):
    """Load an MJCF model and configure the simulation timestep.

    Args:
        xml_path: Path to the MJCF scene.
        seed: Seed for ``random_state``.
        control_dt: Seconds of simulated time per control step.
        physics_dt: MuJoCo integration timestep. The number of substeps per
            control step is ``control_dt // physics_dt``.
        render_spec: Frame height, width, and camera id used by
            ``render()``.
    """
    self._model = mujoco.MjModel.from_xml_path(xml_path.as_posix())
    self._model.vis.global_.offwidth = render_spec.width
    self._model.vis.global_.offheight = render_spec.height
    self._data = mujoco.MjData(self._model)
    self._model.opt.timestep = physics_dt
    self._control_dt = control_dt
    self._n_substeps = int(control_dt // physics_dt)
    self._random = np.random.RandomState(seed)
    self._viewer: Optional[mujoco.Renderer] = None
    self._render_specs = render_spec

render 

render() -> np.ndarray

Render one offscreen frame.

Returns:

Type Description
ndarray

An (height, width, 3) uint8 array from the configured
ndarray

camera.
Source code in packages/lerobot_env_so101/lerobot_env_so101/mujoco_env.py 
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
def render(self) -> np.ndarray:
    """Render one offscreen frame.

    Returns:
        An ``(height, width, 3)`` ``uint8`` array from the configured
        camera.
    """
    if self._viewer is None:
        self._viewer = mujoco.Renderer(
            model=self._model,
            height=self._render_specs.height,
            width=self._render_specs.width,
        )
    self._viewer.update_scene(self._data, camera=self._render_specs.camera_id)
    return self._viewer.render()

close 

close() -> None

Release graphics resources if they exist.

Source code in packages/lerobot_env_so101/lerobot_env_so101/mujoco_env.py 
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def close(self) -> None:
    """Release graphics resources if they exist."""
    viewer = self._viewer
    if viewer is None:
        return

    if hasattr(viewer, "close") and callable(viewer.close):
        try:
            viewer.close()
        except Exception:
            # Teardown failures must not mask the caller's own cleanup, but a
            # leaking GL context only shows up as slow growth over a long run.
            _logger.debug("viewer.close() failed", exc_info=True)

    self._viewer = None

lerobot_env_so101.GymRenderingSpec dataclass 

GymRenderingSpec(
    height: int = 128,
    width: int = 128,
    camera_id: str | int = -1,
    mode: Literal["rgb_array", "human"] = "rgb_array",
)

Offscreen frame size and camera used for rendering.

camera_id of -1 selects MuJoCo's free camera; pass a camera name from the MJCF (for example "front") to render a fixed view.

Wrappers

lerobot_env_so101.SevenDofToFourDofAdapter 

SevenDofToFourDofAdapter(env: Env)

Bases: ActionWrapper

Adapts 7-dim Cartesian actions to this package's native 4-dim action space.

Teleop and HIL tooling built on gym-hil's Panda convention emits 7-dim actions [dx, dy, dz, drx, dry, drz, grasp]. SO-101 environments take a native 4-dim [dx, dy, dz, grasp]; this wrapper discards the orientation deltas (indices 3-5), which the position-only IK controller cannot act on. grasp is passed through unchanged, so its meaning follows the wrapped env's convention (an absolute target in [0, 1], 0=closed/1=open).

Wrap a native SO101GymEnv-based env to accept 7-dim actions from existing teleop/HIL code; new code should prefer the native 4-dim action space.

Wrap a native 4-dim env so it accepts 7-dim actions.

Emits a :class:UserWarning once at construction, rather than per step, so that a caller porting 7-dim code learns why the arm never rotates.

Parameters:

Name Type Description Default
env Env

An environment whose action space is a 4-dim Box.

 required

Raises:

Type Description
ValueError

If the wrapped env's action space is not a 4-dim Box.
Source code in packages/lerobot_env_so101/lerobot_env_so101/wrappers.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
def __init__(self, env: gym.Env):
    """Wrap a native 4-dim env so it accepts 7-dim actions.

    Emits a :class:`UserWarning` once at construction, rather than per
    step, so that a caller porting 7-dim code learns why the arm never
    rotates.

    Args:
        env: An environment whose action space is a 4-dim ``Box``.

    Raises:
        ValueError: If the wrapped env's action space is not a 4-dim
            ``Box``.
    """
    super().__init__(env)
    # Warn once at construction, not per step: a caller porting 7-dim code
    # otherwise gets an arm that never rotates and no indication why.
    warnings.warn(
        "SevenDofToFourDofAdapter discards the rotation deltas (drx, dry, drz) "
        "of 7-dim actions. SO-101's IK controller solves for end-effector "
        "position only, so rotation commands have no effect.",
        stacklevel=2,
    )
    native_space = env.action_space
    if not isinstance(native_space, spaces.Box) or native_space.shape != (FOUR_DOF_DIM,):
        raise ValueError(
            f"SevenDofToFourDofAdapter expects a wrapped env with a {FOUR_DOF_DIM}-dim "
            f"Box action space, got {native_space}"
        )

    insert_at = (3,) * len(_DISCARDED_ROTATION_INDICES)
    low = np.insert(native_space.low, insert_at, -1.0).astype(np.float32)
    high = np.insert(native_space.high, insert_at, 1.0).astype(np.float32)
    self.action_space = spaces.Box(low=low, high=high, dtype=np.float32)

action 

action(action: ndarray) -> np.ndarray

Drop the rotation deltas from a 7-dim action.

Parameters:

Name Type Description Default
action ndarray

[dx, dy, dz, drx, dry, drz, grasp].

 required

Returns:

Type Description
ndarray

[dx, dy, dz, grasp], with indices 3-5 discarded.
Source code in packages/lerobot_env_so101/lerobot_env_so101/wrappers.py 
81
82
83
84
85
86
87
88
89
90
def action(self, action: np.ndarray) -> np.ndarray:
    """Drop the rotation deltas from a 7-dim action.

    Args:
        action: ``[dx, dy, dz, drx, dry, drz, grasp]``.

    Returns:
        ``[dx, dy, dz, grasp]``, with indices 3-5 discarded.
    """
    return np.asarray(action, dtype=np.float32)[_KEPT_INDICES]

Gripper utilities

lerobot_env_so101.gripper.normalized_to_ctrl 

normalized_to_ctrl(norm: float, ctrlrange) -> float

Map a normalized gripper command in [0, 1] (0=closed, 1=open) to ctrl radians.

Out-of-range input is clipped to [0, 1] before mapping.

Source code in packages/lerobot_env_so101/lerobot_env_so101/gripper.py 
31
32
33
34
35
36
37
38
39
40
def normalized_to_ctrl(norm: float, ctrlrange) -> float:
    """Map a normalized gripper command in [0, 1] (0=closed, 1=open) to ctrl radians.

    Out-of-range input is clipped to [0, 1] before mapping.
    """
    norm = np.clip(norm, 0.0, 1.0)
    lo, hi = ctrlrange
    if not GRIPPER_CLOSED_AT_CTRL_LOW:
        norm = 1.0 - norm
    return lo + norm * (hi - lo)

lerobot_env_so101.gripper.ctrl_to_normalized 

ctrl_to_normalized(ctrl: float, ctrlrange) -> float

Inverse of normalized_to_ctrl: ctrl radians -> normalized [0, 1].

Source code in packages/lerobot_env_so101/lerobot_env_so101/gripper.py 
43
44
45
46
47
48
49
def ctrl_to_normalized(ctrl: float, ctrlrange) -> float:
    """Inverse of ``normalized_to_ctrl``: ctrl radians -> normalized [0, 1]."""
    lo, hi = ctrlrange
    norm = (ctrl - lo) / (hi - lo)
    if not GRIPPER_CLOSED_AT_CTRL_LOW:
        norm = 1.0 - norm
    return norm

Inverse kinematics

These back the position-only Cartesian control described in Action space. They are internal to the environment rather than part of the intended public surface, but are documented here because their behaviour explains several of the limitations.

lerobot_env_so101.ik_control.solve_ik 

solve_ik(
    model: MjModel,
    data: MjData,
    site_id: int,
    dof_ids: ndarray,
    target_pos: ndarray,
    ik_method: str = "levenberg_marquardt",
    ik_damping: float = 0.1,
    ik_iterations: int = 20,
    seed_q: ndarray | None = None,
    point_offset: ndarray | None = None,
    approach_axis: int | None = None,
    approach_dir: ndarray = DOWN,
    orientation_weight: float = DEFAULT_ORIENTATION_WEIGHT,
) -> np.ndarray

Solve for the joint angles that place the end-effector at a Cartesian target.

The returned angles are the command for the arm's <position> actuators.

Parameters:

Name Type Description Default
model MjModel

MuJoCo model

 required
data MjData

MuJoCo data

 required
site_id int

Site ID for end-effector

 required
dof_ids ndarray

DOF IDs for controlled joints

 required
target_pos ndarray

Desired Cartesian position (3,)

 required
ik_method str

"levenberg_marquardt" or "pseudoinverse"

 'levenberg_marquardt'
ik_damping float

Damping for LM (if used)

 0.1
ik_iterations int

Max IK iterations

 20
seed_q ndarray | None

Joint angles the iteration starts from. Defaults to the current qpos. The solver converges to the branch nearest the seed, so this is how a caller selects among redundant solutions.

 None
point_offset ndarray | None

Site-frame offset of the controlled point (3,)

 None
approach_axis int | None

Site-frame axis index to point along approach_dir; None solves for position only. Levenberg-Marquardt only.

 None
approach_dir ndarray

Unit world direction for the approach axis (3,)

 DOWN
orientation_weight float

Metres of position error per radian of tilt

 DEFAULT_ORIENTATION_WEIGHT

Returns:

Type Description
ndarray

Target joint angles (len(dof_ids),)

Raises:

Type Description
ValueError

If ik_method is not a known solver, or if a control point offset or approach axis is requested from the pseudoinverse solver.
Source code in packages/lerobot_env_so101/lerobot_env_so101/ik_control.py 
250
251
252
253
254
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
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
def solve_ik(
    model: mujoco.MjModel,
    data: mujoco.MjData,
    site_id: int,
    dof_ids: np.ndarray,
    target_pos: np.ndarray,
    ik_method: str = "levenberg_marquardt",
    ik_damping: float = 0.1,
    ik_iterations: int = 20,
    seed_q: np.ndarray | None = None,
    point_offset: np.ndarray | None = None,
    approach_axis: int | None = None,
    approach_dir: np.ndarray = DOWN,
    orientation_weight: float = DEFAULT_ORIENTATION_WEIGHT,
) -> np.ndarray:
    """
    Solve for the joint angles that place the end-effector at a Cartesian target.

    The returned angles are the command for the arm's ``<position>`` actuators.

    Args:
        model: MuJoCo model
        data: MuJoCo data
        site_id: Site ID for end-effector
        dof_ids: DOF IDs for controlled joints
        target_pos: Desired Cartesian position (3,)
        ik_method: "levenberg_marquardt" or "pseudoinverse"
        ik_damping: Damping for LM (if used)
        ik_iterations: Max IK iterations
        seed_q: Joint angles the iteration starts from. Defaults to the
            current ``qpos``. The solver converges to the branch nearest the
            seed, so this is how a caller selects among redundant solutions.
        point_offset: Site-frame offset of the controlled point (3,)
        approach_axis: Site-frame axis index to point along ``approach_dir``;
            ``None`` solves for position only. Levenberg-Marquardt only.
        approach_dir: Unit world direction for the approach axis (3,)
        orientation_weight: Metres of position error per radian of tilt

    Returns:
        Target joint angles (len(dof_ids),)

    Raises:
        ValueError: If ``ik_method`` is not a known solver, or if a control
            point offset or approach axis is requested from the pseudoinverse
            solver.
    """
    current_q = data.qpos[dof_ids].copy()
    if seed_q is None:
        seed_q = current_q

    # Compute target joint angles via IK
    if ik_method == "levenberg_marquardt":
        target_q = compute_ik_levenberg_marquardt(
            model=model,
            data=data,
            site_id=site_id,
            dof_ids=dof_ids,
            target_pos=target_pos,
            current_q=seed_q,
            damping=ik_damping,
            max_iterations=ik_iterations,
            point_offset=point_offset,
            approach_axis=approach_axis,
            approach_dir=approach_dir,
            orientation_weight=orientation_weight,
        )
    elif ik_method == "pseudoinverse":
        if point_offset is not None or approach_axis is not None:
            raise ValueError("point_offset and approach_axis require ik_method='levenberg_marquardt'")
        target_q = compute_ik_pseudoinverse(
            model=model,
            data=data,
            site_id=site_id,
            dof_ids=dof_ids,
            target_pos=target_pos,
            current_q=seed_q,
            max_iterations=ik_iterations,
        )
    else:
        raise ValueError(f"Unknown IK method: {ik_method}")

    # Restore current state (IK changed qpos for computation)
    data.qpos[dof_ids] = current_q
    mujoco.mj_forward(model, data)

    return target_q

lerobot_env_so101.ik_control.compute_ik_levenberg_marquardt 

compute_ik_levenberg_marquardt(
    model: MjModel,
    data: MjData,
    site_id: int,
    dof_ids: ndarray,
    target_pos: ndarray,
    current_q: ndarray,
    damping: float = 0.1,
    max_iterations: int = 20,
    tolerance: float = 0.0001,
    point_offset: ndarray | None = None,
    approach_axis: int | None = None,
    approach_dir: ndarray = DOWN,
    orientation_weight: float = DEFAULT_ORIENTATION_WEIGHT,
) -> np.ndarray

Compute target joint angles using Levenberg-Marquardt IK.

Parameters:

Name Type Description Default
model MjModel

MuJoCo model

 required
data MjData

MuJoCo data

 required
site_id int

Site ID for end-effector

 required
dof_ids ndarray

DOF IDs for controlled joints

 required
target_pos ndarray

Desired Cartesian position (3,)

 required
current_q ndarray

Current joint angles (n,)

 required
damping float

Damping factor (lambda in LM algorithm)

 0.1
max_iterations int

Maximum IK iterations

 20
tolerance float

Convergence tolerance

 0.0001
point_offset ndarray | None

Offset of the controlled point from the site origin, in the site frame (3,). None controls the site origin.

 None
approach_axis int | None

Index (0, 1, 2) of the site-frame axis to point along approach_dir. None solves for position only.

 None
approach_dir ndarray

Unit world direction for the approach axis (3,)

 DOWN
orientation_weight float

Metres of position error equivalent to one radian of approach-axis tilt

 DEFAULT_ORIENTATION_WEIGHT

Returns:

Type Description
ndarray

Target joint angles (n,)
Source code in packages/lerobot_env_so101/lerobot_env_so101/ik_control.py 
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
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
194
195
196
197
198
199
200
201
202
203
204
205
def compute_ik_levenberg_marquardt(
    model: mujoco.MjModel,
    data: mujoco.MjData,
    site_id: int,
    dof_ids: np.ndarray,
    target_pos: np.ndarray,
    current_q: np.ndarray,
    damping: float = 0.1,
    max_iterations: int = 20,
    tolerance: float = 1e-4,
    point_offset: np.ndarray | None = None,
    approach_axis: int | None = None,
    approach_dir: np.ndarray = DOWN,
    orientation_weight: float = DEFAULT_ORIENTATION_WEIGHT,
) -> np.ndarray:
    """
    Compute target joint angles using Levenberg-Marquardt IK.

    Args:
        model: MuJoCo model
        data: MuJoCo data
        site_id: Site ID for end-effector
        dof_ids: DOF IDs for controlled joints
        target_pos: Desired Cartesian position (3,)
        current_q: Current joint angles (n,)
        damping: Damping factor (lambda in LM algorithm)
        max_iterations: Maximum IK iterations
        tolerance: Convergence tolerance
        point_offset: Offset of the controlled point from the site origin, in
            the site frame (3,). ``None`` controls the site origin.
        approach_axis: Index (0, 1, 2) of the site-frame axis to point along
            ``approach_dir``. ``None`` solves for position only.
        approach_dir: Unit world direction for the approach axis (3,)
        orientation_weight: Metres of position error equivalent to one radian
            of approach-axis tilt

    Returns:
        Target joint angles (n,)
    """
    q = current_q.copy()

    for _iteration in range(max_iterations):
        # Set joint angles and update kinematics
        data.qpos[dof_ids] = q
        mujoco.mj_forward(model, data)

        error, J = _task_error_and_jacobian(
            model,
            data,
            site_id,
            dof_ids,
            target_pos,
            point_offset,
            approach_axis,
            approach_dir,
            orientation_weight,
        )
        error_norm = np.linalg.norm(error)

        # Check convergence
        if error_norm < tolerance:
            break

        # Levenberg-Marquardt update: delta_q = (J^T*J + λI)^-1 * J^T * error
        JTJ = J.T @ J
        lambda_I = damping * np.eye(len(dof_ids))
        delta_q = np.linalg.solve(JTJ + lambda_I, J.T @ error)

        # Update joint angles
        q = q + delta_q

        # Optional: enforce joint limits
        for i, dof_id in enumerate(dof_ids):
            joint_range = model.jnt_range[dof_id]
            if joint_range[0] < joint_range[1]:  # Has limits
                q[i] = np.clip(q[i], joint_range[0], joint_range[1])

    return q

lerobot_env_so101.ik_control.compute_ik_pseudoinverse 

compute_ik_pseudoinverse(
    model: MjModel,
    data: MjData,
    site_id: int,
    dof_ids: ndarray,
    target_pos: ndarray,
    current_q: ndarray,
    max_iterations: int = 20,
    tolerance: float = 0.0001,
    step_size: float = 1.0,
) -> np.ndarray

Compute target joint angles using pseudoinverse (Gauss-Newton) IK.

Faster but less robust than Levenberg-Marquardt.

Source code in packages/lerobot_env_so101/lerobot_env_so101/ik_control.py 
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
def compute_ik_pseudoinverse(
    model: mujoco.MjModel,
    data: mujoco.MjData,
    site_id: int,
    dof_ids: np.ndarray,
    target_pos: np.ndarray,
    current_q: np.ndarray,
    max_iterations: int = 20,
    tolerance: float = 1e-4,
    step_size: float = 1.0,
) -> np.ndarray:
    """
    Compute target joint angles using pseudoinverse (Gauss-Newton) IK.

    Faster but less robust than Levenberg-Marquardt.
    """
    q = current_q.copy()

    for _iteration in range(max_iterations):
        data.qpos[dof_ids] = q
        mujoco.mj_forward(model, data)

        current_pos = data.site_xpos[site_id].copy()
        error = target_pos - current_pos

        if np.linalg.norm(error) < tolerance:
            break

        J_v = np.zeros((3, model.nv))
        J_w = np.zeros((3, model.nv))
        mujoco.mj_jacSite(model, data, J_v, J_w, site_id)
        J = J_v[:, dof_ids]

        # Pseudoinverse update: delta_q = J^† * error
        J_pinv = np.linalg.pinv(J)
        delta_q = J_pinv @ error

        q = q + step_size * delta_q

    return q