diff --git a/tests/test_evaluation.py b/tests/test_evaluation.py index bbb03f8..d8be381 100644 --- a/tests/test_evaluation.py +++ b/tests/test_evaluation.py @@ -14,6 +14,7 @@ from brittle_star_project.environment.env_config import ( ArenaConfig, EnvConfig, ObservationBoundsConfig, + MorphMode, ) from brittle_star_project.environment.env_types import Task @@ -112,3 +113,106 @@ def test_load_metadata_with_override(tmp_path: Path): 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) + + +@pytest.fixture +def mock_training_config(): + return TrainingConfig( + morphology=MorphologyConfig( + segments_per_arm=[1, 1, 1, 1, 1], morph_mode=MorphMode.CENTRALIZED + ), + arena=ArenaConfig(), + environment=EnvConfig(), + obs_bounds=ObservationBoundsConfig(), + ) + + +@pytest.fixture +def mock_metadata(): + return {"architecture": {"message_passing_steps": 2}} + + +def test_build_eval_env_training_morphology(tmp_path, mock_training_config, mock_metadata): + from brittle_star_project.evaluation.eval_env_builder import build_eval_env + from unittest.mock import patch + + model_path = tmp_path / "model.flax" + patch_target = "brittle_star_project.evaluation.eval_env_builder.PolicyAgent.from_checkpoint" + with patch(patch_target) as mock_agent: + mock_agent.return_value = "mock_policy" + bundle = build_eval_env( + model_path=model_path, + training=mock_training_config, + metadata=mock_metadata, + morphology_override_path=None, + ) + assert bundle.segments_per_arm == [1, 1, 1, 1, 1] + assert bundle.num_active_arms == 5 + assert bundle.architecture == "CENTRALIZED" + assert bundle.policy == "mock_policy" + + +def test_build_eval_env_override_morphology(tmp_path, mock_training_config, mock_metadata): + from brittle_star_project.evaluation.eval_env_builder import build_eval_env + from unittest.mock import patch + + model_path = tmp_path / "model.flax" + override_path = tmp_path / "override.yaml" + override_path.write_text(yaml.dump({"segments_per_arm": [1, 0, 1, 0, 1]})) + + with patch("brittle_star_project.evaluation.eval_env_builder.PolicyAgent.from_checkpoint"): + bundle = build_eval_env( + model_path=model_path, + training=mock_training_config, + metadata=mock_metadata, + morphology_override_path=override_path, + ) + assert bundle.segments_per_arm == [1, 0, 1, 0, 1] + assert bundle.num_active_arms == 3 + # Should be smaller than 5*N + assert sum(bundle.action_mask) < len(bundle.action_mask) + + +def test_build_eval_env_action_mask_shape(tmp_path, mock_training_config, mock_metadata): + from brittle_star_project.evaluation.eval_env_builder import build_eval_env + from unittest.mock import patch + + model_path = tmp_path / "model.flax" + override_path = tmp_path / "override.yaml" + override_path.write_text(yaml.dump({"segments_per_arm": [1, 0, 1, 0, 0]})) + + with patch("brittle_star_project.evaluation.eval_env_builder.PolicyAgent.from_checkpoint"): + bundle = build_eval_env( + model_path=model_path, + training=mock_training_config, + metadata=mock_metadata, + morphology_override_path=override_path, + ) + # For each segment with P-control, there's 2 actions (pitch and yaw). + # Total segments = 5 -> 10 actions for training. + assert len(bundle.action_mask) == 10 + # Active segments = 2 -> 4 actions active. + assert sum(bundle.action_mask) == 4 + + +def test_build_eval_env_morph_mode_inherited(tmp_path, mock_training_config, mock_metadata): + from brittle_star_project.evaluation.eval_env_builder import build_eval_env + from brittle_star_project.environment.env_config import MorphMode + from unittest.mock import patch + + model_path = tmp_path / "model.flax" + override_path = tmp_path / "override.yaml" + # No morph_mode in the override YAML + override_path.write_text(yaml.dump({"segments_per_arm": [1, 0, 1, 0, 1]})) + + # Change training config to be RING + mock_training_config.morphology.morph_mode = MorphMode.RING + + with patch("brittle_star_project.evaluation.eval_env_builder.PolicyAgent.from_checkpoint"): + bundle = build_eval_env( + model_path=model_path, + training=mock_training_config, + metadata=mock_metadata, + morphology_override_path=override_path, + ) + assert bundle.architecture == "RING"