#!/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()