feat: allow model metadata cli override

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 `<model_name>_metadata.yaml` alongside the model_path.
+metadata_path: null

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)

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

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(
+        raise FileNotFoundError(f"Could not find metadata YAML at {metadata_path}")
-            f"Could not find metadata YAML for {model_path.name}. Expected it at {metadata_path}"
-        )
     with open(metadata_path, "r") as f:
         return yaml.safe_load(f)
 

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)