diff --git a/scripts/poster_visualisations/render_poster_videos.py b/scripts/poster_visualisations/render_poster_videos.py index b221687..d3cf697 100644 --- a/scripts/poster_visualisations/render_poster_videos.py +++ b/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, diff --git a/scripts/poster_visualisations/render_poster_videos.sh b/scripts/poster_visualisations/render_poster_videos.sh index a6f5b59..10648f3 100755 --- a/scripts/poster_visualisations/render_poster_videos.sh +++ b/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 \ No newline at end of file + "$path"/ring.flax \ \ No newline at end of file diff --git a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py index b514c11..cfb60cd 100644 --- a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py +++ b/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): diff --git a/src/brittle_star_project/environment/env_wrapper.py b/src/brittle_star_project/environment/env_wrapper.py index 8b3e51f..0d2463c 100644 --- a/src/brittle_star_project/environment/env_wrapper.py +++ b/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): diff --git a/src/brittle_star_project/evaluation/video.py b/src/brittle_star_project/evaluation/video.py index 0ba3ce7..a7a629d 100644 --- a/src/brittle_star_project/evaluation/video.py +++ b/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}