⚠️ Running on Google Colab? Execution there can be up to ~50× slower than on a local machine or dedicated GPU, depending on the resources Colab makes available at the time. Use Colab only for testing and following along with the tutorials – not for production runs or benchmarking.
⚠️ Choosing a GPU runtime: The second half of this tutorial uses the GPU-accelerated backend, so you need to pick a GPU runtime on Colab first. You can do this by selecting Runtime > Change runtime type > choose a GPU (e.g., T4 GPU).
# --- Google Colab setup ---
# If you opened this notebook in Google Colab, run this cell first to
# install FlyGym and enable headless (EGL) rendering. It is a no-op when
# the notebook is run anywhere else.
import os
import sys
if "google.colab" in sys.modules:
os.environ["MUJOCO_GL"] = "egl"
os.environ["PYOPENGL_PLATFORM"] = "egl"
%pip install -q tqdm "flygym[warp] @ git+https://github.com/NeLy-EPFL/flygym.git@v2.1.1"
Recording and rendering trajectories¶
So far, we have rendered videos while simulating: a renderer attached to the simulation draws a frame whenever one is due. This is convenient, but it has several drawbacks, especially when many worlds are simulated in parallel on the GPU (Tutorial 3):
- We often don't know in advance which simulations are worth watching. For example, we might want to look at the fastest-walking flies, the ones that fell over, or the episodes in which an RL policy failed. With live rendering, we either render everything or guess in advance.
- Videos are expensive to store. Rendered frames take far more space than the kinematic state (the generalized coordinates,
qpos) that they are rendered from. For many worlds and long simulations, buffering rendered frames quickly exhausts (GPU) memory, while the kinematic state remains small enough to keep for every world. - Simulation and rendering have different optimal numbers of parallel worlds. Physics simulation is memory-light and reaches peak throughput with thousands of worlds. GPU batch rendering needs much more memory per world. When rendering live, the two have to share one setting.
Trajectory recording solves these problems by splitting the two steps. During simulation, a trajectory recorder stores only the kinematic state (qpos, plus mocap body poses if the model has any) at the video frame rate. Afterwards, we can select whichever trajectories we want and render them on the CPU or the GPU, in batches of whatever size fits in memory.
Rendering a recorded trajectory produces exactly the same frames as rendering live with the same type of renderer and settings. Therefore, nothing is lost by recording instead of rendering.
In this tutorial, we will:
- Record a trajectory from a CPU simulation, save it, and render it.
- Record trajectories from 1000 parallel worlds in a GPU simulation, select the most interesting ones, and render them on the GPU.
Recording a trajectory on the CPU¶
As in Tutorial 3, we create a fly in a flat-ground world. We will replay experimentally recorded joint angles using position actuators.
from flygym.compose import (
NeuroMechFly,
ActuatorType,
FlatGroundWorld,
KinematicPosePreset,
)
from flygym.anatomy import Skeleton, JointPreset, ActuatedDOFPreset, AxisOrder
from flygym.utils.math import Rotation3D
def make_model():
fly = NeuroMechFly()
neutral_pose = KinematicPosePreset.NEUTRAL
skeleton = Skeleton(
axis_order=AxisOrder.YAW_PITCH_ROLL, joint_preset=JointPreset.LEGS_ONLY
)
fly.add_joints(skeleton, neutral_pose=neutral_pose)
actuated_dofs = skeleton.get_actuated_dofs_from_preset(
ActuatedDOFPreset.LEGS_ACTIVE_ONLY
)
fly.add_actuators(
actuated_dofs,
actuator_type=ActuatorType.POSITION,
kp=50.0,
neutral_input=neutral_pose,
)
fly.add_leg_adhesion()
fly.colorize()
cam = fly.add_tracking_camera()
world = FlatGroundWorld()
world.add_fly(fly, (0, 0, 0.8), Rotation3D("quat", (1, 0, 0, 0)))
return fly, world, cam
fly, world, cam = make_model()
We load the experimentally recorded joint angles, as in Tutorial 2.
import numpy as np
from flygym_demo.spotlight_data import MotionSnippet
sim_timestep = 1e-4
snippet = MotionSnippet()
dof_angles = snippet.get_joint_angles(
output_timestep=sim_timestep,
output_dof_order=fly.get_actuated_jointdofs_order(ActuatorType.POSITION),
)
Instead of attaching a renderer with set_renderer, we attach a trajectory recorder with set_trajectory_recorder. Its playback_speed and output_fps arguments have the same meaning as for a renderer: they determine when frames are recorded. In the simulation loop, we call record_as_needed() after every step, just as we would call render_as_needed(). (A renderer and a recorder can also be attached at the same time.)
from tqdm import trange
from flygym import Simulation
sim = Simulation(world, timestep=sim_timestep)
recorder = sim.set_trajectory_recorder(playback_speed=0.2, output_fps=25)
sim_steps = 5000 # 0.5 s
sim.reset()
sim.set_leg_adhesion_states(fly.name, np.ones(6))
sim.warmup()
for step in trange(sim_steps):
sim.set_actuator_inputs(fly.name, ActuatorType.POSITION, dof_angles[step])
sim.step()
sim.record_as_needed()
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5000/5000 [00:00<00:00, 12131.39it/s]
The recorded trajectory is a RecordedTrajectory object. It contains the simulation time and qpos of every recorded frame, the mocap body poses (none in this model), and the playback settings.
traj = recorder.trajectory
print(f"Number of frames: {traj.n_frames}")
print(f"Time of each frame (s): {traj.time[:5]} ...")
print(f"qpos shape: {traj.qpos.shape}")
print(f"Size in memory: {traj.qpos.nbytes / 1e3:.0f} KB")
Number of frames: 63 Time of each frame (s): [0.0501 0.0581 0.0661 0.0741 0.0821] ... qpos shape: (63, 73) Size in memory: 37 KB
A trajectory can be saved to and loaded from a compressed .npz file. Note that the model is not saved with the trajectory: to render a trajectory, we need a compiled MuJoCo model with the same qpos layout. If you want to render trajectories in a different session, save the model as well, e.g. with world.save_xml_with_assets(...).
from pathlib import Path
from flygym.recording import RecordedTrajectory
output_dir = Path("demo_output/8_recording_and_rendering_trajectories")
traj.save(output_dir / "cpu_trajectory.npz")
traj = RecordedTrajectory.from_file(output_dir / "cpu_trajectory.npz")
Rendering a trajectory on the CPU¶
To render trajectories on the CPU, we use render_trajectories. We give it a compiled model, one or more trajectories, and the camera(s) to render. It returns the rendered frames: one {camera_name: frames} dictionary per trajectory. If output_path is given, the frames are also saved as videos.
import mediapy
from flygym.recording import render_trajectories
frames = render_trajectories(
sim.mj_model, traj, cam, output_path=output_dir / "cpu_trajectory.mp4"
)
mediapy.show_video(frames[0][cam.name], fps=traj.output_fps)
Recording trajectories on the GPU¶
Trajectory recording is most useful for GPU simulations with many parallel worlds. As in Tutorial 3, each world replays a different part of the experimental recording. Here, world $i$ replays a 0.2 s window starting $i \times$ stride steps into the recording.
from flygym.warp import GPUSimulation
n_worlds = 1000
sim_steps = 2000 # 0.2 s
def make_target_angles_all_worlds(n_worlds, sim_steps):
"""World i replays `sim_steps` steps of the recording, starting at step
`i * stride`."""
n_total_steps, n_dofs = dof_angles.shape
stride = (n_total_steps - sim_steps) // n_worlds
target_angles = np.zeros((n_worlds, sim_steps, n_dofs), dtype=np.float32)
for world_id in range(n_worlds):
start = world_id * stride
target_angles[world_id] = dof_angles[start : start + sim_steps]
return target_angles
target_angles_all_worlds = make_target_angles_all_worlds(n_worlds, sim_steps)
fly, world, cam = make_model()
gpu_sim = GPUSimulation(world, n_worlds, timestep=sim_timestep)
/home/sibwang/Projects/flygym/src/flygym/warp/simulation.py:724: UserWarning: MJWarp does not support noslip iterations. Changing option/noslip_iterations from 5 to 0. warnings.warn(
Warp 1.16.0 initialized:
CUDA Toolkit 12.9, Driver 13.2
Devices:
"cpu" : "CPU"
"cuda:0" : "NVIDIA GeForce RTX 3080 Ti" (12 GiB, sm_86, mempool enabled)
Kernel cache:
/home/sibwang/.cache/warp/1.16.0
On the GPU, set_trajectory_recorder additionally accepts a worlds argument to record only some worlds. Since a trajectory is small, here we record all 1000 worlds.
gpu_recorder = gpu_sim.set_trajectory_recorder(playback_speed=0.2, output_fps=25)
gpu_sim.reset()
gpu_sim.set_leg_adhesion_states(fly.name, np.ones((n_worlds, 6), dtype=np.float32))
gpu_sim.warmup()
for step in trange(sim_steps):
control_inputs = target_angles_all_worlds[:, step, :]
gpu_sim.set_actuator_inputs(fly.name, ActuatorType.POSITION, control_inputs)
gpu_sim.step()
gpu_sim.record_as_needed()
trajectories = gpu_recorder.trajectories
print(f"Recorded {len(trajectories)} trajectories of {trajectories[0].n_frames} frames")
n_bytes = sum(t.qpos.nbytes for t in trajectories)
print(f"Total size of recorded qpos: {n_bytes / 1e6:.1f} MB")
Module reset_data_keyframe__locals__valid_key_mask_07a6dae6 07a6dae load on device 'cuda:0' took 0.34 ms (cached) Module reset_data__locals__reset_xfrc_applied_b4860846 b486084 load on device 'cuda:0' took 0.26 ms (cached) Module reset_data__locals__reset_M_172d3a0a 172d3a0 load on device 'cuda:0' took 0.24 ms (cached) Module reset_data__locals__reset_mocap_ffa9acca ffa9acc load on device 'cuda:0' took 0.22 ms (cached) Module reset_data__locals__reset_contact_28e8b2b7 28e8b2b load on device 'cuda:0' took 0.24 ms (cached) Module reset_data__locals__reset_sleep_104e9a34 104e9a3 load on device 'cuda:0' took 0.22 ms (cached) Module reset_data__locals__reset_nworld_b1bdcc9b b1bdcc9 load on device 'cuda:0' took 0.27 ms (cached) Module reset_data_keyframe__locals__reset_keyframe_data_1fcd7654 1fcd765 load on device 'cuda:0' took 0.29 ms (cached) Module flygym.warp.utils c63eba3 load on device 'cuda:0' took 0.60 ms (cached) Module mujoco_warp._src.smooth 95df6f0 load on device 'cuda:0' took 6.39 ms (cached) Module _nxn_broadphase__locals__kernel_212ba6dc 212ba6d load on device 'cuda:0' took 0.27 ms (cached) Module _nxn_broadphase__locals__kernel_212ba6dc 39e2a53 load on device 'cuda:0' took 0.26 ms (cached) Module _primitive_narrowphase__locals__primitive_narrowphase_cc06ece7 cc06ece load on device 'cuda:0' took 0.49 ms (cached) Module mujoco_warp._src.constraint 5b5fe2d load on device 'cuda:0' took 0.23 ms (cached) Module _friction_dof__locals__kernel_aae08a76 aae08a7 load on device 'cuda:0' took 0.30 ms (cached) Module _efc_contact_init__locals__kernel_86cf00d9 86cf00d load on device 'cuda:0' took 0.28 ms (cached) Module _efc_contact_jac_sparse__locals__kernel_60646703 6064670 load on device 'cuda:0' took 0.26 ms (cached) Module _efc_contact_update__locals__kernel_102b849e 102b849 load on device 'cuda:0' took 0.29 ms (cached) Module mujoco_warp._src.sensor 9a8e696 load on device 'cuda:0' took 1.25 ms (cached) Module mujoco_warp._src.forward 62e76d6 load on device 'cuda:0' took 1.07 ms (cached) Module mujoco_warp._src.passive 5924789 load on device 'cuda:0' took 1.38 ms (cached) Module _qfrc_passive_kernel__locals__kernel_a386b29e a386b29 load on device 'cuda:0' took 0.23 ms (cached) Module mul_m_kernel__locals___mul_m_2d38d8e0 2d38d8e load on device 'cuda:0' took 0.27 ms (cached) Module _energy_vel_kinetic__locals__energy_vel_kinetic_a4f45422 d3e57e0 load on device 'cuda:0' took 0.28 ms (cached) Module _qfrc_smooth__locals__kernel_b98573a2 b98573a load on device 'cuda:0' took 0.24 ms (cached) Module mujoco_warp._src.support 9cbd1f3 load on device 'cuda:0' took 0.28 ms (cached) Module _solve_LD_sparse_fused__locals__kernel_dfb8eec0 8d62dc8 load on device 'cuda:0' took 0.36 ms (cached) Module _solve_init_dof__locals__kernel_14b514c8 14b514c load on device 'cuda:0' took 0.23 ms (cached) Module mujoco_warp._src.solver 1e1f601 load on device 'cuda:0' took 0.52 ms (cached) Module _solve_init_jaref_kernel__locals__kernel_60872e73 60872e7 load on device 'cuda:0' took 0.26 ms (cached) Module mul_m_kernel__locals___mul_m_8efe01a0 8efe01a load on device 'cuda:0' took 0.26 ms (cached) Module _update_constraint_efc__locals__kernel_fc20741b fc20741 load on device 'cuda:0' took 0.25 ms (cached) Module _update_constraint_init_qfrc_constraint_sparse__locals__kernel_75c6b7f4 75c6b7f load on device 'cuda:0' took 0.24 ms (cached) Module _update_gradient_zero_grad_dot__locals__kernel_52bae6b7 52bae6b load on device 'cuda:0' took 0.23 ms (cached) Module _update_gradient_grad__locals__kernel_769b869e 769b869 load on device 'cuda:0' took 0.23 ms (cached) Module _update_gradient_init_h_sparse__locals__kernel_e7a7b918 e7a7b91 load on device 'cuda:0' took 0.25 ms (cached) Module _JTDACJ_sparse__locals__kernel_00f946d4 00f946d load on device 'cuda:0' took 0.27 ms (cached) Module _update_gradient_cholesky_blocked__locals__kernel_8b9b09d2 de3234d load on device 'cuda:0' took 47.45 ms (cached) Module _linesearch_jv_fused_kernel__locals__kernel_a7ff2a68 a7ff2a6 load on device 'cuda:0' took 0.31 ms (cached) Module _linesearch_iterative_kernel__locals__kernel_58e5ed2b f63e89a load on device 'cuda:0' took 0.47 ms (cached) Module _update_constraint_efc__locals__kernel_39bee50d 39bee50 load on device 'cuda:0' took 0.28 ms (cached) Module _update_gradient_zero_grad_dot__locals__kernel_52259719 5225971 load on device 'cuda:0' took 0.25 ms (cached) Module _update_gradient_grad__locals__kernel_11192984 1119298 load on device 'cuda:0' took 0.24 ms (cached) Module _update_gradient_h_incremental_sparse__locals__kernel_75affbaa 75affba load on device 'cuda:0' took 0.26 ms (cached) Module _update_gradient_cholesky_blocked_skip_unchanged__locals__kernel_d4290e10 f762428 load on device 'cuda:0' took 54.74 ms (cached) Module _solve_done__locals__kernel_af8e3fa6 af8e3fa load on device 'cuda:0' took 0.33 ms (cached) Module _contact_match_builder__locals___contact_match_cd7a927f cd7a927 load on device 'cuda:0' took 0.35 ms (cached) Module _contact_sort__locals__contact_sort_3159a2d4 0495cd5 load on device 'cuda:0' took 0.32 ms (cached) Module _next_time_builder__locals___next_time_56ba631e 56ba631 load on device 'cuda:0' took 0.29 ms (cached)
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2000/2000 [00:17<00:00, 114.93it/s]
Recorded 1000 trajectories of 25 frames Total size of recorded qpos: 7.3 MB
gpu_recorder.trajectories is a list of RecordedTrajectory objects, one per recorded world, in the same format as on the CPU. The world_id attribute of each trajectory tells which world it came from.
Now, suppose we want to see how walking differs between flies that walked far and flies that barely moved. We can rank the worlds by the distance walked using the recorded trajectories alone: the position of the fly's base is stored in the qpos entries of its free joint. We then select 6 worlds spread evenly over this ranking, from the farthest to the shortest walker.
import mujoco as mj
free_joint_id = mj.mj_name2id(gpu_sim.mj_model, mj.mjtObj.mjOBJ_JOINT, fly.name)
base_pos_adr = gpu_sim.mj_model.jnt_qposadr[free_joint_id] # x, y, z, then quat
def get_distance_walked(traj):
base_xy = traj.qpos[:, base_pos_adr : base_pos_adr + 2]
return np.linalg.norm(base_xy[-1] - base_xy[0])
distances = np.array([get_distance_walked(traj) for traj in trajectories])
ranking = np.argsort(distances)[::-1] # farthest first
selected_ids = ranking[np.linspace(0, n_worlds - 1, 6).astype(int)]
selected_trajectories = [trajectories[i] for i in selected_ids]
for traj in selected_trajectories:
print(f"World {traj.world_id}: walked {distances[traj.world_id]:.2f} mm")
World 872: walked 2.99 mm World 887: walked 2.56 mm World 784: walked 2.39 mm World 400: walked 1.94 mm World 58: walked 1.20 mm World 243: walked 0.08 mm
Rendering trajectories on the GPU¶
To render trajectories with MuJoCo Warp's GPU batch renderer, we use render_trajectories_gpu. All frames of all given trajectories are rendered in batches of worlds_per_batch parallel worlds. This batch size is independent of the number of simulated worlds, so we can choose it based on the available GPU memory.
As in Tutorial 3, GPU batch rendering requires a slightly modified model (e.g., without textures on the fly). We apply these modifications with modify_world_for_batch_rendering and recompile the model. The modifications don't change the qpos layout, so the recorded trajectories remain compatible with the modified model.
import warnings
from flygym.warp.rendering import modify_world_for_batch_rendering
from flygym.warp.recording import render_trajectories_gpu
with warnings.catch_warnings():
warnings.simplefilter("ignore") # warnings about the modifications
modify_world_for_batch_rendering(world)
batch_rendering_model, _ = world.compile()
frames_gpu = render_trajectories_gpu(
batch_rendering_model, selected_trajectories, cam, worlds_per_batch=32
)
mediapy.show_videos(
{
f"World {traj.world_id}": frames[cam.name]
for traj, frames in zip(selected_trajectories, frames_gpu)
},
fps=selected_trajectories[0].output_fps,
columns=3,
width=240,
)
Module mujoco_warp._src.render_util 0460c49 load on device 'cuda:0' took 0.46 ms (cached) Module mujoco_warp._src.bvh 3a73bda load on device 'cuda:0' took 0.66 ms (cached) Module _build_megakernel__locals___render_megakernel_e2259a22 d1806f2 load on device 'cuda:0' took 0.84 ms (cached)
World 872 | World 887 | World 784 |
World 400 | World 58 | World 243 |
Trajectories recorded on the CPU and on the GPU are both RecordedTrajectory objects in the same format, regardless of how they were recorded. Therefore, a trajectory recorded on the CPU can be rendered on the GPU with render_trajectories_gpu, and a trajectory recorded on the GPU can be rendered on the CPU with render_trajectories, given a model with the same qpos layout. This is useful, for example, to render selected GPU-recorded trajectories with textures or custom scene options, which only the CPU renderer supports.
Summary¶
- Use
sim.set_trajectory_recorder(...)andsim.record_as_needed()to record trajectories during a CPU or GPU simulation. Recording is cheap and can be combined with live rendering. RecordedTrajectoryobjects can be saved withsaveand loaded withfrom_file. Save the model separately if you need to render in another session.- Render trajectories with
render_trajectories(CPU) orrender_trajectories_gpu(GPU). Trajectories recorded on either backend can be rendered on either backend, and the result is identical to rendering live with the same type of renderer.