Merge branch 'dev' into feat/message_passing

tests/test_evaluation.py

@@ -0,0 +1,114 @@
+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,
+)
+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)

tests/test_morphology_render.py

@@ -14,7 +14,7 @@ from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleSta
 
 @pytest.mark.skipif(os.getenv("CI") == "true", reason="No OpenGL display in CI")
 def test_render_morphologies():
-    base_dir = "runs/renders"
+    base_dir = "runs/morphologies"
     os.makedirs(base_dir, exist_ok=True)
 
     # --- 1. Full 5-Arm Morphology ---
@@ -34,7 +34,7 @@ def test_render_morphologies():
     renderer_full = mujoco.Renderer(model=model_full)
     renderer_full.update_scene(data_full, camera=1)
     pixels_full = renderer_full.render()
-    image_path = os.path.join(base_dir, "full_5_arm.png")
+    image_path = os.path.join(base_dir, "5_arm.png")
     Image.fromarray(pixels_full).save(image_path)
     print(f"Generated full morphology render: {image_path}")
 

tests/test_network_shapes.py

@@ -2,8 +2,8 @@ import jax
 import jax.numpy as jnp
 from brittle_star_project.environment.padded_obs_wrapper import (
     compute_padding_masks,
-    pad_observations_batched,
 )
+from brittle_star_project.environment.obs_processing import create_obs_processor
 
 # We use Actor and OneDenseLayerMLP (as the critic) based on your mlps.py
 from brittle_star_project.MLPs.mlps import Actor, OneDenseLayerMLP
@@ -20,21 +20,10 @@ def test_centralized_forward_pass_with_padding():
         "segment_contact": jnp.zeros((batch_size, 14)),
     }
 
-    # 2. Pad Observation using the boolean scattering wrapper
+    # 2. Process and Pad Observation
     masks = compute_padding_masks(segments_per_arm=(4, 0, 4, 2, 4))
-    padded_obs = pad_observations_batched(amputated_obs, masks)
-
-    # Assertions to ensure padding sizes are correct (40 joints, 20 segments)
-    assert padded_obs["joint_position"].shape == (batch_size, 40), "Padding failed for joint keys"
-    assert padded_obs["segment_contact"].shape == (batch_size, 20), (
-        "Padding failed for segment keys"
-    )
-
-    # 3. Concatenate for Centralized MLP (simulating the global state vector)
-    global_state = jnp.concatenate(
-        [padded_obs["joint_position"], padded_obs["joint_velocity"], padded_obs["segment_contact"]],
-        axis=-1,
-    )
+    obs_processor = create_obs_processor(bounds_dict={}, padding_masks=masks)
+    global_state = obs_processor(amputated_obs)
 
     # 40 + 40 + 20 = 100 dimensions
     assert global_state.shape == (batch_size, 100), (

tests/test_target_direction.py

@@ -0,0 +1,78 @@
+import jax.numpy as jnp
+
+from brittle_star_project.configs.main_config import BrittleStarConfig
+from brittle_star_project.environment.env_types import Backend
+from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
+from brittle_star_project.environment.obs_processing import create_obs_processor
+
+
+def test_raw_environment_returns_allocentric_direction():
+    """
+    Verifies that the raw environment returns a GLOBAL (allocentric)
+    direction to the target. If the robot rotates in place,
+    the global vector to the target should remain identical.
+    """
+    env = BrittleStarJaxEnvWrapper.default(num_envs=1, backend=Backend.MJX)
+    env_state = env.reset(seed=42)
+    raw_obs_1 = env_state.observations["unit_xy_direction_to_target"]
+
+    ninety_deg_z_quat = jnp.array([0.7071068, 0.0, 0.0, 0.7071068])
+    new_qpos = env_state.mjx_data.qpos.at[..., 3:7].set(ninety_deg_z_quat)
+    new_data = env_state.mjx_data.replace(qpos=new_qpos)
+    rotated_env_state = env_state.replace(mjx_data=new_data)
+
+    zero_action = jnp.zeros(env.single_action_space.shape)
+    if len(raw_obs_1.shape) > 1:
+        zero_action = jnp.expand_dims(zero_action, 0)
+    final_env_state = env.step(rotated_env_state, zero_action)
+    raw_obs_2 = final_env_state.observations["unit_xy_direction_to_target"]
+
+    # If the vector is allocentric, it should not change when the robot spins.
+    assert jnp.sum(jnp.abs(raw_obs_1 - raw_obs_2)) < 1e-4, (
+        f"The raw environment observation changed when the robot rotated! "
+        f"This means it is already egocentric. "
+        f"Obs 1: {raw_obs_1}, Obs 2: {raw_obs_2}"
+    )
+
+
+def test_processor_converts_to_egocentric_direction():
+    """
+    Verifies that the obs_processor correctly applies a 2D inverse rotation
+    matrix to convert the global target vector into a local (egocentric) vector.
+    """
+    cfg = BrittleStarConfig()
+    env = BrittleStarJaxEnvWrapper.default(num_envs=1, backend=Backend.MJX)
+
+    obs_processor = create_obs_processor(
+        bounds_dict=cfg.obs_bounds.to_bounds_dict(), padding_masks=env.padding_masks
+    )
+
+    env_state = env.reset(seed=42)
+
+    # --- Scenario 1 ---
+    # Robot is rotated 90 degrees Left (facing global Y)
+    # Target is straight ahead on the global X axis [1.0, 0.0]
+    # Because the robot is facing Y, the target on X is to its RIGHT [0.0, -1.0] locally.
+    dummy_obs_1 = dict(env_state.observations)
+    dummy_obs_1["disk_rotation"] = jnp.array([[0.0, 0.0, jnp.pi / 2.0]])
+    dummy_obs_1["unit_xy_direction_to_target"] = jnp.array([[1.0, 0.0]])
+    processed_1 = obs_processor(dummy_obs_1)
+
+    # --- Scenario 2 (used to find the array indices) ---
+    # We change ONLY the target vector so we can isolate it in the final array
+    dummy_obs_2 = dict(env_state.observations)
+    dummy_obs_2["disk_rotation"] = jnp.array([[0.0, 0.0, jnp.pi / 2.0]])
+    dummy_obs_2["unit_xy_direction_to_target"] = jnp.array([[0.0, 1.0]])
+    processed_2 = obs_processor(dummy_obs_2)
+
+    # Find the indices of the elements that changed
+    diff_array = jnp.abs(processed_1[0] - processed_2[0])
+    changed_indices = jnp.where(diff_array > 1e-4)[0]
+
+    local_target = processed_1[0, changed_indices]
+    expected_local_target = jnp.array([0.0, -1.0])
+
+    assert jnp.sum(jnp.abs(local_target - expected_local_target)) < 1e-4, (
+        f"The obs_processor did not correctly rotate the vector to egocentric. "
+        f"Expected {expected_local_target}, but got {local_target}."
+    )