Kinematic trajectory recording + post-hoc video replay (#296)

[sibocw]

Closes #296.

Decouples how many worlds are simulated from how many are rendered. Instead of rasterizing frames during simulation, a recorder stores only generalized coordinates (qpos, plus mocap poses) at the render cadence. The trajectory is replayed to video later — on CPU or GPU. This sidesteps buffering full (n_worlds, n_cams, H, W, 3) RGB tensors, which is hopeless at thousands of worlds: simulate many in parallel, then sub-select and render a few.

Architecture

Recording — record_trajectory_only=True on set_renderer swaps in a recorder (a recorder is just a renderer whose per-frame output is a qpos slice, not pixels):

• RecordedTrajectory — backend-agnostic dataclass (qpos + optional mocap + replay metadata). .save()/.load() are one self-describing .npz each.
• TrajectoryRecorder(Renderer) — CPU single-world (Simulation).
• WarpTrajectoryRecorder(_BaseWarpRenderer) — GPU multi-world (GPUSimulation), one RecordedTrajectory per selected world, identical format to the CPU recorder.

Replay — set qpos, run position-only kinematics (mj_kinematics + mj_camlight, not a full forward step), rasterize:

• flygym.rendering.render_trajectories(mj_model, trajectories, ...) — CPU.
• flygym.warp.rendering.render_trajectories_gpu(mj_model, trajectories, ...) — GPU batch replay; bin-packs (trajectory, frame) work into mjw.Data batches of worlds_per_batch.

A trajectory recorded on either backend replays on either backend (all four combinations), because it carries no trace of how it was produced.

Serialization — save_trajectories/load_trajectories write/read a folder of individual traj_XXXX.npz files.

Deviations from the design notes in #296

The implementation intentionally diverges from the serialization design sketched in the issue thread; see the issue comment for the full rationale. In short:

• The model is not saved with the trajectory. Trajectories are pure kinematic state; the caller persists the model separately (world.save_xml_with_assets(...)) and passes a compiled MjModel into the replay functions. This is the headline signature change vs. the issue's save_trajectories(trajectories, mj_model, dir) / render_trajectories(source, ...).
• No model hash / no MJB pathway / no manifest. A compiled-model hash can't survive the XML round-trip or the GPU batch-render edits, so it had no working consumer once the model was decoupled; an nq/nmocap shape check guards mismatches instead (robust across both).
• CPU and GPU are separate functions (render_trajectories vs render_trajectories_gpu) rather than one function with a use_gpu flag — matching the flygym / flygym.warp split elsewhere.
• 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 use_gpu_batch_rendering=True GPUSimulation already provides one.
• n_workers CPU multiprocessing was dropped (untested, and re-introduced a model-transport pathway); easy to restore if wanted.

Refactor

rendering.py and warp/rendering.py become packages, each split into live_rendering and recorded_trajectory submodules, with backward-compatible re-exports (from flygym.rendering import Renderer etc. unchanged). docs/SUMMARY.md updated for the new module layout.

Tests

tests/core/test_trajectory.py and tests/warp/test_trajectory.py: CPU/GPU recording, npz round-trip, CPU and GPU replay, incompatible-model guard, and a cross-backend kinematics-identity check (geom_xpos from CPU mj_kinematics vs GPU mjw.kinematics). The CPU replay test asserts a bit-exact match against a qpos-consistent inline render. Full suite: 647 passed, 1 skipped; ruff clean.

🤖 Generated with Claude Code