diff --git a/configs/simulation/default.yaml b/configs/simulation/default.yaml index 390cd46..b2ca27a 100644 --- a/configs/simulation/default.yaml +++ b/configs/simulation/default.yaml @@ -13,3 +13,8 @@ max_steps: null # Points to a morphology config YAML file (e.g. configs/morphology/3_arms.yaml). # If null, the training morphology from the model's metadata is used. morphology_override: null + +# Video recording (requires [evaluation] extra) +record_video: false +# When null, video is saved in a per-model evaluation folder alongside the model. +video_output_path: null diff --git a/pyproject.toml b/pyproject.toml index 607e704..df3c6a5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,10 @@ cuda = [ analysis = [ "tensorboard", ] +evaluation = [ + "imageio>=2.35.0", + "imageio-ffmpeg>=0.5.1", +] [dependency-groups] dev = [ diff --git a/scripts/simulate.py b/scripts/simulate.py index 7bb2800..7f27509 100644 --- a/scripts/simulate.py +++ b/scripts/simulate.py @@ -27,6 +27,11 @@ from brittle_star_project.environment.env_config import MorphologyConfig from brittle_star_project.evaluation.checkpoint import load_metadata, metadata_to_configs from brittle_star_project.evaluation.policy import PolicyAgent from brittle_star_project.evaluation.rollout import rollout_headless, rollout_viewer +from brittle_star_project.evaluation.video import ( + record_episode, + create_evaluation_dir, + save_evaluation_metadata, +) @hydra.main(config_path="../configs", config_name="main_config", version_base="1.3") @@ -118,7 +123,48 @@ def main(dict_cfg: DictConfig) -> None: headless = bool(sim_cfg.headless) max_steps = sim_cfg.max_steps - if headless: + if sim_cfg.record_video: + if max_steps is None: + raise ValueError("simulation.max_steps is required when simulation.record_video=true") + + max_steps_i = int(max_steps) + if max_steps_i <= 0: + raise ValueError("simulation.max_steps must be > 0") + + if sim_cfg.video_output_path is None: + eval_dir = create_evaluation_dir(model_path) + output_path = eval_dir / "simulation.mp4" + else: + output_path = Path(hydra.utils.to_absolute_path(sim_cfg.video_output_path)) + eval_dir = output_path.parent + eval_dir.mkdir(parents=True, exist_ok=True) + + result = record_episode( + env=env, + policy=policy, + seed=seed, + max_steps=max_steps_i, + action_low=action_low, + action_high=action_high, + action_mask=action_mask, + output_path=output_path, + ) + + save_evaluation_metadata( + eval_dir=eval_dir, + morphology_override_path=sim_cfg.morphology_override, + seed=seed, + max_steps=max_steps_i, + result=result, + ) + final_dist_str = "n/a" if result.final_xy_dist is None else f"{result.final_xy_dist:.3f}" + print(f"Video saved to {output_path}") + print( + "episode done: " + f"return={result.return_:.6f}, len={result.length}, " + f"target_reached={result.reached_target}, final_xy_dist={final_dist_str}" + ) + elif headless: if max_steps is None: raise ValueError("simulation.max_steps is required when simulation.headless=true") diff --git a/src/brittle_star_project/configs/config_simulation.py b/src/brittle_star_project/configs/config_simulation.py index a10682d..0b93fd8 100644 --- a/src/brittle_star_project/configs/config_simulation.py +++ b/src/brittle_star_project/configs/config_simulation.py @@ -19,3 +19,8 @@ class SimulationSettings: # Observations are padded from the override morphology UP TO the training # morphology's shape via compute_padding_masks(override, reference=training). morphology_override: Optional[str] = None + + # Video recording (requires [evaluation] extra) + record_video: bool = False + # When None, video is saved in a per-model evaluation folder alongside the model. + video_output_path: Optional[str] = None diff --git a/src/brittle_star_project/evaluation/__init__.py b/src/brittle_star_project/evaluation/__init__.py index 73719e6..38a1c8b 100644 --- a/src/brittle_star_project/evaluation/__init__.py +++ b/src/brittle_star_project/evaluation/__init__.py @@ -3,6 +3,7 @@ from __future__ import annotations from .checkpoint import load_metadata, load_params, metadata_to_configs, TrainingConfig from .policy import PolicyAgent, ControlPolicy from .rollout import rollout_headless, rollout_viewer, EpisodeResult +from .video import record_episode, create_evaluation_dir, save_evaluation_metadata __all__ = [ "load_metadata", @@ -14,4 +15,7 @@ __all__ = [ "rollout_headless", "rollout_viewer", "EpisodeResult", + "record_episode", + "create_evaluation_dir", + "save_evaluation_metadata", ] diff --git a/src/brittle_star_project/evaluation/video.py b/src/brittle_star_project/evaluation/video.py new file mode 100644 index 0000000..726edcb --- /dev/null +++ b/src/brittle_star_project/evaluation/video.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +import datetime +from pathlib import Path + +import numpy as np +import yaml + +from brittle_star_project import BrittleStarEnv +from brittle_star_project.evaluation.policy import ControlPolicy +from brittle_star_project.evaluation.rollout import ( + EpisodeResult, + _get_observations, + _get_xy_distance_to_target, + _target_reached, + _maybe_clip_action, +) + + +def create_evaluation_dir(model_path: Path) -> Path: + """Create a unique timestamped directory for saving evaluation results.""" + timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + eval_dir = model_path.parent / f"{model_path.stem}_evaluations" / f"eval_{timestamp}" + eval_dir.mkdir(parents=True, exist_ok=True) + return eval_dir + + +def save_evaluation_metadata( + eval_dir: Path, + *, + morphology_override_path: str | None, + seed: int, + max_steps: int | None, + result: EpisodeResult, +) -> None: + """Save metadata about the evaluation run.""" + metadata = { + "timestamp": datetime.datetime.now().isoformat(), + "morphology_override": morphology_override_path, + "seed": seed, + "max_steps": max_steps, + "result": { + "return": float(result.return_), + "length": int(result.length), + "reached_target": bool(result.reached_target), + "final_xy_dist": float(result.final_xy_dist) + if result.final_xy_dist is not None + else None, + }, + } + with open(eval_dir / "evaluation_metadata.yaml", "w") as f: + yaml.safe_dump(metadata, f, sort_keys=False) + + +def record_episode( + *, + env: BrittleStarEnv, + policy: ControlPolicy, + seed: int, + max_steps: int, + action_low: np.ndarray | None, + action_high: np.ndarray | None, + action_mask: np.ndarray | None = None, + output_path: Path, + fps: int = 60, + width: int = 640, + height: int = 480, +) -> EpisodeResult: + """Run an episode headlessly and record a video using MuJoCo's Renderer and imageio. + + Args: + env: The environment. + policy: The policy agent. + seed: Random seed. + max_steps: Maximum number of steps. + action_low: Minimum action values. + action_high: Maximum action values. + action_mask: Boolean mask for the actions. + output_path: Where to save the .mp4 file. + fps: Frames per second for the video. + width: Video width. + height: Video height. + """ + 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 + + state = env.reset(seed=seed) + model = state.mj_model + data = state.mj_data + + # Use the first camera defined in the environment config, or default to 0 + camera_id = env._config.camera_ids[0] if env._config.camera_ids else 0 + renderer = mujoco.Renderer(model, width=width, height=height) + + ep_return = 0.0 + observations = _get_observations(state) + prev_dist = _get_xy_distance_to_target(observations) if observations else None + reached_target = _target_reached(state=state) + + frames = [] + steps = 0 + + for _ in range(int(max_steps)): + # Capture frame + renderer.update_scene(data, camera=camera_id) + frames.append(renderer.render()) + + # Step environment + 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 + + # Capture final frame + renderer.update_scene(data, camera=camera_id) + frames.append(renderer.render()) + renderer.close() + + # Save video + imageio.mimsave(str(output_path), frames, fps=fps) + + 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, + ) diff --git a/uv.lock b/uv.lock index bca7772..86d6ed9 100644 --- a/uv.lock +++ b/uv.lock @@ -42,6 +42,10 @@ analysis = [ cuda = [ { name = "jax", extra = ["cuda13"] }, ] +evaluation = [ + { name = "imageio" }, + { name = "imageio-ffmpeg" }, +] [package.dev-dependencies] dev = [ @@ -58,6 +62,8 @@ requires-dist = [ { name = "flax", specifier = ">=0.12.2" }, { name = "gymnasium", specifier = ">=1.2.3" }, { name = "hydra-core", specifier = ">=1.3.2" }, + { name = "imageio", marker = "extra == 'evaluation'", specifier = ">=2.35.0" }, + { name = "imageio-ffmpeg", marker = "extra == 'evaluation'", specifier = ">=0.5.1" }, { name = "ipykernel", specifier = "==7.2.0" }, { name = "jax", specifier = "==0.9.0.1" }, { name = "jax", extras = ["cuda13"], marker = "extra == 'cuda'", specifier = "==0.9.0.1" }, @@ -75,7 +81,7 @@ requires-dist = [ { name = "wandb", specifier = "==0.24.2" }, { name = "warp-lang" }, ] -provides-extras = ["cuda", "analysis"] +provides-extras = ["cuda", "analysis", "evaluation"] [package.metadata.requires-dev] dev = [ @@ -795,6 +801,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/49/fa/391e437a34e55095173dca5f24070d89cbc233ff85bf1c29c93248c6588d/imageio-2.37.3-py3-none-any.whl", hash = "sha256:46f5bb8522cd421c0f5ae104d8268f569d856b29eb1a13b92829d1970f32c9f0", size = 317646, upload-time = "2026-03-09T11:31:10.771Z" }, ] +[[package]] +name = "imageio-ffmpeg" +version = "0.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/44/bd/c3343c721f2a1b0c9fc71c1aebf1966a3b7f08c2eea8ed5437a2865611d6/imageio_ffmpeg-0.6.0.tar.gz", hash = "sha256:e2556bed8e005564a9f925bb7afa4002d82770d6b08825078b7697ab88ba1755", size = 25210, upload-time = "2025-01-16T21:34:32.747Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/da/58/87ef68ac83f4c7690961bce288fd8e382bc5f1513860fc7f90a9c1c1c6bf/imageio_ffmpeg-0.6.0-py3-none-macosx_10_9_intel.macosx_10_9_x86_64.whl", hash = "sha256:9d2baaf867088508d4a3458e61eeb30e945c4ad8016025545f66c4b5aaef0a61", size = 24932969, upload-time = "2025-01-16T21:34:20.464Z" }, + { url = "https://files.pythonhosted.org/packages/40/5c/f3d8a657d362cc93b81aab8feda487317da5b5d31c0e1fdfd5e986e55d17/imageio_ffmpeg-0.6.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:b1ae3173414b5fc5f538a726c4e48ea97edc0d2cdc11f103afee655c463fa742", size = 21113891, upload-time = "2025-01-16T21:34:00.277Z" }, + { url = "https://files.pythonhosted.org/packages/33/e7/1925bfbc563c39c1d2e82501d8372734a5c725e53ac3b31b4c2d081e895b/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:1d47bebd83d2c5fc770720d211855f208af8a596c82d17730aa51e815cdee6dc", size = 25632706, upload-time = "2025-01-16T21:33:53.475Z" }, + { url = "https://files.pythonhosted.org/packages/a0/2d/43c8522a2038e9d0e7dbdf3a61195ecc31ca576fb1527a528c877e87d973/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:c7e46fcec401dd990405049d2e2f475e2b397779df2519b544b8aab515195282", size = 29498237, upload-time = "2025-01-16T21:34:13.726Z" }, + { url = "https://files.pythonhosted.org/packages/a0/13/59da54728351883c3c1d9fca1710ab8eee82c7beba585df8f25ca925f08f/imageio_ffmpeg-0.6.0-py3-none-win32.whl", hash = "sha256:196faa79366b4a82f95c0f4053191d2013f4714a715780f0ad2a68ff37483cc2", size = 19652251, upload-time = "2025-01-16T21:34:06.812Z" }, + { url = "https://files.pythonhosted.org/packages/2c/c6/fa760e12a2483469e2bf5058c5faff664acf66cadb4df2ad6205b016a73d/imageio_ffmpeg-0.6.0-py3-none-win_amd64.whl", hash = "sha256:02fa47c83703c37df6bfe4896aab339013f62bf02c5ebf2dce6da56af04ffc0a", size = 31246824, upload-time = "2025-01-16T21:34:28.6Z" }, +] + [[package]] name = "iniconfig" version = "2.3.0"