feat: customizable camera & target

scripts/poster_visualisations/render_poster_videos.py

@@ -40,14 +40,45 @@ def main() -> None:
     )
     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("--max-steps", type=int, default=2000)
     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("--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=None)
+    parser.add_argument("--topdown-camera-fovy", type=float, default=50.0)
+    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=1920)
+    parser.add_argument("--height", type=int, default=1088)
     parser.add_argument("--fps", type=int, default=60)
     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)
 
@@ -83,6 +114,9 @@ def main() -> None:
             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,
             width=args.width,
             height=args.height,
             fps=args.fps,

scripts/poster_visualisations/render_poster_videos.sh

@@ -7,5 +7,4 @@ 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
+  "$path"/ring.flax \

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,44 @@ 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 save_evaluation_metadata(
     eval_dir: Path,
     *,
@@ -66,6 +104,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 +131,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)
@@ -160,6 +201,9 @@ def record_episode_multi_camera(
     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,
     fps: int = 60,
     width: int = 640,
     height: int = 480,
@@ -186,10 +230,14 @@ def record_episode_multi_camera(
     for path in output_paths.values():
         path.parent.mkdir(parents=True, exist_ok=True)
 
-    state = env.reset(seed=seed)
+    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)
+
+    _ensure_offscreen_size(model, width, height)
+
     renderer = mujoco.Renderer(model, width=width, height=height)
     frames = {cam_id: [] for cam_id in camera_ids}