218 lines
7.9 KiB
Python
218 lines
7.9 KiB
Python
import numpy as np
|
|
import pytest
|
|
import yaml
|
|
from pathlib import Path
|
|
|
|
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,
|
|
ArenaConfig,
|
|
EnvConfig,
|
|
ObservationBoundsConfig,
|
|
MorphMode,
|
|
)
|
|
from brittle_star_project.environment.env_types import Task
|
|
|
|
|
|
def test_metadata_to_configs():
|
|
"""Test that a raw metadata dictionary correctly instantiates the typed configs."""
|
|
mock_metadata = {
|
|
"morphology": {
|
|
"segments_per_arm": [4, 0, 4, 0, 0],
|
|
"use_p_control": False,
|
|
},
|
|
"arena": {"sand_ground_color": False, "size": [15.0, 10.0]},
|
|
"environment": {
|
|
"task": "LIGHT_ESCAPE",
|
|
"simulation_time": 5000.0,
|
|
},
|
|
"obs_bounds": {"joint_velocity": [-10.0, 10.0]},
|
|
}
|
|
|
|
config = metadata_to_configs(mock_metadata)
|
|
|
|
assert isinstance(config, TrainingConfig)
|
|
|
|
# Check MorphologyConfig
|
|
assert isinstance(config.morphology, MorphologyConfig)
|
|
assert config.morphology.segments_per_arm == [4, 0, 4, 0, 0]
|
|
assert config.morphology.use_p_control is False
|
|
assert config.morphology.use_torque_control is False # default
|
|
|
|
# Check ArenaConfig
|
|
assert isinstance(config.arena, ArenaConfig)
|
|
assert config.arena.sand_ground_color is False
|
|
assert config.arena.size == [15.0, 10.0]
|
|
assert config.arena.wall_height == 1.5 # default
|
|
|
|
# Check EnvConfig
|
|
assert isinstance(config.environment, EnvConfig)
|
|
assert config.environment.task == Task.LIGHT_ESCAPE
|
|
assert config.environment.simulation_time == 5000.0
|
|
assert config.environment.time_scale == 2 # default
|
|
|
|
# Check ObservationBoundsConfig
|
|
assert isinstance(config.obs_bounds, ObservationBoundsConfig)
|
|
assert config.obs_bounds.joint_velocity == [-10.0, 10.0]
|
|
assert config.obs_bounds.segment_contact == [0.0, 1.0] # default
|
|
|
|
|
|
def test_maybe_clip_action():
|
|
"""Test action clipping against boundaries."""
|
|
# Test valid clipping
|
|
action = np.array([1.5, -2.5, 0.0])
|
|
low = np.array([-1.0, -1.0, -1.0])
|
|
high = np.array([1.0, 1.0, 1.0])
|
|
|
|
clipped = _maybe_clip_action(action, low, high)
|
|
np.testing.assert_array_equal(clipped, np.array([1.0, -1.0, 0.0]))
|
|
|
|
# Test skipping when bounds are None
|
|
unclipped_1 = _maybe_clip_action(action, None, high)
|
|
np.testing.assert_array_equal(unclipped_1, action)
|
|
|
|
unclipped_2 = _maybe_clip_action(action, low, None)
|
|
np.testing.assert_array_equal(unclipped_2, action)
|
|
|
|
# Test skipping on shape mismatch
|
|
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)
|
|
|
|
|
|
@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"
|