test: additional testing and visual confirmations

tests/test_config.py

@@ -1,33 +0,0 @@
-"""Tests for YAML config loading."""
-
-import sys
-from pathlib import Path
-
-import pytest
-
-# Ensure src is on the path when running from the project root
-sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
-
-CONFIGS_DIR = Path(__file__).parent.parent / "configs"
-
-
-class TestYamlConfig:
-    def test_load_yaml_config(self):
-        from experiment_logger.config_utils import load_yaml_config
-
-        config = load_yaml_config(str(CONFIGS_DIR / "default_ppo.yaml"))
-        assert isinstance(config, dict)
-        assert "total_timesteps" in config
-        assert "learning_rate" in config
-
-    def test_load_dev_test_config(self):
-        from experiment_logger.config_utils import load_yaml_config
-
-        config = load_yaml_config(str(CONFIGS_DIR / "dev_test.yaml"))
-        assert config["total_timesteps"] == 100000
-
-    def test_missing_config_raises(self):
-        from experiment_logger.config_utils import load_yaml_config
-
-        with pytest.raises(FileNotFoundError):
-            load_yaml_config("nonexistent.yaml")

tests/test_configs.py

@@ -0,0 +1,42 @@
+import pytest
+from hydra import compose, initialize
+from omegaconf import OmegaConf
+from brittle_star_project.configs.main_config import BrittleStarConfig
+from brittle_star_project.configs.register_configs import register_configs
+
+# Registration must happen before composition to enable validation against schemas
+register_configs()
+
+
+def test_config_composition_centralized():
+    """Test that the centralized configuration composes and validates correctly."""
+    with initialize(version_base="1.3", config_path="../configs"):
+        # We compose the config; it follows main_config.yaml
+        cfg = compose(config_name="main_config", overrides=["architecture=centralized"])
+
+        # Merge with the dataclass class to get a structured DictConfig,
+        # then convert to a real dataclass instance to verify validation.
+        structured_cfg = OmegaConf.to_object(OmegaConf.merge(BrittleStarConfig, cfg))
+
+        # Basic assertions
+        assert "CentralizedConfig" in str(type(structured_cfg.architecture))
+        assert isinstance(structured_cfg.ppo.learning_rate, float)
+        assert structured_cfg.ppo.learning_rate > 0
+
+
+def test_config_composition_decentralized():
+    """Test that the decentralized configuration composes and validates correctly."""
+    with initialize(version_base="1.3", config_path="../configs"):
+        cfg = compose(config_name="main_config", overrides=["architecture=decentralized"])
+
+        # Merge and convert to dataclass instance
+        structured_cfg = OmegaConf.to_object(OmegaConf.merge(BrittleStarConfig, cfg))
+
+        # Basic assertions
+        assert "DecentralizedConfig" in str(type(structured_cfg.architecture))
+        assert isinstance(structured_cfg.ppo.learning_rate, float)
+        assert structured_cfg.ppo.learning_rate > 0
+
+        # Decentralized specifics
+        assert hasattr(structured_cfg.architecture, "message_passing_steps")
+        assert structured_cfg.architecture.message_passing_steps > 0

tests/test_morphology_render.py

@@ -0,0 +1,58 @@
+import jax
+import os
+import mujoco
+from PIL import Image
+from brittle_star_project.environment.env_config import EnvConfig, MorphologyConfig, ArenaConfig
+from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
+
+
+def test_render_morphologies():
+    base_dir = "runs/renders"
+    os.makedirs(base_dir, exist_ok=True)
+
+    # --- 1. Full 5-Arm Morphology ---
+    morph_full = MorphologyConfig(segments_per_arm=[4, 4, 4, 4, 4])
+    env_full = BrittleStarJaxEnvWrapper(
+        morphology=morph_full, arena=ArenaConfig(), env_config=EnvConfig(), num_envs=1
+    )
+    state_full = env_full.reset(seed=0)
+
+    model_full = state_full.mj_model
+    data_full = state_full.mj_data
+
+    # 1. Compute forward kinematics so geoms are correctly positioned
+    mujoco.mj_forward(model_full, data_full)
+
+    # 2. Render using the environment's primary camera (camera=0)
+    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.fromarray(pixels_full).save(image_path)
+    print(f"Generated full morphology render: {image_path}")
+
+    # --- 2. Partially Amputated Morphology ---
+    morph_amp = MorphologyConfig(segments_per_arm=[4, 0, 4, 2, 4])
+    env_amp = BrittleStarJaxEnvWrapper(
+        morphology=morph_amp, arena=ArenaConfig(), env_config=EnvConfig(), num_envs=1
+    )
+    state_amp = env_amp.reset(seed=0)
+
+    model_amp = state_amp.mj_model
+    data_amp = state_amp.mj_data
+
+    # Compute forward kinematics
+    mujoco.mj_forward(model_amp, data_amp)
+
+    renderer_amp = mujoco.Renderer(model=model_amp)
+    renderer_amp.update_scene(data_amp, camera=1)
+    pixels_amp = renderer_amp.render()
+    image_path = os.path.join(base_dir, "amputated_arm.png")
+    Image.fromarray(pixels_amp).save(image_path)
+    print(f"Generated amputated morphology render: {image_path}")
+
+    print("Morphology render test successful!")
+
+
+if __name__ == "__main__":
+    test_render_morphologies()

tests/test_network_shapes.py

@@ -0,0 +1,67 @@
+import jax
+import jax.numpy as jnp
+from brittle_star_project.environment.padded_obs_wrapper import (
+    compute_padding_masks,
+    pad_observations_batched,
+)
+
+# We use Actor and OneDenseLayerMLP (as the critic) based on your mlps.py
+from brittle_star_project.MLPs.mlps import Actor, OneDenseLayerMLP
+
+
+def test_centralized_forward_pass_with_padding():
+    batch_size = 2
+
+    # 1. Simulate Amputated Observation [4, 0, 4, 2, 4] -> 14 segments total
+    # 14 segments * 2 = 28 joints
+    amputated_obs = {
+        "joint_position": jnp.zeros((batch_size, 28)),
+        "joint_velocity": jnp.zeros((batch_size, 28)),
+        "segment_contact": jnp.zeros((batch_size, 14)),
+    }
+
+    # 2. Pad Observation using the boolean scattering wrapper
+    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,
+    )
+
+    # 40 + 40 + 20 = 100 dimensions
+    assert global_state.shape == (batch_size, 100), (
+        f"Expected global state shape (2, 100), got {global_state.shape}"
+    )
+
+    # 4. Initialize dummy networks (40 actuators for the max morphology output)
+    actor = Actor(action_dim=40)
+    critic = OneDenseLayerMLP()  # Acts as the centralized critic
+
+    rng = jax.random.PRNGKey(0)
+    rng_a, rng_c = jax.random.split(rng)
+
+    # Initialize Flax variables
+    actor_params = actor.init(rng_a, global_state)
+    critic_params = critic.init(rng_c, global_state)
+
+    # 5. Forward Pass Assertions
+    action_mean, action_log_std = actor.apply(actor_params, global_state)
+    value = critic.apply(critic_params, global_state)
+
+    assert action_mean.shape == (batch_size, 40), f"Actor mean shape mismatch: {action_mean.shape}"
+    assert action_log_std.shape == (40,), f"Actor log_std shape mismatch: {action_log_std.shape}"
+    assert value.shape == (batch_size, 1) or value.shape == (batch_size,), (
+        f"Critic value shape mismatch: {value.shape}"
+    )
+
+
+if __name__ == "__main__":
+    test_centralized_forward_pass_with_padding()