Export runs as LeRobot datasets

POST /datasets/export turns finished runs of one robot into a LeRobot dataset:
observation.state and .velocity (the robot's joints), action (the commanded
controls), observation.images.main (the video, 640x360) and the scenario text as
the task. With HF_TOKEN in the Modal secret it is pushed to the Hub as a private
dataset; otherwise the zip is stored with the job. Verified end to end with
lerobot 0.4.4 (write, then load back); CI repeats that with a Go2. Trajectories
now also carry every body's world pose, for replay renders.

.github/workflows/ci.yml

@@ -142,3 +142,53 @@ jobs:
           path: ros-out
           if-no-files-found: warn
 
+  lerobot-export:
+    name: LeRobot dataset export
+    runs-on: ubuntu-24.04
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      - name: Install (CPU PyTorch, LeRobot, the simulator)
+        run: |
+          sudo apt-get update && sudo apt-get install -y --no-install-recommends libosmesa6 libgl1 ffmpeg
+          pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+          pip install lerobot==0.4.4 -r worker/requirements.txt
+      - name: Simulate two runs, export them, load the dataset back
+        env:
+          MUJOCO_GL: osmesa
+        run: |
+          git clone -q --depth 1 --filter=blob:none --sparse https://github.com/google-deepmind/mujoco_menagerie.git /tmp/men
+          git -C /tmp/men sparse-checkout set unitree_go2
+          python - <<'EOF'
+          from pathlib import Path
+          from api.mjcf_check import prepare_scene_dir
+          from worker.lerobot_export import run_export_job
+          from worker.sim_runner import config_for, simulate
+
+          scene = """<mujoco><include file="go2.xml"/><worldbody><light pos="0 0 4" dir="0 0 -1"/>
+            <geom name="floor" type="plane" size="0 0 0.05"/></worldbody></mujoco>"""
+          episodes = []
+          for i, controller in enumerate(("hold", "passive")):
+              d = prepare_scene_dir(scene, "unitree_go2", Path("/tmp/men"))
+              cfg = config_for("unitree_go2", 1.0, controller)
+              cfg.width, cfg.height = 320, 180
+              simulate(d / "scene.xml", Path(f"/tmp/run{i}"), cfg)
+              episodes.append({"run_id": f"r{i}", "task": f"Go2 {controller}",
+                               "trajectory_url": f"/tmp/run{i}/trajectory.npz", "video_url": f"/tmp/run{i}/video.mp4"})
+          out = run_export_job({"repo_id": "ci/rosdiff-go2", "robot": "unitree_go2", "episodes": episodes})
+          print(out["summary"])
+          assert out["summary"]["episodes"] == 2, out["summary"]
+          Path("/tmp/ds.zip").write_bytes(out["artifacts"]["dataset.zip"])
+
+          import shutil
+          shutil.unpack_archive("/tmp/ds.zip", "/tmp/ds")
+          from lerobot.datasets.lerobot_dataset import LeRobotDataset
+          ds = LeRobotDataset("ci/rosdiff-go2", root="/tmp/ds")
+          item = ds[len(ds) - 1]
+          print({k: tuple(v.shape) if hasattr(v, "shape") else v for k, v in item.items()})
+          assert ds.num_episodes == 2 and item["task"] == "Go2 passive"
+          assert tuple(item["observation.images.main"].shape) == (3, 360, 640)
+          assert tuple(item["observation.state"].shape) == (19,) and tuple(item["action"].shape) == (12,)
+          EOF

api/config.py

@@ -56,6 +56,7 @@ class Settings:
     job_overhead_s: float  # billed seconds per job beyond the worker's own timing (start-up, idle shutdown)
     isaac_daily_limit: int  # Isaac Sim renders per day, across all testers
     auto_review: bool  # review every finished simulation video with Cosmos Reason
+    hf_namespace: str  # Hugging Face user/org that dataset exports are pushed to (with the worker's HF_TOKEN)
     foxglove_org: str  # optional Foxglove organization slug for the embedded viewer
 
     @property
@@ -113,5 +114,6 @@ def get_settings() -> Settings:
         job_overhead_s=float(_env("JOB_OVERHEAD_S", "45")),
         isaac_daily_limit=int(_env("ISAAC_DAILY_LIMIT", "10")),
         auto_review=_env("AUTO_REVIEW", "true").lower() in ("1", "true", "yes"),
+        hf_namespace=_env("HF_NAMESPACE", "rosdiff"),
         foxglove_org=_env("FOXGLOVE_ORG", ""),
     )

api/main.py

@@ -150,6 +150,12 @@ class SweepRequest(BaseModel):
     seed: int = 0
 
 
+class ExportRequest(BaseModel):
+    run_ids: list[str] = Field(..., min_length=1, max_length=100, description="Your finished runs, one robot")
+    name: str = Field(..., pattern=r"^[a-z0-9][a-z0-9-]{1,60}$", description="Dataset name, e.g. g1-table-pushes")
+    private: bool = True
+
+
 class CheckRequest(BaseModel):
     generation_id: str = Field(..., description="A package from /code/generate (yours)")
 
@@ -436,6 +442,45 @@ def create_app(services: Services | None = None) -> FastAPI:
             raise HTTPException(status_code=e.status, detail=str(e)) from e
         return _job_response(request, s, job)
 
+    @app.post("/datasets/export", response_model=JobResponse)
+    def export_dataset(req: ExportRequest, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)):
+        """Turn finished runs into a LeRobot dataset (state, velocity, action, video, task) on the Hugging Face Hub."""
+        episodes, robots = [], set()
+        for run_id in dict.fromkeys(req.run_ids):
+            run = s.logs.get_run(run_id)
+            if not run or run["tester"] != t.name or (run.get("kind") or "sim") != "sim":
+                raise HTTPException(status_code=404, detail=f"no simulation run {run_id} for this key")
+            run = refresh(s, run)
+            files = (run.get("output") or {}).get("files") or {}
+            if run["status"] != "COMPLETED" or "trajectory" not in files:
+                raise HTTPException(status_code=422, detail=f"run {run_id} is not finished or has no trajectory")
+            robots.add(run.get("robot"))
+            task = ""
+            if run.get("generation_id"):
+                task = ((s.logs.get(run["generation_id"]) or {}).get("request") or {}).get("description", "")
+            episodes.append(
+                {
+                    "run_id": run_id,
+                    "task": task,
+                    "trajectory_url": s.storage.url(files["trajectory"]),
+                    "video_url": s.storage.url(files["video"]) if "video" in files else None,
+                }
+            )
+        if len(robots) > 1:
+            raise HTTPException(status_code=422, detail="all runs in one dataset must use the same robot")
+        robot = robots.pop()
+        repo_id = f"{s.settings.hf_namespace}/rosdiff-{req.name}"
+        job = _submit(
+            s,
+            t.name,
+            "export",
+            {"repo_id": repo_id, "robot": robot, "episodes": episodes, "private": req.private},
+            run_id=new_run_id(),
+            robot=robot,
+            params={"repo_id": repo_id, "runs": [e["run_id"] for e in episodes]},
+        )
+        return _job_response(request, s, job)
+
     @app.post("/code/check", response_model=JobResponse)
     def code_check(req: CheckRequest, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)):
         """Build, lint, test and launch a generated package in the ROS 2 sandbox. Returns at once; see report_url."""

api/runs.py

@@ -412,3 +412,34 @@ def _review_body(job: dict, storage: Storage | None) -> list[str]:
 
 
 KIND_BODIES["review"] = _review_body
+
+
+def _export_body(job: dict, storage: Storage | None) -> list[str]:
+    output = job.get("output") or {}
+    r = output.get("summary") or {}
+    files = output.get("files") or {}
+    link = ""
+    if r.get("hub_url"):
+        link = f' · <a href="{html.escape(r["hub_url"])}">open on Hugging Face</a>'
+    cards = [f'<div class="card"><b>{html.escape(r.get("headline", ""))}</b>{link}</div>']
+    cells = [
+        ("Dataset", r.get("repo_id")),
+        ("Episodes", r.get("episodes")),
+        ("Frames", r.get("frames")),
+        ("FPS", r.get("fps")),
+        ("Runs", ", ".join(r.get("runs") or [])),
+    ]
+    grid = "".join(f"<div><span>{html.escape(k)}</span>{_fmt(v)}</div>" for k, v in cells)
+    cards.append(f'<div class="card"><div class="grid">{grid}</div></div>')
+    repo = html.escape(r.get("repo_id", "<repo_id>"))
+    load = f'LeRobotDataset("{repo}")' if r.get("hub_url") else f'LeRobotDataset("{repo}", root="path/to/unzipped")'
+    cards.append(
+        '<div class="card"><b>Use it</b><pre>pip install lerobot\n'
+        f"from lerobot.datasets.lerobot_dataset import LeRobotDataset\nds = {load}\n"
+        "ds[0].keys()  # observation.state, observation.velocity, action, observation.images.main, task</pre></div>"
+    )
+    cards.extend(file_links(files, storage))
+    return cards
+
+
+KIND_BODIES["export"] = _export_body

tests/test_export.py

@@ -0,0 +1,45 @@
+from conftest import complete
+
+from worker.lerobot_export import features
+
+
+def finished_run(client, backend, h, gen_id, n):
+    run_id = client.post("/simulate/run", json={"generation_id": gen_id}, headers=h).json()["run_id"]
+    complete(backend, f"job-{n}", run_id)
+    backend.jobs[f"job-{n}"]["output"]["files"]["trajectory"] = f"runs/{run_id}/trajectory.npz"
+    return run_id
+
+
+def test_export_endpoint(run_env):
+    client, keys_, backend, services, gen_id = run_env
+    h = {"X-API-Key": keys_["alice"]}
+    runs = [finished_run(client, backend, h, gen_id, n) for n in (1, 2)]
+    r = client.post("/datasets/export", json={"run_ids": runs, "name": "g1-crates"}, headers=h)
+    assert r.status_code == 200, r.text
+    assert r.json()["kind"] == "export" and backend.functions[-1] == "export_lerobot"
+    sent = backend.submitted[-1]
+    assert sent["repo_id"] == "rosdiff/rosdiff-g1-crates" and sent["robot"] == "unitree_g1" and sent["private"]
+    assert [e["run_id"] for e in sent["episodes"]] == runs
+    assert sent["episodes"][0]["task"] == "G1 near a slope and a crate"
+    assert sent["episodes"][0]["trajectory_url"].endswith("trajectory.npz")
+
+
+def test_export_rejects_unfinished_or_foreign_runs(run_env):
+    client, keys_, backend, services, gen_id = run_env
+    h = {"X-API-Key": keys_["alice"]}
+    pending = client.post("/simulate/run", json={"generation_id": gen_id}, headers=h).json()["run_id"]
+    r = client.post("/datasets/export", json={"run_ids": [pending], "name": "x1"}, headers=h)
+    assert r.status_code == 422 and "not finished" in r.text
+    r = client.post("/datasets/export", json={"run_ids": [pending], "name": "x1"}, headers={"X-API-Key": keys_["bob"]})
+    assert r.status_code == 404
+    assert (
+        client.post("/datasets/export", json={"run_ids": [pending], "name": "Bad Name"}, headers=h).status_code == 422
+    )
+
+
+def test_features():
+    f = features(["a", "b"], 2, ["m"], with_video=True)
+    assert f["observation.state"]["shape"] == (2,) and f["action"]["names"] == ["m"]
+    assert f["observation.images.main"]["dtype"] == "video" and "observation.images.main" not in features(
+        ["a"], 1, [], with_video=False
+    )

tests/test_runs.py

@@ -34,6 +34,8 @@ class TestWorker:
         }
         traj = np.load(storage.root / "runs/abc123/trajectory.npz")
         assert traj["qpos"].shape[0] == len(traj["time"]) == 16 and traj["ctrl"].shape == (16, 1)
+        nbody = len(traj["body_names"])
+        assert traj["body_pos"].shape == (16, nbody, 3) and traj["body_quat"].shape == (16, nbody, 4)
         assert (
             list(traj["joint_names"]) == [f"floating_base_joint[{i}]" for i in range(7)] + ["left_knee_joint"]
             and int(traj["fps"]) == 30

worker/lerobot_export.py

@@ -0,0 +1,145 @@
+"""Export simulation runs as a LeRobot dataset (Hugging Face's format for robot learning data).
+
+Each run becomes one episode:
+    observation.state     the robot's joint positions (its free joint first, for mobile robots)
+    observation.velocity  the robot's joint velocities
+    action                the controls sent to the actuators (what the hold controller commanded)
+    observation.images.main  the run's video, downscaled
+    task                  the scenario description the scene was generated from
+
+The dataset is pushed to the Hugging Face Hub as a private dataset when HF_TOKEN is set (Modal secret
+rosdiff-hf); otherwise it comes back as a zip stored next to the job. Load it with
+`LeRobotDataset("<repo_id>")`, or from the unzipped folder with `LeRobotDataset("<repo_id>", root=...)`.
+"""
+
+from __future__ import annotations
+
+import os
+import shutil
+import tempfile
+import time
+import urllib.request
+from pathlib import Path
+
+import numpy as np
+
+IMAGE_WIDTH, IMAGE_HEIGHT = 640, 360
+MAX_ZIP_BYTES = 200_000_000
+
+
+def _fetch(src: str, dest: Path) -> Path:
+    if src.startswith(("https://", "http://")):
+        with urllib.request.urlopen(src, timeout=300) as r:  # noqa: S310 (presigned R2 links)
+            dest.write_bytes(r.read())
+    else:
+        shutil.copy(src.removeprefix("file://"), dest)
+    return dest
+
+
+def _video_frames(path: Path, count: int) -> list[np.ndarray]:
+    import imageio_ffmpeg
+
+    reader = imageio_ffmpeg.read_frames(
+        str(path), output_params=["-vf", f"scale={IMAGE_WIDTH}:{IMAGE_HEIGHT}"], pix_fmt="rgb24"
+    )
+    meta = next(reader)
+    w, h = meta.get("size", (IMAGE_WIDTH, IMAGE_HEIGHT))
+    frames = [np.frombuffer(f, np.uint8).reshape(h, w, 3) for f in reader]
+    if not frames:
+        raise ValueError("the run's video has no frames")
+    # The video and the trajectory are captured at the same instants; pad or trim if an encoder dropped one.
+    frames = frames[:count]
+    while len(frames) < count:
+        frames.append(frames[-1])
+    return frames
+
+
+def features(joint_names: list[str], dof: int, actuator_names: list[str], with_video: bool) -> dict:
+    f = {
+        "observation.state": {"dtype": "float32", "shape": (len(joint_names),), "names": joint_names},
+        "observation.velocity": {"dtype": "float32", "shape": (dof,), "names": None},
+        "action": {"dtype": "float32", "shape": (len(actuator_names),), "names": actuator_names},
+    }
+    if with_video:
+        f["observation.images.main"] = {
+            "dtype": "video",
+            "shape": (IMAGE_HEIGHT, IMAGE_WIDTH, 3),
+            "names": ["height", "width", "channels"],
+        }
+    return f
+
+
+def build_dataset(episodes: list[dict], repo_id: str, root: Path, robot: str | None, with_video: bool = True):
+    """episodes: [{"trajectory": Path, "video": Path | None, "task": str}] -> LeRobotDataset written to root."""
+    from lerobot.datasets.lerobot_dataset import LeRobotDataset
+
+    first = np.load(episodes[0]["trajectory"])
+    joint_names = [str(n) for n in first["joint_names"]]
+    actuator_names = [str(n) for n in first["actuator_names"]]
+    fps = int(first["fps"])
+    dof = int(first["qvel"].shape[1])
+    with_video = with_video and all(e.get("video") for e in episodes)
+    ds = LeRobotDataset.create(
+        repo_id=repo_id,
+        fps=fps,
+        features=features(joint_names, dof, actuator_names, with_video),
+        root=root,
+        robot_type=robot or "none",
+        use_videos=with_video,
+        vcodec="h264",
+    )
+    for ep in episodes:
+        traj = np.load(ep["trajectory"])
+        if [str(n) for n in traj["joint_names"]] != joint_names or int(traj["fps"]) != fps:
+            raise ValueError("all runs in one dataset must use the same robot (same joints) and frame rate")
+        n = len(traj["time"])
+        images = _video_frames(ep["video"], n) if with_video else None
+        for i in range(n):
+            frame = {
+                "observation.state": traj["qpos"][i].astype(np.float32),
+                "observation.velocity": traj["qvel"][i].astype(np.float32),
+                "action": traj["ctrl"][i].astype(np.float32),
+                "task": ep.get("task") or "simulation run",
+            }
+            if with_video:
+                frame["observation.images.main"] = images[i]
+            ds.add_frame(frame)
+        ds.save_episode()
+    ds.finalize()
+    return ds
+
+
+def run_export_job(job_input: dict) -> dict:
+    """Worker entry point: download the runs, build the dataset, push it (or zip it)."""
+    started = time.monotonic()
+    repo_id = job_input["repo_id"]
+    with tempfile.TemporaryDirectory() as tmp:
+        tmp = Path(tmp)
+        episodes = []
+        for i, ep in enumerate(job_input["episodes"]):
+            traj = _fetch(ep["trajectory_url"], tmp / f"traj_{i}.npz")
+            video = _fetch(ep["video_url"], tmp / f"video_{i}.mp4") if ep.get("video_url") else None
+            episodes.append({"trajectory": traj, "video": video, "task": ep.get("task", "")})
+        root = tmp / "dataset"
+        ds = build_dataset(episodes, repo_id, root, job_input.get("robot"))
+        summary = {
+            "repo_id": repo_id,
+            "episodes": ds.meta.total_episodes,
+            "frames": ds.meta.total_frames,
+            "fps": ds.fps,
+            "features": sorted(ds.features),
+            "runs": [ep.get("run_id") for ep in job_input["episodes"]],
+        }
+        artifacts = {}
+        if os.environ.get("HF_TOKEN"):
+            ds.push_to_hub(private=job_input.get("private", True), tags=["rosdiff", "mujoco", "simulation"])
+            summary["hub_url"] = f"https://huggingface.co/datasets/{repo_id}"
+            summary["headline"] = f"{summary['episodes']} episode(s), {summary['frames']} frames → {repo_id} (private)"
+        else:
+            archive = shutil.make_archive(str(tmp / "dataset"), "zip", root)
+            size = Path(archive).stat().st_size
+            if size > MAX_ZIP_BYTES:
+                raise ValueError(f"dataset zip is {size / 1e6:.0f} MB; set HF_TOKEN to push it to the Hub instead")
+            artifacts["dataset.zip"] = Path(archive).read_bytes()
+            summary["headline"] = f"{summary['episodes']} episode(s), {summary['frames']} frames (zip; no HF_TOKEN)"
+    return {"summary": summary, "artifacts": artifacts, "worker_seconds": round(time.monotonic() - started, 1)}

worker/modal_app.py

@@ -189,3 +189,34 @@ class Reviewer:
 def review_video(job_input: dict) -> dict:
     """Entry point the API spawns; hands the job to the (GPU, model-loaded) Reviewer."""
     return Reviewer().review.remote(job_input)
+
+
+# --------------------------------------------------------------------------- LeRobot dataset export (CPU)
+
+lerobot_image = (
+    modal.Image.debian_slim(python_version="3.11")
+    .apt_install("ffmpeg")
+    .pip_install("lerobot==0.4.4", "imageio-ffmpeg==0.6.0")
+    .add_local_python_source("worker")
+)
+
+
+@app.function(
+    image=lerobot_image,
+    cpu=2.0,
+    memory=4096,
+    timeout=900,
+    max_containers=1,
+    scaledown_window=20,
+    secrets=[modal.Secret.from_name(os.environ.get("MODAL_HF_SECRET", "rosdiff-hf"))],
+)
+def export_lerobot(job_input: dict) -> dict:
+    """Runs -> LeRobot dataset, pushed to the Hub with the secret's HF_TOKEN (else returned as a zip)."""
+    import traceback
+
+    from worker.lerobot_export import run_export_job
+
+    try:
+        return run_export_job(job_input)
+    except Exception as e:
+        return {"error": f"{type(e).__name__}: {e}", "traceback": traceback.format_exc()[-2000:]}

worker/sim_runner.py

@@ -10,7 +10,8 @@ Outputs in OUT_DIR:
                   /state (JSON: time, root height, fallen flag, contacts)
     video.mp4     H.264 render from a camera tracking the robot (skipped if rendering is unavailable)
     video.webm    the same video as VP9, for browsers without H.264
-    trajectory.npz  the robot's joint positions/velocities and the controls at `fps` (for dataset export)
+    trajectory.npz  at `fps`: the robot's joint positions/velocities and the controls (dataset export), and
+                  every body's world pose (replay renders)
     summary.json  what happened: duration, real-time factor, fall detection, warnings
 """
 
@@ -369,7 +370,7 @@ def simulate(scene: Path, out_dir: Path, config: RunConfig) -> dict:
 
     scene_ch.log(_scene_update(model), log_time=0)
     traj_q, traj_v, traj_names = _robot_state_index(model, root)
-    traj: dict[str, list] = {"time": [], "qpos": [], "qvel": [], "ctrl": []}
+    traj: dict[str, list] = {"time": [], "qpos": [], "qvel": [], "ctrl": [], "body_pos": [], "body_quat": []}
     steps = max(1, int(math.ceil(config.duration_s / model.opt.timestep)))
     frame_dt = 1.0 / config.fps
     next_frame = 0.0
@@ -391,6 +392,8 @@ def simulate(scene: Path, out_dir: Path, config: RunConfig) -> dict:
             traj["qpos"].append(data.qpos[traj_q].copy())
             traj["qvel"].append(data.qvel[traj_v].copy())
             traj["ctrl"].append(data.ctrl.copy())
+            traj["body_pos"].append(data.xpos.astype(np.float32))
+            traj["body_quat"].append(data.xquat.astype(np.float32))
             if videos:
                 try:
                     renderer.update_scene(data, cam)
@@ -427,6 +430,12 @@ def simulate(scene: Path, out_dir: Path, config: RunConfig) -> dict:
         actuator_names=np.asarray(
             [mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, a) or f"actuator_{a}" for a in range(model.nu)]
         ),
+        body_pos=np.asarray(traj["body_pos"], np.float32),  # world pose of every body (for replay renders)
+        body_quat=np.asarray(traj["body_quat"], np.float32),
+        body_names=np.asarray(
+            [mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, b) or f"body_{b}" for b in range(model.nbody)]
+        ),
+        body_parents=np.asarray(model.body_parentid, np.int32),
         fps=np.asarray(config.fps),
     )
     _close_videos(videos)