diff --git a/configs/simulation/default.yaml b/configs/simulation/default.yaml index e689c3b..61599a4 100644 --- a/configs/simulation/default.yaml +++ b/configs/simulation/default.yaml @@ -20,3 +20,7 @@ record_video: false video_output_path: null # Camera ID to use for video recording (1 is usually the close-up camera) camera_id: 1 + +# 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/scripts/simulate.py b/scripts/simulate.py index f17d8ef..a2b121e 100644 --- a/scripts/simulate.py +++ b/scripts/simulate.py @@ -51,7 +51,11 @@ def main(dict_cfg: DictConfig) -> None: raise ValueError(f"Expected a '.flax' checkpoint, got '{model_path.name}'.") # 2. Discover + load sidecar metadata YAML - metadata = load_metadata(model_path) + metadata_override = None + if sim_cfg.metadata_path is not None: + metadata_override = Path(hydra.utils.to_absolute_path(sim_cfg.metadata_path)) + + metadata = load_metadata(model_path, metadata_override) # 3. Reconstruct typed configs from metadata training = metadata_to_configs(metadata) diff --git a/src/brittle_star_project/configs/config_simulation.py b/src/brittle_star_project/configs/config_simulation.py index e15d94e..9fd1507 100644 --- a/src/brittle_star_project/configs/config_simulation.py +++ b/src/brittle_star_project/configs/config_simulation.py @@ -26,3 +26,7 @@ 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 + + # Optional override for the sidecar metadata YAML file. + # If None, it defaults to the model_path with a `_metadata.yaml` suffix. + metadata_path: Optional[str] = None diff --git a/src/brittle_star_project/evaluation/checkpoint.py b/src/brittle_star_project/evaluation/checkpoint.py index ffbc27a..d10046b 100644 --- a/src/brittle_star_project/evaluation/checkpoint.py +++ b/src/brittle_star_project/evaluation/checkpoint.py @@ -56,13 +56,15 @@ def load_params(path: Path) -> dict: } -def load_metadata(model_path: Path) -> dict: +def load_metadata(model_path: Path, metadata_override_path: Path | None = None) -> dict: """Discover and load the sidecar metadata YAML file.""" - metadata_path = model_path.with_name(model_path.stem + "_metadata.yaml") + if metadata_override_path is not None: + metadata_path = metadata_override_path + else: + metadata_path = model_path.with_name(model_path.stem + "_metadata.yaml") + if not metadata_path.exists(): - raise FileNotFoundError( - f"Could not find metadata YAML for {model_path.name}. Expected it at {metadata_path}" - ) + raise FileNotFoundError(f"Could not find metadata YAML at {metadata_path}") with open(metadata_path, "r") as f: return yaml.safe_load(f) diff --git a/tests/test_evaluation.py b/tests/test_evaluation.py index c8bffa9..bbb03f8 100644 --- a/tests/test_evaluation.py +++ b/tests/test_evaluation.py @@ -1,6 +1,13 @@ import numpy as np +import pytest +import yaml +from pathlib import Path -from brittle_star_project.evaluation.checkpoint import metadata_to_configs, TrainingConfig +from brittle_star_project.evaluation.checkpoint import ( + metadata_to_configs, + TrainingConfig, + load_metadata, +) from brittle_star_project.evaluation.rollout import _maybe_clip_action from brittle_star_project.environment.env_config import ( MorphologyConfig, @@ -75,3 +82,33 @@ def test_maybe_clip_action(): wrong_low = np.array([-1.0, -1.0]) # Shape mismatch unclipped_3 = _maybe_clip_action(action, wrong_low, high) np.testing.assert_array_equal(unclipped_3, action) + + +def test_load_metadata_with_override(tmp_path: Path): + """Test that metadata can be loaded from both default and override paths.""" + # 1. Setup + model_path = tmp_path / "model.flax" + model_path.write_bytes(b"dummy") + + default_metadata_path = tmp_path / "model_metadata.yaml" + default_content = {"version": "default", "seed": 42} + with open(default_metadata_path, "w") as f: + yaml.dump(default_content, f) + + override_path = tmp_path / "custom_metadata.yaml" + override_content = {"version": "override", "seed": 1337} + with open(override_path, "w") as f: + yaml.dump(override_content, f) + + # 2. Test default behavior + loaded_default = load_metadata(model_path) + assert loaded_default == default_content + + # 3. Test override behavior + loaded_override = load_metadata(model_path, metadata_override_path=override_path) + assert loaded_override == override_content + + # 4. Test Error Case + non_existent = tmp_path / "missing.yaml" + with pytest.raises(FileNotFoundError, match="Could not find metadata YAML at"): + load_metadata(model_path, metadata_override_path=non_existent)