admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-31 21:49:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

docs(openarm): add episode-replay example

Browse Source 
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>
This commit is contained in:
Martino Russi
2026-07-31 15:42:32 +02:00
parent ea5cabe6ff
commit 2e37cb22e8
2 changed files with 292 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/openarm/README.md
+75
View File
@@ -0,0 +1,75 @@
# OpenArm — Episode Replay in Simulation
Replay a recorded bimanual-[OpenArm](https://openarm.dev) episode into an mp4 by driving the
official OpenArm MuJoCo model directly from a LeRobot dataset's recorded joint states. This is
a visual sanity check for recorded/commanded trajectories and for the end-effector kinematics
exposed by `OpenArmFollower.make_kinematics()` (see the [OpenArm docs](../../docs/source/openarm.mdx)).
## Model provenance
Everything is pulled from Enactic's official, Apache-2.0 OpenArm repositories — nothing is
vendored into LeRobot:
| Asset | Source | License |
| ------------------------------------ | ------------------------------------------------------------------------------- | ---------- |
| MuJoCo MJCF (used here) | [`enactic/openarm_mujoco`](https://github.com/enactic/openarm_mujoco) | Apache-2.0 |
| URDF / xacro (for `RobotKinematics`) | [`enactic/openarm_description`](https://github.com/enactic/openarm_description) | Apache-2.0 |
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 of a local LeRobot v3.0 dataset
LD_LIBRARY_PATH=$CONDA_PREFIX/lib MUJOCO_GL=egl \
python -m examples.openarm.render_episode \
--dataset data/folding_src_meta \
--episode 1 \
--out openarm_ep1.mp4
# No dataset handy? Smoke-test with a synthetic wave:
LD_LIBRARY_PATH=$CONDA_PREFIX/lib MUJOCO_GL=egl \
python -m examples.openarm.render_episode --demo --out openarm_demo.mp4
```
Useful flags: `--fps` (default 30), `--width` / `--height` (default 960×720), `--mjcf` to point
at an explicit model file.
## Troubleshooting
| Symptom | Fix |
| ------------------------------------- | ---------------------------------------------------------------------- |
| `Could not find the OpenArm v1 MJCF` | Pass `--mjcf`, set `$OPENARM_MJCF`, or install/clone `openarm_mujoco`. |
| `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. |
examples/openarm/render_episode.py
+217
View File
@@ -0,0 +1,217 @@
#!/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) directly
from a LeRobot dataset's recorded ``observation.state`` and renders a headless video. It is
a visual sanity check that a recorded/commanded joint trajectory is what you think it is --
the same forward-kinematics path used to expose end-effector poses (see the OpenArm
end-effector kinematics helper ``OpenArmFollower.make_kinematics``).
The ``observation.state`` is expected in the 16-D bimanual layout (degrees):
right_joint_1..7, right_gripper, left_joint_1..7, left_gripper
Model (Apache-2.0): 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).
Examples:
# replay episode 1 of a local LeRobot v3.0 dataset
LD_LIBRARY_PATH=$CONDA_PREFIX/lib MUJOCO_GL=egl python -m examples.openarm.render_episode \
--dataset data/folding_src_meta --episode 1 --out openarm_ep1.mp4
# smoke test with a synthetic wave (no dataset / model joints only)
LD_LIBRARY_PATH=$CONDA_PREFIX/lib MUJOCO_GL=egl python -m examples.openarm.render_episode \
--demo --out openarm_demo.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
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 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("--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]
print(f"driving {n_frames} frames")
# 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()