feat: render seperate topdown and follow camera video

scripts/poster_visualisations/render_poster_videos.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from brittle_star_project.evaluation.checkpoint import load_metadata, metadata_to_configs
+from brittle_star_project.evaluation.eval_env_builder import build_eval_env
+from brittle_star_project.evaluation.video import record_episode_multi_camera
+
+_ARCH_DIR_MAP = {
+    "CENTRALIZED": "centralized",
+    "FULLY_CONNECTED": "fully-connected",
+    "RING": "ring",
+    "SEGMENT": "segment",
+}
+
+
+def _arch_dir(name: str) -> str:
+    return _ARCH_DIR_MAP.get(name, name.lower())
+
+
+def _resolve_overrides(overrides: list[str], count: int) -> list[str | None]:
+    if not overrides:
+        return [None] * count
+    if len(overrides) == 1 and count > 1:
+        return overrides * count
+    if len(overrides) != count:
+        raise ValueError("morphology overrides must match the number of models")
+    return overrides
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Render top-down and follow videos for poster.")
+    parser.add_argument("models", nargs="+", help="Paths to .flax checkpoints")
+    parser.add_argument(
+        "--morphology-override",
+        action="append",
+        default=[],
+        help="Override morphology YAML path (repeat to match models)",
+    )
+    parser.add_argument("--output-root", default="vids/poster")
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--max-steps", type=int, required=True)
+    parser.add_argument("--topdown-camera", type=int, default=0)
+    parser.add_argument("--follow-camera", type=int, default=1)
+    parser.add_argument("--width", type=int, default=640)
+    parser.add_argument("--height", type=int, default=480)
+    parser.add_argument("--fps", type=int, default=60)
+    args = parser.parse_args()
+
+    overrides = _resolve_overrides(args.morphology_override, len(args.models))
+    output_root = Path(args.output_root)
+
+    for model_path_str, override in zip(args.models, overrides):
+        model_path = Path(model_path_str)
+        metadata = load_metadata(model_path, None)
+        training = metadata_to_configs(metadata)
+
+        bundle = build_eval_env(
+            model_path=model_path,
+            training=training,
+            metadata=metadata,
+            morphology_override_path=override,
+        )
+
+        arch_dir = _arch_dir(bundle.architecture)
+        arms_dir = f"{bundle.num_active_arms}arms"
+        out_dir = output_root / arms_dir / arch_dir
+        out_dir.mkdir(parents=True, exist_ok=True)
+
+        output_paths = {
+            args.topdown_camera: out_dir / "topdown.mp4",
+            args.follow_camera: out_dir / "follow.mp4",
+        }
+
+        result = record_episode_multi_camera(
+            env=bundle.env,
+            policy=bundle.policy,
+            seed=args.seed,
+            max_steps=args.max_steps,
+            action_low=bundle.action_low,
+            action_high=bundle.action_high,
+            action_mask=bundle.action_mask,
+            output_paths=output_paths,
+            camera_ids=[args.topdown_camera, args.follow_camera],
+            width=args.width,
+            height=args.height,
+            fps=args.fps,
+        )
+
+        final_dist = (
+            "n/a" if result.final_xy_dist is None else f"{result.final_xy_dist:.3f}"
+        )
+        print(
+            f"{arms_dir}/{arch_dir}: return={result.return_:.6f}, len={result.length}, "
+            f"target_reached={result.reached_target}, final_xy_dist={final_dist}"
+        )
+
+        bundle.env.close()
+
+
+if __name__ == "__main__":
+    main()

scripts/poster_visualisations/render_poster_videos.sh

@@ -0,0 +1,11 @@
+#!/usr/bin/env bash
+
+# Multi-camera renders per model -> vids/poster/{arms}arms/{arch}/topdown.mp4 + follow.mp4
+
+path=$1
+
+uv run scripts/poster_visualisations/render_poster_videos.py \
+  "$path"/centralized.flax \
+  "$path"/fully-connected.flax \
+  "$path"/ring.flax \
+  --max-steps 10000 --width 640 --height 480 --fps 60

src/brittle_star_project/evaluation/video.py

@@ -147,3 +147,97 @@ def record_episode(
         final_xy_dist=final_dist,
         initial_target_distance=initial_dist,
     )
+
+
+def record_episode_multi_camera(
+    *,
+    env: BrittleStarEnv,
+    policy: ControlPolicy,
+    seed: int,
+    max_steps: int,
+    action_low: np.ndarray | None,
+    action_high: np.ndarray | None,
+    output_paths: dict[int, Path],
+    action_mask: np.ndarray | None = None,
+    camera_ids: list[int] | None = None,
+    fps: int = 60,
+    width: int = 640,
+    height: int = 480,
+) -> EpisodeResult:
+    """Run one episode and render multiple camera views to separate files."""
+    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
+
+    if camera_ids is None:
+        camera_ids = list(output_paths.keys())
+
+    for cam_id in camera_ids:
+        if cam_id not in output_paths:
+            raise ValueError(f"Missing output path for camera {cam_id}")
+
+    output_paths = {cam_id: output_paths[cam_id] for cam_id in camera_ids}
+
+    for path in output_paths.values():
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+    state = env.reset(seed=seed)
+    model = state.mj_model
+    data = state.mj_data
+
+    renderer = mujoco.Renderer(model, width=width, height=height)
+    frames = {cam_id: [] for cam_id in camera_ids}
+
+    ep_return = 0.0
+    observations = _get_observations(state)
+    prev_dist = _get_xy_distance_to_target(observations) if observations else None
+    initial_dist = prev_dist
+    reached_target = _target_reached(state=state)
+
+    steps = 0
+    for _ in range(int(max_steps)):
+        for cam_id in camera_ids:
+            renderer.update_scene(data, camera=cam_id)
+            frames[cam_id].append(renderer.render())
+
+        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
+
+    for cam_id in camera_ids:
+        renderer.update_scene(data, camera=cam_id)
+        frames[cam_id].append(renderer.render())
+
+    renderer.close()
+
+    for cam_id, path in output_paths.items():
+        imageio.mimsave(str(path), frames[cam_id], 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,
+        initial_target_distance=initial_dist,
+    )