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/scripts/record_replay_trajectories_gpu.py b/scripts/record_replay_trajectories_gpu.py new file mode 100644 index 00000000..99c97cab --- /dev/null +++ b/scripts/record_replay_trajectories_gpu.py @@ -0,0 +1,204 @@ +"""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*, 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 +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 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, 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. + +Configure the run by editing the constants below, then:: + + uv run python scripts/record_replay_trajectories_gpu.py +""" + +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 RecordedTrajectory, render_trajectories +from flygym.compose import ActuatorType +from flygym_demo.benchmark import ( + make_model, + ReplayTargetData, + update_target_angles_kernel, + 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 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). + """ + 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 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, + 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 (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 every world 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: + check_gpu() + + 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() + + # --- 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 " + f"{model_dir}." + ) + + # --- Replay post-hoc, reloading the trajectories from disk and sub-selecting --- + 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.") + + if CPU_REPLAY: + # CPU replay needs no special model prep; reuse the unmodified compiled model. + 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 = OUTPUT_DIR / "replay_gpu" + print( + f"Rendering {len(trajectories)} worlds on GPU to {gpu_out} " + 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=WORLDS_PER_BATCH + ) + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/src/flygym/__init__.py b/src/flygym/__init__.py index 9676d5cc..d3cf8bc2 100644 --- a/src/flygym/__init__.py +++ b/src/flygym/__init__.py @@ -9,7 +9,14 @@ 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, + render_trajectories, + launch_interactive_viewer, + preview_model, +) __all__ = [ "assets_dir", @@ -18,6 +25,9 @@ "flybody", "Simulation", "Renderer", + "TrajectoryRecorder", + "RecordedTrajectory", + "render_trajectories", "launch_interactive_viewer", "preview_model", ] 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 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/rendering/__init__.py b/src/flygym/rendering/__init__.py new file mode 100644 index 00000000..06f401b1 --- /dev/null +++ b/src/flygym/rendering/__init__.py @@ -0,0 +1,31 @@ +"""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, + render_trajectories, +) + +__all__ = [ + "Renderer", + "TrajectoryRecorder", + "RecordedTrajectory", + "render_trajectories", + "launch_interactive_viewer", + "preview_model", +] diff --git a/src/flygym/rendering.py b/src/flygym/rendering/live_rendering.py similarity index 94% rename from src/flygym/rendering.py rename to src/flygym/rendering/live_rendering.py index d8c709e3..6cc44c7c 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 @@ -83,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]: @@ -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..ca682c37 --- /dev/null +++ b/src/flygym/rendering/recorded_trajectory.py @@ -0,0 +1,405 @@ +"""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", + "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 `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) + 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 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 + 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 " + "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 " + "RecordedTrajectory.save and replay with render_trajectories." + ) + + +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 (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 + ``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..069eddef 100644 --- a/src/flygym/simulation.py +++ b/src/flygym/simulation.py @@ -7,9 +7,9 @@ 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 +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: @@ -760,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 diff --git a/src/flygym/warp/__init__.py b/src/flygym/warp/__init__.py index 427a6d78..0b848720 100644 --- a/src/flygym/warp/__init__.py +++ b/src/flygym/warp/__init__.py @@ -1,4 +1,19 @@ from .simulation import GPUSimulation -from .rendering import WarpGPUBatchRenderer, WarpCPURenderer +from .rendering import ( + RendererType, + WarpGPUBatchRenderer, + WarpCPURenderer, + WarpTrajectoryRecorder, + modify_world_for_batch_rendering, + render_trajectories_gpu, +) -__all__ = ["GPUSimulation", "WarpGPUBatchRenderer", "WarpCPURenderer"] +__all__ = [ + "GPUSimulation", + "RendererType", + "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..8e368b15 --- /dev/null +++ b/src/flygym/warp/rendering/__init__.py @@ -0,0 +1,26 @@ +"""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.base import RendererType +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__ = [ + "RendererType", + "WarpGPUBatchRenderer", + "WarpCPURenderer", + "WarpTrajectoryRecorder", + "modify_world_for_batch_rendering", + "render_trajectories_gpu", +] diff --git a/src/flygym/warp/rendering.py b/src/flygym/warp/rendering/base.py similarity index 59% rename from src/flygym/warp/rendering.py rename to src/flygym/warp/rendering/base.py index be88be53..f9350fa2 100644 --- a/src/flygym/warp/rendering.py +++ b/src/flygym/warp/rendering/base.py @@ -1,7 +1,7 @@ -import warnings from typing import Any, override from os import PathLike from abc import ABC, abstractmethod +from enum import Enum import mediapy import mujoco as mj @@ -10,12 +10,18 @@ import numpy as np from PIL import Image, ImageDraw, ImageFont -from flygym.compose import BaseWorld -from flygym.rendering import Renderer -from flygym.warp.utils import get_rgb_selected_worlds_and_cameras +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 @@ -23,6 +29,7 @@ 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, @@ -50,6 +57,8 @@ def __init__( 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: @@ -92,15 +101,16 @@ def __init__( @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 + 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) - return True + rendered = True else: - return False + rendered = False + self.curr_time += self.sim_timestep + return rendered @override def reset(self): @@ -275,169 +285,3 @@ 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.""" - - def _render_setup_impl(self, **kwargs: Any) -> None: - if not self._is_scene_option_default(self.scene_option): - raise RuntimeError( - "Custom scene options are not supported with WarpGPUBatchRenderer " - "because it is not implemented in MJWarp batch rendering." - ) - - self._world_ids_gpu = wp.array(self.world_ids, dtype=wp.int32) - self._enabled_cam_ids_gpu = wp.array( - [self._cameras_names2id[n] for n in self.enabled_cam_names], dtype=wp.int32 - ) - cam_mask = [ - self._cameras_id2name[cid] in self.enabled_cam_names - for cid in range(self.mj_model.ncam) - ] - - # Create batch rendering context - self._rendering_context = mjw.create_render_context( - mjm=self.mj_model, - nworld=self._n_worlds_total, - cam_active=cam_mask, - cam_res=self.camera_res[::-1], # MJWarp expects (W, H); we use (H, W) - **kwargs, - ) - - # Remove normal MjRenderer inherited from CPU Renderer - self.scene_option = None - self.mj_renderer = None - - def _render_impl(self, mjw_data: mjw.Data) -> np.ndarray | wp.array: - mjw.refit_bvh(self.mjw_model, mjw_data, self._rendering_context) - mjw.render(self.mjw_model, mjw_data, self._rendering_context) - rgb_out = wp.zeros(self._buf_dim_per_frame, dtype=wp.vec3f) - get_rgb_selected_worlds_and_cameras( - self._rendering_context, - self._world_ids_gpu, - self._enabled_cam_ids_gpu, - rgb_out, - ) - return rgb_out - - def _fetch_frames_to_cpu_impl( - self, world_id_among_rendered: int, cam_id_among_rendered: int - ) -> list[np.ndarray]: - frames = [] - for frame_buffer in self._frames: - frame = frame_buffer[world_id_among_rendered, cam_id_among_rendered, :, :] - frame = (frame * 255.0).numpy().astype(np.uint8) - frames.append(frame) - return frames - - @override - def close(self): - return # nothing to do since we are not using a mj.Renderer context - - @staticmethod - def _is_scene_option_default(scene_option: mj.MjvOption) -> bool: - default_option = mj.MjvOption() - mj.mjv_defaultOption(default_option) - return scene_option == default_option - - -class WarpCPURenderer(_BaseWarpRenderer): - """CPU-side renderer for multi-world MJWarp simulation.""" - - def _render_setup_impl(self, **kwargs: Any) -> None: - self._mj_data_buffer = mj.MjData(self.mj_model) - # Nothing else to do - just use mjRenderer inherited from CPU Renderer - - def _render_impl(self, mjw_data: mjw.Data) -> np.ndarray | wp.array: - rendered_images = np.zeros((*self._buf_dim_per_frame, 3), dtype=np.uint8) - - for world_id in self.world_ids: - wid_among_rendered = self.world_ids.index(world_id) - - # Copy data into CPU MjData struct - mj.mj_resetData(self.mj_model, self._mj_data_buffer) - mjw.get_data_into(self._mj_data_buffer, self.mj_model, mjw_data, world_id) - - # Render each enabled camera and store frames - for cam_name, internal_cam_id in self._cameras_names2id.items(): - cid_among_rendered = self.enabled_cam_names.index(cam_name) - - self.mj_renderer.update_scene( - self._mj_data_buffer, internal_cam_id, self.scene_option - ) - frame = self.mj_renderer.render() - - if self.buffer_frames: - rendered_images[wid_among_rendered, cid_among_rendered] = frame - - return rendered_images - - def _fetch_frames_to_cpu_impl( - self, world_id_among_rendered: int, cam_id_among_rendered: int - ) -> list[np.ndarray]: - frames = [] - for rendered_images in self._frames: - frame = rendered_images[world_id_among_rendered, cam_id_among_rendered, ...] - frames.append(frame) - return frames - - -def modify_world_for_batch_rendering(world: BaseWorld) -> bool: - """Modify world MJCF model to make it compatible with MJWarp's GPU batch rendering. - - This may reduce texture and lighting realism. - - Modification happens in place. Returns True if any modifications were made, False - otherwise. - - Note for developers: Check if anything here can be dropped upon new MJWarp releases. - """ - is_modified = False - - rgb_role = int(mj.mjtTextureRole.mjTEXROLE_RGB) - - # Strip textures from fly body materials - # (rendering textures on complex meshes causes MJWarp memory corruption) - for material in world.mjcf_root.materials: - # Don't touch things that are not part of a Fly - if material.name.split("/")[0] not in world.fly_lookup: - continue - # Make wings half transparent - if "wing" in material.name: - material.rgba[3] = 0.5 - # If material has a texture, remove it to reduce memory use - texture_name = material.textures[rgb_role] - if texture_name: - texture_element = world.mjcf_root.texture(texture_name) - primary_color_rgb = texture_element.rgb1 - material.textures[rgb_role] = "" - material.rgba[:3] = primary_color_rgb - is_modified = True - - # Adjust scale of checker materials (e.g., ground): texrepeat needs to be scaled - # down by 1000x to get the same pattern - unclear why. Only materials that still - # reference a texture (e.g. the ground checker) need this. - for material in world.mjcf_root.materials: - if material.textures[rgb_role]: - 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 - - return is_modified diff --git a/src/flygym/warp/rendering/live_rendering.py b/src/flygym/warp/rendering/live_rendering.py new file mode 100644 index 00000000..4aaa7ec9 --- /dev/null +++ b/src/flygym/warp/rendering/live_rendering.py @@ -0,0 +1,198 @@ +"""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 + +import mujoco as mj +import mujoco_warp as mjw +import warp as wp +import numpy as np + +from flygym.compose import BaseWorld +from flygym.warp.rendering.base import _BaseWarpRenderer +from flygym.warp.utils import get_rgb_selected_worlds_and_cameras + +__all__ = [ + "WarpGPUBatchRenderer", + "WarpCPURenderer", + "modify_world_for_batch_rendering", +] + + +class WarpGPUBatchRenderer(_BaseWarpRenderer): + """GPU-side renderer using MJWarp's GPU batch rendering functionality.""" + + def _render_setup_impl(self, **kwargs: Any) -> None: + if not self._is_scene_option_default(self.scene_option): + raise RuntimeError( + "Custom scene options are not supported with WarpGPUBatchRenderer " + "because it is not implemented in MJWarp batch rendering." + ) + + self._world_ids_gpu = wp.array(self.world_ids, dtype=wp.int32) + self._enabled_cam_ids_gpu = wp.array( + [self._cameras_names2id[n] for n in self.enabled_cam_names], dtype=wp.int32 + ) + cam_mask = [ + self._cameras_id2name[cid] in self.enabled_cam_names + for cid in range(self.mj_model.ncam) + ] + + # Create batch rendering context + self._rendering_context = mjw.create_render_context( + mjm=self.mj_model, + nworld=self._n_worlds_total, + cam_active=cam_mask, + cam_res=self.camera_res[::-1], # MJWarp expects (W, H); we use (H, W) + **kwargs, + ) + + # Remove normal MjRenderer inherited from CPU Renderer + self.scene_option = None + self.mj_renderer = None + + def _render_impl(self, mjw_data: mjw.Data) -> np.ndarray | wp.array: + mjw.refit_bvh(self.mjw_model, mjw_data, self._rendering_context) + mjw.render(self.mjw_model, mjw_data, self._rendering_context) + rgb_out = wp.zeros(self._buf_dim_per_frame, dtype=wp.vec3f) + get_rgb_selected_worlds_and_cameras( + self._rendering_context, + self._world_ids_gpu, + self._enabled_cam_ids_gpu, + rgb_out, + ) + return rgb_out + + def _fetch_frames_to_cpu_impl( + self, world_id_among_rendered: int, cam_id_among_rendered: int + ) -> list[np.ndarray]: + frames = [] + for frame_buffer in self._frames: + frame = frame_buffer[world_id_among_rendered, cam_id_among_rendered, :, :] + frame = (frame * 255.0).numpy().astype(np.uint8) + frames.append(frame) + return frames + + @override + def close(self): + return # nothing to do since we are not using a mj.Renderer context + + @staticmethod + def _is_scene_option_default(scene_option: mj.MjvOption) -> bool: + default_option = mj.MjvOption() + mj.mjv_defaultOption(default_option) + return scene_option == default_option + + +class WarpCPURenderer(_BaseWarpRenderer): + """CPU-side renderer for multi-world MJWarp simulation.""" + + def _render_setup_impl(self, **kwargs: Any) -> None: + self._mj_data_buffer = mj.MjData(self.mj_model) + # Nothing else to do - just use mjRenderer inherited from CPU Renderer + + def _render_impl(self, mjw_data: mjw.Data) -> np.ndarray | wp.array: + rendered_images = np.zeros((*self._buf_dim_per_frame, 3), dtype=np.uint8) + + for world_id in self.world_ids: + wid_among_rendered = self.world_ids.index(world_id) + + # Copy data into CPU MjData struct + mj.mj_resetData(self.mj_model, self._mj_data_buffer) + mjw.get_data_into(self._mj_data_buffer, self.mj_model, mjw_data, world_id) + + # Render each enabled camera and store frames + for cam_name, internal_cam_id in self._cameras_names2id.items(): + cid_among_rendered = self.enabled_cam_names.index(cam_name) + + self.mj_renderer.update_scene( + self._mj_data_buffer, internal_cam_id, self.scene_option + ) + frame = self.mj_renderer.render() + + if self.buffer_frames: + rendered_images[wid_among_rendered, cid_among_rendered] = frame + + return rendered_images + + def _fetch_frames_to_cpu_impl( + self, world_id_among_rendered: int, cam_id_among_rendered: int + ) -> list[np.ndarray]: + frames = [] + for rendered_images in self._frames: + frame = rendered_images[world_id_among_rendered, cam_id_among_rendered, ...] + frames.append(frame) + return frames + + +def modify_world_for_batch_rendering(world: BaseWorld) -> bool: + """Modify world MJCF model to make it compatible with MJWarp's GPU batch rendering. + + This may reduce texture and lighting realism. + + 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. + """ + is_modified = False + + rgb_role = int(mj.mjtTextureRole.mjTEXROLE_RGB) + + # Strip textures from fly body materials + # (rendering textures on complex meshes causes MJWarp memory corruption) + for material in world.mjcf_root.materials: + # Don't touch things that are not part of a Fly + if material.name.split("/")[0] not in world.fly_lookup: + continue + # Make wings half transparent + if "wing" in material.name: + material.rgba[3] = 0.5 + # If material has a texture, remove it to reduce memory use + texture_name = material.textures[rgb_role] + if texture_name: + texture_element = world.mjcf_root.texture(texture_name) + primary_color_rgb = texture_element.rgb1 + material.textures[rgb_role] = "" + material.rgba[:3] = primary_color_rgb + is_modified = True + + # Adjust scale of checker materials (e.g., ground): texrepeat needs to be scaled + # down by 1000x to get the same pattern - unclear why. Only materials that still + # reference a texture (e.g. the ground checker) need this. + for material in world.mjcf_root.materials: + if material.textures[rgb_role]: + material.texrepeat = tuple(tr / 1000 for tr in material.texrepeat) + 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 diff --git a/src/flygym/warp/rendering/recorded_trajectory.py b/src/flygym/warp/rendering/recorded_trajectory.py new file mode 100644 index 00000000..5b7809d9 --- /dev/null +++ b/src/flygym/warp/rendering/recorded_trajectory.py @@ -0,0 +1,296 @@ +"""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: + # 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 = 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) + + @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 + # 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].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 world_id in self.world_ids: + trajectories.append( + RecordedTrajectory( + 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[:, world_id] if self._nmocap > 0 else None, + mocap_quat=( + mocap_quat_all[:, world_id] 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 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 each with " + "RecordedTrajectory.save 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 (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 + 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..1719115e 100644 --- a/src/flygym/warp/simulation.py +++ b/src/flygym/warp/simulation.py @@ -12,8 +12,10 @@ from flygym.simulation import Simulation from flygym.utils.profiling import print_perf_report_parallel from flygym.warp.rendering import ( + RendererType, WarpGPUBatchRenderer, WarpCPURenderer, + WarpTrajectoryRecorder, modify_world_for_batch_rendering, ) from flygym.warp.utils import ( @@ -77,19 +79,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), @@ -99,19 +104,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), @@ -121,19 +129,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), @@ -143,19 +154,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), @@ -165,19 +179,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), @@ -192,22 +209,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), @@ -272,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, @@ -279,34 +298,32 @@ def set_renderer( buffer_frames: bool = True, scene_option: mj.MjvOption | None = None, worlds: list[int] | None = None, - use_gpu_batch_rendering: bool = False, **kwargs: Any, - ) -> WarpGPUBatchRenderer | WarpCPURenderer: + ) -> 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`. - **kwargs: Passed to the renderer. + **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)) - 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, @@ -315,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( @@ -332,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/core/test_trajectory.py b/tests/core/test_trajectory.py new file mode 100644 index 00000000..9b2a735f --- /dev/null +++ b/tests/core/test_trajectory.py @@ -0,0 +1,255 @@ +"""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, + 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( + 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_save_from_file_roundtrip(self, recorded, tmp_path): + traj, _, _ = recorded + 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 + 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_load_multiple(self, recorded, tmp_path): + traj, _, _ = recorded + _save_all([traj, traj], tmp_path) + assert (tmp_path / "traj_0001.npz").exists() + trajs = _load_all(tmp_path) + assert len(trajs) == 2 + + +# --------------------------------------------------------------------------- +# 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_saved_file_writes_video(self, recorded, tmp_path): + traj, sim, _ = recorded + 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 + + 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_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 new file mode 100644 index 00000000..a2847d6b --- /dev/null +++ b/tests/warp/test_trajectory.py @@ -0,0 +1,183 @@ +"""Tests for GPU trajectory recording (WarpTrajectoryRecorder) and GPU replay.""" + +import warnings + +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 ( + RecordedTrajectory, + render_trajectories, +) +from flygym.warp import ( + RendererType, + 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], + renderer_type=RendererType.RECORDED_TRAJECTORY, + ) + 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).""" + 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] + + +# --------------------------------------------------------------------------- +# 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 + 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) + 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)