From 8ba6217ce99e6fe112b4fd0ddd4b2c30f5fbe398 Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Mon, 29 Jun 2026 21:46:23 +0200 Subject: [PATCH 01/11] Add kinematic trajectory recording + post-hoc video replay (#296) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Decouple how many worlds are *simulated* from how many are *rendered*: a recorder stores only qpos (plus mocap poses) at the render cadence instead of rasterizing frames, and the trajectory is replayed to video later on CPU or GPU. This lets e.g. thousands of GPU-simulated worlds be sub-selected and rendered with a small batch, sidestepping the full RGB-tensor buffering that is hopeless at scale. Recording: - RecordedTrajectory: backend-agnostic dataclass (qpos + mocap + replay metadata); RecordedTrajectory.save/.load are one self-describing .npz each. - TrajectoryRecorder(Renderer): CPU single-world recorder. - WarpTrajectoryRecorder(_BaseWarpRenderer): GPU multi-world recorder, one RecordedTrajectory per selected world; same format as the CPU recorder. - Simulation/GPUSimulation.set_renderer(..., record_trajectory_only=True) swaps in the recorder (a recorder is a kind of renderer). Replay (model supplied by the caller, persisted separately via save_xml_with_assets): - 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. - Both set qpos and run position-only kinematics (mj_kinematics + mj_camlight) rather than a full forward step. Serialization: save_trajectories/load_trajectories read/write a folder of individual traj_XXXX.npz files; the model is not stored with the trajectory. Refactor rendering.py and warp/rendering.py into packages, each split into live_rendering and recorded_trajectory submodules, with backward-compatible re-exports. Tests cover CPU/GPU recording, npz round-trip, CPU and GPU replay, and cross-backend kinematics identity. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/SUMMARY.md | 8 +- src/flygym/__init__.py | 16 +- src/flygym/rendering/__init__.py | 35 ++ .../live_rendering.py} | 34 +- src/flygym/rendering/recorded_trajectory.py | 437 ++++++++++++++++++ src/flygym/simulation.py | 43 +- src/flygym/warp/__init__.py | 17 +- src/flygym/warp/rendering/__init__.py | 24 + .../live_rendering.py} | 26 +- .../warp/rendering/recorded_trajectory.py | 287 ++++++++++++ src/flygym/warp/simulation.py | 33 +- tests/core/test_trajectory.py | 254 ++++++++++ tests/warp/test_trajectory.py | 174 +++++++ 13 files changed, 1356 insertions(+), 32 deletions(-) create mode 100644 src/flygym/rendering/__init__.py rename src/flygym/{rendering.py => rendering/live_rendering.py} (95%) create mode 100644 src/flygym/rendering/recorded_trajectory.py create mode 100644 src/flygym/warp/rendering/__init__.py rename src/flygym/warp/{rendering.py => rendering/live_rendering.py} (94%) create mode 100644 src/flygym/warp/rendering/recorded_trajectory.py create mode 100644 tests/core/test_trajectory.py create mode 100644 tests/warp/test_trajectory.py diff --git a/docs/SUMMARY.md b/docs/SUMMARY.md index 08b3bd53..4e96fa49 100644 --- a/docs/SUMMARY.md +++ b/docs/SUMMARY.md @@ -35,7 +35,9 @@ * [flat_ground](api_reference/flygym/compose/world/flat_ground.md) * [musculoskeletal](api_reference/flygym/compose/world/musculoskeletal.md) * [tethered_world](api_reference/flygym/compose/world/tethered_world.md) - * [rendering](api_reference/flygym/rendering.md) + * rendering + * [live_rendering](api_reference/flygym/rendering/live_rendering.md) + * [recorded_trajectory](api_reference/flygym/rendering/recorded_trajectory.md) * [simulation](api_reference/flygym/simulation.md) * utils * [api1to2](api_reference/flygym/utils/api1to2.md) @@ -47,7 +49,9 @@ * [profiling](api_reference/flygym/utils/profiling.md) * [video](api_reference/flygym/utils/video.md) * warp - * [rendering](api_reference/flygym/warp/rendering.md) + * rendering + * [live_rendering](api_reference/flygym/warp/rendering/live_rendering.md) + * [recorded_trajectory](api_reference/flygym/warp/rendering/recorded_trajectory.md) * [simulation](api_reference/flygym/warp/simulation.md) * [utils](api_reference/flygym/warp/utils.md) * [NeuroMechFly Game & Outreach](outreach.md) diff --git a/src/flygym/__init__.py b/src/flygym/__init__.py index 9676d5cc..728e113e 100644 --- a/src/flygym/__init__.py +++ b/src/flygym/__init__.py @@ -9,7 +9,16 @@ from . import compose # noqa: E402 from . import flybody # noqa: E402 from .simulation import Simulation # noqa: E402 -from .rendering import Renderer, launch_interactive_viewer, preview_model # noqa: E402 +from .rendering import ( # noqa: E402 + Renderer, + TrajectoryRecorder, + RecordedTrajectory, + save_trajectories, + load_trajectories, + render_trajectories, + launch_interactive_viewer, + preview_model, +) __all__ = [ "assets_dir", @@ -18,6 +27,11 @@ "flybody", "Simulation", "Renderer", + "TrajectoryRecorder", + "RecordedTrajectory", + "save_trajectories", + "load_trajectories", + "render_trajectories", "launch_interactive_viewer", "preview_model", ] diff --git a/src/flygym/rendering/__init__.py b/src/flygym/rendering/__init__.py new file mode 100644 index 00000000..f02a96b3 --- /dev/null +++ b/src/flygym/rendering/__init__.py @@ -0,0 +1,35 @@ +"""MuJoCo rendering: live rasterization and recorded-trajectory replay. + +This package is split into: + +- `flygym.rendering.live_rendering`: the real-time `Renderer` and viewer helpers. +- `flygym.rendering.recorded_trajectory`: recording ``qpos`` trajectories and + replaying them to video on the CPU. + +All public names are re-exported here for backward compatibility, so +``from flygym.rendering import Renderer`` (etc.) keeps working. +""" + +from flygym.rendering.live_rendering import ( + Renderer, + launch_interactive_viewer, + preview_model, +) +from flygym.rendering.recorded_trajectory import ( + RecordedTrajectory, + TrajectoryRecorder, + save_trajectories, + load_trajectories, + render_trajectories, +) + +__all__ = [ + "Renderer", + "TrajectoryRecorder", + "RecordedTrajectory", + "save_trajectories", + "load_trajectories", + "render_trajectories", + "launch_interactive_viewer", + "preview_model", +] diff --git a/src/flygym/rendering.py b/src/flygym/rendering/live_rendering.py similarity index 95% rename from src/flygym/rendering.py rename to src/flygym/rendering/live_rendering.py index d8c709e3..ad1dc2a2 100644 --- a/src/flygym/rendering.py +++ b/src/flygym/rendering/live_rendering.py @@ -1,3 +1,9 @@ +"""Live (real-time) MuJoCo rendering: the `Renderer` and viewer helpers. + +For recording kinematic trajectories instead of rasterizing frames, and replaying +them to video, see `flygym.rendering.recorded_trajectory`. +""" + import warnings from multiprocessing import Process from pathlib import Path @@ -65,7 +71,7 @@ def __init__( nrows, ncols = camera_res self.buffer_frames = buffer_frames - self.mj_renderer = mj.Renderer(mj_model, nrows, ncols, **kwargs) + self.mj_renderer = self._build_mj_renderer(mj_model, nrows, ncols, **kwargs) # RGB / depth / segmentation are independent and may be enabled in any # combination. Each is produced by its own render() pass (mujoco renders # one output type per call), so we don't enable a mode here -- the mode is @@ -114,9 +120,28 @@ def __init__( # Avoid floating point issues when comparing times self.rendering_rounding_tolerance = mj_model.opt.timestep * 0.5 + def _build_mj_renderer( + self, mj_model: mj.MjModel, nrows: int, ncols: int, **kwargs: Any + ) -> mj.Renderer: + """Construct the underlying ``mujoco.Renderer``. + + Factored out so subclasses that never rasterize (e.g. `TrajectoryRecorder`) + can skip allocating a GL/EGL context by overriding this to return None. + """ + return mj.Renderer(mj_model, nrows, ncols, **kwargs) + def _new_frame_buffer(self) -> dict[str, list]: return {cam_name: [] for cam_name in self._cameras_names2id} + def _due_for_render(self, time: float) -> bool: + """Whether enough time has elapsed since the last render for the next one.""" + min_next_render_time = ( + self._last_render_time_sec + + self._secs_between_renders + - self.rendering_rounding_tolerance + ) + return time >= min_next_render_time + def get_camera_matrix( self, camera: str | mj.MjsCamera, mj_data: mj.MjData, mj_model: mj.MjModel ) -> np.ndarray: @@ -159,12 +184,7 @@ def render_as_needed(self, mj_data: mj.MjData) -> bool: Returns: True if frames were rendered, False otherwise. """ - min_next_render_time = ( - self._last_render_time_sec - + self._secs_between_renders - - self.rendering_rounding_tolerance - ) - if mj_data.time < min_next_render_time: + if not self._due_for_render(mj_data.time): return False self._last_render_time_sec = float(mj_data.time) diff --git a/src/flygym/rendering/recorded_trajectory.py b/src/flygym/rendering/recorded_trajectory.py new file mode 100644 index 00000000..25eb6651 --- /dev/null +++ b/src/flygym/rendering/recorded_trajectory.py @@ -0,0 +1,437 @@ +"""Record minimal kinematic state during simulation and replay it as video (CPU). + +This decouples how many worlds are *simulated* in parallel from how many are +*rendered*. During simulation a `TrajectoryRecorder` stores only the generalized +coordinates (``qpos``, plus mocap poses if used) at the render cadence, instead of +rasterizing frames. The result is a backend-agnostic `RecordedTrajectory`, which can +also be produced on GPU (`flygym.warp.rendering.WarpTrajectoryRecorder`) and replayed +on either backend (`render_trajectories` here, or +`flygym.warp.rendering.render_trajectories_gpu`). + +A trajectory stores kinematic state only -- never a model. To replay one you pass a +compiled `mujoco.MjModel` yourself. Persist the model separately (e.g. with +`BaseCompositionElement.save_xml_with_assets`, a self-contained ``model.xml`` plus +bundled meshes) and recompile it when you need it. The recorded ``qpos`` layout is +preserved across an XML round trip and across `modify_world_for_batch_rendering` +(material/texture/light edits only), so one trajectory replays against the saved model +on either backend. + +At replay time we set ``qpos`` and run *position-only* kinematics (``mj_kinematics`` + +``mj_camlight``) -- not a full ``mj_forward`` -- which deterministically regenerates +every geom/site/camera/light transform the renderer reads. See issue #296. + +Note on faithfulness: a replayed frame shows geometry *consistent with the recorded +qpos*. A frame rendered live during simulation instead shows geometry from the +forward-kinematics pass at the *start* of the last step (one integration step stale -- +a standard MuJoCo quirk), so replay and live render differ by at most one timestep of +motion; the replay is the qpos-faithful one. +""" + +from dataclasses import dataclass +from pathlib import Path +from typing import Any +from os import PathLike + +import mujoco as mj +import numpy as np + +from flygym.rendering.live_rendering import Renderer +from flygym.utils.video import write_video_from_frames + + +__all__ = [ + "RecordedTrajectory", + "TrajectoryRecorder", + "save_trajectories", + "load_trajectories", + "render_trajectories", +] + + +@dataclass +class RecordedTrajectory: + """A single world's recorded kinematic trajectory. + + Holds the per-frame generalized coordinates needed to re-render a world, plus the + metadata describing how it should be rendered. It carries no trace of which backend + produced it nor of the model it came from, so a trajectory recorded on CPU + (`TrajectoryRecorder`) and one recorded on GPU + (`flygym.warp.rendering.WarpTrajectoryRecorder`) are interchangeable inputs to + `render_trajectories` / `flygym.warp.rendering.render_trajectories_gpu`, given a + compatible compiled model. + + Attributes: + qpos: ``(n_frames, nq)`` generalized coordinates, one row per recorded frame. + output_fps: Frame rate the frames were sampled at / should be encoded at. + playback_speed: Playback speed relative to real time (metadata only). + camera_names: Cameras to render by default at replay time. + camera_res: ``(height, width)`` in pixels. + world_id: Index of the world this trajectory came from (0 for CPU). + mocap_pos: ``(n_frames, nmocap, 3)`` mocap positions, or None if the model + has no mocap bodies. + mocap_quat: ``(n_frames, nmocap, 4)`` mocap quaternions, or None. + """ + + qpos: np.ndarray + output_fps: int + playback_speed: float + camera_names: list[str] + camera_res: tuple[int, int] + world_id: int = 0 + mocap_pos: np.ndarray | None = None + mocap_quat: np.ndarray | None = None + + @property + def n_frames(self) -> int: + return int(self.qpos.shape[0]) + + @property + def has_mocap(self) -> bool: + return self.mocap_pos is not None + + def save(self, path: PathLike) -> None: + """Save this trajectory to a single self-describing ``.npz`` file. + + Both the per-frame state and the replay metadata are stored, so the file can be + read back with `load` without any side information. + """ + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + arrays: dict[str, np.ndarray] = { + "qpos": self.qpos, + "output_fps": np.asarray(self.output_fps), + "playback_speed": np.asarray(self.playback_speed), + "camera_names": np.asarray(self.camera_names, dtype=np.str_), + "camera_res": np.asarray(self.camera_res), + "world_id": np.asarray(self.world_id), + } + if self.has_mocap: + arrays["mocap_pos"] = self.mocap_pos + arrays["mocap_quat"] = self.mocap_quat + np.savez_compressed(path, **arrays) + + @classmethod + def load(cls, path: PathLike) -> "RecordedTrajectory": + """Load a trajectory from a ``.npz`` file written by `save`.""" + with np.load(path, allow_pickle=False) as data: + has_mocap = "mocap_pos" in data + return cls( + qpos=data["qpos"], + output_fps=int(data["output_fps"]), + playback_speed=float(data["playback_speed"]), + camera_names=[str(c) for c in data["camera_names"]], + camera_res=tuple(int(x) for x in data["camera_res"]), + world_id=int(data["world_id"]), + mocap_pos=data["mocap_pos"] if has_mocap else None, + mocap_quat=data["mocap_quat"] if has_mocap else None, + ) + + +class TrajectoryRecorder(Renderer): + """Records ``qpos`` (and mocap poses) instead of rasterizing frames. + + A drop-in replacement for `Renderer` on the CPU single-world `Simulation`: it + reuses the same render cadence so recorded frames land at exactly the timepoints + that would have been rendered, but stores only generalized coordinates. The result + is exposed as `recorded_trajectory` and can be replayed to video later by + `render_trajectories` (you supply the compiled model). + + Args: + mj_model: Compiled MuJoCo model. + cameras: Camera(s) to record as the default render cameras. Recording itself + does not depend on the cameras; they are stored as replay metadata. + camera_res: ``(height, width)`` in pixels (replay metadata). + playback_speed: Video playback speed relative to real time. + output_fps: Output video frame rate (also sets the recording cadence). + + Attributes: + recorded_trajectory: The `RecordedTrajectory` accumulated so far. + """ + + def __init__( + self, + mj_model: mj.MjModel, + cameras: str | mj.MjsCamera | list[str | mj.MjsCamera], + *, + camera_res: tuple[int, int] = (240, 320), + playback_speed: float = 0.2, + output_fps: int = 25, + scene_option: mj.MjvOption | None = None, + **kwargs: Any, + ): + # buffer_frames=False so the base class allocates no pixel buffers; we keep + # our own qpos buffer instead. + super().__init__( + mj_model, + cameras, + camera_res=camera_res, + playback_speed=playback_speed, + output_fps=output_fps, + buffer_frames=False, + scene_option=scene_option, + **kwargs, + ) + self._nq = mj_model.nq + self._nmocap = mj_model.nmocap + self._qpos_buf: list[np.ndarray] = [] + self._mocap_pos_buf: list[np.ndarray] = [] + self._mocap_quat_buf: list[np.ndarray] = [] + + def _build_mj_renderer( + self, mj_model: mj.MjModel, nrows: int, ncols: int, **kwargs: Any + ) -> None: + # No rasterization: skip the GL/EGL context allocation entirely. + return None + + def render_as_needed(self, mj_data: mj.MjData) -> bool: + """Record the current state if enough simulation time has elapsed. + + Returns: + True if a frame of state was recorded, False otherwise. + """ + if not self._due_for_render(mj_data.time): + return False + self._last_render_time_sec = float(mj_data.time) + self._qpos_buf.append(mj_data.qpos.copy()) + if self._nmocap > 0: + self._mocap_pos_buf.append(mj_data.mocap_pos.copy()) + self._mocap_quat_buf.append(mj_data.mocap_quat.copy()) + return True + + def reset(self) -> None: + """Clear the recorded state and reset the render timer.""" + self._last_render_time_sec = -np.inf + self._qpos_buf = [] + self._mocap_pos_buf = [] + self._mocap_quat_buf = [] + + @property + def recorded_trajectory(self) -> RecordedTrajectory: + """The trajectory recorded so far (one world).""" + if len(self._qpos_buf) == 0: + raise RuntimeError("No frames have been recorded yet.") + return RecordedTrajectory( + qpos=np.asarray(self._qpos_buf), + output_fps=self.output_fps, + playback_speed=self.playback_speed, + camera_names=list(self._cameras_names2id.keys()), + camera_res=self.camera_res, + world_id=0, + mocap_pos=np.asarray(self._mocap_pos_buf) if self._nmocap > 0 else None, + mocap_quat=np.asarray(self._mocap_quat_buf) if self._nmocap > 0 else None, + ) + + def close(self) -> None: + """No-op: no renderer resources are held.""" + return + + def show_in_notebook(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError( + "TrajectoryRecorder records state, not frames. Save it with " + "save_trajectories and replay with render_trajectories." + ) + + def save_video(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError( + "TrajectoryRecorder records state, not frames. Save it with " + "save_trajectories and replay with render_trajectories." + ) + + +def save_trajectories( + trajectories: RecordedTrajectory | list[RecordedTrajectory], + output_dir: PathLike, +) -> None: + """Save trajectories as individual ``.npz`` files in a folder. + + Writes one ``traj_XXXX.npz`` per trajectory (see `RecordedTrajectory.save`); each + file is self-describing. The model is intentionally *not* saved here -- persist it + yourself (e.g. ``world.save_xml_with_assets(...)``) and recompile it at replay time. + + Args: + trajectories: One trajectory or a list of them. + output_dir: Destination folder (created if needed). + """ + trajectories = _as_trajectory_list(trajectories) + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + for i, traj in enumerate(trajectories): + traj.save(output_dir / f"traj_{i:04d}.npz") + + +def load_trajectories(source: PathLike) -> list[RecordedTrajectory]: + """Load all trajectories from a folder written by `save_trajectories`. + + Returns the trajectories only; supply the compiled model yourself at replay time. + """ + source = Path(source) + files = sorted(source.glob("traj_*.npz")) + if len(files) == 0: + raise ValueError(f"No trajectory files (traj_*.npz) found in {source}.") + return [RecordedTrajectory.load(f) for f in files] + + +def _as_trajectory_list( + trajectories: RecordedTrajectory | list[RecordedTrajectory], +) -> list[RecordedTrajectory]: + if isinstance(trajectories, RecordedTrajectory): + return [trajectories] + trajectories = list(trajectories) + if len(trajectories) == 0: + raise ValueError("No trajectories given.") + return trajectories + + +def _check_model_compatible( + mj_model: mj.MjModel, trajectories: list[RecordedTrajectory] +) -> None: + """Raise if a trajectory's ``qpos``/mocap layout does not fit ``mj_model``. + + A cheap, round-trip-robust guard (unlike a compiled-model hash, the ``qpos`` layout + survives an XML round trip and the batch-render edits): it catches replaying against + a structurally different model before any rendering work. + """ + for i, traj in enumerate(trajectories): + if traj.qpos.shape[1] != mj_model.nq: + raise ValueError( + f"Trajectory {i} (world {traj.world_id}) has nq={traj.qpos.shape[1]} " + f"but the model has nq={mj_model.nq}; it was recorded against a " + "different model." + ) + if traj.has_mocap and traj.mocap_pos.shape[1] != mj_model.nmocap: + raise ValueError( + f"Trajectory {i} (world {traj.world_id}) has " + f"nmocap={traj.mocap_pos.shape[1]} but the model has " + f"nmocap={mj_model.nmocap}; it was recorded against a different model." + ) + + +def _resolve_render_output_paths( + trajectories: list[RecordedTrajectory], + camera_names: list[str], + output_path: PathLike, +) -> list[dict[str, Path]]: + """Map each (trajectory, camera) to an output mp4 path. + + Returns a list (parallel to ``trajectories``) of ``{camera_name: path}`` dicts. + + A single trajectory with a single camera and an ``output_path`` ending in + ``.mp4`` is written directly to that file. Otherwise ``output_path`` is treated + as a directory: each trajectory gets a ``world_{id}`` subfolder (suffixed with its + list index to avoid collisions when world ids repeat), with one ``{camera}.mp4`` + per camera. + """ + output_path = Path(output_path) + single = len(trajectories) == 1 and len(camera_names) == 1 + + paths: list[dict[str, Path]] = [] + for i, traj in enumerate(trajectories): + if single and output_path.suffix == ".mp4": + paths.append({camera_names[0]: output_path}) + continue + traj_dir = ( + output_path + if len(trajectories) == 1 + else output_path / f"world_{traj.world_id}_idx{i:04d}" + ) + paths.append( + {cam: traj_dir / f"{cam.replace('/', '_')}.mp4" for cam in camera_names} + ) + return paths + + +def _resolve_cameras( + cameras: str | list[str] | None, trajectories: list[RecordedTrajectory] +) -> list[str]: + """Camera names to render: the explicit argument, or the recorded set if None. + + When falling back to the recorded set, all trajectories must agree -- the single + CPU renderer (and the GPU batch render context) is built once for a shared set. + """ + if cameras is not None: + return [cameras] if isinstance(cameras, str) else list(cameras) + names = list(trajectories[0].camera_names) + for traj in trajectories[1:]: + if list(traj.camera_names) != names: + raise ValueError( + "Trajectories were recorded with different cameras; pass an explicit " + "`cameras` argument to choose a common set." + ) + return names + + +def _render_trajectory_frames( + mj_model: mj.MjModel, + mj_data: mj.MjData, + renderer: mj.Renderer, + traj: RecordedTrajectory, + camera_ids: dict[str, int], + scene_option: mj.MjvOption | None, +) -> dict[str, list[np.ndarray]]: + """Re-render every frame of one trajectory for the requested cameras. + + Sets ``qpos`` (and mocap) per frame, runs position-only kinematics, and + rasterizes each camera. Returns ``{camera_name: [frame, ...]}``. + """ + frames: dict[str, list[np.ndarray]] = {cam: [] for cam in camera_ids} + for f in range(traj.n_frames): + mj_data.qpos[:] = traj.qpos[f] + if traj.has_mocap: + mj_data.mocap_pos[:] = traj.mocap_pos[f] + mj_data.mocap_quat[:] = traj.mocap_quat[f] + # Position-only kinematics regenerate all geom/site/camera/light transforms + # that the renderer reads -- no need for velocities/forces/contacts. + mj.mj_kinematics(mj_model, mj_data) + mj.mj_camlight(mj_model, mj_data) + for cam_name, cam_id in camera_ids.items(): + renderer.update_scene(mj_data, cam_id, scene_option) + frames[cam_name].append(renderer.render()) + return frames + + +def render_trajectories( + mj_model: mj.MjModel, + trajectories: RecordedTrajectory | list[RecordedTrajectory], + output_path: PathLike, + *, + cameras: str | list[str] | None = None, + scene_option: mj.MjvOption | None = None, +) -> None: + """Replay recorded trajectories to video on the CPU. + + For GPU batch replay, use `flygym.warp.rendering.render_trajectories_gpu`. + + Args: + mj_model: Compiled model the trajectories were recorded against (load it from + wherever you persisted it, e.g. a folder written by + `BaseCompositionElement.save_xml_with_assets`). Its ``qpos`` layout must + match the trajectories. + trajectories: One trajectory or a list of them (e.g. from `load_trajectories`). + output_path: Where to write videos. See `_resolve_render_output_paths` for the + file/directory layout. + cameras: Camera name(s) to render. Defaults to each trajectory's recorded + ``camera_names``. + scene_option: MuJoCo scene options applied at render time. + """ + trajectories = _as_trajectory_list(trajectories) + _check_model_compatible(mj_model, trajectories) + camera_names = _resolve_cameras(cameras, trajectories) + camera_ids = { + c: mj.mj_name2id(mj_model, mj.mjtObj.mjOBJ_CAMERA, c) for c in camera_names + } + out_paths = _resolve_render_output_paths(trajectories, camera_names, output_path) + + height, width = trajectories[0].camera_res + mj_data = mj.MjData(mj_model) + renderer = mj.Renderer(mj_model, height, width) + try: + for traj, traj_out in zip(trajectories, out_paths): + frames = _render_trajectory_frames( + mj_model, mj_data, renderer, traj, camera_ids, scene_option + ) + for cam, cam_frames in frames.items(): + write_video_from_frames( + traj_out[cam], cam_frames, fps=traj.output_fps, codec="libx264" + ) + finally: + renderer.close() diff --git a/src/flygym/simulation.py b/src/flygym/simulation.py index d8565dce..00f510e0 100644 --- a/src/flygym/simulation.py +++ b/src/flygym/simulation.py @@ -9,7 +9,7 @@ from flygym.anatomy import BodySegment from flygym.compose.fly import ActuatorType from flygym.compose.world import BaseWorld -from flygym.rendering import Renderer +from flygym.rendering import Renderer, TrajectoryRecorder from flygym.utils.profiling import print_perf_report @@ -105,6 +105,7 @@ def set_renderer( output_fps: int = 25, buffer_frames: bool = True, scene_option: mj.MjvOption | None = None, + record_trajectory_only: bool = False, **kwargs: Any, ) -> Renderer: """Attach a renderer to this simulation. @@ -117,21 +118,37 @@ def set_renderer( output_fps: Output video frame rate. buffer_frames: If True, store rendered frames in memory. scene_option: MuJoCo scene options. Uses defaults if None. - **kwargs: Passed to ``mujoco.Renderer``. + record_trajectory_only: If True, attach a `TrajectoryRecorder` instead of + a `Renderer`: it records ``qpos`` (and mocap poses) at the render + cadence instead of rasterizing frames. Read the result from + ``self.renderer.recorded_trajectory`` and replay it later with + `render_trajectories`. ``buffer_frames`` is ignored in this mode. + **kwargs: Passed to ``mujoco.Renderer`` (ignored when recording only). Returns: - The created `Renderer` instance. + The created `Renderer` (or `TrajectoryRecorder`) instance. """ - self.renderer = Renderer( - self.mj_model, - cameras, - camera_res=camera_res, - playback_speed=playback_speed, - output_fps=output_fps, - buffer_frames=buffer_frames, - scene_option=scene_option, - **kwargs, - ) + if record_trajectory_only: + self.renderer = TrajectoryRecorder( + self.mj_model, + cameras, + camera_res=camera_res, + playback_speed=playback_speed, + output_fps=output_fps, + scene_option=scene_option, + **kwargs, + ) + else: + self.renderer = Renderer( + self.mj_model, + cameras, + camera_res=camera_res, + playback_speed=playback_speed, + output_fps=output_fps, + buffer_frames=buffer_frames, + scene_option=scene_option, + **kwargs, + ) return self.renderer def render_as_needed(self) -> bool: diff --git a/src/flygym/warp/__init__.py b/src/flygym/warp/__init__.py index 427a6d78..6a8abfeb 100644 --- a/src/flygym/warp/__init__.py +++ b/src/flygym/warp/__init__.py @@ -1,4 +1,17 @@ from .simulation import GPUSimulation -from .rendering import WarpGPUBatchRenderer, WarpCPURenderer +from .rendering import ( + WarpGPUBatchRenderer, + WarpCPURenderer, + WarpTrajectoryRecorder, + modify_world_for_batch_rendering, + render_trajectories_gpu, +) -__all__ = ["GPUSimulation", "WarpGPUBatchRenderer", "WarpCPURenderer"] +__all__ = [ + "GPUSimulation", + "WarpGPUBatchRenderer", + "WarpCPURenderer", + "WarpTrajectoryRecorder", + "modify_world_for_batch_rendering", + "render_trajectories_gpu", +] diff --git a/src/flygym/warp/rendering/__init__.py b/src/flygym/warp/rendering/__init__.py new file mode 100644 index 00000000..85bf23a0 --- /dev/null +++ b/src/flygym/warp/rendering/__init__.py @@ -0,0 +1,24 @@ +"""Multi-world MuJoCo-Warp rendering: live rasterization and trajectory replay. + +Split into `live_rendering` (real-time GPU/CPU renderers) and `recorded_trajectory` +(GPU recording + batch replay). Public names are re-exported for backward +compatibility. +""" + +from flygym.warp.rendering.live_rendering import ( + WarpGPUBatchRenderer, + WarpCPURenderer, + modify_world_for_batch_rendering, +) +from flygym.warp.rendering.recorded_trajectory import ( + WarpTrajectoryRecorder, + render_trajectories_gpu, +) + +__all__ = [ + "WarpGPUBatchRenderer", + "WarpCPURenderer", + "WarpTrajectoryRecorder", + "modify_world_for_batch_rendering", + "render_trajectories_gpu", +] diff --git a/src/flygym/warp/rendering.py b/src/flygym/warp/rendering/live_rendering.py similarity index 94% rename from src/flygym/warp/rendering.py rename to src/flygym/warp/rendering/live_rendering.py index be88be53..351bf8d5 100644 --- a/src/flygym/warp/rendering.py +++ b/src/flygym/warp/rendering/live_rendering.py @@ -1,3 +1,9 @@ +"""Live multi-world MuJoCo-Warp rendering (GPU batch and per-world CPU). + +For recording qpos trajectories on GPU and replaying them, see +`flygym.warp.rendering.recorded_trajectory`. +""" + import warnings from typing import Any, override from os import PathLike @@ -11,12 +17,19 @@ from PIL import Image, ImageDraw, ImageFont from flygym.compose import BaseWorld -from flygym.rendering import Renderer +from flygym.rendering.live_rendering import Renderer from flygym.warp.utils import get_rgb_selected_worlds_and_cameras from flygym.utils.video import write_video_from_frames from flygym.utils.plot import find_font_path +__all__ = [ + "WarpGPUBatchRenderer", + "WarpCPURenderer", + "modify_world_for_batch_rendering", +] + + class _BaseWarpRenderer(Renderer, ABC): @override def __init__( @@ -388,8 +401,15 @@ def modify_world_for_batch_rendering(world: BaseWorld) -> bool: This may reduce texture and lighting realism. - Modification happens in place. Returns True if any modifications were made, False - otherwise. + Modification happens in place on ``world.mjcf_root``. Returns True if any + modifications were made, False otherwise. Only ``world.mjcf_root`` (the `MjSpec`) + and ``world.fly_lookup`` (the fly names) are used, so this can be called on any + object exposing those two attributes -- see `render_trajectories_gpu`, which + applies it to a spec reconstructed from a saved trajectory. + + Note: these are material/texture/light edits only -- they do not change the joint + structure, so a model recompiled afterward keeps the same ``qpos`` layout and a + recorded trajectory stays valid against it. Note for developers: Check if anything here can be dropped upon new MJWarp releases. """ diff --git a/src/flygym/warp/rendering/recorded_trajectory.py b/src/flygym/warp/rendering/recorded_trajectory.py new file mode 100644 index 00000000..06331595 --- /dev/null +++ b/src/flygym/warp/rendering/recorded_trajectory.py @@ -0,0 +1,287 @@ +"""Record qpos trajectories on GPU and replay them via MuJoCo-Warp batch rendering.""" + +from typing import Any, override +from os import PathLike + +import mujoco as mj +import mujoco_warp as mjw +import warp as wp +import numpy as np + +from flygym.rendering.recorded_trajectory import ( + RecordedTrajectory, + _as_trajectory_list, + _check_model_compatible, + _resolve_cameras, + _resolve_render_output_paths, +) +from flygym.warp.rendering.live_rendering import _BaseWarpRenderer +from flygym.warp.utils import get_rgb_selected_worlds_and_cameras +from flygym.utils.video import write_video_from_frames + + +__all__ = ["WarpTrajectoryRecorder", "render_trajectories_gpu"] + + +class WarpTrajectoryRecorder(_BaseWarpRenderer): + """Records ``qpos`` (and mocap poses) per world instead of rasterizing frames. + + The GPU counterpart of `flygym.rendering.TrajectoryRecorder`: a drop-in renderer + swap for `GPUSimulation` that, at the render cadence, copies the generalized + coordinates of the selected worlds off the GPU instead of running the batch + renderer. Each selected world becomes one `RecordedTrajectory`, exposed as + `recorded_trajectories`, in the same format the CPU recorder produces. + """ + + def _build_mj_renderer(self, mj_model, nrows, ncols, **kwargs): + # No rasterization: skip the GL/EGL context allocation entirely. + return None + + def _render_setup_impl(self, **kwargs: Any) -> None: + self._nq = self.mj_model.nq + self._nmocap = self.mj_model.nmocap + # The CPU-side renderer/scene_option inherited from Renderer are unused. + self.mj_renderer = None + self.scene_option = None + + def _render_impl(self, mjw_data: mjw.Data) -> tuple: + # One host transfer per recorded frame: (n_worlds, nq) is tiny next to the + # (n_worlds, n_cams, H, W, 3) RGB tensor the batch renderer would buffer. + qpos = mjw_data.qpos.numpy()[self.world_ids].copy() + if self._nmocap > 0: + mocap_pos = mjw_data.mocap_pos.numpy()[self.world_ids].copy() + mocap_quat = mjw_data.mocap_quat.numpy()[self.world_ids].copy() + else: + mocap_pos = mocap_quat = None + return (qpos, mocap_pos, mocap_quat) + + @property + def recorded_trajectories(self) -> list[RecordedTrajectory]: + """One `RecordedTrajectory` per recorded world.""" + if not self.buffer_frames: + raise RuntimeError( + "Frame buffering was disabled for this recorder, so recorded " + "trajectories are not available." + ) + if len(self._frames) == 0: + raise RuntimeError("No frames have been recorded yet.") + + # self._frames is a list (over time) of (qpos, mocap_pos, mocap_quat) tuples, + # each batched over the recorded worlds along axis 0. + qpos_all = np.stack([f[0] for f in self._frames], axis=0) # (T, n_worlds, nq) + if self._nmocap > 0: + mocap_pos_all = np.stack([f[1] for f in self._frames], axis=0) + mocap_quat_all = np.stack([f[2] for f in self._frames], axis=0) + + trajectories = [] + for w, world_id in enumerate(self.world_ids): + trajectories.append( + RecordedTrajectory( + qpos=qpos_all[:, w, :], + output_fps=self.output_fps, + playback_speed=self.playback_speed, + camera_names=list(self.enabled_cam_names), + camera_res=self.camera_res, + world_id=world_id, + mocap_pos=mocap_pos_all[:, w] if self._nmocap > 0 else None, + mocap_quat=mocap_quat_all[:, w] if self._nmocap > 0 else None, + ) + ) + return trajectories + + def _fetch_frames_to_cpu_impl(self, world_id_among_rendered, cam_id_among_rendered): + raise RuntimeError( + "WarpTrajectoryRecorder records state, not frames. Use " + "`recorded_trajectories` and replay with render_trajectories_gpu." + ) + + @override + def save_video(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError( + "WarpTrajectoryRecorder records state, not frames. Save it with " + "flygym.rendering.save_trajectories and replay with " + "render_trajectories_gpu / flygym.rendering.render_trajectories." + ) + + @override + def show_in_notebook(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError( + "WarpTrajectoryRecorder records state, not frames. Save it with " + "flygym.rendering.save_trajectories and replay with " + "render_trajectories_gpu / flygym.rendering.render_trajectories." + ) + + @override + def close(self) -> None: + return # no mj.Renderer context held + + +# Default number of frames staged into one batched GPU render. Tunable via the +# ``worlds_per_batch`` argument; larger uses more GPU memory. +_DEFAULT_WORLDS_PER_BATCH = 32 + + +def render_trajectories_gpu( + mj_model: mj.MjModel, + trajectories: RecordedTrajectory | list[RecordedTrajectory], + output_path: PathLike, + *, + cameras: str | list[str] | None = None, + worlds_per_batch: int | None = None, +) -> None: + """Replay recorded trajectories to video using MuJoCo-Warp GPU batch rendering. + + The GPU counterpart of `flygym.rendering.render_trajectories`. All ``(trajectory, + frame)`` pairs are flattened into one work-list and bin-packed into batched + ``mjw.Data`` renders: each batch stages a block of recorded ``qpos`` rows into the + parallel worlds, runs position-only kinematics, rasterizes with the same path as + `WarpGPUBatchRenderer`, then scatters the frames back to their trajectories. This + decouples the batch size used for *rendering* (``worlds_per_batch``) from the + number of worlds that were *simulated*. + + ``mj_model`` must already be prepared for GPU batch rendering -- i.e. have textures + stripped and overhead lights added by + `flygym.warp.rendering.modify_world_for_batch_rendering` (textured complex meshes + otherwise corrupt MJWarp memory). The model held by a `GPUSimulation` configured + with ``use_gpu_batch_rendering=True`` already satisfies this; to replay from a model + you persisted yourself, call ``modify_world_for_batch_rendering(world)`` and + recompile before passing it here. Those edits do not change the ``qpos`` layout, so + the recorded trajectories stay valid. + + Args: + mj_model: Compiled, batch-render-ready model (see note above). Its ``qpos`` + layout must match the trajectories. + trajectories: One trajectory or a list of them (e.g. from + `flygym.rendering.load_trajectories`). + output_path: Where to write videos (see `_resolve_render_output_paths`). + cameras: Camera name(s) to render. Defaults to each trajectory's recorded + ``camera_names``. All trajectories must share the same camera set and + resolution (the batch render context is built once). Lengths may differ. + worlds_per_batch: Number of frames staged into one batched render. Defaults to + a heuristic; larger uses more GPU memory. + """ + trajectories = _as_trajectory_list(trajectories) + _check_model_compatible(mj_model, trajectories) + _render_trajectories_gpu( + trajectories, + mj_model, + output_path, + cameras=cameras, + worlds_per_batch=worlds_per_batch, + ) + + +def _render_trajectories_gpu( + trajectories: list[RecordedTrajectory], + mj_model: mj.MjModel, + output_path: PathLike, + *, + cameras: str | list[str] | None, + worlds_per_batch: int | None, +) -> None: + camera_names = _resolve_cameras(cameras, trajectories) + height, width = _validate_shared_camera_res(trajectories) + + cam_ids = [ + mj.mj_name2id(mj_model, mj.mjtObj.mjOBJ_CAMERA, name) for name in camera_names + ] + for name, cid in zip(camera_names, cam_ids): + if cid == -1: + raise ValueError(f"Camera {name!r} not found in the model.") + + # Flatten all (trajectory, frame) pairs into one work-list. + work = [ + (ti, fi) for ti, traj in enumerate(trajectories) for fi in range(traj.n_frames) + ] + if len(work) == 0: + raise ValueError("Trajectories contain no frames.") + + batch = min(worlds_per_batch or _DEFAULT_WORLDS_PER_BATCH, len(work)) + nmocap = mj_model.nmocap + nq = mj_model.nq + + # Pre-allocate output, one (n_frames, H, W, 3) array per (trajectory, camera). + results: list[dict[str, np.ndarray]] = [ + { + cam: np.zeros((traj.n_frames, height, width, 3), dtype=np.uint8) + for cam in camera_names + } + for traj in trajectories + ] + + # GPU-side setup: model, a batched Data of size `batch`, and a render context. + mjw_model = mjw.put_model(mj_model) + mjw_data = mjw.put_data(mj_model, mj.MjData(mj_model), nworld=batch) + cam_mask = [cid in cam_ids for cid in range(mj_model.ncam)] + render_context = mjw.create_render_context( + mjm=mj_model, + nworld=batch, + cam_active=cam_mask, + cam_res=(width, height), # MJWarp expects (W, H); we use (H, W) + ) + world_ids_gpu = wp.array(list(range(batch)), dtype=wp.int32) + cam_ids_gpu = wp.array(cam_ids, dtype=wp.int32) + + for base in range(0, len(work), batch): + chunk = work[base : base + batch] + chunk_len = len(chunk) + + # Stage qpos (and mocap) for this chunk into the parallel worlds. Padding + # slots (last partial batch) repeat the last real frame and are never read. + qpos_batch = np.zeros((batch, nq), dtype=np.float32) + for j, (ti, fi) in enumerate(chunk): + qpos_batch[j] = trajectories[ti].qpos[fi] + qpos_batch[chunk_len:] = qpos_batch[chunk_len - 1] + mjw_data.qpos.assign(qpos_batch) + + if nmocap > 0: + mocap_pos_batch = np.zeros((batch, nmocap, 3), dtype=np.float32) + mocap_quat_batch = np.zeros((batch, nmocap, 4), dtype=np.float32) + for j, (ti, fi) in enumerate(chunk): + mocap_pos_batch[j] = trajectories[ti].mocap_pos[fi] + mocap_quat_batch[j] = trajectories[ti].mocap_quat[fi] + mocap_pos_batch[chunk_len:] = mocap_pos_batch[chunk_len - 1] + mocap_quat_batch[chunk_len:] = mocap_quat_batch[chunk_len - 1] + mjw_data.mocap_pos.assign(mocap_pos_batch) + mjw_data.mocap_quat.assign(mocap_quat_batch) + + # Position-only kinematics regenerate all geom/site/camera/light transforms. + mjw.kinematics(mjw_model, mjw_data) + mjw.camlight(mjw_model, mjw_data) + + mjw.refit_bvh(mjw_model, mjw_data, render_context) + mjw.render(mjw_model, mjw_data, render_context) + + rgb_out = wp.zeros((batch, len(cam_ids), height, width), dtype=wp.vec3f) + get_rgb_selected_worlds_and_cameras( + render_context, world_ids_gpu, cam_ids_gpu, rgb_out + ) + rgb_np = (rgb_out.numpy() * 255.0).astype(np.uint8) # (batch, ncam, H, W, 3) + + # Scatter rendered frames back to their (trajectory, camera, frame) slot. + for j, (ti, fi) in enumerate(chunk): + for ci, cam in enumerate(camera_names): + results[ti][cam][fi] = rgb_np[j, ci] + + out_paths = _resolve_render_output_paths(trajectories, camera_names, output_path) + for ti, traj in enumerate(trajectories): + for cam in camera_names: + write_video_from_frames( + out_paths[ti][cam], + list(results[ti][cam]), + fps=traj.output_fps, + codec="libx264", + ) + + +def _validate_shared_camera_res( + trajectories: list[RecordedTrajectory], +) -> tuple[int, int]: + res = tuple(trajectories[0].camera_res) + for traj in trajectories[1:]: + if tuple(traj.camera_res) != res: + raise ValueError( + "GPU batch rendering requires all trajectories to share one camera " + f"resolution, but found {res} and {tuple(traj.camera_res)}." + ) + return res diff --git a/src/flygym/warp/simulation.py b/src/flygym/warp/simulation.py index 36c46351..d8cfba11 100644 --- a/src/flygym/warp/simulation.py +++ b/src/flygym/warp/simulation.py @@ -14,6 +14,7 @@ from flygym.warp.rendering import ( WarpGPUBatchRenderer, WarpCPURenderer, + WarpTrajectoryRecorder, modify_world_for_batch_rendering, ) from flygym.warp.utils import ( @@ -280,8 +281,9 @@ def set_renderer( scene_option: mj.MjvOption | None = None, worlds: list[int] | None = None, use_gpu_batch_rendering: bool = False, + record_trajectory_only: bool = False, **kwargs: Any, - ) -> WarpGPUBatchRenderer | WarpCPURenderer: + ) -> WarpGPUBatchRenderer | WarpCPURenderer | WarpTrajectoryRecorder: """Attach a renderer to this GPU simulation. Args: @@ -293,14 +295,37 @@ def set_renderer( scene_option: MuJoCo scene options. Uses defaults if None. worlds: Indices of worlds to render. Defaults to all worlds. use_gpu_batch_rendering: If True, use `WarpGPUBatchRenderer`; - otherwise use `WarpCPURenderer`. - **kwargs: Passed to the renderer. + otherwise use `WarpCPURenderer`. Ignored when + ``record_trajectory_only`` is True. + record_trajectory_only: If True, attach a `WarpTrajectoryRecorder` instead + of a renderer: it records ``qpos`` (and mocap poses) for the selected + worlds at the render cadence instead of rasterizing frames. Read the + result from ``self.renderer.recorded_trajectories`` (one per recorded + world) and replay it later with + `flygym.warp.rendering.render_trajectories_gpu` or + `flygym.rendering.render_trajectories`. + **kwargs: Passed to the renderer (ignored when recording only). Returns: - The created renderer instance. + The created renderer (or `WarpTrajectoryRecorder`) instance. """ if worlds is None: worlds = list(range(self.n_worlds)) + + if record_trajectory_only: + self.renderer = WarpTrajectoryRecorder( + self.mj_model, + cameras, + n_worlds_total=self.n_worlds, + worlds=worlds, + camera_res=camera_res, + playback_speed=playback_speed, + output_fps=output_fps, + buffer_frames=True, + **kwargs, + ) + return self.renderer + self.use_gpu_batch_rendering = use_gpu_batch_rendering renderer_kwargs = { diff --git a/tests/core/test_trajectory.py b/tests/core/test_trajectory.py new file mode 100644 index 00000000..15b29895 --- /dev/null +++ b/tests/core/test_trajectory.py @@ -0,0 +1,254 @@ +"""Tests for trajectory recording, serialization, and CPU replay (flygym.rendering).""" + +import os + +import numpy as np +import mujoco as mj +import pytest + +from flygym.anatomy import AxisOrder, JointPreset, Skeleton +from flygym.compose.fly import NeuroMechFly +from flygym.compose.pose import KinematicPosePreset +from flygym.compose.world import TetheredWorld +from flygym.utils.math import Rotation3D +from flygym.simulation import Simulation +from flygym.rendering import ( + RecordedTrajectory, + save_trajectories, + load_trajectories, + render_trajectories, +) +from flygym.rendering.recorded_trajectory import _render_trajectory_frames + +# Rendering (rasterization) needs a headless GL context; skip those on CI runners +# that set SKIP_RENDERING_TESTS=1. Recording and serialization need no GL. +needs_gl = pytest.mark.skipif( + os.environ.get("SKIP_RENDERING_TESTS") == "1", + reason="SKIP_RENDERING_TESTS=1 (headless GL unavailable on this CI runner)", +) + + +# --------------------------------------------------------------------------- +# Helpers / fixtures +# --------------------------------------------------------------------------- + + +def make_sim(name: str) -> tuple[Simulation, str]: + """Build a Simulation whose fly has a tracking camera; return (sim, cam_name).""" + pose = KinematicPosePreset.NEUTRAL.get_pose_by_axis_order(AxisOrder.YAW_PITCH_ROLL) + skeleton = Skeleton( + axis_order=AxisOrder.YAW_PITCH_ROLL, joint_preset=JointPreset.LEGS_ONLY + ) + fly = NeuroMechFly(name=f"{name}_fly") + fly.add_joints(skeleton, neutral_pose=pose) + fly.add_tracking_camera(name="trackcam") + world = TetheredWorld(name=f"{name}_world") + world.add_fly( + fly, + spawn_position=[0, 0, 1.5], + spawn_rotation=Rotation3D("quat", [1, 0, 0, 0]), + ) + sim = Simulation(world) + return sim, fly.cameraname_to_mjcfcamera["trackcam"].name + + +@pytest.fixture(scope="module") +def sim_with_camera(): + return make_sim("traj") + + +@pytest.fixture(scope="module") +def recorded(sim_with_camera): + """Record a short trajectory (no GL needed -- the recorder stores qpos).""" + sim, cam_name = sim_with_camera + rec = sim.set_renderer( + cam_name, camera_res=(64, 64), output_fps=100, record_trajectory_only=True + ) + sim.reset() + for _ in range(200): + sim.step() + sim.render_as_needed() + return rec.recorded_trajectory, sim, cam_name + + +# --------------------------------------------------------------------------- +# Recording +# --------------------------------------------------------------------------- + + +class TestTrajectoryRecorder: + def test_records_frames(self, recorded): + traj, sim, _ = recorded + assert traj.n_frames > 0 + assert traj.qpos.shape == (traj.n_frames, sim.mj_model.nq) + + def test_world_id_is_zero(self, recorded): + traj, _, _ = recorded + assert traj.world_id == 0 + + def test_metadata_captured(self, recorded): + traj, _, cam_name = recorded + assert traj.camera_res == (64, 64) + assert traj.output_fps == 100 + assert cam_name in traj.camera_names + + def test_mocap_recorded_when_present(self, recorded): + # TetheredWorld tethers the fly via a mocap body, so nmocap > 0 and the + # recorder must capture mocap poses alongside qpos. + traj, sim, _ = recorded + if sim.mj_model.nmocap > 0: + assert traj.has_mocap + assert traj.mocap_pos.shape == (traj.n_frames, sim.mj_model.nmocap, 3) + assert traj.mocap_quat.shape == (traj.n_frames, sim.mj_model.nmocap, 4) + else: + assert not traj.has_mocap + + def test_reset_clears_buffer(self, sim_with_camera): + sim, cam_name = sim_with_camera + rec = sim.set_renderer( + cam_name, camera_res=(64, 64), output_fps=100, record_trajectory_only=True + ) + sim.reset() + for _ in range(50): + sim.step() + sim.render_as_needed() + assert rec.recorded_trajectory.n_frames > 0 + rec.reset() + with pytest.raises(RuntimeError): + _ = rec.recorded_trajectory + + +# --------------------------------------------------------------------------- +# Serialization (individual self-describing npz files; no model) +# --------------------------------------------------------------------------- + + +class TestSaveLoad: + def test_roundtrip(self, recorded, tmp_path): + traj, _, _ = recorded + save_trajectories(traj, tmp_path) + assert (tmp_path / "traj_0000.npz").exists() + trajs = load_trajectories(tmp_path) + assert len(trajs) == 1 + assert np.array_equal(trajs[0].qpos, traj.qpos) + assert trajs[0].camera_res == traj.camera_res + assert trajs[0].camera_names == traj.camera_names + assert trajs[0].output_fps == traj.output_fps + + def test_mocap_roundtrip(self, recorded, tmp_path): + traj, _, _ = recorded + save_trajectories(traj, tmp_path) + loaded = load_trajectories(tmp_path)[0] + assert loaded.has_mocap == traj.has_mocap + if traj.has_mocap: + assert np.array_equal(loaded.mocap_pos, traj.mocap_pos) + assert np.array_equal(loaded.mocap_quat, traj.mocap_quat) + + def test_save_list(self, recorded, tmp_path): + traj, _, _ = recorded + save_trajectories([traj, traj], tmp_path) + assert (tmp_path / "traj_0001.npz").exists() + trajs = load_trajectories(tmp_path) + assert len(trajs) == 2 + + def test_load_empty_folder_raises(self, tmp_path): + with pytest.raises(ValueError, match="No trajectory files"): + load_trajectories(tmp_path) + + def test_single_trajectory_save_load(self, recorded, tmp_path): + traj, _, _ = recorded + path = tmp_path / "one.npz" + traj.save(path) + loaded = RecordedTrajectory.load(path) + assert np.array_equal(loaded.qpos, traj.qpos) + assert loaded.world_id == traj.world_id + + +# --------------------------------------------------------------------------- +# Replay validation (no GL: the guard runs before any rasterization) +# --------------------------------------------------------------------------- + + +class TestRenderValidation: + def test_incompatible_model_raises(self, recorded, tmp_path): + traj, sim, _ = recorded + # Drop a qpos column so the trajectory no longer fits the model's nq. + bad = RecordedTrajectory( + qpos=traj.qpos[:, :-1], + output_fps=traj.output_fps, + playback_speed=traj.playback_speed, + camera_names=traj.camera_names, + camera_res=traj.camera_res, + ) + with pytest.raises(ValueError, match="different model"): + render_trajectories(sim.mj_model, bad, tmp_path / "x.mp4") + + +# --------------------------------------------------------------------------- +# CPU replay +# --------------------------------------------------------------------------- + + +@needs_gl +class TestRenderCPU: + def test_render_from_memory_writes_video(self, recorded, tmp_path): + traj, sim, _ = recorded + out = tmp_path / "video.mp4" + render_trajectories(sim.mj_model, traj, out) + assert out.exists() and out.stat().st_size > 0 + + def test_render_from_folder_writes_video(self, recorded, tmp_path): + traj, sim, _ = recorded + folder = tmp_path / "folder" + save_trajectories(traj, folder) + trajs = load_trajectories(folder) + out = tmp_path / "out.mp4" + render_trajectories(sim.mj_model, trajs, out) + assert out.exists() and out.stat().st_size > 0 + + def test_replay_matches_qpos_consistent_render(self, sim_with_camera): + """The production CPU replay reproduces a qpos-consistent render bit-for-bit. + + Reference frames are captured in the same run at the same render timepoints, + rendering after refreshing position kinematics from the post-step qpos (so + the geometry is consistent with the recorded qpos rather than one step + stale). The replay path builds its own MjData independently, so a bit-exact + match validates the full record -> replay round trip. + """ + sim, cam_name = sim_with_camera + cam_id = mj.mj_name2id(sim.mj_model, mj.mjtObj.mjOBJ_CAMERA, cam_name) + + rec = sim.set_renderer( + cam_name, camera_res=(64, 64), output_fps=100, record_trajectory_only=True + ) + sim.reset() + ref_renderer = mj.Renderer(sim.mj_model, 64, 64) + ref_frames = [] + for _ in range(150): + sim.step() + if sim.render_as_needed(): # records qpos at the render cadence + d = sim.mj_data + mj.mj_kinematics(sim.mj_model, d) + mj.mj_camlight(sim.mj_model, d) + ref_renderer.update_scene(d, cam_id) + ref_frames.append(ref_renderer.render().copy()) + ref_renderer.close() + traj = rec.recorded_trajectory + + from flygym.rendering import Renderer + + replay_renderer = Renderer(sim.mj_model, cam_name, camera_res=(64, 64)) + mj_data = mj.MjData(sim.mj_model) + replay = _render_trajectory_frames( + sim.mj_model, + mj_data, + replay_renderer.mj_renderer, + traj, + {cam_name: cam_id}, + replay_renderer.scene_option, + )[cam_name] + replay_renderer.close() + + assert len(replay) == len(ref_frames) > 0 + for r_frame, ref_frame in zip(replay, ref_frames): + assert np.array_equal(r_frame, ref_frame) diff --git a/tests/warp/test_trajectory.py b/tests/warp/test_trajectory.py new file mode 100644 index 00000000..d021e8ea --- /dev/null +++ b/tests/warp/test_trajectory.py @@ -0,0 +1,174 @@ +"""Tests for GPU trajectory recording (WarpTrajectoryRecorder) and GPU replay.""" + +import numpy as np +import mujoco as mj +import pytest + +# These tests require the optional warp (GPU) extra; tag them so they can be +# excluded with ``-m "not warp"``, and skip the whole module if warp is absent. +pytestmark = pytest.mark.warp +pytest.importorskip("warp") + +from flygym.rendering import ( + save_trajectories, + load_trajectories, + render_trajectories, +) +from flygym.warp import ( + WarpTrajectoryRecorder, + render_trajectories_gpu, + modify_world_for_batch_rendering, +) + + +# --------------------------------------------------------------------------- +# Module-scoped fixture: GPU simulation + recorded trajectories +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def recorded_gpu(gpu_sim_factory): + """Record a short trajectory from a 4-world GPU sim, sub-selecting 2 worlds.""" + sim, fly, cam = gpu_sim_factory(n_worlds=4, fly_name="rec_gpu_fly") + rec = sim.set_renderer( + cam, + camera_res=(64, 64), + output_fps=100, + worlds=[0, 2], + record_trajectory_only=True, + ) + sim.reset() + for _ in range(200): + sim.step() + sim.render_as_needed() + return rec, sim, cam + + +def _batch_render_model(sim) -> mj.MjModel: + """Compile a batch-render-ready model from the sim's world (caller's job now).""" + modify_world_for_batch_rendering(sim.world) + return sim.world.compile()[0] + + +# --------------------------------------------------------------------------- +# Recording +# --------------------------------------------------------------------------- + + +class TestWarpTrajectoryRecorder: + def test_is_recorder_not_renderer(self, recorded_gpu): + rec, _, _ = recorded_gpu + assert isinstance(rec, WarpTrajectoryRecorder) + # No rasterization: the recorder holds no mj.Renderer / scene_option. + assert rec.mj_renderer is None + assert rec.scene_option is None + + def test_one_trajectory_per_selected_world(self, recorded_gpu): + rec, _, _ = recorded_gpu + trajs = rec.recorded_trajectories + assert len(trajs) == 2 + assert [t.world_id for t in trajs] == [0, 2] + + def test_trajectory_shapes(self, recorded_gpu): + rec, sim, _ = recorded_gpu + for t in rec.recorded_trajectories: + assert t.n_frames > 0 + assert t.qpos.shape == (t.n_frames, sim.mj_model.nq) + + def test_save_video_raises(self, recorded_gpu): + rec, _, _ = recorded_gpu + with pytest.raises(RuntimeError): + rec.save_video(0, "x.mp4") + + +# --------------------------------------------------------------------------- +# Serialization (backend-agnostic format) +# --------------------------------------------------------------------------- + + +class TestSaveLoadGPU: + def test_roundtrip(self, recorded_gpu, tmp_path): + rec, _, _ = recorded_gpu + trajs = rec.recorded_trajectories + save_trajectories(trajs, tmp_path) + loaded = load_trajectories(tmp_path) + assert len(loaded) == 2 + for a, b in zip(loaded, trajs): + assert np.array_equal(a.qpos, b.qpos) + assert a.world_id == b.world_id + + +# --------------------------------------------------------------------------- +# Replay +# --------------------------------------------------------------------------- + + +class TestRenderGPU: + def test_gpu_replay_writes_videos(self, recorded_gpu, tmp_path): + rec, sim, _ = recorded_gpu + trajs = rec.recorded_trajectories + mj_model = _batch_render_model(sim) + out = tmp_path / "gpu_out" + render_trajectories_gpu(mj_model, trajs, out, worlds_per_batch=8) + videos = sorted(out.rglob("*.mp4")) + assert len(videos) == 2 + assert all(v.stat().st_size > 0 for v in videos) + + def test_gpu_recorded_renders_on_cpu(self, recorded_gpu, tmp_path): + """A GPU-recorded trajectory is backend-agnostic: it replays on CPU too.""" + rec, sim, _ = recorded_gpu + trajs = rec.recorded_trajectories + out = tmp_path / "cpu_out" + render_trajectories(sim.mj_model, trajs, out) + videos = sorted(out.rglob("*.mp4")) + assert len(videos) == 2 + + +class TestKinematicsIdentityAcrossBackends: + def test_cpu_and_gpu_kinematics_match_for_same_qpos(self, recorded_gpu): + """Replay kinematics are identical across backends (not pixels, geometry). + + Pixel output differs stylistically (different rasterizers), but for the same + recorded qpos the CPU (``mj_kinematics``) and GPU (``mjw.kinematics``) passes + must place the geometry in the same poses. We compare ``geom_xpos`` for a few + recorded frames. + """ + import mujoco_warp as mjw + import warp as wp + + rec, sim, _ = recorded_gpu + traj = rec.recorded_trajectories[0] + mj_model = sim.mj_model + + # Sample a few frames spread across the trajectory. + idxs = [0, traj.n_frames // 2, traj.n_frames - 1] + + # CPU geom_xpos for each sampled qpos. + cpu_geom = [] + d = mj.MjData(mj_model) + for f in idxs: + d.qpos[:] = traj.qpos[f] + if traj.has_mocap: + d.mocap_pos[:] = traj.mocap_pos[f] + d.mocap_quat[:] = traj.mocap_quat[f] + mj.mj_kinematics(mj_model, d) + cpu_geom.append(d.geom_xpos.copy()) + + # GPU geom_xpos: stage the same qpos rows into a batched mjw.Data. + mjw_model = mjw.put_model(mj_model) + mjw_data = mjw.put_data(mj_model, mj.MjData(mj_model), nworld=len(idxs)) + qpos_batch = np.stack([traj.qpos[f] for f in idxs]).astype(np.float32) + mjw_data.qpos.assign(qpos_batch) + if traj.has_mocap: + mjw_data.mocap_pos.assign( + np.stack([traj.mocap_pos[f] for f in idxs]).astype(np.float32) + ) + mjw_data.mocap_quat.assign( + np.stack([traj.mocap_quat[f] for f in idxs]).astype(np.float32) + ) + mjw.kinematics(mjw_model, mjw_data) + wp.synchronize() + gpu_geom = mjw_data.geom_xpos.numpy() # (len(idxs), ngeom, 3) + + for i in range(len(idxs)): + assert np.allclose(cpu_geom[i], gpu_geom[i], atol=1e-5) From 424d8723e31d9450ee1e7d17c06c8f8437c2d0a4 Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Mon, 29 Jun 2026 21:49:58 +0200 Subject: [PATCH 02/11] Silence expected batch-render warning in GPU replay test modify_world_for_batch_rendering warns as it adds overhead lights / strips textures. The GPU replay test calls it deliberately, so wrap it in warnings.catch_warnings() like the existing test_rendering.py does. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/warp/test_trajectory.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/warp/test_trajectory.py b/tests/warp/test_trajectory.py index d021e8ea..5518f1cb 100644 --- a/tests/warp/test_trajectory.py +++ b/tests/warp/test_trajectory.py @@ -1,5 +1,7 @@ """Tests for GPU trajectory recording (WarpTrajectoryRecorder) and GPU replay.""" +import warnings + import numpy as np import mujoco as mj import pytest @@ -46,7 +48,11 @@ def recorded_gpu(gpu_sim_factory): def _batch_render_model(sim) -> mj.MjModel: """Compile a batch-render-ready model from the sim's world (caller's job now).""" - modify_world_for_batch_rendering(sim.world) + with warnings.catch_warnings(): + # modify_world_for_batch_rendering warns as it adds overhead lights/strips + # textures; that is expected here, so silence it (matches test_rendering.py). + warnings.simplefilter("ignore") + modify_world_for_batch_rendering(sim.world) return sim.world.compile()[0] From 2805e9eddce59eca4729bc424f8918aa24822d7c Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Mon, 29 Jun 2026 22:34:13 +0200 Subject: [PATCH 03/11] Add GPU trajectory record/replay demo script scripts/record_replay_trajectories_gpu.py mirrors replay_behavior_gpu.py's GPU-resident captured loop but swaps the live batch renderer for a WarpTrajectoryRecorder: it simulates many worlds, records qpos for a small subset, saves the trajectories and model to disk as independent artifacts, then reloads and renders them post-hoc on GPU (and optionally CPU). This exercises the full decoupled record -> save -> replay pipeline end to end. Co-Authored-By: Claude Opus 4.8 (1M context) --- scripts/record_replay_trajectories_gpu.py | 245 ++++++++++++++++++++++ 1 file changed, 245 insertions(+) create mode 100644 scripts/record_replay_trajectories_gpu.py diff --git a/scripts/record_replay_trajectories_gpu.py b/scripts/record_replay_trajectories_gpu.py new file mode 100644 index 00000000..8403b48d --- /dev/null +++ b/scripts/record_replay_trajectories_gpu.py @@ -0,0 +1,245 @@ +"""Demo: record kinematic trajectories on GPU, then render them to video post-hoc. + +This demonstrates the trajectory recording / replay feature (issue #296): it +decouples how many worlds are *simulated* from how many are *rendered*. Buffering +full ``(n_worlds, n_cams, H, W, 3)`` RGB tensors during a large parallel run is +hopeless at thousands of worlds, so instead we record only the generalized +coordinates (``qpos``, plus mocap poses) of a small selected subset of worlds at the +render cadence, and rasterize them afterwards. + +It is the trajectory-recording counterpart of ``replay_behavior_gpu.py``: the same +Spotlight kinematic recording is replayed across many parallel worlds with the same +fully GPU-resident, CUDA-graph-captured inner loop, but the live batch renderer is +swapped for a ``WarpTrajectoryRecorder`` (via ``set_renderer(..., +record_trajectory_only=True)``). The recorder copies ``qpos`` off the GPU at the +render cadence instead of rendering frames. + +The script then shows the full decoupled pipeline: + +1. Simulate ``--n-worlds`` worlds, recording trajectories for ``--record-worlds`` of + them. +2. Save the trajectories (one ``.npz`` each) and the model (``save_xml_with_assets``) + to disk -- they are independent artifacts; a trajectory carries no model. +3. Reload the trajectories from disk and render them to video, on GPU + (``render_trajectories_gpu``) and optionally on CPU (``--cpu-replay``) to show that + a GPU-recorded trajectory is backend-agnostic. + +Example: + uv run python scripts/record_replay_trajectories_gpu.py --output outputs/traj_demo + uv run python scripts/record_replay_trajectories_gpu.py --output outputs/traj_demo --cpu-replay +""" + +import argparse +from pathlib import Path +from time import perf_counter_ns + +import numpy as np +import warp as wp + +from flygym.warp import ( + GPUSimulation, + render_trajectories_gpu, + modify_world_for_batch_rendering, +) +from flygym.warp.utils import check_gpu +from flygym.rendering import save_trajectories, load_trajectories, render_trajectories +from flygym.compose import ActuatorType +from flygym_demo.benchmark import ( + make_model, + ReplayTargetData, + update_target_angles_kernel, + increment_counter_kernel, +) + +_MODEL_SUBDIR = "model" +_TRAJ_SUBDIR = "trajectories" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--output", + type=Path, + default=Path("outputs/traj_demo"), + metavar="DIR", + help="Directory for the saved trajectories, model, and rendered videos " + "(default: outputs/traj_demo). Reused if it already exists.", + ) + parser.add_argument( + "--n-worlds", + type=int, + default=1000, + help="Number of parallel worlds to simulate (default: 1000).", + ) + parser.add_argument( + "--record-worlds", + type=int, + default=5, + help="Number of worlds to record trajectories for and render (default: 5). " + "This is the whole point: simulate many, record/render few.", + ) + parser.add_argument( + "--sim-steps", + type=int, + default=2000, + help="Number of steps to simulate per world (default: 2000 = 0.2 s).", + ) + parser.add_argument( + "--timestep", + type=float, + default=1e-4, + help="Simulation timestep in seconds (default: 1e-4).", + ) + parser.add_argument( + "--worlds-per-batch", + type=int, + default=None, + help="GPU batch size for post-hoc rendering (default: library heuristic). " + "Decoupled from --n-worlds; larger uses more GPU memory.", + ) + parser.add_argument( + "--cpu-replay", + action="store_true", + help="Additionally replay the GPU-recorded trajectories on the CPU, to " + "demonstrate that the recorded format is backend-agnostic.", + ) + return parser.parse_args() + + +def record_trajectories(args: argparse.Namespace): + """Run the GPU simulation, recording qpos for a subset of worlds. + + Returns ``(trajectories, world, sim)``: the recorded trajectories (one per + recorded world), the world (kept so we can persist / re-compile the model), and + the simulation (kept for its unmodified ``mj_model``, used for CPU replay). + """ + n_worlds = args.n_worlds + sim_steps = args.sim_steps + timestep = args.timestep + n_record = min(args.record_worlds, n_worlds) + actuator_type = ActuatorType.POSITION + + fly, world, cam = make_model() + fly_name = fly.name + + # Build per-world target angle slices (world 0 -> first slice, world 1 -> next...). + replay_data = ReplayTargetData( + timestep, fly.get_actuated_jointdofs_order(actuator_type) + ) + target_angles_all_worlds = replay_data.make_target_angles_all_worlds( + n_worlds, sim_steps + ) + n_dofs = target_angles_all_worlds.shape[-1] + + sim = GPUSimulation(world, n_worlds, timestep=timestep) + + # Swap the live batch renderer for a recorder: it stores qpos for the selected + # worlds at the render cadence instead of rasterizing frames. + recorder = sim.set_renderer( + cam, + playback_speed=0.2, + output_fps=25, + worlds=list(range(n_record)), + record_trajectory_only=True, + ) + + # Reset to the neutral keyframe and settle. Must happen *before* the graph + # capture, since `reset` reallocates `mjw_data` (which the captured graph holds). + sim.reset() + sim.set_leg_adhesion_states(fly_name, np.ones((n_worlds, 6), dtype=np.float32)) + sim.warmup() + + # GPU-resident buffers for the captured loop. + target_angles_gpu = wp.array(target_angles_all_worlds) + curr_target_angles_gpu = wp.zeros((n_worlds, n_dofs), dtype=wp.float32) + step_counter = wp.array([0], dtype=wp.int32) + + # Capture the whole GPU-resident step body once (this triggers JIT). The recorder + # reads qpos *outside* the graph (a host transfer), so it is not captured here. + with wp.ScopedCapture() as advance_sim_capture: + wp.launch( + update_target_angles_kernel, + dim=(n_worlds, n_dofs), + inputs=[target_angles_gpu, step_counter], + outputs=[curr_target_angles_gpu], + ) + sim.set_actuator_inputs(fly_name, actuator_type, curr_target_angles_gpu) + sim.step() + wp.launch(increment_counter_kernel, dim=1, outputs=[step_counter]) + + # Untimed warm-up: force any remaining JIT, then reset the counter and recorder so + # recording starts cleanly from step 0. + print(f"Warming up (JIT compilation) {n_worlds} worlds...") + wp.capture_launch(advance_sim_capture.graph) + sim.render_as_needed() + wp.synchronize() + step_counter.zero_() + recorder.reset() + + print( + f"Simulating {sim_steps} steps across {n_worlds} worlds, " + f"recording {n_record}..." + ) + wp.synchronize() + start_time = perf_counter_ns() + for _ in range(sim_steps): + wp.capture_launch(advance_sim_capture.graph) + sim.render_as_needed() # records qpos for the selected worlds at the cadence + wp.synchronize() + walltime_s = (perf_counter_ns() - start_time) / 1e9 + + throughput = n_worlds * sim_steps / walltime_s + trajectories = recorder.recorded_trajectories + print( + f"Simulated {sim_steps} steps * {n_worlds} worlds in {walltime_s:.2f}s " + f"({throughput:.0f} steps/s, {throughput * timestep:.1f}x realtime).\n" + f"Recorded {len(trajectories)} trajectories of " + f"{trajectories[0].n_frames} frames each." + ) + return trajectories, world, sim + + +def main() -> None: + args = parse_args() + check_gpu() + + out = args.output + out.mkdir(parents=True, exist_ok=True) + model_dir = out / _MODEL_SUBDIR + traj_dir = out / _TRAJ_SUBDIR + + # --- Record --- + trajectories, world, sim = record_trajectories(args) + + # --- Persist: trajectories and model are independent artifacts --- + save_trajectories(trajectories, traj_dir) + world.save_xml_with_assets(model_dir, "model.xml") + print( + f"Saved {len(trajectories)} trajectories to {traj_dir} and the model to " + f"{model_dir}." + ) + + # --- Replay post-hoc, reloading the trajectories from disk --- + trajectories = load_trajectories(traj_dir) + + if args.cpu_replay: + # CPU replay needs no special model prep; reuse the unmodified compiled model. + cpu_out = out / "replay_cpu" + print(f"Rendering on CPU to {cpu_out}...") + render_trajectories(sim.mj_model, trajectories, cpu_out) + + # GPU batch rendering needs a batch-ready model (textures stripped, overhead + # lights added). The recorder ran against the unmodified model, but those edits + # don't change the qpos layout, so the trajectories stay valid. + gpu_out = out / "replay_gpu" + print(f"Rendering on GPU to {gpu_out}...") + modify_world_for_batch_rendering(world) + batch_model = world.compile()[0] + render_trajectories_gpu( + batch_model, trajectories, gpu_out, worlds_per_batch=args.worlds_per_batch + ) + print("Done.") + + +if __name__ == "__main__": + main() From bfab1fcf689b715ecb0832652fc8a01be4f40a22 Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Mon, 29 Jun 2026 22:42:55 +0200 Subject: [PATCH 04/11] Demo: record all worlds, render them in GPU batches of 10 Record a trajectory for every simulated world (drop the recorded-subset flag / `worlds=` selection) and render all of them post-hoc in GPU batches of --worlds-per-batch (default 10) -- highlighting that the render batch size is decoupled from the simulated world count. Lower the default --n-worlds to 50 since every world is now rendered. Co-Authored-By: Claude Opus 4.8 (1M context) --- scripts/record_replay_trajectories_gpu.py | 66 +++++++++++------------ 1 file changed, 30 insertions(+), 36 deletions(-) diff --git a/scripts/record_replay_trajectories_gpu.py b/scripts/record_replay_trajectories_gpu.py index 8403b48d..0a918ced 100644 --- a/scripts/record_replay_trajectories_gpu.py +++ b/scripts/record_replay_trajectories_gpu.py @@ -1,11 +1,12 @@ """Demo: record kinematic trajectories on GPU, then render them to video post-hoc. This demonstrates the trajectory recording / replay feature (issue #296): it -decouples how many worlds are *simulated* from how many are *rendered*. Buffering -full ``(n_worlds, n_cams, H, W, 3)`` RGB tensors during a large parallel run is -hopeless at thousands of worlds, so instead we record only the generalized -coordinates (``qpos``, plus mocap poses) of a small selected subset of worlds at the -render cadence, and rasterize them afterwards. +decouples how many worlds are *simulated* from how many are *rendered in one batch*. +Buffering full ``(n_worlds, n_cams, H, W, 3)`` RGB tensors during a large parallel run +is hopeless at thousands of worlds, so instead we record only the generalized +coordinates (``qpos``, plus mocap poses) at the render cadence, and rasterize them +afterwards in small GPU batches whose size is independent of the simulation's world +count. It is the trajectory-recording counterpart of ``replay_behavior_gpu.py``: the same Spotlight kinematic recording is replayed across many parallel worlds with the same @@ -16,13 +17,12 @@ The script then shows the full decoupled pipeline: -1. Simulate ``--n-worlds`` worlds, recording trajectories for ``--record-worlds`` of - them. +1. Simulate ``--n-worlds`` worlds, recording a trajectory for *every* world. 2. Save the trajectories (one ``.npz`` each) and the model (``save_xml_with_assets``) to disk -- they are independent artifacts; a trajectory carries no model. -3. Reload the trajectories from disk and render them to video, on GPU - (``render_trajectories_gpu``) and optionally on CPU (``--cpu-replay``) to show that - a GPU-recorded trajectory is backend-agnostic. +3. Reload the trajectories from disk and render every world to video, on GPU + (``render_trajectories_gpu``, in batches of ``--worlds-per-batch``) and optionally + on CPU (``--cpu-replay``) to show that a GPU-recorded trajectory is backend-agnostic. Example: uv run python scripts/record_replay_trajectories_gpu.py --output outputs/traj_demo @@ -68,15 +68,10 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--n-worlds", type=int, - default=1000, - help="Number of parallel worlds to simulate (default: 1000).", - ) - parser.add_argument( - "--record-worlds", - type=int, - default=5, - help="Number of worlds to record trajectories for and render (default: 5). " - "This is the whole point: simulate many, record/render few.", + default=50, + help="Number of parallel worlds to simulate; a trajectory is recorded for " + "every one (default: 50). Recording is cheap, but every world is rendered " + "afterwards, so keep this modest unless you want many output videos.", ) parser.add_argument( "--sim-steps", @@ -93,9 +88,10 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--worlds-per-batch", type=int, - default=None, - help="GPU batch size for post-hoc rendering (default: library heuristic). " - "Decoupled from --n-worlds; larger uses more GPU memory.", + default=10, + help="Number of frames staged into one GPU render batch (default: 10). This " + "is the render-time parallelism, decoupled from --n-worlds; larger uses more " + "GPU memory.", ) parser.add_argument( "--cpu-replay", @@ -107,16 +103,15 @@ def parse_args() -> argparse.Namespace: def record_trajectories(args: argparse.Namespace): - """Run the GPU simulation, recording qpos for a subset of worlds. + """Run the GPU simulation, recording qpos for every world. - Returns ``(trajectories, world, sim)``: the recorded trajectories (one per - recorded world), the world (kept so we can persist / re-compile the model), and - the simulation (kept for its unmodified ``mj_model``, used for CPU replay). + Returns ``(trajectories, world, sim)``: the recorded trajectories (one per world), + the world (kept so we can persist / re-compile the model), and the simulation + (kept for its unmodified ``mj_model``, used for CPU replay). """ n_worlds = args.n_worlds sim_steps = args.sim_steps timestep = args.timestep - n_record = min(args.record_worlds, n_worlds) actuator_type = ActuatorType.POSITION fly, world, cam = make_model() @@ -133,13 +128,12 @@ def record_trajectories(args: argparse.Namespace): sim = GPUSimulation(world, n_worlds, timestep=timestep) - # Swap the live batch renderer for a recorder: it stores qpos for the selected - # worlds at the render cadence instead of rasterizing frames. + # Swap the live batch renderer for a recorder: it stores qpos for every world at + # the render cadence instead of rasterizing frames. Omitting `worlds` records all. recorder = sim.set_renderer( cam, playback_speed=0.2, output_fps=25, - worlds=list(range(n_record)), record_trajectory_only=True, ) @@ -176,15 +170,12 @@ def record_trajectories(args: argparse.Namespace): step_counter.zero_() recorder.reset() - print( - f"Simulating {sim_steps} steps across {n_worlds} worlds, " - f"recording {n_record}..." - ) + print(f"Simulating {sim_steps} steps across {n_worlds} worlds (recording all)...") wp.synchronize() start_time = perf_counter_ns() for _ in range(sim_steps): wp.capture_launch(advance_sim_capture.graph) - sim.render_as_needed() # records qpos for the selected worlds at the cadence + sim.render_as_needed() # records qpos for every world at the cadence wp.synchronize() walltime_s = (perf_counter_ns() - start_time) / 1e9 @@ -232,7 +223,10 @@ def main() -> None: # lights added). The recorder ran against the unmodified model, but those edits # don't change the qpos layout, so the trajectories stay valid. gpu_out = out / "replay_gpu" - print(f"Rendering on GPU to {gpu_out}...") + print( + f"Rendering {len(trajectories)} worlds on GPU to {gpu_out} " + f"(batches of {args.worlds_per_batch})..." + ) modify_world_for_batch_rendering(world) batch_model = world.compile()[0] render_trajectories_gpu( From df9092887c52988bdb4d7682b8a711144d9f3982 Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Mon, 29 Jun 2026 23:36:31 +0200 Subject: [PATCH 05/11] fix #298 --- src/flygym/warp/rendering/live_rendering.py | 34 +++++++++++---------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/src/flygym/warp/rendering/live_rendering.py b/src/flygym/warp/rendering/live_rendering.py index 351bf8d5..58e3e578 100644 --- a/src/flygym/warp/rendering/live_rendering.py +++ b/src/flygym/warp/rendering/live_rendering.py @@ -443,21 +443,23 @@ def modify_world_for_batch_rendering(world: BaseWorld) -> bool: material.texrepeat = tuple(tr / 1000 for tr in material.texrepeat) is_modified = True - # Add light above each fly explicitly - for body in world.mjcf_root.bodies: - if body.name.split("/")[-1] == "c_thorax": - warnings.warn(f"Adding overhead light for body {body.name}") - body.add_light( - name=body.name.replace("/", "-") + "-overheadlight", - mode=mj.mjtCamLight.mjCAMLIGHT_TRACK, - targetbody=body.name, - pos=(0, 0, 30), - dir=(0, 0, -1), - type=mj.mjtLightType.mjLIGHT_DIRECTIONAL, - ambient=(10, 10, 10), - diffuse=(10, 10, 10), - specular=(0.3, 0.3, 0.3), - ) - is_modified = True + # Add light above each fly explicitly (only until MuJoCo Warp 3.9) + mujoco_warp_version = tuple(int(x) for x in mjw.__version__.split(".")[:2]) + if mujoco_warp_version < (3, 10): + for body in world.mjcf_root.bodies: + if body.name.split("/")[-1] == "c_thorax": + warnings.warn(f"Adding overhead light for body {body.name}") + body.add_light( + name=body.name.replace("/", "-") + "-overheadlight", + mode=mj.mjtCamLight.mjCAMLIGHT_TRACK, + targetbody=body.name, + pos=(0, 0, 30), + dir=(0, 0, -1), + type=mj.mjtLightType.mjLIGHT_DIRECTIONAL, + ambient=(10, 10, 10), + diffuse=(10, 10, 10), + specular=(0.3, 0.3, 0.3), + ) + is_modified = True return is_modified From ad8053874580fcddcd1ef10f404386b10ac073eb Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Mon, 29 Jun 2026 23:36:48 +0200 Subject: [PATCH 06/11] add test script for rendering via recorded trajectories --- scripts/record_replay_trajectories_gpu.py | 139 ++++++++-------------- 1 file changed, 50 insertions(+), 89 deletions(-) diff --git a/scripts/record_replay_trajectories_gpu.py b/scripts/record_replay_trajectories_gpu.py index 0a918ced..e526e862 100644 --- a/scripts/record_replay_trajectories_gpu.py +++ b/scripts/record_replay_trajectories_gpu.py @@ -1,12 +1,12 @@ """Demo: record kinematic trajectories on GPU, then render them to video post-hoc. This demonstrates the trajectory recording / replay feature (issue #296): it -decouples how many worlds are *simulated* from how many are *rendered in one batch*. -Buffering full ``(n_worlds, n_cams, H, W, 3)`` RGB tensors during a large parallel run -is hopeless at thousands of worlds, so instead we record only the generalized -coordinates (``qpos``, plus mocap poses) at the render cadence, and rasterize them -afterwards in small GPU batches whose size is independent of the simulation's world -count. +decouples how many worlds are *simulated* from how many are *rendered*, and from how +many render *in one GPU batch*. Buffering full ``(n_worlds, n_cams, H, W, 3)`` RGB +tensors during a large parallel run is hopeless at thousands of worlds, so instead we +record only the generalized coordinates (``qpos``, plus mocap poses) at the render +cadence, then afterwards sub-select however many worlds we actually want on video and +rasterize them in small GPU batches whose size is independent of the world count. It is the trajectory-recording counterpart of ``replay_behavior_gpu.py``: the same Spotlight kinematic recording is replayed across many parallel worlds with the same @@ -17,19 +17,19 @@ The script then shows the full decoupled pipeline: -1. Simulate ``--n-worlds`` worlds, recording a trajectory for *every* world. +1. Simulate ``N_WORLDS`` worlds, recording a trajectory for *every* world. 2. Save the trajectories (one ``.npz`` each) and the model (``save_xml_with_assets``) to disk -- they are independent artifacts; a trajectory carries no model. -3. Reload the trajectories from disk and render every world to video, on GPU - (``render_trajectories_gpu``, in batches of ``--worlds-per-batch``) and optionally - on CPU (``--cpu-replay``) to show that a GPU-recorded trajectory is backend-agnostic. +3. Reload the trajectories from disk, sub-select ``RENDER_WORLDS`` of them, and render + those to video on GPU (``render_trajectories_gpu``, in batches of + ``WORLDS_PER_BATCH``) and optionally on CPU (``CPU_REPLAY``) to show that a + GPU-recorded trajectory is backend-agnostic. -Example: - uv run python scripts/record_replay_trajectories_gpu.py --output outputs/traj_demo - uv run python scripts/record_replay_trajectories_gpu.py --output outputs/traj_demo --cpu-replay +Configure the run by editing the constants below, then:: + + uv run python scripts/record_replay_trajectories_gpu.py """ -import argparse from pathlib import Path from time import perf_counter_ns @@ -51,67 +51,27 @@ increment_counter_kernel, ) +# --- Configuration (edit these) ------------------------------------------------- +OUTPUT_DIR = Path("outputs/traj_demo") # trajectories, model, and rendered videos +N_WORLDS = 1000 # parallel worlds to simulate; a trajectory is recorded for each +RENDER_WORLDS = 50 # how many recorded worlds to render in the second stage +SIM_STEPS = 2000 # steps to simulate per world (2000 * 1e-4 s = 0.2 s) +TIMESTEP = 1e-4 # simulation timestep in seconds +WORLDS_PER_BATCH = 10 # GPU render batch size, decoupled from N_WORLDS +CPU_REPLAY = False # also replay on CPU (shows the format is backend-agnostic) +# -------------------------------------------------------------------------------- + _MODEL_SUBDIR = "model" _TRAJ_SUBDIR = "trajectories" -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--output", - type=Path, - default=Path("outputs/traj_demo"), - metavar="DIR", - help="Directory for the saved trajectories, model, and rendered videos " - "(default: outputs/traj_demo). Reused if it already exists.", - ) - parser.add_argument( - "--n-worlds", - type=int, - default=50, - help="Number of parallel worlds to simulate; a trajectory is recorded for " - "every one (default: 50). Recording is cheap, but every world is rendered " - "afterwards, so keep this modest unless you want many output videos.", - ) - parser.add_argument( - "--sim-steps", - type=int, - default=2000, - help="Number of steps to simulate per world (default: 2000 = 0.2 s).", - ) - parser.add_argument( - "--timestep", - type=float, - default=1e-4, - help="Simulation timestep in seconds (default: 1e-4).", - ) - parser.add_argument( - "--worlds-per-batch", - type=int, - default=10, - help="Number of frames staged into one GPU render batch (default: 10). This " - "is the render-time parallelism, decoupled from --n-worlds; larger uses more " - "GPU memory.", - ) - parser.add_argument( - "--cpu-replay", - action="store_true", - help="Additionally replay the GPU-recorded trajectories on the CPU, to " - "demonstrate that the recorded format is backend-agnostic.", - ) - return parser.parse_args() - - -def record_trajectories(args: argparse.Namespace): +def record_trajectories(): """Run the GPU simulation, recording qpos for every world. Returns ``(trajectories, world, sim)``: the recorded trajectories (one per world), the world (kept so we can persist / re-compile the model), and the simulation (kept for its unmodified ``mj_model``, used for CPU replay). """ - n_worlds = args.n_worlds - sim_steps = args.sim_steps - timestep = args.timestep actuator_type = ActuatorType.POSITION fly, world, cam = make_model() @@ -119,14 +79,14 @@ def record_trajectories(args: argparse.Namespace): # Build per-world target angle slices (world 0 -> first slice, world 1 -> next...). replay_data = ReplayTargetData( - timestep, fly.get_actuated_jointdofs_order(actuator_type) + TIMESTEP, fly.get_actuated_jointdofs_order(actuator_type) ) target_angles_all_worlds = replay_data.make_target_angles_all_worlds( - n_worlds, sim_steps + N_WORLDS, SIM_STEPS ) n_dofs = target_angles_all_worlds.shape[-1] - sim = GPUSimulation(world, n_worlds, timestep=timestep) + sim = GPUSimulation(world, N_WORLDS, timestep=TIMESTEP) # Swap the live batch renderer for a recorder: it stores qpos for every world at # the render cadence instead of rasterizing frames. Omitting `worlds` records all. @@ -140,12 +100,12 @@ def record_trajectories(args: argparse.Namespace): # Reset to the neutral keyframe and settle. Must happen *before* the graph # capture, since `reset` reallocates `mjw_data` (which the captured graph holds). sim.reset() - sim.set_leg_adhesion_states(fly_name, np.ones((n_worlds, 6), dtype=np.float32)) + sim.set_leg_adhesion_states(fly_name, np.ones((N_WORLDS, 6), dtype=np.float32)) sim.warmup() # GPU-resident buffers for the captured loop. target_angles_gpu = wp.array(target_angles_all_worlds) - curr_target_angles_gpu = wp.zeros((n_worlds, n_dofs), dtype=wp.float32) + curr_target_angles_gpu = wp.zeros((N_WORLDS, n_dofs), dtype=wp.float32) step_counter = wp.array([0], dtype=wp.int32) # Capture the whole GPU-resident step body once (this triggers JIT). The recorder @@ -153,7 +113,7 @@ def record_trajectories(args: argparse.Namespace): with wp.ScopedCapture() as advance_sim_capture: wp.launch( update_target_angles_kernel, - dim=(n_worlds, n_dofs), + dim=(N_WORLDS, n_dofs), inputs=[target_angles_gpu, step_counter], outputs=[curr_target_angles_gpu], ) @@ -163,27 +123,27 @@ def record_trajectories(args: argparse.Namespace): # Untimed warm-up: force any remaining JIT, then reset the counter and recorder so # recording starts cleanly from step 0. - print(f"Warming up (JIT compilation) {n_worlds} worlds...") + print(f"Warming up (JIT compilation) {N_WORLDS} worlds...") wp.capture_launch(advance_sim_capture.graph) sim.render_as_needed() wp.synchronize() step_counter.zero_() recorder.reset() - print(f"Simulating {sim_steps} steps across {n_worlds} worlds (recording all)...") + print(f"Simulating {SIM_STEPS} steps across {N_WORLDS} worlds (recording all)...") wp.synchronize() start_time = perf_counter_ns() - for _ in range(sim_steps): + for _ in range(SIM_STEPS): wp.capture_launch(advance_sim_capture.graph) sim.render_as_needed() # records qpos for every world at the cadence wp.synchronize() walltime_s = (perf_counter_ns() - start_time) / 1e9 - throughput = n_worlds * sim_steps / walltime_s + throughput = N_WORLDS * SIM_STEPS / walltime_s trajectories = recorder.recorded_trajectories print( - f"Simulated {sim_steps} steps * {n_worlds} worlds in {walltime_s:.2f}s " - f"({throughput:.0f} steps/s, {throughput * timestep:.1f}x realtime).\n" + f"Simulated {SIM_STEPS} steps * {N_WORLDS} worlds in {walltime_s:.2f}s " + f"({throughput:.0f} steps/s, {throughput * TIMESTEP:.1f}x realtime).\n" f"Recorded {len(trajectories)} trajectories of " f"{trajectories[0].n_frames} frames each." ) @@ -191,16 +151,14 @@ def record_trajectories(args: argparse.Namespace): def main() -> None: - args = parse_args() check_gpu() - out = args.output - out.mkdir(parents=True, exist_ok=True) - model_dir = out / _MODEL_SUBDIR - traj_dir = out / _TRAJ_SUBDIR + OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + model_dir = OUTPUT_DIR / _MODEL_SUBDIR + traj_dir = OUTPUT_DIR / _TRAJ_SUBDIR # --- Record --- - trajectories, world, sim = record_trajectories(args) + trajectories, world, sim = record_trajectories() # --- Persist: trajectories and model are independent artifacts --- save_trajectories(trajectories, traj_dir) @@ -210,27 +168,30 @@ def main() -> None: f"{model_dir}." ) - # --- Replay post-hoc, reloading the trajectories from disk --- + # --- Replay post-hoc, reloading the trajectories from disk and sub-selecting --- trajectories = load_trajectories(traj_dir) + n_render = min(RENDER_WORLDS, len(trajectories)) + trajectories = trajectories[:n_render] + print(f"Reloaded trajectories; rendering {n_render} of them.") - if args.cpu_replay: + if CPU_REPLAY: # CPU replay needs no special model prep; reuse the unmodified compiled model. - cpu_out = out / "replay_cpu" + cpu_out = OUTPUT_DIR / "replay_cpu" print(f"Rendering on CPU to {cpu_out}...") render_trajectories(sim.mj_model, trajectories, cpu_out) # GPU batch rendering needs a batch-ready model (textures stripped, overhead # lights added). The recorder ran against the unmodified model, but those edits # don't change the qpos layout, so the trajectories stay valid. - gpu_out = out / "replay_gpu" + gpu_out = OUTPUT_DIR / "replay_gpu" print( f"Rendering {len(trajectories)} worlds on GPU to {gpu_out} " - f"(batches of {args.worlds_per_batch})..." + f"(batches of {WORLDS_PER_BATCH})..." ) modify_world_for_batch_rendering(world) batch_model = world.compile()[0] render_trajectories_gpu( - batch_model, trajectories, gpu_out, worlds_per_batch=args.worlds_per_batch + batch_model, trajectories, gpu_out, worlds_per_batch=WORLDS_PER_BATCH ) print("Done.") From 72a1cf3e25faa9d2eb983235d602be9007868f7d Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Mon, 29 Jun 2026 23:53:44 +0200 Subject: [PATCH 07/11] Replace save_trajectories/load_trajectories with RecordedTrajectory methods Now that the model is no longer saved alongside trajectories, the free functions earned their keep only as thin folder loops. Replace them with RecordedTrajectory.save() and RecordedTrajectory.from_file() (classmethod); each trajectory is one self-describing .npz. Callers that want a folder of many just loop. Updates the package re-exports, the GPU recorder error messages/docstrings, the demo script, and the tests accordingly. Co-Authored-By: Claude Opus 4.8 (1M context) --- scripts/record_replay_trajectories_gpu.py | 12 ++-- src/flygym/__init__.py | 4 -- src/flygym/rendering/__init__.py | 4 -- src/flygym/rendering/recorded_trajectory.py | 48 +++----------- .../warp/rendering/recorded_trajectory.py | 12 ++-- tests/core/test_trajectory.py | 65 ++++++++++--------- tests/warp/test_trajectory.py | 10 +-- 7 files changed, 61 insertions(+), 94 deletions(-) diff --git a/scripts/record_replay_trajectories_gpu.py b/scripts/record_replay_trajectories_gpu.py index e526e862..99c97cab 100644 --- a/scripts/record_replay_trajectories_gpu.py +++ b/scripts/record_replay_trajectories_gpu.py @@ -42,7 +42,7 @@ modify_world_for_batch_rendering, ) from flygym.warp.utils import check_gpu -from flygym.rendering import save_trajectories, load_trajectories, render_trajectories +from flygym.rendering import RecordedTrajectory, render_trajectories from flygym.compose import ActuatorType from flygym_demo.benchmark import ( make_model, @@ -160,8 +160,11 @@ def main() -> None: # --- Record --- trajectories, world, sim = record_trajectories() - # --- Persist: trajectories and model are independent artifacts --- - save_trajectories(trajectories, traj_dir) + # --- Persist: each trajectory is one self-describing .npz; the model is a + # separate artifact (a trajectory carries no model). --- + traj_dir.mkdir(parents=True, exist_ok=True) + for i, traj in enumerate(trajectories): + traj.save(traj_dir / f"traj_{i:04d}.npz") world.save_xml_with_assets(model_dir, "model.xml") print( f"Saved {len(trajectories)} trajectories to {traj_dir} and the model to " @@ -169,7 +172,8 @@ def main() -> None: ) # --- Replay post-hoc, reloading the trajectories from disk and sub-selecting --- - trajectories = load_trajectories(traj_dir) + traj_files = sorted(traj_dir.glob("traj_*.npz")) + trajectories = [RecordedTrajectory.from_file(p) for p in traj_files] n_render = min(RENDER_WORLDS, len(trajectories)) trajectories = trajectories[:n_render] print(f"Reloaded trajectories; rendering {n_render} of them.") diff --git a/src/flygym/__init__.py b/src/flygym/__init__.py index 728e113e..d3cf8bc2 100644 --- a/src/flygym/__init__.py +++ b/src/flygym/__init__.py @@ -13,8 +13,6 @@ Renderer, TrajectoryRecorder, RecordedTrajectory, - save_trajectories, - load_trajectories, render_trajectories, launch_interactive_viewer, preview_model, @@ -29,8 +27,6 @@ "Renderer", "TrajectoryRecorder", "RecordedTrajectory", - "save_trajectories", - "load_trajectories", "render_trajectories", "launch_interactive_viewer", "preview_model", diff --git a/src/flygym/rendering/__init__.py b/src/flygym/rendering/__init__.py index f02a96b3..06f401b1 100644 --- a/src/flygym/rendering/__init__.py +++ b/src/flygym/rendering/__init__.py @@ -18,8 +18,6 @@ from flygym.rendering.recorded_trajectory import ( RecordedTrajectory, TrajectoryRecorder, - save_trajectories, - load_trajectories, render_trajectories, ) @@ -27,8 +25,6 @@ "Renderer", "TrajectoryRecorder", "RecordedTrajectory", - "save_trajectories", - "load_trajectories", "render_trajectories", "launch_interactive_viewer", "preview_model", diff --git a/src/flygym/rendering/recorded_trajectory.py b/src/flygym/rendering/recorded_trajectory.py index 25eb6651..ca682c37 100644 --- a/src/flygym/rendering/recorded_trajectory.py +++ b/src/flygym/rendering/recorded_trajectory.py @@ -42,8 +42,6 @@ __all__ = [ "RecordedTrajectory", "TrajectoryRecorder", - "save_trajectories", - "load_trajectories", "render_trajectories", ] @@ -93,7 +91,9 @@ def save(self, path: PathLike) -> None: """Save this trajectory to a single self-describing ``.npz`` file. Both the per-frame state and the replay metadata are stored, so the file can be - read back with `load` without any side information. + read back with `from_file` without any side information. The model is + intentionally not saved -- persist it yourself (e.g. + ``world.save_xml_with_assets(...)``) and recompile it at replay time. """ path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) @@ -111,7 +111,7 @@ def save(self, path: PathLike) -> None: np.savez_compressed(path, **arrays) @classmethod - def load(cls, path: PathLike) -> "RecordedTrajectory": + def from_file(cls, path: PathLike) -> "RecordedTrajectory": """Load a trajectory from a ``.npz`` file written by `save`.""" with np.load(path, allow_pickle=False) as data: has_mocap = "mocap_pos" in data @@ -228,49 +228,16 @@ def close(self) -> None: def show_in_notebook(self, *args: Any, **kwargs: Any) -> None: raise RuntimeError( "TrajectoryRecorder records state, not frames. Save it with " - "save_trajectories and replay with render_trajectories." + "RecordedTrajectory.save and replay with render_trajectories." ) def save_video(self, *args: Any, **kwargs: Any) -> None: raise RuntimeError( "TrajectoryRecorder records state, not frames. Save it with " - "save_trajectories and replay with render_trajectories." + "RecordedTrajectory.save and replay with render_trajectories." ) -def save_trajectories( - trajectories: RecordedTrajectory | list[RecordedTrajectory], - output_dir: PathLike, -) -> None: - """Save trajectories as individual ``.npz`` files in a folder. - - Writes one ``traj_XXXX.npz`` per trajectory (see `RecordedTrajectory.save`); each - file is self-describing. The model is intentionally *not* saved here -- persist it - yourself (e.g. ``world.save_xml_with_assets(...)``) and recompile it at replay time. - - Args: - trajectories: One trajectory or a list of them. - output_dir: Destination folder (created if needed). - """ - trajectories = _as_trajectory_list(trajectories) - output_dir = Path(output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - for i, traj in enumerate(trajectories): - traj.save(output_dir / f"traj_{i:04d}.npz") - - -def load_trajectories(source: PathLike) -> list[RecordedTrajectory]: - """Load all trajectories from a folder written by `save_trajectories`. - - Returns the trajectories only; supply the compiled model yourself at replay time. - """ - source = Path(source) - files = sorted(source.glob("traj_*.npz")) - if len(files) == 0: - raise ValueError(f"No trajectory files (traj_*.npz) found in {source}.") - return [RecordedTrajectory.load(f) for f in files] - - def _as_trajectory_list( trajectories: RecordedTrajectory | list[RecordedTrajectory], ) -> list[RecordedTrajectory]: @@ -406,7 +373,8 @@ def render_trajectories( wherever you persisted it, e.g. a folder written by `BaseCompositionElement.save_xml_with_assets`). Its ``qpos`` layout must match the trajectories. - trajectories: One trajectory or a list of them (e.g. from `load_trajectories`). + trajectories: One trajectory or a list of them (load saved ones with + `RecordedTrajectory.from_file`). output_path: Where to write videos. See `_resolve_render_output_paths` for the file/directory layout. cameras: Camera name(s) to render. Defaults to each trajectory's recorded diff --git a/src/flygym/warp/rendering/recorded_trajectory.py b/src/flygym/warp/rendering/recorded_trajectory.py index 06331595..2433f2ca 100644 --- a/src/flygym/warp/rendering/recorded_trajectory.py +++ b/src/flygym/warp/rendering/recorded_trajectory.py @@ -98,16 +98,16 @@ def _fetch_frames_to_cpu_impl(self, world_id_among_rendered, cam_id_among_render @override def save_video(self, *args: Any, **kwargs: Any) -> None: raise RuntimeError( - "WarpTrajectoryRecorder records state, not frames. Save it with " - "flygym.rendering.save_trajectories and replay with " + "WarpTrajectoryRecorder records state, not frames. Save each with " + "RecordedTrajectory.save and replay with " "render_trajectories_gpu / flygym.rendering.render_trajectories." ) @override def show_in_notebook(self, *args: Any, **kwargs: Any) -> None: raise RuntimeError( - "WarpTrajectoryRecorder records state, not frames. Save it with " - "flygym.rendering.save_trajectories and replay with " + "WarpTrajectoryRecorder records state, not frames. Save each with " + "RecordedTrajectory.save and replay with " "render_trajectories_gpu / flygym.rendering.render_trajectories." ) @@ -151,8 +151,8 @@ def render_trajectories_gpu( Args: mj_model: Compiled, batch-render-ready model (see note above). Its ``qpos`` layout must match the trajectories. - trajectories: One trajectory or a list of them (e.g. from - `flygym.rendering.load_trajectories`). + trajectories: One trajectory or a list of them (load saved ones with + `flygym.rendering.RecordedTrajectory.from_file`). output_path: Where to write videos (see `_resolve_render_output_paths`). cameras: Camera name(s) to render. Defaults to each trajectory's recorded ``camera_names``. All trajectories must share the same camera set and diff --git a/tests/core/test_trajectory.py b/tests/core/test_trajectory.py index 15b29895..9b2a735f 100644 --- a/tests/core/test_trajectory.py +++ b/tests/core/test_trajectory.py @@ -14,12 +14,23 @@ from flygym.simulation import Simulation from flygym.rendering import ( RecordedTrajectory, - save_trajectories, - load_trajectories, render_trajectories, ) from flygym.rendering.recorded_trajectory import _render_trajectory_frames + +def _save_all(trajectories, folder): + """Save a list of trajectories as traj_XXXX.npz in a folder (test helper).""" + folder.mkdir(parents=True, exist_ok=True) + for i, traj in enumerate(trajectories): + traj.save(folder / f"traj_{i:04d}.npz") + + +def _load_all(folder): + """Load all trajectories from a folder of traj_*.npz (test helper).""" + return [RecordedTrajectory.from_file(p) for p in sorted(folder.glob("traj_*.npz"))] + + # Rendering (rasterization) needs a headless GL context; skip those on CI runners # that set SKIP_RENDERING_TESTS=1. Recording and serialization need no GL. needs_gl = pytest.mark.skipif( @@ -124,45 +135,35 @@ def test_reset_clears_buffer(self, sim_with_camera): class TestSaveLoad: - def test_roundtrip(self, recorded, tmp_path): + def test_save_from_file_roundtrip(self, recorded, tmp_path): traj, _, _ = recorded - save_trajectories(traj, tmp_path) - assert (tmp_path / "traj_0000.npz").exists() - trajs = load_trajectories(tmp_path) - assert len(trajs) == 1 - assert np.array_equal(trajs[0].qpos, traj.qpos) - assert trajs[0].camera_res == traj.camera_res - assert trajs[0].camera_names == traj.camera_names - assert trajs[0].output_fps == traj.output_fps + path = tmp_path / "one.npz" + traj.save(path) + assert path.exists() + loaded = RecordedTrajectory.from_file(path) + assert np.array_equal(loaded.qpos, traj.qpos) + assert loaded.camera_res == traj.camera_res + assert loaded.camera_names == traj.camera_names + assert loaded.output_fps == traj.output_fps + assert loaded.world_id == traj.world_id def test_mocap_roundtrip(self, recorded, tmp_path): traj, _, _ = recorded - save_trajectories(traj, tmp_path) - loaded = load_trajectories(tmp_path)[0] + path = tmp_path / "m.npz" + traj.save(path) + loaded = RecordedTrajectory.from_file(path) assert loaded.has_mocap == traj.has_mocap if traj.has_mocap: assert np.array_equal(loaded.mocap_pos, traj.mocap_pos) assert np.array_equal(loaded.mocap_quat, traj.mocap_quat) - def test_save_list(self, recorded, tmp_path): + def test_save_load_multiple(self, recorded, tmp_path): traj, _, _ = recorded - save_trajectories([traj, traj], tmp_path) + _save_all([traj, traj], tmp_path) assert (tmp_path / "traj_0001.npz").exists() - trajs = load_trajectories(tmp_path) + trajs = _load_all(tmp_path) assert len(trajs) == 2 - def test_load_empty_folder_raises(self, tmp_path): - with pytest.raises(ValueError, match="No trajectory files"): - load_trajectories(tmp_path) - - def test_single_trajectory_save_load(self, recorded, tmp_path): - traj, _, _ = recorded - path = tmp_path / "one.npz" - traj.save(path) - loaded = RecordedTrajectory.load(path) - assert np.array_equal(loaded.qpos, traj.qpos) - assert loaded.world_id == traj.world_id - # --------------------------------------------------------------------------- # Replay validation (no GL: the guard runs before any rasterization) @@ -197,11 +198,11 @@ def test_render_from_memory_writes_video(self, recorded, tmp_path): render_trajectories(sim.mj_model, traj, out) assert out.exists() and out.stat().st_size > 0 - def test_render_from_folder_writes_video(self, recorded, tmp_path): + def test_render_from_saved_file_writes_video(self, recorded, tmp_path): traj, sim, _ = recorded - folder = tmp_path / "folder" - save_trajectories(traj, folder) - trajs = load_trajectories(folder) + path = tmp_path / "traj.npz" + traj.save(path) + trajs = [RecordedTrajectory.from_file(path)] out = tmp_path / "out.mp4" render_trajectories(sim.mj_model, trajs, out) assert out.exists() and out.stat().st_size > 0 diff --git a/tests/warp/test_trajectory.py b/tests/warp/test_trajectory.py index 5518f1cb..06e51cef 100644 --- a/tests/warp/test_trajectory.py +++ b/tests/warp/test_trajectory.py @@ -12,8 +12,7 @@ pytest.importorskip("warp") from flygym.rendering import ( - save_trajectories, - load_trajectories, + RecordedTrajectory, render_trajectories, ) from flygym.warp import ( @@ -96,8 +95,11 @@ class TestSaveLoadGPU: def test_roundtrip(self, recorded_gpu, tmp_path): rec, _, _ = recorded_gpu trajs = rec.recorded_trajectories - save_trajectories(trajs, tmp_path) - loaded = load_trajectories(tmp_path) + for i, traj in enumerate(trajs): + traj.save(tmp_path / f"traj_{i:04d}.npz") + loaded = [ + RecordedTrajectory.from_file(p) for p in sorted(tmp_path.glob("traj_*.npz")) + ] assert len(loaded) == 2 for a, b in zip(loaded, trajs): assert np.array_equal(a.qpos, b.qpos) From 51f64c666f947861dd3742a9e8169a99c19ac1db Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Thu, 2 Jul 2026 00:15:30 +0200 Subject: [PATCH 08/11] add aliases for fly --- src/flygym/compose/world/base_world.py | 18 ++++++++++++++++-- src/flygym/simulation.py | 12 +++++++++++- 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/src/flygym/compose/world/base_world.py b/src/flygym/compose/world/base_world.py index 25b2bad7..22317041 100644 --- a/src/flygym/compose/world/base_world.py +++ b/src/flygym/compose/world/base_world.py @@ -81,12 +81,26 @@ def mjcf_root(self) -> mj.MjSpec: def fly_lookup(self) -> dict[str, BaseFly]: """Lookup for `Fly` objects in the world, keyed by fly name.""" return self._fly_lookup + + @property + def fly(self) -> BaseFly: + """Get the single fly in the world. + + Raises: + ValueError: If there is not exactly one fly in the world. + """ + if len(self.fly_lookup) != 1: + raise ValueError( + "World contains multiple flies. " + "`.fly` is ambiguous; use `.fly_lookup` instead." + ) + return next(iter(self.fly_lookup.values())) @abstractmethod def _attach_fly_mjcf( self, fly: BaseFly, - spawn_position: Vec3, + spawn_position: Vec3 | tuple[float, float, float], spawn_rotation: Rotation3D, *args, **kwargs, @@ -125,7 +139,7 @@ def _add_skybox(self): def add_fly( self, fly: BaseFly, - spawn_position: Vec3, + spawn_position: Vec3 | tuple[float, float, float], spawn_rotation: Rotation3D, *args: Any, **kwargs: Any, diff --git a/src/flygym/simulation.py b/src/flygym/simulation.py index 00f510e0..069eddef 100644 --- a/src/flygym/simulation.py +++ b/src/flygym/simulation.py @@ -7,7 +7,7 @@ from jaxtyping import Float from flygym.anatomy import BodySegment -from flygym.compose.fly import ActuatorType +from flygym.compose.fly import BaseFly, ActuatorType from flygym.compose.world import BaseWorld from flygym.rendering import Renderer, TrajectoryRecorder from flygym.utils.profiling import print_perf_report @@ -777,3 +777,13 @@ def close(self): self.eye_renderer = None # Don't destruct self.retina and self.eye_renderer_scene_option: they can be # reused and retina init requires some IO ops. + + @property + def fly(self) -> BaseFly: + """Return the single fly in the world, or raise an error if there are multiple.""" + return self.world.fly + + @property + def fly_lookup(self) -> dict[str, BaseFly]: + """Return the fly lookup dictionary from the world.""" + return self.world.fly_lookup \ No newline at end of file From 1c27d7c70e1713051e3e5f5c98a91e1f32fe6fa5 Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Thu, 2 Jul 2026 00:16:07 +0200 Subject: [PATCH 09/11] add quick access for n_dofs, etc --- src/flygym/compose/fly/base_fly.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/src/flygym/compose/fly/base_fly.py b/src/flygym/compose/fly/base_fly.py index 3d63c183..9f8bf3d2 100644 --- a/src/flygym/compose/fly/base_fly.py +++ b/src/flygym/compose/fly/base_fly.py @@ -269,10 +269,20 @@ def get_bodysegs_order(self) -> list[BodySegment]: """ return list(self.bodyseg_to_mjcfbody.keys()) + @property + def n_bodysegs(self) -> int: + """Number of body segments in this fly.""" + return len(self.bodyseg_to_mjcfbody) + def get_jointdofs_order(self) -> list[JointDOF]: """Same as `get_bodysegs_order()`, but for joint DoFs instead of body segments.""" return list(self.jointdof_to_mjcfjoint.keys()) + @property + def n_jointdofs(self) -> int: + """Number of joint DoFs in this fly.""" + return len(self.jointdof_to_mjcfjoint) + def get_actuated_jointdofs_order( self, actuator_type: "ActuatorType | str" ) -> list[JointDOF]: @@ -282,6 +292,19 @@ def get_actuated_jointdofs_order( actuator_type = ActuatorType(actuator_type) return list(self.jointdof_to_mjcfactuator_by_type[actuator_type].keys()) + @property + def n_actuated_jointdofs(self) -> int: + raise RuntimeError( + "`n_actuated_jointdofs` is ambiguous because there might be different " + "actuator types. Use `get_n_actuated_jointdofs(actuator_type)` instead, " + "similar to `fly.get_actuated_jointdofs_order(actuator_type)`." + ) + + def get_n_actuated_jointdofs(self, actuator_type: "ActuatorType | str") -> int: + """Number of joint DoFs actuated by the specified actuator type.""" + actuator_type = ActuatorType(actuator_type) + return len(self.jointdof_to_mjcfactuator_by_type[actuator_type]) + def get_legs_order(self) -> list[str]: """Get the ordered list of leg position identifiers (same as `anatomy.LEGS`).""" return LEGS From 9fc5f02d9e525c4494a3d826b2fc7928a5e02989 Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Thu, 2 Jul 2026 00:18:04 +0200 Subject: [PATCH 10/11] allow lazy getters to load data into existing buffer --- src/flygym/warp/simulation.py | 44 ++++++++++++++++++++++++----------- 1 file changed, 30 insertions(+), 14 deletions(-) diff --git a/src/flygym/warp/simulation.py b/src/flygym/warp/simulation.py index d8cfba11..84cddd5e 100644 --- a/src/flygym/warp/simulation.py +++ b/src/flygym/warp/simulation.py @@ -78,19 +78,22 @@ def reset(self) -> None: @override def get_joint_angles( - self, fly_name: str + self, fly_name: str, dst: wp.array | None = None ) -> Float[wp.array, "n_worlds n_jointdofs"]: """Get joint angles for all parallel worlds. Args: fly_name: Name of the fly. + dst: Optional warp array to store the result. If not specified, a new array + is allocated. Returns: Warp array of shape ``(n_worlds, n_jointdofs)`` in radians, ordered as in ``fly.get_jointdofs_order()``. """ indices = self._wp_intern_qposadrs_by_fly[fly_name] - dst = wp.zeros((self.n_worlds, indices.size), dtype=wp.float32) + if dst is None: + dst = wp.zeros((self.n_worlds, indices.size), dtype=wp.float32) wp.launch( wp_gather_indexed_cols_2d, dim=(self.n_worlds, indices.size), @@ -100,19 +103,22 @@ def get_joint_angles( @override def get_joint_velocities( - self, fly_name: str + self, fly_name: str, dst: wp.array | None = None ) -> Float[wp.array, "n_worlds n_jointdofs"]: """Get joint velocities for all parallel worlds. Args: fly_name: Name of the fly. + dst: Optional warp array to store the result. If not specified, a new array + is allocated. Returns: Warp array of shape ``(n_worlds, n_jointdofs)`` in radians per second, ordered as in ``fly.get_jointdofs_order()``. """ indices = self._wp_intern_qveladrs_by_fly[fly_name] - dst = wp.zeros((self.n_worlds, indices.size), dtype=wp.float32) + if dst is None: + dst = wp.zeros((self.n_worlds, indices.size), dtype=wp.float32) wp.launch( wp_gather_indexed_cols_2d, dim=(self.n_worlds, indices.size), @@ -122,19 +128,22 @@ def get_joint_velocities( @override def get_body_positions( - self, fly_name: str + self, fly_name: str, dst: wp.array | None = None ) -> Float[wp.array, "n_worlds n_bodies 3"]: """Get global body positions for all parallel worlds. Args: fly_name: Name of the fly. + dst: Optional warp array to store the result. If not specified, a new array + is allocated. Returns: Warp array of shape ``(n_worlds, n_bodies, 3)`` in mm, ordered as in ``fly.get_bodysegs_order()``. """ indices = self._wp_internal_bodyids_by_fly[fly_name] - dst = wp.zeros((self.n_worlds, indices.size, 3), dtype=wp.float32) + if dst is None: + dst = wp.zeros((self.n_worlds, indices.size, 3), dtype=wp.float32) wp.launch( wp_gather_indexed_rows_vec3f, dim=(self.n_worlds, indices.size), @@ -144,19 +153,22 @@ def get_body_positions( @override def get_body_rotations( - self, fly_name: str + self, fly_name: str, dst: wp.array | None = None ) -> Float[wp.array, "n_worlds n_bodies 4"]: """Get global body orientations as quaternions for all parallel worlds. Args: fly_name: Name of the fly. + dst: Optional warp array to store the result. If not specified, a new array + is allocated. Returns: Warp array of shape ``(n_worlds, n_bodies, 4)`` (w, x, y, z), ordered as in ``fly.get_bodysegs_order()``. """ indices = self._wp_internal_bodyids_by_fly[fly_name] - dst = wp.zeros((self.n_worlds, indices.size, 4), dtype=wp.float32) + if dst is None: + dst = wp.zeros((self.n_worlds, indices.size, 4), dtype=wp.float32) wp.launch( wp_gather_indexed_rows_quatf, dim=(self.n_worlds, indices.size), @@ -166,19 +178,22 @@ def get_body_rotations( @override def get_site_positions( - self, fly_name: str + self, fly_name: str, dst: wp.array | None = None ) -> Float[wp.array, "n_worlds n_sites 3"]: """Get global anatomical-joint site positions for all parallel worlds. Args: fly_name: Name of the fly. + dst: Optional warp array to store the result. If not specified, a new array + is allocated. Returns: Warp array of shape ``(n_worlds, n_sites, 3)`` in mm, ordered as in ``fly.get_sites_order()``. """ indices = self._wp_internal_siteids_by_fly[fly_name] - dst = wp.zeros((self.n_worlds, indices.size, 3), dtype=wp.float32) + if dst is None: + dst = wp.zeros((self.n_worlds, indices.size, 3), dtype=wp.float32) wp.launch( wp_gather_indexed_rows_vec3f, dim=(self.n_worlds, indices.size), @@ -193,22 +208,23 @@ def time(self) -> float: @override def get_actuator_forces( - self, - fly_name: str, - actuator_type: ActuatorType, + self, fly_name: str, actuator_type: ActuatorType, dst: wp.array | None = None ) -> Float[wp.array, "n_worlds n_actuators"]: """Get actuator forces for all parallel worlds. Args: fly_name: Name of the fly. actuator_type: Type of actuator to query. + dst: Optional warp array to store the result. If not specified, a new array + is allocated. Returns: Warp array of shape ``(n_worlds, n_actuators)``, ordered as in ``fly.get_actuated_jointdofs_order(actuator_type)``. """ indices = self._wp_intern_actuatorids_by_type_by_fly[actuator_type][fly_name] - dst = wp.zeros((self.n_worlds, indices.size), dtype=wp.float32) + if dst is None: + dst = wp.zeros((self.n_worlds, indices.size), dtype=wp.float32) wp.launch( wp_gather_indexed_cols_2d, dim=(self.n_worlds, indices.size), From 075d294c84533d1a7347d106716a4283868f70a0 Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Thu, 2 Jul 2026 00:19:24 +0200 Subject: [PATCH 11/11] wip - refactor rendering --- src/flygym/rendering/live_rendering.py | 2 +- src/flygym/warp/__init__.py | 2 + src/flygym/warp/rendering/__init__.py | 2 + src/flygym/warp/rendering/base.py | 287 ++++++++++++++++++ src/flygym/warp/rendering/live_rendering.py | 269 +--------------- .../warp/rendering/recorded_trajectory.py | 37 ++- src/flygym/warp/simulation.py | 43 +-- tests/warp/test_rendering.py | 30 +- tests/warp/test_simulation.py | 8 +- tests/warp/test_trajectory.py | 3 +- 10 files changed, 353 insertions(+), 330 deletions(-) create mode 100644 src/flygym/warp/rendering/base.py diff --git a/src/flygym/rendering/live_rendering.py b/src/flygym/rendering/live_rendering.py index ad1dc2a2..6cc44c7c 100644 --- a/src/flygym/rendering/live_rendering.py +++ b/src/flygym/rendering/live_rendering.py @@ -89,7 +89,7 @@ def __init__( self.scene_option = mj.MjvOption() else: self.scene_option = scene_option - mj.mjv_defaultOption(self.scene_option) # this sets default scene options + mj.mjv_defaultOption(self.scene_option) self._cameras_names2id = {} for spec in cameras if isinstance(cameras, list) else [cameras]: diff --git a/src/flygym/warp/__init__.py b/src/flygym/warp/__init__.py index 6a8abfeb..0b848720 100644 --- a/src/flygym/warp/__init__.py +++ b/src/flygym/warp/__init__.py @@ -1,5 +1,6 @@ from .simulation import GPUSimulation from .rendering import ( + RendererType, WarpGPUBatchRenderer, WarpCPURenderer, WarpTrajectoryRecorder, @@ -9,6 +10,7 @@ __all__ = [ "GPUSimulation", + "RendererType", "WarpGPUBatchRenderer", "WarpCPURenderer", "WarpTrajectoryRecorder", diff --git a/src/flygym/warp/rendering/__init__.py b/src/flygym/warp/rendering/__init__.py index 85bf23a0..8e368b15 100644 --- a/src/flygym/warp/rendering/__init__.py +++ b/src/flygym/warp/rendering/__init__.py @@ -5,6 +5,7 @@ compatibility. """ +from flygym.warp.rendering.base import RendererType from flygym.warp.rendering.live_rendering import ( WarpGPUBatchRenderer, WarpCPURenderer, @@ -16,6 +17,7 @@ ) __all__ = [ + "RendererType", "WarpGPUBatchRenderer", "WarpCPURenderer", "WarpTrajectoryRecorder", diff --git a/src/flygym/warp/rendering/base.py b/src/flygym/warp/rendering/base.py new file mode 100644 index 00000000..f9350fa2 --- /dev/null +++ b/src/flygym/warp/rendering/base.py @@ -0,0 +1,287 @@ +from typing import Any, override +from os import PathLike +from abc import ABC, abstractmethod +from enum import Enum + +import mediapy +import mujoco as mj +import mujoco_warp as mjw +import warp as wp +import numpy as np +from PIL import Image, ImageDraw, ImageFont + +from flygym.rendering.live_rendering import Renderer +from flygym.utils.video import write_video_from_frames +from flygym.utils.plot import find_font_path + +__all__ = ["RendererType", "_BaseWarpRenderer"] + + +class RendererType(Enum): + CPU = "cpu" + GPU_BATCH = "gpu_batch" + RECORDED_TRAJECTORY = "recorded_trajectory" + + +class _BaseWarpRenderer(Renderer, ABC): + @override + def __init__( + self, + mj_model: mj.MjModel, + cameras: str | mj.MjsCamera | list[str | mj.MjsCamera], + sim_timestep: float, + n_worlds_total: int | None = None, + *, + worlds: list[int] | None = None, + camera_res: tuple[int, int] = (240, 320), + playback_speed: float = 0.2, + output_fps: int = 25, + buffer_frames: bool = True, + scene_option: mj.MjvOption | None = None, + **kwargs: Any, + ): + # Common setup using MjRenderer + super().__init__( + mj_model, + cameras, + camera_res=camera_res, + playback_speed=playback_speed, + output_fps=output_fps, + buffer_frames=buffer_frames, + scene_option=scene_option, + **kwargs, + ) + + # Warp-specific setup + self.mjw_model = mjw.put_model(mj_model) + self._n_worlds_total = n_worlds_total + self.camera_res = camera_res + self.buffer_frames = buffer_frames + self.curr_time = 0.0 + self.sim_timestep = sim_timestep + + # Figure out which worlds should be rendered + if worlds is None: + if n_worlds_total is None: + raise ValueError( + "If 'worlds' is not specified, all worlds are rendered. " + "In that case, 'n_worlds_total' must be specified." + ) + worlds = list(range(n_worlds_total)) + self.world_ids = worlds + if len(self.world_ids) == 0: + raise ValueError("At least one valid world must be specified.") + + # Figure out which cameras should be rendered + if not isinstance(cameras, list): + cameras = [cameras] + self.enabled_cam_names = [] + for cam in cameras if isinstance(cameras, list) else [cameras]: + _, cam_name = self._resolve_camera_id_and_name(cam) + self.enabled_cam_names.append(cam_name) + if len(self.enabled_cam_names) == 0: + raise ValueError("At least one valid camera must be specified.") + + # Set up parameters for automatically determining when to render frames + self.playback_speed = playback_speed + self.output_fps = output_fps + self._secs_between_renders = 1 / (output_fps / playback_speed) + self._last_render_time_sec = -np.inf + + # Buffer to store rendered images + n_worlds_to_render = len(self.world_ids) + n_cams_to_render = len(self.enabled_cam_names) + self._buf_dim_per_frame = (n_worlds_to_render, n_cams_to_render, *camera_res) + if buffer_frames: + self._frames: list[np.ndarray | wp.array] = [] + else: + self._frames = None + + self._render_setup_impl(**kwargs) + + @override + def render_as_needed(self, mjw_data: mjw.Data) -> bool: + if self.curr_time >= self._last_render_time_sec + self._secs_between_renders: + self._last_render_time_sec = self.curr_time + rendered_images = self._render_impl(mjw_data) + if self.buffer_frames: + self._frames.append(rendered_images) + rendered = True + else: + rendered = False + self.curr_time += self.sim_timestep + return rendered + + @override + def reset(self): + self._last_render_time_sec = -np.inf + if self.buffer_frames: + self._frames = [] + + @override + def show_in_notebook( + self, + world_id: int, + camera: str | mj.MjsCamera | list[str | mj.MjsCamera] | None = None, + scale: float | None = None, + **kwargs, + ): + """Display recorded frames in a Jupyter notebook. + + Args: + world_id: Which parallel world to display frames for + camera: Camera(s) to display. If None, displays all enabled cameras. + scale: Optional factor by which to rescale frames before display. + **kwargs: Additional arguments passed to mediapy.show_video + """ + camera_names = self._normalize_camera_spec(camera) + + for cam_name in camera_names: + cam_id, _ = self._resolve_camera_id_and_name(cam_name) + if isinstance(world_id, int): + frames = self._fetch_frames_to_cpu_oneworld(world_id, cam_id, scale) + else: + frames = self._fetch_frames_to_cpu_multipleworlds( + world_id, cam_id, scale + ) + title = f"world {world_id}, camera {cam_name}" + mediapy.show_video(frames, fps=self.output_fps, title=title, **kwargs) + + @override + def save_video( + self, + world_id: int | list[int], + output_path: dict[str | mj.MjsCamera, PathLike] | PathLike, + scale: float | None = None, + **kwargs, + ) -> None: + """Save recorded frames as video files. + + Args: + world_id: Which parallel world to save frames for + output_path: Either a dict mapping camera specs to file paths, or: + - If single camera: a file path to save to + - If multiple cameras: a directory path to save all videos to + scale: Optional factor by which to rescale frames before saving. + **kwargs: Additional arguments passed to imageio.imwrite + """ + path_by_camera = self._resolve_output_paths(output_path) + + for cam_name, path in path_by_camera.items(): + cam_id, _ = self._resolve_camera_id_and_name(cam_name) + if isinstance(world_id, int): + frames = self._fetch_frames_to_cpu_oneworld(world_id, cam_id, scale) + else: + frames = self._fetch_frames_to_cpu_multipleworlds( + world_id, cam_id, scale + ) + path.parent.mkdir(parents=True, exist_ok=True) + write_video_from_frames( + path, frames, fps=self.output_fps, codec="libx264", **kwargs + ) + + def _fetch_frames_to_cpu_oneworld( + self, world_id: int, cam_id: int, scale: float | None = None + ) -> list[np.ndarray]: + if not self.buffer_frames: + raise RuntimeError( + "Frame buffering was disabled for this renderer, so recorded frames " + "are not available for saving or display." + ) + + if len(self._frames) == 0: + raise RuntimeError("No frames have been recorded yet.") + + if world_id not in self.world_ids: + raise ValueError( + f"world_id {world_id} was not among the rendered worlds: " + f"{self.world_ids}" + ) + world_id_among_rendered = self.world_ids.index(world_id) + + cam_name = self._cameras_id2name.get(cam_id, None) + if cam_name is None: + raise ValueError(f"Camera ID '{cam_id}' not found.") + if cam_name not in self.enabled_cam_names: + raise ValueError( + f"Camera '{cam_name}' (ID {cam_id}) was not among the rendered cameras." + ) + cam_id_among_rendered = self.enabled_cam_names.index(cam_name) + + frames = self._fetch_frames_to_cpu_impl( + world_id_among_rendered, cam_id_among_rendered + ) + if scale is not None: + render_res = tuple(int(x * scale) for x in self.camera_res) + for i, frame in enumerate(frames): + pil_frame = Image.fromarray(frame) + pil_frame_resized = pil_frame.resize( + render_res[::-1], # PIL expects (W, H); we use (H, W) + resample=Image.Resampling.LANCZOS, + ) + frame_resized = np.array(pil_frame_resized) + frames[i] = frame_resized + return frames + + def _fetch_frames_to_cpu_multipleworlds( + self, world_ids: list[int], cam_id: int, scale: float | None + ) -> dict[int, list[np.ndarray]]: + # Set up canvas for displaying frames from multiple worlds in a grid + n_worlds = len(world_ids) + n_rows = int(np.ceil(np.sqrt(n_worlds))) + n_cols = int(np.ceil(n_worlds / n_rows)) + # If scale is unspecified, make the output resolution roughly matches the + # resolution of a single world/camera - this avoids creating an excessively + # large canvas when there are many worlds. + if scale is None: + scale = 1 / n_cols + render_res = tuple(int(x * scale) for x in self.camera_res) + canvas_shape = (render_res[0] * n_rows, render_res[1] * n_cols) + n_frames = len(self._frames) + merged_frames = [ + np.zeros((*canvas_shape, 3), dtype=np.uint8) for _ in range(n_frames) + ] + + # Set up font for overlaying world IDs + FONT_FAMILY = "Arial" + FONT_SIZE_RATIO_OF_HEIGHT = 0.07 + MIN_FONT_SIZE = 7 + font_path = find_font_path(FONT_FAMILY) + font_size = max(MIN_FONT_SIZE, int(render_res[0] * FONT_SIZE_RATIO_OF_HEIGHT)) + font = ImageFont.truetype(font_path, font_size) + + # Fetch frames for each world and paste them onto the canvas + for i, wid in enumerate(world_ids): + row = i // n_cols + col = i % n_cols + row_slice = slice(row * render_res[0], (row + 1) * render_res[0]) + col_slice = slice(col * render_res[1], (col + 1) * render_res[1]) + + world_frames = self._fetch_frames_to_cpu_oneworld(wid, cam_id, scale) + assert len(world_frames) == n_frames, "inconsistent frame counts" + for j, world_frame in enumerate(world_frames): + # Overlay world ID text + pil_frame = Image.fromarray(world_frame) + draw = ImageDraw.Draw(pil_frame) + text = f"World {wid}" + text_pos = (0.03 * render_res[1], 0.02 * render_res[0]) # (x, y) + draw.text(text_pos, text, font=font, fill=(255, 255, 255)) + world_frame = np.array(pil_frame) + # Paste onto canvas + merged_frames[j][row_slice, col_slice] = world_frame + + return merged_frames + + @abstractmethod + def _render_setup_impl(self, **kwargs: Any) -> None: + pass + + @abstractmethod + def _render_impl(self, mjw_data: mjw.Data) -> np.ndarray | wp.array: + pass + + @abstractmethod + def _fetch_frames_to_cpu_impl( + self, world_id_among_rendered: int, cam_id_among_rendered: int + ) -> list[np.ndarray]: + pass diff --git a/src/flygym/warp/rendering/live_rendering.py b/src/flygym/warp/rendering/live_rendering.py index 58e3e578..4aaa7ec9 100644 --- a/src/flygym/warp/rendering/live_rendering.py +++ b/src/flygym/warp/rendering/live_rendering.py @@ -6,22 +6,15 @@ import warnings from typing import Any, override -from os import PathLike -from abc import ABC, abstractmethod -import mediapy import mujoco as mj import mujoco_warp as mjw import warp as wp import numpy as np -from PIL import Image, ImageDraw, ImageFont from flygym.compose import BaseWorld -from flygym.rendering.live_rendering import Renderer +from flygym.warp.rendering.base import _BaseWarpRenderer from flygym.warp.utils import get_rgb_selected_worlds_and_cameras -from flygym.utils.video import write_video_from_frames -from flygym.utils.plot import find_font_path - __all__ = [ "WarpGPUBatchRenderer", @@ -30,266 +23,6 @@ ] -class _BaseWarpRenderer(Renderer, ABC): - @override - def __init__( - self, - mj_model: mj.MjModel, - cameras: str | mj.MjsCamera | list[str | mj.MjsCamera], - n_worlds_total: int | None = None, - *, - worlds: list[int] | None = None, - camera_res: tuple[int, int] = (240, 320), - playback_speed: float = 0.2, - output_fps: int = 25, - buffer_frames: bool = True, - scene_option: mj.MjvOption | None = None, - **kwargs: Any, - ): - # Common setup using MjRenderer - super().__init__( - mj_model, - cameras, - camera_res=camera_res, - playback_speed=playback_speed, - output_fps=output_fps, - buffer_frames=buffer_frames, - scene_option=scene_option, - **kwargs, - ) - - # Warp-specific setup - self.mjw_model = mjw.put_model(mj_model) - self._n_worlds_total = n_worlds_total - self.camera_res = camera_res - self.buffer_frames = buffer_frames - - # Figure out which worlds should be rendered - if worlds is None: - if n_worlds_total is None: - raise ValueError( - "If 'worlds' is not specified, all worlds are rendered. " - "In that case, 'n_worlds_total' must be specified." - ) - worlds = list(range(n_worlds_total)) - self.world_ids = worlds - if len(self.world_ids) == 0: - raise ValueError("At least one valid world must be specified.") - - # Figure out which cameras should be rendered - if not isinstance(cameras, list): - cameras = [cameras] - self.enabled_cam_names = [] - for cam in cameras if isinstance(cameras, list) else [cameras]: - _, cam_name = self._resolve_camera_id_and_name(cam) - self.enabled_cam_names.append(cam_name) - if len(self.enabled_cam_names) == 0: - raise ValueError("At least one valid camera must be specified.") - - # Set up parameters for automatically determining when to render frames - self.playback_speed = playback_speed - self.output_fps = output_fps - self._secs_between_renders = 1 / (output_fps / playback_speed) - self._last_render_time_sec = -np.inf - - # Buffer to store rendered images - n_worlds_to_render = len(self.world_ids) - n_cams_to_render = len(self.enabled_cam_names) - self._buf_dim_per_frame = (n_worlds_to_render, n_cams_to_render, *camera_res) - if buffer_frames: - self._frames: list[np.ndarray | wp.array] = [] - else: - self._frames = None - - self._render_setup_impl(**kwargs) - - @override - def render_as_needed(self, mjw_data: mjw.Data) -> bool: - curr_time = mjw_data.time.numpy()[0] # assume all worlds have the same time - if curr_time >= self._last_render_time_sec + self._secs_between_renders: - self._last_render_time_sec = curr_time - rendered_images = self._render_impl(mjw_data) - if self.buffer_frames: - self._frames.append(rendered_images) - return True - else: - return False - - @override - def reset(self): - self._last_render_time_sec = -np.inf - if self.buffer_frames: - self._frames = [] - - @override - def show_in_notebook( - self, - world_id: int, - camera: str | mj.MjsCamera | list[str | mj.MjsCamera] | None = None, - scale: float | None = None, - **kwargs, - ): - """Display recorded frames in a Jupyter notebook. - - Args: - world_id: Which parallel world to display frames for - camera: Camera(s) to display. If None, displays all enabled cameras. - scale: Optional factor by which to rescale frames before display. - **kwargs: Additional arguments passed to mediapy.show_video - """ - camera_names = self._normalize_camera_spec(camera) - - for cam_name in camera_names: - cam_id, _ = self._resolve_camera_id_and_name(cam_name) - if isinstance(world_id, int): - frames = self._fetch_frames_to_cpu_oneworld(world_id, cam_id, scale) - else: - frames = self._fetch_frames_to_cpu_multipleworlds( - world_id, cam_id, scale - ) - title = f"world {world_id}, camera {cam_name}" - mediapy.show_video(frames, fps=self.output_fps, title=title, **kwargs) - - @override - def save_video( - self, - world_id: int | list[int], - output_path: dict[str | mj.MjsCamera, PathLike] | PathLike, - scale: float | None = None, - **kwargs, - ) -> None: - """Save recorded frames as video files. - - Args: - world_id: Which parallel world to save frames for - output_path: Either a dict mapping camera specs to file paths, or: - - If single camera: a file path to save to - - If multiple cameras: a directory path to save all videos to - scale: Optional factor by which to rescale frames before saving. - **kwargs: Additional arguments passed to imageio.imwrite - """ - path_by_camera = self._resolve_output_paths(output_path) - - for cam_name, path in path_by_camera.items(): - cam_id, _ = self._resolve_camera_id_and_name(cam_name) - if isinstance(world_id, int): - frames = self._fetch_frames_to_cpu_oneworld(world_id, cam_id, scale) - else: - frames = self._fetch_frames_to_cpu_multipleworlds( - world_id, cam_id, scale - ) - path.parent.mkdir(parents=True, exist_ok=True) - write_video_from_frames( - path, frames, fps=self.output_fps, codec="libx264", **kwargs - ) - - def _fetch_frames_to_cpu_oneworld( - self, world_id: int, cam_id: int, scale: float | None = None - ) -> list[np.ndarray]: - if not self.buffer_frames: - raise RuntimeError( - "Frame buffering was disabled for this renderer, so recorded frames " - "are not available for saving or display." - ) - - if len(self._frames) == 0: - raise RuntimeError("No frames have been recorded yet.") - - if world_id not in self.world_ids: - raise ValueError( - f"world_id {world_id} was not among the rendered worlds: " - f"{self.world_ids}" - ) - world_id_among_rendered = self.world_ids.index(world_id) - - cam_name = self._cameras_id2name.get(cam_id, None) - if cam_name is None: - raise ValueError(f"Camera ID '{cam_id}' not found.") - if cam_name not in self.enabled_cam_names: - raise ValueError( - f"Camera '{cam_name}' (ID {cam_id}) was not among the rendered cameras." - ) - cam_id_among_rendered = self.enabled_cam_names.index(cam_name) - - frames = self._fetch_frames_to_cpu_impl( - world_id_among_rendered, cam_id_among_rendered - ) - if scale is not None: - render_res = tuple(int(x * scale) for x in self.camera_res) - for i, frame in enumerate(frames): - pil_frame = Image.fromarray(frame) - pil_frame_resized = pil_frame.resize( - render_res[::-1], # PIL expects (W, H); we use (H, W) - resample=Image.Resampling.LANCZOS, - ) - frame_resized = np.array(pil_frame_resized) - frames[i] = frame_resized - return frames - - def _fetch_frames_to_cpu_multipleworlds( - self, world_ids: list[int], cam_id: int, scale: float | None - ) -> dict[int, list[np.ndarray]]: - # Set up canvas for displaying frames from multiple worlds in a grid - n_worlds = len(world_ids) - n_rows = int(np.ceil(np.sqrt(n_worlds))) - n_cols = int(np.ceil(n_worlds / n_rows)) - # If scale is unspecified, make the output resolution roughly matches the - # resolution of a single world/camera - this avoids creating an excessively - # large canvas when there are many worlds. - if scale is None: - scale = 1 / n_cols - render_res = tuple(int(x * scale) for x in self.camera_res) - canvas_shape = (render_res[0] * n_rows, render_res[1] * n_cols) - n_frames = len(self._frames) - merged_frames = [ - np.zeros((*canvas_shape, 3), dtype=np.uint8) for _ in range(n_frames) - ] - - # Set up font for overlaying world IDs - FONT_FAMILY = "Arial" - FONT_SIZE_RATIO_OF_HEIGHT = 0.07 - MIN_FONT_SIZE = 7 - font_path = find_font_path(FONT_FAMILY) - font_size = max(MIN_FONT_SIZE, int(render_res[0] * FONT_SIZE_RATIO_OF_HEIGHT)) - font = ImageFont.truetype(font_path, font_size) - - # Fetch frames for each world and paste them onto the canvas - for i, wid in enumerate(world_ids): - row = i // n_cols - col = i % n_cols - row_slice = slice(row * render_res[0], (row + 1) * render_res[0]) - col_slice = slice(col * render_res[1], (col + 1) * render_res[1]) - - world_frames = self._fetch_frames_to_cpu_oneworld(wid, cam_id, scale) - assert len(world_frames) == n_frames, "inconsistent frame counts" - for j, world_frame in enumerate(world_frames): - # Overlay world ID text - pil_frame = Image.fromarray(world_frame) - draw = ImageDraw.Draw(pil_frame) - text = f"World {wid}" - text_pos = (0.03 * render_res[1], 0.02 * render_res[0]) # (x, y) - draw.text(text_pos, text, font=font, fill=(255, 255, 255)) - world_frame = np.array(pil_frame) - # Paste onto canvas - merged_frames[j][row_slice, col_slice] = world_frame - - return merged_frames - - @abstractmethod - def _render_setup_impl(self, **kwargs: Any) -> None: - pass - - @abstractmethod - def _render_impl(self, mjw_data: mjw.Data) -> np.ndarray | wp.array: - pass - - @abstractmethod - def _fetch_frames_to_cpu_impl( - self, world_id_among_rendered: int, cam_id_among_rendered: int - ) -> list[np.ndarray]: - pass - - class WarpGPUBatchRenderer(_BaseWarpRenderer): """GPU-side renderer using MJWarp's GPU batch rendering functionality.""" diff --git a/src/flygym/warp/rendering/recorded_trajectory.py b/src/flygym/warp/rendering/recorded_trajectory.py index 2433f2ca..5b7809d9 100644 --- a/src/flygym/warp/rendering/recorded_trajectory.py +++ b/src/flygym/warp/rendering/recorded_trajectory.py @@ -45,12 +45,16 @@ def _render_setup_impl(self, **kwargs: Any) -> None: self.scene_option = None def _render_impl(self, mjw_data: mjw.Data) -> tuple: - # One host transfer per recorded frame: (n_worlds, nq) is tiny next to the - # (n_worlds, n_cams, H, W, 3) RGB tensor the batch renderer would buffer. - qpos = mjw_data.qpos.numpy()[self.world_ids].copy() + # Device-to-device clones only -- no `.numpy()`, so this issues no host + # transfer and forces no CUDA sync. That keeps the recorder graph-capturable + # (and lets it run in an async/captured step loop without stalling it). World + # selection and the single host transfer are deferred to + # `recorded_trajectories`. Cloning the full (n_worlds, nq) state is cheap next + # to the (n_worlds, n_cams, H, W, 3) RGB tensor the batch renderer would buffer. + qpos = wp.clone(mjw_data.qpos) if self._nmocap > 0: - mocap_pos = mjw_data.mocap_pos.numpy()[self.world_ids].copy() - mocap_quat = mjw_data.mocap_quat.numpy()[self.world_ids].copy() + mocap_pos = wp.clone(mjw_data.mocap_pos) + mocap_quat = wp.clone(mjw_data.mocap_quat) else: mocap_pos = mocap_quat = None return (qpos, mocap_pos, mocap_quat) @@ -66,25 +70,30 @@ def recorded_trajectories(self) -> list[RecordedTrajectory]: if len(self._frames) == 0: raise RuntimeError("No frames have been recorded yet.") - # self._frames is a list (over time) of (qpos, mocap_pos, mocap_quat) tuples, - # each batched over the recorded worlds along axis 0. - qpos_all = np.stack([f[0] for f in self._frames], axis=0) # (T, n_worlds, nq) + # self._frames is a list (over time) of (qpos, mocap_pos, mocap_quat) tuples + # of warp arrays, each batched over *all* worlds along axis 0. The single host + # transfer (`.numpy()`) and the per-world selection happen here, not per frame. + qpos_all = np.stack( + [f[0].numpy() for f in self._frames], axis=0 + ) # (T, n_worlds_total, nq) if self._nmocap > 0: - mocap_pos_all = np.stack([f[1] for f in self._frames], axis=0) - mocap_quat_all = np.stack([f[2] for f in self._frames], axis=0) + mocap_pos_all = np.stack([f[1].numpy() for f in self._frames], axis=0) + mocap_quat_all = np.stack([f[2].numpy() for f in self._frames], axis=0) trajectories = [] - for w, world_id in enumerate(self.world_ids): + for world_id in self.world_ids: trajectories.append( RecordedTrajectory( - qpos=qpos_all[:, w, :], + qpos=qpos_all[:, world_id, :], output_fps=self.output_fps, playback_speed=self.playback_speed, camera_names=list(self.enabled_cam_names), camera_res=self.camera_res, world_id=world_id, - mocap_pos=mocap_pos_all[:, w] if self._nmocap > 0 else None, - mocap_quat=mocap_quat_all[:, w] if self._nmocap > 0 else None, + mocap_pos=mocap_pos_all[:, world_id] if self._nmocap > 0 else None, + mocap_quat=( + mocap_quat_all[:, world_id] if self._nmocap > 0 else None + ), ) ) return trajectories diff --git a/src/flygym/warp/simulation.py b/src/flygym/warp/simulation.py index 84cddd5e..1719115e 100644 --- a/src/flygym/warp/simulation.py +++ b/src/flygym/warp/simulation.py @@ -12,6 +12,7 @@ from flygym.simulation import Simulation from flygym.utils.profiling import print_perf_report_parallel from flygym.warp.rendering import ( + RendererType, WarpGPUBatchRenderer, WarpCPURenderer, WarpTrajectoryRecorder, @@ -289,6 +290,7 @@ def step(self) -> None: def set_renderer( self, cameras: str | mj.MjsCamera | list[str | mj.MjsCamera], + renderer_type: RendererType = RendererType.GPU_BATCH, *, camera_res: tuple[int, int] = (240, 320), playback_speed: float = 0.2, @@ -296,30 +298,19 @@ def set_renderer( buffer_frames: bool = True, scene_option: mj.MjvOption | None = None, worlds: list[int] | None = None, - use_gpu_batch_rendering: bool = False, - record_trajectory_only: bool = False, **kwargs: Any, ) -> WarpGPUBatchRenderer | WarpCPURenderer | WarpTrajectoryRecorder: """Attach a renderer to this GPU simulation. Args: cameras: Camera(s) to render. - camera_res: ``(height, width)`` in pixels. + renderer_type: Renderer type. Defaults to `RendererType.GPU_BATCH`. + camera_res: `(height, width)` in pixels. playback_speed: Video playback speed relative to real time. output_fps: Output video frame rate. buffer_frames: If True, store rendered frames in memory. scene_option: MuJoCo scene options. Uses defaults if None. worlds: Indices of worlds to render. Defaults to all worlds. - use_gpu_batch_rendering: If True, use `WarpGPUBatchRenderer`; - otherwise use `WarpCPURenderer`. Ignored when - ``record_trajectory_only`` is True. - record_trajectory_only: If True, attach a `WarpTrajectoryRecorder` instead - of a renderer: it records ``qpos`` (and mocap poses) for the selected - worlds at the render cadence instead of rasterizing frames. Read the - result from ``self.renderer.recorded_trajectories`` (one per recorded - world) and replay it later with - `flygym.warp.rendering.render_trajectories_gpu` or - `flygym.rendering.render_trajectories`. **kwargs: Passed to the renderer (ignored when recording only). Returns: @@ -328,26 +319,11 @@ def set_renderer( if worlds is None: worlds = list(range(self.n_worlds)) - if record_trajectory_only: - self.renderer = WarpTrajectoryRecorder( - self.mj_model, - cameras, - n_worlds_total=self.n_worlds, - worlds=worlds, - camera_res=camera_res, - playback_speed=playback_speed, - output_fps=output_fps, - buffer_frames=True, - **kwargs, - ) - return self.renderer - - self.use_gpu_batch_rendering = use_gpu_batch_rendering - renderer_kwargs = { "mj_model": self.mj_model, "n_worlds_total": self.n_worlds, "cameras": cameras, + "sim_timestep": self.timestep, "camera_res": camera_res, "playback_speed": playback_speed, "output_fps": output_fps, @@ -356,7 +332,10 @@ def set_renderer( "worlds": worlds, **kwargs, } - if use_gpu_batch_rendering: + self.renderer_type = renderer_type + if renderer_type == RendererType.CPU: + self.renderer = WarpCPURenderer(**renderer_kwargs) + elif renderer_type == RendererType.GPU_BATCH: is_model_modified = modify_world_for_batch_rendering(self.world) if is_model_modified: warnings.warn( @@ -373,8 +352,10 @@ def set_renderer( self.mjw_model, self.mjw_data = self._mj_structs_to_mjw_structs() renderer_kwargs["mj_model"] = self.mj_model self.renderer = WarpGPUBatchRenderer(**renderer_kwargs) + elif renderer_type == RendererType.RECORDED_TRAJECTORY: + self.renderer = WarpTrajectoryRecorder(**renderer_kwargs) else: - self.renderer = WarpCPURenderer(**renderer_kwargs) + raise ValueError(f"Unsupported renderer type: {renderer_type}") return self.renderer diff --git a/tests/warp/test_rendering.py b/tests/warp/test_rendering.py index f92c0d71..0ad17e3f 100644 --- a/tests/warp/test_rendering.py +++ b/tests/warp/test_rendering.py @@ -13,7 +13,7 @@ from flygym.anatomy import Skeleton, JointPreset, AxisOrder from flygym.compose import NeuroMechFly, FlatGroundWorld, KinematicPosePreset from flygym.utils.math import Rotation3D -from flygym.warp import WarpCPURenderer +from flygym.warp import WarpCPURenderer, RendererType from flygym.warp.rendering import modify_world_for_batch_rendering @@ -36,7 +36,7 @@ def render_bundle(gpu_sim_factory): playback_speed=0.001, # tiny interval: 1/(10/0.001) = 0.0001 s ≈ 1 step output_fps=10, worlds=[0, 1], - use_gpu_batch_rendering=False, + renderer_type=RendererType.CPU, buffer_frames=True, ) yield sim, fly, cam, renderer @@ -79,7 +79,7 @@ def test_empty_worlds_raises(self, gpu_sim_factory): cam, camera_res=(64, 64), worlds=[], - use_gpu_batch_rendering=False, + renderer_type=RendererType.CPU, ) @@ -108,12 +108,20 @@ def test_returns_false_too_soon(self, render_bundle): sim, fly, cam, renderer = render_bundle sim.reset() renderer.reset() - _advance_past_render_interval(sim, renderer) - renderer.render_as_needed(sim.mjw_data) # first render - - # Immediately call again — not enough time has elapsed - did_render = renderer.render_as_needed(sim.mjw_data) - assert did_render is False + # render_as_needed advances its own clock by one sim timestep per call (it no + # longer reads sim.time, to stay graph-capturable). The "too soon" guard is + # only meaningful when the render interval spans more than one step, so widen + # it for this test (the fixture's interval is ≈ one step). + saved_interval = renderer._secs_between_renders + renderer._secs_between_renders = 100 * sim.mj_model.opt.timestep + try: + assert renderer.render_as_needed(sim.mjw_data) is True # first render fires + + # Call again one timestep later — not enough time has elapsed. + did_render = renderer.render_as_needed(sim.mjw_data) + assert did_render is False + finally: + renderer._secs_between_renders = saved_interval def test_frame_buffered_after_render(self, render_bundle): sim, fly, cam, renderer = render_bundle @@ -220,7 +228,7 @@ def test_only_specified_worlds_in_world_ids(self, gpu_sim_factory): cam, camera_res=(64, 64), worlds=[1], - use_gpu_batch_rendering=False, + renderer_type=RendererType.CPU, ) assert renderer.world_ids == [1] @@ -233,7 +241,7 @@ def test_default_worlds_is_all_worlds(self, gpu_sim_factory): cam, camera_res=(64, 64), worlds=None, - use_gpu_batch_rendering=False, + renderer_type=RendererType.CPU, ) assert renderer.world_ids == [0, 1, 2] diff --git a/tests/warp/test_simulation.py b/tests/warp/test_simulation.py index 822e86ae..360e1c98 100644 --- a/tests/warp/test_simulation.py +++ b/tests/warp/test_simulation.py @@ -17,7 +17,7 @@ KinematicPosePreset, ) from flygym.utils.math import Rotation3D -from flygym.warp import GPUSimulation +from flygym.warp import GPUSimulation, RendererType from flygym.warp.rendering import WarpCPURenderer @@ -340,7 +340,7 @@ def test_set_renderer_returns_warp_cpu_renderer(self, gpu_bundle): cam, camera_res=(64, 64), worlds=[0, 1], - use_gpu_batch_rendering=False, + renderer_type=RendererType.CPU, ) assert isinstance(renderer, WarpCPURenderer) @@ -352,7 +352,7 @@ def test_set_renderer_world_ids(self, gpu_bundle): cam, camera_res=(64, 64), worlds=[0, 2], - use_gpu_batch_rendering=False, + renderer_type=RendererType.CPU, ) assert renderer.world_ids == [0, 2] @@ -369,7 +369,7 @@ def test_print_performance_report(self, gpu_bundle, capsys): cam, camera_res=(64, 64), worlds=[0], - use_gpu_batch_rendering=False, + renderer_type=RendererType.CPU, ) sim.print_performance_report() captured = capsys.readouterr() diff --git a/tests/warp/test_trajectory.py b/tests/warp/test_trajectory.py index 06e51cef..a2847d6b 100644 --- a/tests/warp/test_trajectory.py +++ b/tests/warp/test_trajectory.py @@ -16,6 +16,7 @@ render_trajectories, ) from flygym.warp import ( + RendererType, WarpTrajectoryRecorder, render_trajectories_gpu, modify_world_for_batch_rendering, @@ -36,7 +37,7 @@ def recorded_gpu(gpu_sim_factory): camera_res=(64, 64), output_fps=100, worlds=[0, 2], - record_trajectory_only=True, + renderer_type=RendererType.RECORDED_TRAJECTORY, ) sim.reset() for _ in range(200):