diff --git a/.gitignore b/.gitignore index b40f5c8..d23008f 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,7 @@ outputs/ multirun/ metrics/ adjacency_debug.txt +vids/ # Python-generated files __pycache__/ diff --git a/configs/simulation/default.yaml b/configs/simulation/default.yaml index 61599a4..1ebc263 100644 --- a/configs/simulation/default.yaml +++ b/configs/simulation/default.yaml @@ -21,6 +21,10 @@ video_output_path: null # Camera ID to use for video recording (1 is usually the close-up camera) camera_id: 1 +video_width: 640 +video_height: 80 +video_fps: 60 + # Optional override for the metadata YAML file path. # If null, the script looks for `_metadata.yaml` alongside the model_path. metadata_path: null diff --git a/docs/api/simulation.md b/docs/api/simulation.md index 1e288b9..45d7878 100644 --- a/docs/api/simulation.md +++ b/docs/api/simulation.md @@ -38,4 +38,17 @@ uv run scripts/simulate.py \ Videos and evaluation metadata are stored in timestamped folders alongside the model: `runs/your_run/final_model_evaluations/eval_/simulation.mp4` +### Top-Down and Follow Cameras + +Using the following script, you can render a top-down and follow camera view for multiple models at once: + +```bash +uv run scripts/poster_visualisations/render_poster_videos.py \ + runs/final-models/centralized/.../final_model.flax \ + runs/final-models/fully-connected/.../final_model.flax \ + runs/final-models/ring/.../final_model.flax \ + --max-steps 10000 --width 640 --height 480 --fps 60 \ + --output-root vids/poster/ +``` + For batch evaluation and cross-model comparison, see the **[Evaluation Guide](./evaluation.md)**. diff --git a/scripts/poster_visualisations/render_poster_videos.py b/scripts/poster_visualisations/render_poster_videos.py new file mode 100644 index 0000000..eca60fc --- /dev/null +++ b/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() diff --git a/scripts/poster_visualisations/render_poster_videos.sh b/scripts/poster_visualisations/render_poster_videos.sh new file mode 100755 index 0000000..10648f3 --- /dev/null +++ b/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 \ \ No newline at end of file diff --git a/scripts/poster_visualisations/render_static_path_image.py b/scripts/poster_visualisations/render_static_path_image.py new file mode 100644 index 0000000..8229f52 --- /dev/null +++ b/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() diff --git a/scripts/poster_visualisations/render_static_path_image.sh b/scripts/poster_visualisations/render_static_path_image.sh new file mode 100755 index 0000000..ab21ca0 --- /dev/null +++ b/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 diff --git a/scripts/simulate.py b/scripts/simulate.py index 9f3e550..db8f70d 100644 --- a/scripts/simulate.py +++ b/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( diff --git a/src/brittle_star_project/configs/config_simulation.py b/src/brittle_star_project/configs/config_simulation.py index 9fd1507..0a75b8d 100644 --- a/src/brittle_star_project/configs/config_simulation.py +++ b/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. 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 b4b44db..cf2e5c2 100644 --- a/src/brittle_star_project/evaluation/video.py +++ b/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, + )