Merge branch 'dev' into simulate-results

.gitignore

@@ -2,6 +2,9 @@
 artifacts/*
 runs/*
 wandb/
+outputs/
+multirun/
+metrics/
 
 # Python-generated files
 __pycache__/

configs/README.md

@@ -1,23 +1,54 @@
-# Configuration Files
+# Brittle Star Configuration System
 
-This directory contains configuration files for training experiments.
+This project uses **Hydra** for a modular, hierarchical, and strictly-typed configuration system.
 
-## Usage
+## Core Concepts
 
-Use `--config` with `scripts/train.py` to run an experiment:
+1.  **Composition over Inheritance**: Instead of one giant config file, the configuration is composed of small, domain-specific modules (PPO settings, architecture, morphology, etc.).
+2.  **Strict Typing**: Every configuration is validated against a Python dataclass schema (`ConfigStore`). Misspelled keys throw a `ConfigAttributeError` immediately.
+3.  **CLI Swapping**: You can swap entire modules or override individual values from the command line without touching code.
 
+## Directory Structure
+
+- `main_config.yaml`: The root entry point defining the default composition.
+- `experiment/`: High-level experiment settings (seed, device).
+- `logging/`: WandB and checkpointing configuration.
+- `ppo/`: PPO training hyperparameters.
+- `architecture/`: Polymorphic network architectures (centralized vs. decentralized).
+- `morphology/`: Physical robot definitions (number of segments, amputations).
+- `arena/`: Environment physics and visual settings.
+- `environment/`: Task-specific settings (Directed Locomotion, Light Escape).
+
+## Common Commands
+
+### Local Debugging
+Run a quick test with minimal iterations:
 ```bash
-python scripts/train.py --config configs/default_ppo.yaml
+python scripts/train.py experiment=dev_test ppo=fast
 ```
 
-You can overriding settings via CLI:
+### Swapping Architectures or Morphologies
+Test a decentralized controller on a 3-arm robot:
 ```bash
-python scripts/train.py --config configs/default_ppo.yaml --learning-rate 0.001
+python scripts/train.py architecture=decentralized morphology=3_arms
 ```
 
-## Available Configurations
+### HPC Production
+Run stable PPO with WandB enabled (HPC submission scripts handle the `hydra.run.dir` redirection):
+```bash
+python scripts/train.py ppo=stable logging=wandb_enabled
+```
 
-- `default_ppo.yaml`: Baseline config.
-- `dev_test.yaml`: Fast iteration for development.
-- `production_training.yaml`: Full-scale training.
-- `personal_template.yaml`: Template for team members to customize.
+### Dry-Run Validation
+Check if your configuration is valid without starting the simulation:
+```bash
+python scripts/train.py --cfg job
+```
+
+## Developer Notes
+
+- **Adding a new group**: Create a subdirectory in `configs/` and register the new dataclass in `src/brittle_star_project/configs/register_configs.py`.
+- **Typo Catching**: If you see a `ConfigAttributeError`, check for typos in your YAML keys or CLI overrides.
+- **Output Redirection**: We use `experiment.base_run_dir` to configure where logs and models are stored (defaults to `runs/`).
+  - To change it locally: `python scripts/train.py experiment.base_run_dir=/path/to/custom/dir`
+  - On HPC, ensure this points to a fast scratch storage.

configs/architecture/centralized.yaml

@@ -0,0 +1,22 @@
+# Centralized Actor-Critic Architecture
+# Baseline configuration with a single global sensor and motor.
+
+# Default values are defined in CentralizedConfig dataclass.
+# Use this configuration for standard PPO experiments.
+
+name: "centralized"
+sensor:
+  hidden_dims: [300, 300, 300]
+  activation: "tanh"
+
+motor:
+  hidden_dims: []
+  activation: "tanh"
+
+feature_extractor:
+  hidden_dims: [300, 300, 300]
+  activation: "tanh"
+
+critic:
+  hidden_dims: []
+  activation: "tanh"

configs/architecture/decentralized.yaml

@@ -0,0 +1,32 @@
+# Decentralized Actor Architecture (NerveNet-MLP variant)
+# Multi-agent/distributed configuration using local sensors, propagators, and motors.
+
+# Default values are defined in DecentralizedConfig dataclass.
+# Use this configuration for decentralized execution experiments.
+
+name: "decentralized"
+sensor:
+  hidden_dims: [300, 300, 300]
+  activation: "tanh"
+
+propagator:
+  hidden_dims: [300, 300, 300]
+  activation: "tanh"
+
+motor:
+  hidden_dims: []
+  activation: "tanh"
+
+feature_extractor:
+  hidden_dims: [300, 300, 300]
+  activation: "tanh"
+
+critic:
+  hidden_dims: []
+  activation: "tanh"
+
+# Synchronous message-passing rounds per control step
+message_passing_steps: 1
+
+# Connectivity topology (e.g., ring, fully_connected)
+topology_type: "ring"

configs/arena/default.yaml

@@ -0,0 +1,8 @@
+# Default Arena Configuration
+# Base aquarium environment settings.
+
+size: [10.0, 5.0]
+sand_ground_color: true
+attach_target: true
+wall_height: 1.5
+wall_thickness: 0.1

configs/default_ppo.yaml

@@ -1,48 +0,0 @@
-# PPO Training Configuration Template
-#
-# This file provides an example configuration for PPO training.
-# Copy this file and modify it for your specific experiments.
-#
-# Usage:
-#   python src/train.py --config-path configs/my_config.yaml
-# Or override specific parameters:
-#   python src/train.py --learning-rate 0.001 --num-envs 32
-
-# Experiment settings
-exp_name: "brittle_star_ppo"
-seed: 1
-
-# Tracking settings
-track: false  # Set to true to enable WandB logging
-wandb_project_name: "PPO-Modularity"
-wandb_entity: "SEL3-2026-Groep-4"  # Set to your WandB username or team name
-
-# Model saving
-save_model: true
-checkpoint_frequency: 100  # Save checkpoint every N iterations (0 = no checkpoints)
-
-# Environment settings
-num_envs: 16
-
-# Training hyperparameters
-total_timesteps: 10000000
-learning_rate: 0.00025
-num_steps: 128
-anneal_lr: true
-
-# PPO specific
-gamma: 0.99
-gae_lambda: 0.95
-num_minibatches: 4
-update_epochs: 4
-norm_adv: true
-clip_coef: 0.1
-clip_vloss: true
-ent_coef: 0.01
-vf_coef: 0.5
-max_grad_norm: 0.5
-target_kl: null
-
-# Hardware
-cuda: true
-torch_deterministic: true

configs/dev_test.yaml

@@ -1,42 +0,0 @@
-# Quick Development/Testing Configuration
-#
-# Fast configuration for development and testing with short runs.
-
-# Experiment settings
-exp_name: "brittle_star_dev_test"
-seed: 123
-
-# Tracking settings - IMPORTANT: Set your own wandb_entity!
-track: true
-wandb_project_name: "PPO-Modularity-Dev"
-wandb_entity: "SEL3-2026-Groep-4"  # ⚠️  SET THIS TO YOUR WANDB USERNAME OR TEAM
-
-# Model saving
-save_model: true
-checkpoint_frequency: 10  # More frequent checkpoints for testing
-
-# Environment settings
-num_envs: 4  # Smaller for faster iteration
-
-# Training hyperparameters - Fast/testing
-total_timesteps: 100000  # Short run for testing
-learning_rate: 0.001  # Higher learning rate for faster learning
-num_steps: 64  # Shorter rollouts
-anneal_lr: true
-
-# PPO specific - Optimized for quick results
-gamma: 0.99
-gae_lambda: 0.95
-num_minibatches: 2
-update_epochs: 2  # Fewer epochs for speed
-norm_adv: true
-clip_coef: 0.1
-clip_vloss: true
-ent_coef: 0.02  # Higher entropy for exploration
-vf_coef: 0.5
-max_grad_norm: 0.5
-target_kl: null
-
-# Hardware
-cuda: true
-torch_deterministic: true

configs/environment/directed_locomotion.yaml

@@ -0,0 +1,12 @@
+# Directed Locomotion Environment
+# Baseline task setting.
+
+task: DIRECTED_LOCOMOTION
+simulation_time: 5.0
+num_physics_steps_per_control_step: 10
+time_scale: 2
+camera_ids: [0, 1]
+render_size: [480, 640]
+joint_randomization_noise_scale: 0.0
+target_distance: 3.0
+light_perlin_noise_scale: 0

configs/environment/light_escape.yaml

@@ -0,0 +1,12 @@
+# Light Escape Environment
+# Advanced task requiring movement away from light source.
+
+task: LIGHT_ESCAPE
+simulation_time: 5.0
+num_physics_steps_per_control_step: 10
+time_scale: 2
+camera_ids: [0, 1]
+render_size: [480, 640]
+joint_randomization_noise_scale: 0.0
+target_distance: 3.0
+light_perlin_noise_scale: 200 # Must be integer factor of 200

configs/example.yaml

@@ -1,5 +0,0 @@
-morphology:
-  num_arms: 2
-  num_segments_per_arm: 4
-  use_p_control: true
-  use_torque_control: false

configs/experiment/base.yaml

@@ -0,0 +1,7 @@
+# Base Experiment Configuration
+# Default values align with ExperimentConfig dataclass.
+
+exp_name: "brittle_star_ppo"
+seed: 1
+torch_deterministic: true
+cuda: true

configs/experiment/dev_test.yaml

@@ -0,0 +1,7 @@
+# Testing Experiment Configuration
+# Quick experiment for local development/testing.
+
+exp_name: "dev_test_brittle_star"
+seed: 42
+torch_deterministic: true
+cuda: true

configs/experiment/hpc_smoke_test.yaml

@@ -0,0 +1,7 @@
+# HPC Smoke Test Configuration
+# Uses minimal settings but simulates HPC environment.
+
+exp_name: "hpc_smoke_test"
+seed: 123
+torch_deterministic: true
+cuda: true

configs/hpc/smoke_test.yaml

@@ -1,11 +0,0 @@
-# Minimal config to verify HPC setup is functional.
-# Run with: python scripts/train.py --config-path configs/hpc/smoke_test.yaml
-exp_name: "hpc_smoke_test"
-seed: 0
-track: false           # Test WandB integration
-capture_video: false  # No rendering for smoke test
-save_model: true      # Test the end-of-training save routine
-num_envs: 512
-total_timesteps: 65536
-num_steps: 128
-cuda: true

configs/hpc/wandb_expand.yaml

@@ -1,10 +0,0 @@
-exp_name: "explained_var_fun_more_steps"
-seed: 42
-track: true
-wandb_project_name: "LET-THERE-BE-MORE-LOGGING"
-wandb_entity: "SEL3-2026-Groep-4"
-
-num_envs: 16
-num_steps: 256
-total_timesteps: 50000
-cuda: true

configs/hpc/wandb_test.yaml

@@ -1,11 +0,0 @@
-# Configuration to verify WandB online tracking
-exp_name: "hpc_wandb_verification"
-seed: 42
-track: true               # Enabled for testing WandB
-wandb_project_name: "PPO-Modularity"
-wandb_entity: "SEL3-2026-Groep-4"
-
-num_envs: 128
-total_timesteps: 50000    # Short run for quick verification
-num_steps: 128
-cuda: true

configs/logging/default.yaml

@@ -0,0 +1,11 @@
+# Default Logging Configuration
+# Offline local-only setup (WandB disabled).
+
+track: false
+wandb_project_name: "PPO-Modularity"
+wandb_entity: "SEL3-2026-Groep-4"
+capture_video: false
+save_model: true
+checkpoint_frequency: 100
+upload_model: false
+hf_entity: ""

configs/logging/wandb_enabled.yaml

@@ -0,0 +1,11 @@
+# WandB Enabled Logging Configuration
+# For production/cloud experiments with weights synced.
+
+track: true
+wandb_project_name: "PPO-Modularity"
+wandb_entity: "SEL3-2026-Groep-4"
+capture_video: false
+save_model: true
+checkpoint_frequency: 100
+upload_model: false
+hf_entity: ""

configs/main_config.yaml

@@ -0,0 +1,21 @@
+# Brittle Star Project - Main Configuration
+# This file defines the default composition of the hierarchical configuration.
+# Sub-configs are loaded from the relative directories.
+
+defaults:
+  - brittle_star_config
+  - experiment: base
+  - logging: default
+  - ppo: default
+  - architecture: centralized
+  - morphology: 5_arms_full
+  - arena: default
+  - environment: directed_locomotion
+  - simulation: default
+  - _self_
+
+hydra:
+  job:
+    chdir: True
+  run:
+    dir: ${experiment.base_run_dir}/${experiment.exp_name}/${now:%Y-%m-%d}/${now:%H-%M-%S}

configs/morphology/3_arms.yaml

@@ -0,0 +1,6 @@
+# 3 Arms Morphology Configuration
+# Symmetric amputation (arms 1 and 3 removed).
+
+segments_per_arm: [4, 0, 4, 0, 4]
+use_p_control: true
+use_torque_control: false

configs/morphology/5_arms_full.yaml

@@ -0,0 +1,6 @@
+# 5 Arms Full Morphology Configuration
+# Baseline 5-arm brittle star.
+
+segments_per_arm: [4, 4, 4, 4, 4]
+use_p_control: true
+use_torque_control: false

configs/morphology/partial_amputation.yaml

@@ -0,0 +1,6 @@
+# Partial Amputation Configuration
+# Random partial amputation for robustness testing.
+
+segments_per_arm: [4, 2, 4, 4, 4]
+use_p_control: true
+use_torque_control: false

configs/personal_template.yaml

@@ -1,40 +0,0 @@
-# Personal Configuration Example for Team Member
-#
-# Copy this template and customize for your personal experiments
-
-# Experiment settings - PERSONALIZE THESE
-exp_name: "YOUR_NAME_experiment_v1"  # ⚠️ Change YOUR_NAME 
-seed: 42
-
-# WandB settings - ⚠️ IMPORTANT: Set your credentials!
-track: true  # Enable WandB tracking
-wandb_project_name: "PPO-Modularity"
-wandb_entity: "SEL3-2026-Groep-4"  # ⚠️ CHANGE THIS to your WandB username/team
-
-# Quick experiment settings (modify as needed)
-total_timesteps: 500000  # 500K for quick results
-num_envs: 8
-learning_rate: 0.0005
-num_steps: 128
-
-# Model saving
-save_model: true  
-checkpoint_frequency: 25  # Save checkpoints frequently
-
-# Standard PPO settings (usually don't need to change)
-gamma: 0.99
-gae_lambda: 0.95
-num_minibatches: 4
-update_epochs: 4
-norm_adv: true
-clip_coef: 0.2
-clip_vloss: true
-ent_coef: 0.01
-vf_coef: 0.5
-max_grad_norm: 0.5
-target_kl: null
-anneal_lr: true
-
-# Hardware
-cuda: true
-torch_deterministic: true

configs/ppo/debug.yaml

@@ -0,0 +1,16 @@
+learning_rate: 0.0003
+total_timesteps: 409600 
+num_envs: 32
+num_steps: 32
+anneal_lr: true
+gamma: 0.99
+gae_lambda: 0.95
+num_minibatches: 32
+update_epochs: 4
+norm_adv: true
+clip_coef: 0.2
+clip_vloss: true
+ent_coef: 0.005
+vf_coef: 1.0
+max_grad_norm: 0.5
+target_kl: null

configs/ppo/default.yaml

@@ -0,0 +1,19 @@
+# Default PPO Configuration
+# Standard hyperparams from original codebase.
+
+learning_rate: 0.00025
+total_timesteps: 10000000
+num_envs: 100
+num_steps: 128
+anneal_lr: true
+gamma: 0.99
+gae_lambda: 0.95
+num_minibatches: 4
+update_epochs: 4
+norm_adv: true
+clip_coef: 0.1
+clip_vloss: true
+ent_coef: 0.01
+vf_coef: 0.5
+max_grad_norm: 0.5
+target_kl: null

configs/ppo/fast.yaml

@@ -0,0 +1,19 @@
+# Fast PPO Configuration
+# Lower timestep count for quick iterations/testing.
+
+learning_rate: 0.0005
+total_timesteps: 500000
+num_envs: 8
+num_steps: 128
+anneal_lr: true
+gamma: 0.99
+gae_lambda: 0.95
+num_minibatches: 4
+update_epochs: 4
+norm_adv: true
+clip_coef: 0.2
+clip_vloss: true
+ent_coef: 0.01
+vf_coef: 0.5
+max_grad_norm: 0.5
+target_kl: null

configs/ppo/smoke_test.yaml

@@ -0,0 +1,19 @@
+# Fast PPO Configuration
+# Lower timestep count for quick iterations/testing.
+
+learning_rate: 0.0005
+total_timesteps: 65536
+num_envs: 512
+num_steps: 128
+anneal_lr: true
+gamma: 0.99
+gae_lambda: 0.95
+num_minibatches: 4
+update_epochs: 4
+norm_adv: true
+clip_coef: 0.2
+clip_vloss: true
+ent_coef: 0.01
+vf_coef: 0.5
+max_grad_norm: 0.5
+target_kl: null

configs/ppo/stable.yaml

@@ -0,0 +1,19 @@
+# Stable PPO Configuration
+# Standard hyperparams with lower LR and larger batch.
+
+learning_rate: 0.0001
+total_timesteps: 10000000
+num_envs: 100
+num_steps: 256
+anneal_lr: true
+gamma: 0.99
+gae_lambda: 0.95
+num_minibatches: 8
+update_epochs: 4
+norm_adv: true
+clip_coef: 0.1
+clip_vloss: true
+ent_coef: 0.01
+vf_coef: 0.5
+max_grad_norm: 0.5
+target_kl: null

configs/production_training.yaml

@@ -1,42 +0,0 @@
-# Production Training Configuration
-#
-# Full-scale training configuration for production runs
-# with wandb logging enabled.
-
-# Experiment settings
-exp_name: "brittle_star_production_training"
-seed: 42
-
-# Tracking
-track: true
-capture_video: false
-wandb_project_name: "PPO-Modularity"
-wandb_entity: "SEL3-2026-Groep-4"
-
-# Model saving
-save_model: true
-checkpoint_frequency: 100  # Save checkpoint every 100 iterations
-
-# Environment settings
-num_envs: 512
-
-# Training hyperparameters
-total_timesteps: 50000000
-num_steps: 256
-num_minibatches: 4
-update_epochs: 4
-
-learning_rate: 2.5e-4
-anneal_lr: true
-gamma: 0.99
-gae_lambda: 0.95
-clip_coef: 0.1
-clip_vloss: true
-ent_coef: 0.01
-vf_coef: 0.5
-max_grad_norm: 0.5
-target_kl: null
-
-# Hardware
-cuda: true
-torch_deterministic: true

configs/simulation/default.yaml

@@ -0,0 +1,11 @@
+# Default Simulation Settings
+# These values are used by scripts/simulate.py
+
+# Path to the trained model (optional)
+model_path: null
+
+# Type of model to use if no path is provided (e.g., random)
+model_type: "random"
+
+# Execution backend (MJX or BRAX)
+backend: "MJX"

env/hpc/modules.txt

@@ -1,3 +1,4 @@
 GCCcore/13.3.0
 Python/3.12.3-GCCcore-13.3.0
 FFmpeg/7.0.2-GCCcore-13.3.0
+Hydra/1.3.2-GCCcore-13.3.0

experiments/debug-experiment-10042026/used_variables.md

@@ -0,0 +1,107 @@
+## Default envconfig
+task: Task = Task.DIRECTED_LOCOMOTION
+simulation_time: float = 500.0
+num_physics_steps_per_control_step: int = 10
+time_scale: int = 2
+camera_ids: list[int] = field(default_factory=lambda: [0, 1])
+render_size: tuple[int, int] = (480, 640)
+joint_randomization_noise_scale: float = 0.0
+target_distance: float = 3.0
+light_perlin_noise_scale: int = 0
+
+
+## Default ppoargs
+seed: int = 1
+torch_deterministic: bool = True
+cuda: bool = True
+track: bool = False
+checkpoint_frequency: int = 100
+learning_rate: float = 2.5e-4
+anneal_lr: bool = True
+gamma: float = 0.99
+gae_lambda: float = 0.95
+update_epochs: int = 4
+norm_adv: bool = True
+clip_vloss: bool = True
+max_grad_norm: float = 0.5
+target_kl: float | None = None
+batch_size: int = 0
+minibatch_size: int = 0
+num_iterations: int = 0
+
+## Used config file: (hpc/debug.yaml)
+exp_name: "debug-experiment"
+seed: 42
+track: true
+wandb_project_name: "Let's-find-that-bug"
+wandb_entity: "SEL3-2026-Groep-4"
+run_dir: "/data/gent/465/vsc46589"
+num_envs: 32
+num_steps: 32
+num_minibatches: 32
+total_timesteps: 409600 
+num_arms: 2
+cuda: true
+
+ent_coef: 0.005
+vf_coef: 1.0
+clip_coef: 0.2
+
+anneal_lr: true
+learning_rate: 0.0003
+
+## Arena config:
+size: tuple[float, float] = (10.0, 5.0)
+sand_ground_color: bool = True
+attach_target: bool = True
+wall_height: float = 1.5
+wall_thickness: float = 0.1
+
+## Morphology:
+num_segments_per_arm: int = 4
+use_p_control: bool = True
+use_torque_control: bool = False
+
+## MLPs:
+### Sensor & Feature_extractor:
+Both with 3 layers of 300 neurons per layer.
+
+class GenericDenseLayersWithActivation(nn.Module):
+    layer_sizes: Sequence[int] = field(default_factory=lambda: [64, 64])
+    activation: Callable = nn.tanh
+
+    @nn.compact
+    def __call__(self, x):
+        for size in self.layer_sizes:
+            x = nn.Dense(size, kernel_init=orthogonal(jnp.sqrt(2)))(x)
+            x = self.activation(x)
+        return x
+
+### Actor:
+class Actor(nn.Module):
+    action_dim: int
+    @nn.compact
+    def __call__(self, x):
+        mean = nn.Dense(self.action_dim, kernel_init=orthogonal(0.01), bias_init=constant(0.0))(x)
+        log_std = self.param("log_std", nn.initializers.zeros, (self.action_dim,))
+        return mean, log_std
+
+### Critic:
+class OneDenseLayerMLP(nn.Module):
+    @nn.compact
+    def __call__(self, x):
+        return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x)
+
+### Observations:
+_ALLOWED_OBS_KEYS = {
+    "joint_position",
+    "joint_velocity",
+    "joint_actuator_force",
+    "actuator_force",
+    "disk_position",
+    "disk_rotation",
+    "disk_linear_velocity",
+    "disk_angular_velocity",
+    "unit_xy_direction_to_target",
+    "xy_distance_to_target",
+}

pyproject.toml

@@ -26,7 +26,7 @@ dependencies = [
     "pyopengl>=3.1.10",
     "pyopengl-accelerate>=3.1.10",
     "pyyaml>=6.0",
-    "tyro>=1.0.10",
+    "hydra-core>=1.3.2",
     "wandb==0.24.2",
     "torch>=2.4.0",
 ]

scripts/hpc/train.pbs

@@ -63,11 +63,11 @@ elif [ -f "$PBS_O_WORKDIR/.env" ]; then
     export $(grep -v '^#' "$PBS_O_WORKDIR/.env" | xargs)
 fi
 
-# TODO Once experiments get serious, change the config
+# Run training using Hydra overrides
 python scripts/train.py \
-    --env-config-path configs/hpc/wandb_expand.yaml \
-    --hyperparameter-config-path configs/hpc/wandb_expand.yaml \
-    --run-dir "$SCRATCH_RUNDIR"
+    hydra.run.dir="$SCRATCH_RUNDIR" \
+    ppo=stable \
+    logging=wandb_enabled
 
 echo ">>> Staging out results to $DATA_RUNDIR..."
 cp -r "$SCRATCH_RUNDIR/." "$DATA_RUNDIR/"

scripts/simulate.py

@@ -1,30 +1,54 @@
+"""Simulate a trained policy in the MuJoCo viewer.
+
+Uses Hydra to load the same BrittleStarConfig that was used during training.
+Override settings via CLI, e.g.:
+    python scripts/simulate.py morphology=3_arms
+"""
+
 from __future__ import annotations
 
-import argparse
 import itertools
 import time
 from pathlib import Path
 from typing import Any
+
 import flax
+import hydra
 import jax
 import jax.numpy as jnp
 import numpy as np
+from omegaconf import DictConfig, OmegaConf
 
-from brittle_star_project import (
-    Backend,
-)
-from brittle_star_project.environment import ArenaConfig, EnvConfig, MorphologyConfig, from_file
+from brittle_star_project import BrittleStarEnv, BrittleStarEnvFactory
+from brittle_star_project.configs.main_config import BrittleStarConfig
+from brittle_star_project.configs.register_configs import register_configs
 
+_ALLOWED_OBS_KEYS = {
+    "joint_position",
+    "joint_velocity",
+    "joint_actuator_force",
+    "actuator_force",
+    "disk_position",
+    "disk_rotation",
+    "disk_linear_velocity",
+    "disk_angular_velocity",
+    "unit_xy_direction_to_target",
+    "xy_distance_to_target",
+}
 
-def _flatten_obs_dict(obs_dict: dict[str, Any]) -> jnp.ndarray:
+def _transform_obs_dict(obs_dict: dict[str, Any]) -> jnp.ndarray:
     """Flatten the env's observation dict into a 1D vector.
 
-    concatenates values in the dict's iteration order and skips empty arrays.
+    Matches training behavior:
+    - only includes keys in _ALLOWED_OBS_KEYS
+    - iterates keys in sorted order for stable layout
+    - skips empty arrays
     """
-
     parts: list[jnp.ndarray] = []
-    for v in obs_dict.values():
-        arr = jnp.asarray(v)
+    for key in sorted(obs_dict.keys()):
+        if key not in _ALLOWED_OBS_KEYS:
+            continue
+        arr = jnp.asarray(obs_dict[key])
         if arr.size == 0:
             continue
         parts.append(arr.reshape((-1,)))
@@ -45,7 +69,9 @@ class CleanRLPPOPolicy:
     ) -> None:
         from brittle_star_project.MLPs.mlps import Actor, GenericDenseLayersWithActivation
 
-        self._sensor = GenericDenseLayersWithActivation()
+        hidden_dim = int(sensor_params["params"]["Dense_0"]["kernel"].shape[1])
+
+        self._sensor = GenericDenseLayersWithActivation(layer_sizes=[hidden_dim, hidden_dim])
         self._actor = Actor(action_dim=action_dim)
         self._sensor_apply = jax.jit(self._sensor.apply)
         self._actor_apply = jax.jit(self._actor.apply)
@@ -77,28 +103,29 @@ class CleanRLPPOPolicy:
         def _parse_checkpoint(restored_obj: Any) -> tuple[Any, Any, Any, Any, Any]:
             """Extract checkpoint parts.
 
-            Returns (args_dict, sensor_params, actor_params, critic_params,
+            Returns (config_dict, sensor_params, actor_params, critic_params,
             feature_extractor_params).
 
-            `PPOTrainer` saves:
-              flax.serialization.to_bytes(
-                [vars(args), [sensor, actor, critic, feature_extractor]]
-              )
+            PPOTrainer saves:
+              flax.serialization.to_bytes([
+                config_dict,
+                [sensor_params, actor_params, critic_params, feature_extractor_params],
+              ])
 
-            `msgpack_restore()` may restore lists as dicts keyed by string
-            indices ("0", "1", ...), so we accept both shapes.
+            msgpack_restore() may restore lists as dicts keyed by string indices
+            ("0", "1", ...), so we accept both shapes.
             """
 
-            args_part: Any | None = None
+            cfg_part: Any | None = None
             params_part: Any = restored_obj
 
             if isinstance(restored_obj, (list, tuple)) and len(restored_obj) >= 2:
-                args_part = restored_obj[0]
+                cfg_part = restored_obj[0]
                 params_part = restored_obj[1]
             elif _looks_like_indexed_dict(restored_obj) and (
                 "0" in restored_obj or "1" in restored_obj
             ):
-                args_part = restored_obj.get("0", restored_obj.get(0))
+                cfg_part = restored_obj.get("0", restored_obj.get(0))
                 params_part = restored_obj.get("1", restored_obj.get(1))
 
             if _looks_like_indexed_dict(params_part):
@@ -109,7 +136,7 @@ class CleanRLPPOPolicy:
                 if sensor_params is None or actor_params is None:
                     raise ValueError("Missing required params in checkpoint")
                 return (
-                    args_part,
+                    cfg_part,
                     sensor_params,
                     actor_params,
                     critic_params,
@@ -122,7 +149,7 @@ class CleanRLPPOPolicy:
                 critic_params = params_part[2] if len(params_part) >= 3 else None
                 feature_extractor_params = params_part[3] if len(params_part) >= 4 else None
                 return (
-                    args_part,
+                    cfg_part,
                     sensor_params,
                     actor_params,
                     critic_params,
@@ -131,14 +158,13 @@ class CleanRLPPOPolicy:
 
             raise ValueError(
                 f"Unexpected checkpoint structure in {path}. "
-                "Expected [args_dict, [sensor_params, actor_params, critic_params, "
-                "feature_extractor_params]] "
-                "or an equivalent dict-indexed variant."
+                "Expected [config_dict, [sensor_params, actor_params, critic_params, "
+                "feature_extractor_params]] or an equivalent dict-indexed variant."
             )
 
         payload = path.read_bytes()
         restored = flax.serialization.msgpack_restore(payload)
-        _args_dict, sensor_params, actor_params, _critic_params, _feature_extractor_params = (
+        _cfg_dict, sensor_params, actor_params, _critic_params, _feature_extractor_params = (
             _parse_checkpoint(restored)
         )
 
@@ -149,7 +175,7 @@ class CleanRLPPOPolicy:
         )
 
     def act(self, *, observations: dict[str, Any]) -> np.ndarray:
-        obs = _flatten_obs_dict(observations)
+        obs = _transform_obs_dict(observations)
         hidden = self._sensor_apply(self._params["sensor_params"], obs)
         mean, _log_std = self._actor_apply(self._params["actor_params"], hidden)
 
@@ -172,27 +198,26 @@ def _target_reached(*, state: Any) -> bool:
 
 def _rollout_one_episode_headless(
     *,
-    env: Any,
+    env: BrittleStarEnv,
     policy: CleanRLPPOPolicy,
     seed: int,
     max_steps: int,
 ) -> tuple[float, int, bool, float | None]:
-    """Run one rollout up to `max_steps`.
+    """Run one rollout up to max_steps.
 
     Returns (return, length, reached_target, final_xy_dist).
+
+    Note: In the MJC backend, the raw env reward can be 0.0; we compute a simple
+    progress reward based on xy_distance_to_target.
     """
+
     state = env.reset(seed=seed)
 
     ep_return = 0.0
-
     observations = _get_observations(state)
     prev_dist = _get_xy_distance_to_target(observations)
     reached_target = _target_reached(state=state)
 
-    # NOTE: In the MJC backend, `state.reward` is always 0.0.
-    # To get a meaningful return, we compute a simple progress reward:
-    #   r_t = d_{t-1} - d_t
-    # where d is `xy_distance_to_target`.
     steps = 0
     for _ in range(int(max_steps)):
         action = policy.act(observations=observations)
@@ -220,7 +245,7 @@ def _rollout_one_episode_headless(
 
 def _run_one_episode_viewer(
     *,
-    env: Any,
+    env: BrittleStarEnv,
     policy: CleanRLPPOPolicy,
     seed: int,
     state: Any,
@@ -232,15 +257,14 @@ def _run_one_episode_viewer(
     model = state.mj_model
     data = state.mj_data
 
-    seed = int(seed)
+    _ = int(seed)
     episode_return = 0.0
     observations = _get_observations(state)
     prev_dist = _get_xy_distance_to_target(observations)
     reached_target = _target_reached(state=state)
 
     steps = 0
-    # Use the viewer as a context manager to avoid GLX teardown races
-    # (e.g. GLXBadDrawable from X_GLXSwapBuffers after a window is destroyed).
+    # Use the viewer as a context manager to avoid GLX teardown races.
     with mujoco.viewer.launch_passive(model, data) as viewer:
         step_iter = range(int(max_steps)) if max_steps is not None else itertools.count()
         for _step_idx in step_iter:
@@ -248,7 +272,7 @@ def _run_one_episode_viewer(
                 break
             step_start = time.time()
 
-            action = policy.act(observations=observations)
+            action = policy.act(observations=observations or {})
             if model.nu > 0 and action.shape != (int(model.nu),):
                 raise ValueError(
                     f"Policy returned action shape {action.shape}, expected ({int(model.nu)},)"
@@ -287,119 +311,86 @@ def _run_one_episode_viewer(
     )
 
 
-def parse_args() -> argparse.Namespace:
-    p = argparse.ArgumentParser(
-        description="Run a trained policy for exactly one episode (viewer or headless)."
-    )
-    p.add_argument(
-        "--config-path",
-        type=str,
-        default=None,
-        help=(
-            "Path to an environment YAML config (morphology/arena/env). "
-            "If omitted, uses the environment defaults. "
-            "Relative paths are resolved from the repository root."
-        ),
-    )
-    p.add_argument(
-        "--model",
-        type=str,
-        required=True,
-        help=("Path to the Flax checkpoint saved by scripts/train.py (final_model.flax)."),
-    )
-    p.add_argument(
-        "--headless",
-        action="store_true",
-        help="Run without the MuJoCo viewer (still exactly one episode).",
-    )
-    p.add_argument(
-        "--max-steps",
-        type=int,
-        default=None,
-        help=(
-            "Number of control steps to run. "
-            "In --headless mode this is required and acts as a fixed horizon. "
-            "In viewer mode the default is infinite (run until window closed or target reached)."
-        ),
-    )
-    p.add_argument(
-        "--backend",
-        choices=[b.value for b in Backend],
-        default=Backend.MJC.value,
-    )
-    p.add_argument("--seed", type=int, default=0)
-    return p.parse_args()
+def _infer_checkpoint_obs_dim(policy: CleanRLPPOPolicy) -> int | None:
+    """Best-effort read of the first Dense kernel input dim (obs dim)."""
+
+    try:
+        kernel = policy._params["sensor_params"]["params"]["Dense_0"]["kernel"]
+        return int(getattr(kernel, "shape")[0])
+    except Exception:
+        return None
 
 
-def main() -> None:
-    from brittle_star_project.environment import (
-        BrittleStarEnv,
-        BrittleStarEnvFactory,
+@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
+def main(dict_cfg: DictConfig) -> None:
+    # Convert DictConfig to structured dataclass, ensuring the root schema is applied.
+    config: BrittleStarConfig = OmegaConf.to_object(
+        OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg)
     )
 
-    args = parse_args()
+    backend = config.simulation.backend
+    seed = int(config.experiment.seed)
 
-    if args.config_path is None:
-        morphology_cfg = MorphologyConfig()
-        arena_cfg = ArenaConfig()
-        env_cfg = EnvConfig()
-    else:
-        repo_root = Path(__file__).resolve().parents[1]
-        config_path = Path(args.config_path)
-        if not config_path.is_absolute():
-            config_path = repo_root / config_path
-        morphology_cfg, arena_cfg, env_cfg = from_file(str(config_path))
+    model_path_str = config.simulation.model_path
+    if model_path_str is None:
+        raise ValueError(
+            "simulation.model_path must be set to a .flax checkpoint (e.g. final_model.flax)"
+        )
+
+    # Hydra chdir changes CWD; resolve relative paths relative to the invocation.
+    model_path = Path(hydra.utils.to_absolute_path(model_path_str))
+    if model_path.suffix != ".flax":
+        raise ValueError(f"Expected a '.flax' checkpoint, got '{model_path.name}'.")
 
     # ======= ENVIRONMENT SETUP =======
-
-    backend = Backend(args.backend)
-
     factory = BrittleStarEnvFactory()
-    raw_env = factory.create_environment(backend, morphology_cfg, arena_cfg, env_cfg)
+    raw_env = factory.create_environment(
+        backend,
+        config.morphology,
+        config.arena,
+        config.environment,
+    )
     env = BrittleStarEnv(
         raw_env,
         backend=backend,
-        config=env_cfg,
-        morphology_config=morphology_cfg,
+        config=config.environment,
+        morphology_config=config.morphology,
     )
 
-    seed_for_env = int(args.seed) if args.seed is not None else 0
-    state = env.reset(seed=seed_for_env)
+    state0 = env.reset(seed=seed)
 
     # ======= MODEL SETUP =======
+    nu = int(state0.mj_model.nu)
+    policy = CleanRLPPOPolicy.load(model_path, action_dim=nu)
 
-    # Extract the number of actuators (nu) from the environment's model, so we can pass it to the
-    # policy/model.
-    nu = int(state.mj_model.nu)
-
-    model_path = Path(args.model)
-    if model_path.suffix != ".flax":
+    # Helpful early failure when configs don't match the checkpoint.
+    observations0 = _get_observations(state0)
+    env_obs_dim = int(_transform_obs_dict(observations0 or {}).shape[0])
+    ckpt_obs_dim = _infer_checkpoint_obs_dim(policy)
+    if ckpt_obs_dim is not None and ckpt_obs_dim != env_obs_dim:
         raise ValueError(
-            f"Expected the training artifact '.flax', got '{model_path.name}'."
+            "Checkpoint/env mismatch: "
+            f"checkpoint expects obs_dim={ckpt_obs_dim}, env provides obs_dim={env_obs_dim}. "
+            "Use the same Hydra config (morphology/arena/environment) "
+            "that was used during training."
         )
 
-    policy = CleanRLPPOPolicy.load(
-        model_path,
-        action_dim=nu,
-    )
-
-    default_seed = seed_for_env
-
     # ======= SIMULATION =======
+    headless = bool(config.simulation.headless)
+    max_steps = config.simulation.max_steps
 
-    if args.headless:
-        if args.max_steps is None:
-            raise ValueError("--max-steps is required in --headless mode")
-        max_steps = int(args.max_steps)
-        if max_steps <= 0:
-            raise ValueError("--max-steps must be > 0")
+    if headless:
+        if max_steps is None:
+            raise ValueError("simulation.max_steps is required when simulation.headless=true")
+        max_steps_i = int(max_steps)
+        if max_steps_i <= 0:
+            raise ValueError("simulation.max_steps must be > 0")
 
-        ep_seed = int(args.seed) if args.seed is not None else default_seed
         ep_return, ep_len, reached_target, final_dist = _rollout_one_episode_headless(
             env=env,
             policy=policy,
-            seed=ep_seed,
-            max_steps=max_steps,
+            seed=seed,
+            max_steps=max_steps_i,
         )
         final_dist_str = "n/a" if final_dist is None else f"{final_dist:.3f}"
         print(
@@ -408,27 +399,29 @@ def main() -> None:
             f"target_reached={reached_target}, final_xy_dist={final_dist_str}"
         )
     else:
-        max_steps: int | None
-        if args.max_steps is None:
-            max_steps = None
+        if max_steps is not None:
+            max_steps_i = int(max_steps)
+            if max_steps_i <= 0:
+                raise ValueError("simulation.max_steps must be > 0")
+            max_steps_val: int | None = max_steps_i
         else:
-            max_steps = int(args.max_steps)
-            if max_steps <= 0:
-                raise ValueError("--max-steps must be > 0")
+            max_steps_val = None
+
+        model_dt = float(state0.mj_model.opt.timestep)
+        control_dt = model_dt * float(config.environment.num_physics_steps_per_control_step)
 
-        model_dt = float(state.mj_model.opt.timestep)
-        control_dt = model_dt * float(env_cfg.num_physics_steps_per_control_step)
         _run_one_episode_viewer(
             env=env,
             policy=policy,
-            seed=int(args.seed) if args.seed is not None else default_seed,
-            state=state,
+            seed=seed,
+            state=state0,
             control_dt=control_dt,
-            max_steps=max_steps,
+            max_steps=max_steps_val,
         )
 
     env.close()
 
 
 if __name__ == "__main__":
+    register_configs()
     main()

scripts/train.py

@@ -1,75 +1,60 @@
-import subprocess
-import time
-
-import torch
 import os
+import torch
+import hydra
+from omegaconf import DictConfig, OmegaConf
 
-from brittle_star_project.dataclasses import PPOArgs
+from brittle_star_project.configs.main_config import BrittleStarConfig
+from brittle_star_project.configs.register_configs import register_configs
 from brittle_star_project.trainers.PPOTrainer import PPOTrainer
 from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
-
-from experiment_logger import UnifiedLogger
-from experiment_logger.config_utils import merge_config_with_cli, print_config
+from experiment_logger import init_logger, get_logger
 
 
-def make_env(config_path: str | None, num_envs: int) -> BrittleStarJaxEnvWrapper:
-    if config_path is None:
-        return BrittleStarJaxEnvWrapper.default(num_envs=num_envs)
-    return BrittleStarJaxEnvWrapper.from_config(config_path, num_envs=num_envs)
+def make_env(cfg: BrittleStarConfig) -> BrittleStarJaxEnvWrapper:
+    """Create the environment using the structured configuration."""
+    return BrittleStarJaxEnvWrapper(
+        morphology=cfg.morphology,
+        arena=cfg.arena,
+        env_config=cfg.environment,
+        num_envs=cfg.ppo.num_envs,
+    )
 
 
-def parse_args() -> PPOArgs:
-    import argparse
+@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
+def main(dict_cfg: DictConfig):
+    # 1. Convert DictConfig to structured dataclass, ensuring the root schema is applied correctly.
+    config: BrittleStarConfig = OmegaConf.to_object(
+        OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg)
+    )
 
-    # Use argparse to reliably extract just the config path without swallowing --help
-    parser = argparse.ArgumentParser(add_help=False)
-    parser.add_argument("--hyperparameter-config-path", type=str, default=None)
-    known_args, _ = parser.parse_known_args()
+    # 2. Setup run metadata
+    # Hydra changes CWD to the output directory by default.
+    run_dir = os.getcwd()
+    run_name = os.path.basename(run_dir)
 
-    args = merge_config_with_cli(PPOArgs, config_file=known_args.hyperparameter_config_path)
-    return args
+    # 3. Initialize Logger
+    cfg_dict = OmegaConf.to_container(dict_cfg, resolve=True, throw_on_missing=True)
+    init_logger(
+        run_name=run_name,
+        config=cfg_dict,
+        project_name=config.logging.wandb_project_name,
+        entity=config.logging.wandb_entity,
+        base_dir=os.path.dirname(run_dir),
+        use_wandb=config.logging.track,
+    )
+    logger = get_logger()
+    logger.info(f"Hydra-initialized run: {run_name}")
+    logger.info(f"Output directory: {run_dir}")
 
+    # 4. Setup Environment and Torch
+    env = make_env(config)
+    torch.backends.cudnn.deterministic = config.experiment.torch_deterministic
 
-def get_git_hash() -> str:
-    try:
-        return (
-            subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode("ascii").strip()
-        )
-    except (subprocess.CalledProcessError, UnicodeDecodeError):
-        return "none"
+    # 5. Train - pass structured config directly
+    ppo_trainer = PPOTrainer(config, env, run_dir, run_name)
+    ppo_trainer.train()
 
 
 if __name__ == "__main__":
-    args = parse_args()
-
-    args.batch_size = args.num_envs * args.num_steps
-    args.minibatch_size = args.batch_size // args.num_minibatches
-    args.num_iterations = args.total_timesteps // args.batch_size
-
-    git_hash = get_git_hash()
-    run_name = f"{args.exp_name}__seed_{args.seed}__{git_hash}__{int(time.time())}"
-
-    if args.run_dir is None:
-        run_dir = f"runs/{run_name}"
-    else:
-        run_dir = args.run_dir
-
-    os.makedirs(run_dir, exist_ok=True)
-
-    # Initialize Global Logger
-    logger = UnifiedLogger(
-        config=vars(args),
-        project_name=args.wandb_project_name,  # or default PPO-Modularity if missing
-        run_name=run_name,
-        base_dir=os.path.dirname(run_dir),
-        use_wandb=args.track,
-    )
-
-    print_config(args, title="PPO Training Configuration")
-
-    env = make_env(args.env_config_path, args.num_envs)
-
-    torch.backends.cudnn.deterministic = args.torch_deterministic
-
-    ppo_trainer = PPOTrainer(args, env, run_dir, run_name)
-    ppo_trainer.train()
+    register_configs()
+    main()

src/brittle_star_project/MLPs/mlps.py

@@ -58,6 +58,10 @@ class Storage:
     returns: jnp.array
     rewards: jnp.array
 
+    raw_actions: jnp.ndarray = None  # before clipping
+    means: jnp.ndarray = None  # policy mean
+    stds: jnp.ndarray = None  # policy std
+
     def replace(self, **kwargs) -> "Storage":
         fs = fields(self)
         return Storage(**{f.name: kwargs.get(f.name, getattr(self, f.name)) for f in fs})

src/brittle_star_project/configs/config_architecture.py

@@ -0,0 +1,66 @@
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+
+@dataclass
+class LayerConfig:
+    hidden_dims: List[int] = field(default_factory=lambda: [64, 64])
+    activation: str = "tanh"
+
+
+@dataclass
+class ArchitectureConfig:
+    """Base class for actor-critic network configurations.
+
+    Both centralized and decentralized architectures share a centralized critic
+    composed of a feature extractor followed by a shallow output layer.
+
+    See docs/design/actor-critic.md for the full design rationale.
+    """
+
+    name: str = "base"
+
+    # Actor pipeline
+    sensor: Optional[LayerConfig] = None
+    propagator: Optional[LayerConfig] = None
+    motor: Optional[LayerConfig] = None
+
+    # Critic pipeline
+    feature_extractor: Optional[LayerConfig] = None
+    critic: Optional[LayerConfig] = None
+
+    # Decentralized
+    message_passing_steps: Optional[int] = None
+    topology_type: Optional[str] = None  # Supported values: "ring", "fully_connected"
+
+
+@dataclass
+class CentralizedConfig(ArchitectureConfig):
+    """Centralized actor-critic architecture (baseline).
+
+    The actor is a single global policy composed of a sensor (input network)
+    and a motor (output network). The sensor receives the full concatenated
+    global observation; the motor projects the hidden state to all joint actions.
+
+    See docs/design/actor-critic.md for the full design rationale.
+    """
+
+    name: str = "centralized"
+
+
+@dataclass
+class DecentralizedConfig(ArchitectureConfig):
+    """Decentralized actor architecture (NerveNet-MLP variant).
+
+    Each node runs a local sensor, exchanges messages with neighbours via a
+    propagator for a fixed number of steps, and then a local motor produces
+    the joint offset for that node only.
+
+    The critic remains centralized (shared with the base class): it receives the
+    full concatenated global observation and outputs a single scalar.
+
+    See docs/design/actor-critic.md and docs/design/communication.md for the
+    full design rationale.
+    """
+
+    name: str = "decentralized"

src/brittle_star_project/configs/config_experiment.py

@@ -0,0 +1,11 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class ExperimentConfig:
+    exp_name: str = "brittle_star_ppo"
+    seed: int = 1
+    torch_deterministic: bool = True
+    cuda: bool = True
+    debug_sanity: bool = False
+    base_run_dir: str = "runs"

src/brittle_star_project/configs/config_ppo.py

@@ -0,0 +1,22 @@
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class PPOConfig:
+    learning_rate: float = 2.5e-4
+    total_timesteps: int = 10000000
+    num_envs: int = 100
+    num_steps: int = 128
+    anneal_lr: bool = True
+    gamma: float = 0.99
+    gae_lambda: float = 0.95
+    num_minibatches: int = 4
+    update_epochs: int = 4
+    norm_adv: bool = True
+    clip_coef: float = 0.1
+    clip_vloss: bool = True
+    ent_coef: float = 0.01
+    vf_coef: float = 0.5
+    max_grad_norm: float = 0.5
+    target_kl: Optional[float] = None

src/brittle_star_project/configs/config_simulation.py

@@ -0,0 +1,12 @@
+from dataclasses import dataclass
+from typing import Optional
+from brittle_star_project.environment.env_types import Backend
+
+
+@dataclass
+class SimulationSettings:
+    """Settings for the simulation script."""
+
+    model_path: Optional[str] = None
+    model_type: str = "random"
+    backend: Backend = Backend.MJX

src/brittle_star_project/configs/main_config.py

@@ -0,0 +1,28 @@
+from dataclasses import dataclass, field
+
+from experiment_logger.config_logger import LoggingConfig
+from brittle_star_project.configs.config_experiment import ExperimentConfig
+from brittle_star_project.configs.config_ppo import PPOConfig
+from brittle_star_project.configs.config_architecture import ArchitectureConfig
+from brittle_star_project.configs.config_simulation import SimulationSettings
+from brittle_star_project.environment.env_config import MorphologyConfig, ArenaConfig, EnvConfig
+
+
+@dataclass
+class BrittleStarConfig:
+    """Root configuration for a brittle star training run.
+
+    Composed of strictly separated sub-configs. Each sub-config can be swapped
+    independently via CLI or a different YAML file. See configs/README.md.
+    """
+
+    experiment: ExperimentConfig = field(default_factory=ExperimentConfig)
+    logging: LoggingConfig = field(default_factory=LoggingConfig)
+    ppo: PPOConfig = field(default_factory=PPOConfig)
+    # This field is polymorphic; defaults to the base class to allow subclasses
+    # (CentralizedConfig, DecentralizedConfig) to be merged in via Hydra.
+    architecture: ArchitectureConfig = field(default_factory=ArchitectureConfig)
+    morphology: MorphologyConfig = field(default_factory=MorphologyConfig)
+    arena: ArenaConfig = field(default_factory=ArenaConfig)
+    environment: EnvConfig = field(default_factory=EnvConfig)
+    simulation: SimulationSettings = field(default_factory=SimulationSettings)

src/brittle_star_project/configs/register_configs.py

@@ -0,0 +1,40 @@
+from hydra.core.config_store import ConfigStore
+
+from experiment_logger.config_logger import LoggingConfig
+from brittle_star_project.configs.config_experiment import ExperimentConfig
+from brittle_star_project.configs.config_ppo import PPOConfig
+from brittle_star_project.configs.config_architecture import (
+    CentralizedConfig,
+    DecentralizedConfig,
+)
+from brittle_star_project.configs.config_simulation import SimulationSettings
+from brittle_star_project.environment.env_config import MorphologyConfig, ArenaConfig, EnvConfig
+from brittle_star_project.configs.main_config import BrittleStarConfig
+
+
+def register_configs() -> None:
+    """Register all dataclasses with Hydra's ConfigStore.
+
+    This must be called before hydra.main() processes the config, ensuring
+    every structured config is validated against its Python schema. Typos in
+    YAML keys will raise ConfigAttributeError at startup.
+    """
+    cs = ConfigStore.instance()
+
+    # Root schema
+    cs.store(name="brittle_star_config", node=BrittleStarConfig)
+
+    # Sub-config groups — each group corresponds to a configs/ subdirectory.
+    cs.store(group="experiment", name="base_experiment", node=ExperimentConfig)
+    cs.store(group="logging", name="base_logging", node=LoggingConfig)
+    cs.store(group="ppo", name="base_ppo", node=PPOConfig)
+
+    # Architecture variants — swap via CLI: architecture=decentralized
+    cs.store(group="architecture", name="centralized_schema", node=CentralizedConfig)
+    cs.store(group="architecture", name="decentralized_schema", node=DecentralizedConfig)
+
+    # Environment configs
+    cs.store(group="morphology", name="base_morphology", node=MorphologyConfig)
+    cs.store(group="arena", name="base_arena", node=ArenaConfig)
+    cs.store(group="environment", name="base_environment", node=EnvConfig)
+    cs.store(group="simulation", name="base_simulation", node=SimulationSettings)

src/brittle_star_project/dataclasses/PPOArgs.py

@@ -1,116 +0,0 @@
-from dataclasses import dataclass
-
-import jax
-
-
-@jax.tree_util.register_dataclass
-@dataclass
-class PPOArgs:
-    """
-    source: https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/ppo_atari_envpool_xla_jax_scan.py
-    """
-
-    # path to environment config file, if None, use default config
-    env_config_path: str | None = None
-
-    # path to hyperparameter config file (yaml), if None, use default config
-    hyperparameter_config_path: str | None = None
-
-    # the name of this experiment
-    exp_name: str = "brittle_star_ppo"
-
-    # the directory to save the experiment results
-    run_dir: str | None = None
-
-    # seed of the experiment
-    seed: int = 1
-
-    # if toggled, `torch.backends.cudnn.deterministic=False`
-    torch_deterministic: bool = True
-
-    # if toggled, cuda will be enabled by default
-    cuda: bool = True
-
-    # if toggled, this experiment will be tracked with Weights and Biases
-    track: bool = False
-
-    # the wandb's project name
-    wandb_project_name: str = "PPO-Modularity"
-
-    # the entity (team) of wandb's project
-    wandb_entity: str | None = "SEL3-2026-Groep-4"
-
-    # whether to capture videos of the agent performances (check out `videos` folder)
-    capture_video: bool = False
-
-    # whether to save model into the `runs/{run_name}` folder
-    save_model: bool = True
-
-    # checkpoint frequency (in iterations, 0 = no intermediate checkpoints)
-    checkpoint_frequency: int = 100
-
-    # whether to upload the saved model to huggingface
-    upload_model: bool = False
-
-    # the user or org name of the model repository from the Hugging Face Hub
-    hf_entity: str = ""
-
-    # ==== Algorithm specific dataclasses ====
-
-    # total timesteps of the experiments
-    total_timesteps: int = 10000000
-
-    # the learning rate of the optimizer
-    learning_rate: float = 2.5e-4
-
-    # the number of parallel game environments
-    num_envs: int = 100
-
-    # the number of steps to run in each environment per policy rollout
-    num_steps: int = 128
-
-    # Toggle learning rate annealing for policy and value networks
-    anneal_lr: bool = True
-
-    # the discount factor gamma
-    gamma: float = 0.99
-
-    # the lambda for the general advantage estimation
-    gae_lambda: float = 0.95
-
-    # the number of mini-batches
-    num_minibatches: int = 4
-
-    # the K epochs to update the policy
-    update_epochs: int = 4
-
-    # Toggles advantages normalization
-    norm_adv: bool = True
-
-    # the surrogate clipping coefficient
-    clip_coef: float = 0.1
-
-    # Toggles whether or not to use a clipped loss for the value function, as per the paper.
-    clip_vloss: bool = True
-
-    # coefficient of the entropy
-    ent_coef: float = 0.01
-
-    # coefficient of the value function
-    vf_coef: float = 0.5
-
-    # the maximum norm for the gradient clipping
-    max_grad_norm: float = 0.5
-
-    # the target KL divergence threshold
-    target_kl: float | None = None
-
-    # ==== to be filled in runtime ====
-    # the batch size (computed in runtime)
-    batch_size: int = 0
-
-    # the mini-batch size (computed in runtime)
-    minibatch_size: int = 0
-
-    # the number of iterations (computed in runtime)
-    num_iterations: int = 0

src/brittle_star_project/dataclasses/__init__.py

@@ -1,8 +1,6 @@
-from .PPOArgs import PPOArgs
 from .EpisodeStatistics import EpisodeStatistics
 
 
 __all__ = [
-    "PPOArgs",
     "EpisodeStatistics",
 ]

src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py

@@ -1,14 +1,11 @@
 import jax
 import jax.numpy as jnp
 
-from brittle_star_project import (
-    EnvConfig,
-    BrittleStarEnvFactory,
-    MorphologyConfig,
-    ArenaConfig,
-    Backend,
-)
-from brittle_star_project.environment import from_file
+from experiment_logger import get_logger
+from .env_config import EnvConfig, MorphologyConfig, ArenaConfig
+from .env_types import Backend
+from .factory import BrittleStarEnvFactory
+from .padded_obs_wrapper import compute_padding_masks, pad_observations_batched
 
 
 class BrittleStarJaxEnvWrapper:
@@ -29,14 +26,15 @@ class BrittleStarJaxEnvWrapper:
             self._backend, self._morphology, self._arena, self._env_config
         )
 
+        # Pre-compute masks for observation padding
+        self._padding_masks = compute_padding_masks(self._morphology.segments_per_arm)
+
         self._vectorized_reset = jax.jit(jax.vmap(self._env.reset))
         self._vectorized_step = jax.jit(jax.vmap(self._env.step))
         self._vectorized_action_sample = jax.jit(jax.vmap(self._env.action_space.sample))
 
         self._action_rng = None
 
-        from experiment_logger import get_logger
-
         self.logger = get_logger()
         self.logger.info(
             f"Initialized BrittleStarJaxEnvWrapper with {num_envs} envs on {backend.value}"
@@ -62,7 +60,12 @@ class BrittleStarJaxEnvWrapper:
         self.logger.info(f"Resetting vectorized environment environments with seed {seed}")
         self._action_rng, env_rng = jax.random.split(jax.random.PRNGKey(seed), 2)
         env_rngs = jnp.array(jax.random.split(env_rng, self._num_envs))
-        return self._vectorized_reset(rng=env_rngs)
+        state = self._vectorized_reset(rng=env_rngs)
+
+        state = state.replace(
+            observations=pad_observations_batched(state.observations, self._padding_masks)
+        )
+        return state
 
     def sample_actions(self):
         assert self._action_rng is not None, "Call reset() before sample_actions()"
@@ -72,7 +75,12 @@ class BrittleStarJaxEnvWrapper:
         return self._vectorized_action_sample(rng=jnp.array(sub_rngs))
 
     def step(self, state, action):
-        return self._vectorized_step(state=state, action=action)
+        next_state = self._vectorized_step(state=state, action=action)
+
+        next_state = next_state.replace(
+            observations=pad_observations_batched(next_state.observations, self._padding_masks)
+        )
+        return next_state
 
     def close(self):
         self._env.close()
@@ -86,15 +94,6 @@ class BrittleStarJaxEnvWrapper:
             morphology, arena, env_config, num_envs=num_envs, backend=backend
         )
 
-    @staticmethod
-    def from_config(
-        config_path: str, num_envs: int, backend: Backend = Backend.MJX
-    ) -> "BrittleStarJaxEnvWrapper":
-        morphology_cfg, arena_cfg, env_cfg = from_file(config_path)
-        return BrittleStarJaxEnvWrapper(
-            morphology_cfg, arena_cfg, env_cfg, num_envs=num_envs, backend=backend
-        )
-
     def __str__(self):
         morphology_str = str(self._morphology)
         arena_str = str(self._arena)

src/brittle_star_project/environment/__init__.py

@@ -1,4 +1,4 @@
-from .env_config import ArenaConfig, EnvConfig, MorphologyConfig, from_file
+from .env_config import ArenaConfig, EnvConfig, MorphologyConfig
 from .env_types import Backend, Task
 from .env_wrapper import BrittleStarEnv, StepResult
 from .factory import BrittleStarEnvFactory
@@ -12,5 +12,4 @@ __all__ = [
     "BrittleStarEnv",
     "StepResult",
     "BrittleStarEnvFactory",
-    "from_file",
 ]

src/brittle_star_project/environment/env_config.py

@@ -5,24 +5,37 @@ from dataclasses import dataclass, field
 from .env_types import Task
 
 
-@dataclass(frozen=True, slots=True)
+@dataclass
 class MorphologyConfig:
-    num_arms: int = 5
-    num_segments_per_arm: int = 4
+    """Brittle star morphology configuration.
+
+    segments_per_arm defines the number of segments for each arm. The length of
+    this list implicitly sets the number of arms. Use 0 segments to represent
+    a fully amputated arm (e.g., [4, 0, 4, 2, 4] for a 5-arm morphology with
+    arm 1 removed and arm 3 shortened).
+
+    The upstream biorobot library natively supports per-arm segment counts.
+    """
+
+    segments_per_arm: list[int] = field(default_factory=lambda: [4, 4, 4, 4, 4])
     use_p_control: bool = True
     use_torque_control: bool = False
 
+    @property
+    def num_arms(self) -> int:
+        return len(self.segments_per_arm)
 
-@dataclass(frozen=True, slots=True)
+
+@dataclass
 class ArenaConfig:
-    size: tuple[float, float] = (10.0, 5.0)
+    size: list[float] = field(default_factory=lambda: [10.0, 5.0])
     sand_ground_color: bool = True
     attach_target: bool = True
     wall_height: float = 1.5
     wall_thickness: float = 0.1
 
 
-@dataclass(frozen=True, slots=True)
+@dataclass
 class EnvConfig:
     """Shared environment settings.
 
@@ -31,13 +44,13 @@ class EnvConfig:
 
     task: Task = Task.DIRECTED_LOCOMOTION
 
-    simulation_time: float = 5.0
+    simulation_time: float = 10000.0
     num_physics_steps_per_control_step: int = 10
     time_scale: int = 2
 
     camera_ids: list[int] = field(default_factory=lambda: [0, 1])
     # (height, width)
-    render_size: tuple[int, int] = (480, 640)
+    render_size: list[int] = field(default_factory=lambda: [480, 640])
 
     joint_randomization_noise_scale: float = 0.0
 
@@ -47,16 +60,3 @@ class EnvConfig:
     # Light escape
     # Per docs in upstream env config: integer factors of 200.
     light_perlin_noise_scale: int = 0
-
-
-def from_file(path: str) -> tuple[MorphologyConfig, ArenaConfig, EnvConfig]:
-    """Load configurations from a YAML file."""
-    import yaml
-
-    with open(path, "r") as f:
-        config_dict = yaml.safe_load(f)
-
-        morphology = MorphologyConfig(**config_dict.get("morphology", {}))
-        arena = ArenaConfig(**config_dict.get("arena", {}))
-        env = EnvConfig(**config_dict.get("env", {}))
-        return morphology, arena, env

src/brittle_star_project/environment/factory.py

@@ -22,7 +22,7 @@ class BrittleStarEnvFactory:
 
         spec = default_brittle_star_morphology_specification(
             num_arms=config.num_arms,
-            num_segments_per_arm=config.num_segments_per_arm,
+            num_segments_per_arm=list(config.segments_per_arm),
             use_p_control=config.use_p_control,
             use_torque_control=config.use_torque_control,
         )

src/brittle_star_project/environment/padded_obs_wrapper.py

@@ -0,0 +1,108 @@
+"""Observation padding wrapper for amputated brittle star morphologies.
+
+When using a centralized controller, the global observation vector must remain
+a constant size regardless of how many segments are amputated. This wrapper pads
+the observation dictionary values with zeros using spatial insertion so that the
+flattened observation maintains the correct physical mapping to the neural network.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+import jax.numpy as jnp
+
+# Observation keys whose size scales with the number of joints (2 per segment).
+_JOINT_SCALED_KEYS = frozenset(
+    {
+        "joint_position",
+        "joint_velocity",
+        "joint_actuator_force",
+        "actuator_force",
+    }
+)
+
+# Observation keys whose size scales with the number of segments (1 per segment).
+_SEGMENT_SCALED_KEYS = frozenset(
+    {
+        "segment_contact",
+    }
+)
+
+
+def compute_padding_masks(
+    segments_per_arm: tuple[int, ...],
+    reference_segments_per_arm: tuple[int, ...] = (4, 4, 4, 4, 4),
+) -> dict[str, Any]:
+    """Pre-compute boolean masks for spatial insertion of observations.
+
+    Args:
+        segments_per_arm: The current (possibly amputated) morphology.
+        reference_segments_per_arm: The full morphology that defines the expected size.
+
+    Returns:
+        A dict containing 1D boolean masks and target sizes.
+    """
+    if len(segments_per_arm) != len(reference_segments_per_arm):
+        raise ValueError(
+            f"Morphology mismatch: current has {len(segments_per_arm)} arms, "
+            f"but reference requires {len(reference_segments_per_arm)} arms."
+        )
+
+    mask_1x = []
+    mask_2x = []
+
+    for arm_idx, (actual, ref) in enumerate(zip(segments_per_arm, reference_segments_per_arm)):
+        if not (0 <= actual <= ref):
+            raise ValueError(
+                f"Invalid amputation at arm {arm_idx}: "
+                f"actual segments ({actual}) must be between 0 and reference ({ref})."
+            )
+        # 1x scaling (e.g., contacts: 1 value per segment)
+        # 1x scaling (e.g., contacts: 1 value per segment)
+        mask_1x.extend([True] * actual + [False] * (ref - actual))
+        # 2x scaling (e.g., joints: 2 values per segment)
+        mask_2x.extend([True] * (actual * 2) + [False] * ((ref - actual) * 2))
+
+    return {
+        "mask_1x": jnp.array(mask_1x, dtype=bool),
+        "mask_2x": jnp.array(mask_2x, dtype=bool),
+        "target_size_1x": sum(reference_segments_per_arm),
+        "target_size_2x": sum(reference_segments_per_arm) * 2,
+    }
+
+
+def pad_observation(
+    obs: dict[str, Any],
+    masks: dict[str, Any],
+) -> dict[str, Any]:
+    """Pad an observation dict using spatial insertion."""
+    padded = {}
+    for key, value in obs.items():
+        if key in _JOINT_SCALED_KEYS:
+            out = jnp.zeros(masks["target_size_2x"], dtype=value.dtype)
+            padded[key] = out.at[masks["mask_2x"]].set(value)
+        elif key in _SEGMENT_SCALED_KEYS:
+            out = jnp.zeros(masks["target_size_1x"], dtype=value.dtype)
+            padded[key] = out.at[masks["mask_1x"]].set(value)
+        else:
+            padded[key] = value
+    return padded
+
+
+def pad_observations_batched(
+    obs: dict[str, Any],
+    masks: dict[str, Any],
+) -> dict[str, Any]:
+    """Pad a batched observation dict (leading batch dimension) using spatial insertion."""
+    padded = {}
+    for key, value in obs.items():
+        batch_size = value.shape[0]
+        if key in _JOINT_SCALED_KEYS:
+            out = jnp.zeros((batch_size, masks["target_size_2x"]), dtype=value.dtype)
+            padded[key] = out.at[:, masks["mask_2x"]].set(value)
+        elif key in _SEGMENT_SCALED_KEYS:
+            out = jnp.zeros((batch_size, masks["target_size_1x"]), dtype=value.dtype)
+            padded[key] = out.at[:, masks["mask_1x"]].set(value)
+        else:
+            padded[key] = value
+    return padded

src/brittle_star_project/ppo.py

@@ -99,6 +99,7 @@ def get_action_and_value(
     hidden_critic = feature_extractor_apply(params["feature_extractor_params"], x)
     hidden_sensor = message_passer(hidden_sensor)
     mean, log_std = actor_apply(params["actor_params"], hidden_sensor)
+    log_std = jnp.clip(log_std, -5, 2)
     std = jnp.exp(log_std)
 
     logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
@@ -142,6 +143,7 @@ def ppo_loss(
     pg_loss1 = -mb_advantages * ratio
     pg_loss2 = -mb_advantages * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef)
     pg_loss = jnp.maximum(pg_loss1, pg_loss2).mean()
+
     v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean()
     entropy_loss = entropy.mean()
     loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef

src/brittle_star_project/trainers/PPOTrainer.py

@@ -13,7 +13,8 @@ from flax.training.train_state import TrainState
 
 from experiment_logger import get_logger
 
-from brittle_star_project.dataclasses import EpisodeStatistics, PPOArgs
+from brittle_star_project.configs.main_config import BrittleStarConfig
+from brittle_star_project.dataclasses import EpisodeStatistics
 from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
 from brittle_star_project.MLPs.mlps import (
     Actor,
@@ -24,6 +25,33 @@ from brittle_star_project.MLPs.mlps import (
 )
 from brittle_star_project.ppo import PPO
 
+# TODO: move to config
+_ALLOWED_OBS_KEYS = {
+    "joint_position",
+    "joint_velocity",
+    "joint_actuator_force",
+    "actuator_force",
+    "disk_position",
+    "disk_rotation",
+    "disk_linear_velocity",
+    "disk_angular_velocity",
+    "unit_xy_direction_to_target",
+    "xy_distance_to_target",
+}
+# TODO: clip scaled reward?
+
+
+@jax.jit
+def _get_xy_distance_to_target(obs_dict: dict) -> jnp.ndarray:
+    """Extract xy_distance_to_target for all environments."""
+    # obs_dict is a dict of arrays with leading batch dimension (num_envs, ...)
+    return obs_dict["xy_distance_to_target"].squeeze(-1)  # shape: (num_envs,)
+
+
+@jax.jit
+def _clip_action(action: jnp.ndarray, low: jnp.ndarray, high: jnp.ndarray) -> jnp.ndarray:
+    return jnp.clip(action, low, high)
+
 
 def _compute_explained_variance(values: jnp.ndarray, returns: jnp.ndarray) -> float:
     var_returns = jnp.var(returns)
@@ -37,11 +65,25 @@ def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, lear
     return learning_rate * frac
 
 
+@jax.jit
+def _normalize_obs(obs, mean, var, eps=1e-8):
+    return jnp.clip((obs - mean) / jnp.sqrt(var + eps), -10.0, 10.0)
+
+
 @jax.jit
 def _convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray:
-    return jax.vmap(lambda o: jnp.concatenate([v.flatten() for v in o.values() if v.size > 0]))(
-        obs_dict
-    )
+    """Convert the raw observation dict → flat array, filtering unwanted keys."""
+
+    def _filter_and_flatten(o: dict) -> jnp.ndarray:
+        values = []
+        for key in sorted(o.keys()):
+            if key in _ALLOWED_OBS_KEYS:  # TODO: NORMALIZATION or .. of observations??
+                v = o[key]
+                if v.size > 0:
+                    values.append(jnp.asarray(v).flatten())
+        return jnp.concatenate(values)
+
+    return jax.vmap(_filter_and_flatten)(obs_dict)
 
 
 def _get_action_and_value_noise(
@@ -52,6 +94,8 @@ def _get_action_and_value_noise(
     agent_state: TrainState,
     next_obs: jnp.ndarray,
     key: jax.random.PRNGKey,
+    action_low,
+    action_high,
 ):
     hidden = sensor.apply(agent_state.params["sensor_params"], next_obs)
     hidden_critic = feature_extractor.apply(
@@ -59,13 +103,16 @@ def _get_action_and_value_noise(
     )
 
     mean, log_std = actor.apply(agent_state.params["actor_params"], hidden)
+    log_std = jnp.clip(log_std, -5, 2)
     key, subkey = jax.random.split(key)
     noise = jax.random.normal(subkey, shape=mean.shape)
     std = jnp.exp(log_std)
-    action = mean + noise * std
-    logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
+    raw_action = mean + noise * std
+    clipped_action = _clip_action(raw_action, action_low, action_high)
+    logprob = -0.5 * (((raw_action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
     value = critic.apply(agent_state.params["critic_params"], hidden_critic)
-    return action, logprob, value.squeeze(-1), key
+
+    return clipped_action, raw_action, logprob, value.squeeze(-1), mean, std, key
 
 
 def _step_once(
@@ -76,23 +123,28 @@ def _step_once(
     feature_extractor: GenericDenseLayersWithActivation,
     actor: Actor,
     critic: OneDenseLayerMLP,
+    action_low,
+    action_high,
 ):
     agent_state, episode_stats, obs, done, key, env_state = carry
-    action, logprob, value, key = _get_action_and_value_noise(
-        sensor, feature_extractor, actor, critic, agent_state, obs, key
+    clipped_action, raw_action, logprob, value, mean, std, key = _get_action_and_value_noise(
+        sensor, feature_extractor, actor, critic, agent_state, obs, key, action_low, action_high
     )
 
     episode_stats, env_state, (next_obs, reward, next_done) = env_step_fn(
-        episode_stats, env_state, action
+        episode_stats, env_state, clipped_action
     )
 
     storage = Storage(
         obs=obs,
-        actions=action,
+        actions=raw_action,
+        raw_actions=raw_action,
         logprobs=logprob,
         dones=done,
         values=value,
         rewards=reward,
+        means=mean,
+        stds=std,
         returns=jnp.zeros_like(reward),
         advantages=jnp.zeros_like(reward),
     )
@@ -103,6 +155,8 @@ def _step_env_wrapped(episode_stats, env_state, action, env_step_fn):
     next_env_state = env_step_fn(env_state, action)
 
     reward = next_env_state.reward
+    reward *= 20000
+    reward = jnp.clip(reward, -10, 10)
     terminated = next_env_state.terminated
     truncated = next_env_state.truncated
     done = terminated | truncated
@@ -140,6 +194,8 @@ def _rollout_jit(
     feature_extractor: GenericDenseLayersWithActivation,
     actor: Actor,
     critic: OneDenseLayerMLP,
+    action_low,
+    action_high,
 ):
     (agent_state, episode_stats, next_obs, next_done, key, env_state), storage = jax.lax.scan(
         partial(
@@ -149,6 +205,8 @@ def _rollout_jit(
             actor=actor,
             critic=critic,
             env_step_fn=step_env_fn,
+            action_low=action_low,
+            action_high=action_high,
         ),
         (agent_state, episode_stats, next_obs, next_done, key, env_state),
         (),
@@ -191,7 +249,9 @@ def _compute_gae_jit(
         (dones[1:], values[1:], values[:-1], storage.rewards),
         reverse=True,
     )
-    return storage.replace(advantages=advantages, returns=advantages + storage.values)
+    returns = advantages + storage.values
+    advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
+    return storage.replace(advantages=advantages, returns=returns)
 
 
 @dataclass
@@ -210,14 +270,23 @@ class TrainingMeasurements:
 
 
 class PPOTrainer:
-    def __init__(self, args: PPOArgs, env: BrittleStarJaxEnvWrapper, run_dir: str, run_name: str):
-        self.args = args
+    def __init__(
+        self, cfg: BrittleStarConfig, env: BrittleStarJaxEnvWrapper, run_dir: str, run_name: str
+    ):
+        self.cfg = cfg
+        self.ppo = cfg.ppo
+        self.experiment = cfg.experiment
+        self.logging_cfg = cfg.logging
         self.env = env
         self.run_dir = run_dir
         self.run_name = run_name
         self.logger = get_logger()
 
-        self.key = jax.random.PRNGKey(args.seed)
+        # Derived runtime fields
+        self.batch_size = self.ppo.num_envs * self.ppo.num_steps
+        self.num_iterations = self.ppo.total_timesteps // self.batch_size
+
+        self.key = jax.random.PRNGKey(self.experiment.seed)
 
         self.sensor, self.feature_extractor, self.actor, self.critic = self._init_agent()
         self.sensor.apply = jax.jit(self.sensor.apply)
@@ -225,29 +294,34 @@ class PPOTrainer:
         self.actor.apply = jax.jit(self.actor.apply)
         self.critic.apply = jax.jit(self.critic.apply)
 
+        action_low = jnp.asarray(self.env.single_action_space.low, dtype=jnp.float32)
+        action_high = jnp.asarray(self.env.single_action_space.high, dtype=jnp.float32)
+
         self._rollout_jit = jax.jit(
             partial(
                 _rollout_jit,
-                max_steps=self.args.num_steps,
+                max_steps=self.ppo.num_steps,
                 step_env_fn=partial(_step_env_wrapped, env_step_fn=self.env.step),
                 sensor=self.sensor,
                 feature_extractor=self.feature_extractor,
                 actor=self.actor,
                 critic=self.critic,
+                action_low=action_low,
+                action_high=action_high,
             )
         )
         self._compute_gae_jit = jax.jit(
             partial(
                 _compute_gae_jit,
-                num_envs=self.args.num_envs,
-                gamma=self.args.gamma,
-                gae_lambda=self.args.gae_lambda,
+                num_envs=self.ppo.num_envs,
+                gamma=self.ppo.gamma,
+                gae_lambda=self.ppo.gae_lambda,
                 feature_extractor=self.feature_extractor,
                 critic=self.critic,
             )
         )
 
-        self._ppo = PPO(self.args, self.sensor, self.actor, self.critic, self.feature_extractor)
+        self._ppo = PPO(self.ppo, self.sensor, self.actor, self.critic, self.feature_extractor)
 
         self.agent_state = self._init_agent_state()
 
@@ -256,16 +330,16 @@ class PPOTrainer:
         self._init_random()
 
     def _init_random(self):
-        self.logger.info(f"[RANDOM]: Setting random seed to {self.args.seed}")
+        self.logger.info(f"[RANDOM]: Setting random seed to {self.experiment.seed}")
 
-        random.seed(self.args.seed)
-        np.random.seed(self.args.seed)
+        random.seed(self.experiment.seed)
+        np.random.seed(self.experiment.seed)
 
     def _init_agent(self):
         self.logger.info("[AGENT]: Initializing agent...")
 
-        sensor = GenericDenseLayersWithActivation()
-        feature_extractor = GenericDenseLayersWithActivation()
+        sensor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300])
+        feature_extractor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300])
         actor = Actor(action_dim=self.env.single_action_space.shape[0])
         critic = OneDenseLayerMLP()
         return sensor, feature_extractor, actor, critic
@@ -277,15 +351,11 @@ class PPOTrainer:
             self.key, 5
         )
 
-        sample_obs = jnp.concatenate(
-            [
-                v.flatten()
-                for v in self.env.single_observation_space.sample(
-                    rng=jax.random.PRNGKey(0)
-                ).values()
-                if v.size > 0
-            ]
-        )
+        dummy_reset = self.env.reset(seed=0)
+        sample_obs = _convert_obs_dict_to_array(dummy_reset.observations)[0]  # take first env
+        self.obs_mean = jnp.zeros((len(sample_obs),))
+        self.obs_var = jnp.ones((len(sample_obs),))
+        self.obs_count = 1e-4
         sensor_params = self.sensor.init(sensor_key, sample_obs)
         feature_extractor_params = self.feature_extractor.init(feature_extractor_key, sample_obs)
         actor_params = self.actor.init(actor_key, self.sensor.apply(sensor_params, sample_obs))
@@ -299,17 +369,17 @@ class PPOTrainer:
                 AgentParams(sensor_params, actor_params, critic_params, feature_extractor_params)
             ),
             tx=optax.chain(
-                optax.clip_by_global_norm(self.args.max_grad_norm),
+                optax.clip_by_global_norm(self.ppo.max_grad_norm),
                 optax.inject_hyperparams(optax.adam)(
                     learning_rate=partial(
                         _linear_schedule,
-                        minibatch_count=self.args.num_minibatches,
-                        update_epochs=self.args.update_epochs,
-                        num_iterations=self.args.num_iterations,
-                        learning_rate=self.args.learning_rate,
+                        minibatch_count=self.ppo.num_minibatches,
+                        update_epochs=self.ppo.update_epochs,
+                        num_iterations=self.num_iterations,
+                        learning_rate=self.ppo.learning_rate,
                     )
-                    if self.args.anneal_lr
-                    else self.args.learning_rate,
+                    if self.ppo.anneal_lr
+                    else self.ppo.learning_rate,
                     eps=1e-5,
                 ),
             ),
@@ -319,12 +389,31 @@ class PPOTrainer:
         self.logger.info("[EPISODE STATS]: Initializing episode stats...")
 
         return EpisodeStatistics(
-            episode_returns=jnp.zeros(self.args.num_envs, dtype=jnp.float32),
-            episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32),
-            returned_episode_returns=jnp.zeros(self.args.num_envs, jnp.float32),
-            returned_episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32),
+            episode_returns=jnp.zeros(self.ppo.num_envs, dtype=jnp.float32),
+            episode_lengths=jnp.zeros(self.ppo.num_envs, dtype=jnp.int32),
+            returned_episode_returns=jnp.zeros(self.ppo.num_envs, jnp.float32),
+            returned_episode_lengths=jnp.zeros(self.ppo.num_envs, dtype=jnp.int32),
         )
 
+    def _update_obs_stats(self, obs: jnp.ndarray):
+        batch_mean = jnp.mean(obs, axis=0)
+        batch_var = jnp.var(obs, axis=0)
+        batch_count = obs.shape[0]
+
+        delta = batch_mean - self.obs_mean
+        total_count = self.obs_count + batch_count
+
+        new_mean = self.obs_mean + delta * batch_count / total_count
+
+        m_a = self.obs_var * self.obs_count
+        m_b = batch_var * batch_count
+        M2 = m_a + m_b + delta**2 * self.obs_count * batch_count / total_count
+        new_var = M2 / total_count
+
+        self.obs_mean = new_mean
+        self.obs_var = new_var
+        self.obs_count = total_count
+
     def _rollout(self, env_state, next_obs, next_done) -> tuple[Any, ...]:
         return self._rollout_jit(
             self.agent_state,
@@ -350,7 +439,39 @@ class PPOTrainer:
         start_time,
         iteration_time_start,
         training_measurements,
+        storage,
+        next_obs,
+        xy_distance,
     ):
+        data = jax.device_get(
+            {
+                "rewards": storage.rewards[0],
+                "values": storage.values[0],
+                "returns": storage.returns[0],
+                "advantages": storage.advantages[0],
+                "actions": storage.actions[0],
+                "raw_actions": storage.raw_actions[0],
+                "means": storage.means[0],
+                "stds": storage.stds[0],
+                "logprobs": storage.logprobs[0],
+            }
+        )
+
+        storage_metrics = {
+            "rollout/env0/return_mean": float(np.mean(data["returns"])),
+            "rollout/env0/advantage_mean": float(np.mean(data["advantages"])),
+            "rollout/env0/value_mean": float(np.mean(data["values"])),
+            "rollout/env0/value_vs_return_diff": float(np.mean(data["values"] - data["returns"])),
+            "rollout/env0/reward_mean": float(np.mean(data["rewards"])),
+            "rollout/env0/mean_mean": float(np.mean(data["means"])),
+            "rollout/env0/logprob_mean": float(np.mean(data["logprobs"])),
+            "rollout/env0/action_mean": float(np.mean(data["actions"])),
+            "rollout/env0/raw_action_mean": float(np.mean(data["raw_actions"])),
+        }
+
+        for i in range(len(xy_distance)):
+            storage_metrics[f"env_data/env{i}_xy_dist_target"] = float(xy_distance[i])
+
         metrics = {
             "charts/avg_episodic_return": training_measurements.avg_episodic_return,
             "charts/avg_episodic_length": np.mean(
@@ -371,8 +492,9 @@ class PPOTrainer:
             "losses/loss": training_measurements.loss[-1, -1].item(),
             "charts/SPS": int(global_step / (time.time() - start_time)),
             "charts/SPS_update": int(
-                self.args.num_envs * self.args.num_steps / (time.time() - iteration_time_start)
+                self.ppo.num_envs * self.ppo.num_steps / (time.time() - iteration_time_start)
             ),
+            **storage_metrics,
         }
         self.logger.log(metrics, step=global_step)
 
@@ -443,6 +565,7 @@ class PPOTrainer:
                 avg_terminated_length=avg_terminated_length,
                 avg_truncated_length=avg_truncated_length,
             ),
+            storage,
         )
 
     def _close(self):
@@ -451,8 +574,14 @@ class PPOTrainer:
     def _save_model(self, model_path: str):
         self.logger.info("[SAVE]: Saving the final model...")
 
+        from dataclasses import asdict as _asdict
+
+        config_dict = {
+            "experiment": _asdict(self.experiment),
+            "ppo": _asdict(self.ppo),
+        }
         params = [
-            vars(self.args),
+            config_dict,
             [
                 self.agent_state.params["sensor_params"],
                 self.agent_state.params["actor_params"],
@@ -464,8 +593,7 @@ class PPOTrainer:
 
     def train(self):
         """
-        Train the PPO agent for a specified number of iterations
-        (passed through PPOArgs in constructor).
+        Train the PPO agent for a specified number of iterations.
         Closes the environment at the end of training.
         """
         self.logger.info(f"running name: {self.run_name}")
@@ -473,47 +601,58 @@ class PPOTrainer:
         self.logger.info("[TRAIN]: Resetting environment...")
         self.logger.log_non_interactive(f"Initial reset started: {time.ctime()}")
 
-        env_state = self.env.reset(seed=self.args.seed)
+        env_state = self.env.reset(seed=self.experiment.seed)
         next_obs = _convert_obs_dict_to_array(env_state.observations)
-        next_done = jnp.zeros(self.args.num_envs, dtype=jnp.bool_)
+        next_done = jnp.zeros(self.ppo.num_envs, dtype=jnp.bool_)
 
         self.logger.log_non_interactive(f"Initial reset completed: {time.ctime()}")
 
         global_step = 0
         start_time = time.time()
 
-        iter_bar = self.logger.progress_bar(range(1, self.args.num_iterations + 1))
+        iter_bar = self.logger.progress_bar(range(1, self.num_iterations + 1))
         for iteration in iter_bar:
             iteration_time_start = time.time()
 
-            env_state, next_obs, next_done, training_measurements = self._step(
+            env_state, next_obs, next_done, training_measurements, storage = self._step(
                 env_state, next_obs, next_done, iteration=iteration
             )
+            self._update_obs_stats(next_obs)
+            next_obs = _normalize_obs(next_obs, self.obs_mean, self.obs_var)
 
-            global_step += self.args.num_steps * self.args.num_envs
+            xy_distance = _get_xy_distance_to_target(env_state.observations)
+
+            global_step += self.ppo.num_steps * self.ppo.num_envs
             self._log(
                 global_step,
                 self.episode_stats,
                 start_time,
                 iteration_time_start,
                 training_measurements,
+                storage,
+                next_obs,
+                xy_distance,
             )
 
             sps = int(global_step / (time.time() - start_time))
-            remaining_steps = self.args.total_timesteps - global_step
+            remaining_steps = self.ppo.total_timesteps - global_step
             eta_seconds = int(remaining_steps / sps) if sps > 0 else 0
             eta_str = str(datetime.timedelta(seconds=eta_seconds))
 
             self.logger.log_non_interactive(
-                f"Iteration {iteration}/{self.args.num_iterations} | "
-                f"Step {global_step}/{self.args.total_timesteps} | "
+                f"Iteration {iteration}/{self.num_iterations} | "
+                f"Step {global_step}/{self.ppo.total_timesteps} | "
                 f"SPS {sps} | "
                 f"Return {training_measurements.avg_episodic_return:.4f} | "
                 f"ETA {eta_str}"
             )
 
-        if self.args.save_model:
-            model_path = f"{self.run_dir}/{self.args.exp_name}.cleanrl_model"
+            if getattr(self.cfg.experiment, "debug_sanity", False):
+                self.logger.info("\n[SANITY CHECK] Successfully completed 1 epoch")
+                break
+
+        if self.logging_cfg.save_model:
+            model_path = f"{self.run_dir}/{self.experiment.exp_name}.cleanrl_model"
             self._save_model(model_path=model_path)
 
         self._close()

src/experiment_logger/__init__.py

@@ -4,8 +4,8 @@ This package provides a unified interface for logging to multiple backends
 (WandB, disk, stdout) simultaneously, ensuring no data loss.
 """
 
-from experiment_logger.config_utils import load_yaml_config, merge_config_with_cli
-from experiment_logger.unified_logger import UnifiedLogger, get_logger
+from experiment_logger.config_utils import load_yaml_config
+from experiment_logger.unified_logger import UnifiedLogger, get_logger, init_logger
 from experiment_logger.simple_logger import SimpleLogger
 from experiment_logger.wandb_utils import finish_wandb, init_wandb
 
@@ -13,9 +13,9 @@ __all__ = [
     "UnifiedLogger",
     "SimpleLogger",
     "get_logger",
+    "init_logger",
     "init_wandb",
     "finish_wandb",
     "load_yaml_config",
-    "merge_config_with_cli",
 ]
 __version__ = "0.1.0"

src/experiment_logger/config_logger.py

@@ -0,0 +1,14 @@
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class LoggingConfig:
+    track: bool = False
+    wandb_project_name: str = "PPO-Modularity"
+    wandb_entity: Optional[str] = "SEL3-2026-Groep-4"
+    capture_video: bool = False
+    save_model: bool = True
+    checkpoint_frequency: int = 100
+    upload_model: bool = False
+    hf_entity: str = ""

src/experiment_logger/config_utils.py

@@ -1,15 +1,12 @@
 """Configuration utilities for loading YAML configs and merging with CLI args."""
 
 import os
-import sys
 from typing import Dict, Any, Type, TypeVar
 import yaml
 from dataclasses import fields, is_dataclass
 
 from experiment_logger.unified_logger import get_logger
 
-log = get_logger()
-
 T = TypeVar("T")
 
 
@@ -24,7 +21,7 @@ def load_yaml_config(config_path: str) -> Dict[str, Any]:
     if config is None:
         return {}
 
-    log.info(f"Loaded configuration from: {config_path}")
+    get_logger().info(f"Loaded configuration from: {config_path}")
     return config
 
 
@@ -35,7 +32,7 @@ def save_yaml_config(config: Dict[str, Any], config_path: str):
     with open(config_path, "w") as f:
         yaml.dump(config, f, default_flow_style=False, indent=2, sort_keys=False)
 
-    log.info(f"Saved configuration to: {config_path}")
+    get_logger().info(f"Saved configuration to: {config_path}")
 
 
 def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T:
@@ -66,82 +63,21 @@ def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T:
                     else:
                         filtered_config[key] = field.type(value) if value is not None else None  # type: ignore
             except (ValueError, TypeError) as e:
-                log.warning(f"Could not convert {key}={value} to {field.type}: {e}")
+                get_logger().warning(f"Could not convert {key}={value} to {field.type}: {e}")
                 filtered_config[key] = value
         else:
-            log.warning(f"Unknown configuration parameter: {key}")
+            get_logger().warning(f"Unknown configuration parameter: {key}")
 
     return cls(**filtered_config)
 
 
-def merge_config_with_cli(config_class: Type[T], config_file: str | None = None) -> T:
-    """Merge YAML config with CLI arguments, with CLI taking precedence.
-
-    Args:
-        config_class: Dataclass type to create
-        config_file: Path to YAML config file (optional)
-
-    Returns:
-        Instance of config_class with merged configuration
-    """
-    # Parse CLI args first to get the default/CLI values
-    import tyro
-
-    # Check if --config is in sys.argv and extract it
-    extracted_config_file = config_file
-    if "--config" in sys.argv:
-        config_idx = sys.argv.index("--config")
-        if config_idx + 1 < len(sys.argv):
-            extracted_config_file = sys.argv[config_idx + 1]
-            # Remove from sys.argv so tyro doesn't see it
-            sys.argv.pop(config_idx)  # Remove --config
-            sys.argv.pop(config_idx)  # Remove config file path
-
-    # Load YAML config if available
-    yaml_config = {}
-    if extracted_config_file and os.path.exists(extracted_config_file):
-        yaml_config = load_yaml_config(extracted_config_file)
-        log.info(f"Merging YAML config from {extracted_config_file} with CLI args")
-    elif extracted_config_file:
-        log.warning(f"Config file not found: {extracted_config_file}, using CLI args only")
-
-    # Create default instance to know what the defaults are
-    default_instance = config_class()
-    default_dict = {f.name: getattr(default_instance, f.name) for f in fields(config_class)}  # type: ignore
-
-    # Parse CLI args
-    cli_instance = tyro.cli(config_class)
-    cli_dict = {f.name: getattr(cli_instance, f.name) for f in fields(config_class)}  # type: ignore
-
-    # Merge configs: YAML as base, CLI overrides non-default values
-    final_config = {}
-
-    for field in fields(config_class):  # type: ignore
-        field_name = field.name
-        default_value = default_dict[field_name]
-        yaml_value = yaml_config.get(field_name, default_value)
-        cli_value = cli_dict[field_name]
-
-        # Use CLI value if it's different from default, otherwise use YAML value
-        if cli_value != default_value:
-            final_config[field_name] = cli_value
-            if yaml_value != default_value and yaml_value != cli_value:
-                log.info(f"CLI override: {field_name}={cli_value} (YAML had {yaml_value})")
-        else:
-            final_config[field_name] = yaml_value
-            if yaml_value != default_value:
-                log.info(f"YAML config: {field_name}={yaml_value}")
-
-    return config_class(**final_config)
-
-
 def print_config(config: Any, title: str = "Configuration"):
     """Pretty print configuration."""
-    log.info(f"{title}:")
+    get_logger().info(f"{title}:")
     if is_dataclass(config):
         for field in fields(config):
             value = getattr(config, field.name)
-            log.info(f"  {field.name}: {value}")
+            get_logger().info(f"  {field.name}: {value}")
     else:
         for key, value in vars(config).items():
-            log.info(f"  {key}: {value}")
+            get_logger().info(f"  {key}: {value}")

src/experiment_logger/unified_logger.py

@@ -6,9 +6,7 @@ This logger ensures all experimental data is preserved by writing to:
 3. stdout (for real-time monitoring)
 """
 
-import datetime
 import logging
-import subprocess
 import yaml
 import sys
 import time
@@ -21,38 +19,67 @@ import numpy as np
 
 from experiment_logger.wandb_utils import finish_wandb, init_wandb
 
-# Global singleton storage
-_global_logger = None
+# Global storage for the active logger and the proxy singleton
+_active_logger: Optional[Any] = None
+_proxy_instance: Optional["LoggerProxy"] = None
 
 
-def get_logger() -> "UnifiedLogger":
-    """Retrieve the global UnifiedLogger. If not initialized, fallback to auto-initialization."""
-    global _global_logger
-    if _global_logger is None:
-        try:
-            commit_hash = (
-                subprocess.check_output(
-                    ["git", "rev-parse", "--short", "HEAD"], stderr=subprocess.STDOUT
-                )
-                .decode("utf-8")
-                .strip()
-            )
-        except Exception:
-            commit_hash = "unknown"
+def get_logger() -> "LoggerProxy":
+    """Retrieve the global LoggerProxy.
 
-        timestamp_str = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-        generic_name = f"{timestamp_str}_{commit_hash}_brittle_star"
+    This should be used for all logging calls. It returns a proxy that
+    delegates to the active logger (defaulting to a SimpleLogger until
+    init_logger is called).
+    """
+    global _proxy_instance, _active_logger
+    if _proxy_instance is None:
+        if _active_logger is None:
+            # Fallback to SimpleLogger to avoid premature directory creation
+            from experiment_logger.simple_logger import SimpleLogger
 
-        # Initialize generic fallback logger without WandB
-        _global_logger = UnifiedLogger(
-            run_name=generic_name,
-            config={"auto_initialized": True},
-            use_wandb=False,
-            _set_as_global=False,  # Prevent recursive call inside __init__
-        )
-        _global_logger.warning(f"UnifiedLogger auto-initialized with name: {generic_name}")
+            _active_logger = SimpleLogger(run_name="pre_init")
 
-    return _global_logger
+        _proxy_instance = LoggerProxy()
+
+    return _proxy_instance
+
+
+def init_logger(**kwargs) -> "UnifiedLogger":
+    """Initialize the full UnifiedLogger and set it as the active logger.
+
+    This should be called once the configuration is ready. It will create
+    the output directories and set up all logging backends.
+    """
+    global _active_logger
+    logger = UnifiedLogger(**kwargs)
+    _active_logger = logger
+    return logger
+
+
+class LoggerProxy:
+    """Proxy that delegates all method calls to the active logger instance.
+
+    This allows the logger to be swapped out (e.g., from a SimpleLogger to
+    a UnifiedLogger) without any clients needing to update their references.
+    """
+
+    def _get_logger(self) -> Any:
+        global _active_logger
+        if _active_logger is None:
+            # This shouldn't normally happen since get_logger handles it
+            from experiment_logger.simple_logger import SimpleLogger
+
+            _active_logger = SimpleLogger(run_name="pre_init_fallback")
+        return _active_logger
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._get_logger(), name)
+
+    def __enter__(self):
+        return self._get_logger().__enter__()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return self._get_logger().__exit__(exc_type, exc_val, exc_tb)
 
 
 class UnifiedLogger:
@@ -68,7 +95,6 @@ class UnifiedLogger:
         use_wandb: bool = True,
         save_code: bool = True,
         log_level: int = logging.INFO,
-        _set_as_global: bool = True,
     ):
         """Initialize the unified logger.
 
@@ -80,7 +106,6 @@ class UnifiedLogger:
             base_dir: Base directory for local storage
             use_wandb: Whether to use WandB logging
             save_code: Whether to save code to WandB
-            _set_as_global: Internal flag to override the global singleton
         """
         self.run_name = run_name
         self.config = config
@@ -119,11 +144,6 @@ class UnifiedLogger:
             self._text_logger.addHandler(fh)
             self._text_logger.addHandler(ch)
 
-        # Set as global singleton
-        global _global_logger
-        if _set_as_global:
-            _global_logger = self
-
         # Save config to disk
         self._save_config()
 

tests/test_config.py

@@ -1,38 +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")
-
-    def test_merge_config_with_cli_is_callable(self):
-        from experiment_logger.config_utils import merge_config_with_cli
-
-        assert callable(merge_config_with_cli)

tests/test_configs.py

@@ -0,0 +1,48 @@
+from pathlib import Path
+from hydra import compose, initialize_config_dir
+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."""
+    config_dir = str(Path(__file__).parent.parent / "configs")
+    with initialize_config_dir(version_base="1.3", config_dir=config_dir):
+        # We compose the config; it follows main_config.yaml
+        cfg = compose(config_name="main_config", overrides=["architecture=centralized"])
+
+        # Merge with the structured schema and convert to a real dataclass instance
+        structured_cfg = OmegaConf.to_object(
+            OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), cfg)
+        )
+
+        # Basic assertions
+        assert structured_cfg.architecture.name == "centralized"
+        assert structured_cfg.architecture.propagator is None
+        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."""
+    config_dir = str(Path(__file__).parent.parent / "configs")
+    with initialize_config_dir(version_base="1.3", config_dir=config_dir):
+        cfg = compose(config_name="main_config", overrides=["architecture=decentralized"])
+
+        # Merge and convert to dataclass instance
+        structured_cfg = OmegaConf.to_object(
+            OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), cfg)
+        )
+
+        # Basic assertions
+        assert structured_cfg.architecture.name == "decentralized"
+        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,65 @@
+import pytest
+import os
+import sys
+
+# CRITICAL for headless cross-platform testing (devcontainers etc)
+if sys.platform == "linux" and "DISPLAY" not in os.environ and "WAYLAND_DISPLAY" not in os.environ:
+    os.environ.setdefault("MUJOCO_GL", "egl")
+
+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
+
+
+@pytest.mark.skipif(os.getenv("CI") == "true", reason="No OpenGL display in CI")
+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()

uv.lock

@@ -18,6 +18,7 @@ dependencies = [
     { name = "evosax" },
     { name = "flax" },
     { name = "gymnasium" },
+    { name = "hydra-core" },
     { name = "ipykernel" },
     { name = "jax" },
     { name = "matplotlib" },
@@ -30,7 +31,6 @@ dependencies = [
     { name = "pyopengl-accelerate" },
     { name = "pyyaml" },
     { name = "torch" },
-    { name = "tyro" },
     { name = "wandb" },
     { name = "warp-lang" },
 ]
@@ -57,6 +57,7 @@ requires-dist = [
     { name = "evosax", specifier = "==0.2.0" },
     { name = "flax", specifier = ">=0.12.2" },
     { name = "gymnasium", specifier = ">=1.2.3" },
+    { name = "hydra-core", specifier = ">=1.3.2" },
     { name = "ipykernel", specifier = "==7.2.0" },
     { name = "jax", specifier = "==0.9.0.1" },
     { name = "jax", extras = ["cuda13"], marker = "extra == 'cuda'", specifier = "==0.9.0.1" },
@@ -71,7 +72,6 @@ requires-dist = [
     { name = "pyyaml", specifier = ">=6.0" },
     { name = "tensorboard", marker = "extra == 'analysis'" },
     { name = "torch", specifier = ">=2.4.0" },
-    { name = "tyro", specifier = ">=1.0.10" },
     { name = "wandb", specifier = "==0.24.2" },
     { name = "warp-lang" },
 ]
@@ -111,6 +111,12 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" },
 ]
 
+[[package]]
+name = "antlr4-python3-runtime"
+version = "4.9.3"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/3e/38/7859ff46355f76f8d19459005ca000b6e7012f2f1ca597746cbcd1fbfe5e/antlr4-python3-runtime-4.9.3.tar.gz", hash = "sha256:f224469b4168294902bb1efa80a8bf7855f24c99aef99cbefc1bcd3cce77881b", size = 117034, upload-time = "2021-11-06T17:52:23.524Z" }
+
 [[package]]
 name = "appnope"
 version = "0.1.4"
@@ -492,15 +498,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ea/97/4f78412f73a9350bc8f934441bae5b68b102c8f4240a7f06b4114b51d6de/dm_tree-0.1.9-cp312-cp312-win_amd64.whl", hash = "sha256:9020a5ce256fcc83aa4bc190cc96dd66e87685db0a6e501b0c06aa492c2e38fc", size = 102022, upload-time = "2025-01-30T20:45:28.701Z" },
 ]
 
-[[package]]
-name = "docstring-parser"
-version = "0.17.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/b2/9d/c3b43da9515bd270df0f80548d9944e389870713cc1fe2b8fb35fe2bcefd/docstring_parser-0.17.0.tar.gz", hash = "sha256:583de4a309722b3315439bb31d64ba3eebada841f2e2cee23b99df001434c912", size = 27442, upload-time = "2025-07-21T07:35:01.868Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/55/e2/2537ebcff11c1ee1ff17d8d0b6f4db75873e3b0fb32c2d4a2ee31ecb310a/docstring_parser-0.17.0-py3-none-any.whl", hash = "sha256:cf2569abd23dce8099b300f9b4fa8191e9582dda731fd533daf54c4551658708", size = 36896, upload-time = "2025-07-21T07:35:00.684Z" },
-]
-
 [[package]]
 name = "dotmap"
 version = "1.3.30"
@@ -753,6 +750,20 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/c5/7b/bca5613a0c3b542420cf92bd5e5fb8ebd5435ce1011a091f66bb7693285e/humanize-4.15.0-py3-none-any.whl", hash = "sha256:b1186eb9f5a9749cd9cb8565aee77919dd7c8d076161cf44d70e59e3301e1769", size = 132203, upload-time = "2025-12-20T20:16:11.67Z" },
 ]
 
+[[package]]
+name = "hydra-core"
+version = "1.3.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "antlr4-python3-runtime" },
+    { name = "omegaconf" },
+    { name = "packaging" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/6d/8e/07e42bc434a847154083b315779b0a81d567154504624e181caf2c71cd98/hydra-core-1.3.2.tar.gz", hash = "sha256:8a878ed67216997c3e9d88a8e72e7b4767e81af37afb4ea3334b269a4390a824", size = 3263494, upload-time = "2023-02-23T18:33:43.03Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/c6/50/e0edd38dcd63fb26a8547f13d28f7a008bc4a3fd4eb4ff030673f22ad41a/hydra_core-1.3.2-py3-none-any.whl", hash = "sha256:fa0238a9e31df3373b35b0bfb672c34cc92718d21f81311d8996a16de1141d8b", size = 154547, upload-time = "2023-02-23T18:33:40.801Z" },
+]
+
 [[package]]
 name = "identify"
 version = "2.6.18"
@@ -1531,6 +1542,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ec/11/3f1ee9dce24b41812dd572a037c4436d4d21f759fbe373cc271b0ce98805/nvidia_nvvm-13.2.51-py3-none-win_amd64.whl", hash = "sha256:a4809baaa5429eabe1878853761ce31f0ba15216e2348710b7898dc591f5fc14", size = 56751075, upload-time = "2026-03-09T10:11:09.994Z" },
 ]
 
+[[package]]
+name = "omegaconf"
+version = "2.3.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "antlr4-python3-runtime" },
+    { name = "pyyaml" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/09/48/6388f1bb9da707110532cb70ec4d2822858ddfb44f1cdf1233c20a80ea4b/omegaconf-2.3.0.tar.gz", hash = "sha256:d5d4b6d29955cc50ad50c46dc269bcd92c6e00f5f90d23ab5fee7bfca4ba4cc7", size = 3298120, upload-time = "2022-12-08T20:59:22.753Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/e3/94/1843518e420fa3ed6919835845df698c7e27e183cb997394e4a670973a65/omegaconf-2.3.0-py3-none-any.whl", hash = "sha256:7b4df175cdb08ba400f45cae3bdcae7ba8365db4d165fc65fd04b050ab63b46b", size = 79500, upload-time = "2022-12-08T20:59:19.686Z" },
+]
+
 [[package]]
 name = "opencv-python"
 version = "4.13.0.92"
@@ -2315,18 +2339,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74caf5e34b66d9f3a429af689c1c7128daba1d8208df60e81106b115c00d6fca", size = 188266850, upload-time = "2026-01-20T16:00:43.041Z" },
 ]
 
-[[package]]
-name = "typeguard"
-version = "4.5.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "typing-extensions" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/2b/e8/66e25efcc18542d58706ce4e50415710593721aae26e794ab1dec34fb66f/typeguard-4.5.1.tar.gz", hash = "sha256:f6f8ecbbc819c9bc749983cc67c02391e16a9b43b8b27f15dc70ed7c4a007274", size = 80121, upload-time = "2026-02-19T16:09:03.392Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/91/88/b55b3117287a8540b76dbdd87733808d4d01c8067a3b339408c250bb3600/typeguard-4.5.1-py3-none-any.whl", hash = "sha256:44d2bf329d49a244110a090b55f5f91aa82d9a9834ebfd30bcc73651e4a8cc40", size = 36745, upload-time = "2026-02-19T16:09:01.6Z" },
-]
-
 [[package]]
 name = "typing-extensions"
 version = "4.15.0"
@@ -2361,20 +2373,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" },
 ]
 
-[[package]]
-name = "tyro"
-version = "1.0.10"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "docstring-parser" },
-    { name = "typeguard" },
-    { name = "typing-extensions" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/30/c1/0a5850badd3f18373d6a0366091638674cec6780b558c1c5b846adea938b/tyro-1.0.10.tar.gz", hash = "sha256:2822eacac963a4922bf7eafe3b156a1f0f7fe8e34148202987581224f25565c2", size = 481084, upload-time = "2026-03-18T08:24:17.307Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/05/be/a0b4c9fa64999a2e337cbefcdedd2e101e8dd88a84e4fa497bd0e4531dc1/tyro-1.0.10-py3-none-any.whl", hash = "sha256:8de87a3a40c8a91f10831f8f0638cd0eed00f0e4de9cd3d561e967f407477210", size = 183433, upload-time = "2026-03-18T08:24:16.012Z" },
-]
-
 [[package]]
 name = "tzdata"
 version = "2025.3"