Compare commits

...

4 Commits

Author SHA1 Message Date
Martino Russi 166713561d (add) EE-kinematics retargeting for cross-morphology policy rollout: 2026-08-06 19:34:27 +02:00
Pepijn 506d16c7cd Merge branch 'main' into feat/openarm_ee_kinematics 2026-08-03 12:41:57 +02:00
Martino Russi 2e37cb22e8 docs(openarm): add episode-replay example
Add examples/openarm with a self-contained script that replays a recorded
bimanual-OpenArm episode into an mp4 by driving the official OpenArm MuJoCo
model (enactic/openarm_mujoco, v1) from a LeRobot dataset's observation.state,
plus a README covering model provenance, setup, and troubleshooting.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-31 15:42:32 +02:00
Martino Russi ea5cabe6ff feat(openarm): add end-effector kinematics support
Expose optional URDF-based forward/inverse kinematics for the OpenArm
follower so it can be recorded/commanded in end-effector (Cartesian)
space, mirroring the SO-100 kinematics pattern.

- config: add optional urdf_path / target_frame_name fields
- robot: add arm_motor_names property and make_kinematics() helper that
  builds a RobotKinematics solver (lazy import; placo only needed when used)

No behavior change when urdf_path is unset (joint-space only).

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-31 15:24:25 +02:00
8 changed files with 1895 additions and 1 deletions
+95
View File
@@ -0,0 +1,95 @@
# OpenArm — Episode Replay in Simulation
Replay a recorded bimanual-[OpenArm](https://openarm.dev) episode into an mp4 by driving the
official OpenArm MuJoCo model from a LeRobot dataset's recorded joint states.
By default the replay goes **through end-effector (Cartesian) space**: for every frame and
every arm it runs forward kinematics (recorded joints → EE pose) and then inverse kinematics
(EE pose → joints), and drives the simulator with the IK-recovered joints. This exercises the
exact `RobotKinematics` solver that `OpenArmFollower.make_kinematics()` builds (see the
[OpenArm docs](../../docs/source/openarm.mdx)), so the video is a visual sanity check of the
end-effector kinematics — not just of the raw recording. It also prints the FK→IK round-trip
error (mean joint error in degrees and mean EE-position error in mm). Pass `--joint-space` to
bypass kinematics and replay the raw recorded joints directly.
## Model provenance
Everything is pulled from Enactic's official, Apache-2.0 OpenArm repositories — nothing is
vendored into LeRobot:
| Asset | Source | License |
| ------------------------------------ | ------------------------------------------------------------------------------- | ---------- |
| MuJoCo MJCF (rendering) | [`enactic/openarm_mujoco`](https://github.com/enactic/openarm_mujoco) | Apache-2.0 |
| URDF / xacro (for `RobotKinematics`, the FK/IK round-trip) | [`enactic/openarm_description`](https://github.com/enactic/openarm_description) | Apache-2.0 |
End-effector kinematics use [`placo`](https://github.com/Rhoban/placo) under the hood (install
LeRobot with the `placo-dep` extra). The script auto-locates the URDF via `--urdf`
`$OPENARM_URDF`; set the tip link with `--ee-frame``$OPENARM_EE_FRAME`.
Use the **v1** MuJoCo revision (`v1/openarm_bimanual.xml`). v2 is a different wrist hardware
revision (DM3507) and will look sign-flipped when replaying v1 recordings.
## Setup
```bash
# LeRobot in your env (see https://huggingface.co/docs/lerobot/installation)
# Plus the sim/replay deps:
pip install mujoco av pandas
# Get the OpenArm MuJoCo model (either works):
pip install openarm-mujoco # installs models under <prefix>/share/openarm_mujoco/
# or
git clone https://github.com/enactic/openarm_mujoco.git # then pass --mjcf .../v1/openarm_bimanual.xml
```
The script auto-locates the model in this order: `--mjcf` arg → `$OPENARM_MJCF`
`<sys.prefix>/share/openarm_mujoco/v1/openarm_bimanual.xml`.
## Dataset layout
`observation.state` must be the 16-D bimanual vector (degrees):
```
right_joint_1..7, right_gripper, left_joint_1..7, left_gripper
```
Only the 14 arm joints affect the rendered pose; the two gripper scalars drive the fingers.
## Run
Headless rendering needs `MUJOCO_GL=egl`, and MuJoCo's GL libs on `LD_LIBRARY_PATH`
(in conda: `$CONDA_PREFIX/lib`).
```bash
# Replay episode 1 through end-effector kinematics (default)
LD_LIBRARY_PATH=$CONDA_PREFIX/lib MUJOCO_GL=egl \
python -m examples.openarm.render_episode \
--dataset data/folding_src_meta \
--episode 1 \
--urdf /path/to/openarm.urdf \
--ee-frame openarm_finger_tip_link \
--out openarm_ep1.mp4
# Bypass kinematics and replay the raw recorded joints directly
LD_LIBRARY_PATH=$CONDA_PREFIX/lib MUJOCO_GL=egl \
python -m examples.openarm.render_episode \
--dataset data/folding_src_meta --episode 1 --joint-space --out openarm_ep1_raw.mp4
# No dataset handy? Smoke-test with a synthetic wave (add --urdf to also exercise the kinematics):
LD_LIBRARY_PATH=$CONDA_PREFIX/lib MUJOCO_GL=egl \
python -m examples.openarm.render_episode --demo --joint-space --out openarm_demo.mp4
```
Useful flags: `--fps` (default 30), `--width` / `--height` (default 960×720), `--mjcf` to point
at an explicit model file, `--urdf` / `--ee-frame` for the kinematics, `--joint-space` to skip it.
## Troubleshooting
| Symptom | Fix |
| ------------------------------------- | ---------------------------------------------------------------------- |
| `Could not find the OpenArm v1 MJCF` | Pass `--mjcf`, set `$OPENARM_MJCF`, or install/clone `openarm_mujoco`. |
| `End-effector replay needs the OpenArm URDF` | Pass `--urdf` / set `$OPENARM_URDF`, or use `--joint-space`. |
| Large FK→IK round-trip error reported | Wrong `--ee-frame` link name, or the URDF doesn't match the recording. |
| `libEGL`/`GLEW` / blank window errors | Ensure `MUJOCO_GL=egl` and `LD_LIBRARY_PATH=$CONDA_PREFIX/lib`. |
| Wrists look mirrored / flipped | You are on the v2 model; switch to **v1**. |
| `KeyError: 'observation.state'` | Dataset isn't in the expected 16-D bimanual layout. |
+68
View File
@@ -0,0 +1,68 @@
#!/usr/bin/env python
"""Grab one frame from each rollout camera and save labeled PNGs to confirm mapping.
Uses the *same* V4L2 + MJPG + resolution settings as the rollout command so what you
see is what the policy sees. Run in the lerobot312 env (has cv2):
python examples/openarm/check_cameras.py
Then open the printed PNG paths and check that left_wrist / right_wrist are correct.
"""
from __future__ import annotations
import os
import cv2
# label -> (device, width, height) — matches the rollout --robot.cameras block
CAMERAS = {
"left_wrist": ("/dev/video8", 1280, 720),
"base": ("/dev/video6", 640, 480),
"right_wrist": ("/dev/video4", 1280, 720),
}
OUT_DIR = os.path.dirname(os.path.abspath(__file__))
WARMUP_FRAMES = 12 # let auto-exposure/white-balance settle
def fourcc_str(cap) -> str:
v = int(cap.get(cv2.CAP_PROP_FOURCC))
return "".join(chr((v >> (8 * i)) & 0xFF) for i in range(4))
def main() -> None:
for label, (dev, w, h) in CAMERAS.items():
cap = cv2.VideoCapture(dev, cv2.CAP_V4L2)
if not cap.isOpened():
print(f"[{label}] {dev}: FAILED to open")
continue
cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*"MJPG"))
cap.set(cv2.CAP_PROP_FRAME_WIDTH, w)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, h)
cap.set(cv2.CAP_PROP_FPS, 30)
frame = None
for _ in range(WARMUP_FRAMES):
ok, f = cap.read()
if ok:
frame = f
aw = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
ah = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fc = fourcc_str(cap)
cap.release()
if frame is None:
print(f"[{label}] {dev}: opened ({aw}x{ah} {fc}) but no frame read")
continue
# Burn the label into the image so the saved file is self-identifying.
cv2.putText(frame, f"{label} {dev}", (20, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.2,
(0, 255, 0), 3, cv2.LINE_AA)
out = os.path.join(OUT_DIR, f"cam_{label}.png")
cv2.imwrite(out, frame)
print(f"[{label}] {dev}: {aw}x{ah} {fc} -> {out}")
if __name__ == "__main__":
main()
+607
View File
@@ -0,0 +1,607 @@
<?xml version="1.0" ?>
<!-- =================================================================================== -->
<!-- | This document was autogenerated by xacro from openarm_v10.urdf.xacro | -->
<!-- | EDITING THIS FILE BY HAND IS NOT RECOMMENDED | -->
<!-- =================================================================================== -->
<robot name="openarm">
<link name="world"/>
<joint name="openarm_body_world_joint" type="fixed">
<parent link="world"/>
<child link="openarm_body_link0"/>
<origin rpy="0 0 0" xyz="0 0 0"/>
</joint>
<link name="openarm_body_link0">
<visual name="openarm_body_link0_visual">
<origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/body/visual/body_link0.dae" scale="0.001 0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_body_link0_collision">
<origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/body/collision/body_link0_symp.stl" scale="0.001 0.001 0.001"/>
<!-- <mesh filename="package://openarm_description/meshes/arm/${arm_type}/collision/${name}.stl" scale="0.001 0.001 0.001" /> -->
</geometry>
</collision>
<inertial>
<origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
<mass value="13.89"/>
<inertia ixx="1.653" ixy="0.0" ixz="0.0" iyy="1.653" iyz="0.0" izz="0.051"/>
</inertial>
</link>
<joint name="openarm_left_openarm_body_link0_joint" type="fixed">
<parent link="openarm_body_link0"/>
<child link="openarm_left_link0"/>
<origin rpy="-1.5708 0 0" xyz="0 0.0410356 0.746983"/>
</joint>
<link name="openarm_left_link0">
<visual name="openarm_left_link0_visual">
<origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/visual/link0.dae" scale="0.001 -0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_left_link0_collision">
<origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/collision/link0_symp.stl" scale="0.001 -0.001 0.001"/>
<!-- <mesh filename="package://openarm_description/meshes/arm/${arm_type}/collision/${name}.stl" scale="0.001 0.001 0.001" /> -->
</geometry>
</collision>
<inertial>
<origin rpy="0.0 0.0 0.0" xyz="-0.0009483362816297526 -0.0001580207020448382 0.03076860287587199"/>
<mass value="1.1432284943239561"/>
<inertia ixx="0.001128" ixy="4e-06" ixz="-3.3e-05" iyy="0.000962" iyz="7e-06" izz="0.00147"/>
</inertial>
</link>
<link name="openarm_left_link1">
<visual name="openarm_left_link1_visual">
<origin rpy="0.0 0.0 0.0" xyz="-0.0 0.0 -0.0625"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/visual/link1.dae" scale="0.001 -0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_left_link1_collision">
<origin rpy="0.0 0.0 0.0" xyz="-0.0 0.0 -0.0625"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/collision/link1_symp.stl" scale="0.001 -0.001 0.001"/>
<!-- <mesh filename="package://openarm_description/meshes/arm/${arm_type}/collision/${name}.stl" scale="0.001 0.001 0.001" /> -->
</geometry>
</collision>
<inertial>
<origin rpy="0.0 0.0 0.0" xyz="0.0011467657911800769 -3.319987657026362e-05 0.05395284380736254"/>
<mass value="1.1416684646202298"/>
<inertia ixx="0.001567" ixy="1e-06" ixz="-2.9e-05" iyy="0.001273" iyz="-1e-06" izz="0.001016"/>
</inertial>
</link>
<joint name="openarm_left_joint1" type="revolute">
<origin rpy="0 0 0" xyz="0.0 0.0 0.0625"/>
<parent link="openarm_left_link0"/>
<child link="openarm_left_link1"/>
<axis xyz="0 0 1"/>
<limit effort="40" lower="-3.490659" upper="1.3962629999999998" velocity="16.754666"/>
</joint>
<link name="openarm_left_link2">
<visual name="openarm_left_link2_visual">
<origin rpy="0.0 0.0 0.0" xyz="0.0301 0.0 -0.1225"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/visual/link2.dae" scale="0.001 -0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_left_link2_collision">
<origin rpy="0.0 0.0 0.0" xyz="0.0301 0.0 -0.1225"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/collision/link2_symp.stl" scale="0.001 -0.001 0.001"/>
<!-- <mesh filename="package://openarm_description/meshes/arm/${arm_type}/collision/${name}.stl" scale="0.001 0.001 0.001" /> -->
</geometry>
</collision>
<inertial>
<origin rpy="0.0 0.0 0.0" xyz="0.00839629182351943 2.0145102027597523e-08 0.03256649300522363"/>
<mass value="0.2775092746011571"/>
<inertia ixx="0.000359" ixy="-1e-06" ixz="-0.000109" iyy="0.000376" iyz="-1e-06" izz="0.000232"/>
</inertial>
</link>
<joint name="openarm_left_joint2" type="revolute">
<origin rpy="-1.57079632679 0 0" xyz="-0.0301 0.0 0.06"/>
<parent link="openarm_left_link1"/>
<child link="openarm_left_link2"/>
<axis xyz="-1 0 0"/>
<limit effort="40" lower="-3.3161253267948965" upper="0.17453267320510335" velocity="16.754666"/>
</joint>
<link name="openarm_left_link3">
<visual name="openarm_left_link3_visual">
<origin rpy="0.0 0.0 0.0" xyz="-0.0 -0.0 -0.18875"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/visual/link3.dae" scale="0.001 0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_left_link3_collision">
<origin rpy="0.0 0.0 0.0" xyz="-0.0 -0.0 -0.18875"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/collision/link3_symp.stl" scale="0.001 0.001 0.001"/>
<!-- <mesh filename="package://openarm_description/meshes/arm/${arm_type}/collision/${name}.stl" scale="0.001 0.001 0.001" /> -->
</geometry>
</collision>
<inertial>
<origin rpy="0.0 0.0 0.0" xyz="-0.002104752099628911 0.0005549085042607548 0.09047470545721961"/>
<mass value="1.073863338202347"/>
<inertia ixx="0.004372" ixy="1e-06" ixz="1.1e-05" iyy="0.004319" iyz="-3.6e-05" izz="0.000661"/>
</inertial>
</link>
<joint name="openarm_left_joint3" type="revolute">
<origin rpy="0 0 0" xyz="0.0301 0.0 0.06625"/>
<parent link="openarm_left_link2"/>
<child link="openarm_left_link3"/>
<axis xyz="0 0 1"/>
<limit effort="27" lower="-1.570796" upper="1.570796" velocity="5.445426"/>
</joint>
<link name="openarm_left_link4">
<visual name="openarm_left_link4_visual">
<origin rpy="0.0 0.0 0.0" xyz="0.0 -0.0315 -0.3425"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/visual/link4.dae" scale="0.001 0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_left_link4_collision">
<origin rpy="0.0 0.0 0.0" xyz="0.0 -0.0315 -0.3425"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/collision/link4_symp.stl" scale="0.001 0.001 0.001"/>
<!-- <mesh filename="package://openarm_description/meshes/arm/${arm_type}/collision/${name}.stl" scale="0.001 0.001 0.001" /> -->
</geometry>
</collision>
<inertial>
<origin rpy="0.0 0.0 0.0" xyz="-0.0029006831074562967 -0.03030575826634669 0.06339637422196209"/>
<mass value="0.6348534566833373"/>
<inertia ixx="0.000623" ixy="-1e-06" ixz="-1.9e-05" iyy="0.000511" iyz="3.8e-05" izz="0.000334"/>
</inertial>
</link>
<joint name="openarm_left_joint4" type="revolute">
<origin rpy="0 0 0" xyz="-0 0.0415354 0.202733"/>
<parent link="openarm_left_link3"/>
<child link="openarm_left_link4"/>
<axis xyz="0 1 0"/>
<limit effort="27" lower="0.0" upper="2.443461" velocity="5.445426"/>
</joint>
<link name="openarm_left_link5">
<visual name="openarm_left_link5_visual">
<origin rpy="0.0 0.0 0.0" xyz="-0.0 -0.0 -0.438"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/visual/link5.dae" scale="0.001 -0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_left_link5_collision">
<origin rpy="0.0 0.0 0.0" xyz="-0.0 -0.0 -0.438"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/collision/link5_symp.stl" scale="0.001 -0.001 0.001"/>
<!-- <mesh filename="package://openarm_description/meshes/arm/${arm_type}/collision/${name}.stl" scale="0.001 0.001 0.001" /> -->
</geometry>
</collision>
<inertial>
<origin rpy="0.0 0.0 0.0" xyz="-0.003049665024221911 -0.0008866902457326625 0.043079803024980934"/>
<mass value="0.6156588026168502"/>
<inertia ixx="0.000423" ixy="8e-06" ixz="6e-06" iyy="0.000445" iyz="6e-06" izz="0.000324"/>
</inertial>
</link>
<joint name="openarm_left_joint5" type="revolute">
<origin rpy="0 0 0" xyz="0.0 -0.0315 0.0955"/>
<parent link="openarm_left_link4"/>
<child link="openarm_left_link5"/>
<axis xyz="0 0 1"/>
<limit effort="7" lower="-1.570796" upper="1.570796" velocity="20.943946"/>
</joint>
<link name="openarm_left_link6">
<visual name="openarm_left_link6_visual">
<origin rpy="0.0 0.0 0.0" xyz="-0.0375 -0.0 -0.5585"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/visual/link6.dae" scale="0.001 -0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_left_link6_collision">
<origin rpy="0.0 0.0 0.0" xyz="-0.0375 -0.0 -0.5585"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/collision/link6_symp.stl" scale="0.001 -0.001 0.001"/>
<!-- <mesh filename="package://openarm_description/meshes/arm/${arm_type}/collision/${name}.stl" scale="0.001 0.001 0.001" /> -->
</geometry>
</collision>
<inertial>
<origin rpy="0.0 0.0 0.0" xyz="-0.037136587005447405 -0.00033230528343419053 -9.498374522309838e-05"/>
<mass value="0.475202773187987"/>
<inertia ixx="0.000143" ixy="-1e-06" ixz="1e-06" iyy="0.000157" iyz="-1e-06" izz="0.000159"/>
</inertial>
</link>
<joint name="openarm_left_joint6" type="revolute">
<origin rpy="0 0 0" xyz="0.0375 0.0 0.1205"/>
<parent link="openarm_left_link5"/>
<child link="openarm_left_link6"/>
<axis xyz="1 0 0"/>
<limit effort="7" lower="-0.785398" upper="0.785398" velocity="20.943946"/>
</joint>
<link name="openarm_left_link7">
<visual name="openarm_left_link7_visual">
<origin rpy="0.0 0.0 0.0" xyz="0.0 -0.0 -0.5585"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/visual/link7.dae" scale="0.001 -0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_left_link7_collision">
<origin rpy="0.0 0.0 0.0" xyz="0.0 -0.0 -0.5585"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/collision/link7_symp.stl" scale="0.001 -0.001 0.001"/>
<!-- <mesh filename="package://openarm_description/meshes/arm/${arm_type}/collision/${name}.stl" scale="0.001 0.001 0.001" /> -->
</geometry>
</collision>
<inertial>
<origin rpy="0.0 0.0 0.0" xyz="6.875510271106056e-05 -0.01266175250761268 0.06951945409987448"/>
<mass value="0.4659771327380578"/>
<inertia ixx="0.000639" ixy="-1e-06" ixz="1e-06" iyy="0.000497" iyz="-8.9e-05" izz="0.000342"/>
</inertial>
</link>
<joint name="openarm_left_joint7" type="revolute">
<origin rpy="0 0 0" xyz="-0.0375 0.0 0.0"/>
<parent link="openarm_left_link6"/>
<child link="openarm_left_link7"/>
<axis xyz="0 -1 0"/>
<limit effort="7" lower="-1.570796" upper="1.570796" velocity="20.943946"/>
</joint>
<joint name="openarm_right_openarm_body_link0_joint" type="fixed">
<parent link="openarm_body_link0"/>
<child link="openarm_right_link0"/>
<origin rpy="1.5708 0 0" xyz="0 -0.0209647 0.746983"/>
</joint>
<link name="openarm_right_link0">
<visual name="openarm_right_link0_visual">
<origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/visual/link0.dae" scale="0.001 0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_right_link0_collision">
<origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/collision/link0_symp.stl" scale="0.001 0.001 0.001"/>
<!-- <mesh filename="package://openarm_description/meshes/arm/${arm_type}/collision/${name}.stl" scale="0.001 0.001 0.001" /> -->
</geometry>
</collision>
<inertial>
<origin rpy="0.0 0.0 0.0" xyz="-0.0009483362816297526 0.0001580207020448382 0.03076860287587199"/>
<mass value="1.1432284943239561"/>
<inertia ixx="0.001128" ixy="-4e-06" ixz="-3.3e-05" iyy="0.000962" iyz="-7e-06" izz="0.00147"/>
</inertial>
</link>
<link name="openarm_right_link1">
<visual name="openarm_right_link1_visual">
<origin rpy="0.0 0.0 0.0" xyz="-0.0 0.0 -0.0625"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/visual/link1.dae" scale="0.001 0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_right_link1_collision">
<origin rpy="0.0 0.0 0.0" xyz="-0.0 0.0 -0.0625"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/collision/link1_symp.stl" scale="0.001 0.001 0.001"/>
<!-- <mesh filename="package://openarm_description/meshes/arm/${arm_type}/collision/${name}.stl" scale="0.001 0.001 0.001" /> -->
</geometry>
</collision>
<inertial>
<origin rpy="0.0 0.0 0.0" xyz="0.0011467657911800769 3.319987657026362e-05 0.05395284380736254"/>
<mass value="1.1416684646202298"/>
<inertia ixx="0.001567" ixy="-1e-06" ixz="-2.9e-05" iyy="0.001273" iyz="1e-06" izz="0.001016"/>
</inertial>
</link>
<joint name="openarm_right_joint1" type="revolute">
<origin rpy="0 0 0" xyz="0.0 0.0 0.0625"/>
<parent link="openarm_right_link0"/>
<child link="openarm_right_link1"/>
<axis xyz="0 0 1"/>
<limit effort="40" lower="-1.396263" upper="3.490659" velocity="16.754666"/>
</joint>
<link name="openarm_right_link2">
<visual name="openarm_right_link2_visual">
<origin rpy="0.0 0.0 0.0" xyz="0.0301 0.0 -0.1225"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/visual/link2.dae" scale="0.001 0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_right_link2_collision">
<origin rpy="0.0 0.0 0.0" xyz="0.0301 0.0 -0.1225"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/collision/link2_symp.stl" scale="0.001 0.001 0.001"/>
<!-- <mesh filename="package://openarm_description/meshes/arm/${arm_type}/collision/${name}.stl" scale="0.001 0.001 0.001" /> -->
</geometry>
</collision>
<inertial>
<origin rpy="0.0 0.0 0.0" xyz="0.00839629182351943 -2.0145102027597523e-08 0.03256649300522363"/>
<mass value="0.2775092746011571"/>
<inertia ixx="0.000359" ixy="1e-06" ixz="-0.000109" iyy="0.000376" iyz="1e-06" izz="0.000232"/>
</inertial>
</link>
<joint name="openarm_right_joint2" type="revolute">
<origin rpy="1.57079632679 0 0" xyz="-0.0301 0.0 0.06"/>
<parent link="openarm_right_link1"/>
<child link="openarm_right_link2"/>
<axis xyz="-1 0 0"/>
<limit effort="40" lower="-0.17453267320510335" upper="3.3161253267948965" velocity="16.754666"/>
</joint>
<link name="openarm_right_link3">
<visual name="openarm_right_link3_visual">
<origin rpy="0.0 0.0 0.0" xyz="-0.0 -0.0 -0.18875"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/visual/link3.dae" scale="0.001 0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_right_link3_collision">
<origin rpy="0.0 0.0 0.0" xyz="-0.0 -0.0 -0.18875"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/collision/link3_symp.stl" scale="0.001 0.001 0.001"/>
<!-- <mesh filename="package://openarm_description/meshes/arm/${arm_type}/collision/${name}.stl" scale="0.001 0.001 0.001" /> -->
</geometry>
</collision>
<inertial>
<origin rpy="0.0 0.0 0.0" xyz="-0.002104752099628911 0.0005549085042607548 0.09047470545721961"/>
<mass value="1.073863338202347"/>
<inertia ixx="0.004372" ixy="1e-06" ixz="1.1e-05" iyy="0.004319" iyz="-3.6e-05" izz="0.000661"/>
</inertial>
</link>
<joint name="openarm_right_joint3" type="revolute">
<origin rpy="0 0 0" xyz="0.0301 0.0 0.06625"/>
<parent link="openarm_right_link2"/>
<child link="openarm_right_link3"/>
<axis xyz="0 0 1"/>
<limit effort="27" lower="-1.570796" upper="1.570796" velocity="5.445426"/>
</joint>
<link name="openarm_right_link4">
<visual name="openarm_right_link4_visual">
<origin rpy="0.0 0.0 0.0" xyz="0.0 -0.0315 -0.3425"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/visual/link4.dae" scale="0.001 0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_right_link4_collision">
<origin rpy="0.0 0.0 0.0" xyz="0.0 -0.0315 -0.3425"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/collision/link4_symp.stl" scale="0.001 0.001 0.001"/>
<!-- <mesh filename="package://openarm_description/meshes/arm/${arm_type}/collision/${name}.stl" scale="0.001 0.001 0.001" /> -->
</geometry>
</collision>
<inertial>
<origin rpy="0.0 0.0 0.0" xyz="-0.0029006831074562967 -0.03030575826634669 0.06339637422196209"/>
<mass value="0.6348534566833373"/>
<inertia ixx="0.000623" ixy="-1e-06" ixz="-1.9e-05" iyy="0.000511" iyz="3.8e-05" izz="0.000334"/>
</inertial>
</link>
<joint name="openarm_right_joint4" type="revolute">
<origin rpy="0 0 0" xyz="-0 0.0415354 0.202733"/>
<parent link="openarm_right_link3"/>
<child link="openarm_right_link4"/>
<axis xyz="0 1 0"/>
<limit effort="27" lower="0.0" upper="2.443461" velocity="5.445426"/>
</joint>
<link name="openarm_right_link5">
<visual name="openarm_right_link5_visual">
<origin rpy="0.0 0.0 0.0" xyz="-0.0 -0.0 -0.438"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/visual/link5.dae" scale="0.001 0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_right_link5_collision">
<origin rpy="0.0 0.0 0.0" xyz="-0.0 -0.0 -0.438"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/collision/link5_symp.stl" scale="0.001 0.001 0.001"/>
<!-- <mesh filename="package://openarm_description/meshes/arm/${arm_type}/collision/${name}.stl" scale="0.001 0.001 0.001" /> -->
</geometry>
</collision>
<inertial>
<origin rpy="0.0 0.0 0.0" xyz="-0.003049665024221911 0.0008866902457326625 0.043079803024980934"/>
<mass value="0.6156588026168502"/>
<inertia ixx="0.000423" ixy="-8e-06" ixz="6e-06" iyy="0.000445" iyz="-6e-06" izz="0.000324"/>
</inertial>
</link>
<joint name="openarm_right_joint5" type="revolute">
<origin rpy="0 0 0" xyz="0.0 -0.0315 0.0955"/>
<parent link="openarm_right_link4"/>
<child link="openarm_right_link5"/>
<axis xyz="0 0 1"/>
<limit effort="7" lower="-1.570796" upper="1.570796" velocity="20.943946"/>
</joint>
<link name="openarm_right_link6">
<visual name="openarm_right_link6_visual">
<origin rpy="0.0 0.0 0.0" xyz="-0.0375 -0.0 -0.5585"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/visual/link6.dae" scale="0.001 0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_right_link6_collision">
<origin rpy="0.0 0.0 0.0" xyz="-0.0375 -0.0 -0.5585"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/collision/link6_symp.stl" scale="0.001 0.001 0.001"/>
<!-- <mesh filename="package://openarm_description/meshes/arm/${arm_type}/collision/${name}.stl" scale="0.001 0.001 0.001" /> -->
</geometry>
</collision>
<inertial>
<origin rpy="0.0 0.0 0.0" xyz="-0.037136587005447405 0.00033230528343419053 -9.498374522309838e-05"/>
<mass value="0.475202773187987"/>
<inertia ixx="0.000143" ixy="1e-06" ixz="1e-06" iyy="0.000157" iyz="1e-06" izz="0.000159"/>
</inertial>
</link>
<joint name="openarm_right_joint6" type="revolute">
<origin rpy="0 0 0" xyz="0.0375 0.0 0.1205"/>
<parent link="openarm_right_link5"/>
<child link="openarm_right_link6"/>
<axis xyz="1 0 0"/>
<limit effort="7" lower="-0.785398" upper="0.785398" velocity="20.943946"/>
</joint>
<link name="openarm_right_link7">
<visual name="openarm_right_link7_visual">
<origin rpy="0.0 0.0 0.0" xyz="0.0 -0.0 -0.5585"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/visual/link7.dae" scale="0.001 0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_right_link7_collision">
<origin rpy="0.0 0.0 0.0" xyz="0.0 -0.0 -0.5585"/>
<geometry>
<mesh filename="package://openarm_description/assets/robot/openarm_v1.0/mesh/arm/collision/link7_symp.stl" scale="0.001 0.001 0.001"/>
<!-- <mesh filename="package://openarm_description/meshes/arm/${arm_type}/collision/${name}.stl" scale="0.001 0.001 0.001" /> -->
</geometry>
</collision>
<inertial>
<origin rpy="0.0 0.0 0.0" xyz="6.875510271106056e-05 0.01266175250761268 0.06951945409987448"/>
<mass value="0.4659771327380578"/>
<inertia ixx="0.000639" ixy="1e-06" ixz="1e-06" iyy="0.000497" iyz="8.9e-05" izz="0.000342"/>
</inertial>
</link>
<joint name="openarm_right_joint7" type="revolute">
<origin rpy="0 0 0" xyz="-0.0375 0.0 0.0"/>
<parent link="openarm_right_link6"/>
<child link="openarm_right_link7"/>
<axis xyz="0 1 0"/>
<limit effort="7" lower="-1.570796" upper="1.570796" velocity="20.943946"/>
</joint>
<link name="openarm_left_hand_tcp"/>
<joint name="openarm_left_hand_tcp_joint" type="fixed">
<origin rpy="0 0 0" xyz="0 0 0"/>
<parent link="openarm_left_link7"/>
<child link="openarm_left_hand_tcp"/>
</joint>
<link name="openarm_left_left_finger">
<visual name="openarm_left_left_finger_visual">
<origin rpy="0.0 0.0 0.0" xyz="0.0 -0.05 -0.673001"/>
<geometry>
<mesh filename="package://openarm_description/assets/end_effector/parallel_link/meshes/visual/finger.dae" scale="0.001 0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_left_left_finger_collision">
<origin rpy="0.0 0.0 0.0" xyz="0.0 -0.05 -0.673001"/>
<geometry>
<mesh filename="package://openarm_description/assets/end_effector/parallel_link/meshes/collision/finger.stl" scale="0.001 0.001 0.001"/>
</geometry>
</collision>
<inertial>
<origin rpy="0 0 0" xyz="0.0064528 0.01702 0.0219685"/>
<mass value="0.03602545343277134"/>
<inertia ixx="2.3749999999999997e-06" ixy="1e-06" ixz="1e-06" iyy="2.3749999999999997e-06" iyz="1e-06" izz="7.5e-07"/>
</inertial>
</link>
<link name="openarm_left_right_finger">
<visual name="openarm_left_right_finger_visual">
<origin rpy="0.0 0.0 0.0" xyz="0.0 0.05 -0.673001"/>
<geometry>
<mesh filename="package://openarm_description/assets/end_effector/parallel_link/meshes/visual/finger.dae" scale="0.001 -0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_left_right_finger_collision">
<origin rpy="0.0 0.0 0.0" xyz="0.0 0.05 -0.673001"/>
<geometry>
<mesh filename="package://openarm_description/assets/end_effector/parallel_link/meshes/collision/finger.stl" scale="0.001 -0.001 0.001"/>
</geometry>
</collision>
<inertial>
<origin rpy="0 0 0" xyz="0.0064528 -0.01702 0.0219685"/>
<mass value="0.03602545343277134"/>
<inertia ixx="2.3749999999999997e-06" ixy="1e-06" ixz="1e-06" iyy="2.3749999999999997e-06" iyz="1e-06" izz="7.5e-07"/>
</inertial>
</link>
<joint name="openarm_left_finger_joint1" type="prismatic">
<parent link="openarm_left_link7"/>
<child link="openarm_left_right_finger"/>
<origin rpy="0 0 0" xyz="0.0 0.0 0.1025"/>
<axis xyz="0 -1 0"/>
<limit effort="333" lower="0.0" upper="0.044" velocity="10.0"/>
</joint>
<joint name="openarm_left_finger_joint2" type="prismatic">
<parent link="openarm_left_link7"/>
<child link="openarm_left_left_finger"/>
<origin rpy="0 0 0" xyz="0.0 -0.0 0.1025"/>
<axis xyz="0 1 0"/>
<limit effort="333" lower="0.0" upper="0.044" velocity="10.0"/>
<mimic joint="openarm_left_finger_joint1"/>
</joint>
<!-- <joint name="${ee_prefix}finger_joint1" type="prismatic">
<parent link="${connected_to}" /> <child link="${ee_prefix}right_finger" />
<origin xyz="0 -0.006 0.115" rpy="0 0 0" />
<axis xyz="0 -1 0" />
<limit effort="333" lower="0.0" upper="0.044" velocity="10.0" />
</joint>
<joint name="${ee_prefix}finger_joint2" type="prismatic">
<parent link="${connected_to}" /> <child link="${ee_prefix}left_finger" />
<origin xyz="0 0.006 0.115" rpy="0 0 0" />
<axis xyz="0 1 0" />
<limit effort="333" lower="0.0" upper="0.044" velocity="10.0" />
<mimic joint="${ee_prefix}finger_joint1" />
</joint> -->
<link name="openarm_right_hand_tcp"/>
<joint name="openarm_right_hand_tcp_joint" type="fixed">
<origin rpy="0 0 0" xyz="0 0 0"/>
<parent link="openarm_right_link7"/>
<child link="openarm_right_hand_tcp"/>
</joint>
<link name="openarm_right_left_finger">
<visual name="openarm_right_left_finger_visual">
<origin rpy="0.0 0.0 0.0" xyz="0.0 -0.05 -0.673001"/>
<geometry>
<mesh filename="package://openarm_description/assets/end_effector/parallel_link/meshes/visual/finger.dae" scale="0.001 0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_right_left_finger_collision">
<origin rpy="0.0 0.0 0.0" xyz="0.0 -0.05 -0.673001"/>
<geometry>
<mesh filename="package://openarm_description/assets/end_effector/parallel_link/meshes/collision/finger.stl" scale="0.001 0.001 0.001"/>
</geometry>
</collision>
<inertial>
<origin rpy="0 0 0" xyz="0.0064528 0.01702 0.0219685"/>
<mass value="0.03602545343277134"/>
<inertia ixx="2.3749999999999997e-06" ixy="1e-06" ixz="1e-06" iyy="2.3749999999999997e-06" iyz="1e-06" izz="7.5e-07"/>
</inertial>
</link>
<link name="openarm_right_right_finger">
<visual name="openarm_right_right_finger_visual">
<origin rpy="0.0 0.0 0.0" xyz="0.0 0.05 -0.673001"/>
<geometry>
<mesh filename="package://openarm_description/assets/end_effector/parallel_link/meshes/visual/finger.dae" scale="0.001 -0.001 0.001"/>
</geometry>
</visual>
<collision name="openarm_right_right_finger_collision">
<origin rpy="0.0 0.0 0.0" xyz="0.0 0.05 -0.673001"/>
<geometry>
<mesh filename="package://openarm_description/assets/end_effector/parallel_link/meshes/collision/finger.stl" scale="0.001 -0.001 0.001"/>
</geometry>
</collision>
<inertial>
<origin rpy="0 0 0" xyz="0.0064528 -0.01702 0.0219685"/>
<mass value="0.03602545343277134"/>
<inertia ixx="2.3749999999999997e-06" ixy="1e-06" ixz="1e-06" iyy="2.3749999999999997e-06" iyz="1e-06" izz="7.5e-07"/>
</inertial>
</link>
<joint name="openarm_right_finger_joint1" type="prismatic">
<parent link="openarm_right_link7"/>
<child link="openarm_right_right_finger"/>
<origin rpy="0 0 0" xyz="0.0 0.0 0.1025"/>
<axis xyz="0 -1 0"/>
<limit effort="333" lower="0.0" upper="0.044" velocity="10.0"/>
</joint>
<joint name="openarm_right_finger_joint2" type="prismatic">
<parent link="openarm_right_link7"/>
<child link="openarm_right_left_finger"/>
<origin rpy="0 0 0" xyz="0.0 -0.0 0.1025"/>
<axis xyz="0 1 0"/>
<limit effort="333" lower="0.0" upper="0.044" velocity="10.0"/>
<mimic joint="openarm_right_finger_joint1"/>
</joint>
<!-- <joint name="${ee_prefix}finger_joint1" type="prismatic">
<parent link="${connected_to}" /> <child link="${ee_prefix}right_finger" />
<origin xyz="0 -0.006 0.115" rpy="0 0 0" />
<axis xyz="0 -1 0" />
<limit effort="333" lower="0.0" upper="0.044" velocity="10.0" />
</joint>
<joint name="${ee_prefix}finger_joint2" type="prismatic">
<parent link="${connected_to}" /> <child link="${ee_prefix}left_finger" />
<origin xyz="0 0.006 0.115" rpy="0 0 0" />
<axis xyz="0 1 0" />
<limit effort="333" lower="0.0" upper="0.044" velocity="10.0" />
<mimic joint="${ee_prefix}finger_joint1" />
</joint> -->
</robot>
+344
View File
@@ -0,0 +1,344 @@
#!/usr/bin/env python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Replay a recorded bimanual-OpenArm episode into an mp4, in simulation.
This drives the official OpenArm MuJoCo model (``enactic/openarm_mujoco``, v1) from a
LeRobot dataset's recorded ``observation.state`` and renders a headless video.
By default the replay goes *through end-effector (Cartesian) space* rather than pushing the
recorded joint angles straight into the simulator: for every frame and every arm it runs
forward kinematics (recorded joints -> EE pose) and then inverse kinematics (EE pose ->
joints), and drives MuJoCo with the IK-recovered joints. This exercises the exact
``lerobot.model.RobotKinematics`` solver that ``OpenArmFollower.make_kinematics()`` builds,
so the rendered video is a visual sanity check of the OpenArm end-effector kinematics -- not
just of the raw recording. Pass ``--joint-space`` to bypass kinematics and replay the raw
recorded joints directly (the previous behaviour).
The ``observation.state`` is expected in the 16-D bimanual layout (degrees):
right_joint_1..7, right_gripper, left_joint_1..7, left_gripper
The single published OpenArm URDF is *bimanual*, so end-effector kinematics build one solver
per arm (right/left), each keyed to that arm's joints (``openarm_<side>_joint1..7``) and
tool-center frame (``openarm_<side>_hand_tcp``). Solving in the full bimanual frame is what
keeps both arms at their correct, matching heights.
Assets (both Apache-2.0, nothing vendored into LeRobot):
MuJoCo MJCF: https://github.com/enactic/openarm_mujoco (use the **v1** revision; v2 is a
different wrist hardware revision and will look sign-flipped on v1 recordings).
URDF (for RobotKinematics): https://github.com/enactic/openarm_description
Examples:
# replay episode 1 of a local LeRobot v3.0 dataset through EE kinematics (bimanual URDF)
LD_LIBRARY_PATH=$CONDA_PREFIX/lib MUJOCO_GL=egl python -m examples.openarm.render_episode \
--dataset data/folding_src_meta --episode 1 \
--urdf /path/to/openarm_bimanual.urdf \
--out openarm_ep1.mp4
# smoke test with a synthetic wave (no dataset; add --urdf to also exercise the kinematics)
LD_LIBRARY_PATH=$CONDA_PREFIX/lib MUJOCO_GL=egl python -m examples.openarm.render_episode \
--demo --out openarm_demo.mp4
# bypass kinematics and replay the raw recorded joints directly
LD_LIBRARY_PATH=$CONDA_PREFIX/lib MUJOCO_GL=egl python -m examples.openarm.render_episode \
--dataset data/folding_src_meta --episode 1 --joint-space --out openarm_ep1_raw.mp4
"""
from __future__ import annotations
import argparse
import os
import sys
from pathlib import Path
os.environ.setdefault("MUJOCO_GL", "egl") # headless GPU rendering
import numpy as np
# Policy/dataset state layout: 16-D, degrees.
POLICY_ORDER = [
*(f"right_joint_{i}" for i in range(1, 8)),
"right_gripper",
*(f"left_joint_{i}" for i in range(1, 8)),
"left_gripper",
]
# MuJoCo joint name for each of the 14 arm entries (grippers handled separately).
ARM_MAP = {
**{f"right_joint_{i}": f"openarm_right_joint{i}" for i in range(1, 8)},
**{f"left_joint_{i}": f"openarm_left_joint{i}" for i in range(1, 8)},
}
RIGHT_GRIPPER_IDX, LEFT_GRIPPER_IDX = 7, 15
GRIP_FULL_DEG = 65.0 # follower gripper limit magnitude -> fully open
# Per-arm slices into the 16-D state (7 arm joints each; grippers handled separately).
ARM_JOINT_SLICES = {"right": slice(0, 7), "left": slice(8, 15)}
# The only published OpenArm URDF is *bimanual*, so each arm has its own joint names and
# end-effector (tool-center-point) frame. Using the bimanual URDF -- rather than a single
# arm placed at the origin -- is what makes FK/IK return correct world poses for each arm
# (it encodes the shoulder-mount transform), so the two arms land at the right heights.
ARM_SIDE_JOINTS = {
side: [f"openarm_{side}_joint{i}" for i in range(1, 8)] for side in ("right", "left")
}
DEFAULT_EE_FRAMES = {side: f"openarm_{side}_hand_tcp" for side in ("right", "left")}
def locate_mjcf(explicit: str | None) -> str:
"""Resolve the OpenArm v1 bimanual MJCF path.
Priority: --mjcf arg, then $OPENARM_MJCF, then the file installed by the
``openarm_mujoco`` pip package under ``<prefix>/share/openarm_mujoco/v1``.
"""
if explicit:
return explicit
if os.environ.get("OPENARM_MJCF"):
return os.environ["OPENARM_MJCF"]
for prefix in (sys.prefix, os.environ.get("CONDA_PREFIX", "")):
if not prefix:
continue
cand = Path(prefix) / "share" / "openarm_mujoco" / "v1" / "openarm_bimanual.xml"
if cand.exists():
return str(cand)
raise SystemExit(
"Could not find the OpenArm v1 MJCF. Pass --mjcf /path/to/v1/openarm_bimanual.xml, "
"set $OPENARM_MJCF, or clone https://github.com/enactic/openarm_mujoco."
)
def locate_urdf(explicit: str | None) -> str | None:
"""Resolve the single-arm OpenArm URDF path for ``RobotKinematics`` (may be None)."""
if explicit:
return explicit
if os.environ.get("OPENARM_URDF"):
return os.environ["OPENARM_URDF"]
return None
def make_kinematics(urdf_path: str, ee_frames: dict[str, str]) -> dict:
"""Build one ``RobotKinematics`` EE solver per arm from the bimanual OpenArm URDF.
This mirrors what ``OpenArmFollower.make_kinematics()`` does per arm, but we construct
``RobotKinematics`` directly (instead of instantiating an ``OpenArmFollower``) so this
render-only example does not pull in the CAN/motor hardware stack. Each side gets its own
solver keyed to that arm's URDF joint names (``openarm_<side>_joint1..7``) and tool-center
frame, so forward/inverse kinematics resolve in the full bimanual (world) frame -- that is
what keeps the two arms at their correct, matching heights.
"""
from lerobot.model import RobotKinematics
return {
side: RobotKinematics(
urdf_path=urdf_path,
target_frame_name=ee_frames[side],
joint_names=ARM_SIDE_JOINTS[side],
)
for side in ("right", "left")
}
def ee_roundtrip(kins: dict, traj: np.ndarray) -> tuple[np.ndarray, float, float]:
"""Route a recorded joint trajectory through end-effector space, per arm.
For each frame and each arm, run forward kinematics (recorded joints -> EE pose) then
inverse kinematics (EE pose -> joints) with that arm's own solver, replacing the arm joints
with the IK-recovered ones. Grippers pass through unchanged. Returns the new trajectory plus
the mean joint and EE-position round-trip errors (a sanity check that the solver tracks the
recording).
"""
out = traj.copy()
joint_err = []
pos_err = []
n_arm = len(ARM_SIDE_JOINTS["right"])
for t in range(traj.shape[0]):
for side, sl in ARM_JOINT_SLICES.items():
kin = kins[side]
recorded = traj[t, sl].astype(np.float64)
ee_pose = kin.forward_kinematics(recorded)
recovered = kin.inverse_kinematics(recorded, ee_pose)[:n_arm]
out[t, sl] = recovered
joint_err.append(np.abs(recovered - recorded).mean())
# EE position after IK vs. the FK target, to catch non-converged solves.
pos_err.append(np.linalg.norm(kin.forward_kinematics(recovered)[:3, 3] - ee_pose[:3, 3]))
return out, float(np.mean(joint_err)), float(np.mean(pos_err))
def load_state_from_dataset(root: str, ep: int) -> np.ndarray:
"""Read one episode's recorded observation.state (N, 16; degrees) from a LeRobot v3.0 root."""
import pandas as pd
root = Path(root)
ep_meta = pd.read_parquet(root / "meta" / "episodes" / "chunk-000" / "file-000.parquet")
row = ep_meta[ep_meta["episode_index"] == ep].iloc[0]
a, b = int(row["dataset_from_index"]), int(row["dataset_to_index"])
dchunk, dfile = int(row["data/chunk_index"]), int(row["data/file_index"])
df = pd.read_parquet(root / "data" / f"chunk-{dchunk:03d}" / f"file-{dfile:03d}.parquet")
df = df[(df["index"] >= a) & (df["index"] < b)].sort_values("frame_index")
return np.stack(df["observation.state"].to_numpy()).astype(np.float32)
def encode_mp4(frames: list[np.ndarray], path: str, fps: int) -> None:
import av
h, w = frames[0].shape[:2]
container = av.open(path, mode="w")
stream = container.add_stream("libx264", rate=fps)
stream.width, stream.height, stream.pix_fmt = w, h, "yuv420p"
for f in frames:
frame = av.VideoFrame.from_ndarray(np.ascontiguousarray(f), format="rgb24")
for pkt in stream.encode(frame):
container.mux(pkt)
for pkt in stream.encode():
container.mux(pkt)
container.close()
def main() -> None:
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawTextHelpFormatter)
ap.add_argument("--dataset", default=None, help="LeRobot v3.0 dataset root (meta/ + data/)")
ap.add_argument("--episode", type=int, default=0)
ap.add_argument("--demo", action="store_true", help="drive a synthetic wave (no dataset)")
ap.add_argument("--mjcf", default=None, help="path to v1/openarm_bimanual.xml (see locate_mjcf)")
ap.add_argument(
"--urdf",
default=None,
help="bimanual OpenArm URDF for end-effector kinematics (or set $OPENARM_URDF)",
)
ap.add_argument(
"--right-ee-frame",
default=os.environ.get("OPENARM_RIGHT_EE_FRAME", DEFAULT_EE_FRAMES["right"]),
help="right-arm end-effector link name in the URDF (or set $OPENARM_RIGHT_EE_FRAME)",
)
ap.add_argument(
"--left-ee-frame",
default=os.environ.get("OPENARM_LEFT_EE_FRAME", DEFAULT_EE_FRAMES["left"]),
help="left-arm end-effector link name in the URDF (or set $OPENARM_LEFT_EE_FRAME)",
)
ap.add_argument(
"--joint-space",
action="store_true",
help="bypass kinematics and replay the raw recorded joints directly",
)
ap.add_argument("--out", default="openarm_episode.mp4")
ap.add_argument("--fps", type=int, default=30)
ap.add_argument("--width", type=int, default=960)
ap.add_argument("--height", type=int, default=720)
args = ap.parse_args()
import mujoco
model = mujoco.MjModel.from_xml_path(locate_mjcf(args.mjcf))
model.vis.global_.offwidth = max(model.vis.global_.offwidth, args.width)
model.vis.global_.offheight = max(model.vis.global_.offheight, args.height)
data = mujoco.MjData(model)
qadr = {
pk: int(model.jnt_qposadr[mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, mj)])
for pk, mj in ARM_MAP.items()
}
def finger_adr(names):
out = []
for nm in names:
jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, nm)
out.append((int(model.jnt_qposadr[jid]), model.jnt_range[jid].copy(), int(model.jnt_type[jid])))
return out
right_fingers = finger_adr(["openarm_right_finger_joint1", "openarm_right_finger_joint2"])
left_fingers = finger_adr(["openarm_left_finger_joint1", "openarm_left_finger_joint2"])
hinge_type = int(mujoco.mjtJoint.mjJNT_HINGE)
def finger_target(gripper_deg, rng, jtype):
opening = min(1.0, abs(gripper_deg) / GRIP_FULL_DEG) # 0=closed .. 1=open
if jtype == hinge_type: # hinge in radians; sign encodes side via range direction
lo, hi = rng
mag = np.deg2rad(min(abs(gripper_deg), GRIP_FULL_DEG))
return np.clip(-mag if lo < 0 else mag, lo, hi)
return rng[0] + opening * (rng[1] - rng[0]) # slide (v1): lo=closed .. hi=open
if args.demo:
n_frames = 120
traj = np.zeros((n_frames, 16), np.float32)
wave = 40.0 * np.sin(np.linspace(0, 2 * np.pi, n_frames))
for i, pk in enumerate(POLICY_ORDER):
if pk in qadr:
traj[:, i] = wave * (0.5 + 0.5 * (i % 3))
elif args.dataset:
traj = load_state_from_dataset(args.dataset, args.episode)
else:
raise SystemExit("provide --dataset <root> (with --episode) or --demo")
n_frames = traj.shape[0]
# Route the recorded joints through end-effector (Cartesian) space via the same solver
# OpenArmFollower.make_kinematics() builds, unless the user opted for raw joint replay.
urdf_path = None if args.joint_space else locate_urdf(args.urdf)
if args.joint_space:
print(f"driving {n_frames} frames (raw joint space)")
elif urdf_path is None:
raise SystemExit(
"End-effector replay needs the OpenArm URDF: pass --urdf /path/to/openarm.urdf "
"(and --ee-frame if the tip link differs), set $OPENARM_URDF, or use --joint-space "
"to replay the raw recorded joints. URDF: https://github.com/enactic/openarm_description"
)
else:
ee_frames = {"right": args.right_ee_frame, "left": args.left_ee_frame}
kins = make_kinematics(urdf_path, ee_frames)
traj, joint_err, pos_err = ee_roundtrip(kins, traj)
print(
f"driving {n_frames} frames (end-effector kinematics via frames "
f"right='{ee_frames['right']}', left='{ee_frames['left']}'; "
f"FK->IK round-trip: {joint_err:.3f}° mean joint error, {pos_err * 1e3:.2f} mm mean EE error)"
)
# Auto-frame the arms (exclude pedestal/world) from body positions at the mid pose.
mid = n_frames // 2
for i, pk in enumerate(POLICY_ORDER):
if pk in qadr:
data.qpos[qadr[pk]] = np.deg2rad(traj[mid, i])
mujoco.mj_forward(model, data)
arm_pts = [
data.xpos[b].copy()
for b in range(model.nbody)
if (mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, b) or "").startswith("openarm")
and "base" not in (mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, b) or "")
]
arm_pts = np.array(arm_pts) if arm_pts else data.xpos[1:]
lo, hi = arm_pts.min(0), arm_pts.max(0)
cam = mujoco.MjvCamera()
mujoco.mjv_defaultCamera(cam)
cam.azimuth, cam.elevation = 150.0, -20.0
cam.distance = max(0.8, float(np.linalg.norm(hi - lo)) * 1.3)
cam.lookat[:] = (lo + hi) / 2.0
renderer = mujoco.Renderer(model, height=args.height, width=args.width)
frames = []
for t in range(n_frames):
for i, pk in enumerate(POLICY_ORDER):
if pk in qadr:
data.qpos[qadr[pk]] = np.deg2rad(traj[t, i])
for adr, rng, jt in right_fingers:
data.qpos[adr] = finger_target(float(traj[t, RIGHT_GRIPPER_IDX]), rng, jt)
for adr, rng, jt in left_fingers:
data.qpos[adr] = finger_target(float(traj[t, LEFT_GRIPPER_IDX]), rng, jt)
mujoco.mj_forward(model, data)
renderer.update_scene(data, camera=cam)
frames.append(renderer.render())
renderer.close()
encode_mp4(frames, args.out, args.fps)
print(f"wrote {args.out} ({n_frames} frames @ {args.fps} fps, {args.width}x{args.height})")
if __name__ == "__main__":
main()
+127
View File
@@ -0,0 +1,127 @@
#!/usr/bin/env python
"""Replay a precomputed retargeted SHORT-arm trajectory on the real bimanual OpenArm.
The trajectory (N,16 degrees, POLICY order: right_joint_1..7,right_gripper,left_joint_1..7,
left_gripper) is produced offline by .precompute_short_traj.py as:
FK on the LONG model (recorded data) -> gripper-tip pose -> IK on the SHORT model,
clamped to the real per-arm joint limits. This is the deterministic, no-policy validation of
the morphology retarget before running the live policy rollout.
Run in the deployment env (lerobot312), with CAN up (can0/can1):
python examples/openarm/retarget_replay.py \
--traj openarm_ep1_short_retargeted.npy \
--left-port can1 --right-port can0 --id openarms \
--max-relative-target 8.0 --fps 30 --ramp-seconds 4.0
Use --dry-run first to print the ramp/first/last targets without touching the robot.
"""
from __future__ import annotations
import argparse
import time
import numpy as np
# POLICY column -> robot action key (bimanual, ".pos", degrees)
JMAP: list[tuple[str, int]] = (
[(f"right_joint_{i}.pos", i - 1) for i in range(1, 8)]
+ [("right_gripper.pos", 7)]
+ [(f"left_joint_{i}.pos", 8 + i - 1) for i in range(1, 8)]
+ [("left_gripper.pos", 15)]
)
def action_of(row: np.ndarray) -> dict[str, float]:
return {k: float(row[c]) for k, c in JMAP}
def present_row(obs: dict) -> np.ndarray:
row = np.zeros(16, dtype=np.float64)
for k, c in JMAP:
row[c] = float(obs[k])
return row
def build_robot(args):
from lerobot.robots.bi_openarm_follower import BiOpenArmFollower, BiOpenArmFollowerConfig
from lerobot.robots.openarm_follower import OpenArmFollowerConfigBase
common = dict(
can_interface="socketcan",
disable_torque_on_disconnect=True,
max_relative_target=args.max_relative_target,
cameras={},
)
cfg = BiOpenArmFollowerConfig(
id=args.id,
left_arm_config=OpenArmFollowerConfigBase(port=args.left_port, side="left", **common),
right_arm_config=OpenArmFollowerConfigBase(port=args.right_port, side="right", **common),
cameras={},
)
return BiOpenArmFollower(cfg)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--traj", required=True, help="(N,16) npy in degrees, POLICY order")
ap.add_argument("--left-port", default="can1")
ap.add_argument("--right-port", default="can0")
ap.add_argument("--id", default="openarms")
ap.add_argument("--max-relative-target", type=float, default=8.0)
ap.add_argument("--fps", type=float, default=30.0)
ap.add_argument("--ramp-seconds", type=float, default=4.0,
help="time to smoothly move from current pose to the first frame")
ap.add_argument("--start-index", type=int, default=0)
ap.add_argument("--end-index", type=int, default=-1)
ap.add_argument("--dry-run", action="store_true", help="print plan, do not connect")
args = ap.parse_args()
traj = np.load(args.traj).astype(np.float64)
assert traj.ndim == 2 and traj.shape[1] == 16, f"expected (N,16), got {traj.shape}"
end = traj.shape[0] if args.end_index < 0 else args.end_index
traj = traj[args.start_index:end]
n = traj.shape[0]
dt = 1.0 / args.fps
print(f"loaded {args.traj}: {n} frames @ {args.fps} fps (~{n*dt:.1f}s)")
print("first target:", np.round(traj[0], 1))
print("last target:", np.round(traj[-1], 1))
if args.dry_run:
print("[dry-run] not connecting.")
return
robot = build_robot(args)
print("connecting... (ensure CAN is up and arms are clear)")
robot.connect()
try:
# --- gentle ramp from current pose to first frame ---
present = present_row(robot.get_observation())
n_ramp = max(1, int(round(args.ramp_seconds * args.fps)))
print(f"ramping to first frame over {args.ramp_seconds:.1f}s ({n_ramp} steps)...")
for k in range(1, n_ramp + 1):
a = k / n_ramp
row = (1.0 - a) * present + a * traj[0]
t0 = time.perf_counter()
robot.send_action(action_of(row))
time.sleep(max(0.0, dt - (time.perf_counter() - t0)))
# --- stream the trajectory ---
print("replaying...")
start = time.perf_counter()
for t in range(n):
robot.send_action(action_of(traj[t]))
target = start + (t + 1) * dt
time.sleep(max(0.0, target - time.perf_counter()))
if t % 60 == 0:
print(f" frame {t}/{n} ({t/args.fps:.1f}s)")
print("done.")
except KeyboardInterrupt:
print("\ninterrupted by user.")
finally:
print("disconnecting (torque off)...")
robot.disconnect()
if __name__ == "__main__":
main()
+615
View File
@@ -0,0 +1,615 @@
#!/usr/bin/env python
"""lerobot-rollout with gold-standard MuJoCo EE retargeting at the robot boundary.
The policy was trained on the LONG arm (upper arm +5 cm). The real robot is the
SHORT (stock) arm. We bridge the morphology gap with the *exact* MuJoCo FK/IK that
produced the validated cyan/red overlay video, wrapping the raw robot so that:
* get_observation(): SHORT joints --FK(short)--> gripper-tip pose --IK(long)-->
LONG joints (the state the policy expects)
* send_action(): LONG joint targets --FK(long)--> pose --IK(short)-->
SHORT joint targets (clamped to the real per-arm limits)
Between those two boundaries everything (state, relative-action anchor, policy
output) lives consistently in LONG space, so no other part of the rollout stack
needs to change. We inject the wrapper by patching
``lerobot.rollout.context.make_robot_from_config`` and then hand off to the normal
``lerobot-rollout`` entry point, so every CLI flag behaves identically.
Usage: same args as ``lerobot-rollout``, e.g.
python examples/openarm/rollout_retarget.py \
--policy.path=/home/yope/Documents/sonic/data/folding_latest \
--robot.type=bi_openarm_follower --robot.id=openarms \
--robot.cameras='{ ... }' \
--robot.left_arm_config.port=can1 ... --task="Fold the T-shirt properly" \
--fps=30 --duration=2000 --device=cuda --display_data=true
Env toggles (optional):
RETARGET_OBS=0 disable observation retargeting (short->long)
RETARGET_ACT=0 disable action retargeting (long->short)
RETARGET_ITERS=25 IK iterations per tick (warm-started)
RETARGET_ITERS0=80 IK iterations on the very first tick (cold seed)
RETARGET_NULL_GAIN=0.3 nullspace bias: pull short joints toward long joints (0 = off,
EE-only). Higher keeps the elbow closer to the long pose.
Joint-space smoothing streamer (decouples motor rate from the slow control loop):
STREAM=1 enable the background smoothing streamer
STREAM_HZ=60 motor command rate of the streamer thread (Hz)
STREAM_SMOOTH_TIME=0.10 SmoothDamp time constant (s); larger = smoother/laggier
STREAM_MAX_SPEED=150 per-joint speed cap (deg/s); 0 disables the cap
"""
from __future__ import annotations
import importlib.util
import logging
import os
import threading
import time
os.environ.setdefault("MUJOCO_GL", "egl")
import mujoco
import numpy as np
logger = logging.getLogger("rollout_retarget")
_HERE = os.path.dirname(os.path.abspath(__file__))
_REPO = os.path.abspath(os.path.join(_HERE, "..", ".."))
def _load(name: str, path: str):
spec = importlib.util.spec_from_file_location(name, path)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
return mod
# Reuse the *validated* geometry + FK/IK helpers (same code that made the overlay video).
ov = _load("ov", os.path.join(_REPO, ".overlay_rest.py"))
rt = _load("rt", os.path.join(_REPO, ".roundtrip_overlay.py"))
from lerobot.robots.openarm_follower.config_openarm_follower import ( # noqa: E402
LEFT_DEFAULT_JOINTS_LIMITS,
RIGHT_DEFAULT_JOINTS_LIMITS,
)
SIDES = ("right", "left")
LIMITS = {"right": RIGHT_DEFAULT_JOINTS_LIMITS, "left": LEFT_DEFAULT_JOINTS_LIMITS}
# POLICY_ORDER index of each gripper in the 16-vector (right block 0..7, left block 8..15).
GRIPPER_IDX = {"right": 7, "left": 15}
def smooth_damp(
current: np.ndarray,
target: np.ndarray,
velocity: np.ndarray,
smooth_time: float,
dt: float,
max_speed: float = 0.0,
) -> tuple[np.ndarray, np.ndarray]:
"""Critically-damped 2nd-order smoothing toward ``target`` (vectorized SmoothDamp).
Gives C1-continuous position + velocity with no overshoot, and re-plans every tick,
so it degrades gracefully when targets arrive irregularly or late. Returns the new
(position, velocity). ``max_speed <= 0`` disables the per-joint speed cap.
"""
smooth_time = max(1e-4, smooth_time)
omega = 2.0 / smooth_time
x = omega * dt
exp = 1.0 / (1.0 + x + 0.48 * x * x + 0.235 * x * x * x)
change = current - target
original_to = target.copy()
if max_speed and max_speed > 0.0:
max_change = max_speed * smooth_time
change = np.clip(change, -max_change, max_change)
shifted_target = current - change
temp = (velocity + omega * change) * dt
velocity = (velocity - omega * temp) * exp
output = shifted_target + (change + temp) * exp
# Kill overshoot: if we crossed the original target, snap to it and match velocity.
overshoot = (original_to - current > 0.0) == (output > original_to)
output = np.where(overshoot, original_to, output)
velocity = np.where(overshoot, (output - original_to) / dt, velocity)
return output, velocity
class RetargetRobot:
"""Boundary wrapper that retargets state (short->long) and actions (long->short).
All non-overridden attributes/methods proxy transparently to the wrapped robot
(``connect``, ``disconnect``, ``observation_features``, ``cameras``, ...), so the
rollout stack treats it exactly like the underlying ``bi_openarm_follower``.
"""
def __init__(
self,
robot,
obs_iters: int = 25,
act_iters: int = 25,
iters0: int = 80,
retarget_obs: bool = True,
retarget_act: bool = True,
null_gain: float = 0.3,
stream: bool = False,
stream_hz: float = 60.0,
smooth_time: float = 0.10,
max_speed: float = 150.0,
) -> None:
self._robot = robot
self._lock = threading.Lock()
# Serialises *all* real-robot bus I/O (obs reads + streamer writes) so the
# observation thread and the streamer thread never touch the CAN bus at once.
self._io_lock = threading.Lock()
self._obs_iters = obs_iters
self._act_iters = act_iters
self._iters0 = iters0
self._do_obs = retarget_obs
self._do_act = retarget_act
# Nullspace secondary-task gain: pull the redundant DOF toward the reference
# (other-arm) joints so the two arms match in joint space without moving the EE.
self._null_gain = null_gain
# EE-priority guard for the nullspace bias (metres of 6-vector residual norm).
self._ee_tol = 0.01 # only consider a fallback when biased residual exceeds this
self._ee_slack = 0.005 # ...and EE-only beats it by more than this
# Task-only iterations appended after the biased iterations (re-tightens EE a bit
# without fully washing out the elbow bias; the fallback is the real EE guarantee).
self._final_task_iters = int(os.environ.get("RETARGET_FINAL_TASK_ITERS", "3"))
# --- joint-space smoothing streamer ---
self._stream = stream
self._stream_hz = stream_hz
self._smooth_time = smooth_time
self._max_speed = max_speed
self._goal_lock = threading.Lock()
self._goal: np.ndarray | None = None # latest short target (16-vec, deg), POLICY order
self._stream_current: np.ndarray | None = None # smoothed setpoint (16-vec, deg)
self._stream_vel = np.zeros(16)
self._stream_stop = threading.Event()
self._stream_thread: threading.Thread | None = None
# Full last real short pose (16-vec, deg) incl grippers, for streamer seeding.
self._last_short_full: np.ndarray | None = None
# Present short pose captured at connect() -> exact home to return to on shutdown.
self._home: np.ndarray | None = None
m_short = mujoco.MjModel.from_xml_path(ov.MJCF)
m_long, _ = ov.make_long(m_short)
self.m_short, self.m_long = m_short, m_long
# Separate MjData per direction so obs- and action-side solves never share buffers
# (the ThreadSafeRobot lock already serialises calls, but this is belt-and-braces).
self.d_short_o = mujoco.MjData(m_short) # obs: FK on short
self.d_long_o = mujoco.MjData(m_long) # obs: IK on long
self.d_long_a = mujoco.MjData(m_long) # action: FK on long
self.d_short_a = mujoco.MjData(m_short) # action: IK on short
self.qadr_s = rt.joint_adr(m_short)
self.qadr_l = rt.joint_adr(m_long)
self.tcp_s = {s: mujoco.mj_name2id(m_short, mujoco.mjtObj.mjOBJ_BODY, rt.TCP[s]) for s in SIDES}
self.tcp_l = {s: mujoco.mj_name2id(m_long, mujoco.mjtObj.mjOBJ_BODY, rt.TCP[s]) for s in SIDES}
self.dofs_s = {s: rt.arm_dofs(m_short, s) for s in SIDES} # (qposadr, dofadr, range)
self.dofs_l = {s: rt.arm_dofs(m_long, s) for s in SIDES}
# Real per-arm limits (radians) so retargeted commands are executable without clipping.
self.rng_real = {
s: np.deg2rad(np.array([LIMITS[s][f"joint_{i}"] for i in range(1, 8)], float)) for s in SIDES
}
self._obs_seed: dict[str, np.ndarray | None] = {s: None for s in SIDES}
self._act_seed: dict[str, np.ndarray | None] = {s: None for s in SIDES}
# Last *real* short-arm joints (rad) seen in get_observation. Used to seed the
# first action-side IK so the initial command lands in the null-space branch
# nearest the robot's actual pose (the 7-DOF arm is redundant, so the same EE
# admits many joint configs; without this the first tick could command a large
# elbow-swivel reconfiguration toward an arbitrary branch).
self._last_short: dict[str, np.ndarray] = {}
self._first_obs = True
self._first_act = True
if self._stream:
logger.info(
"Smoothing streamer ENABLED (hz=%.0f, smooth_time=%.3fs, max_speed=%.0f deg/s)",
self._stream_hz,
self._smooth_time,
self._max_speed,
)
logger.info(
"RetargetRobot ready (obs=%s, act=%s, iters=%d, iters0=%d)",
self._do_obs,
self._do_act,
self._obs_iters,
self._iters0,
)
# -- transparent proxy for everything else (connect/disconnect/features/...) --
def __getattr__(self, name):
return getattr(object.__getattribute__(self, "_robot"), name)
# -- helpers -----------------------------------------------------------------
@staticmethod
def _read_arm_deg(d: dict, side: str) -> np.ndarray:
return np.array([d[f"{side}_joint_{i}.pos"] for i in range(1, 8)], float)
def _has_full_arms(self, d: dict) -> bool:
return all(f"{s}_joint_{i}.pos" in d for s in SIDES for i in range(1, 8))
def _ik(
self,
m,
d,
tcp_id: int,
dofadr: np.ndarray,
qadr7: np.ndarray,
rng7: np.ndarray,
pt: np.ndarray,
Rt: np.ndarray,
seed: np.ndarray,
q_ref: np.ndarray,
iters: int,
use_null: bool = True,
lam: float = 0.06,
step: float = 0.25,
null_step: float = 0.08,
) -> tuple[np.ndarray, float]:
"""Damped least-squares IK for the 6-DOF EE pose with a nullspace bias toward ``q_ref``.
Primary task: reach (pt, Rt). Secondary task (projected into the task nullspace so it
never disturbs the EE): minimise ||q - q_ref||, using the 7-DOF arm's 1 redundant DOF
to keep the elbow close to the reference (long/short) configuration.
"""
q = seed.copy()
jacp = np.zeros((3, m.nv))
jacr = np.zeros((3, m.nv))
eye6 = np.eye(6)
eye7 = np.eye(7)
last = 1e9
for it in range(iters):
d.qpos[qadr7] = q
mujoco.mj_kinematics(m, d)
mujoco.mj_comPos(m, d) # required by mj_jac
p = d.xpos[tcp_id].copy()
R = d.xmat[tcp_id].reshape(3, 3)
e = np.concatenate([pt - p, rt.rot_err(R, Rt)])
last = float(np.linalg.norm(e))
mujoco.mj_jac(m, d, jacp, jacr, p, tcp_id)
J = np.vstack([jacp[:, dofadr], jacr[:, dofadr]]) # 6x7
Jt = J.T
# Task step: DAMPED pinv for stability near singularities. Clipped on its own so
# the primary always gets its full authority.
dq = np.clip(Jt @ np.linalg.solve(J @ Jt + (lam**2) * eye6, e), -step, step)
# Apply the nullspace bias only in the early iterations; the last
# ``final_task_iters`` are task-only so the EE is always re-tightened from the
# biased configuration (strict EE priority even if q_ref is far/infeasible).
null_active = use_null and self._null_gain > 0.0 and (iters - it) > self._final_task_iters
if null_active:
# Nullspace projector from the TRUE pinv so J @ N == 0 exactly: the secondary
# (elbow-toward-q_ref) task lives purely in the redundant DOF and never moves
# the EE. Clipped small so it stays strictly secondary to the task step.
Jpinv_true = np.linalg.pinv(J) # 7x6
nullproj = eye7 - Jpinv_true @ J # 7x7
dq_null = np.clip(nullproj @ (self._null_gain * (q_ref - q)), -null_step, null_step)
dq = dq + dq_null
q = q + dq
q = np.clip(q, rng7[:, 0], rng7[:, 1])
return q, last
def _solve(self, m, d, tcp_id, dofadr, qadr7, rng7, pt, Rt, seed, q_ref, iters):
"""Nullspace-biased IK with strict EE priority.
Solve with the elbow-toward-``q_ref`` bias; if that leaves an EE residual that a
pure EE-only solve would beat by more than ``_ee_slack``, fall back to EE-only. So the
bias is applied only when it is (nearly) free in EE terms — a weird/infeasible target
can never trade away end-effector accuracy for elbow matching.
"""
q, res = self._ik(m, d, tcp_id, dofadr, qadr7, rng7, pt, Rt, seed, q_ref, iters)
if self._null_gain > 0.0 and res > self._ee_tol:
q0, res0 = self._ik(
m, d, tcp_id, dofadr, qadr7, rng7, pt, Rt, seed, q_ref, iters, use_null=False
)
if res0 + self._ee_slack < res:
return q0, res0
return q, res
# -- observation: SHORT joints -> LONG joints (FK short, IK long) ------------
def get_observation(self) -> dict:
with self._io_lock:
obs = self._robot.get_observation()
# Record the full real short pose (arms + grippers) for streamer seeding.
if self._has_full_arms(obs):
full = np.zeros(16)
for s in SIDES:
full[rt.ARM_JOINT_SLICES[s]] = self._read_arm_deg(obs, s)
gk = f"{s}_gripper.pos"
if gk in obs:
full[GRIPPER_IDX[s]] = float(obs[gk])
self._last_short_full = full
if not self._do_obs or not self._has_full_arms(obs):
return obs
with self._lock:
state = np.zeros(16)
for s in SIDES:
arm = self._read_arm_deg(obs, s)
state[rt.ARM_JOINT_SLICES[s]] = arm
self._last_short[s] = np.deg2rad(arm)
rt.set_arms(self.m_short, self.d_short_o, self.qadr_s, state)
mujoco.mj_forward(self.m_short, self.d_short_o)
out = dict(obs)
for s in SIDES:
pt = self.d_short_o.xpos[self.tcp_s[s]].copy()
Rt = self.d_short_o.xmat[self.tcp_s[s]].reshape(3, 3)
q_ref = np.deg2rad(state[rt.ARM_JOINT_SLICES[s]]) # bias long pose toward short obs
seed = self._obs_seed[s]
if seed is None:
seed = q_ref.copy()
iters = self._iters0 if self._first_obs else self._obs_iters
q7, _ = self._solve(
self.m_long,
self.d_long_o,
self.tcp_l[s],
self.dofs_l[s][1],
self.dofs_l[s][0],
self.dofs_l[s][2], # clamp to LONG model ranges (policy's training space)
pt,
Rt,
seed,
q_ref,
iters,
)
self._obs_seed[s] = q7
deg = np.rad2deg(q7)
for i in range(1, 8):
out[f"{s}_joint_{i}.pos"] = float(deg[i - 1])
self._first_obs = False
return out
# -- action: LONG joint targets -> SHORT joint targets (FK long, IK short) ---
def _retarget_action(self, action: dict) -> dict:
"""Return the short-arm action dict (IK-retargeted if enabled, else a copy)."""
if not self._do_act or not isinstance(action, dict) or not self._has_full_arms(action):
return action
with self._lock:
long_q = np.zeros(16)
for s in SIDES:
long_q[rt.ARM_JOINT_SLICES[s]] = self._read_arm_deg(action, s)
rt.set_arms(self.m_long, self.d_long_a, self.qadr_l, long_q)
mujoco.mj_forward(self.m_long, self.d_long_a)
out = dict(action)
for s in SIDES:
pt = self.d_long_a.xpos[self.tcp_l[s]].copy()
Rt = self.d_long_a.xmat[self.tcp_l[s]].reshape(3, 3)
q_ref = np.deg2rad(long_q[rt.ARM_JOINT_SLICES[s]]) # bias short pose toward long target
seed = self._act_seed[s]
if seed is None:
# Prefer the robot's real current short pose (nearest branch, minimal
# startup motion); fall back to the long target angles if unseen.
seed = (
self._last_short[s].copy()
if s in self._last_short
else q_ref.copy()
)
iters = self._iters0 if self._first_act else self._act_iters
q7, _ = self._solve(
self.m_short,
self.d_short_a,
self.tcp_s[s],
self.dofs_s[s][1],
self.dofs_s[s][0],
self.rng_real[s], # clamp to REAL limits -> safe on hardware
pt,
Rt,
seed,
q_ref,
iters,
)
self._act_seed[s] = q7
deg = np.rad2deg(q7)
for i in range(1, 8):
out[f"{s}_joint_{i}.pos"] = float(deg[i - 1])
self._first_act = False
return out
def send_action(self, action: dict):
short = self._retarget_action(action)
# Streaming path: just publish the target; the streamer thread writes to the bus.
if self._stream and isinstance(short, dict) and self._has_full_arms(short):
self._set_goal(short)
return short
with self._io_lock:
return self._robot.send_action(short)
# -- streamer: joint-space smoothing at a fixed high rate ---------------------
@staticmethod
def _dict_to_vec(d: dict, fallback: np.ndarray | None = None) -> np.ndarray:
vec = np.zeros(16) if fallback is None else fallback.copy()
for s in SIDES:
base = rt.ARM_JOINT_SLICES[s].start
for i in range(1, 8):
vec[base + i - 1] = float(d[f"{s}_joint_{i}.pos"])
gk = f"{s}_gripper.pos"
if gk in d:
vec[GRIPPER_IDX[s]] = float(d[gk])
return vec
@staticmethod
def _vec_to_action(vec: np.ndarray) -> dict:
out = {}
for s in SIDES:
base = rt.ARM_JOINT_SLICES[s].start
for i in range(1, 8):
out[f"{s}_joint_{i}.pos"] = float(vec[base + i - 1])
out[f"{s}_gripper.pos"] = float(vec[GRIPPER_IDX[s]])
return out
def _set_goal(self, short: dict) -> None:
with self._goal_lock:
base = self._goal if self._goal is not None else self._last_short_full
self._goal = self._dict_to_vec(short, fallback=base)
def _seed_current(self) -> None:
"""Seed the streamer setpoint from the robot's actual present pose (ramp start)."""
try:
with self._io_lock:
obs = self._robot.get_observation()
if self._has_full_arms(obs):
self._stream_current = self._dict_to_vec(obs)
self._stream_vel = np.zeros(16)
if self._home is None:
self._home = self._stream_current.copy()
logger.info("Streamer seeded from present robot pose")
return
except Exception as e: # noqa: BLE001
logger.warning("Streamer seed failed (%s); will seed from first goal", e)
self._stream_current = None
def _stream_loop(self) -> None:
dt = 1.0 / self._stream_hz
pos_eps = 0.05 # deg: below this distance to goal we consider the axis settled
vel_eps = 0.5 # deg/s: below this speed we consider motion stopped
while not self._stream_stop.is_set():
t0 = time.perf_counter()
with self._goal_lock:
goal = None if self._goal is None else self._goal.copy()
if goal is None:
time.sleep(dt)
continue
if self._stream_current is None:
self._stream_current = goal.copy()
self._stream_vel = np.zeros(16)
self._stream_current, self._stream_vel = smooth_damp(
self._stream_current, goal, self._stream_vel, self._smooth_time, dt, self._max_speed
)
# Damiao MIT mode holds the last command, so when we've converged on the
# goal and stopped moving we skip the CAN write entirely. This keeps the
# bus from saturating (Errno 105) during the pauses when the policy goal
# is constant, and frees the _io_lock so the main loop's reads stay fast.
settled = (
float(np.max(np.abs(self._stream_current - goal))) < pos_eps
and float(np.max(np.abs(self._stream_vel))) < vel_eps
)
if not settled:
act = self._vec_to_action(self._stream_current)
try:
with self._io_lock:
self._robot.send_action(act)
except Exception as e: # noqa: BLE001
logger.warning("Streamer send_action failed: %s", e)
sleep_t = dt - (time.perf_counter() - t0)
if sleep_t > 0:
time.sleep(sleep_t)
# -- lifecycle (start/stop the streamer around the real connect/disconnect) ---
def connect(self, *args, **kwargs):
result = self._robot.connect(*args, **kwargs)
if self._stream and self._stream_thread is None:
self._seed_current()
self._stream_stop.clear()
self._stream_thread = threading.Thread(
target=self._stream_loop, name="RetargetStreamer", daemon=True
)
self._stream_thread.start()
logger.info("Smoothing streamer thread started")
return result
def _drain_home(self, timeout_s: float = 6.0, tol_deg: float = 0.7) -> None:
"""Command the captured home pose and keep the streamer running until it
actually converges (or times out), so the arms fully reach home before we
cut the streamer thread. Only meaningful when streaming is enabled."""
if self._home is None or self._stream_thread is None:
return
logger.info("Returning arms to home pose (draining streamer)...")
with self._goal_lock:
self._goal = self._home.copy()
t_start = time.perf_counter()
while time.perf_counter() - t_start < timeout_s:
cur = self._stream_current
if cur is not None and float(np.max(np.abs(cur - self._home))) < tol_deg:
logger.info("Home pose reached")
return
time.sleep(0.05)
logger.warning("Home drain timed out after %.1fs; stopping streamer anyway", timeout_s)
def disconnect(self, *args, **kwargs):
if self._stream_thread is not None:
try:
self._drain_home()
except Exception as e: # noqa: BLE001
logger.warning("Home drain failed: %s", e)
self._stream_stop.set()
self._stream_thread.join(timeout=2.0)
self._stream_thread = None
logger.info("Smoothing streamer thread stopped")
return self._robot.disconnect(*args, **kwargs)
def _wrap_factory(orig):
obs_iters = int(os.environ.get("RETARGET_ITERS", "25"))
act_iters = int(os.environ.get("RETARGET_ITERS", "25"))
iters0 = int(os.environ.get("RETARGET_ITERS0", "80"))
do_obs = os.environ.get("RETARGET_OBS", "1") != "0"
do_act = os.environ.get("RETARGET_ACT", "1") != "0"
null_gain = float(os.environ.get("RETARGET_NULL_GAIN", "0.3"))
stream = os.environ.get("STREAM", "0") != "0"
stream_hz = float(os.environ.get("STREAM_HZ", "40"))
smooth_time = float(os.environ.get("STREAM_SMOOTH_TIME", "0.10"))
max_speed = float(os.environ.get("STREAM_MAX_SPEED", "150"))
def factory(cfg):
real = orig(cfg)
logger.info("Wrapping %s with RetargetRobot (MuJoCo long<->short EE retarget)", type(real).__name__)
return RetargetRobot(
real,
obs_iters=obs_iters,
act_iters=act_iters,
iters0=iters0,
retarget_obs=do_obs,
retarget_act=do_act,
null_gain=null_gain,
stream=stream,
stream_hz=stream_hz,
smooth_time=smooth_time,
max_speed=max_speed,
)
return factory
def _patch_rtc_realtime():
"""Make RTC re-anchor new chunks on the *actual* number of actions consumed
during inference instead of the fps-based estimate.
The stock ActionQueue computes ``real_delay = ceil(latency * fps)`` and discards
that many actions from every new chunk. When the control loop runs below --fps
(e.g. 20 Hz while fps=30), real_delay (18) overshoots the actions that were truly
consumed (indexes_diff=12), so it skips ~6 extra actions per inference and the
trajectory plays faster than real time. ``indexes_diff = last_index - idx_before``
is the ground-truth count of what the robot actually executed during inference, so
skipping exactly that many gives real-time playback independent of the loop rate.
"""
try:
from lerobot.policies.rtc.action_queue import ActionQueue
except Exception as e: # noqa: BLE001
logger.warning("Could not patch RTC action queue for real-time playback: %s", e)
return
def _resolve(self, real_delay, action_index_before_inference=None):
if action_index_before_inference is not None:
return max(0, self.last_index - action_index_before_inference)
return max(0, real_delay)
ActionQueue._check_and_resolve_delays = _resolve
logger.info("Patched RTC ActionQueue: re-anchor on real consumed index (no over-skip)")
def main():
import lerobot.rollout.context as context
context.make_robot_from_config = _wrap_factory(context.make_robot_from_config)
_patch_rtc_realtime()
from lerobot.scripts.lerobot_rollout import main as rollout_main
rollout_main()
if __name__ == "__main__":
main()
@@ -74,6 +74,14 @@ class OpenArmFollowerConfigBase:
# Set to a positive scalar for all motors, or a dict mapping motor names to limits
max_relative_target: float | dict[str, float] | None = None
# End-effector kinematics (optional). Point `urdf_path` at the OpenArm URDF and set
# `target_frame_name` to the end-effector link in that URDF to enable forward/inverse
# kinematics in Cartesian (EE) space -- via `lerobot.model.RobotKinematics` and the
# shared FK/IK processor steps. Left as None, the robot behaves exactly as before
# (joint space only). The URDF is user-supplied, so no large asset is vendored here.
urdf_path: str | None = None
target_frame_name: str | None = None
# Camera configurations
cameras: dict[str, CameraConfig] = field(default_factory=dict)
@@ -17,7 +17,7 @@
import logging
import time
from functools import cached_property
from typing import Any
from typing import TYPE_CHECKING, Any
from lerobot.cameras import make_cameras_from_configs
from lerobot.lerobot_types import RobotAction, RobotObservation
@@ -25,6 +25,9 @@ from lerobot.motors import Motor, MotorCalibration, MotorNormMode
from lerobot.motors.damiao import DamiaoMotorsBus
from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
if TYPE_CHECKING:
from lerobot.model import RobotKinematics
from ..robot import Robot
from ..utils import ensure_safe_goal_position
from .config_openarm_follower import (
@@ -120,6 +123,33 @@ class OpenArmFollower(Robot):
"""Action features."""
return self._motors_ft
@property
def arm_motor_names(self) -> list[str]:
"""Arm joints forming the kinematic chain to the end-effector (excludes the gripper)."""
return [motor for motor in self.bus.motors if motor != "gripper"]
def make_kinematics(self) -> "RobotKinematics":
"""Build a solver for end-effector forward/inverse kinematics.
Requires ``config.urdf_path`` (path to the OpenArm URDF) and
``config.target_frame_name`` (the end-effector link in that URDF). Pair the returned
solver with the shared FK/IK processor steps in
``lerobot.robots.so_follower.robot_kinematic_processor`` to record or command the arm
in end-effector (Cartesian) space instead of raw joint angles.
"""
from lerobot.model import RobotKinematics
if self.config.urdf_path is None or self.config.target_frame_name is None:
raise ValueError(
"OpenArm end-effector kinematics require config.urdf_path and "
"config.target_frame_name (the OpenArm URDF and its end-effector link)."
)
return RobotKinematics(
urdf_path=self.config.urdf_path,
target_frame_name=self.config.target_frame_name,
joint_names=self.arm_motor_names,
)
@property
def is_connected(self) -> bool:
"""Check if robot is connected."""