Pseudo-rendering: Save only state data while "rendering" and generate videos post-hoc

[sibocw]

When simulating many worlds in parallel, we can very easily run out of memory, especially in GPU simulations.

One solution is to add a "pseudo-renderer," which records only xpos, xquat, cam_xpos, cam_xmat, and other related data when render_as_needed is called. Then, instead of saving frames to an actual MP4 file, this render generates a PseudoRenderOutput object, which can be serialized/deserialized into/from NPZ files.

Then there should a separate functions, materialize_pseudo_render_output(output_objects: list[PseudoRenderOutput] | PseudoRenderOutput, n_workers: int, output_path: Path, use_gpu: bool) that runs the "simulation" again, this time with mocap only (no physics), and actually generates the videos. This decouples the number of worlds simulated in parallel from the number of worlds rendered in parallel: for example, we can simulate 8000 worlds in parallel for model training, but select a smaller number of them posthoc and materialize the videos by rendering only 10 worlds in parallel.

Notes:

1. materialize_pseudo_render_output should parallelize over both PseudoRenderOutput instances and time: the total number of frames is sum(x.n_frames for x in output_objects) (the output objects don't have to share the same length). The function should parallelize over these many steps.
2. This should support both CPU and GPU rendering.
3. Probably the mjmodel also needs to be saved, so probably a compressed Pickle file is better than NPZ because we can save the model and the history in one file...
4. I'm not sure if "pseudo-renderer" and "materialize" are good names for this... open to suggestions.

---

Update (design conclusions)

After some discussion, here are the refined conclusions. The core idea is unchanged and confirmed worthwhile; the memory win is the whole point (today the Warp renderer buffers full (n_worlds, n_cams, H, W, 3) RGB tensors per frame, which is hopeless at thousands of worlds). The changes below concern what we record, how we serialize, and naming.

1. Record qpos, not derived pose fields

Record the generalized coordinates (qpos, plus mocap_pos/mocap_quat if mocap bodies are used) — not xpos/xquat/cam_xpos/cam_xmat.

• xpos/cam_xpos/etc. are derived and an incomplete basis for re-rendering: mjv_updateScene builds the scene from geom_xpos/geom_xmat/site_xpos/light poses, not body xpos/xquat.
• At materialize time, set qpos and run position-only kinematics (mj_kinematics + mj_camlight, or mj_fwdPosition; not full mj_forward — we don't need velocities/accelerations/contacts). This deterministically regenerates all geom/site/camera/light transforms.
• Why not store geom poses and skip kinematics? Writing geom_xpos/geom_xmat directly into mjData does work (mjv_updateScene and Warp's refit_bvh read those fields directly), and it skips the kinematics traversal — but the compute saved (microseconds of kinematics) is negligible next to rasterization + mp4 encode (milliseconds+), while the storage cost balloons: qpos is ~O(100) floats/frame, vs. 12 × ngeom ≈ thousands of floats/frame for the fly. qpos is smaller, complete, and deterministic.
• Note on terminology: "mocap only, no physics" really means "run mj_forward/position-kinematics but never mj_step." Setting mocap_pos/mocap_quat is itself an input to mj_kinematics; it doesn't bypass forward kinematics, it requires it.

Precondition / main risk: all parallel worlds must share one compiled model. This holds today (no per-world domain randomization in the codebase). If per-world morphology randomization is ever added, a single saved model silently renders every replay with the wrong geometry — so we guard this with a model hash (below) and document the assumption.

2. Serialization: a self-describing folder, model saved once, guarded by a hash

Drop the single-pickle idea. Pickling a live MjModel is fragile across MuJoCo versions and forces reading the whole blob to get anything — which fights the parallel-materialize goal (workers want to lazily load only the worlds/frame-ranges they need; NPZ is mmap-able/per-array, a monolithic pickle is not).

Instead, an output folder containing:

• model as mj_saveModel MJB bytes (version-portable, explicit) — saved once per folder, regardless of how many trajectories.
• state as one .npz per trajectory (qpos arrays; ragged lengths across trajectories are fine — one file each; optional mocap_*).
• manifest as manifest.json: output_fps, playback_speed, camera names, world_ids, camera_res, scene_option, and the model hash.

Model hash: each RecordedTrajectory stores a fingerprint of the model it was recorded against (hash of the mj_saveModel MJB buffer, cached at recorder construction so it isn't recomputed per frame). save_trajectories asserts all trajectories' hashes match the passed model; render_trajectories re-checks against the model file in the folder — so hand-mixing files from different runs fails loudly instead of rendering garbage.

API (free functions, not .save() methods) — the model is the shared invariant across N trajectories, so a free function is the natural owner of the "one model + many trajectories" relationship:

def save_trajectories(
    trajectories: RecordedTrajectory | list[RecordedTrajectory],
    mj_model: mj.MjModel,
    output_dir: Path,
) -> None: ...

def render_trajectories(
    source: Path | tuple[list[RecordedTrajectory], mj.MjModel],
    output_path: Path,
    *,
    n_workers: int = 1,
    use_gpu: bool = False,
    worlds_per_batch: int | None = None,  # GPU batch size
) -> None: ...

render_trajectories accepts either a self-describing folder (primary, ergonomic, symmetric with save_trajectories) or an in-memory (trajectories, mj_model) pair (no disk round-trip). We drop the "loose list of files, check they're in the same dir" variant — it's error-prone and dominated by the folder form.

3. Naming

• TrajectoryRecorder — records state instead of pixels for CPU single-world Simulation (subclass of Renderer).
• WarpTrajectoryRecorder — same, for GPU multi-world GPUSimulation (subclass of _BaseWarpRenderer).
• RecordedTrajectory — the output object (one per world, holds qpos history + model hash + metadata).
• save_trajectories() / render_trajectories() — the free functions.

Drops the "pseudo"/"materialize" jargon in favor of what it actually is: recording a kinematic trajectory and replaying it.

4. Recording source and render backend are fully decoupled

Both recorders emit the same RecordedTrajectory format (CPU produces one; GPU produces a list, one per world). A RecordedTrajectory is just qpos history + model hash + metadata — it carries no trace of which backend produced it. So all four combinations work:

Recorded by	Rendered by render_trajectories
TrajectoryRecorder (CPU Simulation)	CPU or GPU
WarpTrajectoryRecorder (GPU GPUSimulation)	CPU or GPU

This is the point of recording qpos against a shared model: the trajectory is backend-agnostic, so e.g. a cheap CPU run can be replayed with GPU batch rendering, or thousands of GPU-simulated worlds can be sub-selected and replayed on a few CPU workers.

---

Concrete implementation plan

Where it slots in. Both renderer hierarchies already funnel everything through render_as_needed → buffer → save_video. A recorder is just a renderer whose per-frame output is a qpos slice instead of a rasterized image. So these are new subclasses, not a new pipeline — one per simulation backend, mirroring the existing Renderer / _BaseWarpRenderer split.

Class structure

Renderer (rendering.py, CPU/single-world)
├── TrajectoryRecorder        # NEW — records qpos (+mocap) for CPU `Simulation`, no pixels
└── _BaseWarpRenderer (warp/rendering.py, ABC; render_as_needed/buffer/save plumbing)
    ├── WarpGPUBatchRenderer   # existing — renders pixels via mjw batch render
    ├── WarpCPURenderer        # existing — renders pixels via per-world mj.Renderer
    └── WarpTrajectoryRecorder # NEW — records qpos (+mocap) for GPU `GPUSimulation`, no pixels

Both recorders override the render hook to buffer state instead of pixels, and both expose the recorded state as attributes (in the same RecordedTrajectory format), differing only in cardinality to match the simulation backend:

Recorder	Backend	Attributes
TrajectoryRecorder	CPU single-world Simulation	mj_model: MjModel, recorded_trajectory: RecordedTrajectory
WarpTrajectoryRecorder	GPU multi-world GPUSimulation	mj_model: MjModel, recorded_trajectories: list[RecordedTrajectory] (one per world)

(mj_model is already inherited from Renderer.) The singular-vs-list attribute mirrors single-world vs. multi-world; otherwise the contained RecordedTrajectory objects are identical, so downstream save_trajectories / render_trajectories treat both uniformly (the CPU recorder's single trajectory is just passed as a 1-element list).

TrajectoryRecorder(Renderer) (CPU, single-world, used by Simulation):

• __init__: skip/short-circuit the mj.Renderer allocation (no rasterization); cache nq, the qpos layout, and the model hash.
• override render_as_needed(mj_data): at the inherited cadence, append mj_data.qpos.copy() (+ mocap_* if present) to the buffer backing recorded_trajectory.

WarpTrajectoryRecorder(_BaseWarpRenderer) (GPU, multi-world, used by GPUSimulation):

• _render_setup_impl: tear down the inherited mj.Renderer/scene_option; cache nq, layout, model hash.
• _render_impl(mjw_data): return mjw_data.qpos.numpy()[self.world_ids] (shape (n_worlds, nq)), plus mocap_* if present — the buffer entry (replacing the RGB tensor) that is split per world into recorded_trajectories.

Both reuse the inherited render_as_needed cadence verbatim, so recorded frames are pre-thinned to exactly the render timepoints (store the fps/cadence in the manifest so replay timing matches). The model hash is computed identically (hash(mj_saveModel bytes)) so trajectories from either backend are interchangeable downstream.

RecordedTrajectory (plain dataclass):

• qpos: np.ndarray (n_frames, nq), optional mocap_pos/mocap_quat, model_hash: str, world_id: int, and render metadata (fps, cameras, camera_res).

Materialize (render_trajectories)

Flatten all (trajectory, frame) pairs into one work-list (handles variable lengths via note 1), then by backend:

• GPU (use_gpu=True): bin-pack the work-list into batches of worlds_per_batch states. Per batch: stage the qpos rows into an mjw.Data of nworld=batch_size, run Warp position-kinematics (mjw.kinematics + mjw.camlight), refit_bvh, mjw.render, then scatter frames back to their (trajectory, frame) slots. Reuses WarpGPUBatchRenderer's render path.
• CPU (use_gpu=False): embarrassingly parallel over frames/trajectories via multiprocessing (n_workers); each worker owns an mj.Renderer + MjData, sets qpos, runs mj_kinematics/mj_camlight, renders. Mirrors WarpCPURenderer._render_impl.
• Encode stage: reassemble frames in order per trajectory, write mp4 (libx264) — sequential per video, parallel across videos. Reuse write_video_from_frames.

Build order

1. RecordedTrajectory dataclass + model-hash helper (mj_saveModel → bytes → hash).
2. WarpTrajectoryRecorder(_BaseWarpRenderer) — _render_impl returns qpos; verify it drops into existing GPU training loops as a renderer swap.
3. TrajectoryRecorder(Renderer) — CPU single-world recorder; same RecordedTrajectory output format.
4. save_trajectories / loader (folder: MJB + per-trajectory npz + manifest.json; hash checks).
5. render_trajectories CPU path (simpler, correctness baseline).
6. render_trajectories GPU batch path (bin-packing + Warp kinematics + batch render).
7. Tests: round-trip short recordings from both backends; assert replayed frames ≈ frames rendered inline at sim time (same camera/cadence), and that a CPU-recorded trajectory and a GPU-recorded one render identically; assert hash mismatch raises.

[sibocw]

Implemented in #297.

The core idea landed as designed (record qpos at the render cadence, replay with position-only kinematics, backend-agnostic trajectories), but the serialization/API converged somewhere different from the design notes above after building it out. Summary of what shipped and why it differs:

What shipped

• RecordedTrajectory — qpos (+ optional mocap) + replay metadata. .save()/.load() are one self-describing .npz each.
• TrajectoryRecorder(Renderer) (CPU) and WarpTrajectoryRecorder(_BaseWarpRenderer) (GPU, one trajectory per world) — same format, swapped in via set_renderer(..., record_trajectory_only=True).
• flygym.rendering.render_trajectories(mj_model, trajectories, ...) (CPU) and flygym.warp.rendering.render_trajectories_gpu(mj_model, trajectories, ...) (GPU batch, worlds_per_batch).
• save_trajectories/load_trajectories over a folder of traj_XXXX.npz files.
• All four record/replay backend combinations work.

Deviations from the design conclusions above

1. The model is no longer saved with the trajectory. A trajectory is pure kinematic state. The caller persists the model separately with save_xml_with_assets (self-contained model.xml + bundled meshes) and passes a compiled MjModel into replay. So the signatures became save_trajectories(trajectories, output_dir) and render_trajectories(mj_model, trajectories, output_path, ...) — dropping the mj_model arg from save and the folder-or-tuple source union from render.

2. No model hash, no MJB, no manifest. The hash was meant to guard "wrong model", but a compiled-model hash doesn't survive the XML round-trip or the GPU batch-render edits, so it had no working consumer once the model was decoupled. Replaced by an nq/nmocap shape check, which is robust across both. (This also retires the model_hash/model_dir_hash question from earlier in the thread — now zero hashes.)

3. CPU and GPU are separate functions (render_trajectories / render_trajectories_gpu) instead of one use_gpu flag, keeping flygym and flygym.warp cleanly delineated like the rest of the codebase.

4. GPU batch-render prep is the caller's responsibility. render_trajectories_gpu expects a batch-ready model (modify_world_for_batch_rendering, now exported from flygym.warp); a GPUSimulation built with use_gpu_batch_rendering=True already provides one. The old hidden reload→modify→recompile is gone.

5. n_workers CPU multiprocessing dropped — untested, and it reintroduced a model-transport pathway now that the model isn't on disk next to the trajectories. Trivial to add back with an MjModel-bytes transport if we want it.

The naming (TrajectoryRecorder / WarpTrajectoryRecorder / RecordedTrajectory / save_trajectories / render_trajectories) stuck.

Also refactored rendering.py and warp/rendering.py into packages split into live_rendering + recorded_trajectory, with backward-compatible re-exports. Tests cover both backends, npz round-trip, CPU+GPU replay, and cross-backend kinematics identity. Full suite: 647 passed, 1 skipped.