Merge pull request #57 from SELab-3-2026/feat/video-simulation

Simulation visualisations for poster

scripts/poster_visualisations/render_poster_videos.py

@@ -0,0 +1,150 @@
+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",
+}
+
+ROBOT_COLOR_MAP = {
+    "CENTRALIZED": "#0D567C",  # Blue
+    "FULLY_CONNECTED": "#8C0E0F",  # Reddish
+    "RING": "#FCB304",  # Pale Yellow
+}
+
+
+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, default=5000)
+    parser.add_argument("--topdown-camera", type=int, default=0)
+    parser.add_argument("--follow-camera", type=int, default=1)
+    parser.add_argument("--topdown-camera-x", type=float, default=-3.0)
+    parser.add_argument("--topdown-camera-y", type=float, default=0.0)
+    parser.add_argument("--topdown-camera-z", type=float, default=4.5)
+    parser.add_argument("--topdown-camera-fovy", type=float, default=None)
+    parser.add_argument("--target-x", type=float, default=-6.0)
+    parser.add_argument("--target-y", type=float, default=0.0)
+    parser.add_argument("--width", type=int, default=2160)
+    parser.add_argument("--height", type=int, default=960)
+    parser.add_argument("--fps", type=int, default=60)
+    parser.add_argument(
+        "--robot-color",
+        default="#2B4162",
+        help="Hex color for the brittle star robot",
+    )
+    args = parser.parse_args()
+
+    if (args.target_x is None) != (args.target_y is None):
+        raise ValueError("target-x and target-y must be provided together")
+
+    target_xy = None
+    if args.target_x is not None:
+        target_xy = (float(args.target_x), float(args.target_y))
+
+    camera_fovy = None
+    if args.topdown_camera_fovy is not None:
+        camera_fovy = {args.topdown_camera: float(args.topdown_camera_fovy)}
+
+    camera_x = None
+    if args.topdown_camera_x is not None:
+        camera_x = {args.topdown_camera: float(args.topdown_camera_x)}
+
+    camera_y = None
+    if args.topdown_camera_y is not None:
+        camera_y = {args.topdown_camera: float(args.topdown_camera_y)}
+
+    camera_z = None
+    if args.topdown_camera_z is not None:
+        camera_z = {args.topdown_camera: float(args.topdown_camera_z)}
+
+    camera_xyz = (camera_x, camera_y, camera_z)
+
+    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",
+        }
+
+        print(bundle.architecture)
+        color = ROBOT_COLOR_MAP.get(bundle.architecture, args.robot_color)
+
+        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],
+            camera_fovy=camera_fovy,
+            camera_xyz=camera_xyz,
+            target_xy=target_xy,
+            robot_color=color,
+            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,10 @@
+#!/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 \

scripts/poster_visualisations/render_static_path_image.py

@@ -0,0 +1,241 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+
+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.rollout import (
+    _get_observations,
+    _maybe_clip_action,
+    _target_reached,
+)
+from brittle_star_project.evaluation.video import (
+    _apply_camera_overrides,
+    _ensure_offscreen_size,
+    hex_to_rgba,
+)
+
+ROBOT_COLOR_MAP = {
+    "CENTRALIZED": "#0D567C",  # Blue
+    "FULLY_CONNECTED": "#8C0E0F",  # Reddish
+    "RING": "#FCB304",  # Pale Yellow
+}
+
+
+def _enum_value(enum_obj, *names: str) -> int:
+    for name in names:
+        if hasattr(enum_obj, name):
+            return int(getattr(enum_obj, name))
+    raise AttributeError(f"Could not find any of {names!r} on {enum_obj!r}")
+
+
+def _append_sphere(scene, mujoco, center: np.ndarray, radius: float, rgba: np.ndarray) -> None:
+    geom = scene.geoms[scene.ngeom]
+    mujoco.mjv_initGeom(
+        geom,
+        mujoco.mjtGeom.mjGEOM_SPHERE,
+        np.asarray([radius, 0.0, 0.0], dtype=np.float32),
+        center,
+        np.eye(3, dtype=np.float32).reshape(-1),
+        rgba,
+    )
+    scene.ngeom += 1
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Render a static path image from a rollout.")
+    parser.add_argument("model", help="Path to .flax checkpoint")
+    parser.add_argument("--morphology-override", default=None)
+    parser.add_argument("--output-path", required=True)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--max-steps", type=int, default=5000)
+    parser.add_argument("--body-name", default="BrittleStarMorphology/central_disk")
+    parser.add_argument("--camera-id", type=int, default=0)
+    parser.add_argument("--camera-x", type=float, default=-3.0)
+    parser.add_argument("--camera-y", type=float, default=0.0)
+    parser.add_argument("--camera-z", type=float, default=4.5)
+    parser.add_argument("--camera-fovy", type=float, default=None)
+    parser.add_argument("--target-x", type=float, default=-6.0)
+    parser.add_argument("--target-y", type=float, default=0.0)
+    parser.add_argument("--width", type=int, default=2160)
+    parser.add_argument("--height", type=int, default=960)
+    parser.add_argument("--frame-stride", type=int, default=15)
+    parser.add_argument(
+        "--path-color", default="#FA9F42"
+    )  # ring = #888888, centralized = #2B4162, fully connected = FA9F42
+    parser.add_argument(
+        "--robot-color",
+        default="#FA9F42",
+        help="Hex color for brittle star robot (e.g. #ff0000)",
+    )
+    args = parser.parse_args()
+
+    model_path = Path(args.model)
+    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=args.morphology_override,
+    )
+
+    if (args.target_x is None) != (args.target_y is None):
+        raise ValueError("target-x and target-y must be provided together")
+
+    target_xy = None
+    if args.target_x is not None:
+        target_xy = (float(args.target_x), float(args.target_y))
+
+    try:
+        import imageio
+        import mujoco
+    except ImportError as e:
+        raise ImportError(
+            "Static image rendering requires 'mujoco' and 'imageio'. "
+            "Please install the evaluation dependencies: `uv pip install .[evaluation]`"
+        ) from e
+
+    reset_kwargs = {}
+    if target_xy is not None:
+        reset_kwargs["target_position"] = (target_xy[0], target_xy[1], 0.0)
+
+    state = bundle.env.reset(seed=args.seed, **reset_kwargs)
+    model = state.mj_model
+    data = state.mj_data
+
+    _apply_camera_overrides(
+        model,
+        camera_fovy={args.camera_id: float(args.camera_fovy)}
+        if args.camera_fovy is not None
+        else None,
+        camera_xyz=(
+            {args.camera_id: float(args.camera_x)} if args.camera_x is not None else None,
+            {args.camera_id: float(args.camera_y)} if args.camera_y is not None else None,
+            {args.camera_id: float(args.camera_z)} if args.camera_z is not None else None,
+        ),
+    )
+    _ensure_offscreen_size(model, args.width, args.height)
+
+    body_id = mujoco.mj_name2id(
+        model, mujoco.mjtObj.mjOBJ_BODY, "BrittleStarMorphology/central_disk"
+    )
+
+    robot_rgba = hex_to_rgba(ROBOT_COLOR_MAP.get(bundle.architecture, args.robot_color), 1.0)
+
+    # Optionally override robot color by recoloring geoms belonging to the robot's body subtree.
+    if args.robot_color is not None:
+        # Collect body IDs in the subtree rooted at `body_id` by walking parent links.
+        nbody = int(model.nbody)
+        body_parent = model.body_parentid
+        robot_body_ids = set([int(body_id)])
+        for i in range(1, nbody):
+            cur = int(i)
+            # walk up until root (0) or until we hit the robot root
+            while cur not in (-1, 0, int(body_id)):
+                cur = int(body_parent[cur])
+            if cur == int(body_id):
+                robot_body_ids.add(i)
+
+        # Recolor geoms whose body id is in the robot subtree
+        for g in range(int(model.ngeom)):
+            if int(model.geom_bodyid[g]) in robot_body_ids:
+                model.geom_rgba[g][:] = robot_rgba
+
+    # Optionally override robot color by recoloring geoms belonging to the robot's body subtree.
+    if args.robot_color is not None:
+        # Collect body IDs in the subtree rooted at `body_id` by walking parent links.
+        nbody = int(model.nbody)
+        body_parent = model.body_parentid
+        robot_body_ids = set([int(body_id)])
+        for i in range(1, nbody):
+            cur = int(i)
+            # walk up until root (0) or until we hit the robot root
+            while cur not in (-1, 0, int(body_id)):
+                cur = int(body_parent[cur])
+            if cur == int(body_id):
+                robot_body_ids.add(i)
+
+        # Recolor geoms whose body id is in the robot subtree
+        for g in range(int(model.ngeom)):
+            if int(model.geom_bodyid[g]) in robot_body_ids:
+                model.geom_rgba[g][:] = robot_rgba
+
+    positions = []
+    observations = _get_observations(state)
+
+    for _ in range(int(args.max_steps)):
+        positions.append(np.asarray(data.xpos[body_id], dtype=np.float32))
+
+        obs_dict = observations or {}
+        action = bundle.policy.act(observations=obs_dict)
+        if bundle.action_mask is not None:
+            action = action[bundle.action_mask]
+        action = _maybe_clip_action(action, bundle.action_low, bundle.action_high)
+
+        state = bundle.env.step(state=state, action=action)
+        data = state.mj_data
+        observations = _get_observations(state)
+
+        if _target_reached(state=state):
+            break
+
+    positions_arr = np.vstack(positions)
+    if len(positions_arr) < 2:
+        raise ValueError("Need at least two rollout positions to render a path")
+
+    path_points = positions_arr.copy()
+    path_points[:, 2] -= 0.02
+
+    path_step = max(1, int(args.frame_stride))
+    path_points_visible = path_points[::path_step]
+    path_rgba = hex_to_rgba(ROBOT_COLOR_MAP.get(bundle.architecture, args.path_color), 0.92)
+
+    ctx = mujoco.GLContext(args.width, args.height)
+    ctx.make_current()
+    try:
+        catmask = _enum_value(mujoco.mjtCatBit, "mjCAT_ALL")
+        camera_type = _enum_value(mujoco.mjtCamera, "mjCAMERA_FIXED")
+        font_scale = _enum_value(mujoco.mjtFontScale, "mjFONTSCALE_100")
+
+        maxgeom = int(model.ngeom + len(path_points_visible) + 8)
+        scene = mujoco.MjvScene(model, maxgeom=maxgeom)
+        option = mujoco.MjvOption()
+        perturb = mujoco.MjvPerturb()
+        camera = mujoco.MjvCamera()
+        mujoco.mjv_defaultOption(option)
+        mujoco.mjv_defaultPerturb(perturb)
+        mujoco.mjv_defaultCamera(camera)
+        camera.type = camera_type
+        camera.fixedcamid = int(args.camera_id)
+        if hasattr(camera, "trackbodyid"):
+            camera.trackbodyid = -1
+
+        context = mujoco.MjrContext(model, font_scale)
+        viewport = mujoco.MjrRect(0, 0, args.width, args.height)
+
+        mujoco.mjv_updateScene(model, data, option, perturb, camera, catmask, scene)
+
+        for idx, path_point in enumerate(path_points_visible):
+            path_rgba[3] = 0.10 + 0.70 * (idx / max(len(path_points_visible) - 1, 1))
+            _append_sphere(scene, mujoco, path_point, 0.03, path_rgba)
+
+        rgb = np.empty((args.height, args.width, 3), dtype=np.uint8)
+        depth = np.empty((args.height, args.width), dtype=np.float32)
+        mujoco.mjr_render(viewport, scene, context)
+        mujoco.mjr_readPixels(rgb, depth, viewport, context)
+        imageio.imwrite(args.output_path, np.flipud(rgb))
+
+        context.free()
+    finally:
+        ctx.free()
+
+    bundle.env.close()
+
+
+if __name__ == "__main__":
+    main()

scripts/poster_visualisations/render_static_path_image.sh

@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+
+# Static path image + optional ghost render
+
+path=$1 # path to .flax model with metadata.yaml alongside it
+
+uv run scripts/poster_visualisations/render_static_path_image.py \
+  "$path" \
+  --output-path vids/poster/5arms/centralized/path.png \
+  --ghost-overlay

scripts/simulate.py

@@ -107,6 +107,9 @@ def main(dict_cfg: DictConfig) -> None:
             action_mask=action_mask,
             output_path=output_path,
             camera_id=sim_cfg.camera_id,
+            width=sim_cfg.video_width,
+            height=sim_cfg.video_height,
+            fps=sim_cfg.video_fps,
         )
 
         save_evaluation_metadata(