Add Isaac Sim RTX renders of finished runs

POST /simulate/run/{id}/isaac replays a run in Isaac Sim: worker/usd_scene.py
builds a USD stage from the MuJoCo model with every body's pose time-sampled
from the run (and a camera that tracks the robot like the MuJoCo video), and
Isaac Sim only renders it frame by frame with RTX - no second physics engine, so
both videos show the same motion. Deployed as its own Modal app on an L4
(worker/modal_isaac.py) so its large image never slows the main deploy; limited
to ISAAC_DAILY_LIMIT renders a day (default 10). The USD stage is tested against
MuJoCo's geom positions frame by frame; the Isaac Sim step itself needs the NGC
image and a GPU and has not been run yet.

api/jobs.py

@@ -258,5 +258,21 @@ def _auto_review(s: Services, run: dict) -> None:
         log.info("automatic review of run %s skipped: %s", run["id"], e)
 
 
+def start_isaac(s: Services, run: dict) -> dict:
+    """Render a finished run again with Isaac Sim RTX (a kinematic replay of MuJoCo's motion)."""
+    from .runs import new_run_id
+
+    files = (run.get("output") or {}).get("files") or {}
+    if run["status"] != "COMPLETED" or "trajectory" not in files or "scene" not in files:
+        raise JobError(422, "only a completed run (with its scene and trajectory) can be rendered")
+    mjcf = s.storage.read_text(files["scene"])
+    if not mjcf:
+        raise JobError(422, "the run's scene file is no longer available")
+    job_input = {"mjcf": mjcf, "robot": run.get("robot"), "trajectory_url": s.storage.url(files["trajectory"])}
+    return submit(
+        s, run["tester"], "isaac", job_input, run_id=new_run_id(), parent_id=run["id"], robot=run.get("robot")
+    )
+
+
 # Follow-ups run once when a job of that kind completes.
 ON_COMPLETE: dict[str, list] = {"sim": [_auto_review]}

api/main.py

@@ -28,7 +28,7 @@ from rag.store import RagStore
 from .auth import Allowlist, Tester
 from .codegen import DirectBackend, OpenCodeBackend, generate_code
 from .config import Settings, get_settings
-from .jobs import JobError, month_spend, refresh, start_review
+from .jobs import JobError, month_spend, refresh, start_isaac, start_review
 from .jobs import submit as submit_job
 from .llm import LLM, ChatClient, LLMError
 from .logstore import LogStore
@@ -481,6 +481,19 @@ def create_app(services: Services | None = None) -> FastAPI:
         )
         return _job_response(request, s, job)
 
+    @app.post("/simulate/run/{run_id}/isaac", response_model=JobResponse)
+    def isaac_render(run_id: str, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)):
+        """Re-render a finished run photorealistically with Isaac Sim (limited to ISAAC_DAILY_LIMIT per day)."""
+        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="no simulation run with that id for this key")
+        run = refresh(s, run)
+        try:
+            job = start_isaac(s, run)
+        except JobError as e:
+            raise HTTPException(status_code=e.status, detail=str(e)) from e
+        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/modal_client.py

@@ -10,6 +10,8 @@ from .sim_backend import BackendError
 
 class ModalClient:
     name = "modal"
+    # Functions deployed as their own app (worker/modal_isaac.py), so the big image never slows the main deploy.
+    APP_SUFFIX = {"isaac_render": "-isaac"}
 
     def __init__(self, app_name: str, modal_module: ModuleType | None = None):
         if modal_module is None:
@@ -26,7 +28,7 @@ class ModalClient:
     def submit(self, job_input: dict, execution_timeout_s: int, function: str) -> dict:
         """Start the job and return at once. Time limits are set on the deployed functions (see modal_app.py)."""
         try:
-            fn = self._modal.Function.from_name(self.app_name, function)
+            fn = self._modal.Function.from_name(self.app_name + self.APP_SUFFIX.get(function, ""), function)
             call = fn.spawn(job_input)
         except Exception as e:
             raise BackendError(f"Modal could not start {self.app_name}/{function}: {type(e).__name__}: {e}") from e

api/runs.py

@@ -443,3 +443,27 @@ def _export_body(job: dict, storage: Storage | None) -> list[str]:
 
 
 KIND_BODIES["export"] = _export_body
+
+
+def _isaac_body(job: dict, storage: Storage | None) -> list[str]:
+    output = job.get("output") or {}
+    r = output.get("summary") or {}
+    files = output.get("files") or {}
+    cards = [f'<div class="card"><b>{html.escape(r.get("headline", ""))}</b></div>']
+    if storage is not None and "isaac.mp4" in files:
+        url = html.escape(storage.url(files["isaac.mp4"]))
+        cards.append(
+            f'<div class="card"><b>Isaac Sim render</b><video controls autoplay muted loop playsinline>'
+            f'<source src="{url}" type="video/mp4"></video></div>'
+        )
+    cells = [("Renderer", r.get("renderer")), ("Frames", r.get("frames")), ("Length", f"{r.get('seconds')} s")]
+    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>')
+    cards.append(
+        '<div class="card muted">A replay: MuJoCo simulated the run, Isaac Sim draws the same motion with RTX '
+        "lighting (no second physics engine), so the two videos show the same thing.</div>"
+    )
+    return cards
+
+
+KIND_BODIES["isaac"] = _isaac_body

requirements-dev.txt

@@ -4,3 +4,4 @@
 pytest>=8
 ruff>=0.6
 mcap>=1.1
+usd-core>=25.2  # tests for the Isaac Sim replay stage (worker/usd_scene.py)

tests/test_jobs.py

@@ -95,3 +95,46 @@ def test_follow_up_jobs_listed_on_parent_report(run_env):
     url = client.get(f"/jobs/{run_id}", headers={"X-API-Key": keys_["alice"]}).json()["report_url"]
     page = client.get(url.split("testserver")[-1]).text
     assert "Follow-up jobs" in page and "Robustness sweep" in page and "sw1" in page
+
+
+def test_isaac_render_endpoint(run_env):
+    from dataclasses import replace
+
+    client, keys_, backend, services, gen_id = run_env
+    h = {"X-API-Key": keys_["alice"]}
+    run_id = start_run(client, keys_["alice"], gen_id)
+    assert client.post(f"/simulate/run/{run_id}/isaac", headers=h).status_code == 422  # not finished yet
+    complete(backend, "job-1", run_id)
+    out = backend.jobs["job-1"]["output"]["files"]
+    out["trajectory"] = f"runs/{run_id}/trajectory.npz"
+    out["scene"] = f"runs/{run_id}/scene.xml"
+    (services.storage.root / out["scene"]).parent.mkdir(parents=True, exist_ok=True)
+    (services.storage.root / out["scene"]).write_text("<mujoco/>")
+    r = client.post(f"/simulate/run/{run_id}/isaac", headers=h)
+    assert r.status_code == 200 and r.json()["kind"] == "isaac" and backend.functions[-1] == "isaac_render"
+    assert backend.submitted[-1]["mjcf"] == "<mujoco/>" and backend.submitted[-1]["trajectory_url"].endswith(".npz")
+    services.settings = replace(services.settings, isaac_daily_limit=1)
+    assert client.post(f"/simulate/run/{run_id}/isaac", headers=h).status_code == 429
+
+
+def test_modal_routes_isaac_to_its_own_app():
+    from api.modal_client import ModalClient
+
+    seen = []
+
+    class Fake:
+        class Function:
+            @staticmethod
+            def from_name(app, fn):
+                seen.append((app, fn))
+
+                class F:
+                    def spawn(self, _):
+                        return type("C", (), {"object_id": "fc"})()
+
+                return F()
+
+    client = ModalClient("rosdiff-sim", modal_module=Fake)
+    client.submit({}, 60, "isaac_render")
+    client.submit({}, 60, "sweep")
+    assert seen == [("rosdiff-sim-isaac", "isaac_render"), ("rosdiff-sim", "sweep")]

tests/test_runs.py

@@ -524,3 +524,17 @@ def test_modal_app_definition(monkeypatch):
     assert mod.GPU is None and mod.GL == "osmesa" and mod.app.name == "rosdiff-sim"
     assert isinstance(mod.run_simulation, modal.Function)
     sys.modules.pop("worker.modal_app", None)
+
+
+def test_modal_isaac_app_definition():
+    import importlib
+    import sys
+
+    import modal
+
+    sys.modules.pop("worker.modal_isaac", None)
+    mod = importlib.import_module("worker.modal_isaac")
+    assert mod.app.name == "rosdiff-sim-isaac" and isinstance(mod.isaac_render, modal.Function)
+    main = importlib.import_module("worker.modal_app")
+    for fn in ("run_simulation", "sweep", "ros_check_humble", "ros_check_jazzy", "review_video", "export_lerobot"):
+        assert isinstance(getattr(main, fn), modal.Function), fn

tests/test_usd.py

@@ -0,0 +1,49 @@
+import mujoco
+import numpy as np
+import pytest
+from conftest import GOOD_SCENE
+
+pytest.importorskip("pxr")
+from pxr import Usd, UsdGeom  # noqa: E402
+
+from api.mjcf_check import prepare_scene_dir  # noqa: E402
+from worker.usd_scene import build_stage  # noqa: E402
+
+
+def test_replay_stage_matches_mujoco(tmp_path, menagerie):
+    """Every drawn geom's world position in the USD equals MuJoCo's geom_xpos, frame by frame."""
+    scene_dir = prepare_scene_dir(GOOD_SCENE, "unitree_g1", menagerie)
+    model = mujoco.MjModel.from_xml_path(str(scene_dir / "scene.xml"))
+    data = mujoco.MjData(model)
+    data.qvel[:] = np.random.default_rng(0).normal(0, 1, model.nv)  # tumble things around
+    pos, quat, geom_xpos = [], [], []
+    for _ in range(20):
+        for _ in range(10):
+            mujoco.mj_step(model, data)
+        pos.append(data.xpos.copy())
+        quat.append(data.xquat.copy())
+        geom_xpos.append(data.geom_xpos.copy())
+    names = [mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, b) or f"body_{b}" for b in range(model.nbody)]
+    np.savez(tmp_path / "traj.npz", body_pos=pos, body_quat=quat, body_names=names, fps=30)
+
+    info = build_stage(scene_dir / "scene.xml", tmp_path / "traj.npz", tmp_path / "scene.usdc", "pelvis")
+    assert info["frames"] == 20 and info["geoms_drawn"] >= 3
+    stage = Usd.Stage.Open(str(tmp_path / "scene.usdc"))
+    checked = 0
+    for prim in stage.Traverse():
+        name = prim.GetName()
+        if not (name.startswith("g") and name[1:].isdigit()):
+            continue
+        g = int(name[1:])
+        if model.geom_type[g] == mujoco.mjtGeom.mjGEOM_PLANE:
+            continue
+        for i in (0, 9, 19):
+            m = UsdGeom.Xformable(prim).ComputeLocalToWorldTransform(Usd.TimeCode(i))
+            centre = np.array(m.Transform((0, 0, 0))) if model.geom_type[g] != mujoco.mjtGeom.mjGEOM_MESH else None
+            if centre is not None:
+                assert np.allclose(centre, geom_xpos[i][g], atol=1e-4), (g, i)
+                checked += 1
+    assert checked >= 6
+    cam = UsdGeom.Xformable(stage.GetPrimAtPath("/World/Camera"))
+    eye = np.array(cam.ComputeLocalToWorldTransform(Usd.TimeCode(0)).ExtractTranslation())
+    assert eye[2] > pos[0][1][2]  # looking down at the robot

worker/isaac_render.py

@@ -0,0 +1,63 @@
+"""Render a replay stage (worker/usd_scene.py) with Isaac Sim's RTX renderer, one PNG per frame.
+
+Runs inside Isaac Sim's own Python (the container's /isaac-sim/python.sh), not ours:
+
+    /isaac-sim/python.sh worker/isaac_render.py --usd scene.usdc --out frames/ --frames 300 --fps 30
+
+Nothing is simulated here: the stage already holds every body's pose per frame, so each frame only sets the
+timeline to that frame's time and renders the camera. Standard library + Isaac Sim + Pillow only.
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+
+def main(argv):
+    p = argparse.ArgumentParser()
+    p.add_argument("--usd", required=True)
+    p.add_argument("--out", required=True)
+    p.add_argument("--frames", type=int, required=True)
+    p.add_argument("--fps", type=int, default=30)
+    p.add_argument("--width", type=int, default=1280)
+    p.add_argument("--height", type=int, default=720)
+    p.add_argument("--subframes", type=int, default=4, help="RTX accumulation per frame (quality vs time)")
+    args = p.parse_args(argv)
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+
+    from isaacsim import SimulationApp
+
+    app = SimulationApp({"headless": True, "width": args.width, "height": args.height})
+    try:
+        import omni.replicator.core as rep
+        import omni.timeline
+        import omni.usd
+        from PIL import Image
+
+        if not omni.usd.get_context().open_stage(args.usd):
+            raise RuntimeError(f"Isaac Sim could not open {args.usd}")
+        for _ in range(30):  # let the stage, materials and shaders load
+            app.update()
+        product = rep.create.render_product("/World/Camera", (args.width, args.height))
+        rgb = rep.AnnotatorRegistry.get_annotator("rgb")
+        rgb.attach([product])
+        timeline = omni.timeline.get_timeline_interface()
+        timeline.set_current_time(0.0)
+        rep.orchestrator.step(rt_subframes=16, delta_time=0.0, pause_timeline=True)  # warm-up frame
+        for i in range(args.frames):
+            timeline.set_current_time(i / args.fps)
+            rep.orchestrator.step(rt_subframes=args.subframes, delta_time=0.0, pause_timeline=True)
+            data = rgb.get_data()
+            Image.fromarray(data[:, :, :3]).save(out / f"f{i:05d}.png")
+            if i % 30 == 0:
+                print(f"[isaac_render] frame {i}/{args.frames}", flush=True)
+        rgb.detach()
+        product.destroy()
+    finally:
+        app.close()
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))

worker/modal_isaac.py

@@ -0,0 +1,130 @@
+"""Isaac Sim renders on Modal (L4, RTX), deployed as its own app so the ~20 GB image never slows down the
+main worker deploy:
+
+    modal deploy -m worker.modal_isaac          # app "<MODAL_APP>-isaac", function isaac_render
+
+Isaac Sim's container comes from NVIDIA NGC (ISAAC_IMAGE). If your NGC account needs a login to pull it,
+create a Modal secret with REGISTRY_USERNAME=$oauthtoken and REGISTRY_PASSWORD=<NGC API key> and name it in
+MODAL_NGC_SECRET. Using the container means accepting the NVIDIA Isaac Sim license (ACCEPT_EULA below).
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import tempfile
+import time
+import urllib.request
+from pathlib import Path
+
+import modal
+
+APP_NAME = os.environ.get("MODAL_APP", "rosdiff-sim") + "-isaac"
+ISAAC_IMAGE = os.environ.get("ISAAC_IMAGE", "nvcr.io/nvidia/isaac-sim:5.0.0")
+MENAGERIE_DIR = "/opt/mujoco_menagerie"
+ROBOT_DIRS = ("unitree_g1", "unitree_h1", "unitree_go2", "franka_emika_panda")  # api.robots.ROBOTS
+NGC_SECRET = os.environ.get("MODAL_NGC_SECRET")
+
+image = (
+    modal.Image.from_registry(
+        ISAAC_IMAGE, secret=modal.Secret.from_name(NGC_SECRET) if NGC_SECRET else None, add_python="3.11"
+    )
+    .entrypoint([])
+    .env({"ACCEPT_EULA": "Y", "PRIVACY_CONSENT": "Y", "OMNI_KIT_ACCEPT_EULA": "YES", "MENAGERIE_DIR": MENAGERIE_DIR})
+    .run_commands(
+        "git clone --depth 1 --filter=blob:none --sparse https://github.com/google-deepmind/mujoco_menagerie.git "
+        f"{MENAGERIE_DIR} && cd {MENAGERIE_DIR} && git sparse-checkout set {' '.join(ROBOT_DIRS)} && rm -rf .git"
+    )
+    .pip_install("mujoco==3.14.0", "numpy>=1.26", "usd-core>=25.2", "imageio-ffmpeg==0.6.0")
+    .add_local_python_source("api", "worker")
+)
+
+app = modal.App(APP_NAME)
+
+MAX_SECONDS = 20.0  # longest replay rendered (cost control)
+
+
+@app.function(image=image, gpu="L4", timeout=1800, max_containers=1, scaledown_window=20)
+def isaac_render(job_input: dict) -> dict:
+    import traceback
+
+    try:
+        return _render(job_input)
+    except Exception as e:
+        return {"error": f"{type(e).__name__}: {e}", "traceback": traceback.format_exc()[-2000:]}
+
+
+def _render(job_input: dict) -> dict:
+    import imageio_ffmpeg
+
+    from api.mjcf_check import SCENE_FILE, prepare_scene_dir, static_check
+    from worker.sim_runner import config_for
+    from worker.usd_scene import build_stage
+
+    started = time.monotonic()
+    mjcf, robot = job_input["mjcf"], job_input.get("robot")
+    failed = static_check(mjcf, robot)
+    if failed:
+        raise ValueError(f"scene rejected: {failed.error}")
+    with tempfile.TemporaryDirectory() as tmp:
+        tmp = Path(tmp)
+        scene_dir = prepare_scene_dir(mjcf, robot, Path(MENAGERIE_DIR))
+        traj = tmp / "trajectory.npz"
+        with urllib.request.urlopen(job_input["trajectory_url"], timeout=120) as r:  # noqa: S310
+            traj.write_bytes(r.read())
+        cfg = config_for(robot, 1.0, "hold", render=False)
+        info = build_stage(
+            scene_dir / SCENE_FILE, traj, tmp / "scene.usdc", cfg.root_body, cfg.camera_distance, MAX_SECONDS
+        )
+        frames = tmp / "frames"
+        cmd = [
+            "/isaac-sim/python.sh",
+            str(Path(__file__).with_name("isaac_render.py")),
+            "--usd",
+            str(tmp / "scene.usdc"),
+            "--out",
+            str(frames),
+            "--frames",
+            str(info["frames"]),
+            "--fps",
+            str(info["fps"]),
+        ]
+        proc = subprocess.run(cmd, capture_output=True, text=True, timeout=1500)
+        pngs = sorted(frames.glob("f*.png")) if frames.is_dir() else []
+        if len(pngs) < info["frames"]:
+            raise RuntimeError(f"Isaac Sim rendered {len(pngs)}/{info['frames']} frames: {proc.stderr[-1500:]}")
+        video = tmp / "isaac.mp4"
+        subprocess.run(
+            [
+                imageio_ffmpeg.get_ffmpeg_exe(),
+                "-loglevel",
+                "error",
+                "-framerate",
+                str(info["fps"]),
+                "-i",
+                str(frames / "f%05d.png"),
+                "-c:v",
+                "libx264",
+                "-pix_fmt",
+                "yuv420p",
+                "-crf",
+                "20",
+                "-movflags",
+                "+faststart",
+                str(video),
+            ],
+            check=True,
+            timeout=600,
+        )
+        summary = {
+            "renderer": f"Isaac Sim RTX ({ISAAC_IMAGE.rsplit(':', 1)[-1]})",
+            "frames": info["frames"],
+            "seconds": round(info["frames"] / info["fps"], 2),
+            "geoms_drawn": info["geoms_drawn"],
+            "headline": f"{info['frames']} frames rendered with Isaac Sim RTX",
+        }
+        return {
+            "summary": summary,
+            "artifacts": {"isaac.mp4": video.read_bytes()},
+            "worker_seconds": round(time.monotonic() - started, 1),
+        }

worker/usd_scene.py

@@ -0,0 +1,200 @@
+"""Build an animated USD stage from a MuJoCo scene and a run's trajectory, for rendering in Isaac Sim.
+
+The stage is a kinematic replay: every MuJoCo body becomes an Xform under /World/Bodies whose world pose is
+time-sampled at the trajectory's frame rate, with its visible geoms (MuJoCo's viewer groups 0-2) as children.
+A camera, time-sampled to follow the robot like the MuJoCo video, sits at /World/Camera. No physics runs in
+Isaac Sim; it only renders what MuJoCo simulated, so both videos show exactly the same motion.
+
+Pure `pxr` (usd-core) + mujoco, so it can be built and tested anywhere; Isaac Sim just opens the file.
+"""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import mujoco
+import numpy as np
+
+VISIBLE_GROUPS = (0, 1, 2)
+PLANE_SIZE = 40.0  # metres, for MuJoCo's infinite planes
+
+
+def _name(s: str) -> str:
+    """A valid USD prim name."""
+    out = "".join(c if c.isalnum() or c == "_" else "_" for c in s)
+    return out if out and not out[0].isdigit() else f"_{out}"
+
+
+def _quat(q):
+    from pxr import Gf
+
+    return Gf.Quatf(float(q[0]), Gf.Vec3f(float(q[1]), float(q[2]), float(q[3])))
+
+
+def _material(stage, rgba, cache: dict):
+    from pxr import Gf, Sdf, UsdShade
+
+    key = tuple(round(float(c), 3) for c in rgba)
+    if key in cache:
+        return cache[key]
+    path = f"/World/Looks/m{len(cache)}"
+    mat = UsdShade.Material.Define(stage, path)
+    shader = UsdShade.Shader.Define(stage, f"{path}/Shader")
+    shader.CreateIdAttr("UsdPreviewSurface")
+    shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set(Gf.Vec3f(*key[:3]))
+    shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(key[3])
+    shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(0.6)
+    mat.CreateSurfaceOutput().ConnectToSource(shader.ConnectableAPI(), "surface")
+    cache[key] = mat
+    return mat
+
+
+def _geom_prim(stage, model: mujoco.MjModel, g: int, parent: str):
+    """One USD prim for MuJoCo geom g, placed in its body's frame. Returns None for types we do not draw."""
+    from pxr import Gf, UsdGeom, Vt
+
+    T = mujoco.mjtGeom
+    gtype, size = model.geom_type[g], model.geom_size[g]
+    path = f"{parent}/g{g}"
+    if gtype == T.mjGEOM_PLANE:
+        prim = UsdGeom.Cube.Define(stage, path)
+        prim.CreateSizeAttr(1.0)
+        sx = size[0] * 2 if size[0] > 0 else PLANE_SIZE
+        sy = size[1] * 2 if size[1] > 0 else PLANE_SIZE
+        scale = Gf.Vec3f(sx, sy, 0.002)
+        offset = Gf.Vec3d(0, 0, -0.001)
+    elif gtype == T.mjGEOM_BOX:
+        prim = UsdGeom.Cube.Define(stage, path)
+        prim.CreateSizeAttr(2.0)
+        scale, offset = Gf.Vec3f(*map(float, size)), None
+    elif gtype in (T.mjGEOM_SPHERE, T.mjGEOM_ELLIPSOID):
+        prim = UsdGeom.Sphere.Define(stage, path)
+        prim.CreateRadiusAttr(1.0)
+        r = size if gtype == T.mjGEOM_ELLIPSOID else (size[0],) * 3
+        scale, offset = Gf.Vec3f(*map(float, r)), None
+    elif gtype == T.mjGEOM_CAPSULE:
+        prim = UsdGeom.Capsule.Define(stage, path)
+        prim.CreateRadiusAttr(float(size[0]))
+        prim.CreateHeightAttr(float(2 * size[1]))
+        prim.CreateAxisAttr("Z")
+        scale, offset = None, None
+    elif gtype == T.mjGEOM_CYLINDER:
+        prim = UsdGeom.Cylinder.Define(stage, path)
+        prim.CreateRadiusAttr(float(size[0]))
+        prim.CreateHeightAttr(float(2 * size[1]))
+        prim.CreateAxisAttr("Z")
+        scale, offset = None, None
+    elif gtype == T.mjGEOM_MESH:
+        m = int(model.geom_dataid[g])
+        va, vn = model.mesh_vertadr[m], model.mesh_vertnum[m]
+        fa, fn = model.mesh_faceadr[m], model.mesh_facenum[m]
+        prim = UsdGeom.Mesh.Define(stage, path)
+        prim.CreatePointsAttr(Vt.Vec3fArray.FromNumpy(model.mesh_vert[va : va + vn].astype(np.float32)))
+        prim.CreateFaceVertexCountsAttr(Vt.IntArray([3] * int(fn)))
+        prim.CreateFaceVertexIndicesAttr(
+            Vt.IntArray.FromNumpy(model.mesh_face[fa : fa + fn].reshape(-1).astype(np.int32))
+        )
+        prim.CreateSubdivisionSchemeAttr("none")
+        scale, offset = None, None
+    else:  # heightfields, SDFs: not drawn (they are in the MuJoCo video)
+        return None
+    xf = UsdGeom.Xformable(prim)
+    pos = Gf.Vec3d(*map(float, model.geom_pos[g]))
+    if offset is not None:
+        pos += offset
+    xf.AddTranslateOp().Set(pos)
+    xf.AddOrientOp().Set(_quat(model.geom_quat[g]))
+    if scale is not None:
+        xf.AddScaleOp().Set(scale)
+    return prim
+
+
+def _camera_pose(target: np.ndarray, azimuth: float, elevation: float, distance: float):
+    """Eye position for MuJoCo's free/tracking camera convention (degrees)."""
+    az, el = math.radians(azimuth), math.radians(elevation)
+    forward = np.array([math.cos(el) * math.cos(az), math.cos(el) * math.sin(az), math.sin(el)])
+    return target - distance * forward
+
+
+def build_stage(
+    scene: Path,
+    trajectory: Path,
+    out: Path,
+    root_body: str | None,
+    camera_distance: float = 4.5,
+    max_seconds: float | None = None,
+) -> dict:
+    """Write out (a .usda/.usdc) and return a short description of what is in it."""
+    from pxr import Gf, Usd, UsdGeom, UsdLux, UsdShade
+
+    model = mujoco.MjModel.from_xml_path(str(scene))
+    traj = np.load(trajectory)
+    fps = int(traj["fps"])
+    pos, quat = traj["body_pos"], traj["body_quat"]
+    names = [str(n) for n in traj["body_names"]]
+    if len(names) != model.nbody:
+        raise ValueError("the trajectory does not belong to this scene (body count differs)")
+    frames = len(pos) if max_seconds is None else min(len(pos), int(max_seconds * fps) + 1)
+
+    stage = Usd.Stage.CreateNew(str(out))  # .usdc: binary crate file (much smaller than .usda)
+    UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.z)
+    UsdGeom.SetStageMetersPerUnit(stage, 1.0)
+    stage.SetTimeCodesPerSecond(fps)
+    stage.SetFramesPerSecond(fps)
+    stage.SetStartTimeCode(0)
+    stage.SetEndTimeCode(frames - 1)
+    world = UsdGeom.Xform.Define(stage, "/World")
+    stage.SetDefaultPrim(world.GetPrim())
+    UsdGeom.Scope.Define(stage, "/World/Looks")
+    UsdGeom.Scope.Define(stage, "/World/Bodies")
+
+    materials: dict = {}
+    drawn = 0
+    for b in range(model.nbody):
+        geoms = [g for g in range(model.ngeom) if model.geom_bodyid[g] == b and model.geom_group[g] in VISIBLE_GROUPS]
+        if not geoms:
+            continue
+        path = f"/World/Bodies/{_name(names[b])}_{b}"
+        body = UsdGeom.Xform.Define(stage, path)
+        t_op, o_op = body.AddTranslateOp(), body.AddOrientOp()
+        if b == 0:  # the world body never moves
+            t_op.Set(Gf.Vec3d(0, 0, 0))
+            o_op.Set(Gf.Quatf(1, 0, 0, 0))
+        else:
+            for i in range(frames):
+                t_op.Set(Gf.Vec3d(*map(float, pos[i, b])), Usd.TimeCode(i))
+                o_op.Set(_quat(quat[i, b]), Usd.TimeCode(i))
+        for g in geoms:
+            prim = _geom_prim(stage, model, g, path)
+            if prim is None:
+                continue
+            rgba = model.mat_rgba[model.geom_matid[g]] if model.geom_matid[g] >= 0 else model.geom_rgba[g]
+            UsdShade.MaterialBindingAPI.Apply(prim.GetPrim()).Bind(_material(stage, rgba, materials))
+            drawn += 1
+
+    sun = UsdLux.DistantLight.Define(stage, "/World/Sun")
+    sun.CreateIntensityAttr(3000)
+    sun.CreateAngleAttr(1.0)
+    UsdGeom.Xformable(sun).AddRotateXYZOp().Set(Gf.Vec3f(35, 0, 20))
+    dome = UsdLux.DomeLight.Define(stage, "/World/Sky")
+    dome.CreateIntensityAttr(600)
+
+    cam = UsdGeom.Camera.Define(stage, "/World/Camera")
+    cam.CreateFocalLengthAttr(18.0)
+    cam.CreateClippingRangeAttr(Gf.Vec2f(0.05, 200.0))
+    xf = cam.AddTransformOp()
+    root = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, root_body) if root_body else -1
+    for i in range(frames):
+        target = pos[i, root] if root > 0 else np.asarray(model.stat.center)
+        distance = camera_distance if root > 0 else 1.5 * float(model.stat.extent)
+        eye = _camera_pose(target, 20.0, -22.0, distance)
+        view = Gf.Matrix4d().SetLookAt(Gf.Vec3d(*map(float, eye)), Gf.Vec3d(*map(float, target)), Gf.Vec3d(0, 0, 1))
+        xf.Set(view.GetInverse(), Usd.TimeCode(i))
+        if root <= 0:
+            break  # static camera
+    stage.GetRootLayer().Save()
+    return {"frames": frames, "fps": fps, "bodies": model.nbody, "geoms_drawn": drawn, "materials": len(materials)}
+
+
+__all__ = ["build_stage"]