Deployed 6a66208 with MkDocs version: 1.6.1

src/brittle_star_project/configs/config_simulation.py

@@ -26,6 +26,9 @@ class SimulationSettings:
     video_output_path: Optional[str] = None
     # Camera ID to use for video recording (1 is usually the close-up camera)
     camera_id: int = 1
+    video_width: int = 640
+    video_height: int = 480
+    video_fps: int = 60
 
     # Optional override for the sidecar metadata YAML file.
     # If None, it defaults to the model_path with a `_metadata.yaml` suffix.

src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py

@@ -65,11 +65,21 @@ class BrittleStarJaxEnvWrapper:
     def single_observation_space(self):
         return self._env.observation_space
 
-    def reset(self, seed: int = 0):
+    def reset(self, seed: int = 0, target_position: tuple[float, float] | None = None):
         self.logger.info(f"Resetting vectorized environment environments with seed {seed}")
         self._action_rng, env_rng = jax.random.split(jax.random.PRNGKey(seed), 2)
         env_rngs = jnp.array(jax.random.split(env_rng, self._num_envs))
-        state = self._vectorized_reset(rng=env_rngs)
+
+        # If a target_position is provided, pass it through to the underlying env.reset
+        if target_position is None:
+            state = jax.jit(jax.vmap(lambda rng: self._env.reset(rng=rng)))(env_rngs)
+        else:
+            tp = jnp.asarray(target_position, dtype=jnp.float32)
+            tp_batched = jnp.tile(tp[None, :], (self._num_envs, 1))
+            state = jax.jit(jax.vmap(lambda rng, t: self._env.reset(rng=rng, target_position=t)))(
+                env_rngs, tp_batched
+            )
+
         return state
 
     def sample_actions(self):

src/brittle_star_project/environment/env_wrapper.py

@@ -62,9 +62,12 @@ class BrittleStarEnv:
 
         return jax.random.PRNGKey(seed)
 
-    def reset(self, *, seed: int = 0):
+    def reset(self, *, seed: int = 0, target_position: tuple[float, float, float] | None = None):
         rng = self.make_rng(seed)
-        state = self._env.reset(rng=rng)
+        if target_position is not None:
+            state = self._env.reset(rng=rng, target_position=target_position)
+        else:
+            state = self._env.reset(rng=rng)
         return state
 
     def render(self, *, state: Any):

src/brittle_star_project/evaluation/video.py

@@ -25,6 +25,59 @@ def create_evaluation_dir(model_path: Path) -> Path:
     return eval_dir
 
 
+def _ensure_offscreen_size(model, width: int, height: int) -> None:
+    vis_global = getattr(getattr(model, "vis", None), "global_", None)
+    if vis_global is None:
+        return
+    vis_global.offwidth = int(max(width, vis_global.offwidth))
+    vis_global.offheight = int(max(height, vis_global.offheight))
+
+
+def _apply_camera_overrides(
+    model,
+    *,
+    camera_fovy: dict[int, float] | None = None,
+    camera_xyz: tuple[
+        dict[int, float] | None,
+        dict[int, float] | None,
+        dict[int, float] | None,
+    ] = (None, None, None),
+) -> None:
+    if not camera_fovy and not (camera_xyz[0] or camera_xyz[1] or camera_xyz[2]):
+        return
+
+    ncam = int(getattr(model, "ncam", 0))
+    for cam_id, fovy in (camera_fovy or {}).items():
+        if cam_id < 0 or cam_id >= ncam:
+            raise ValueError(f"Camera id {cam_id} is out of range")
+        model.cam_fovy[cam_id] = float(fovy)
+
+    for cam_id, x in (camera_xyz[0] or {}).items():
+        if cam_id < 0 or cam_id >= ncam:
+            raise ValueError(f"Camera id {cam_id} is out of range")
+        model.cam_pos[cam_id][0] = float(x)
+
+    for cam_id, y in (camera_xyz[1] or {}).items():
+        if cam_id < 0 or cam_id >= ncam:
+            raise ValueError(f"Camera id {cam_id} is out of range")
+        model.cam_pos[cam_id][1] = float(y)
+
+    for cam_id, z in (camera_xyz[2] or {}).items():
+        if cam_id < 0 or cam_id >= ncam:
+            raise ValueError(f"Camera id {cam_id} is out of range")
+        model.cam_pos[cam_id][2] = float(z)
+
+
+def hex_to_rgba(hex_color: str, alpha: float) -> np.ndarray:
+    color = hex_color.lstrip("#")
+    if len(color) != 6:
+        raise ValueError(f"Expected a 6-digit hex color, got {hex_color!r}")
+    red = int(color[0:2], 16) / 255.0
+    green = int(color[2:4], 16) / 255.0
+    blue = int(color[4:6], 16) / 255.0
+    return np.asarray([red, green, blue, float(alpha)], dtype=np.float32)
+
+
 def save_evaluation_metadata(
     eval_dir: Path,
     *,
@@ -66,6 +119,7 @@ def record_episode(
     fps: int = 60,
     width: int = 640,
     height: int = 480,
+    target_xy: tuple[float, float] | None = None,
 ) -> EpisodeResult:
     """Run an episode headlessly and record a video using MuJoCo's Renderer and imageio.
 
@@ -92,10 +146,12 @@ def record_episode(
             "Please install the evaluation dependencies: `uv pip install .[evaluation]`"
         ) from e
 
-    state = env.reset(seed=seed)
+    state = env.reset(seed=seed, target_position=target_xy)
     model = state.mj_model
     data = state.mj_data
 
+    _ensure_offscreen_size(model, width, height)
+
     renderer = mujoco.Renderer(model, width=width, height=height)
     ep_return = 0.0
     observations = _get_observations(state)
@@ -147,3 +203,133 @@ 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,
+    camera_fovy: dict[int, float] | None = None,
+    camera_xyz: tuple[
+        dict[int, float] | None,
+        dict[int, float] | None,
+        dict[int, float] | None,
+    ] = (None, None, None),
+    target_xy: tuple[float, float] | None = None,
+    robot_color: str | 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, target_position=(target_xy[0], target_xy[1], 0.0))
+    model = state.mj_model
+    data = state.mj_data
+
+    _apply_camera_overrides(model, camera_fovy=camera_fovy, camera_xyz=camera_xyz)
+
+    if robot_color is not None:
+        robot_body_id = mujoco.mj_name2id(
+            model, mujoco.mjtObj.mjOBJ_BODY, "BrittleStarMorphology/central_disk"
+        )
+        if robot_body_id < 0:
+            raise ValueError("Body 'BrittleStarMorphology/central_disk' not found in the model")
+
+        robot_rgba = hex_to_rgba(robot_color, 1.0)
+        body_parent = model.body_parentid
+        robot_body_ids = {int(robot_body_id)}
+
+        for body_id in range(1, int(model.nbody)):
+            current_body_id = int(body_id)
+            while current_body_id not in (-1, 0, int(robot_body_id)):
+                current_body_id = int(body_parent[current_body_id])
+            if current_body_id == int(robot_body_id):
+                robot_body_ids.add(body_id)
+
+        for geom_id in range(int(model.ngeom)):
+            if int(model.geom_bodyid[geom_id]) in robot_body_ids:
+                model.geom_rgba[geom_id][:] = robot_rgba
+
+    _ensure_offscreen_size(model, width, height)
+
+    renderer = mujoco.Renderer(model, width=width, height=height)
+    writers = {
+        cam_id: imageio.get_writer(str(path), fps=fps) for cam_id, path in output_paths.items()
+    }
+
+    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
+    try:
+        for _ in range(int(max_steps)):
+            for cam_id in camera_ids:
+                renderer.update_scene(data, camera=cam_id)
+                writers[cam_id].append_data(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)
+            writers[cam_id].append_data(renderer.render())
+    finally:
+        renderer.close()
+        for writer in writers.values():
+            writer.close()
+
+    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,
+    )