From 2e37cb22e8f508a77519160b3ada7df7a1545749 Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Fri, 31 Jul 2026 15:42:32 +0200 Subject: [PATCH] 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 --- examples/openarm/README.md | 75 ++++++++++ examples/openarm/render_episode.py | 217 +++++++++++++++++++++++++++++ 2 files changed, 292 insertions(+) create mode 100644 examples/openarm/README.md create mode 100644 examples/openarm/render_episode.py diff --git a/examples/openarm/README.md b/examples/openarm/README.md new file mode 100644 index 000000000..2c5b556a4 --- /dev/null +++ b/examples/openarm/README.md @@ -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 /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` → +`/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. | diff --git a/examples/openarm/render_episode.py b/examples/openarm/render_episode.py new file mode 100644 index 000000000..858a019a5 --- /dev/null +++ b/examples/openarm/render_episode.py @@ -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 ``/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 (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()