Add real-time mp4 recording for offline OCP runs

Adds `--video PATH` to `offline_task.py`, with playback fps defaulted to
`1/config.dt` so trajectories play at wall-clock speed.

The three LIPA configs switch to sequential LQR so offline runs
converge bit-deterministically.

Author: joaospinto

mpx/config/config_aliengo_trot_two_step.py

@@ -78,8 +78,8 @@ def _lipa_settings():
         µ_update_factor=0.9,
         cost_improvement_threshold=1e-3,
         primal_violation_threshold=1e-5,
-        use_parallel_lqr=True,
-        num_parallel_line_search_steps=8,
+        use_parallel_lqr=False,
+        num_parallel_line_search_steps=1,
     )
 
 lipa_settings = _lipa_settings()

mpx/config/config_barrel_roll.py

@@ -122,8 +122,8 @@ def _lipa_settings():
         µ_update_factor=0.9,
         cost_improvement_threshold=1e-3,
         primal_violation_threshold=1e-5,
-        use_parallel_lqr=True,
-        num_parallel_line_search_steps=8,
+        use_parallel_lqr=False,
+        num_parallel_line_search_steps=1,
     )
 
 lipa_settings = _lipa_settings()

mpx/config/config_h1_jump_forward.py

@@ -73,8 +73,8 @@ def _lipa_settings():
         cost_improvement_threshold=1e-3,
         primal_violation_threshold=1e-5,
         num_iterative_refinement_steps=2,
-        use_parallel_lqr=True,
-        num_parallel_line_search_steps=8,
+        use_parallel_lqr=False,
+        num_parallel_line_search_steps=1,
     )
 
 lipa_settings = _lipa_settings()
@@ -89,8 +89,8 @@ def _lipa_settings_enforce():
         cost_improvement_threshold=1e-3,
         primal_violation_threshold=1e-5,
         num_iterative_refinement_steps=2,
-        use_parallel_lqr=True,
-        num_parallel_line_search_steps=8,
+        use_parallel_lqr=False,
+        num_parallel_line_search_steps=1,
     )
 
 lipa_settings_enforce = _lipa_settings_enforce()

mpx/data/acrobot/scene.xml

@@ -3,7 +3,7 @@
   <option timestep="0.01" gravity="0 0 -9.81"/>
 
   <visual>
-    <global azimuth="135" elevation="-20"/>
+    <global azimuth="135" elevation="-20" offwidth="1920" offheight="1080"/>
     <headlight ambient="0.35 0.35 0.35" diffuse="0.75 0.75 0.75" specular="0.2 0.2 0.2"/>
     <rgba haze="1 1 1 1"/>
   </visual>

mpx/data/aliengo/scene_flat.xml

@@ -22,7 +22,7 @@
   <visual>
     <headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0"/>
     <rgba haze="0.15 0.25 0.35 1"/>
-    <global azimuth="-130" elevation="-20"/>
+    <global azimuth="-130" elevation="-20" offwidth="1920" offheight="1080"/>
   </visual>
 
   <asset>

mpx/data/unitree_h1/mjx_scene_h1_walk.xml

@@ -6,7 +6,7 @@
   <visual>
     <headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0"/>
     <rgba haze="0.15 0.25 0.35 1"/>
-    <global azimuth="160" elevation="-20"/>
+    <global azimuth="160" elevation="-20" offwidth="1920" offheight="1080"/>
   </visual>
 
   <asset>

mpx/examples/offline_task.py

@@ -9,6 +9,12 @@
 sys.path.append(os.path.abspath(os.path.join(dir_path, "..")))
 os.environ.setdefault("XLA_FLAGS", "--xla_gpu_enable_command_buffer=")
 
+# Headless video recording uses `mujoco.Renderer`, which requires an OpenGL
+# backend to be configured before the first `import mujoco` in the process.
+if "--video" in sys.argv:
+    os.environ.setdefault("MUJOCO_GL", "egl")
+    os.environ.setdefault("PYOPENGL_PLATFORM", "egl")
+
 import jax
 import jax.numpy as jnp
 import mujoco
@@ -217,7 +223,52 @@ def _predicted_base_positions(config, model, qpos_sequence):
     return base_positions
 
 
-def _play_mujoco_trajectory(result, headless=False, loop=True, ghost_stride=1):
+def _record_offline_trajectory_video(result, video_path, fps=None, loop_count=1, width=1280, height=720):
+    """Render the optimised offline trajectory `X` to an mp4 via offscreen GL.
+
+    No viewer involved. Uses `MUJOCO_GL=egl` (auto-set when the example sees
+    `--video` in argv) and the shared `sim_utils.VideoRecorder`. One frame per
+    state in `X`, optionally repeated `loop_count` times so short trajectories
+    aren't blink-and-miss-it. Playback is real-time by default (fps = 1/config.dt).
+    """
+    config = result["config"]
+    scene_path = result["scene_path"]
+    X = np.asarray(result["X"])
+    model = mujoco.MjModel.from_xml_path(scene_path)
+    data = mujoco.MjData(model)
+    os.makedirs(os.path.dirname(os.path.abspath(video_path)) or ".", exist_ok=True)
+    
+    config_fps = 1.0 / float(config.dt)
+    if fps is None:
+        fps = int(round(config_fps))
+        subsample = 1
+    else:
+        subsample = max(1, int(round(config_fps / fps)))
+        fps = int(round(config_fps / subsample))
+
+    recorder = sim_utils.VideoRecorder(model, video_path, fps=fps, width=width, height=height)
+    try:
+        for _ in range(loop_count):
+            for i, state in enumerate(X):
+                if i % subsample != 0:
+                    continue
+                qpos, qvel = _state_to_mujoco(config, state)
+                data.qpos = np.asarray(qpos)
+                data.qvel = np.asarray(qvel)
+                mujoco.mj_forward(model, data)
+                recorder.capture(data)
+    finally:
+        recorder.close()
+        print(f"Wrote video: {video_path} at {fps} fps (subsample {subsample}), {width}x{height}")
+
+
+def _play_mujoco_trajectory(result, headless=False, loop=True, ghost_stride=1, video=None, fps=None, width=None, height=None):
+    if video is not None:
+        # Recording mode: render trajectory frames to mp4 instead of opening
+        # the viewer. Implies headless.
+        _record_offline_trajectory_video(result, video, fps=fps, width=width or 1280, height=height or 720)
+        return
+
     config = result["config"]
     scene_path = result["scene_path"]
     X = np.asarray(result["X"])
@@ -292,6 +343,10 @@ def run_task(
     verbose=True,
     loop=True,
     lipa_enforce_inequalities=None,
+    video=None,
+    fps=None,
+    width=None,
+    height=None,
 ):
     result = solve_task(
         task_name,
@@ -308,7 +363,7 @@ def run_task(
         f"iterations {stats['n_iterations']} | "
         f"avg iter time {stats['average_iteration_time_ms']:.3f} ms"
     )
-    _play_mujoco_trajectory(result, headless=headless, loop=loop)
+    _play_mujoco_trajectory(result, headless=headless, loop=loop, video=video, fps=fps, width=width, height=height)
 
 
 def build_parser(default_task=None):
@@ -327,6 +382,12 @@ def build_parser(default_task=None):
     parser.add_argument("--max-iter", type=int, default=100)
     parser.add_argument("--quiet", action="store_true")
     parser.add_argument("--no-loop", action="store_true")
+    parser.add_argument("--video", type=str, default=None,
+                        help="Render the optimised trajectory to this mp4 path (forces headless).")
+    parser.add_argument("--fps", type=int, default=None,
+                        help="Target FPS for video recording (e.g. 50 or 60). Subsamples the trajectory.")
+    parser.add_argument("--width", type=int, default=1280, help="Video width (default: 1280).")
+    parser.add_argument("--height", type=int, default=720, help="Video height (default: 720).")
     enforce_group = parser.add_mutually_exclusive_group()
     enforce_group.add_argument(
         "--lipa-enforce-inequalities",
@@ -358,6 +419,10 @@ def main(default_task=None):
         verbose=not args.quiet,
         loop=not args.no_loop,
         lipa_enforce_inequalities=args.lipa_enforce_inequalities,
+        video=args.video,
+        fps=args.fps,
+        width=args.width,
+        height=args.height,
     )
 
 

mpx/utils/sim.py

@@ -418,16 +418,22 @@ def __init__(
         distance: float = 3.0,
         azimuth: float = 90.0,
         elevation: float = -20.0,
+        bit_depth: int = 10,
     ):
         import imageio  # late import: only needed when recording
 
         self._renderer = mujoco.Renderer(model, height=height, width=width)
+        
+        # Pick pixel format based on bit depth
+        pix_fmt = "yuv420p10le" if bit_depth == 10 else "yuv420p"
+        
         self._writer = imageio.get_writer(
             path,
             format="FFMPEG",
             codec="libx264",
             fps=fps,
             macro_block_size=1,
+            output_params=["-pix_fmt", pix_fmt],
         )
         self._cam = mujoco.MjvCamera()
         self._cam.distance = float(distance)
@@ -436,10 +442,23 @@ def __init__(
         self._cam.lookat[:] = [0.0, 0.0, 0.0]
 
     def capture(self, data: mujoco.MjData, lookat: np.ndarray | None = None) -> None:
-        """Render and append one frame; lookat defaults to the floating-base position."""
+        """Render and append one frame.
+
+        Default `lookat` is the floating-base world position (`qpos[:3]`) when
+        the model has at least 3 generalised coords — the typical case for
+        legged-robot configs. For low-dimensional models (e.g. acrobot's
+        2-DOF qpos), falls back to the world body's xpos so the camera stays
+        pointed at something sensible instead of crashing on the reshape.
+        """
 
         if lookat is None:
-            lookat = np.asarray(data.qpos[:3], dtype=np.float64)
+            qpos = np.asarray(data.qpos, dtype=np.float64)
+            if qpos.size >= 3:
+                lookat = qpos[:3]
+            else:
+                # Centre on the first non-world body's world position. Always 3D.
+                lookat = np.asarray(data.xpos[1] if data.xpos.shape[0] > 1 else data.xpos[0],
+                                    dtype=np.float64)
         self._cam.lookat[:] = np.asarray(lookat, dtype=np.float64).reshape(3)
         self._renderer.update_scene(data, self._cam)
         self._writer.append_data(self._renderer.render())