feat(simulate): capture video

configs/simulation/default.yaml

@@ -13,3 +13,8 @@ max_steps: null
 # Points to a morphology config YAML file (e.g. configs/morphology/3_arms.yaml).
 # If null, the training morphology from the model's metadata is used.
 morphology_override: null
+
+# Video recording (requires [evaluation] extra)
+record_video: false
+# When null, video is saved in a per-model evaluation folder alongside the model.
+video_output_path: null

pyproject.toml

@@ -38,6 +38,10 @@ cuda = [
 analysis = [
     "tensorboard",
 ]
+evaluation = [
+    "imageio>=2.35.0",
+    "imageio-ffmpeg>=0.5.1",
+]
 
 [dependency-groups]
 dev = [

scripts/simulate.py

@@ -27,6 +27,11 @@ from brittle_star_project.environment.env_config import MorphologyConfig
 from brittle_star_project.evaluation.checkpoint import load_metadata, metadata_to_configs
 from brittle_star_project.evaluation.policy import PolicyAgent
 from brittle_star_project.evaluation.rollout import rollout_headless, rollout_viewer
+from brittle_star_project.evaluation.video import (
+    record_episode,
+    create_evaluation_dir,
+    save_evaluation_metadata,
+)
 
 
 @hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
@@ -118,7 +123,48 @@ def main(dict_cfg: DictConfig) -> None:
     headless = bool(sim_cfg.headless)
     max_steps = sim_cfg.max_steps
 
-    if headless:
+    if sim_cfg.record_video:
+        if max_steps is None:
+            raise ValueError("simulation.max_steps is required when simulation.record_video=true")
+
+        max_steps_i = int(max_steps)
+        if max_steps_i <= 0:
+            raise ValueError("simulation.max_steps must be > 0")
+
+        if sim_cfg.video_output_path is None:
+            eval_dir = create_evaluation_dir(model_path)
+            output_path = eval_dir / "simulation.mp4"
+        else:
+            output_path = Path(hydra.utils.to_absolute_path(sim_cfg.video_output_path))
+            eval_dir = output_path.parent
+            eval_dir.mkdir(parents=True, exist_ok=True)
+
+        result = record_episode(
+            env=env,
+            policy=policy,
+            seed=seed,
+            max_steps=max_steps_i,
+            action_low=action_low,
+            action_high=action_high,
+            action_mask=action_mask,
+            output_path=output_path,
+        )
+
+        save_evaluation_metadata(
+            eval_dir=eval_dir,
+            morphology_override_path=sim_cfg.morphology_override,
+            seed=seed,
+            max_steps=max_steps_i,
+            result=result,
+        )
+        final_dist_str = "n/a" if result.final_xy_dist is None else f"{result.final_xy_dist:.3f}"
+        print(f"Video saved to {output_path}")
+        print(
+            "episode done: "
+            f"return={result.return_:.6f}, len={result.length}, "
+            f"target_reached={result.reached_target}, final_xy_dist={final_dist_str}"
+        )
+    elif headless:
         if max_steps is None:
             raise ValueError("simulation.max_steps is required when simulation.headless=true")
 

src/brittle_star_project/configs/config_simulation.py

@@ -19,3 +19,8 @@ class SimulationSettings:
     # Observations are padded from the override morphology UP TO the training
     # morphology's shape via compute_padding_masks(override, reference=training).
     morphology_override: Optional[str] = None
+
+    # Video recording (requires [evaluation] extra)
+    record_video: bool = False
+    # When None, video is saved in a per-model evaluation folder alongside the model.
+    video_output_path: Optional[str] = None

src/brittle_star_project/evaluation/__init__.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from .checkpoint import load_metadata, load_params, metadata_to_configs, TrainingConfig
 from .policy import PolicyAgent, ControlPolicy
 from .rollout import rollout_headless, rollout_viewer, EpisodeResult
+from .video import record_episode, create_evaluation_dir, save_evaluation_metadata
 
 __all__ = [
     "load_metadata",
@@ -14,4 +15,7 @@ __all__ = [
     "rollout_headless",
     "rollout_viewer",
     "EpisodeResult",
+    "record_episode",
+    "create_evaluation_dir",
+    "save_evaluation_metadata",
 ]

src/brittle_star_project/evaluation/video.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import datetime
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+from brittle_star_project import BrittleStarEnv
+from brittle_star_project.evaluation.policy import ControlPolicy
+from brittle_star_project.evaluation.rollout import (
+    EpisodeResult,
+    _get_observations,
+    _get_xy_distance_to_target,
+    _target_reached,
+    _maybe_clip_action,
+)
+
+
+def create_evaluation_dir(model_path: Path) -> Path:
+    """Create a unique timestamped directory for saving evaluation results."""
+    timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    eval_dir = model_path.parent / f"{model_path.stem}_evaluations" / f"eval_{timestamp}"
+    eval_dir.mkdir(parents=True, exist_ok=True)
+    return eval_dir
+
+
+def save_evaluation_metadata(
+    eval_dir: Path,
+    *,
+    morphology_override_path: str | None,
+    seed: int,
+    max_steps: int | None,
+    result: EpisodeResult,
+) -> None:
+    """Save metadata about the evaluation run."""
+    metadata = {
+        "timestamp": datetime.datetime.now().isoformat(),
+        "morphology_override": morphology_override_path,
+        "seed": seed,
+        "max_steps": max_steps,
+        "result": {
+            "return": float(result.return_),
+            "length": int(result.length),
+            "reached_target": bool(result.reached_target),
+            "final_xy_dist": float(result.final_xy_dist)
+            if result.final_xy_dist is not None
+            else None,
+        },
+    }
+    with open(eval_dir / "evaluation_metadata.yaml", "w") as f:
+        yaml.safe_dump(metadata, f, sort_keys=False)
+
+
+def record_episode(
+    *,
+    env: BrittleStarEnv,
+    policy: ControlPolicy,
+    seed: int,
+    max_steps: int,
+    action_low: np.ndarray | None,
+    action_high: np.ndarray | None,
+    action_mask: np.ndarray | None = None,
+    output_path: Path,
+    fps: int = 60,
+    width: int = 640,
+    height: int = 480,
+) -> EpisodeResult:
+    """Run an episode headlessly and record a video using MuJoCo's Renderer and imageio.
+
+    Args:
+        env: The environment.
+        policy: The policy agent.
+        seed: Random seed.
+        max_steps: Maximum number of steps.
+        action_low: Minimum action values.
+        action_high: Maximum action values.
+        action_mask: Boolean mask for the actions.
+        output_path: Where to save the .mp4 file.
+        fps: Frames per second for the video.
+        width: Video width.
+        height: Video height.
+    """
+    try:
+        import imageio
+        import mujoco
+    except ImportError as e:
+        raise ImportError(
+            "Video recording requires 'imageio' and 'mujoco'. "
+            "Please install the evaluation dependencies: `uv pip install .[evaluation]`"
+        ) from e
+
+    state = env.reset(seed=seed)
+    model = state.mj_model
+    data = state.mj_data
+
+    # Use the first camera defined in the environment config, or default to 0
+    camera_id = env._config.camera_ids[0] if env._config.camera_ids else 0
+    renderer = mujoco.Renderer(model, width=width, height=height)
+
+    ep_return = 0.0
+    observations = _get_observations(state)
+    prev_dist = _get_xy_distance_to_target(observations) if observations else None
+    reached_target = _target_reached(state=state)
+
+    frames = []
+    steps = 0
+
+    for _ in range(int(max_steps)):
+        # Capture frame
+        renderer.update_scene(data, camera=camera_id)
+        frames.append(renderer.render())
+
+        # Step environment
+        obs_dict = observations or {}
+        action = policy.act(observations=obs_dict)
+        if action_mask is not None:
+            action = action[action_mask]
+        action = _maybe_clip_action(action, action_low, action_high)
+
+        state = env.step(state=state, action=action)
+        steps += 1
+
+        observations = _get_observations(state)
+        cur_dist = _get_xy_distance_to_target(observations) if observations else None
+        if prev_dist is not None and cur_dist is not None:
+            ep_return += prev_dist - cur_dist
+        prev_dist = cur_dist
+
+        reached_target = _target_reached(state=state)
+        if reached_target:
+            break
+
+    # Capture final frame
+    renderer.update_scene(data, camera=camera_id)
+    frames.append(renderer.render())
+    renderer.close()
+
+    # Save video
+    imageio.mimsave(str(output_path), frames, fps=fps)
+
+    final_dist = _get_xy_distance_to_target(observations) if observations else None
+    return EpisodeResult(
+        return_=ep_return,
+        length=steps,
+        reached_target=reached_target,
+        final_xy_dist=final_dist,
+    )

uv.lock

@@ -42,6 +42,10 @@ analysis = [
 cuda = [
     { name = "jax", extra = ["cuda13"] },
 ]
+evaluation = [
+    { name = "imageio" },
+    { name = "imageio-ffmpeg" },
+]
 
 [package.dev-dependencies]
 dev = [
@@ -58,6 +62,8 @@ requires-dist = [
     { name = "flax", specifier = ">=0.12.2" },
     { name = "gymnasium", specifier = ">=1.2.3" },
     { name = "hydra-core", specifier = ">=1.3.2" },
+    { name = "imageio", marker = "extra == 'evaluation'", specifier = ">=2.35.0" },
+    { name = "imageio-ffmpeg", marker = "extra == 'evaluation'", specifier = ">=0.5.1" },
     { name = "ipykernel", specifier = "==7.2.0" },
     { name = "jax", specifier = "==0.9.0.1" },
     { name = "jax", extras = ["cuda13"], marker = "extra == 'cuda'", specifier = "==0.9.0.1" },
@@ -75,7 +81,7 @@ requires-dist = [
     { name = "wandb", specifier = "==0.24.2" },
     { name = "warp-lang" },
 ]
-provides-extras = ["cuda", "analysis"]
+provides-extras = ["cuda", "analysis", "evaluation"]
 
 [package.metadata.requires-dev]
 dev = [
@@ -795,6 +801,20 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/49/fa/391e437a34e55095173dca5f24070d89cbc233ff85bf1c29c93248c6588d/imageio-2.37.3-py3-none-any.whl", hash = "sha256:46f5bb8522cd421c0f5ae104d8268f569d856b29eb1a13b92829d1970f32c9f0", size = 317646, upload-time = "2026-03-09T11:31:10.771Z" },
 ]
 
+[[package]]
+name = "imageio-ffmpeg"
+version = "0.6.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/44/bd/c3343c721f2a1b0c9fc71c1aebf1966a3b7f08c2eea8ed5437a2865611d6/imageio_ffmpeg-0.6.0.tar.gz", hash = "sha256:e2556bed8e005564a9f925bb7afa4002d82770d6b08825078b7697ab88ba1755", size = 25210, upload-time = "2025-01-16T21:34:32.747Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/da/58/87ef68ac83f4c7690961bce288fd8e382bc5f1513860fc7f90a9c1c1c6bf/imageio_ffmpeg-0.6.0-py3-none-macosx_10_9_intel.macosx_10_9_x86_64.whl", hash = "sha256:9d2baaf867088508d4a3458e61eeb30e945c4ad8016025545f66c4b5aaef0a61", size = 24932969, upload-time = "2025-01-16T21:34:20.464Z" },
+    { url = "https://files.pythonhosted.org/packages/40/5c/f3d8a657d362cc93b81aab8feda487317da5b5d31c0e1fdfd5e986e55d17/imageio_ffmpeg-0.6.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:b1ae3173414b5fc5f538a726c4e48ea97edc0d2cdc11f103afee655c463fa742", size = 21113891, upload-time = "2025-01-16T21:34:00.277Z" },
+    { url = "https://files.pythonhosted.org/packages/33/e7/1925bfbc563c39c1d2e82501d8372734a5c725e53ac3b31b4c2d081e895b/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:1d47bebd83d2c5fc770720d211855f208af8a596c82d17730aa51e815cdee6dc", size = 25632706, upload-time = "2025-01-16T21:33:53.475Z" },
+    { url = "https://files.pythonhosted.org/packages/a0/2d/43c8522a2038e9d0e7dbdf3a61195ecc31ca576fb1527a528c877e87d973/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:c7e46fcec401dd990405049d2e2f475e2b397779df2519b544b8aab515195282", size = 29498237, upload-time = "2025-01-16T21:34:13.726Z" },
+    { url = "https://files.pythonhosted.org/packages/a0/13/59da54728351883c3c1d9fca1710ab8eee82c7beba585df8f25ca925f08f/imageio_ffmpeg-0.6.0-py3-none-win32.whl", hash = "sha256:196faa79366b4a82f95c0f4053191d2013f4714a715780f0ad2a68ff37483cc2", size = 19652251, upload-time = "2025-01-16T21:34:06.812Z" },
+    { url = "https://files.pythonhosted.org/packages/2c/c6/fa760e12a2483469e2bf5058c5faff664acf66cadb4df2ad6205b016a73d/imageio_ffmpeg-0.6.0-py3-none-win_amd64.whl", hash = "sha256:02fa47c83703c37df6bfe4896aab339013f62bf02c5ebf2dce6da56af04ffc0a", size = 31246824, upload-time = "2025-01-16T21:34:28.6Z" },
+]
+
 [[package]]
 name = "iniconfig"
 version = "2.3.0"