Merge branch 'dev' into docfix

.gitignore

@@ -5,6 +5,7 @@ wandb/
 outputs/
 multirun/
 metrics/
+adjacency_debug.txt
 
 # Python-generated files
 __pycache__/

README.md

@@ -13,44 +13,37 @@ To set up the UV module, you can run the following command:
 uv sync --frozen
 ```
 
-### Configuration
+## Repository Structure
 
-1. **Copy the default configuration:**
+```text
+.
+├── configs/                # Hydra configuration files (YAML)
+├── docs/                   # Comprehensive documentation and API guides
+├── runs/                   # Default output directory for Hydra and training artifacts
+├── scripts/                # High-level entrypoints for training, simulation, and evaluation
+├── src/
+│   └── brittle_star_project/ # Core library and environment logic
+│       ├── evaluation/     # Checkpoint evaluation, rollout logic, and metrics persistence
+│       └── trainers/       # Training loop implementations (e.g., PPO)
+└── tests/                  # Unit and integration tests
+```
+
+## Usage
+
+For detailed instructions on how to use the project, please refer to the **[API Documentation](docs/README.md)**.
+
+### Quick Start
+
+1. **Train a model:**
    ```bash
-   cp configs/default_ppo.yaml configs/my_experiment.yaml
+   uv run python scripts/train.py ppo.learning_rate=0.001 logging.track=true
    ```
 
-2. **Edit `configs/my_experiment.yaml`** to set your WandB credentials:
-   ```yaml
-   track: true  # Enable WandB logging
-   wandb_entity: "your-wandb-username"  # Replace with your username/team
-   wandb_project_name: "PPO-Modularity"
-   ```
+2. **Monitor progress:**
+   See [Tracking & Monitoring](docs/api/tracking.md).
 
-3. **(Optional) Login to WandB:**
-   ```bash
-   uv run wandb login
-   ```
-
-### Training
-
-example command:
-
-```bash
-uv run python scripts/train.py
-```
-
-Or use a custom config file:
-
-```bash
-uv run python scripts/train.py --config configs/my_experiment.yaml
-```
-
-Override specific parameters:
-
-```bash
-uv run python scripts/train.py --learning-rate 0.001 --num-envs 32 --track
-```
+3. **Simulate a trained model:**
+   See [Simulation & Evaluation](docs/api/simulation.md).
 
 ## HPC
 

configs/architecture/decentralized.yaml

@@ -26,7 +26,7 @@ critic:
   activation: "tanh"
 
 # Synchronous message-passing rounds per control step
-message_passing_steps: 1
+message_passing_steps: 4
 
 # Connectivity topology (e.g., ring, fully_connected)
-topology_type: "ring"
+topology_type: "fully_connected"

configs/centralized-final.yaml

@@ -0,0 +1,71 @@
+# Custom Main Configuration
+#
+# Use with:
+#   uv run python scripts/train.py --config-name main_config_custom
+#
+# This keeps the project defaults intact while giving you a single custom
+# training entrypoint you can edit freely.
+
+defaults:
+  - brittle_star_config
+  - experiment: base
+  - logging: default
+  - evaluation: default
+  - ppo: default
+  - architecture: centralized
+  - morphology: 5_arms_full
+  - arena: default
+  - environment: directed_locomotion
+  - simulation: default
+  - _self_
+
+morphology:
+  morph_mode: CENTRALIZED
+
+experiment:
+  exp_name: "final-models/centralized/"
+  seed: 42
+  torch_deterministic: true
+  cuda: true
+
+logging:
+  track: true
+  save_model: true
+  save_checkpoints: true
+  upload_final_model: true
+  upload_checkpoints: true
+  checkpoint_frequency: 20
+  wandb_project_name: "final-models"
+
+evaluation:
+  evaluate_checkpoints: true
+  eval_max_steps: 2000
+  eval_seed: 0
+
+ppo:
+  learning_rate: 0.0001
+  total_timesteps: 16384000
+  num_envs: 128
+  num_steps: 64
+  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.001
+  vf_coef: 1.0
+  max_grad_norm: 0.5
+  target_kl: 0.02
+
+environment:
+  simulation_time: 100000.0
+  target_distance: 3.0
+
+hydra:
+  job:
+    chdir: true
+  run:
+    dir: ${experiment.base_run_dir}/${experiment.exp_name}/${now:%Y-%m-%d}/${now:%H-%M-%S}

configs/environment/dir_loc_further.yaml

@@ -0,0 +1,2 @@
+simulation_time: 50000.0
+target_distance: 3.0

configs/environment/directed_locomotion.yaml

@@ -2,11 +2,11 @@
 # Baseline task setting.
 
 task: DIRECTED_LOCOMOTION
-simulation_time: 5000.0
+simulation_time: 100000.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: 0.6
+target_distance: 3.0
 light_perlin_noise_scale: 0

configs/environment/light_escape.yaml

@@ -2,7 +2,7 @@
 # Advanced task requiring movement away from light source.
 
 task: LIGHT_ESCAPE
-simulation_time: 5.0
+simulation_time: 100000.0
 num_physics_steps_per_control_step: 10
 time_scale: 2
 camera_ids: [0, 1]

configs/evaluation/default.yaml

@@ -0,0 +1,8 @@
+# Default Evaluation Configuration
+# Settings used for checkpoint evaluation during training.
+
+evaluate_checkpoints: false
+# Max number of control steps during evaluation rollout.
+eval_max_steps: 2000
+# Seed for deterministic evaluation reset.
+eval_seed: 0

configs/evaluation/poster.yaml

@@ -0,0 +1,23 @@
+# @package evaluation
+# Configuration for the models used in the poster comparison.
+
+# Standard evaluation settings
+evaluate_checkpoints: false
+eval_max_steps: 5000
+eval_seed: 0
+
+# Cross-model comparison settings
+# We use 10 episodes to get a more robust average for the final poster results.
+comparison_base_seed: 0
+comparison_num_episodes: 2
+comparison_output_csv: "runs/evaluation/comparison.csv"
+
+# Paths to the .cleanrl_model files to be compared (relative to workspace root).
+comparison_models:
+  - "runs/input-space-2-arms/2026-05-02/08-14-58/final_model.flax"
+
+# Path to the morphologies to evaluate against.
+comparison_morphologies:
+  - "configs/morphology/5_arms_full.yaml"
+  - "configs/morphology/3_arms.yaml"
+  - "configs/morphology/2_arms.yaml"

configs/experiment/long_2arm.yaml

@@ -0,0 +1,6 @@
+# Testing chicken dinner 4 but further distance.
+
+exp_name: "long2arm"
+seed: 123
+torch_deterministic: true
+cuda: true

configs/fully-connected-final.yaml

@@ -0,0 +1,74 @@
+# Custom Main Configuration
+#
+# Use with:
+#   uv run python scripts/train.py --config-name main_config_custom
+#
+# This keeps the project defaults intact while giving you a single custom
+# training entrypoint you can edit freely.
+
+defaults:
+  - brittle_star_config
+  - experiment: base
+  - logging: default
+  - evaluation: default
+  - ppo: default
+  - architecture: decentralized
+  - morphology: 5_arms_full
+  - arena: default
+  - environment: directed_locomotion
+  - simulation: default
+  - _self_
+
+architecture:
+  topology_type: "fully_connected"
+
+morphology:
+  morph_mode: FULLY_CONNECTED
+
+experiment:
+  exp_name: "final-models/fully-connected/"
+  seed: 42
+  torch_deterministic: true
+  cuda: true
+
+logging:
+  track: true
+  save_model: true
+  save_checkpoints: true
+  upload_final_model: true
+  upload_checkpoints: true
+  checkpoint_frequency: 20
+  wandb_project_name: "final-models"
+
+evaluation:
+  evaluate_checkpoints: true
+  eval_max_steps: 2000
+  eval_seed: 0
+
+ppo:
+  learning_rate: 0.0001
+  total_timesteps: 16384000
+  num_envs: 128
+  num_steps: 64
+  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.001
+  vf_coef: 1.0
+  max_grad_norm: 0.5
+  target_kl: 0.02
+
+environment:
+  simulation_time: 100000.0
+  target_distance: 3.0
+
+hydra:
+  job:
+    chdir: true
+  run:
+    dir: ${experiment.base_run_dir}/${experiment.exp_name}/${now:%Y-%m-%d}/${now:%H-%M-%S}

configs/logging/default.yaml

@@ -10,4 +10,4 @@ save_checkpoints: true
 checkpoint_frequency: 100
 upload_final_model: false
 upload_checkpoints: false
-hf_entity: ""
+hf_entity: ""

configs/main_config.yaml

@@ -6,11 +6,13 @@ defaults:
   - brittle_star_config
   - experiment: base
   - logging: default
+  - evaluation: default
   - ppo: default
   - architecture: centralized
   - morphology: 5_arms_full
   - arena: default
   - environment: directed_locomotion
+  - obs_bounds: default
   - simulation: default
   - _self_
 

configs/morphology/2_arms_decentralized.yaml

@@ -0,0 +1,6 @@
+# 2 Arms Morphology Configuration
+
+segments_per_arm: [4, 0, 4, 0, 0]
+use_p_control: true
+use_torque_control: false
+morph_mode: FULLY_CONNECTED

configs/morphology/5_arms_damaged.yaml

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

configs/morphology/5_arms_full_fullconnected.yaml

@@ -0,0 +1,7 @@
+# 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
+morph_mode: FULLY_CONNECTED

configs/obs_bounds/default.yaml

@@ -0,0 +1 @@
+# Defaults provided by dataclass

configs/ppo/chickendinnerwinner.yaml

@@ -0,0 +1,16 @@
+anneal_lr: true
+clip_coef: 0.2
+clip_vloss: true
+ent_coef: 0.001
+gae_lambda: 0.95
+gamma: 0.99
+learning_rate: 0.0001
+max_grad_norm: 0.5
+norm_adv: true
+num_envs: 32
+num_minibatches: 32
+num_steps: 64
+target_kl: 0.02
+total_timesteps: 12288000
+update_epochs: 4
+vf_coef: 1.0

configs/ppo/smoke_test.yaml

@@ -2,9 +2,9 @@
 # Lower timestep count for quick iterations/testing.
 
 learning_rate: 0.0005
-total_timesteps: 65536
-num_envs: 512
-num_steps: 128
+total_timesteps: 1024
+num_envs: 32
+num_steps: 32
 anneal_lr: true
 gamma: 0.99
 gae_lambda: 0.95

configs/ring-final.yaml

@@ -0,0 +1,74 @@
+# Custom Main Configuration
+#
+# Use with:
+#   uv run python scripts/train.py --config-name main_config_custom
+#
+# This keeps the project defaults intact while giving you a single custom
+# training entrypoint you can edit freely.
+
+defaults:
+  - brittle_star_config
+  - experiment: base
+  - logging: default
+  - evaluation: default
+  - ppo: default
+  - architecture: decentralized
+  - morphology: 5_arms_full
+  - arena: default
+  - environment: directed_locomotion
+  - simulation: default
+  - _self_
+
+architecture:
+  topology_type: "ring"
+
+morphology:
+  morph_mode: RING
+
+experiment:
+  exp_name: "final-models/ring/"
+  seed: 42
+  torch_deterministic: true
+  cuda: true
+
+logging:
+  track: true
+  save_model: true
+  save_checkpoints: true
+  upload_final_model: true
+  upload_checkpoints: true
+  checkpoint_frequency: 20
+  wandb_project_name: "final-models"
+
+evaluation:
+  evaluate_checkpoints: true
+  eval_max_steps: 2000
+  eval_seed: 0
+
+ppo:
+  learning_rate: 0.0001
+  total_timesteps: 16384000
+  num_envs: 128
+  num_steps: 64
+  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.001
+  vf_coef: 1.0
+  max_grad_norm: 0.5
+  target_kl: 0.02
+
+environment:
+  simulation_time: 100000.0
+  target_distance: 3.0
+
+hydra:
+  job:
+    chdir: true
+  run:
+    dir: ${experiment.base_run_dir}/${experiment.exp_name}/${now:%Y-%m-%d}/${now:%H-%M-%S}

configs/simulation/default.yaml

@@ -9,7 +9,18 @@ headless: false
 # In headless mode this is required; in viewer mode null means "infinite".
 max_steps: null
 
-# Optional: path to the Hydra config.yaml used during training.
-# When set, scripts/simulate.py will use it to default morphology/arena/environment/architecture
-# to match training (unless you explicitly override those keys via CLI).
-trained_config_path: null
+# Optional: override morphology for amputation experiments.
+# Points to a morphology config YAML file (e.g. configs/morphology/3_arms.yaml).
+# If null, the training morphology from the model's metadata is used.
+morphology_override: null
+
+# Video recording (requires [evaluation] extra)
+record_video: false
+# When null, video is saved in a per-model evaluation folder alongside the model.
+video_output_path: null
+# Camera ID to use for video recording (1 is usually the close-up camera)
+camera_id: 1
+
+# Optional override for the metadata YAML file path.
+# If null, the script looks for `<model_name>_metadata.yaml` alongside the model_path.
+metadata_path: null

docs/DEVELOPMENT.md

@@ -62,22 +62,9 @@ In the devcontainer, this will succeed on both CPU and GPU. A `GpuDevice` is exp
 
 ## Logging & Monitoring
 
-This project uses a unified logging system through the `experiment_logger` package. For a full API reference, see the [package README](../src/experiment_logger/README.md).
-
-### Quick Setup
-
-1.  **Authorization**: Export your API key in your terminal to enable WandB synchronization:
-    ```bash
-    export WANDB_API_KEY=your_copied_api_key_here
-    ```
-2.  **Toggle Tracking**: Use the `--track` flag in `scripts/train.py` to enable online sync.
-3.  **Local Monitoring**: All runs are recorded in the `runs/` directory. View scalars with TensorBoard:
-    ```bash
-    tensorboard --logdir runs/
-    ```
-
-### Environment Awareness
-
-The logger automatically detects if it is running in an interactive terminal or a non-interactive environment (like an HPC Slurm job). It will automatically disable progress bars and switch to robust fallback modes (offline logging) to ensure your experiments never hang.
+This project uses a unified logging system through the `experiment_logger` package. 
 
+- **Usage in Code**: To use the logger in your scripts, refer to the [package README](../src/experiment_logger/README.md) for the API reference.
+- **WandB/TensorBoard Setup**: For information on how to configure tracking for experiments, see the [Tracking & Monitoring API Guide](./api/tracking.md).
 
+The logger automatically detects if it is running in an interactive terminal or a non-interactive environment (like an HPC Slurm job), adjusting progress bars and fallback modes accordingly.

docs/README.md

@@ -2,6 +2,8 @@
 
 ## Design & architecture (`/design`)
 
+If you are interested in the "why did you do it like this?"
+
 - [Actor/critic architecture](./design/actor-critic.md): Description of the actor-critic pipeline.
 - [Communication](./design/communication.md): Message propagation, Nerve-Net style.
 - [Controllers](./design/controllers.md): Macroscopig brain toplogy, centralized, arm-level, segment-level.
@@ -11,5 +13,9 @@
 
 ## API reference (`/api`)
 
-- [Environment](./api/environment.md): MuJoCo environment interaction, state retrieval, and configuration.
-- [Simulate](./api/simulate.md): Simulation rendering.
+If you are interested in the "how do I use it?"
+
+- [Training](./api/training.md): How to configure and run experiments.
+- [Tracking & Monitoring](./api/tracking.md): Setting up WandB and TensorBoard to monitor runs.
+- [Simulation](./api/simulation.md): Visualizing and evaluating models.
+- [Environment](./api/environment.md): MuJoCo environment interaction and configuration.

docs/api/analysis.md

@@ -0,0 +1,84 @@
+# Analysis & Plotting Tools
+
+This guide outlines the tools available for analyzing experimental data and generating poster-quality visualizations for the Brittle Star project.
+
+## Shared Configuration
+
+All plotting scripts share a central configuration in `scripts/plots/plot_config.py`. This file defines:
+- **Color Palette:** A color-blind friendly, high-contrast palette for different architectures.
+- **Typography:** Consistent font sizes and styles tailored for A0 posters.
+- **Markers:** Shared visual indicators, such as the ★ used for best performers.
+
+## Comparison Visualization
+
+The `scripts/plots/analyze_comparisons.py` script generates grouped bar charts comparing the performance of different architectures across various morphologies.
+
+### Usage
+
+Run the script from the root of the project, providing the path to your evaluation CSV:
+
+```bash
+# Basic usage (saves PNG and SVG to runs/evaluation/plots/)
+uv run python scripts/plots/analyze_comparisons.py path/to/results.csv
+
+# Advanced usage for Figma/Poster integration
+uv run python scripts/plots/analyze_comparisons.py path/to/results.csv \
+    --output_dir docs/assets/plots/ \
+    --font_size 30 \
+    --fig_width 14 \
+    --fig_height 10
+```
+
+### CLI Arguments
+
+- `input_csv`: (Required) Path to the CSV file containing evaluation results.
+- `--output_dir`, `-o`: Directory where plots will be saved (default: `runs/evaluation/plots`).
+- `--show_titles`: Include titles in the plots. Default is **False**, as titles are typically added natively in design tools like Figma.
+- `--font_size`: Base font size in points (default: 28).
+- `--fig_width` / `--fig_height`: Physical dimensions of the plot in inches. Match these to your Figma layout to maintain exact font sizes.
+
+### Outputs
+
+The script generates four key plots, each saved as both `.png` and `.svg`:
+1.  **Forward Velocity:** Grouped bar chart (cm/s).
+2.  **Accumulated Reward:** Mean cumulative reward.
+3.  **Success Rate:** Target acquisition percentage.
+4.  **Distance Remaining:** Navigational accuracy.
+
+---
+
+## Convergence Analysis
+
+The `scripts/plots/analyze_convergence.py` script determines the convergence point of training runs.
+
+### Usage
+
+```bash
+uv run python scripts/plots/analyze_convergence.py --output_dir runs/convergence/
+```
+
+### Configuration
+
+- **File Mapping:** The script uses hardcoded paths in the `FILE_MAPPING` dictionary. Update these paths to point to your specific run evaluation files.
+- **CLI Arguments:** Supports the same `--show_titles`, `--font_size`, and `--fig_width/height` flags as the comparison script.
+
+### Outputs
+
+Generates three plots (PNG & SVG):
+1. `convergence_comparison`: Grouped horizontal bar chart.
+2. `progress_reward_curves`: Line plots of reward over time.
+3. `progress_velocity_curves`: Line plots of velocity over time.
+
+---
+
+## Poster Integration (Figma)
+
+### SVG & Scaling
+We recommend using the **SVG** outputs for poster design in Figma:
+1. **No Resolution Loss:** SVGs are vector-based and will remain sharp at any size.
+2. **Native Text:** Text in the SVG imports as native text layers in Figma.
+3. **Exact Font Matching:** To ensure a `28pt` font in the plot matches a `28pt` font in your poster, set the `--fig_width` and `--fig_height` to match the physical dimensions of the plot box in your Figma layout.
+4. **Editable:** You can "Ungroup" the SVG in Figma to manually move labels, adjust colors, or tweak individual bars.
+
+### Image Placeholders
+The comparison charts include light-gray square placeholders below the X-axis. These are designed as guides; in Figma, you can drop your morphology renders or illustrations directly on top of these squares.

docs/api/evaluation.md

@@ -0,0 +1,56 @@
+# Checkpoint & Model Evaluation
+
+This guide covers how to evaluate trained brittle star models, with a focus on measuring defect tolerance (amputations) across different controller architectures.
+
+## Checkpoint Evaluation (During Training)
+
+The `PPOTrainer` can automatically evaluate every saved checkpoint using the fast MJX backend. This is enabled via configuration.
+
+### Configuration
+
+In your experiment config or via CLI:
+```bash
+python scripts/train.py evaluation.evaluate_checkpoints=true evaluation.eval_max_steps=5000
+```
+
+Results are saved to `runs/<run_dir>/metrics/checkpoint_evaluation.csv` and synced to Weights & Biases if enabled.
+
+## Cross-Model & Defect Tolerance Analysis
+
+To measure how well different controllers handle damage (amputations), use `scripts/compare_models.py`. This script performs a grid search over models x morphologies.
+
+1. Create or update a YAML file in `configs/evaluation`.
+2. Run the benchmark:
+
+```bash
+python scripts/compare_models.py evaluation=poster
+```
+
+The script will evaluate every combination of model and morphology for the specified number of episodes.
+
+The results are saved to a CSV (default: `metrics/model_comparison.csv`).
+
+### CSV Schema
+
+| Column                | Description                                                  |
+|-----------------------|--------------------------------------------------------------|
+| `model_path`          | Path to the trained weights.                                 |
+| `architecture`        | The `morph_mode` of the model (e.g., `CENTRALIZED`, `RING`). |
+| `arm_0` ... `arm_4`   | Number of segments in each arm slot (0 = amputated).         |
+| `num_active_arms`     | Total number of arms with segments > 0.                      |
+| `seed`                | The episode seed.                                            |
+| `eval_return`         | Accumulated shaped reward.                                   |
+| `approx_max_velocity` | Average velocity: `(initial_dist - final_dist) / steps`.     |
+| `reached_target`      | Whether the robot finished within the success radius.        |
+
+## Post-hoc Checkpoint Scanning
+
+If you need to re-evaluate every saved checkpoint in a run (e.g., to generate a learning curve with different metrics):
+
+```bash
+python scripts/evaluate_checkpoints.py \
+    simulation.model_path=runs/<run_id>/final_model.flax \
+    evaluation.eval_max_steps=2000
+```
+
+This script scans the `checkpoints/` directory of the specified run and evaluates every `.flax` file it finds using the model's training morphology.

docs/api/simulate.md

@@ -1,14 +0,0 @@
-# Training and Simulation for Brittle Star Models
-
-## Simulating a model
-
-In order to simulate and view the behavior of a trained model, you can use the `simulate.py` script. This script allows you to specify the path to a trained model and will launch a simulation using that model. This script has the following parameters:
-
-- `--model`: The path to the trained model artifact to simulate.
-- `--model-type`: The type of model to simulate (e.g., `random`, ...)
-- `--task`: The task to simulate (e.g., `directed_locomotion`, ...)
-- `--seed`: The random seed for reproducibility.
-
-```bash
-python simulate.py --model artifacts/my_model --model-type random --task directed_locomotion --seed 0
-```

docs/api/simulation.md

@@ -0,0 +1,41 @@
+# Simulation & Evaluation
+
+The simulation pipeline allows you to visualize trained models and evaluate their performance under various conditions.
+
+## Overview
+
+The simulation pipeline is metadata-driven. Training-specific configurations (morphology, arena, environment, etc.) are automatically loaded from the `_metadata.yaml` file associated with the model checkpoint.
+
+## Basic Simulation
+
+To simulate a model in the MuJoCo viewer:
+
+```bash
+uv run scripts/simulate.py simulation.model_path=runs/your_run/final_model.flax
+```
+
+## Amputation & Morphology Overrides
+
+You can test trained models on different morphologies (e.g., amputating legs) by providing a morphology override. The observations will be automatically padded up to the training morphology's dimensions:
+
+```bash
+uv run scripts/simulate.py \
+    simulation.model_path=runs/your_run/final_model.flax \
+    simulation.morphology_override=configs/morphology/3_arms.yaml
+```
+
+## Video Recording
+
+Recording videos requires the `[evaluation]` extra:
+
+```bash
+uv run scripts/simulate.py \
+    simulation.model_path=runs/your_run/final_model.flax \
+    simulation.record_video=true \
+    simulation.max_steps=1000
+```
+
+Videos and evaluation metadata are stored in timestamped folders alongside the model:
+`runs/your_run/final_model_evaluations/eval_<timestamp>/simulation.mp4`
+
+For batch evaluation and cross-model comparison, see the **[Evaluation Guide](./evaluation.md)**.

docs/api/tracking.md

@@ -0,0 +1,60 @@
+# Tracking & Monitoring
+
+This guide explains how to monitor your experiments using Weights & Biases (WandB) and TensorBoard.
+
+## Weights & Biases (WandB)
+
+WandB is used for online synchronization and visualization of training metrics.
+
+### Authorization
+
+Export your API key in your terminal to enable WandB synchronization:
+
+```bash
+export WANDB_API_KEY=your_copied_api_key_here
+```
+
+Alternatively, you can log in using the CLI:
+
+```bash
+uv run wandb login
+```
+
+### Enabling Tracking
+
+To enable online sync during a training run, set `logging.track=true` on the command line:
+
+```bash
+uv run python scripts/train.py logging.track=true
+```
+
+You can also configure your project and entity:
+
+```bash
+uv run python scripts/train.py \
+    logging.track=true \
+    logging.wandb_project_name="MyProject" \
+    logging.wandb_entity="my-team"
+```
+
+These can also be set in your configuration YAML file under the `logging` key.
+
+## Local Monitoring with TensorBoard
+
+All runs are recorded locally in the `runs/` directory (or the directory specified in `experiment.base_run_dir`). You can view scalars and other metrics with TensorBoard:
+
+```bash
+tensorboard --logdir runs/
+```
+
+Access the interface at `http://localhost:6006`.
+
+### CLI Exploration Tool
+
+For quick diagnostics or to export data to CSV without launching the full TensorBoard UI, you can use the `explore_tensorboard.py` script:
+
+```bash
+uv run python scripts/analysis/explore_tensorboard.py runs/your_run_name/
+```
+
+See the detailed description in [`/scripts/analysis/README.md`](../../scripts/analysis/README.md).

docs/api/training.md

@@ -0,0 +1,55 @@
+# Training Models
+
+This guide covers how to configure and run training experiments for the Brittle Star project using Hydra-based configurations.
+
+## Configuration
+
+The project uses a modular configuration system powered by [Hydra](https://hydra.cc/). Instead of passing many command-line flags, you select and override configuration groups.
+
+### Creating a Custom Experiment
+
+1. **Create a new experiment file:**
+   Create a file at `configs/experiment/my_experiment.yaml`. You can copy an existing one as a template:
+   ```bash
+   cp configs/experiment/base.yaml configs/experiment/my_experiment.yaml
+   ```
+
+2. **Edit `configs/experiment/my_experiment.yaml`** to set your experiment parameters:
+   ```yaml
+   # @package _global_
+   experiment:
+     exp_name: "my_custom_run"
+     seed: 42
+   ```
+
+## Training Execution
+
+To start a training run with the default settings defined in `configs/main_config.yaml`:
+
+```bash
+uv run python scripts/train.py
+```
+
+### Using a Custom Experiment Configuration
+
+To run with your custom experiment file:
+
+```bash
+uv run python scripts/train.py experiment=my_experiment
+```
+
+```bash
+uv run python scripts/train.py ppo.learning_rate=0.001 ppo.num_envs=32 logging.track=true
+```
+
+## Evaluation During Training
+
+By default, the trainer saves checkpoints but does not evaluate them. To enable automatic headless evaluation of every saved checkpoint, set `evaluation.evaluate_checkpoints=true`:
+
+```bash
+uv run python scripts/train.py evaluation.evaluate_checkpoints=true
+```
+
+For more details on evaluation metrics and comparison tools, see [Evaluation](./evaluation.md).
+
+For more details on tracking your experiments, see [Tracking & Monitoring](./tracking.md).

docs/design/input_action_spaces.md

@@ -5,12 +5,26 @@ space (outputs). The control models map these observations directly to physical
 
 **Inputs (state space)**
 
-The observation space provides the agent with its current physical state and its objective.
+The observation space provides the agent with its current physical state and its navigational objective. With a
+decentralized control architecture in mind, we divide these inputs into global and local states.
 
-- Joint positions: the current angles of all joints in the morphology.
-- Joint velocities: the current moving speed of the joints.
-- Goal vector: instad of just a scalar distance, the goal is represented asa a vector (distance and ange/direction) to
-  the target.
+Global inputs, always broadcasted to all nodes:
+
+- Vertical orientation/tilt: A single, simplified metric representing the tilt/vertical alignment of the agent's
+  central body/disk, a.k.a. the deviation from the global Z-axis. Its value is derived from the environment's raw disk
+  rotation 3D vector $[roll, pitch, yaw]$:
+  $$
+  tilt = sqrt(roll^2 + pitch^2)
+  $$
+- Goal vector: A 2D unit vector representing the *egocentric* direction to the target. A value of $[1.0, 0.0]$
+  indicates that the target is directly in front of the agent (angle 0).
+
+Local inputs, routed directly to specific nodes:
+
+- Joint positions: The current angles of all joints within the morphology.
+- Joint velocities: The current angular velocities of the joints.
+- Joint actuator forces: The physical forces currently exerted at each specific joint.
+- Segment contact: These values indicate whether each physical segment of the agent is currently touching the ground.
 
 **Outputs (action space)**
 
@@ -18,6 +32,16 @@ The action space defines how the agent interacts with the environment.
 
 - Joint offsets: *absolute* target positions (offsets) for the joints, i.e. the exact angle the joint should move to.
 
+## Normalization and Scaling
+
+Both the input (observation) and output (action) spaces are rescaled to the range **$[-1, 1]$**. 
+
+For the input space, all raw physical values (angles, velocities, forces, distances) are normalized based on their
+defined physical bounds. If a value exceeds these bounds during simulation, it is clipped to the $[-1, 1]$ range.
+
+For the output space, the neural network's tanh-activated outputs (which naturally fall in $[-1, 1]$) are linearly
+mapped to the physical joint limits defined in the robot's morphology.
+
 ## Rationale
 
 When designing the state space, we must ask: *Could a human operator perform this task given only these inputs?*
@@ -26,15 +50,50 @@ When designing the state space, we must ask: *Could a human operator perform thi
   providing only the joint position is insufficient to determine the direction a limb is currently moving. By
   explicitly including joint velocities, the agent can immediately infer momentum and movement direction without
   needing to memorize past states.
-- Goal Vector (Distance + Angle): Providing only the scalar "distance to the goal" as an input is akin to blindfolding
-  the robot and asking it to find a target by playing "hot or cold." By providing a full vector, the agent knows
-  exactly where the target is relative to its current orientation, allowing for directed and efficient locomotion.
 - Absolute Joint Offsets: The physical Brittle Star robot relies on servo motors (if we were to build this simulated
   robot), which are inherently position-controlled devices. (Continuous rotation servos exist, but they are less
   commonly used for joints.) If our network outputted continuous torques (forces), a significant portion of the
   reinforcement learning process would be wasted on learning low-level PID control dynamics (i.e., how much force to
   apply to hold a position). Abstracting this away forces the learning algorithm to focus entirely on higher-level gait
   generation and locomotion.
+- Simplified vertical orientation: We drop the full 3D spatial rotation and angular velocity arrays in favor of a
+  single vertical orientation metric (tilt). For a brittle star moving accross a flat plane, this metric is sufficient
+  for the agent to sense if it is losing balance or flipping over.
+- Force representation: We strictly retain the joint actuator forces and drop the more generic actuator force. Forces
+  that are explicitly tied to individual joints are significantly easier to route into decentralized, local limb nodes,
+  which is necessary for our message-passing architecture.
+- Goal Vector (Distance + Angle): Providing only the scalar "distance to the goal" as an input is akin to blindfolding
+  the robot and asking it to find a target by playing "hot or cold." By providing a full vector, the agent knows
+  exactly where the target is relative to its current orientation, allowing for directed and efficient locomotion.
+
+  The goal representation is explicitly divided into a directional vector and a scalar distance. Keeping the direction
+  as a normalized unit vector bounds the values to the $[-1, 1]$ range, which stabilizes neural network training.
+  Providing only a scalar "distance to the goal" would force the agent to learning localized searching behaviors (e.g.
+  random walks or spiraling) to deduce the direction, drastically increasing the difficulty of the learning task.
+
+  **NOTE:** We later dropped the "distance to vector", switching to only a direction as the input. Our reasoning is
+  the agent should always move towards the goal (it should not learn to stop at the goal), which allows for this
+  simplification that decreases the model input size.
+  
+  The environment provides a raw `unit_xy_direction_to_target` (global), which we transform into a calculated
+  `robot_direction_to_target` (egocentric) before passing it to the MLPs. This vector consists of the X and Y
+  direction, where a value of $[1.0, 0.0]$ (mapping to an angle of $0$) means the robot is facing directly towards the
+  target.
+- Contact sensors: Segment contact detects external ground interaction and is biologically vital for timing gait
+  transitions.
+- Zero-Centered Rescaling ($[-1, 1]$): Using a zero-centered range is standard best practice for continuous control
+  tasks. It provides several mathematical and physical advantages:
+  - Improved Gradient Flow: Neural networks optimize faster when inputs are zero-centered. If all inputs were positive
+    (e.g., $[0, 1]$), the gradients during backpropagation would be forced to the same sign, causing inefficient
+    "zig-zag" weight updates.
+  - Meaningful Neutral State: In robotics, $0.0$ naturally represents a resting state (zero velocity, centered
+    position, no force). In a $[-1, 1]$ system, this physical rest maps to a neutral $0.0$ signal in the network.
+    This also correctly communicaties a "neutral/dead" signal for amputated limbs that are padded with $0.0$ values.
+
+Specifically, we do not include some available inputs:
+
+- Global position: Absolute spatial coordinates can cause the agent to overfit to a specific coordinate frame or map,
+  rather than learning general, adaptable locomotion strategies.
 
 ## Limitations and alternatives
 
@@ -50,3 +109,56 @@ Alternative state and action formulations include:
 - Scalar Goal Distance: Giving the agent only the scalar distance to the target would force it to learn a localized
   searching behavior (e.g., spiraling or random walks) to determine the correct direction. While biologically plausible
   for simpler organisms following chemical gradients, it drastically increases the difficulty of the learning task.
+- **$[0, 1]$ Rescaling:** While some domains (like computer vision) use $[0, 1]$ scaling, it is generally avoided in
+  robotics. Scaling to $[0, 1]$ would mean that a resting joint (velocity = 0) maps to an input of $0.5$. This
+  constant positive bias forces the network to waste capacity learning to ignore or subtract this baseline signal just
+  to stand still. Furthermore, it breaks the "dead signal" interpretation of zero-padding used for amputations.
+
+## MuJoCo
+
+This is what the filtered input vectors look like in MuJoCo, with $J$ joints and $S$ segments:
+
+- `joint_position`: shape=(J,), dtype=float64
+- `joint_velocity`: shape=(J,), dtype=float64
+- `joint_actuator_force`: shape=(J,), dtype=float64
+- `segment_contact`: shape=(S,), dtype=float64
+- `robot_direction_to_target`: shape=(2,), dtype=float64, egocentric
+- `disk_z_tilt`: shape=(1,), dtype=float64, derived from `disk_rotation`
+
+This brings the entire input space down to $3J + S + 4$ float64's, compared to $4J + S + 15$ float64's for the
+unfiltered inputs.
+
+For reference, these are all the inputs that are available in the MuJoCo environment:
+
+```
+obs keys: ['joint_position', 'joint_velocity', 'joint_actuator_force', 'actuator_force', 'disk_position', 'disk_rotation', 'disk_linear_velocity', 'disk_angular_velocity', 'tendon_position', 'tendon_velocity', 'segment_contact', 'unit_xy_direction_to_target', 'xy_distance_to_target']
+
+raw observations dict:
+{'joint_position': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),
+ 'joint_velocity': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),
+ 'joint_actuator_force': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),
+ 'actuator_force': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),
+ 'disk_position': array([0.  , 0.  , 0.11]),
+ 'disk_rotation': (0.0, -0.0, 0.0),
+ 'disk_linear_velocity': array([0., 0., 0.]),
+ 'disk_angular_velocity': array([0., 0., 0.]),
+ 'tendon_position': array([], dtype=float64),
+ 'tendon_velocity': array([], dtype=float64),
+ 'segment_contact': array([0., 0., 0., 0., 0., 0.]),
+ 'unit_xy_direction_to_target': array([-0.95333378, -0.30191837]),
+ 'xy_distance_to_target': array([3.])}
+
+(shapes)
+joint_position: shape=(12,), dtype=float64, size=12
+joint_velocity: shape=(12,), dtype=float64, size=12
+joint_actuator_force: shape=(12,), dtype=float64, size=12
+actuator_force: shape=(12,), dtype=float64, size=12
+disk_position: shape=(3,), dtype=float64, size=3
+disk_rotation: shape=(3,), dtype=float64, size=3
+disk_linear_velocity: shape=(3,), dtype=float64, size=3
+disk_angular_velocity: shape=(3,), dtype=float64, size=3
+tendon_position: shape=(0,), dtype=float64, size=0
+tendon_velocity: shape=(0,), dtype=float64, size=0
+segment_contact: shape=(6,), dtype=float64, size=6
+xy_distance_to_target: shape=(1,), dtype=float64, size=1
+```

experiments/debug-experiment-10042026/used_variables.md

@@ -1,107 +0,0 @@
-## 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

@@ -38,6 +38,10 @@ cuda = [
 analysis = [
     "tensorboard",
 ]
+evaluation = [
+    "imageio>=2.35.0",
+    "imageio-ffmpeg>=0.5.1",
+]
 
 [dependency-groups]
 dev = [

scripts/compare_models.py

@@ -0,0 +1,182 @@
+"""Compare multiple trained policies across shared evaluation conditions.
+
+For each model listed in evaluation.comparison_models, this script runs
+`comparison_num_episodes` headless rollouts (seeded sequentially from
+`comparison_base_seed`) and writes a results CSV to `comparison_output_csv`.
+
+Results include two metrics per episode:
+- `eval_return`    — shaped reward (same function used during training)
+- `max_velocity`  — approximated as initial_xy_dist / steps taken
+
+Usage:
+    # With the default evaluation config
+    python scripts/compare_models.py evaluation=poster
+
+    # Override the output path on the fly
+    python scripts/compare_models.py evaluation=poster \\
+        evaluation.comparison_output_csv=metrics/quick_comparison.csv
+"""
+
+from __future__ import annotations
+
+import csv
+import logging
+import time
+from pathlib import Path
+
+import hydra
+from omegaconf import DictConfig, OmegaConf
+
+from brittle_star_project.configs.main_config import BrittleStarConfig
+from brittle_star_project.configs.register_configs import register_configs
+from brittle_star_project.evaluation import build_eval_env
+from brittle_star_project.evaluation.checkpoint import load_metadata, metadata_to_configs
+from brittle_star_project.evaluation.rollout import rollout_headless
+
+_FIELDNAMES = [
+    "model_path",
+    "architecture",
+    "arm_0",
+    "arm_1",
+    "arm_2",
+    "arm_3",
+    "arm_4",
+    "num_active_arms",
+    "seed",
+    "reached_target",
+    "episode_length",
+    "eval_return",
+    "initial_target_distance",
+    "final_xy_dist",
+    "approx_max_velocity",
+]
+
+
+def _approx_max_velocity(result) -> float | None:
+    """Approximate max velocity as distance covered per step.
+
+    This is a rough upper bound: (initial_dist - final_dist) / steps.
+    """
+    if result.initial_target_distance is None or result.final_xy_dist is None or result.length <= 0:
+        return None
+    dist_covered = result.initial_target_distance - result.final_xy_dist
+    return dist_covered / result.length
+
+
+@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
+def main(dict_cfg: DictConfig) -> None:
+    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
+    logger = logging.getLogger(__name__)
+
+    cfg: BrittleStarConfig = OmegaConf.to_object(
+        OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg)
+    )
+    eval_cfg = cfg.evaluation
+
+    model_paths = [str(p) for p in eval_cfg.comparison_models]
+    if not model_paths:
+        raise ValueError(
+            "evaluation.comparison_models is empty. "
+            "Add at least one model path in your evaluation config."
+        )
+
+    base_seed = int(eval_cfg.comparison_base_seed)
+    num_episodes = int(eval_cfg.comparison_num_episodes)
+    max_steps = int(eval_cfg.eval_max_steps)
+
+    seeds = list(range(base_seed, base_seed + num_episodes))
+
+    output_path = Path(hydra.utils.to_absolute_path(eval_cfg.comparison_output_csv))
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    logger.info(
+        f"Comparing {len(model_paths)} models over {num_episodes} episodes "
+        f"(seeds {seeds[0]}–{seeds[-1]})."
+    )
+    logger.info(f"Results will be written to: {output_path}")
+
+    with open(output_path, "w", newline="") as csv_file:
+        writer = csv.DictWriter(csv_file, fieldnames=_FIELDNAMES)
+        writer.writeheader()
+
+        for model_path_str in model_paths:
+            model_path = Path(hydra.utils.to_absolute_path(model_path_str))
+            logger.info(f"Evaluating model: {model_path.name}")
+
+            try:
+                metadata = load_metadata(model_path)
+            except FileNotFoundError as e:
+                logger.warning(f"Skipping model — {e}")
+                continue
+
+            training = metadata_to_configs(metadata)
+
+            # Determine morphologies to evaluate
+            # If comparison_morphologies is empty, use the model's training morphology
+            morphologies = [None]
+            if eval_cfg.comparison_morphologies:
+                morphologies = [
+                    Path(hydra.utils.to_absolute_path(m)) for m in eval_cfg.comparison_morphologies
+                ]
+
+            for morph_path in morphologies:
+                morph_label = morph_path.name if morph_path else "training"
+                logger.info(f"  Morphology: {morph_label}")
+
+                bundle = build_eval_env(
+                    model_path=model_path,
+                    training=training,
+                    metadata=metadata,
+                    morphology_override_path=morph_path,
+                )
+
+                for seed in seeds:
+                    t0 = time.time()
+                    result = rollout_headless(
+                        env=bundle.env,
+                        policy=bundle.policy,
+                        seed=seed,
+                        max_steps=max_steps,
+                        action_low=bundle.action_low,
+                        action_high=bundle.action_high,
+                        action_mask=bundle.action_mask,
+                    )
+                    elapsed = time.time() - t0
+
+                    velocity = _approx_max_velocity(result)
+
+                    logger.debug(
+                        f"    seed={seed:3d} | "
+                        f"reached={str(result.reached_target):<5} | "
+                        f"return={result.return_:+8.3f} | "
+                        f"steps={result.length:4d} | "
+                        f"({elapsed:.1f}s)"
+                    )
+
+                    row = {
+                        "model_path": model_path_str,
+                        "architecture": bundle.architecture,
+                        "num_active_arms": bundle.num_active_arms,
+                        "seed": seed,
+                        "reached_target": result.reached_target,
+                        "episode_length": result.length,
+                        "eval_return": result.return_,
+                        "initial_target_distance": result.initial_target_distance,
+                        "final_xy_dist": result.final_xy_dist,
+                        "approx_max_velocity": velocity,
+                    }
+                    # Add per-arm segments
+                    for i, segs in enumerate(bundle.segments_per_arm):
+                        row[f"arm_{i}"] = segs
+
+                    writer.writerow(row)
+                    csv_file.flush()
+
+                bundle.env.close()
+
+    logger.info(f"Done. Results saved to {output_path}")
+
+
+if __name__ == "__main__":
+    register_configs()
+    main()

scripts/evaluate_checkpoints.py

@@ -0,0 +1,264 @@
+"""Re-evaluate saved checkpoints from a completed training run using MJX.
+
+This script scans the checkpoint directory of a training run (the `checkpoints/`
+folder inside a Hydra output directory), loads each `.flax` checkpoint, runs
+one deterministic evaluation episode with `build_eval_rollout_fn`, and appends
+the result to the run's `metrics/checkpoint_evaluation.csv`.
+
+It is intended for post-training analysis when per-checkpoint evaluation was not
+enabled during training (`evaluate_checkpoints: false`).
+
+Usage:
+    python scripts/evaluate_checkpoints.py \
+        simulation.model_path=runs/2024-01-01/12-00-00/final_model.flax \
+        evaluation.eval_max_steps=5000 \
+        evaluation.eval_seed=0
+
+The script resolves the run directory from `simulation.model_path`, discovers
+all `*.flax` checkpoints under `checkpoints/`, and evaluates them in order.
+"""
+
+from __future__ import annotations
+from brittle_star_project.MLPs.mlps import (
+    Actor,
+    GenericDenseLayersWithActivation,
+    MessagePasser,
+)
+from brittle_star_project.MLPs.adjancency_builder import build_adjacency
+from brittle_star_project.environment import MorphMode
+from brittle_star_project.MLPs.routing import apply_per_node
+import logging
+import re
+from pathlib import Path
+
+import hydra
+import jax
+import numpy as np
+import jax.numpy as jnp
+
+from omegaconf import DictConfig, OmegaConf
+
+from brittle_star_project.configs.main_config import BrittleStarConfig
+from brittle_star_project.configs.register_configs import register_configs
+from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
+from brittle_star_project.environment.obs_processing import create_obs_processor
+from brittle_star_project.environment.padded_obs_wrapper import compute_padding_masks
+from brittle_star_project.evaluation.checkpoint import (
+    load_metadata,
+    load_params,
+    metadata_to_configs,
+)
+from brittle_star_project.evaluation.evaluate_mjx import (
+    append_checkpoint_eval_row,
+    build_eval_rollout_fn,
+    evaluate_checkpoint_mjx,
+)
+from brittle_star_project.trainers.PPOTrainer import reward_fn
+
+
+def _parse_iteration(checkpoint_path: Path) -> int:
+    """Parse the iteration number from a checkpoint filename like `checkpoint_0042.flax`."""
+    match = re.search(r"(\d+)", checkpoint_path.stem)
+    return int(match.group(1)) if match else -1
+
+
+@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
+def main(dict_cfg: DictConfig) -> None:
+    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
+    logger = logging.getLogger(__name__)
+
+    cfg: BrittleStarConfig = OmegaConf.to_object(
+        OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg)
+    )
+    sim_cfg = cfg.simulation
+    eval_cfg = cfg.evaluation
+
+    # --- Resolve the model path to find the run directory ---
+    model_path_str = sim_cfg.model_path
+    if model_path_str is None:
+        raise ValueError(
+            "simulation.model_path must point to the final_model.flax of a training run."
+        )
+
+    model_path = Path(hydra.utils.to_absolute_path(model_path_str))
+    run_dir = model_path.parent
+
+    checkpoints_dir = run_dir / "checkpoints"
+    if not checkpoints_dir.exists():
+        raise FileNotFoundError(
+            f"No checkpoints/ directory found in run directory: {run_dir}\n"
+            "Make sure simulation.model_path points to a completed training run."
+        )
+
+    checkpoints = sorted(checkpoints_dir.glob("*.flax"), key=_parse_iteration)
+    if not checkpoints:
+        raise FileNotFoundError(f"No .flax checkpoints found in {checkpoints_dir}")
+
+    logger.info(f"Found {len(checkpoints)} checkpoint(s) in {checkpoints_dir}")
+
+    # --- Load sidecar metadata + reconstruct training config ---
+    metadata_override = (
+        Path(hydra.utils.to_absolute_path(sim_cfg.metadata_path))
+        if sim_cfg.metadata_path is not None
+        else None
+    )
+    metadata = load_metadata(model_path, metadata_override)
+    training = metadata_to_configs(metadata)
+
+    padding_masks = compute_padding_masks(
+        segments_per_arm=training.morphology.segments_per_arm,
+        reference_segments_per_arm=training.morphology.segments_per_arm,
+    )
+
+    morph_mode = training.morphology.morph_mode
+
+    segments_per_arm = jnp.asarray(
+        training.morphology.segments_per_arm,
+        dtype=jnp.int32,
+    )
+
+    num_arms = (
+        jnp.where(
+            segments_per_arm > 0,
+            1,
+            0,
+        )
+        .sum()
+        .item()
+    )
+
+    match morph_mode:
+        case MorphMode.CENTRALIZED:
+            needed_copies = 1
+            agent_indices = [0, 1, 2, 3, 4]
+
+        case MorphMode.FULLY_CONNECTED | MorphMode.RING:
+            agent_mask = segments_per_arm > 0
+            agent_indices = jnp.where(agent_mask)[0]
+            needed_copies = num_arms
+
+        case MorphMode.SEGMENT:
+            agent_mask = segments_per_arm > 0
+            agent_indices = jnp.where(agent_mask)[0]
+
+            needed_copies = (segments_per_arm.sum() + num_arms).item()
+
+    obs_processor = create_obs_processor(
+        bounds_dict=training.obs_bounds.to_bounds_dict(),
+        padding_masks=padding_masks,
+        num_arms=num_arms,
+        needed_copies=needed_copies,
+        morph_mode=morph_mode,
+        segments_per_arm=segments_per_arm,
+        agent_indices=agent_indices,
+    )
+
+    env = BrittleStarJaxEnvWrapper(
+        morphology=training.morphology,
+        arena=training.arena,
+        env_config=training.environment,
+        num_envs=1,
+    )
+
+    action_low = np.asarray(env.single_action_space.low, dtype=np.float32)
+    action_high = np.asarray(env.single_action_space.high, dtype=np.float32)
+
+    sensor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300])
+    actor = Actor(action_dim=env.single_action_space.shape[0])
+    sensor.apply = jax.jit(sensor.apply)
+    actor.apply = jax.jit(actor.apply)
+
+    eval_fn = build_eval_rollout_fn(
+        env=env,
+        obs_processor=obs_processor,
+        sensor_apply=sensor.apply,
+        actor_apply=actor.apply,
+        action_low=action_low,
+        action_high=action_high,
+        reward_fn=reward_fn,
+    )
+
+    morph_mode = training.morphology.morph_mode
+
+    segments_per_arm = jnp.asarray(
+        training.morphology.segments_per_arm,
+        dtype=jnp.int32,
+    )
+
+    match morph_mode:
+        case MorphMode.CENTRALIZED:
+            needed_copies = 1
+
+        case MorphMode.FULLY_CONNECTED | MorphMode.RING:
+            needed_copies = jnp.where(segments_per_arm > 0, 1, 0).sum().item()
+
+        case MorphMode.SEGMENT:
+            needed_copies = (
+                segments_per_arm.sum() + jnp.where(segments_per_arm > 0, 1, 0).sum()
+            ).item()
+
+    adj = build_adjacency(
+        training.morphology.segments_per_arm,
+        morph_mode,
+    )
+
+    sensor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300])
+
+    actor = Actor(action_dim=env.single_action_space.shape[0] // needed_copies)
+
+    message_passer = (
+        MessagePasser(
+            hidden_dim=300,
+            num_propagation_steps=4,
+            adj_matrix=adj,
+        )
+        if morph_mode != MorphMode.CENTRALIZED
+        else None
+    )
+
+    eval_fn = build_eval_rollout_fn(
+        env=env,
+        obs_processor=obs_processor,
+        sensor_apply=lambda p, x: apply_per_node(sensor.apply, p, x),
+        actor_apply=lambda p, x: apply_per_node(actor.apply, p, x),
+        message_passer_apply=(None if message_passer is None else message_passer.apply),
+        action_low=action_low,
+        action_high=action_high,
+        reward_fn=reward_fn,
+    )
+    seed = int(eval_cfg.eval_seed)
+    max_steps = int(eval_cfg.eval_max_steps)
+
+    logger.info(f"Evaluating each checkpoint (seed={seed}, max_steps={max_steps}).")
+
+    for checkpoint_path in checkpoints:
+        iteration = _parse_iteration(checkpoint_path)
+        try:
+            params = load_params(checkpoint_path)
+        except Exception as e:
+            logger.warning(f"Could not load {checkpoint_path.name}: {e}")
+            continue
+
+        result = evaluate_checkpoint_mjx(eval_fn, params, seed=seed, max_steps=max_steps)
+        csv_path = append_checkpoint_eval_row(
+            run_dir,
+            iteration=iteration,
+            trained_timesteps=0,  # unknown without training logs
+            result=result,
+        )
+
+        logger.debug(
+            f"checkpoint={iteration:5d} | "
+            f"reached={str(result.reached_target):<5} | "
+            f"return={result.eval_return:+8.3f} | "
+            f"steps={result.steps:4d} | "
+            f"final_dist={result.final_xy_dist:.3f}"
+        )
+
+    logger.info(f"Done. CSV at: {csv_path}")
+    env.close()
+
+
+if __name__ == "__main__":
+    register_configs()
+    main()

scripts/plots/analyze_comparisons.py

@@ -0,0 +1,445 @@
+"""
+Poster Comparison Visualizations
+
+This script generates a Forward Velocity plot and three secondary plots (Accumulated Reward, Success
+Rate, Distance Remaining).
+"""
+
+import os
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+
+from plot_config import (
+    COLORS,
+    apply_style,
+    BEST_PERFORMER_MARKER,
+    BEST_PERFORMER_TEXT,
+    BEST_PERFORMER_COLOR,
+    create_common_parser,
+    LEGEND_KWARGS,
+)
+
+
+def load_and_preprocess_data(filepath):
+    """Loads CSV and prepares the metrics for plotting."""
+    df = pd.read_csv(filepath)
+
+    # Ensure success rate can be averaged numerically
+    if "reached_target" in df.columns:
+        df["reached_target"] = df["reached_target"].astype(int)
+
+    return df
+
+
+def _add_square_placeholders(ax, x_positions, labels):
+    """Adds square placeholders for images below the x-axis."""
+    for x, label in zip(x_positions, labels):
+        # Create a roughly square rectangle in a mix of data/axes coords
+        # Shifted down to avoid overlapping with x-tick labels
+        rect = plt.Rectangle(
+            (x - 0.25, -0.40),
+            0.5,
+            0.18,
+            transform=ax.get_xaxis_transform(),
+            facecolor="#F0F0F0",
+            edgecolor="#A9A9A9",
+            linestyle="--",
+            zorder=1,
+            clip_on=False,
+        )
+        ax.add_patch(rect)
+        ax.text(
+            x,
+            -0.31,
+            f"[ Insert {label}\nImage ]",
+            transform=ax.get_xaxis_transform(),
+            ha="center",
+            va="center",
+            fontsize=10,
+            color="#888888",
+            zorder=2,
+        )
+
+
+def plot_grouped_bar(
+    df,
+    metric_col,
+    ylabel,
+    title,
+    output_filename,
+    output_dir,
+    higher_is_better=True,
+    show_titles=False,
+    figsize=(12, 8),
+):
+    """Generates and saves a highly customized grouped bar chart (grouped by Morphology)."""
+    grouped = (
+        df.groupby(["num_active_arms", "architecture"])[metric_col]
+        .agg(["mean", "std"])
+        .reset_index()
+    )
+    morphologies = sorted(grouped["num_active_arms"].unique(), reverse=True)
+    architectures = grouped["architecture"].unique()
+
+    fig, ax = plt.subplots(figsize=figsize)
+    bar_width = 0.35
+    x_indices = np.arange(len(morphologies))
+    all_bars = {}
+    all_means = []
+
+    for i, arch in enumerate(architectures):
+        arch_data = grouped[grouped["architecture"] == arch]
+        means = [
+            arch_data[arch_data["num_active_arms"] == m]["mean"].values[0]
+            if not arch_data[arch_data["num_active_arms"] == m].empty
+            else 0
+            for m in morphologies
+        ]
+        stds = [
+            arch_data[arch_data["num_active_arms"] == m]["std"].values[0]
+            if not arch_data[arch_data["num_active_arms"] == m].empty
+            else 0
+            for m in morphologies
+        ]
+        all_means.extend(means)
+        x_pos = x_indices + (i * bar_width) - (bar_width / 2 if len(architectures) == 2 else 0)
+        color = COLORS.get(arch, "#888888")
+        clean_label = arch.replace("_", " ").title()
+        bars = ax.bar(
+            x_pos,
+            means,
+            bar_width,
+            yerr=stds,
+            label=clean_label,
+            color=color,
+            capsize=8,
+            error_kw={"elinewidth": 2, "alpha": 0.7},
+        )
+        all_bars[arch] = (x_pos, means, stds, bars)
+
+    for m_idx, m in enumerate(morphologies):
+        m_means = {arch: all_bars[arch][1][m_idx] for arch in architectures}
+        best_arch = (
+            max(m_means, key=m_means.get) if higher_is_better else min(m_means, key=m_means.get)
+        )
+        best_x = all_bars[best_arch][0][m_idx]
+        best_y = all_bars[best_arch][1][m_idx]
+        best_std = all_bars[best_arch][2][m_idx]
+        offset = best_std + (abs(max(m_means.values())) * 0.05) if m_means.values() else 0
+        ax.text(
+            best_x,
+            best_y + offset,
+            BEST_PERFORMER_TEXT,
+            ha="center",
+            va="bottom",
+            fontsize=28,
+            color=BEST_PERFORMER_COLOR,
+        )
+
+    # Aesthetics
+    ax.set_ylabel(ylabel, labelpad=15)
+    if show_titles:
+        ax.set_title(title, pad=25, fontweight="bold")
+
+    x_ticks_pos = (
+        x_indices
+        + (bar_width / 2 if len(architectures) % 2 == 0 else 0)
+        - (bar_width / 2 if len(architectures) == 2 else 0)
+    )
+    ax.set_xticks(x_ticks_pos)
+    ax.set_xticklabels([f"{m} Arms" for m in morphologies])
+    ax.tick_params(axis="x", pad=25)  # More padding for the squares
+
+    # X-axis at zero
+    ax.axhline(0, color="black", linewidth=1.5)
+    ax.spines["bottom"].set_visible(False)
+
+    # Y-axis limits explicitly including 0
+    if all_means:
+        min_val = min([*all_means, 0])
+        max_val = max([*all_means, 0])
+        margin = (max_val - min_val) * 0.15 if max_val != min_val else 0.1
+        ax.set_ylim(min_val - margin, max_val + margin * 1.5)  # Extra top margin for stars
+        # Format y-ticks to not have excessive decimals, include 0
+        ticks = (
+            [min_val, max_val]
+            if min_val == 0 and max_val == 0
+            else sorted(list(set([min_val, 0, max_val])))
+        )
+        ax.set_yticks(ticks)
+        ax.yaxis.set_major_formatter(
+            plt.FuncFormatter(lambda x, _: f"{x:.2f}" if abs(x) < 10 else f"{x:.0f}")
+        )
+
+    _add_square_placeholders(ax, x_ticks_pos, [f"{m} Arms" for m in morphologies])
+
+    # Add custom legend entry for best performer
+    ax.plot(
+        [],
+        [],
+        marker=BEST_PERFORMER_MARKER,
+        color="w",
+        markerfacecolor=BEST_PERFORMER_COLOR,
+        markersize=15,
+        label="Best Performance",
+        ls="",
+    )
+    ax.legend(**LEGEND_KWARGS, ncol=len(architectures) + 1)
+    ax.set_facecolor("white")
+    fig.patch.set_facecolor("white")
+
+    os.makedirs(output_dir, exist_ok=True)
+    base_path = os.path.join(output_dir, os.path.splitext(output_filename)[0])
+    plt.savefig(f"{base_path}.png", dpi=300, bbox_inches="tight")
+    plt.savefig(f"{base_path}.svg", format="svg", bbox_inches="tight")
+    plt.close()
+
+
+def plot_grouped_bar_alt(
+    df,
+    metric_col,
+    ylabel,
+    title,
+    output_filename,
+    output_dir,
+    higher_is_better=True,
+    show_titles=False,
+    figsize=(12, 8),
+):
+    """Generates and saves a highly customized grouped bar chart (grouped by Architecture)."""
+    grouped = (
+        df.groupby(["architecture", "num_active_arms"])[metric_col]
+        .agg(["mean", "std"])
+        .reset_index()
+    )
+    architectures = sorted(grouped["architecture"].unique())
+    morphologies = sorted(grouped["num_active_arms"].unique(), reverse=True)
+
+    fig, ax = plt.subplots(figsize=figsize)
+    bar_width = 0.8 / len(morphologies)
+    x_indices = np.arange(len(architectures))
+    all_bars = {}
+    all_means = []
+
+    for i, m in enumerate(morphologies):
+        m_data = grouped[grouped["num_active_arms"] == m]
+        means = [
+            m_data[m_data["architecture"] == arch]["mean"].values[0]
+            if not m_data[m_data["architecture"] == arch].empty
+            else 0
+            for arch in architectures
+        ]
+        stds = [
+            m_data[m_data["architecture"] == arch]["std"].values[0]
+            if not m_data[m_data["architecture"] == arch].empty
+            else 0
+            for arch in architectures
+        ]
+        all_means.extend(means)
+
+        # Offset bars based on morphology index
+        offset = (i - len(morphologies) / 2 + 0.5) * bar_width
+        x_pos = x_indices + offset
+
+        # We can use a color gradient or different colors for morphologies
+        # For simplicity, using a colormap
+        color = plt.cm.viridis(i / max(1, len(morphologies) - 1))
+
+        bars = ax.bar(
+            x_pos,
+            means,
+            bar_width,
+            yerr=stds,
+            label=f"{m} Arms",
+            color=color,
+            capsize=4,
+            error_kw={"elinewidth": 1.5, "alpha": 0.7},
+        )
+        all_bars[m] = (x_pos, means, stds, bars)
+
+    for a_idx, arch in enumerate(architectures):
+        a_means = {m: all_bars[m][1][a_idx] for m in morphologies}
+        best_m = (
+            max(a_means, key=a_means.get) if higher_is_better else min(a_means, key=a_means.get)
+        )
+        best_x = all_bars[best_m][0][a_idx]
+        best_y = all_bars[best_m][1][a_idx]
+        best_std = all_bars[best_m][2][a_idx]
+        offset = best_std + (abs(max(a_means.values())) * 0.05) if a_means.values() else 0
+        ax.text(
+            best_x,
+            best_y + offset,
+            BEST_PERFORMER_TEXT,
+            ha="center",
+            va="bottom",
+            fontsize=20,
+            color=BEST_PERFORMER_COLOR,
+        )
+
+    # Aesthetics
+    ax.set_ylabel(ylabel, labelpad=15)
+    if show_titles:
+        ax.set_title(title + " (Alt)", pad=25, fontweight="bold")
+
+    ax.set_xticks(x_indices)
+    ax.set_xticklabels([arch.replace("_", " ").title() for arch in architectures])
+    ax.tick_params(axis="x", pad=25)
+
+    # X-axis at zero
+    ax.axhline(0, color="black", linewidth=1.5)
+    ax.spines["bottom"].set_visible(False)
+
+    if all_means:
+        min_val = min([*all_means, 0])
+        max_val = max([*all_means, 0])
+        margin = (max_val - min_val) * 0.15 if max_val != min_val else 0.1
+        ax.set_ylim(min_val - margin, max_val + margin * 1.5)
+        ticks = (
+            [min_val, max_val]
+            if min_val == 0 and max_val == 0
+            else sorted(list(set([min_val, 0, max_val])))
+        )
+        ax.set_yticks(ticks)
+        ax.yaxis.set_major_formatter(
+            plt.FuncFormatter(lambda x, _: f"{x:.2f}" if abs(x) < 10 else f"{x:.0f}")
+        )
+
+    # In this alt plot, placeholders might be per architecture
+    _add_square_placeholders(
+        ax, x_indices, [arch.replace("_", "\n").title() for arch in architectures]
+    )
+
+    ax.plot(
+        [],
+        [],
+        marker=BEST_PERFORMER_MARKER,
+        color="w",
+        markerfacecolor=BEST_PERFORMER_COLOR,
+        markersize=15,
+        label="Best Performance",
+        ls="",
+    )
+    ax.legend(**LEGEND_KWARGS, ncol=len(morphologies) + 1)
+    ax.set_facecolor("white")
+    fig.patch.set_facecolor("white")
+
+    os.makedirs(output_dir, exist_ok=True)
+    base_path = os.path.join(output_dir, os.path.splitext(output_filename)[0])
+    plt.savefig(f"{base_path}.png", dpi=300, bbox_inches="tight")
+    plt.savefig(f"{base_path}.svg", format="svg", bbox_inches="tight")
+    plt.close()
+
+
+if __name__ == "__main__":
+    parser = create_common_parser(description="Generate comparison poster plots.")
+    parser.add_argument(
+        "input_csv", help="Path to the input CSV file containing evaluation results."
+    )
+    args = parser.parse_args()
+
+    INPUT_CSV = args.input_csv
+    OUTPUT_DIR = args.output_dir
+
+    if not os.path.exists(INPUT_CSV):
+        print(f"Error: Could not find {INPUT_CSV}. Please ensure the file exists.")
+    else:
+        df = load_and_preprocess_data(INPUT_CSV)
+        print("Data loaded successfully. Generating poster plots...")
+
+        apply_style(font_size=args.font_size)
+        kwargs = {"show_titles": args.show_titles, "figsize": (args.fig_width, args.fig_height)}
+
+        # Velocity Conversion: m/s to cm/s
+        if "approx_max_velocity" in df.columns:
+            df["approx_max_velocity"] = df["approx_max_velocity"] * 100
+
+        # 1. Primary Plot: Forward Velocity
+        plot_grouped_bar(
+            df=df,
+            metric_col="approx_max_velocity",
+            ylabel="Max Forward Velocity (cm/s)",
+            title="Graceful Degradation: Velocity Across Morphologies",
+            output_filename="poster_plot_velocity.png",
+            output_dir=OUTPUT_DIR,
+            higher_is_better=True,
+            **kwargs,
+        )
+        plot_grouped_bar_alt(
+            df=df,
+            metric_col="approx_max_velocity",
+            ylabel="Max Forward Velocity (cm/s)",
+            title="Graceful Degradation: Velocity Across Morphologies",
+            output_filename="poster_plot_velocity_alt.png",
+            output_dir=OUTPUT_DIR,
+            higher_is_better=True,
+            **kwargs,
+        )
+
+        # 2. Secondary Plot: Accumulated Reward
+        plot_grouped_bar(
+            df=df,
+            metric_col="eval_return",
+            ylabel="Mean Cumulative Reward",
+            title="Overall Efficiency Across Morphologies",
+            output_filename="poster_plot_reward.png",
+            output_dir=OUTPUT_DIR,
+            higher_is_better=True,
+            **kwargs,
+        )
+        plot_grouped_bar_alt(
+            df=df,
+            metric_col="eval_return",
+            ylabel="Mean Cumulative Reward",
+            title="Overall Efficiency Across Morphologies",
+            output_filename="poster_plot_reward_alt.png",
+            output_dir=OUTPUT_DIR,
+            higher_is_better=True,
+            **kwargs,
+        )
+
+        # 3. Secondary Plot: Success Rate
+        plot_grouped_bar(
+            df=df,
+            metric_col="reached_target",
+            ylabel="Success Rate (%)",
+            title="Target Acquisition Consistency",
+            output_filename="poster_plot_success_rate.png",
+            output_dir=OUTPUT_DIR,
+            higher_is_better=True,
+            **kwargs,
+        )
+        plot_grouped_bar_alt(
+            df=df,
+            metric_col="reached_target",
+            ylabel="Success Rate (%)",
+            title="Target Acquisition Consistency",
+            output_filename="poster_plot_success_rate_alt.png",
+            output_dir=OUTPUT_DIR,
+            higher_is_better=True,
+            **kwargs,
+        )
+
+        # 4. Secondary Plot: Final Distance Remaining
+        plot_grouped_bar(
+            df=df,
+            metric_col="final_xy_dist",
+            ylabel="Distance to Target Remaining",
+            title="Navigational Accuracy (Lower is Better)",
+            output_filename="poster_plot_distance.png",
+            output_dir=OUTPUT_DIR,
+            higher_is_better=False,  # For distance, a lower score is better
+            **kwargs,
+        )
+        plot_grouped_bar_alt(
+            df=df,
+            metric_col="final_xy_dist",
+            ylabel="Distance to Target Remaining",
+            title="Navigational Accuracy (Lower is Better)",
+            output_filename="poster_plot_distance_alt.png",
+            output_dir=OUTPUT_DIR,
+            higher_is_better=False,
+            **kwargs,
+        )
+
+        print(f"All plots generated in the '{OUTPUT_DIR}/' directory.")

scripts/plots/analyze_convergence.py

@@ -0,0 +1,335 @@
+"""
+Convergence Analysis Script for Poster Visualizations
+
+This script analyzes evaluation metrics from multiple training runs to determine
+the convergence point of different reinforcement learning architectures.
+
+Workflow:
+1. Loads evaluation data from the CSV files defined in FILE_MAPPING.
+2. Calculates a rolling average of the reward and velocity to smooth noise.
+3. Determines the convergence timestep for each metric (first time 95% of peak is reached).
+4. Generates a grouped bar chart comparing convergence speed and line plots of the raw curves.
+
+Usage:
+    uv run python scripts/analysis/analyze_convergence.py
+
+Note: For these metrics to be valid, the evaluation CSVs must be generated with
+exploration noise strictly disabled (e.g., taking the mean of the action distribution).
+"""
+
+import logging
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+from enum import Enum
+
+from plot_config import COLORS, apply_style, create_common_parser, LEGEND_KWARGS
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+logger = logging.getLogger(__name__)
+
+# --- Globals & Configuration ---
+USING_DUMMY_DATA = False
+SMOOTHING_WINDOW = 3
+CONVERGENCE_THRESHOLD = 0.95
+
+
+class Columns(str, Enum):
+    # ... (rest of the file remains same, just need to update plotting functions and obtain_data)
+    """Column names expected in every evaluation CSV."""
+
+    ARCH = "architecture"
+    TIMESTEPS = "total_trained_timesteps"
+    REWARD = "accumulated_reward"
+    VELOCITY = "velocity"
+
+
+# Maps architecture display names to the path of their evaluation CSV.
+# Update these paths once real evaluation data is available.
+FILE_MAPPING: dict[str, str] = {
+    "centralized 2 arms": "runs/dummy/dummy_centralized_2_arms.csv",
+    "centralized 5 arms": "runs/dummy/dummy_centralized_5_arms.csv",
+    "decentralized fully connected": "runs/dummy/dummy_decentralized_fully_connected.csv",
+    "decentralized ring-level": "runs/dummy/dummy_decentralized_ring-level.csv",
+    "decentralized segment-level": "runs/dummy/dummy_decentralized_segment-level.csv",
+}
+
+# Architecture profiles for dummy data generation: (max_reward, max_velocity, sigmoid_speed)
+_DUMMY_PROFILES: dict[str, tuple[float, float, float]] = {
+    "centralized 2 arms": (300, 0.8, 1.2),
+    "centralized 5 arms": (450, 1.1, 1.0),
+    "decentralized fully connected": (500, 1.3, 0.7),
+    "decentralized ring-level": (480, 1.2, 0.8),
+    "decentralized segment-level": (520, 1.4, 0.6),
+}
+
+
+def generate_dummy_csvs(file_mapping: dict[str, str]):
+    """
+    Generates one dummy CSV per architecture in FILE_MAPPING at their expected locations.
+    Skips any architecture without a defined profile.
+    """
+    checkpoints = list(range(100, 1100, 100))
+    timesteps = [cp * 10_000 for cp in checkpoints]
+
+    for arch, path in file_mapping.items():
+        if arch not in _DUMMY_PROFILES:
+            logger.warning(f"No dummy profile for '{arch}'. Skipping.")
+            continue
+
+        m_reward, m_vel, speed = _DUMMY_PROFILES[arch]
+
+        rows = []
+        for i, ts in enumerate(timesteps):
+            progress = 1 / (1 + np.exp(-speed * (i - 4)))
+            rows.append(
+                {
+                    Columns.TIMESTEPS: ts,
+                    Columns.REWARD: m_reward * progress + np.random.normal(0, 5),
+                    Columns.VELOCITY: m_vel * progress + np.random.normal(0, 0.02),
+                }
+            )
+
+        # Create parent directories if they don't exist
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+
+        pd.DataFrame(rows).to_csv(path, index=False)
+        logger.info(f"Generated dummy CSV at expected path: {path}")
+
+
+def load_metrics(file_mapping: dict[str, str]) -> pd.DataFrame:
+    """
+    Loads one CSV per architecture, injects the architecture name as a column,
+    and returns the combined DataFrame with only the required columns.
+    """
+    required = [Columns.TIMESTEPS, Columns.REWARD, Columns.VELOCITY]
+    dfs = []
+
+    for arch_name, filepath in file_mapping.items():
+        if not os.path.exists(filepath):
+            logger.warning(f"File not found: '{filepath}'. Skipping.")
+            continue
+
+        df = pd.read_csv(filepath)
+
+        missing = [c for c in required if c not in df.columns]
+        if missing:
+            logger.warning(f"Missing columns {missing} in '{filepath}'. Skipping.")
+            continue
+
+        df = df[required].copy()
+        df[Columns.ARCH] = arch_name
+        dfs.append(df)
+
+    return pd.concat(dfs, ignore_index=True) if dfs else pd.DataFrame()
+
+
+def _convergence_timestep(series: pd.Series, timesteps: pd.Series) -> float:
+    """Returns the first timestep where the smoothed series reaches 95% of its peak."""
+    smoothed = series.rolling(window=SMOOTHING_WINDOW, min_periods=1).mean()
+    threshold = smoothed.max() * CONVERGENCE_THRESHOLD
+    return timesteps[smoothed >= threshold].iloc[0]
+
+
+def analyze_convergence(df: pd.DataFrame) -> pd.DataFrame:
+    """
+    For each architecture, determines the convergence timestep based on both
+    reward and velocity, returning one summary row per architecture.
+    """
+    results = []
+
+    for arch in df[Columns.ARCH].unique():
+        arch_data = df[df[Columns.ARCH] == arch].sort_values(Columns.TIMESTEPS)
+
+        results.append(
+            {
+                "Architecture": arch,
+                "Reward_Convergence_Timestep": _convergence_timestep(
+                    arch_data[Columns.REWARD], arch_data[Columns.TIMESTEPS]
+                ),
+                "Velocity_Convergence_Timestep": _convergence_timestep(
+                    arch_data[Columns.VELOCITY], arch_data[Columns.TIMESTEPS]
+                ),
+            }
+        )
+
+    return pd.DataFrame(results)
+
+
+def _add_bar_labels(bars, max_val: float):
+    """Annotates each bar with its value in white bold text, positioned inside."""
+    for bar in bars:
+        width = bar.get_width()
+        label = f"{width / 1e6:.1f}M" if width >= 1e6 else f"{width:,.0f}"
+        plt.text(
+            width - (max_val * 0.02),
+            bar.get_y() + bar.get_height() / 2,
+            label,
+            ha="right",
+            va="center",
+            fontsize=11,
+            color="white",
+            fontweight="bold",
+        )
+
+
+def plot_grouped_convergence_chart(
+    results_df: pd.DataFrame, output_filename: str, output_dir: str, **kwargs
+):
+    """
+    Saves a grouped horizontal bar chart comparing Reward and Velocity convergence timesteps
+    across all architectures.
+    """
+    sorted_df = results_df.sort_values("Reward_Convergence_Timestep", ascending=True)
+    architectures = sorted_df["Architecture"].tolist()
+    y_pos = np.arange(len(architectures))
+    bar_height = 0.35
+    max_val = sorted_df[
+        ["Reward_Convergence_Timestep", "Velocity_Convergence_Timestep"]
+    ].values.max()
+
+    fig, ax = plt.subplots(figsize=kwargs.get("figsize", (12, 8)))
+
+    bars_reward = ax.barh(
+        y_pos + bar_height / 2,
+        sorted_df["Reward_Convergence_Timestep"],
+        height=bar_height,
+        label="Reward Convergence",
+        color="#1f77b4",
+    )
+    bars_velocity = ax.barh(
+        y_pos - bar_height / 2,
+        sorted_df["Velocity_Convergence_Timestep"],
+        height=bar_height,
+        label="Velocity Convergence",
+        color="#ff7f0e",
+    )
+
+    title_suffix = " (DUMMY DATA)" if USING_DUMMY_DATA else ""
+    if kwargs.get("show_titles", True):
+        ax.set_title(
+            f"Comparison of Training Convergence Timesteps{title_suffix}", fontsize=20, pad=20
+        )
+    ax.set_xlabel("Timesteps to Convergence (95% of peak)", fontsize=16)
+    ax.set_ylabel("Architecture", fontsize=16)
+    ax.set_yticks(y_pos)
+    ax.set_yticklabels(architectures, fontsize=14)
+    ax.tick_params(axis="x", labelsize=14)
+    ax.legend(**LEGEND_KWARGS, ncol=2)
+    ax.set_xlim(left=0)
+    ax.spines["top"].set_visible(False)
+    ax.spines["right"].set_visible(False)
+
+    _add_bar_labels(bars_reward, max_val)
+    _add_bar_labels(bars_velocity, max_val)
+
+    plt.tight_layout()
+    os.makedirs(output_dir, exist_ok=True)
+    base_path = os.path.join(output_dir, os.path.splitext(output_filename)[0])
+    plt.savefig(f"{base_path}.png", dpi=300, bbox_inches="tight")
+    plt.savefig(f"{base_path}.svg", format="svg", bbox_inches="tight")
+    plt.close()
+
+
+def plot_metric_curves(
+    df: pd.DataFrame, metric_col: str, title: str, output_filename: str, output_dir: str, **kwargs
+):
+    """
+    Saves a line plot of the given metric over training timesteps for every architecture.
+    """
+    fig, ax = plt.subplots(figsize=kwargs.get("figsize", (12, 7)))
+
+    for arch in df[Columns.ARCH].unique():
+        arch_data = df[df[Columns.ARCH] == arch].sort_values(Columns.TIMESTEPS)
+        color_key = arch.split()[0].upper() if isinstance(arch, str) else "UNKNOWN"
+        color = COLORS.get(color_key, "#888888")
+        ax.plot(
+            arch_data[Columns.TIMESTEPS],
+            arch_data[metric_col],
+            label=arch,
+            marker="o",
+            markersize=4,
+            alpha=0.8,
+            color=color,
+        )
+
+    title_suffix = " (DUMMY DATA)" if USING_DUMMY_DATA else ""
+    if kwargs.get("show_titles", True):
+        ax.set_title(f"{title}{title_suffix}", fontsize=18, pad=20)
+    ax.set_xlabel("Training Timesteps", fontsize=14)
+    ax.set_ylabel(metric_col.replace("_", " ").title(), fontsize=14)
+    ax.legend(**LEGEND_KWARGS, ncol=len(df[Columns.ARCH].unique()))
+    ax.grid(True, linestyle="--", alpha=0.6)
+    ax.set_xlim(left=0)
+    ax.set_ylim(bottom=0)
+
+    plt.tight_layout()
+    os.makedirs(output_dir, exist_ok=True)
+    base_path = os.path.join(output_dir, os.path.splitext(output_filename)[0])
+    plt.savefig(f"{base_path}.png", dpi=300, bbox_inches="tight")
+    plt.savefig(f"{base_path}.svg", format="svg", bbox_inches="tight")
+    plt.close()
+
+
+def plot_results(df: pd.DataFrame, results: pd.DataFrame, output_dir: str, **kwargs):
+    """Generates and saves all analysis plots."""
+    plot_grouped_convergence_chart(
+        results, output_filename="convergence_comparison.png", output_dir=output_dir, **kwargs
+    )
+    plot_metric_curves(
+        df,
+        Columns.REWARD,
+        "Training Progress: Accumulated Reward",
+        "progress_reward_curves.png",
+        output_dir=output_dir,
+        **kwargs,
+    )
+    plot_metric_curves(
+        df,
+        Columns.VELOCITY,
+        "Training Progress: Velocity",
+        "progress_velocity_curves.png",
+        output_dir=output_dir,
+        **kwargs,
+    )
+
+
+def obtain_data() -> pd.DataFrame:
+    """Resolves the file mapping, falling back to generated dummy CSVs if needed."""
+    global USING_DUMMY_DATA
+    if not any(os.path.exists(p) for p in FILE_MAPPING.values()):
+        logger.info("No real evaluation files found. Generating dummy CSVs at expected locations.")
+        generate_dummy_csvs(FILE_MAPPING)
+        USING_DUMMY_DATA = True
+
+    return load_metrics(FILE_MAPPING)
+
+
+def run_analysis(output_dir: str, **kwargs):
+    """Orchestrates data loading, convergence analysis, and plot generation."""
+    df = obtain_data()
+    if df.empty:
+        logger.error("No data found to analyze.")
+        return
+
+    results = analyze_convergence(df)
+    plot_results(df, results, output_dir, **kwargs)
+    logger.info("Analysis complete. Plots saved to disk.")
+
+
+if __name__ == "__main__":
+    parser = create_common_parser(description="Analyze training convergence.")
+    args = parser.parse_args()
+
+    apply_style(font_size=args.font_size)
+    run_analysis(
+        output_dir=args.output_dir,
+        show_titles=args.show_titles,
+        figsize=(args.fig_width, args.fig_height),
+    )

scripts/plots/plot_config.py

@@ -0,0 +1,77 @@
+import argparse
+import matplotlib.pyplot as plt
+
+# Shared Color Palette (Colorblind friendly, high contrast)
+# Matches poster design
+COLORS = {
+    "CENTRALIZED": "#2B4162",  # Deep Slate Blue
+    "FULLY_CONNECTED": "#FA9F42",  # Vibrant Orange
+    "RING_LEVEL": "#4E937A",  # Muted Teal
+    "SEGMENT_LEVEL": "#B4436C",  # Soft Red
+    "DECENTRALIZED": "#4E937A",  # Default decentralized fallback
+}
+
+
+def apply_style(font_size=28):
+    """
+    Applies the shared typography and aesthetic settings to Matplotlib.
+    """
+    plt.rcParams.update(
+        {
+            "font.size": font_size,
+            "axes.labelsize": font_size + 4,
+            "axes.titlesize": font_size + 8,
+            "xtick.labelsize": font_size - 4,
+            "ytick.labelsize": font_size - 4,
+            "legend.fontsize": font_size - 6,
+            "axes.linewidth": 2,
+            "axes.spines.top": False,
+            "axes.spines.right": False,
+            "axes.spines.left": False,
+            "figure.facecolor": "white",
+            "axes.facecolor": "white",
+            "savefig.bbox": "tight",
+            "savefig.dpi": 300,
+        }
+    )
+
+
+# Star marker for best performer
+BEST_PERFORMER_TEXT = "★"
+BEST_PERFORMER_MARKER = "*"
+BEST_PERFORMER_COLOR = "#D4AF37"  # Gold
+
+# Centralized Legend Configuration
+LEGEND_KWARGS = {
+    "loc": "upper center",
+    "bbox_to_anchor": (0.5, -0.5),
+    "frameon": False,
+}
+
+
+def create_common_parser(description: str) -> argparse.ArgumentParser:
+    """
+    Creates an argparse parser with common plotting arguments.
+    """
+    parser = argparse.ArgumentParser(description=description)
+    parser.add_argument(
+        "--output_dir",
+        "-o",
+        default="runs/evaluation/plots",
+        help="Directory to save the generated plots.",
+    )
+    parser.add_argument(
+        "--show_titles",
+        action="store_true",
+        help="Include titles in the plots. Default is False for easier poster integration.",
+    )
+    parser.add_argument(
+        "--font_size", type=int, default=28, help="Base font size in points. Default is 28."
+    )
+    parser.add_argument(
+        "--fig_width", type=float, default=12.0, help="Figure width in inches. Default is 12.0."
+    )
+    parser.add_argument(
+        "--fig_height", type=float, default=8.0, help="Figure height in inches. Default is 8.0."
+    )
+    return parser

scripts/simulate.py

@@ -1,667 +1,163 @@
 """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
-
-To replay a run using the *exact* Hydra config used during training, pass:
-    python scripts/simulate.py simulation.trained_config_path=runs/.../.hydra/config.yaml \
+Automatically extracts the training configuration (morphology, environment, etc.)
+from the sidecar metadata YAML file to ensure simulation perfectly matches training.
+Override simulation settings via CLI, e.g.:
+    uv run scripts/simulate.py \
+        simulation.morphology_override=configs/morphology/3_arms.yaml \
         simulation.model_path=runs/.../final_model.flax
 """
 
 from __future__ import annotations
 
-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
-import yaml
-from omegaconf import DictConfig, OmegaConf, open_dict
+from omegaconf import DictConfig, OmegaConf
+
 
-from brittle_star_project import Backend, BrittleStarEnv, BrittleStarEnvFactory
 from brittle_star_project.configs.main_config import BrittleStarConfig
 from brittle_star_project.configs.register_configs import register_configs
-from brittle_star_project.environment.padded_obs_wrapper import (
-    compute_padding_masks,
-    pad_observation,
+
+from brittle_star_project.evaluation.checkpoint import load_metadata, metadata_to_configs
+from brittle_star_project.evaluation.eval_env_builder import build_eval_env
+from brittle_star_project.evaluation.rollout import rollout_headless, rollout_viewer
+from brittle_star_project.evaluation.video import (
+    record_episode,
+    create_evaluation_dir,
+    save_evaluation_metadata,
 )
 
-_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 _dense_layer_sizes_from_params(params: Any) -> list[int]:
-    """Infer GenericDenseLayersWithActivation.layer_sizes from a Flax params tree."""
-
-    try:
-        dense_params = params["params"]
-    except Exception as exc:
-        raise ValueError("Unexpected sensor params structure (missing 'params')") from exc
-
-    layer_sizes: list[int] = []
-    idx = 0
-    while True:
-        key = f"Dense_{idx}"
-        if key not in dense_params:
-            break
-        kernel = dense_params[key]["kernel"]
-        layer_sizes.append(int(np.asarray(kernel).shape[1]))
-        idx += 1
-
-    if not layer_sizes:
-        raise ValueError("Could not infer Dense_* layers from sensor params")
-    return layer_sizes
-
-
-def _infer_action_dim_from_actor_params(params: Any) -> int | None:
-    """Best-effort infer action_dim from a Flax Actor params tree."""
-
-    try:
-        dense0 = params["params"]["Dense_0"]
-        bias = dense0.get("bias")
-        kernel = dense0.get("kernel")
-    except Exception:
-        return None
-
-    if bias is not None:
-        try:
-            return int(np.asarray(bias).shape[0])
-        except Exception:
-            return None
-
-    if kernel is not None:
-        try:
-            return int(np.asarray(kernel).shape[1])
-        except Exception:
-            return None
-
-    return None
-
-
-def _has_cli_override(overrides: list[str], key: str) -> bool:
-    prefixes = (f"{key}=", f"{key}.", f"+{key}=", f"+{key}.")
-    return any(str(o).startswith(prefixes) for o in overrides)
-
-
-def _maybe_clip_action(
-    action: np.ndarray,
-    low: np.ndarray | None,
-    high: np.ndarray | None,
-) -> np.ndarray:
-    if low is None or high is None:
-        return action
-    low = np.asarray(low, dtype=np.float32).ravel()
-    high = np.asarray(high, dtype=np.float32).ravel()
-    if low.shape != action.shape or high.shape != action.shape:
-        return action
-    return np.clip(action, low, high)
-
-
-def _transform_obs_dict(obs_dict: dict[str, Any]) -> jnp.ndarray:
-    """Flatten the env's observation dict into a 1D vector.
-
-    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 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,)))
-
-    if not parts:
-        return jnp.zeros((0,), dtype=jnp.float32)
-    return jnp.concatenate(parts, axis=0)
-
-
-# A minimal policy class to load a CleanRL/Flax checkpoint and run inference.
-class CleanRLPPOPolicy:
-    def __init__(
-        self,
-        *,
-        sensor_params: Any,
-        actor_params: Any,
-        action_dim: int,
-    ) -> None:
-        from brittle_star_project.MLPs.mlps import Actor, GenericDenseLayersWithActivation
-
-        layer_sizes = _dense_layer_sizes_from_params(sensor_params)
-        self._sensor = GenericDenseLayersWithActivation(layer_sizes=layer_sizes)
-        self._actor = Actor(action_dim=action_dim)
-        self._sensor_apply = jax.jit(self._sensor.apply)
-        self._actor_apply = jax.jit(self._actor.apply)
-        self._params = {
-            "sensor_params": sensor_params,
-            "actor_params": actor_params,
-        }
-
-    @staticmethod
-    def load(
-        path: Path,
-        *,
-        action_dim: int,
-    ) -> "CleanRLPPOPolicy":
-        def _get_index(container: Any, idx: int) -> Any:
-            if isinstance(container, (list, tuple)):
-                return container[idx]
-            if isinstance(container, dict):
-                return container.get(idx, container.get(str(idx)))
-            raise KeyError(idx)
-
-        def _looks_like_indexed_dict(container: Any) -> bool:
-            return (
-                isinstance(container, dict)
-                and container
-                and all(str(k).isdigit() for k in container.keys())
-            )
-
-        def _parse_checkpoint(restored_obj: Any) -> tuple[Any, Any, Any, Any, Any]:
-            """Extract checkpoint parts.
-
-            Returns (config_dict, sensor_params, actor_params, critic_params,
-            feature_extractor_params).
-
-            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.
-            """
-
-            cfg_part: Any | None = None
-            params_part: Any = restored_obj
-
-            if isinstance(restored_obj, (list, tuple)) and len(restored_obj) >= 2:
-                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
-            ):
-                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):
-                sensor_params = _get_index(params_part, 0)
-                actor_params = _get_index(params_part, 1)
-                critic_params = _get_index(params_part, 2)
-                feature_extractor_params = _get_index(params_part, 3)
-                if sensor_params is None or actor_params is None:
-                    raise ValueError("Missing required params in checkpoint")
-                return (
-                    cfg_part,
-                    sensor_params,
-                    actor_params,
-                    critic_params,
-                    feature_extractor_params,
-                )
-
-            if isinstance(params_part, (list, tuple)) and len(params_part) >= 2:
-                sensor_params = params_part[0]
-                actor_params = params_part[1]
-                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 (
-                    cfg_part,
-                    sensor_params,
-                    actor_params,
-                    critic_params,
-                    feature_extractor_params,
-                )
-
-            # Accept a plain dict-shaped Flax params mapping commonly produced
-            # by saving `agent_state.params` directly. Typical keys are
-            # 'sensor_params' and 'actor_params', or sometimes nested under 'params'.
-            if isinstance(restored_obj, dict):
-                # Top-level params dict
-                params_sub = restored_obj.get("params", {})
-                sensor_params = restored_obj.get("sensor_params") or params_sub.get("sensor_params")
-                actor_params = restored_obj.get("actor_params") or params_sub.get("actor_params")
-                critic_params = restored_obj.get("critic_params") or params_sub.get("critic_params")
-                feature_extractor_params = restored_obj.get(
-                    "feature_extractor_params"
-                ) or params_sub.get("feature_extractor_params")
-                # Some checkpoints only save actor+sensor as top-level
-                if sensor_params is not None and actor_params is not None:
-                    return (
-                        cfg_part,
-                        sensor_params,
-                        actor_params,
-                        critic_params,
-                        feature_extractor_params,
-                    )
-
-            raise ValueError(
-                f"Unexpected checkpoint structure in {path}. "
-                "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)
-        _cfg_dict, sensor_params, actor_params, _critic_params, _feature_extractor_params = (
-            _parse_checkpoint(restored)
-        )
-
-        ckpt_action_dim = _infer_action_dim_from_actor_params(actor_params)
-        if ckpt_action_dim is not None and ckpt_action_dim != action_dim:
-            raise ValueError(
-                "Checkpoint/env mismatch: "
-                f"checkpoint expects action_dim={ckpt_action_dim}, "
-                f"env provides action_dim={action_dim}. "
-                "Use the same Hydra config (morphology/arena/environment) "
-                "that was used during training."
-            )
-
-        return CleanRLPPOPolicy(
-            sensor_params=sensor_params,
-            actor_params=actor_params,
-            action_dim=action_dim,
-        )
-
-    def act(self, *, observations: dict[str, Any]) -> np.ndarray:
-        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)
-
-        # Always evaluate with the actor mean.
-        # (Sampling adds exploration noise, which is useful for training but not for evaluation.)
-        return np.asarray(mean, dtype=np.float32).ravel()
-
-
-def _get_observations(state: Any) -> dict[str, Any] | None:
-    return getattr(state, "observations", None)
-
-
-def _get_xy_distance_to_target(observations: dict[str, Any]) -> float | None:
-    return float(np.asarray(observations["xy_distance_to_target"]).reshape(-1)[0])
-
-
-def _target_reached(*, state: Any) -> bool:
-    return bool(getattr(state, "terminated", False) or getattr(state, "truncated", False))
-
-
-def _rollout_one_episode_headless(
-    *,
-    env: BrittleStarEnv,
-    policy: CleanRLPPOPolicy,
-    seed: int,
-    max_steps: int,
-    action_low: np.ndarray | None,
-    action_high: np.ndarray | None,
-    padding_masks: dict[str, Any] | None,
-) -> tuple[float, int, bool, float | None]:
-    """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)
-
-    steps = 0
-    for _ in range(int(max_steps)):
-        obs_dict = observations or {}
-        if padding_masks is not None:
-            obs_dict = pad_observation(obs_dict, padding_masks)
-
-        action = policy.act(observations=obs_dict)
-        action = _maybe_clip_action(action, action_low, action_high)
-
-        nu = int(state.mj_model.nu)
-        if nu > 0 and action.shape != (nu,):
-            raise ValueError(f"Policy returned action shape {action.shape}, expected ({nu},)")
-
-        state = env.step(state=state, action=action)
-        steps += 1
-
-        observations = _get_observations(state)
-        cur_dist = _get_xy_distance_to_target(observations)
-        if prev_dist is not None and cur_dist is not None:
-            ep_return += prev_dist - cur_dist
-        prev_dist = cur_dist
-
-        reached_target = _target_reached(state=state)
-        if reached_target:
-            break
-
-    final_dist = _get_xy_distance_to_target(observations)
-    return ep_return, steps, reached_target, final_dist
-
-
-def _run_one_episode_viewer(
-    *,
-    env: BrittleStarEnv,
-    policy: CleanRLPPOPolicy,
-    seed: int,
-    state: Any,
-    control_dt: float,
-    max_steps: int | None,
-    action_low: np.ndarray | None,
-    action_high: np.ndarray | None,
-    padding_masks: dict[str, Any] | None,
-) -> None:
-    import mujoco.viewer
-
-    model = state.mj_model
-    data = state.mj_data
-
-    _ = 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.
-    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:
-            if not viewer.is_running():
-                break
-            step_start = time.time()
-
-            obs_dict = observations or {}
-            if padding_masks is not None:
-                obs_dict = pad_observation(obs_dict, padding_masks)
-
-            action = policy.act(observations=obs_dict)
-            action = _maybe_clip_action(action, action_low, action_high)
-            if model.nu > 0 and action.shape != (int(model.nu),):
-                raise ValueError(
-                    f"Policy returned action shape {action.shape}, expected ({int(model.nu)},)"
-                )
-
-            # The passive viewer runs a GUI thread; protect MuJoCo state mutation.
-            with viewer.lock():
-                state = env.step(state=state, action=action)
-
-            if not viewer.is_running():
-                break
-            viewer.sync()
-
-            steps += 1
-
-            observations = _get_observations(state)
-            cur_dist = _get_xy_distance_to_target(observations)
-            if prev_dist is not None and cur_dist is not None:
-                episode_return += prev_dist - cur_dist
-            prev_dist = cur_dist
-
-            reached_target = _target_reached(state=state)
-            if reached_target:
-                break
-
-            remaining = control_dt - (time.time() - step_start)
-            if remaining > 0:
-                time.sleep(remaining)
-
-    dist = _get_xy_distance_to_target(observations)
-    dist_str = "n/a" if dist is None else f"{dist:.3f}"
-    print(
-        "episode done: "
-        f"return={episode_return:.6f}, len={steps}, "
-        f"target_reached={reached_target}, final_xy_dist={dist_str}"
-    )
-
-
-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 _load_trained_config(path: Path) -> DictConfig:
-    """Load a trained config YAML.
-
-    Supports both:
-    - Hydra's run config (e.g. runs/.../.hydra/config.yaml)
-    - This project's logger metadata YAMLs, which may contain
-      ``!!python/object/apply:...`` tags for Enums.
-
-    For safety, we *do not* execute Python constructors from YAML; we only
-    treat these tags as data and extract their scalar arguments.
-    """
-
-    if not path.exists():
-        raise FileNotFoundError(f"trained_config_path does not exist: '{path}'.")
-    if not path.is_file():
-        raise ValueError(f"trained_config_path must be a file, got: '{path}'.")
-
-    try:
-        return OmegaConf.load(path)
-    except Exception as exc:
-        python_apply_prefix = "tag:yaml.org,2002:python/object/apply:"
-
-        class _SafeLoaderWithPythonApply(yaml.SafeLoader):
-            pass
-
-        def _construct_python_apply(
-            loader: yaml.SafeLoader,
-            _tag_suffix: str,
-            node: yaml.Node,
-        ) -> Any:
-            if isinstance(node, yaml.SequenceNode):
-                seq = loader.construct_sequence(node)
-                if len(seq) == 1:
-                    return seq[0]
-                return seq
-            if isinstance(node, yaml.MappingNode):
-                return loader.construct_mapping(node)
-            return loader.construct_scalar(node)
-
-        _SafeLoaderWithPythonApply.add_multi_constructor(
-            python_apply_prefix, _construct_python_apply
-        )
-
-        try:
-            data = yaml.load(path.read_text(encoding="utf-8"), Loader=_SafeLoaderWithPythonApply)
-        except Exception as yaml_exc:
-            raise ValueError(
-                "Failed to load trained_config_path as YAML. "
-                "If this is a Hydra run, pass the run's '.hydra/config.yaml' file. "
-                f"Got: '{path}'."
-            ) from yaml_exc
-
-        if not isinstance(data, dict):
-            raise ValueError(
-                "trained_config_path must contain a YAML mapping (dict-like) at the root. "
-                f"Got type={type(data).__name__} from '{path}'."
-            ) from exc
-
-        # Normalize known enum-like strings to their Enum *names* so OmegaConf's
-        # structured config merge behaves like the normal Hydra config.
-        from brittle_star_project.environment.env_types import Task
-
-        env_cfg = data.get("environment")
-        if isinstance(env_cfg, dict) and isinstance(env_cfg.get("task"), str):
-            task_str = str(env_cfg["task"])
-            try:
-                env_cfg["task"] = Task[task_str].name
-            except Exception:
-                try:
-                    env_cfg["task"] = Task(task_str).name
-                except Exception:
-                    pass
-
-        return OmegaConf.create(data)
-
 
 @hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
 def main(dict_cfg: DictConfig) -> None:
-    # Compose against the structured schema first, so missing keys are validated.
-    cfg = OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg)
+    # 1. Hydra composes ONLY SimulationSettings
+    cfg = OmegaConf.to_object(OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg))
+    sim_cfg = cfg.simulation
 
-    # Optional: override env-defining sections (morphology/arena/environment/architecture)
-    # using the exact Hydra config that was used for training.
-    trained_cfg_path = cfg.simulation.trained_config_path
-    if trained_cfg_path:
-        overrides_raw = OmegaConf.select(cfg, "hydra.overrides.task") or []
-        overrides = [str(o) for o in overrides_raw]
-
-        trained_cfg_path_abs = Path(hydra.utils.to_absolute_path(trained_cfg_path))
-        trained_cfg = _load_trained_config(trained_cfg_path_abs)
-        if "hydra" in trained_cfg:
-            with open_dict(trained_cfg):
-                del trained_cfg["hydra"]
-
-        with open_dict(cfg):
-            for key in ("morphology", "arena", "environment", "architecture"):
-                if key in trained_cfg and not _has_cli_override(overrides, key):
-                    base_node = OmegaConf.select(cfg, key)
-                    override_node = OmegaConf.select(trained_cfg, key)
-                    try:
-                        cfg[key] = OmegaConf.merge(base_node, override_node)
-                    except Exception as exc:
-                        raise ValueError(
-                            "Failed to merge trained config into the active Hydra config. "
-                            f"Key={key!r}, trained_config_path='{trained_cfg_path_abs}'."
-                        ) from exc
-
-    # Convert DictConfig to structured dataclass.
-    config: BrittleStarConfig = OmegaConf.to_object(cfg)
-
-    backend = Backend.MJC
-    seed = int(config.experiment.seed)
-
-    if getattr(config.architecture, "name", None) != "centralized":
-        raise ValueError(
-            "simulate.py currently only supports architecture=centralized. "
-            f"Got architecture.name={getattr(config.architecture, 'name', None)!r}. "
-            "(Training supports decentralized, but simulation wiring for it isn't implemented.)"
-        )
-
-    model_path_str = config.simulation.model_path
+    model_path_str = sim_cfg.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 =======
-    factory = BrittleStarEnvFactory()
-    raw_env = factory.create_environment(
-        backend,
-        config.morphology,
-        config.arena,
-        config.environment,
-    )
-    env = BrittleStarEnv(
-        raw_env,
-        backend=backend,
-        config=config.environment,
-        morphology_config=config.morphology,
+    # 2. Discover + load sidecar metadata YAML
+    metadata_override = None
+    if sim_cfg.metadata_path is not None:
+        metadata_override = Path(hydra.utils.to_absolute_path(sim_cfg.metadata_path))
+
+    metadata = load_metadata(model_path, metadata_override)
+
+    # 3. Reconstruct typed configs from metadata
+    training = metadata_to_configs(metadata)
+
+    seed = int(cfg.experiment.seed)
+
+    # 4-7. Build evaluation environment and policy
+    override_path = None
+    if sim_cfg.morphology_override is not None:
+        override_path = Path(hydra.utils.to_absolute_path(sim_cfg.morphology_override))
+
+    bundle = build_eval_env(
+        model_path=model_path,
+        training=training,
+        metadata=metadata,
+        morphology_override_path=override_path,
     )
 
+    env = bundle.env
+    policy = bundle.policy
+    action_low = bundle.action_low
+    action_high = bundle.action_high
+    action_mask = bundle.action_mask
+
     state0 = env.reset(seed=seed)
 
-    # Match training's padded observation layout for amputated morphologies.
-    padding_masks = compute_padding_masks(config.morphology.segments_per_arm)
+    # 8. Run simulation
+    headless = bool(sim_cfg.headless)
+    max_steps = sim_cfg.max_steps
 
-    # Match training's action clipping behavior.
-    action_space = getattr(raw_env, "action_space", None)
-    action_low = (
-        None if action_space is None else np.asarray(action_space.low, dtype=np.float32).ravel()
-    )
-    action_high = (
-        None if action_space is None else np.asarray(action_space.high, dtype=np.float32).ravel()
-    )
-
-    # ======= MODEL SETUP =======
-    nu = int(state0.mj_model.nu)
-    policy = CleanRLPPOPolicy.load(model_path, action_dim=nu)
-
-    # Helpful early failure when configs don't match the checkpoint.
-    observations0 = _get_observations(state0)
-    obs0_dict = pad_observation(observations0 or {}, padding_masks)
-    env_obs_dim = int(_transform_obs_dict(obs0_dict).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(
-            "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."
-        )
-
-    # ======= SIMULATION =======
-    headless = bool(config.simulation.headless)
-    max_steps = config.simulation.max_steps
-
-    if headless:
+    if sim_cfg.record_video:
         if max_steps is None:
-            raise ValueError("simulation.max_steps is required when simulation.headless=true")
+            raise ValueError("simulation.max_steps is required when simulation.record_video=true")
+
         max_steps_i = int(max_steps)
         if max_steps_i <= 0:
             raise ValueError("simulation.max_steps must be > 0")
 
-        ep_return, ep_len, reached_target, final_dist = _rollout_one_episode_headless(
+        if sim_cfg.video_output_path is None:
+            eval_dir = create_evaluation_dir(model_path)
+            output_path = eval_dir / "simulation.mp4"
+        else:
+            output_path = Path(hydra.utils.to_absolute_path(sim_cfg.video_output_path))
+            eval_dir = output_path.parent
+            eval_dir.mkdir(parents=True, exist_ok=True)
+
+        result = record_episode(
             env=env,
             policy=policy,
             seed=seed,
             max_steps=max_steps_i,
             action_low=action_low,
             action_high=action_high,
-            padding_masks=padding_masks,
+            action_mask=action_mask,
+            output_path=output_path,
+            camera_id=sim_cfg.camera_id,
         )
-        final_dist_str = "n/a" if final_dist is None else f"{final_dist:.3f}"
+
+        save_evaluation_metadata(
+            eval_dir=eval_dir,
+            morphology_override_path=sim_cfg.morphology_override,
+            seed=seed,
+            max_steps=max_steps_i,
+            result=result,
+        )
+        final_dist_str = "n/a" if result.final_xy_dist is None else f"{result.final_xy_dist:.3f}"
+        print(f"Video saved to {output_path}")
         print(
             "episode done: "
-            f"return={ep_return:.6f}, len={ep_len}, "
-            f"target_reached={reached_target}, final_xy_dist={final_dist_str}"
+            f"return={result.return_:.6f}, len={result.length}, "
+            f"target_reached={result.reached_target}, final_xy_dist={final_dist_str}"
+        )
+    elif 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")
+
+        result = rollout_headless(
+            env=env,
+            policy=policy,
+            seed=seed,
+            max_steps=max_steps_i,
+            action_low=action_low,
+            action_high=action_high,
+            action_mask=action_mask,
+        )
+        final_dist_str = "n/a" if result.final_xy_dist is None else f"{result.final_xy_dist:.3f}"
+        print(
+            "episode done: "
+            f"return={result.return_:.6f}, len={result.length}, "
+            f"target_reached={result.reached_target}, final_xy_dist={final_dist_str}"
         )
     else:
+        max_steps_val = 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_val = None
+            max_steps_val = max_steps_i
 
         model_dt = float(state0.mj_model.opt.timestep)
-        control_dt = model_dt * float(config.environment.num_physics_steps_per_control_step)
+        control_dt = model_dt * float(training.environment.num_physics_steps_per_control_step)
 
-        _run_one_episode_viewer(
+        rollout_viewer(
             env=env,
             policy=policy,
             seed=seed,
@@ -670,7 +166,7 @@ def main(dict_cfg: DictConfig) -> None:
             max_steps=max_steps_val,
             action_low=action_low,
             action_high=action_high,
-            padding_masks=padding_masks,
+            action_mask=action_mask,
         )
 
     env.close()

scripts/simulate.sh

@@ -0,0 +1,9 @@
+#!/usr/bin/env bash
+
+path=$1
+
+uv run simulate.py \
+    simulation.model_path="$path"/final_model.flax \
+    simulation.record_video=True \
+    simulation.video_output_path=../vids/simulation.mp4 \
+    simulation.max_steps=10000

scripts/tools/dump_mjcf.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+"""
+Dump MJCF XML for a brittle-star morphology using the project's Hydra configs.
+
+Usage examples:
+
+  # Use a named morphology config from configs/morphology (Hydra style)
+  uv run python scripts/analysis/dump_mjcf.py morphology=3_arms
+
+  # Use a morphology override YAML (same key as simulation.morphology_override)
+  uv run python scripts/analysis/dump_mjcf.py \
+    simulation.morphology_override=configs/morphology/3_arms.yaml
+
+Output path:
+    Provide `dump_out=path/to/file.xml` on the command line, otherwise writes `morphology.xml` in
+    current directory or `runs/morphologies/<name>.xml`.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import logging
+import sys
+from pathlib import Path
+from typing import Any, Optional
+
+import hydra
+import yaml
+from omegaconf import DictConfig, OmegaConf
+
+from brittle_star_project.configs.register_configs import register_configs
+from brittle_star_project.environment.env_config import MorphologyConfig
+from brittle_star_project.environment.factory import BrittleStarEnvFactory
+
+logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
+logger = logging.getLogger(__name__)
+
+
+def extract_xml_string(obj: Any) -> Optional[str]:
+    """
+    Attempts to serialize the morphology object to an XML string by checking
+    common dm_control and internal API methods.
+    """
+    serialization_methods = [
+        "to_xml_string",
+        "to_xml",
+        "to_string",
+        "to_mjcf",
+        "to_mjcf_string",
+        "get_mjcf",
+        "get_mjcf_str",
+        "export_to_xml_string",
+    ]
+
+    # If the object itself has an 'mjcf' attribute, try to serialize that instead
+    target_obj = getattr(obj, "mjcf", obj)
+
+    for method_name in serialization_methods:
+        method = getattr(target_obj, method_name, None)
+        if callable(method):
+            try:
+                xml_data = method()
+                # Safely handle both string and byte responses
+                if isinstance(xml_data, str):
+                    return xml_data
+                elif isinstance(xml_data, bytes):
+                    return xml_data.decode("utf-8")
+            except Exception as e:
+                logger.debug(f"Method {method_name}() failed during serialization: {e}")
+
+    return None
+
+
+def resolve_output_path(cfg: DictConfig) -> Path:
+    """Determines the appropriate output path for the MJCF XML."""
+    dump_out = cfg.get("dump_out", None)
+    if dump_out is not None:
+        return Path(hydra.utils.to_absolute_path(str(dump_out)))
+
+    morph_name = "morphology"
+    for arg in sys.argv[1:]:
+        if arg.startswith("morphology="):
+            morph_name = arg.split("=", 1)[1]
+            break
+
+    default_out = (
+        f"runs/morphologies/{morph_name}.xml" if morph_name != "morphology" else "morphology.xml"
+    )
+    return Path(hydra.utils.to_absolute_path(default_out))
+
+
+@hydra.main(config_path="../../configs", config_name="main_config", version_base="1.3")
+def main(cfg: DictConfig) -> None:
+    """Main entry point to construct the morphology and dump its XML."""
+    logger.info("Initializing morphology construction...")
+
+    # Extract morphology config safely using dict `.get()` to avoid OmegaConf AttributeErrors
+    simulation_cfg = cfg.get("simulation", cfg)
+    override_path = simulation_cfg.get("morphology_override", None)
+
+    if override_path:
+        logger.info(f"Using morphology override: {override_path}")
+        with open(hydra.utils.to_absolute_path(override_path), "r") as f:
+            data = yaml.safe_load(f) or {}
+        morph_cfg = MorphologyConfig(**data)
+    else:
+        # Fallback to default simulation morphology, or an empty base config
+        morph_node = simulation_cfg.get("morphology", cfg.get("morphology", None))
+
+        if morph_node is not None:
+            # Convert OmegaConf node to dict and instantiate MorphologyConfig.
+            # This ensures any missing keys gracefully fall back to the dataclass defaults.
+            morph_dict = OmegaConf.to_container(morph_node, resolve=True)
+            if isinstance(morph_dict, dict):
+                # Filter to avoid unexpected kwargs if the dataclass is strictly defined
+                if dataclasses.is_dataclass(MorphologyConfig):
+                    valid_keys = {f.name for f in dataclasses.fields(MorphologyConfig)}
+                    morph_dict = {k: v for k, v in morph_dict.items() if k in valid_keys}
+                morph_cfg = MorphologyConfig(**morph_dict)
+            else:
+                morph_cfg = MorphologyConfig()
+        else:
+            morph_cfg = MorphologyConfig()
+
+    morphology = BrittleStarEnvFactory.create_morphology(morph_cfg)
+
+    xml_text = extract_xml_string(morphology)
+    if not xml_text:
+        raise RuntimeError("Failed to serialize morphology to MJCF/XML. ")
+
+    out_path = resolve_output_path(cfg)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    with out_path.open("w", encoding="utf-8") as f:
+        f.write(xml_text)
+
+    logger.info(f"Successfully exported MJCF XML to: {out_path}")
+
+
+if __name__ == "__main__":
+    register_configs()
+    main()

scripts/tools/extract_observation_bounds.py

@@ -0,0 +1,138 @@
+#!/usr/bin/env python3
+"""Empirically extract observation bounds (focused on joint velocities).
+
+This script creates a MuJoCo environment using the project's factory and
+randomly samples actions to discover observed maxima for selected
+observation keys (joint_velocity, joint_position, joint_actuator_force).
+
+Usage:
+  python scripts/extract_observation_bounds.py \
+    --morphology configs/morphology/3_arms.yaml --num-steps 5000 --seed 42
+
+If `--morphology` is omitted the default `MorphologyConfig()` is used.
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+import yaml
+import numpy as np
+
+from brittle_star_project import BrittleStarEnvFactory, BrittleStarEnv, Backend
+from brittle_star_project.environment.env_config import (
+    MorphologyConfig,
+    ArenaConfig,
+    EnvConfig,
+)
+
+
+def load_morphology(path: str | None) -> MorphologyConfig:
+    if path is None:
+        return MorphologyConfig()
+    p = Path(path)
+    if not p.exists():
+        raise FileNotFoundError(f"Morphology file not found: {p}")
+    with open(p, "r") as f:
+        data = yaml.safe_load(f) or {}
+    return MorphologyConfig(**data)
+
+
+def _extract_observations(state):
+    # Under different backends the returned state may be a dict or an object
+    obs = getattr(state, "observations", None)
+    if obs is None and isinstance(state, dict):
+        obs = state.get("observations", state)
+    return obs
+
+
+def find_empirical_bounds(
+    morph_cfg: MorphologyConfig,
+    arena_cfg: ArenaConfig,
+    env_cfg: EnvConfig,
+    num_steps: int = 5000,
+    seed: int = 42,
+) -> None:
+    factory = BrittleStarEnvFactory()
+    raw_env = factory.create_environment(Backend.MJC, morph_cfg, arena_cfg, env_cfg)
+    env = BrittleStarEnv(raw_env, backend=Backend.MJC, config=env_cfg, morphology_config=morph_cfg)
+
+    # Initial reset
+    state = env.reset(seed=seed)
+
+    # Determine action bounds
+    action_space = getattr(raw_env, "action_space", None)
+    if action_space is None:
+        raise RuntimeError("Environment missing `action_space`; cannot sample actions.")
+
+    action_low = np.asarray(action_space.low, dtype=np.float32)
+    action_high = np.asarray(action_space.high, dtype=np.float32)
+    action_shape = action_low.shape
+
+    # Track maximum absolute observed values
+    tracked_keys = ["joint_velocity", "joint_position", "joint_actuator_force"]
+    max_observed = {k: 0.0 for k in tracked_keys}
+
+    # Include observation at reset
+    obs0 = _extract_observations(state)
+    if isinstance(obs0, dict):
+        for k in tracked_keys:
+            if k in obs0:
+                max_observed[k] = max(max_observed[k], float(np.max(np.abs(np.asarray(obs0[k])))))
+
+    rng = np.random.RandomState(seed)
+    for i in range(num_steps):
+        u = rng.uniform(size=action_shape)
+        action = action_low + (action_high - action_low) * u
+
+        # Provide a numpy RNG to the env step; wrapper will pass it if accepted.
+        step_out = env.step(state=state, action=action, rng=env.make_rng(seed + i + 1))
+
+        # Unpack next state from common return conventions
+        if hasattr(step_out, "state"):
+            next_state = step_out.state
+        elif isinstance(step_out, (tuple, list)) and len(step_out) >= 1:
+            next_state = step_out[0]
+        else:
+            next_state = step_out
+
+        obs = _extract_observations(next_state)
+        if isinstance(obs, dict):
+            for k in tracked_keys:
+                if k in obs:
+                    val = float(np.max(np.abs(np.asarray(obs[k]))))
+                    if val > max_observed[k]:
+                        max_observed[k] = val
+
+        state = next_state
+
+    # Print recommended bounds with a 20% safety margin
+    print("\n--- Recommended Observation Bounds (20% margin) ---")
+    for k, v in max_observed.items():
+        if v == 0.0:
+            print(f"{k}: observed max 0.0 (increase sampling or inspect env)")
+        else:
+            safe = v * 1.2
+            print(f"{k}: [-{safe:.6f}, {safe:.6f}]  (observed max: {v:.6f})")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--morphology", type=str, default=None, help="Path to morphology YAML (optional)"
+    )
+    parser.add_argument(
+        "--num-steps", type=int, default=5000, help="Number of random steps to sample"
+    )
+    parser.add_argument("--seed", type=int, default=42, help="RNG seed")
+    args = parser.parse_args()
+
+    morph_cfg = load_morphology(args.morphology)
+    arena_cfg = ArenaConfig()
+    env_cfg = EnvConfig()
+
+    find_empirical_bounds(morph_cfg, arena_cfg, env_cfg, num_steps=args.num_steps, seed=args.seed)
+
+
+if __name__ == "__main__":
+    main()

src/brittle_star_project/MLPs/__init__.py

@@ -0,0 +1,19 @@
+from .mlps import (
+    GenericDenseLayersWithActivation,
+    OneDenseLayerMLP,
+    Actor,
+    MessagePasser,
+    AgentParams,
+    Storage,
+)
+from .adjancency_builder import build_adjacency
+
+__all__ = [
+    "GenericDenseLayersWithActivation",
+    "OneDenseLayerMLP",
+    "Actor",
+    "MessagePasser",
+    "AgentParams",
+    "Storage",
+    "build_adjacency",
+]

src/brittle_star_project/MLPs/adjancency_builder.py

@@ -0,0 +1,67 @@
+from brittle_star_project.environment.env_config import MorphMode
+import jax.numpy as jnp
+
+
+def build_adjacency(segments_per_arm, mode: MorphMode):
+    num_arms = sum(1 for s in segments_per_arm if s > 0)
+    num_segments = sum(segments_per_arm)
+
+    # FOR NOW SEMI HARDCODE:
+    # CENTRALIZED: 1 agent, no stress, adja = 1,1 = [[1]]
+    # FULLY CONNECTED: 5 agents: adj = alle 1
+    # CENTRAL DISK:#arms= 5 agents, only neighbor as adjacent so diagonal kinda..
+    # ARM = #segments agents: diago kinda, but extra, center ring too, put center mlps first or..
+
+    if mode == MorphMode.CENTRALIZED:
+        return jnp.ones((1, 1))
+
+    if mode == MorphMode.FULLY_CONNECTED:
+        adj = jnp.ones((num_arms, num_arms))  # everybody adjacent everybody
+        return adj
+
+    if mode == MorphMode.RING:  # ring
+        adj = jnp.zeros((num_arms, num_arms))
+        for i in range(num_arms):
+            adj = adj.at[i, i].set(1)  # self
+            adj = adj.at[i, (i - 1) % num_arms].set(1)
+            adj = adj.at[i, (i + 1) % num_arms].set(1)  # left and right..
+        return adj
+
+    if mode == MorphMode.SEGMENT:
+        num_nodes = num_arms + num_segments
+        adj = jnp.zeros((num_nodes, num_nodes))
+
+        # first ring
+        for i in range(num_arms):
+            # self
+            adj = adj.at[i, i].set(1)
+
+            # ring neighbors
+            adj = adj.at[i, (i - 1) % num_arms].set(1)
+            adj = adj.at[i, (i + 1) % num_arms].set(1)
+
+        # then segment chains
+        idx = 0
+        for arm_idx, seg_count in enumerate(segments_per_arm):
+            for i in range(seg_count):
+                seg_node = num_arms + idx + i
+
+                adj = adj.at[seg_node, seg_node].set(1)
+                if i > 0:
+                    adj = adj.at[seg_node, seg_node - 1].set(1)
+                if i < seg_count - 1:
+                    adj = adj.at[seg_node, seg_node + 1].set(1)
+
+            idx += seg_count
+
+        idx = 0
+        for arm_idx, seg_count in enumerate(segments_per_arm):
+            first_seg = num_arms + idx  # first segment of this arm
+
+            # connect ring node first segment
+            adj = adj.at[arm_idx, first_seg].set(1)
+            adj = adj.at[first_seg, arm_idx].set(1)
+
+            idx += seg_count
+
+        return adj

src/brittle_star_project/MLPs/mlps.py

@@ -1,11 +1,11 @@
 from dataclasses import dataclass, fields, field
 
-import flax
 import flax.linen as nn
 import jax.numpy as jnp
 import jax.tree_util
 from typing import Sequence, Callable
 from flax.linen.initializers import constant, orthogonal
+from flax.core import FrozenDict
 
 
 # semi generic so we can easily make a config for it in experiments
@@ -37,30 +37,56 @@ class Actor(nn.Module):
         return mean, log_std
 
 
+class MessagePasser(nn.Module):
+    hidden_dim: int
+    num_propagation_steps: int
+    adj_matrix: jnp.ndarray
+
+    @nn.compact
+    def __call__(self, x: jnp.ndarray):
+        for _ in range(self.num_propagation_steps):
+            # (n_nodes, feat)
+            messages = nn.Dense(self.hidden_dim)(x)
+            messages = nn.tanh(messages)
+
+            # note: if mean is wanted: adj_matrix / (adj.sum(axis=-1, keepdims=True) + 1e-8)
+            agg = self.adj_matrix
+            aggregated = agg @ messages
+
+            x_concat = jnp.concatenate([x, aggregated], axis=-1)
+
+            gate = nn.sigmoid(nn.Dense(self.hidden_dim)(x_concat))
+            candidate = nn.tanh(nn.Dense(self.hidden_dim)(x_concat))
+            x = gate * x + (1 - gate) * candidate
+
+        return x
+
+
 @jax.tree_util.register_dataclass
 @dataclass
 class AgentParams:
-    sensor_params: flax.core.FrozenDict
-    actor_params: flax.core.FrozenDict
-    critic_params: flax.core.FrozenDict
-    feature_extractor_params: flax.core.FrozenDict
+    sensor_params: FrozenDict | dict
+    actor_params: FrozenDict | dict
+    critic_params: FrozenDict | dict
+    feature_extractor_params: FrozenDict | dict
+    message_passer_params: FrozenDict | dict
 
 
 @jax.tree_util.register_dataclass
 @dataclass
 class Storage:
-    obs: jnp.array
-    actions: jnp.array
-    logprobs: jnp.array
-    dones: jnp.array
-    values: jnp.array
-    advantages: jnp.array
-    returns: jnp.array
-    rewards: jnp.array
+    obs: jnp.ndarray
+    actions: jnp.ndarray
+    logprobs: jnp.ndarray
+    dones: jnp.ndarray
+    values: jnp.ndarray
+    advantages: jnp.ndarray
+    returns: jnp.ndarray
+    rewards: jnp.ndarray
 
-    raw_actions: jnp.ndarray = None  # before clipping
-    means: jnp.ndarray = None  # policy mean
-    stds: jnp.ndarray = None  # policy std
+    raw_actions: jnp.ndarray | None = None  # before clipping
+    means: jnp.ndarray | None = None  # policy mean
+    stds: jnp.ndarray | None = None  # policy std
 
     def replace(self, **kwargs) -> "Storage":
         fs = fields(self)

src/brittle_star_project/MLPs/routing.py

@@ -0,0 +1,22 @@
+"""Shared JAX routing utilities for decentralized multi-agent models."""
+
+import jax
+
+
+def apply_per_node(apply_fn, params, x):
+    """Apply a Flax module independently to each node.
+
+    Args:
+        apply_fn: The module's ``apply`` method (e.g. ``sensor.apply``).
+        params: Per-node parameters with shape ``(num_nodes, ...)``.
+        x: Input tensor with shape ``(batch, num_nodes, features)``.
+
+    Returns:
+        Output tensor with shape ``(batch, num_nodes, out_features)``.
+    """
+
+    def apply_single_node(p, x_node):
+        # x_node: (batch, feat) — one node's input across the batch
+        return jax.vmap(lambda xi: apply_fn(p, xi))(x_node)
+
+    return jax.vmap(apply_single_node, in_axes=(0, 1), out_axes=1)(params, x)

src/brittle_star_project/__init__.py

@@ -2,7 +2,14 @@ from .environment.env_types import Backend, Task
 from .environment.env_config import ArenaConfig, EnvConfig, MorphologyConfig
 from .environment.factory import BrittleStarEnvFactory
 from .environment.env_wrapper import BrittleStarEnv
-from .render import simulate_policy, SimulationConfig, ControlPolicy
+from .evaluation import (
+    PolicyAgent,
+    ControlPolicy,
+    load_metadata,
+    rollout_headless,
+    rollout_viewer,
+    EpisodeResult,
+)
 
 __all__ = [
     "ArenaConfig",
@@ -12,7 +19,10 @@ __all__ = [
     "EnvConfig",
     "MorphologyConfig",
     "Task",
-    "simulate_policy",
-    "SimulationConfig",
+    "PolicyAgent",
     "ControlPolicy",
+    "load_metadata",
+    "rollout_headless",
+    "rollout_viewer",
+    "EpisodeResult",
 ]

src/brittle_star_project/configs/config_evaluation.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class EvaluationConfig:
+    """Evaluation settings.
+
+    Currently used for synchronous checkpoint evaluation during training.
+    """
+
+    # When enabled, each saved checkpoint is evaluated headlessly and the results
+    # are appended to a CSV in the run's metrics/ folder.
+    evaluate_checkpoints: bool = False
+    eval_max_steps: int = 5000
+    eval_seed: int = 0
+
+    # Cross-model comparison settings.
+    # comparison_base_seed is the starting seed for generating episode seeds.
+    comparison_base_seed: int = 0
+    # comparison_num_episodes controls how many target positions to evaluate for each model.
+    comparison_num_episodes: int = 5
+    # comparison_models lists the paths (relative to workspace root) to the .cleanrl_model files.
+    comparison_models: list[str] = field(default_factory=list)
+    # Path where the comparison results CSV will be saved (relative to workspace root).
+    comparison_output_csv: str = "metrics/model_comparison.csv"
+    # Morphology override YAML paths for cross-morphology comparison.
+    # Each path points to a file in configs/morphology/ (e.g., "configs/morphology/3_arms.yaml").
+    # When empty, each model is evaluated only on its training morphology.
+    comparison_morphologies: list[str] = field(default_factory=list)
+
+    def __post_init__(self) -> None:
+        if self.evaluate_checkpoints and self.eval_max_steps <= 0:
+            raise ValueError(
+                "Configuration Error: 'eval_max_steps' must be > 0 when "
+                "'evaluate_checkpoints' is enabled."
+            )

src/brittle_star_project/configs/config_simulation.py

@@ -13,7 +13,20 @@ class SimulationSettings:
     # If None, viewer mode runs until window closed or target reached.
     max_steps: Optional[int] = None
 
-    # Optional: point to a Hydra config.yaml from a training run (e.g. runs/.../.hydra/config.yaml).
-    # When set, the simulation script can override
-    # morphology/arena/environment/architecture to match.
-    trained_config_path: Optional[str] = None
+    # Override morphology for amputation experiments.
+    # When set, the environment uses this morphology instead of the trained one.
+    # Points to a morphology config YAML file (e.g. configs/morphology/3_arms.yaml).
+    # Observations are padded from the override morphology UP TO the training
+    # morphology's shape via compute_padding_masks(override, reference=training).
+    morphology_override: Optional[str] = None
+
+    # Video recording (requires [evaluation] extra)
+    record_video: bool = False
+    # When None, video is saved in a per-model evaluation folder alongside the model.
+    video_output_path: Optional[str] = None
+    # Camera ID to use for video recording (1 is usually the close-up camera)
+    camera_id: int = 1
+
+    # Optional override for the sidecar metadata YAML file.
+    # If None, it defaults to the model_path with a `_metadata.yaml` suffix.
+    metadata_path: Optional[str] = None

src/brittle_star_project/configs/main_config.py

@@ -2,10 +2,16 @@ 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_evaluation import EvaluationConfig
 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
+from brittle_star_project.environment.env_config import (
+    MorphologyConfig,
+    ArenaConfig,
+    EnvConfig,
+    ObservationBoundsConfig,
+)
 
 
 @dataclass
@@ -18,6 +24,7 @@ class BrittleStarConfig:
 
     experiment: ExperimentConfig = field(default_factory=ExperimentConfig)
     logging: LoggingConfig = field(default_factory=LoggingConfig)
+    evaluation: EvaluationConfig = field(default_factory=EvaluationConfig)
     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.
@@ -25,4 +32,5 @@ class BrittleStarConfig:
     morphology: MorphologyConfig = field(default_factory=MorphologyConfig)
     arena: ArenaConfig = field(default_factory=ArenaConfig)
     environment: EnvConfig = field(default_factory=EnvConfig)
+    obs_bounds: ObservationBoundsConfig = field(default_factory=ObservationBoundsConfig)
     simulation: SimulationSettings = field(default_factory=SimulationSettings)

src/brittle_star_project/configs/register_configs.py

@@ -2,13 +2,19 @@ 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_evaluation import EvaluationConfig
 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.environment.env_config import (
+    MorphologyConfig,
+    ArenaConfig,
+    EnvConfig,
+    ObservationBoundsConfig,
+)
 from brittle_star_project.configs.main_config import BrittleStarConfig
 
 
@@ -27,6 +33,7 @@ def register_configs() -> None:
     # 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="evaluation", name="base_evaluation", node=EvaluationConfig)
     cs.store(group="ppo", name="base_ppo", node=PPOConfig)
 
     # Architecture variants — swap via CLI: architecture=decentralized
@@ -37,4 +44,5 @@ def register_configs() -> None:
     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="obs_bounds", name="base_obs_bounds", node=ObservationBoundsConfig)
     cs.store(group="simulation", name="base_simulation", node=SimulationSettings)

src/brittle_star_project/dataclasses/EpisodeStatistics.py

@@ -4,7 +4,7 @@ import jax.numpy as jnp
 
 @flax.struct.dataclass
 class EpisodeStatistics:
-    episode_returns: jnp.array
-    episode_lengths: jnp.array
-    returned_episode_returns: jnp.array
-    returned_episode_lengths: jnp.array
+    episode_returns: jnp.ndarray
+    episode_lengths: jnp.ndarray
+    returned_episode_returns: jnp.ndarray
+    returned_episode_lengths: jnp.ndarray

src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py

@@ -5,7 +5,7 @@ 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
+from .padded_obs_wrapper import compute_padding_masks
 
 
 class BrittleStarJaxEnvWrapper:
@@ -48,6 +48,15 @@ class BrittleStarJaxEnvWrapper:
     def raw(self):
         return self._env
 
+    @property
+    def padding_masks(self) -> dict:
+        """Pre-computed boolean masks for amputated limb padding.
+
+        Pass to create_obs_processor so the processor handles padding
+        after normalization in the correct pipeline order.
+        """
+        return self._padding_masks
+
     @property
     def single_action_space(self):
         return self._env.action_space
@@ -61,10 +70,6 @@ class BrittleStarJaxEnvWrapper:
         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))
         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):
@@ -75,12 +80,7 @@ class BrittleStarJaxEnvWrapper:
         return self._vectorized_action_sample(rng=jnp.array(sub_rngs))
 
     def step(self, state, 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
+        return self._vectorized_step(state=state, action=action)
 
     def close(self):
         self._env.close()

src/brittle_star_project/environment/__init__.py

@@ -1,7 +1,9 @@
-from .env_config import ArenaConfig, EnvConfig, MorphologyConfig
+from .env_config import ArenaConfig, EnvConfig, MorphologyConfig, MorphMode
 from .env_types import Backend, Task
-from .env_wrapper import BrittleStarEnv, StepResult
+from .env_wrapper import BrittleStarEnv
 from .factory import BrittleStarEnvFactory
+from .obs_processing import create_obs_processor
+from .padded_obs_wrapper import compute_padding_masks
 
 __all__ = [
     "ArenaConfig",
@@ -10,6 +12,8 @@ __all__ = [
     "Backend",
     "Task",
     "BrittleStarEnv",
-    "StepResult",
     "BrittleStarEnvFactory",
+    "MorphMode",
+    "create_obs_processor",
+    "compute_padding_masks",
 ]

src/brittle_star_project/environment/env_config.py

@@ -1,10 +1,18 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
+from enum import Enum
 
 from .env_types import Task
 
 
+class MorphMode(Enum):
+    CENTRALIZED = 0
+    FULLY_CONNECTED = 1
+    RING = 2
+    SEGMENT = 3
+
+
 @dataclass
 class MorphologyConfig:
     """Brittle star morphology configuration.
@@ -20,6 +28,7 @@ class MorphologyConfig:
     segments_per_arm: list[int] = field(default_factory=lambda: [4, 4, 4, 4, 4])
     use_p_control: bool = True
     use_torque_control: bool = False
+    morph_mode: MorphMode = MorphMode.CENTRALIZED
 
     @property
     def num_arms(self) -> int:
@@ -60,3 +69,31 @@ class EnvConfig:
     # Light escape
     # Per docs in upstream env config: integer factors of 200.
     light_perlin_noise_scale: int = 0
+
+
+@dataclass
+class ObservationBoundsConfig:
+    """Physical observation bounds for deterministic min-max normalization."""
+
+    # Empirical testing based on the extract_observation_bounds.py script run for 1.000.000 steps
+
+    # Based on max. ctrlrange (0.78539816339744828) in XML, but empirical testing went slightly over
+    joint_position: list[float] = field(default_factory=lambda: [-0.8, 0.8])
+    # Empirical testing showed max. 3.22, adding buffer to be safe. Consider higher values "fast".
+    joint_velocity: list[float] = field(default_factory=lambda: [-5.0, 5.0])
+    # Based on max. forceRange in XML, verified with empirical testing
+    joint_actuator_force: list[float] = field(default_factory=lambda: [-3.75, 3.75])
+    # Based on intuition and reasoning
+    segment_contact: list[float] = field(default_factory=lambda: [0.0, 1.0])
+    robot_direction_to_target: list[float] = field(default_factory=lambda: [-1.0, 1.0])
+    disk_z_tilt: list[float] = field(default_factory=lambda: [0.0, 3.141592653589793])
+
+    def to_bounds_dict(self) -> dict[str, tuple[float, float]]:
+        return {
+            "disk_z_tilt": tuple(self.disk_z_tilt),
+            "joint_actuator_force": tuple(self.joint_actuator_force),
+            "joint_position": tuple(self.joint_position),
+            "joint_velocity": tuple(self.joint_velocity),
+            "robot_direction_to_target": tuple(self.robot_direction_to_target),
+            "segment_contact": tuple(self.segment_contact),
+        }

src/brittle_star_project/environment/obs_processing.py

@@ -0,0 +1,192 @@
+import jax
+import jax.numpy as jnp
+from typing import Dict, Tuple, Optional
+
+from brittle_star_project.environment.env_config import MorphMode
+
+from experiment_logger import get_logger
+
+logger = get_logger()
+
+_JOINT_SCALED_KEYS = frozenset(
+    {
+        "joint_position",
+        "joint_velocity",
+        "joint_actuator_force",
+        "actuator_force",
+    }
+)
+
+_SEGMENT_SCALED_KEYS = frozenset(
+    {
+        "segment_contact",
+    }
+)
+
+
+def _build_joint_indices(segments_per_arm, indices_mlp):
+    indices = []
+    start = 0
+    for i, segs in enumerate(segments_per_arm):
+        # 2 joints per segment
+        if i in indices_mlp:
+            count = segs * 2
+            idx = jnp.arange(start, start + count)
+            indices.append(idx)
+            start += count
+    return indices
+
+
+def _build_segment_indices(segments_per_arm, indices_mlp):
+    indices = []
+    start = 0
+    for i, segs in enumerate(segments_per_arm):
+        if i in indices_mlp:
+            idx = jnp.arange(start, start + segs)
+            indices.append(idx)
+            start += segs
+    return indices
+
+
+def create_obs_processor(
+    bounds_dict: Dict[str, Tuple[float, float]],
+    num_arms: int,
+    needed_copies: int,
+    padding_masks: Optional[Dict] = None,
+    morph_mode: MorphMode = MorphMode.CENTRALIZED,
+    segments_per_arm=[4, 4, 4, 4, 4],
+    agent_indices=[0, 1, 2, 3, 4],
+):
+    # made a set to allow O(1) search
+    ordered_keys = frozenset(
+        [
+            "disk_z_tilt",
+            "joint_actuator_force",
+            "joint_position",
+            "joint_velocity",
+            "robot_direction_to_target",
+            "segment_contact",
+        ]
+    )
+    segment_indices = _build_segment_indices(segments_per_arm, agent_indices)
+    joint_indices = _build_joint_indices(segments_per_arm, agent_indices)
+
+    def _add_derived_features(obs: dict) -> dict:
+        new_obs = dict(obs)
+        if "disk_rotation" in new_obs:
+            rot = new_obs["disk_rotation"]
+            new_obs["disk_z_tilt"] = jnp.sqrt(jnp.pow(rot[0], 2) + jnp.pow(rot[1], 2))
+
+            if "unit_xy_direction_to_target" in new_obs:
+                yaw = rot[2]
+                unit_x, unit_y = new_obs["unit_xy_direction_to_target"]
+                cos_yaw, sin_yaw = jnp.cos(yaw), jnp.sin(yaw)
+                new_x = unit_x * cos_yaw + unit_y * sin_yaw
+                new_y = -unit_x * sin_yaw + unit_y * cos_yaw
+                new_obs["robot_direction_to_target"] = jnp.stack([new_x, new_y])
+
+        return new_obs
+
+    def _normalize_features(obs: dict) -> dict:
+        normalized = {}
+        for key, arr in obs.items():
+            if key in bounds_dict:
+                low, high = bounds_dict[key]
+                if low == -1.0 and high == 1.0:
+                    normalized[key] = jnp.clip(arr, -1.0, 1.0)
+                else:
+                    arr_clipped = jnp.clip(arr, low, high)
+                    normalized[key] = 2.0 * (arr_clipped - low) / (high - low) - 1.0
+            else:
+                normalized[key] = arr
+        return normalized
+
+    def _split_to_agents(obs: dict, morph_mode) -> dict:
+        output = {}
+        num_agents = needed_copies  # IMPORTANT: number of MLPs
+
+        segs_per_arm = 4
+        joints_per_segment = 2
+        joints_per_arm = segs_per_arm * joints_per_segment
+        for key, arr in obs.items():
+            arr = jnp.asarray(arr)
+            if arr.size == 0:
+                continue
+
+            if arr.ndim == 0:
+                arr = arr.reshape(1)
+
+            if key in _SEGMENT_SCALED_KEYS:
+                per_agent = []
+                for i, _ in enumerate(agent_indices):
+                    idx = segment_indices[i]
+                    taken = jnp.take(arr, idx, axis=0)
+                    pad_len = segs_per_arm - taken.shape[0]
+                    padded = jnp.pad(taken, [(0, pad_len)] + [(0, 0)] * (taken.ndim - 1))
+
+                    per_agent.append(padded.reshape(-1))
+                arr = jnp.stack(per_agent)
+            elif key in _JOINT_SCALED_KEYS:
+                per_agent = []
+                for i, _ in enumerate(agent_indices):
+                    idx = joint_indices[i]
+                    taken = jnp.take(arr, idx, axis=0)
+                    pad_len = joints_per_arm - taken.shape[0]
+                    padded = jnp.pad(taken, [(0, pad_len)] + [(0, 0)] * (taken.ndim - 1))
+
+                    per_agent.append(padded.reshape(-1))
+                arr = jnp.stack(per_agent)
+            else:
+                arr = jnp.repeat(arr[None, :], num_agents, axis=0)
+
+            if morph_mode == MorphMode.CENTRALIZED:
+                output[key] = arr.reshape(1, -1)
+            elif key in _JOINT_SCALED_KEYS:
+                output[key] = arr.reshape(num_agents, -1)
+            elif key in _SEGMENT_SCALED_KEYS:
+                output[key] = arr[:, None]
+            else:
+                output[key] = arr
+
+        return output
+
+    def _flatten_features(obs: dict) -> jnp.ndarray:
+        """
+        Input:
+            key -> (num_arms, feat_per_key)
+
+        Output:
+            (num_arms, total_features)
+        """
+        values = []
+
+        for key in sorted(ordered_keys):
+            if key not in obs:
+                continue
+
+            arr = jnp.asarray(obs[key])  # (num_arms, feat)
+
+            if arr.size == 0:
+                continue
+
+            if arr.ndim == 1:
+                arr = arr[:, None]
+
+            arr = arr.reshape(arr.shape[0], -1)
+
+            values.append(arr)
+
+        return jnp.concatenate(values, axis=-1)  # (num_arms, total_feat)
+
+    def _process_single(obs_dict: dict) -> jnp.ndarray:
+        processed = _add_derived_features(obs_dict)
+        processed = _normalize_features(processed)
+        processed = _split_to_agents(processed, morph_mode)
+        flat = _flatten_features(processed)  # (num_arms, total_feat)
+
+        logger.debug(f"[FLATTENED FINAL] shape: {flat.shape}")
+        logger.debug(f"[PER AGENT] example row 0 shape: {flat[0].shape}")
+
+        return flat  # (agents, feat)
+
+    return jax.jit(jax.vmap(_process_single))

src/brittle_star_project/environment/padded_obs_wrapper.py

@@ -1,35 +1,11 @@
-"""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.
-"""
+"""Observation padding masks for amputated brittle star morphologies."""
 
 from __future__ import annotations
 
 from typing import Any, Sequence
 
-import jax
 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: Sequence[int],
@@ -54,13 +30,18 @@ def compute_padding_masks(
     mask_2x = []
 
     for arm_idx, (actual, ref) in enumerate(zip(segments_per_arm, reference_segments_per_arm)):
+        if not isinstance(actual, int):
+            actual = actual.item()
+
+        if not isinstance(ref, int):
+            ref = ref.item()
+
         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))
@@ -71,60 +52,3 @@ def compute_padding_masks(
         "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():
-        padded_dtype = _padding_dtype(value)
-        if key in _JOINT_SCALED_KEYS:
-            out = jnp.zeros(masks["target_size_2x"], dtype=padded_dtype)
-            padded[key] = out.at[masks["mask_2x"]].set(value)
-        elif key in _SEGMENT_SCALED_KEYS:
-            out = jnp.zeros(masks["target_size_1x"], dtype=padded_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]
-        padded_dtype = _padding_dtype(value)
-        if key in _JOINT_SCALED_KEYS:
-            out = jnp.zeros((batch_size, masks["target_size_2x"]), dtype=padded_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=padded_dtype)
-            padded[key] = out.at[:, masks["mask_1x"]].set(value)
-        else:
-            padded[key] = value
-    return padded
-
-
-def _padding_dtype(value: Any) -> jnp.dtype:
-    """Choose a JAX-safe dtype for padding arrays.
-
-    When JAX x64 is disabled, allocating float64 zeros emits a warning. We
-    preserve the original dtype whenever it is supported, and otherwise fall
-    back to float32 for padding buffers.
-    """
-
-    dtype = getattr(value, "dtype", None)
-    if dtype is None:
-        dtype = jnp.asarray(value).dtype
-    else:
-        dtype = jnp.dtype(dtype)
-    if dtype == jnp.float64 and not jax.config.read("jax_enable_x64"):
-        return jnp.float32
-    return dtype

src/brittle_star_project/evaluation/__init__.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from .checkpoint import load_metadata, load_params, metadata_to_configs, TrainingConfig
+from .evaluate_mjx import (
+    CheckpointEvalResult,
+    append_checkpoint_eval_row,
+    build_eval_rollout_fn,
+    evaluate_checkpoint_mjx,
+)
+from .evaluate import evaluate_policy
+from .policy import PolicyAgent, ControlPolicy
+from .rollout import rollout_headless, rollout_viewer, EpisodeResult
+from .video import record_episode, create_evaluation_dir, save_evaluation_metadata
+from .eval_env_builder import EvalEnvBundle, build_eval_env
+
+__all__ = [
+    # checkpoint loading
+    "load_metadata",
+    "load_params",
+    "metadata_to_configs",
+    "TrainingConfig",
+    # MJX evaluation
+    "CheckpointEvalResult",
+    "append_checkpoint_eval_row",
+    "build_eval_rollout_fn",
+    "evaluate_checkpoint_mjx",
+    # CPU evaluation
+    "evaluate_policy",
+    # policy
+    "PolicyAgent",
+    "ControlPolicy",
+    # rollout
+    "rollout_headless",
+    "rollout_viewer",
+    "EpisodeResult",
+    # video
+    "record_episode",
+    "create_evaluation_dir",
+    "save_evaluation_metadata",
+    # env builder
+    "EvalEnvBundle",
+    "build_eval_env",
+]

src/brittle_star_project/evaluation/checkpoint.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+import yaml
+from dataclasses import dataclass
+from pathlib import Path
+
+from collections.abc import Mapping
+
+import flax
+from omegaconf import OmegaConf
+
+from brittle_star_project.environment.env_config import (
+    MorphologyConfig,
+    ArenaConfig,
+    EnvConfig,
+    ObservationBoundsConfig,
+)
+
+
+@dataclass
+class TrainingConfig:
+    """Holds typed configurations extracted from a training run's metadata."""
+
+    morphology: MorphologyConfig
+    arena: ArenaConfig
+    environment: EnvConfig
+    obs_bounds: ObservationBoundsConfig
+
+
+def load_params(path: Path) -> dict:
+    """Load model parameters from a .flax checkpoint file."""
+    payload = path.read_bytes()
+    restored = flax.serialization.msgpack_restore(payload)
+
+    sensor_params = None
+    actor_params = None
+    message_passer_params = None
+
+    # Extract params from restored checkpoint
+    if isinstance(restored, Mapping):
+        params_sub = restored.get("params", {})
+        sensor_params = restored.get("sensor_params") or params_sub.get("sensor_params")
+        actor_params = restored.get("actor_params") or params_sub.get("actor_params")
+        message_passer_params = restored.get("message_passer_params") or params_sub.get(
+            "message_passer_params"
+        )
+    elif isinstance(restored, (list, tuple)) and len(restored) >= 2:
+        params_part = restored[1]
+        if isinstance(params_part, Mapping):
+            sensor_params = params_part.get("0", params_part.get(0))
+            actor_params = params_part.get("1", params_part.get(1))
+        elif isinstance(params_part, (list, tuple)) and len(params_part) >= 2:
+            sensor_params = params_part[0]
+            actor_params = params_part[1]
+
+    if sensor_params is None or actor_params is None:
+        raise ValueError(f"Could not extract sensor and actor params from checkpoint: {path}")
+
+    return {
+        "sensor_params": sensor_params,
+        "actor_params": actor_params,
+        "message_passer_params": message_passer_params,
+    }
+
+
+def load_metadata(model_path: Path, metadata_override_path: Path | None = None) -> dict:
+    """Discover and load the sidecar metadata YAML file."""
+    if metadata_override_path is not None:
+        metadata_path = metadata_override_path
+    else:
+        metadata_path = model_path.with_name(model_path.stem + "_metadata.yaml")
+
+    if not metadata_path.exists():
+        raise FileNotFoundError(f"Could not find metadata YAML at {metadata_path}")
+    with open(metadata_path, "r") as f:
+        return yaml.safe_load(f)
+
+
+def metadata_to_configs(metadata: dict) -> TrainingConfig:
+    """Reconstruct typed configuration objects from a metadata dictionary."""
+    trained_morphology = OmegaConf.to_object(
+        OmegaConf.merge(OmegaConf.structured(MorphologyConfig), metadata.get("morphology", {}))
+    )
+    trained_arena = OmegaConf.to_object(
+        OmegaConf.merge(OmegaConf.structured(ArenaConfig), metadata.get("arena", {}))
+    )
+
+    env_dict = metadata.get("environment", {})
+    if isinstance(env_dict.get("task"), str):
+        from brittle_star_project.environment.env_types import Task
+
+        try:
+            env_dict["task"] = Task[env_dict["task"]].name
+        except Exception:
+            try:
+                env_dict["task"] = Task(env_dict["task"]).name
+            except Exception:
+                pass
+
+    trained_environment = OmegaConf.to_object(
+        OmegaConf.merge(OmegaConf.structured(EnvConfig), env_dict)
+    )
+    trained_obs_bounds = OmegaConf.to_object(
+        OmegaConf.merge(
+            OmegaConf.structured(ObservationBoundsConfig), metadata.get("obs_bounds", {})
+        )
+    )
+
+    return TrainingConfig(
+        morphology=trained_morphology,
+        arena=trained_arena,
+        environment=trained_environment,
+        obs_bounds=trained_obs_bounds,
+    )

src/brittle_star_project/evaluation/eval_env_builder.py

@@ -0,0 +1,176 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import jax.numpy as jnp
+import numpy as np
+import yaml
+from omegaconf import OmegaConf
+
+from brittle_star_project import Backend, BrittleStarEnv, BrittleStarEnvFactory
+from brittle_star_project.environment.env_config import MorphMode, MorphologyConfig
+from brittle_star_project.environment.obs_processing import create_obs_processor
+from brittle_star_project.environment.padded_obs_wrapper import compute_padding_masks
+from brittle_star_project.evaluation.checkpoint import TrainingConfig
+from brittle_star_project.evaluation.policy import PolicyAgent
+from brittle_star_project.MLPs.adjancency_builder import build_adjacency
+
+
+@dataclass
+class EvalEnvBundle:
+    """Everything needed to run a headless evaluation episode."""
+
+    env: BrittleStarEnv
+    policy: PolicyAgent
+    action_low: np.ndarray | None
+    action_high: np.ndarray | None
+    action_mask: np.ndarray | None
+    segments_per_arm: list[int]
+    num_active_arms: int
+    architecture: str
+
+
+def build_eval_env(
+    *,
+    model_path: Path,
+    training: TrainingConfig,
+    metadata: dict,
+    morphology_override_path: Path | str | None = None,
+) -> EvalEnvBundle:
+    """Build environment + policy for evaluation, optionally with a morphology override."""
+
+    # 1. Determine environment morphology
+    if morphology_override_path is not None:
+        override_path = Path(morphology_override_path)
+        if not override_path.exists():
+            raise FileNotFoundError(f"Could not find morphology override YAML at {override_path}")
+        with open(override_path, "r") as f:
+            override_dict = yaml.safe_load(f)
+        env_morphology = OmegaConf.to_object(
+            OmegaConf.merge(OmegaConf.structured(MorphologyConfig), override_dict)
+        )
+        # Force morph_mode to be inherited from training since it's baked into weights
+        env_morphology.morph_mode = training.morphology.morph_mode
+    else:
+        env_morphology = training.morphology
+
+    # 2. Build obs_processor with TRAINING morphology padding masks always
+    padding_masks = compute_padding_masks(
+        segments_per_arm=env_morphology.segments_per_arm,
+        reference_segments_per_arm=training.morphology.segments_per_arm,
+    )
+
+    training_segs_per_arm = jnp.array(training.morphology.segments_per_arm)
+
+    needed_copies = 0
+    agent_indices = [0, 1, 2, 3, 4]
+    match training.morphology.morph_mode:
+        case MorphMode.CENTRALIZED:
+            needed_copies = 1
+        case MorphMode.FULLY_CONNECTED | MorphMode.RING:
+            agent_mask = training_segs_per_arm > 0
+            agent_indices = jnp.where(agent_mask)[0].tolist()
+            needed_copies = jnp.where(training_segs_per_arm > 0, 1, 0).sum().item()
+        case MorphMode.SEGMENT:
+            agent_mask = training_segs_per_arm > 0
+            agent_indices = jnp.where(agent_mask)[0].tolist()
+            needed_copies = (
+                training_segs_per_arm.sum() + jnp.where(training_segs_per_arm > 0, 1, 0).sum()
+            ).item()
+
+    num_arms_training = jnp.where(training_segs_per_arm > 0, 1, 0).sum().item()
+
+    obs_processor = create_obs_processor(
+        bounds_dict=training.obs_bounds.to_bounds_dict(),
+        padding_masks=padding_masks,
+        needed_copies=needed_copies,
+        num_arms=num_arms_training,
+        morph_mode=training.morphology.morph_mode,
+        segments_per_arm=env_morphology.segments_per_arm,
+        agent_indices=agent_indices,
+    )
+
+    # 3. Build environment
+    backend = Backend.MJC
+    factory = BrittleStarEnvFactory()
+    raw_env = factory.create_environment(
+        backend,
+        env_morphology,
+        training.arena,
+        training.environment,
+    )
+    env = BrittleStarEnv(
+        raw_env,
+        backend=backend,
+        config=training.environment,
+        morphology_config=env_morphology,
+    )
+
+    # Calculate the action dimension the model was trained with
+    training_total_actions = sum(training.morphology.segments_per_arm) * 2
+    trained_action_dim = training_total_actions // needed_copies
+
+    # 4. Load policy
+    message_passing_steps = (metadata.get("architecture", {}) or {}).get("message_passing_steps")
+    if message_passing_steps is None:
+        message_passing_steps = 4
+    message_passing_steps = int(message_passing_steps)
+
+    adj_matrix = None
+    if training.morphology.morph_mode != MorphMode.CENTRALIZED:
+        adj_matrix = build_adjacency(
+            training.morphology.segments_per_arm, training.morphology.morph_mode
+        )
+
+        override_segs = env_morphology.segments_per_arm
+        if training.morphology.morph_mode in (MorphMode.FULLY_CONNECTED, MorphMode.RING):
+            for i, segs in enumerate(override_segs):
+                if segs == 0 and i < adj_matrix.shape[0]:
+                    adj_matrix = adj_matrix.at[i, :].set(0)
+                    adj_matrix = adj_matrix.at[:, i].set(0)
+        elif training.morphology.morph_mode == MorphMode.SEGMENT:
+            for i, segs in enumerate(override_segs):
+                if segs == 0 and i < num_arms_training:
+                    adj_matrix = adj_matrix.at[i, :].set(0)
+                    adj_matrix = adj_matrix.at[:, i].set(0)
+
+            idx = 0
+            for arm_idx, seg_count in enumerate(training.morphology.segments_per_arm):
+                if override_segs[arm_idx] == 0:
+                    for i in range(seg_count):
+                        seg_node = num_arms_training + idx + i
+                        if seg_node < adj_matrix.shape[0]:
+                            adj_matrix = adj_matrix.at[seg_node, :].set(0)
+                            adj_matrix = adj_matrix.at[:, seg_node].set(0)
+                idx += seg_count
+
+    policy = PolicyAgent.from_checkpoint(
+        model_path,
+        action_dim=trained_action_dim,
+        obs_processor=obs_processor,
+        message_passing_steps=message_passing_steps,
+        adj_matrix=adj_matrix,
+    )
+
+    # 5. Build action clipping and masks
+    action_mask = np.asarray(padding_masks["mask_2x"])
+
+    action_space = getattr(raw_env, "action_space", None)
+    action_low = (
+        None if action_space is None else np.asarray(action_space.low, dtype=np.float32).ravel()
+    )
+    action_high = (
+        None if action_space is None else np.asarray(action_space.high, dtype=np.float32).ravel()
+    )
+
+    return EvalEnvBundle(
+        env=env,
+        policy=policy,
+        action_low=action_low,
+        action_high=action_high,
+        action_mask=action_mask,
+        segments_per_arm=env_morphology.segments_per_arm,
+        num_active_arms=sum(1 for s in env_morphology.segments_per_arm if s > 0),
+        architecture=env_morphology.morph_mode.name,
+    )

src/brittle_star_project/evaluation/evaluate.py

@@ -0,0 +1,58 @@
+"""MJC-based (CPU) checkpoint evaluation.
+
+This module provides the CPU-bound evaluation path using the standard MJC backend.
+It is primarily used by the `evaluate_checkpoints` CLI to compute metrics and
+render videos.
+"""
+
+from pathlib import Path
+
+import numpy as np
+
+from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
+from brittle_star_project.environment.obs_processing import create_obs_processor
+from brittle_star_project.evaluation.policy import PolicyAgent
+from brittle_star_project.evaluation.rollout import EpisodeResult, rollout_headless
+
+
+def evaluate_policy(
+    env: BrittleStarJaxEnvWrapper,
+    policy_path: str | Path,
+    seed: int,
+    max_steps: int,
+) -> EpisodeResult:
+    """Evaluate a trained policy in a CPU-bound environment.
+
+    Args:
+        env: Initialised CPU environment (MJC backend).
+        policy_path: Path to the `.cleanrl_model` weights file.
+        seed: Random seed for environment reset.
+        max_steps: Maximum number of control steps.
+
+    Returns:
+        Structured result containing return, length, and distance metrics.
+    """
+    obs_processor = create_obs_processor(
+        bounds_dict=env.cfg.obs_bounds.to_bounds_dict(),
+        padding_masks=env.padding_masks,
+    )
+
+    action_dim = env.single_action_space.shape[0]
+
+    policy = PolicyAgent.from_checkpoint(
+        model_path=Path(policy_path),
+        action_dim=action_dim,
+        obs_processor=obs_processor,
+    )
+
+    action_low = np.asarray(env.single_action_space.low, dtype=np.float32)
+    action_high = np.asarray(env.single_action_space.high, dtype=np.float32)
+
+    return rollout_headless(
+        env=env,
+        policy=policy,
+        seed=seed,
+        max_steps=max_steps,
+        action_low=action_low,
+        action_high=action_high,
+    )

src/brittle_star_project/evaluation/evaluate_mjx.py

@@ -0,0 +1,258 @@
+"""MJX-based headless checkpoint evaluation.
+
+This module provides a fast, JIT-compiled evaluation path using the MJX
+(JAX-accelerated MuJoCo) backend. It is intended for evaluating checkpoints
+*during* or *after* a training run, where the environment and policy are
+already fully initialised.
+
+The key functions are:
+
+- `build_eval_rollout_fn` — builds and JIT-compiles a single-episode rollout function from the
+  training environment and policy components.
+- `evaluate_checkpoint_mjx` — runs that function for a given set of parameters and returns a typed
+  `CheckpointEvalResult`.
+- `append_checkpoint_eval_row` — persists the result to the run's
+  `metrics/checkpoint_evaluation.csv`, migrating old schemas automatically.
+"""
+
+from __future__ import annotations
+
+import csv
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable
+
+import jax
+import jax.numpy as jnp
+
+
+@dataclass
+class CheckpointEvalResult:
+    """Structured result from a single MJX checkpoint evaluation episode."""
+
+    steps: int
+    """Number of control steps taken (≤ max_steps)."""
+
+    reached_target: bool
+    """Whether the robot reached the target (terminated) before max_steps."""
+
+    eval_return: float
+    """Accumulated shaped reward over the episode."""
+
+    final_xy_dist: float
+    """XY distance to target at episode end. 0.0 when ``reached_target`` is True."""
+
+    initial_xy_dist: float
+    """XY distance to target at episode start."""
+
+
+def build_eval_rollout_fn(
+    *,
+    env: Any,
+    obs_processor: Callable,
+    sensor_apply: Callable,
+    actor_apply: Callable,
+    message_passer_apply: Callable | None = None,
+    action_low: jnp.ndarray,
+    action_high: jnp.ndarray,
+    reward_fn: Callable,
+) -> Callable:
+    """Build and JIT-compile a single-episode MJX evaluation rollout.
+
+    All outputs are JAX arrays. Convert to Python scalars before logging.
+
+    Args:
+        env: The training environment wrapper. Must expose `env.raw` with
+            `reset` and `step` methods compatible with `jax.vmap`.
+        obs_processor: Observation normalisation / padding callable, as
+            returned by `create_obs_processor`.
+        sensor_apply: The sensor network's `apply` method (JIT-compiled).
+        actor_apply: The actor network's `apply` method (JIT-compiled).
+        message_passer_apply: Optional message-passing module apply method.
+            When provided, it is applied between the sensor and actor, using
+            `params["message_passer_params"]`.
+        action_low: Per-joint action lower bound (JAX array, shape `(action_dim,)`).
+        action_high: Per-joint action upper bound (JAX array, shape `(action_dim,)`).
+        reward_fn: Shaped reward function with signature
+            `reward_fn(env_state, next_env_state) -> jnp.ndarray`.
+            Typically, the module-level `reward_fn` from `PPOTrainer`.
+
+    Returns:
+        A JIT-compiled callable that runs one deterministic evaluation episode.
+    """
+    # vmap over a batch of 1 so the MJX API is satisfied without any
+    # extra bookkeeping in the caller.
+    reset_1 = jax.vmap(env.raw.reset)
+    step_1 = jax.vmap(env.raw.step)
+
+    def _eval_rollout(params: dict, seed: int, max_steps: int):
+        rng = jax.random.PRNGKey(seed)
+        rngs = jnp.asarray(jax.random.split(rng, 1))
+        state = reset_1(rng=rngs)
+
+        initial_xy_dist = jnp.squeeze(state.observations["xy_distance_to_target"])
+
+        t0 = jnp.asarray(0, dtype=jnp.int32)
+        done0 = jnp.squeeze(state.terminated | state.truncated)
+        return0 = jnp.asarray(0.0, dtype=jnp.float32)
+
+        def cond(carry):
+            t, _state, done, _return_ = carry
+            return jnp.logical_and(t < max_steps, jnp.logical_not(done))
+
+        def body(carry):
+            t, state, _done, return_ = carry
+
+            obs = obs_processor(state.observations)
+            hidden = sensor_apply(params["sensor_params"], obs)
+            if message_passer_apply is not None:
+                mp_params = params["message_passer_params"]
+                hidden = jax.vmap(lambda x: message_passer_apply(mp_params, x))(hidden)
+            mean, _log_std = actor_apply(params["actor_params"], hidden)
+
+            # Deterministic action: use the actor mean, no exploration noise.
+            flat_mean = mean.reshape(mean.shape[0], -1)
+            action = jnp.clip(flat_mean, action_low, action_high)
+            next_state = step_1(state=state, action=action)
+
+            shaped_reward = reward_fn(state, next_state)
+            return_ = return_ + jnp.squeeze(shaped_reward)
+
+            done_next = jnp.squeeze(next_state.terminated | next_state.truncated)
+            return (t + 1, next_state, done_next, return_)
+
+        t, final_state, _done, return_ = jax.lax.while_loop(cond, body, (t0, state, done0, return0))
+
+        reached_target = jnp.squeeze(final_state.terminated)
+        final_xy_dist_raw = jnp.squeeze(final_state.observations["xy_distance_to_target"])
+        # Clamp to 0 when the target was reached so downstream consumers
+        # don't have to special-case "terminated" themselves.
+        final_xy_dist = jnp.where(reached_target, 0.0, final_xy_dist_raw)
+
+        return t, reached_target, return_, final_xy_dist, initial_xy_dist
+
+    return jax.jit(_eval_rollout)
+
+
+def evaluate_checkpoint_mjx(
+    eval_fn: Callable,
+    params: dict,
+    *,
+    seed: int,
+    max_steps: int,
+) -> CheckpointEvalResult:
+    """Run one deterministic evaluation episode and return typed metrics.
+
+    Args:
+        eval_fn: A JIT-compiled function as returned by `build_eval_rollout_fn`.
+        params: Agent parameter dict (e.g. ``agent_state.params``).
+        seed: Random seed for environment reset (controls target placement).
+        max_steps: Maximum number of control steps before the episode is cut off.
+
+    Returns:
+        A `CheckpointEvalResult` with all JAX arrays converted to
+        plain Python scalars.
+    """
+    steps, reached, eval_return, final_xy_dist, initial_xy_dist = eval_fn(params, seed, max_steps)
+    return CheckpointEvalResult(
+        steps=int(steps),
+        reached_target=bool(reached),
+        eval_return=float(eval_return),
+        final_xy_dist=float(final_xy_dist),
+        initial_xy_dist=float(initial_xy_dist),
+    )
+
+
+_FIELDNAMES = [
+    "checkpoint",
+    "trained_timesteps",
+    "eval_steps",
+    "eval_return",
+    "final_xy_dist",
+    "initial_xy_dist",
+    "reached_target",
+]
+
+
+def _migrate_csv_if_needed(csv_path: Path) -> None:
+    """Rewrite the CSV with the canonical field names if the schema changed.
+
+    Best-effort: any exception is silently swallowed so that a schema mismatch
+    never causes a training crash.
+    """
+    try:
+        with open(csv_path, "r", newline="") as f:
+            header = next(csv.reader(f), None)
+
+        if header is None or list(header) == _FIELDNAMES:
+            return  # Nothing to migrate.
+
+        migrated_rows: list[dict[str, Any]] = []
+        with open(csv_path, "r", newline="") as f:
+            for row in csv.DictReader(f):
+                migrated_rows.append(
+                    {
+                        "checkpoint": row.get("checkpoint", row.get("iteration")),
+                        "trained_timesteps": row.get("trained_timesteps"),
+                        "eval_steps": row.get("eval_steps", row.get("steps_to_target")),
+                        "eval_return": row.get("eval_return"),
+                        "final_xy_dist": row.get("final_xy_dist"),
+                        "initial_xy_dist": row.get("initial_xy_dist"),
+                        "reached_target": row.get("reached_target"),
+                    }
+                )
+
+        with open(csv_path, "w", newline="") as f:
+            writer = csv.DictWriter(f, fieldnames=_FIELDNAMES)
+            writer.writeheader()
+            writer.writerows(migrated_rows)
+    except Exception:
+        pass  # Never crash training on a migration issue.
+
+
+def append_checkpoint_eval_row(
+    run_dir: str | Path,
+    *,
+    iteration: int,
+    trained_timesteps: int,
+    result: CheckpointEvalResult,
+) -> Path:
+    """Append one evaluation row to `<run_dir>/metrics/checkpoint_evaluation.csv`.
+
+    Creates the file (including the `metrics/` directory) if it does not yet
+    exist. Migrates the file to the current schema if the header has changed.
+
+    Args:
+        run_dir: Root directory of the training run (Hydra's output dir).
+        iteration: Training iteration number, used as the checkpoint identifier.
+        trained_timesteps: Total environment steps taken at this checkpoint.
+        result: Evaluation result as returned by `evaluate_checkpoint_mjx`.
+
+    Returns:
+        Absolute path to the CSV file (useful for W&B sync).
+    """
+    metrics_dir = Path(run_dir) / "metrics"
+    metrics_dir.mkdir(parents=True, exist_ok=True)
+    csv_path = metrics_dir / "checkpoint_evaluation.csv"
+
+    if csv_path.exists():
+        _migrate_csv_if_needed(csv_path)
+
+    file_exists = csv_path.exists()
+    with open(csv_path, "a", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=_FIELDNAMES)
+        if not file_exists:
+            writer.writeheader()
+        writer.writerow(
+            {
+                "checkpoint": int(iteration),
+                "trained_timesteps": int(trained_timesteps),
+                "eval_steps": result.steps,
+                "eval_return": result.eval_return,
+                "final_xy_dist": result.final_xy_dist,
+                "initial_xy_dist": result.initial_xy_dist,
+                "reached_target": result.reached_target,
+            }
+        )
+
+    return csv_path

src/brittle_star_project/evaluation/policy.py

@@ -0,0 +1,168 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Protocol
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from brittle_star_project.MLPs.routing import apply_per_node
+from brittle_star_project.evaluation.checkpoint import load_params
+
+
+class ControlPolicy(Protocol):
+    """Protocol for any policy that can produce actions from observations."""
+
+    def act(self, *, observations: dict[str, Any]) -> np.ndarray: ...
+
+
+class PolicyAgent:
+    """Wraps a trained Flax actor for deterministic inference."""
+
+    def __init__(
+        self,
+        *,
+        sensor_params: Any,
+        actor_params: Any,
+        message_passer_params: Any | None = None,
+        message_passing_steps: int | None = None,
+        adj_matrix: Any | None = None,
+        action_dim: int,
+        obs_processor: Any,
+    ) -> None:
+        from brittle_star_project.MLPs.mlps import (
+            Actor,
+            GenericDenseLayersWithActivation,
+            MessagePasser,
+        )
+
+        # Infer layer sizes from params
+        try:
+            dense_params = (
+                sensor_params.get("params", {})
+                if isinstance(sensor_params, dict)
+                else sensor_params["params"]
+            )
+        except Exception:
+            dense_params = sensor_params
+
+        layer_sizes = []
+        idx = 0
+        while True:
+            key = f"Dense_{idx}"
+            if key not in dense_params:
+                break
+
+            layer_sizes.append(int(np.asarray(dense_params[key]["kernel"]).shape[-1]))
+            idx += 1
+
+        if not layer_sizes:
+            raise ValueError("Could not infer Dense_* layers from sensor params")
+
+        self._sensor = GenericDenseLayersWithActivation(layer_sizes=layer_sizes)
+        self._actor = Actor(action_dim=action_dim)
+
+        self._message_passer = None
+        if message_passer_params is not None and not (
+            isinstance(message_passer_params, dict) and len(message_passer_params) == 0
+        ):
+            if message_passing_steps is None or adj_matrix is None:
+                raise ValueError(
+                    "Checkpoint contains message_passer_params but PolicyAgent was not given "
+                    "message_passing_steps and adj_matrix. Pass these when constructing the agent "
+                    "so decentralized evaluation matches training."
+                )
+
+            hidden_dim = int(layer_sizes[-1])
+            self._message_passer = MessagePasser(
+                hidden_dim=hidden_dim,
+                num_propagation_steps=int(message_passing_steps),
+                adj_matrix=jnp.asarray(adj_matrix),
+            )
+            self._message_passer.apply = jax.jit(self._message_passer.apply)
+        self._sensor.apply = jax.jit(self._sensor.apply)
+        self._actor.apply = jax.jit(self._actor.apply)
+        self._params = {
+            "sensor_params": sensor_params,
+            "actor_params": actor_params,
+            "message_passer_params": message_passer_params,
+        }
+        self._obs_processor = obs_processor
+
+    @classmethod
+    def from_params(
+        cls,
+        *,
+        sensor_params: Any,
+        actor_params: Any,
+        message_passer_params: Any | None = None,
+        message_passing_steps: int | None = None,
+        adj_matrix: Any | None = None,
+        action_dim: int,
+        obs_processor: Any,
+    ) -> "PolicyAgent":
+        """Construct a PolicyAgent directly from in-memory parameters."""
+        return cls(
+            sensor_params=sensor_params,
+            actor_params=actor_params,
+            message_passer_params=message_passer_params,
+            message_passing_steps=message_passing_steps,
+            adj_matrix=adj_matrix,
+            action_dim=action_dim,
+            obs_processor=obs_processor,
+        )
+
+    def set_params(
+        self,
+        *,
+        sensor_params: Any,
+        actor_params: Any,
+        message_passer_params: Any | None = None,
+    ) -> None:
+        """Update parameters for evaluation without rebuilding the model."""
+        self._params["sensor_params"] = sensor_params
+        self._params["actor_params"] = actor_params
+        self._params["message_passer_params"] = message_passer_params
+
+    @classmethod
+    def from_checkpoint(
+        cls,
+        model_path: Path,
+        *,
+        action_dim: int,
+        obs_processor: Any,
+        message_passing_steps: int | None = None,
+        adj_matrix: Any | None = None,
+    ) -> "PolicyAgent":
+        """Load params from .flax and construct the agent."""
+        params = load_params(model_path)
+
+        return cls(
+            sensor_params=params["sensor_params"],
+            actor_params=params["actor_params"],
+            message_passer_params=params.get("message_passer_params"),
+            message_passing_steps=message_passing_steps,
+            adj_matrix=adj_matrix,
+            action_dim=action_dim,
+            obs_processor=obs_processor,
+        )
+
+    def act(self, *, observations: dict[str, Any]) -> np.ndarray:
+        """Return deterministic action (actor mean, no exploration noise)."""
+        batched_obs = jax.tree.map(lambda x: jnp.asarray(x)[None, ...], observations)
+        obs = self._obs_processor(batched_obs)
+
+        hidden = apply_per_node(self._sensor.apply, self._params["sensor_params"], obs)
+
+        if self._message_passer is not None:
+            mp_params = self._params.get("message_passer_params")
+            if mp_params is None or (isinstance(mp_params, dict) and len(mp_params) == 0):
+                raise ValueError(
+                    "PolicyAgent has a message passer but message_passer_params are missing/empty."
+                )
+            hidden = jax.vmap(lambda x: self._message_passer.apply(mp_params, x))(hidden)
+
+        mean, _log_std = apply_per_node(self._actor.apply, self._params["actor_params"], hidden)
+
+        return np.asarray(mean, dtype=np.float32).ravel()

src/brittle_star_project/evaluation/rollout.py

@@ -0,0 +1,168 @@
+from __future__ import annotations
+
+import itertools
+import time
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+from brittle_star_project import BrittleStarEnv
+from brittle_star_project.evaluation.policy import ControlPolicy
+
+
+@dataclass
+class EpisodeResult:
+    return_: float
+    length: int
+    reached_target: bool
+    final_xy_dist: float | None
+    initial_target_distance: float | None
+
+
+def _get_observations(state: Any) -> dict[str, Any] | None:
+    return getattr(state, "observations", None)
+
+
+def _get_xy_distance_to_target(observations: dict[str, Any]) -> float | None:
+    return float(np.asarray(observations["xy_distance_to_target"]).reshape(-1)[0])
+
+
+def _target_reached(*, state: Any) -> bool:
+    return bool(getattr(state, "terminated", False) or getattr(state, "truncated", False))
+
+
+def _maybe_clip_action(
+    action: np.ndarray,
+    low: np.ndarray | None,
+    high: np.ndarray | None,
+) -> np.ndarray:
+    if low is None or high is None:
+        return action
+    low = np.asarray(low, dtype=np.float32).ravel()
+    high = np.asarray(high, dtype=np.float32).ravel()
+    if low.shape != action.shape or high.shape != action.shape:
+        return action
+    return np.clip(action, low, high)
+
+
+def rollout_headless(
+    *,
+    env: BrittleStarEnv,
+    policy: ControlPolicy,
+    seed: int,
+    max_steps: int,
+    action_low: np.ndarray | None,
+    action_high: np.ndarray | None,
+    action_mask: np.ndarray | None = None,
+) -> EpisodeResult:
+    """Run an episode headlessly and return the result."""
+    state = env.reset(seed=seed)
+
+    ep_return = 0.0
+    observations = _get_observations(state)
+    prev_dist = _get_xy_distance_to_target(observations) if observations else None
+    initial_target_distance = prev_dist
+    reached_target = _target_reached(state=state)
+
+    steps = 0
+    for _ in range(int(max_steps)):
+        obs_dict = observations or {}
+
+        action = policy.act(observations=obs_dict)
+        if action_mask is not None:
+            action = action[action_mask]
+        action = _maybe_clip_action(action, action_low, action_high)
+
+        state = env.step(state=state, action=action)
+        steps += 1
+
+        observations = _get_observations(state)
+        cur_dist = _get_xy_distance_to_target(observations) if observations else None
+        if prev_dist is not None and cur_dist is not None:
+            ep_return += prev_dist - cur_dist
+        prev_dist = cur_dist
+
+        reached_target = _target_reached(state=state)
+        if reached_target:
+            break
+
+    final_dist = _get_xy_distance_to_target(observations) if observations else None
+    return EpisodeResult(
+        return_=ep_return,
+        length=steps,
+        reached_target=reached_target,
+        final_xy_dist=final_dist,
+        initial_target_distance=initial_target_distance,
+    )
+
+
+def rollout_viewer(
+    *,
+    env: BrittleStarEnv,
+    policy: ControlPolicy,
+    seed: int,
+    state: Any,
+    control_dt: float,
+    max_steps: int | None,
+    action_low: np.ndarray | None,
+    action_high: np.ndarray | None,
+    action_mask: np.ndarray | None = None,
+) -> None:
+    """Run an episode using the interactive MuJoCo viewer."""
+    import mujoco.viewer
+
+    model = state.mj_model
+    data = state.mj_data
+
+    episode_return = 0.0
+    observations = _get_observations(state)
+
+    prev_dist = _get_xy_distance_to_target(observations) if observations else None
+    reached_target = _target_reached(state=state)
+
+    steps = 0
+    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 _ in step_iter:
+            if not viewer.is_running():
+                break
+            step_start = time.time()
+
+            obs_dict = observations or {}
+
+            action = policy.act(observations=obs_dict)
+            if action_mask is not None:
+                action = action[action_mask]
+            action = _maybe_clip_action(action, action_low, action_high)
+
+            with viewer.lock():
+                state = env.step(state=state, action=action)
+
+            if not viewer.is_running():
+                break
+            viewer.sync()
+
+            steps += 1
+
+            observations = _get_observations(state)
+            cur_dist = _get_xy_distance_to_target(observations) if observations else None
+            if prev_dist is not None and cur_dist is not None:
+                episode_return += prev_dist - cur_dist
+            prev_dist = cur_dist
+
+            reached_target = _target_reached(state=state)
+            if reached_target:
+                break
+
+            remaining = control_dt - (time.time() - step_start)
+            if remaining > 0:
+                time.sleep(remaining)
+
+    dist = _get_xy_distance_to_target(observations) if observations else None
+    dist_str = "n/a" if dist is None else f"{dist:.3f}"
+    print(
+        "episode done: "
+        f"return={episode_return:.6f}, len={steps}, "
+        f"target_reached={reached_target}, final_xy_dist={dist_str}"
+    )

src/brittle_star_project/evaluation/video.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+import datetime
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+from brittle_star_project import BrittleStarEnv
+from brittle_star_project.evaluation.policy import ControlPolicy
+from brittle_star_project.evaluation.rollout import (
+    EpisodeResult,
+    _get_observations,
+    _get_xy_distance_to_target,
+    _target_reached,
+    _maybe_clip_action,
+)
+
+
+def create_evaluation_dir(model_path: Path) -> Path:
+    """Create a unique timestamped directory for saving evaluation results."""
+    timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    eval_dir = model_path.parent / f"{model_path.stem}_evaluations" / f"eval_{timestamp}"
+    eval_dir.mkdir(parents=True, exist_ok=True)
+    return eval_dir
+
+
+def save_evaluation_metadata(
+    eval_dir: Path,
+    *,
+    morphology_override_path: str | None,
+    seed: int,
+    max_steps: int | None,
+    result: EpisodeResult,
+) -> None:
+    """Save metadata about the evaluation run."""
+    metadata = {
+        "timestamp": datetime.datetime.now().isoformat(),
+        "morphology_override": morphology_override_path,
+        "seed": seed,
+        "max_steps": max_steps,
+        "result": {
+            "return": float(result.return_),
+            "length": int(result.length),
+            "reached_target": bool(result.reached_target),
+            "final_xy_dist": float(result.final_xy_dist)
+            if result.final_xy_dist is not None
+            else None,
+        },
+    }
+    with open(eval_dir / "evaluation_metadata.yaml", "w") as f:
+        yaml.safe_dump(metadata, f, sort_keys=False)
+
+
+def record_episode(
+    *,
+    env: BrittleStarEnv,
+    policy: ControlPolicy,
+    seed: int,
+    max_steps: int,
+    action_low: np.ndarray | None,
+    action_high: np.ndarray | None,
+    action_mask: np.ndarray | None = None,
+    output_path: Path,
+    camera_id: int = 1,
+    fps: int = 60,
+    width: int = 640,
+    height: int = 480,
+) -> EpisodeResult:
+    """Run an episode headlessly and record a video using MuJoCo's Renderer and imageio.
+
+    Args:
+        env: The environment.
+        policy: The policy agent.
+        seed: Random seed.
+        max_steps: Maximum number of steps.
+        action_low: Minimum action values.
+        action_high: Maximum action values.
+        action_mask: Boolean mask for the actions.
+        output_path: Where to save the .mp4 file.
+        camera_id: Camera index to use for rendering (1 is usually close-up).
+        fps: Frames per second for the video.
+        width: Video width.
+        height: Video height.
+    """
+    try:
+        import imageio
+        import mujoco
+    except ImportError as e:
+        raise ImportError(
+            "Video recording requires 'imageio' and 'mujoco'. "
+            "Please install the evaluation dependencies: `uv pip install .[evaluation]`"
+        ) from e
+
+    state = env.reset(seed=seed)
+    model = state.mj_model
+    data = state.mj_data
+
+    renderer = mujoco.Renderer(model, width=width, height=height)
+    ep_return = 0.0
+    observations = _get_observations(state)
+    prev_dist = _get_xy_distance_to_target(observations) if observations else None
+    initial_dist = prev_dist
+    reached_target = _target_reached(state=state)
+
+    frames = []
+    steps = 0
+
+    for _ in range(int(max_steps)):
+        # Capture frame
+        renderer.update_scene(data, camera=camera_id)
+        frames.append(renderer.render())
+
+        # Step environment
+        obs_dict = observations or {}
+        action = policy.act(observations=obs_dict)
+        if action_mask is not None:
+            action = action[action_mask]
+        action = _maybe_clip_action(action, action_low, action_high)
+
+        state = env.step(state=state, action=action)
+        steps += 1
+
+        observations = _get_observations(state)
+        cur_dist = _get_xy_distance_to_target(observations) if observations else None
+        if prev_dist is not None and cur_dist is not None:
+            ep_return += prev_dist - cur_dist
+        prev_dist = cur_dist
+
+        reached_target = _target_reached(state=state)
+        if reached_target:
+            break
+
+    # Capture final frame
+    renderer.update_scene(data, camera=camera_id)
+    frames.append(renderer.render())
+    renderer.close()
+
+    # Save video
+    imageio.mimsave(str(output_path), frames, fps=fps)
+
+    final_dist = _get_xy_distance_to_target(observations) if observations else None
+    return EpisodeResult(
+        return_=ep_return,
+        length=steps,
+        reached_target=reached_target,
+        final_xy_dist=final_dist,
+        initial_target_distance=initial_dist,
+    )

src/brittle_star_project/ppo.py

@@ -1,14 +1,27 @@
 from functools import partial
 
-import flax
 import jax
 import jax.numpy as jnp
+from jax import debug
+from flax.core import FrozenDict
+from experiment_logger import get_logger
+from brittle_star_project.utils import logged_jit
+
+logger = get_logger()
 
 
 # Chose to use a class as it seemed the easiest way to integrate the CleanRL code style
 # with our need to seperate concerns
 class PPO:
-    def __init__(self, args, sensor, actor, critic, feature_extractor, message_passer=None):
+    def __init__(
+        self,
+        args,
+        sensor_apply,
+        actor_apply,
+        critic_apply,
+        feature_extractor_apply,
+        message_passer=None,
+    ):
         self.args = args
 
         if not message_passer:
@@ -18,10 +31,10 @@ class PPO:
             partial(
                 ppo_loss,
                 args=args,
-                sensor_apply=sensor.apply,
-                actor_apply=actor.apply,
-                critic_apply=critic.apply,
-                feature_extractor_apply=feature_extractor.apply,
+                sensor_apply=sensor_apply,
+                actor_apply=actor_apply,
+                critic_apply=critic_apply,
+                feature_extractor_apply=feature_extractor_apply,
                 message_passer=message_passer,
             ),
             has_aux=True,
@@ -29,8 +42,14 @@ class PPO:
 
     # This PPO class should be initialized only once,
     # or this function will need to recompile
-    @partial(jax.jit, static_argnums=0)
+    @partial(logged_jit, static_argnums=0)
     def update_ppo(self, agent_state, storage, key):
+        debug.callback(logger.debug, f"[PPO] storage.obs shape: {storage.obs.shape}")
+        debug.callback(logger.debug, f"[PPO] storage.actions shape: {storage.actions.shape}")
+        debug.callback(logger.debug, f"[PPO] storage.logprobs shape: {storage.logprobs.shape}")
+        debug.callback(logger.debug, f"[PPO] storage.advantages shape: {storage.advantages.shape}")
+        debug.callback(logger.debug, f"[PPO] storage.returns shape: {storage.returns.shape}")
+
         args = self.args
         ppo_loss_grad_fn = self.ppo_loss_grad_fn
 
@@ -49,6 +68,16 @@ class PPO:
             shuffled_storage = jax.tree.map(convert_data, flatten_storage)
 
             def update_minibatch(agent_state, minibatch):
+                debug.callback(logger.debug, f"[PPO] minibatch.obs: {minibatch.obs.shape}")
+                debug.callback(logger.debug, f"[PPO] minibatch.actions: {minibatch.actions.shape}")
+                debug.callback(
+                    logger.debug, f"[PPO] minibatch.logprobs: {minibatch.logprobs.shape}"
+                )
+                debug.callback(
+                    logger.debug, f"[PPO] minibatch.advantages: {minibatch.advantages.shape}"
+                )
+                debug.callback(logger.debug, f"[PPO] minibatch.returns: {minibatch.returns.shape}")
+
                 (loss, (pg_loss, v_loss, entropy_loss, approx_kl)), grads = ppo_loss_grad_fn(
                     agent_state.params,
                     minibatch.obs,
@@ -58,19 +87,12 @@ class PPO:
                     minibatch.returns,
                 )
                 agent_state = agent_state.apply_gradients(grads=grads)
-                return agent_state, (
-                    loss,
-                    pg_loss,
-                    v_loss,
-                    entropy_loss,
-                    approx_kl,
-                    grads,
-                )
+                return agent_state, (loss, pg_loss, v_loss, entropy_loss, approx_kl)
 
             agent_state, metrics = jax.lax.scan(update_minibatch, agent_state, shuffled_storage)
             return (agent_state, key), metrics
 
-        (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads) = jax.lax.scan(
+        (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl) = jax.lax.scan(
             update_epoch, (agent_state, key), (), length=args.update_epochs
         )
         return agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key
@@ -84,27 +106,45 @@ that are now not in the same scope
 """
 
 
-@partial(jax.jit, static_argnums=(0, 1, 2, 3, 4))
+@partial(logged_jit, static_argnums=(0, 1, 2, 3, 4))
 def get_action_and_value(
     sensor_apply,
     actor_apply,
     message_passer,
     critic_apply,
     feature_extractor_apply,
-    params: flax.core.FrozenDict,
+    params: FrozenDict,
     x: jnp.ndarray,
     action: jnp.ndarray,
 ):
     hidden_sensor = sensor_apply(params["sensor_params"], x)
     hidden_critic = feature_extractor_apply(params["feature_extractor_params"], x)
-    hidden_sensor = message_passer(hidden_sensor)
+
+    # only apply message passing in decentralized context
+    if message_passer is not None:
+        hidden_sensor = message_passer(params["message_passer_params"], hidden_sensor)
+
+    debug.callback(logger.debug, f"[SHAPE] hidden_sensor: {hidden_sensor.shape}")
+    debug.callback(logger.debug, f"[SHAPE] hidden_critic: {hidden_critic.shape}")
+
     mean, log_std = actor_apply(params["actor_params"], hidden_sensor)
+
+    debug.callback(logger.debug, f"[SHAPE] mean: {mean.shape}")
+    debug.callback(logger.debug, f"[SHAPE] log_std: {log_std.shape}")
+    debug.callback(logger.debug, f"[SHAPE] action: {action.shape}")
+
     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)
-    entropy = (0.5 + 0.5 * jnp.log(2 * jnp.pi) + log_std).sum(-1)
+    logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi))
+    debug.callback(logger.debug, f"[SHAPE] logprob pre-sum: {logprob.shape}")
+
+    logprob = logprob.sum(axis=(-2, -1))
+    debug.callback(logger.debug, f"[SHAPE] logprob final: {logprob.shape}")
+
+    entropy = (0.5 + 0.5 * jnp.log(2 * jnp.pi) + log_std).sum(axis=(-2, -1))
     value = critic_apply(params["critic_params"], hidden_critic).squeeze(-1)
+    debug.callback(logger.debug, f"[SHAPE] value: {value.shape}")
 
     return logprob, entropy, value
 
@@ -150,7 +190,7 @@ def ppo_loss(
     return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl))
 
 
-def identity(hidden):
+def identity(_, hidden):
     """
     Used for seamless jax integration,
     avoids having branching inside jitted function,

src/brittle_star_project/render/__init__.py

@@ -1,3 +0,0 @@
-from .renderer import simulate_policy, SimulationConfig, ControlPolicy
-
-__all__ = ["simulate_policy", "SimulationConfig", "ControlPolicy"]

src/brittle_star_project/render/renderer.py

@@ -1,78 +0,0 @@
-from __future__ import annotations
-
-import time
-from dataclasses import dataclass
-from typing import Any, Protocol
-
-import numpy as np
-
-
-@dataclass
-class SimulationConfig:
-    realtime: bool = True
-    seed: int = 0
-
-
-class ControlPolicy(Protocol):
-    def act(self, *, obs: np.ndarray | None = None, t: float = 0.0) -> np.ndarray: ...
-
-
-def _default_observations(data: Any) -> np.ndarray:
-    qpos = np.asarray(data.qpos, dtype=np.float32).ravel()
-    qvel = np.asarray(data.qvel, dtype=np.float32).ravel()
-    return np.concatenate([qpos, qvel], axis=0)
-
-
-def simulate_policy(
-    policy: ControlPolicy,
-    config: SimulationConfig,
-    state: Any | None = None,
-) -> None:
-    """Open MuJoCo's native viewer and step using actions from a policy.
-
-    This path drives MuJoCo physics directly (mj_step) and uses the policy output
-    as `data.ctrl`.
-    """
-
-    import mujoco.viewer
-
-    if state is None:
-        raise ValueError("A valid environment state must be provided.")
-
-    model = state.mj_model
-    data = state.mj_data
-
-    start = time.time()
-    with mujoco.viewer.launch_passive(model, data) as viewer:
-        while viewer.is_running():
-            step_start = time.time()
-
-            t = time.time() - start
-
-            # Input vector for the policy
-            # TODO: custom input
-            obs = _default_observations(data)
-
-            # Policy action
-            ctrl = policy.act(obs=obs, t=t)
-
-            # Check if the policy output vector give an input for each actuator (nu)
-            # TODO: what if model trained on full morphology but we want to test on a damaged one?
-            # (nu mismatch)
-            if model.nu > 0:
-                ctrl = np.asarray(ctrl, dtype=np.float32).ravel()
-                if ctrl.shape != (model.nu,):
-                    raise ValueError(
-                        f"Policy returned ctrl shape {ctrl.shape}, expected ({model.nu},)"
-                    )
-                data.ctrl[:] = ctrl
-
-            # Step the simulation and update the viewer
-            mujoco.mj_step(model, data)
-            viewer.sync()
-
-            # If we're running in realtime mode, sleep to maintain real-time pacing.
-            if config.realtime:
-                remaining = model.opt.timestep - (time.time() - step_start)
-                if remaining > 0:
-                    time.sleep(remaining)

src/brittle_star_project/trainers/PPOTrainer.py

@@ -3,12 +3,13 @@ import random
 import time
 from dataclasses import asdict, dataclass
 from functools import partial
-from typing import Any
+from typing import Any, Optional
 
 import jax
 import jax.numpy as jnp
 import numpy as np
 import optax
+import flax.linen as nn
 from flax.training.train_state import TrainState
 
 from experiment_logger import get_logger
@@ -16,39 +17,32 @@ from experiment_logger import get_logger
 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.environment.obs_processing import create_obs_processor
+from brittle_star_project.evaluation.evaluate_mjx import (
+    append_checkpoint_eval_row,
+    build_eval_rollout_fn,
+    evaluate_checkpoint_mjx,
+)
+from brittle_star_project.MLPs.routing import apply_per_node
 from brittle_star_project.MLPs.mlps import (
     Actor,
     AgentParams,
     GenericDenseLayersWithActivation,
+    MessagePasser,
     OneDenseLayerMLP,
     Storage,
 )
+from brittle_star_project.MLPs.adjancency_builder import build_adjacency
 from brittle_star_project.ppo import PPO
+from brittle_star_project.environment import MorphMode
+from brittle_star_project.utils import logged_jit
+
+from brittle_star_project.environment.env_types import Backend
 
-# 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
+@logged_jit
 def _clip_action(action: jnp.ndarray, low: jnp.ndarray, high: jnp.ndarray) -> jnp.ndarray:
     return jnp.clip(action, low, high)
 
@@ -59,81 +53,113 @@ def _compute_explained_variance(values: jnp.ndarray, returns: jnp.ndarray) -> fl
     return float(explained_var)
 
 
-@jax.jit
+@logged_jit
 def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, learning_rate):
     frac = 1.0 - (count // (minibatch_count * update_epochs)) / num_iterations
     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:
-    """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(
-    sensor: GenericDenseLayersWithActivation,
-    feature_extractor: GenericDenseLayersWithActivation,
-    actor: Actor,
-    critic: OneDenseLayerMLP,
+    sensor: nn.Module,
+    feature_extractor: nn.Module,
+    actor: nn.Module,
+    critic: nn.Module,
+    message_passer: Optional[nn.Module],
     agent_state: TrainState,
     next_obs: jnp.ndarray,
-    key: jax.random.PRNGKey,
+    key,
     action_low,
     action_high,
 ):
-    hidden = sensor.apply(agent_state.params["sensor_params"], next_obs)
-    hidden_critic = feature_extractor.apply(
-        agent_state.params["feature_extractor_params"], next_obs
+    # (B, n_nodes, feat)
+    hidden = apply_per_node(sensor.apply, agent_state.params["sensor_params"], next_obs)
+
+    if message_passer is not None:
+        params = agent_state.params["message_passer_params"]
+        # (n_nodes, feat) --> let each node talk with its neighbours ==> vmap over B dimension
+        hidden = jax.vmap(lambda x: message_passer.apply(params, x))(hidden)
+
+    hidden_critic = apply_shared(
+        feature_extractor, agent_state.params["feature_extractor_params"], next_obs
     )
 
-    mean, log_std = actor.apply(agent_state.params["actor_params"], hidden)
+    mean, log_std = apply_per_node(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)
-    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 clipped_action, raw_action, logprob, value.squeeze(-1), mean, std, key
+    raw_action = mean + noise * std
+    flat_action = raw_action.reshape(
+        raw_action.shape[0], -1
+    )  # concat the per agent, keep the envs dim (batch, agent * action)
+    flat_clipped_action = _clip_action(flat_action, action_low, action_high)
+
+    logprob = -0.5 * (((raw_action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(
+        axis=(-2, -1)
+    )
+    value = apply_shared(critic, agent_state.params["critic_params"], hidden_critic)
+
+    return flat_clipped_action, raw_action, logprob, value.squeeze(-1), mean, std, key
 
 
 def _step_once(
     carry,
     _,
     env_step_fn,
-    sensor: GenericDenseLayersWithActivation,
-    feature_extractor: GenericDenseLayersWithActivation,
-    actor: Actor,
-    critic: OneDenseLayerMLP,
+    num_envs: int,
+    sensor: nn.Module,
+    feature_extractor: nn.Module,
+    actor: nn.Module,
+    critic: nn.Module,
+    message_passer: Optional[nn.Module],
     action_low,
     action_high,
 ):
-    agent_state, episode_stats, obs, done, key, env_state = carry
-    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
+    agent_state, episode_stats, obs, done, key, env_state, terminated_any, truncated_any = carry
+    flat_clipped_action, raw_action, logprob, value, mean, std, key = _get_action_and_value_noise(
+        sensor,
+        feature_extractor,
+        actor,
+        critic,
+        message_passer,
+        agent_state,
+        obs,
+        key,
+        action_low,
+        action_high,
+    )
+    logger = get_logger()
+
+    logger.debug(f"[_step_once] raw_action: {raw_action.shape}")
+    logger.debug(f"[_step_once] clipped_action: {flat_clipped_action.shape}")
+
+    # Supporting signals (often where mismatch originates)
+    logger.debug(f"[_step_once] logprob: {logprob.shape}")
+    logger.debug(f"[_step_once] value: {value.shape}")
+    logger.debug(f"[_step_once] mean: {mean.shape}")
+    logger.debug(f"[_step_once] std: {std.shape}")
+
+    key, reset_key = jax.random.split(key)
+    reset_rngs = jax.random.split(reset_key, num_envs)
+
+    # ---- ENV STEP ----
+    key, reset_key = jax.random.split(key)
+    reset_rngs = jax.random.split(reset_key, num_envs)
+
+    episode_stats, env_state, (next_obs, reward, next_done, terminated, truncated) = env_step_fn(
+        episode_stats,
+        env_state,
+        flat_clipped_action,
+        reset_rngs,
     )
 
-    episode_stats, env_state, (next_obs, reward, next_done) = env_step_fn(
-        episode_stats, env_state, clipped_action
-    )
+    terminated_any = terminated_any | terminated
+    truncated_any = truncated_any | truncated
+
+    logger.debug(f"[_step_once] next_obs: {next_obs.shape}")
+    logger.debug(f"[_step_once] reward: {reward.shape}")
+    logger.debug(f"[_step_once] next_done: {next_done.shape}")
 
     storage = Storage(
         obs=obs,
@@ -148,11 +174,25 @@ def _step_once(
         returns=jnp.zeros_like(reward),
         advantages=jnp.zeros_like(reward),
     )
-    return (agent_state, episode_stats, next_obs, next_done, key, env_state), storage
+    return (
+        agent_state,
+        episode_stats,
+        next_obs,
+        next_done,
+        key,
+        env_state,
+        terminated_any,
+        truncated_any,
+    ), storage
 
 
-def _reward_fn(env_state, next_env_state):
-    # if delta distance positive ==> brittle star walking away from target
+def reward_fn(env_state, next_env_state):
+    """Shaped reward used during training and checkpoint evaluation.
+
+    Public so that ``evaluation.evaluate_mjx`` can import it and produce
+    metrics that are directly comparable to training-time returns.
+    """
+    # Positive delta_distance means the brittle star is moving *away* from target.
     delta_distance = (
         next_env_state.observations["xy_distance_to_target"]
         - env_state.observations["xy_distance_to_target"]
@@ -168,12 +208,20 @@ def _reward_fn(env_state, next_env_state):
     return jnp.where(next_env_state.terminated, 50.0, clipped_env_reward - penalty)
 
 
-def _step_env_wrapped(episode_stats, env_state, action, env_step_fn):
-    next_env_state = env_step_fn(env_state, action)
+def _step_env_wrapped(
+    episode_stats,
+    env_state,
+    action,
+    reset_rngs,
+    env_step_fn,
+    reset_single_fn,
+    obs_processor,
+):
+    next_env_state_pre_reset = env_step_fn(env_state, action)
 
-    reward = _reward_fn(env_state, next_env_state)
-    terminated = next_env_state.terminated
-    truncated = next_env_state.truncated
+    reward = reward_fn(env_state, next_env_state_pre_reset)
+    terminated = next_env_state_pre_reset.terminated
+    truncated = next_env_state_pre_reset.truncated
     done = terminated | truncated
 
     new_episode_return = episode_stats.episode_returns + reward
@@ -189,13 +237,50 @@ def _step_env_wrapped(episode_stats, env_state, action, env_step_fn):
             done, new_episode_length, episode_stats.returned_episode_lengths
         ),
     )
+
+    def _maybe_reset(state_i, rng_i, do_reset_i):
+        def _do(_):
+            reset_state = reset_single_fn(rng=rng_i)
+
+            def _cast_leaf(new_leaf, like_leaf):
+                if like_leaf is None or new_leaf is None:
+                    return new_leaf
+
+                # Use jnp.asarray(...) to robustly get dtype for both JAX arrays and Python scalars.
+                like_dtype = jnp.asarray(like_leaf).dtype
+
+                # Avoid unnecessary work when already matching.
+                if hasattr(new_leaf, "dtype") and new_leaf.dtype == like_dtype:
+                    return new_leaf
+
+                return jnp.asarray(new_leaf, dtype=like_dtype)
+
+            # `lax.cond` requires both branches to return identical PyTree types/dtypes.
+            return jax.tree_util.tree_map(_cast_leaf, reset_state, state_i)
+
+        def _dont(_):
+            return state_i
+
+        return jax.lax.cond(do_reset_i, _do, _dont, operand=None)
+
+    # Auto-reset done envs so rollouts continue with fresh episode initial states.
+    next_env_state = jax.vmap(_maybe_reset)(next_env_state_pre_reset, reset_rngs, done)
+
     return (
         episode_stats,
         next_env_state,
-        (_convert_obs_dict_to_array(next_env_state.observations), reward, done),
+        (obs_processor(next_env_state.observations), reward, done, terminated, truncated),
     )
 
 
+def apply_shared(net, params, x):
+    # x: (batch, nodes, feat)
+    # If the critic expects a single vector per environment:
+    batch_size = x.shape[0]
+    x_flattened = x.reshape(batch_size, -1)
+    return jax.vmap(lambda xi: net.apply(params, xi))(x_flattened)
+
+
 def _rollout_jit(
     agent_state,
     episode_stats,
@@ -205,29 +290,67 @@ def _rollout_jit(
     key,
     max_steps,
     step_env_fn,
-    sensor: GenericDenseLayersWithActivation,
-    feature_extractor: GenericDenseLayersWithActivation,
-    actor: Actor,
-    critic: OneDenseLayerMLP,
+    num_envs: int,
+    sensor: nn.Module,
+    feature_extractor: nn.Module,
+    actor: nn.Module,
+    critic: nn.Module,
+    message_passer: Optional[nn.Module],
     action_low,
     action_high,
 ):
-    (agent_state, episode_stats, next_obs, next_done, key, env_state), storage = jax.lax.scan(
+    terminated_any0 = jnp.zeros((num_envs,), dtype=jnp.bool_)
+    truncated_any0 = jnp.zeros((num_envs,), dtype=jnp.bool_)
+
+    (
+        (
+            agent_state,
+            episode_stats,
+            next_obs,
+            next_done,
+            key,
+            env_state,
+            terminated_any,
+            truncated_any,
+        ),
+        storage,
+    ) = jax.lax.scan(
         partial(
             _step_once,
             sensor=sensor,
             feature_extractor=feature_extractor,
             actor=actor,
             critic=critic,
+            message_passer=message_passer,
             env_step_fn=step_env_fn,
+            num_envs=num_envs,
             action_low=action_low,
             action_high=action_high,
         ),
-        (agent_state, episode_stats, next_obs, next_done, key, env_state),
+        (
+            agent_state,
+            episode_stats,
+            next_obs,
+            next_done,
+            key,
+            env_state,
+            terminated_any0,
+            truncated_any0,
+        ),
         (),
         max_steps,
     )
-    return agent_state, episode_stats, next_obs, next_done, storage, key, env_state
+    return (
+        agent_state,
+        episode_stats,
+        next_obs,
+        next_done,
+        storage,
+        key,
+        env_state,
+        terminated_any,
+        truncated_any,
+    )
 
 
 def _compute_gae_once(carry, inp, gamma, gae_lambda):
@@ -250,9 +373,10 @@ def _compute_gae_jit(
     feature_extractor,
     critic,
 ):
-    next_value = critic.apply(
+    next_value = apply_shared(
+        critic,
         agent_state.params["critic_params"],
-        feature_extractor.apply(agent_state.params["feature_extractor_params"], next_obs),
+        apply_shared(feature_extractor, agent_state.params["feature_extractor_params"], next_obs),
     ).squeeze(-1)
 
     advantages = jnp.zeros((num_envs,))
@@ -286,12 +410,17 @@ class TrainingMeasurements:
 
 class PPOTrainer:
     def __init__(
-        self, cfg: BrittleStarConfig, env: BrittleStarJaxEnvWrapper, run_dir: str, run_name: str
+        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.evaluation_cfg = cfg.evaluation
         self.env = env
         self.run_dir = run_dir
         self.run_name = run_name
@@ -303,29 +432,69 @@ class PPOTrainer:
 
         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)
-        self.feature_extractor.apply = jax.jit(self.feature_extractor.apply)
-        self.actor.apply = jax.jit(self.actor.apply)
-        self.critic.apply = jax.jit(self.critic.apply)
+        self.morph_mode = self.cfg.morphology.morph_mode
+
+        self.segments_per_arm = jnp.asarray(self.cfg.morphology.segments_per_arm, dtype=jnp.int32)
+        self.num_segments = self.segments_per_arm.sum().item()
+        self.num_arms = jnp.where(self.segments_per_arm > 0, 1, 0).sum().item()
+
+        self.logger.info(f"[INIT]: Used morphology mode {self.morph_mode}")
+        self.adj = build_adjacency(cfg.morphology.segments_per_arm, self.morph_mode)
+
+        (
+            self.sensor,
+            self.message_passer,
+            self.actor,
+            self.feature_extractor,
+            self.critic,
+            self.needed_copies,
+            self.agent_indices,
+        ) = self._init_agent()
+
+        self.sensor.apply = logged_jit(self.sensor.apply)
+        self.feature_extractor.apply = logged_jit(self.feature_extractor.apply)
+        self.actor.apply = logged_jit(self.actor.apply)
+        self.critic.apply = logged_jit(self.critic.apply)
+
+        # Build the centralized observation processor: derive -> normalize -> pad -> flatten.
+        self.obs_processor = create_obs_processor(
+            bounds_dict=self.cfg.obs_bounds.to_bounds_dict(),
+            needed_copies=self.needed_copies,
+            num_arms=self.num_arms,
+            morph_mode=self.morph_mode,
+            padding_masks=self.env.padding_masks,
+            segments_per_arm=self.segments_per_arm,
+            agent_indices=self.agent_indices,
+        )
+
+        self.logger.debug(f"needed copies = {self.needed_copies}")
 
         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._action_low = action_low
+        self._action_high = action_high
 
-        self._rollout_jit = jax.jit(
+        self._rollout_jit = logged_jit(
             partial(
                 _rollout_jit,
                 max_steps=self.ppo.num_steps,
-                step_env_fn=partial(_step_env_wrapped, env_step_fn=self.env.step),
+                step_env_fn=partial(
+                    _step_env_wrapped,
+                    env_step_fn=self.env.step,
+                    reset_single_fn=self.env.raw.reset,
+                    obs_processor=self.obs_processor,
+                ),
+                num_envs=self.ppo.num_envs,
                 sensor=self.sensor,
                 feature_extractor=self.feature_extractor,
                 actor=self.actor,
                 critic=self.critic,
+                message_passer=self.message_passer,
                 action_low=action_low,
                 action_high=action_high,
             )
         )
-        self._compute_gae_jit = jax.jit(
+        self._compute_gae_jit = logged_jit(
             partial(
                 _compute_gae_jit,
                 num_envs=self.ppo.num_envs,
@@ -336,13 +505,38 @@ class PPOTrainer:
             )
         )
 
-        self._ppo = PPO(self.ppo, self.sensor, self.actor, self.critic, self.feature_extractor)
+        def apply_sensor(p, x):
+            return apply_per_node(self.sensor.apply, p, x)
+
+        def apply_actor(p, x):
+            return apply_per_node(self.actor.apply, p, x)
+
+        def apply_critic(p, x):
+            return apply_shared(self.critic, p, x)
+
+        def apply_feature(p, x):
+            return apply_shared(self.feature_extractor, p, x)
+
+        def apply_message_passer(p, x):
+            assert self.message_passer is not None
+            return jax.vmap(lambda x_in: self.message_passer.apply(p, x_in))(x)
+
+        self._ppo = PPO(
+            self.ppo,
+            apply_sensor,
+            apply_actor,
+            apply_critic,
+            apply_feature,
+            apply_message_passer if self.message_passer is not None else None,
+        )
 
         self.agent_state = self._init_agent_state()
 
         self.episode_stats = self._init_episode_stats()
 
         self._init_random()
+        # Lazily-built JIT-compiled MJX eval rollout, created on first evaluation.
+        self._eval_fn = None
 
     def _init_random(self):
         self.logger.info(f"[RANDOM]: Setting random seed to {self.experiment.seed}")
@@ -352,36 +546,136 @@ class PPOTrainer:
 
     def _init_agent(self):
         self.logger.info("[AGENT]: Initializing agent...")
+        agent_indices = [0, 1, 2, 3, 4]
+        match self.morph_mode:
+            case MorphMode.CENTRALIZED:
+                needed_copies = 1
+            case MorphMode.FULLY_CONNECTED | MorphMode.RING:
+                agent_mask = self.segments_per_arm > 0
+                agent_indices = jnp.where(agent_mask)[0]
+                needed_copies = jnp.where(self.segments_per_arm > 0, 1, 0).sum().item()
+            case MorphMode.SEGMENT:
+                agent_mask = self.segments_per_arm > 0
+                agent_indices = jnp.where(agent_mask)[0]
+                needed_copies = (
+                    self.segments_per_arm.sum() + jnp.where(self.segments_per_arm > 0, 1, 0).sum()
+                ).item()
 
+        # scale actor output with size of model --> more models ==> less actions needed per model
+        actor = Actor(action_dim=self.env.single_action_space.shape[0] // needed_copies)
         sensor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300])
+        message_passer: Optional[nn.Module] = (
+            MessagePasser(
+                hidden_dim=300,
+                num_propagation_steps=self.cfg.architecture.message_passing_steps or 4,
+                adj_matrix=self.adj,
+            )
+            if self.morph_mode != MorphMode.CENTRALIZED
+            else None
+        )
+
         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
+        return (
+            sensor,
+            message_passer,
+            actor,
+            feature_extractor,
+            critic,
+            needed_copies,
+            agent_indices,
+        )
 
     def _init_agent_state(self) -> TrainState:
         self.logger.info("[AGENT STATE]: Initializing agent state...")
 
-        self.key, sensor_key, actor_key, critic_key, feature_extractor_key = jax.random.split(
-            self.key, 5
+        self.key, sensor_key, actor_key, critic_key, feature_extractor_key, message_passer_key = (
+            jax.random.split(self.key, 6)
         )
 
         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),))
+
+        for k, v in dummy_reset.observations.items():
+            self.logger.debug(k, v.shape)
+
+        sample_obs = self.obs_processor(dummy_reset.observations)[0]  # take first env
+
+        self.logger.debug(f"[_init_agent_state] sample_obs: {sample_obs.shape}")
+        self.obs_mean = jnp.zeros((sample_obs.shape[-1],))
+        self.obs_var = jnp.ones((sample_obs.shape[-1],))
         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))
-        critic_params = self.critic.init(
-            critic_key, self.feature_extractor.apply(feature_extractor_params, sample_obs)
+        self.logger.debug(f"[_init_agent_state] obs_mean: {self.obs_mean.shape}")
+        self.logger.debug(f"[_init_agent_state] obs_var: {self.obs_var.shape}")
+
+        self.logger.debug(f"[_init_agent_state]: Needed copies: {self.needed_copies}")
+        sensor_keys = jax.random.split(sensor_key, self.needed_copies)
+        actor_keys = jax.random.split(actor_key, self.needed_copies)
+
+        # (needed_copies, X)
+        sensor_params = jax.vmap(lambda k: self.sensor.init(k, sample_obs))(sensor_keys)
+        self.logger.debug(
+            f"[_init_agent_state] sensor_params: {jax.tree.map(lambda x: x.shape, sensor_params)}"
+        )
+
+        single_sensor_param = jax.tree.map(lambda x: x[0], sensor_params)
+        self.logger.debug(
+            f"[_init_agent_state] single_sensor_param: {
+                jax.tree.map(lambda x: x.shape, single_sensor_param)
+            }"
+        )
+
+        sensor_params_sample = self.sensor.apply(single_sensor_param, sample_obs)
+        self.logger.debug(
+            f"[_init_agent_state] sensor_params_sample shape: {sensor_params_sample.shape}"
+        )
+
+        actor_params = jax.vmap(lambda k: self.actor.init(k, sensor_params_sample))(actor_keys)
+        self.logger.debug(
+            f"[_init_agent_state] actor_params: {jax.tree.map(lambda x: x.shape, actor_params)}"
+        )
+
+        message_passer_params = {}
+        if self.morph_mode != MorphMode.CENTRALIZED:
+            assert self.message_passer is not None, "decentralized modes require a message passer"
+
+            message_passer_params = self.message_passer.init(
+                message_passer_key,
+                self.sensor.apply(single_sensor_param, sample_obs),
+            )
+            self.logger.debug(
+                f"[_init_agent_state] message_passer_params: {
+                    jax.tree.map(lambda x: x.shape, message_passer_params)
+                }"
+            )
+
+        flat_obs = sample_obs.reshape(-1)  # BECAUSE 1 centralized critic
+        self.logger.debug(f"[_init_agent_state] flat_obs: {flat_obs.shape}")
+
+        feature_extractor_params = self.feature_extractor.init(feature_extractor_key, flat_obs)
+        self.logger.debug(
+            f"[_init_agent_state] feature_extractor_params: {
+                jax.tree.map(lambda x: x.shape, feature_extractor_params)
+            }"
+        )
+
+        critic_input = self.feature_extractor.apply(feature_extractor_params, flat_obs)
+        self.logger.debug(f"[_init_agent_state] critic_input: {critic_input.shape}")
+
+        critic_params = self.critic.init(critic_key, critic_input)
+        self.logger.debug(
+            f"[_init_agent_state] critic_params: {jax.tree.map(lambda x: x.shape, critic_params)}"
         )
 
         return TrainState.create(
             apply_fn=None,
             params=asdict(
-                AgentParams(sensor_params, actor_params, critic_params, feature_extractor_params)
+                AgentParams(
+                    sensor_params,
+                    actor_params,
+                    critic_params,
+                    feature_extractor_params,
+                    message_passer_params,
+                )
             ),
             tx=optax.chain(
                 optax.clip_by_global_norm(self.ppo.max_grad_norm),
@@ -410,25 +704,6 @@ class PPOTrainer:
             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,
@@ -503,7 +778,7 @@ class PPOTrainer:
     def _step(self, env_state, next_obs, next_done, iteration: int) -> tuple:
         if iteration == 1:
             self.logger.log_non_interactive(f"Starting first rollout (JIT): {time.ctime()}")
-
+        self.logger.debug(f"[_step] next_obs (in): {next_obs.shape}")
         (
             self.agent_state,
             self.episode_stats,
@@ -512,13 +787,15 @@ class PPOTrainer:
             storage,
             self.key,
             next_env_state,
+            terminated_any,
+            truncated_any,
         ) = self._rollout(env_state, next_obs, next_done)
-
+        self.logger.debug(f"[_step] next_obs (post-rollout): {next_obs.shape}")
         if iteration == 1:
             self.logger.log_non_interactive(f"First rollout completed: {time.ctime()}")
 
         storage = self._compute_gae(storage, next_obs, next_done)
-
+        self.logger.debug(f"[_step] storage.obs (post-gae): {storage.obs.shape}")
         if iteration == 1:
             self.logger.log_non_interactive(f"Starting first PPO update (JIT): {time.ctime()}")
 
@@ -535,8 +812,8 @@ class PPOTrainer:
 
         explained_var = _compute_explained_variance(storage.values, storage.returns)
 
-        terminated = next_env_state.terminated
-        truncated = next_env_state.truncated
+        terminated = terminated_any
+        truncated = truncated_any
         episode_lengths = self.episode_stats.returned_episode_lengths
 
         num_terminated = int(jnp.sum(terminated).item())
@@ -583,6 +860,65 @@ class PPOTrainer:
             params=self.agent_state.params, step=iteration, metadata=asdict(self.cfg)
         )
 
+    def _evaluate_checkpoint(self, iteration: int, *, trained_timesteps: int) -> None:
+        """Evaluate the current checkpoint and persist metrics to CSV.
+
+        Delegates all evaluation logic to `evaluation.evaluate_mjx`.
+        Best-effort: a failure here must never abort training.
+        """
+        if not self.evaluation_cfg.evaluate_checkpoints:
+            return
+
+        max_steps = int(self.evaluation_cfg.eval_max_steps)
+        seed = int(self.evaluation_cfg.eval_seed)
+
+        if max_steps <= 0:
+            self.logger.warning("[EVAL]: eval_max_steps must be > 0; skipping evaluation")
+            return
+
+        if not self.logging_cfg.save_checkpoints or self.logging_cfg.checkpoint_frequency <= 0:
+            self.logger.warning(
+                "[EVAL]: evaluate_checkpoints is enabled but checkpoint saving is disabled; "
+                "skipping evaluation"
+            )
+            return
+
+        try:
+            if self._eval_fn is None:
+                if getattr(self.env, "backend", None) != Backend.MJX:
+                    self.logger.warning(
+                        f"[EVAL]: Training env backend is {self.env.backend}; "
+                        "MJX evaluation may be unavailable/slow."
+                    )
+                self._eval_fn = build_eval_rollout_fn(
+                    env=self.env,
+                    obs_processor=self.obs_processor,
+                    sensor_apply=lambda p, x: apply_per_node(self.sensor.apply, p, x),
+                    actor_apply=lambda p, x: apply_per_node(self.actor.apply, p, x),
+                    message_passer_apply=(
+                        None if self.message_passer is None else self.message_passer.apply
+                    ),
+                    action_low=self._action_low,
+                    action_high=self._action_high,
+                    reward_fn=reward_fn,
+                )
+
+            result = evaluate_checkpoint_mjx(
+                self._eval_fn,
+                self.agent_state.params,
+                seed=seed,
+                max_steps=max_steps,
+            )
+            csv_path = append_checkpoint_eval_row(
+                self.run_dir,
+                iteration=iteration,
+                trained_timesteps=int(trained_timesteps),
+                result=result,
+            )
+            self.logger.sync_file(csv_path)
+        except Exception as e:
+            self.logger.warning(f"[EVAL]: Checkpoint evaluation failed: {e}")
+
     def train(self):
         """
         Train the PPO agent for a specified number of iterations.
@@ -594,7 +930,10 @@ class PPOTrainer:
         self.logger.log_non_interactive(f"Initial reset started: {time.ctime()}")
 
         env_state = self.env.reset(seed=self.experiment.seed)
-        next_obs = _convert_obs_dict_to_array(env_state.observations)
+
+        next_obs = self.obs_processor(env_state.observations)
+        self.logger.debug(f"[train] next_obs: {next_obs.shape}")
+
         next_done = jnp.zeros(self.ppo.num_envs, dtype=jnp.bool_)
 
         self.logger.log_non_interactive(f"Initial reset completed: {time.ctime()}")
@@ -609,8 +948,6 @@ class PPOTrainer:
             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.ppo.num_steps * self.ppo.num_envs
             self._log(
@@ -638,6 +975,7 @@ class PPOTrainer:
             if self.logging_cfg.save_checkpoints and self.logging_cfg.checkpoint_frequency > 0:
                 if iteration % self.logging_cfg.checkpoint_frequency == 0:
                     self._save_checkpoint(iteration)
+                    self._evaluate_checkpoint(iteration, trained_timesteps=global_step)
 
             if getattr(self.cfg.experiment, "debug_sanity", False):
                 self.logger.info("\n[SANITY CHECK] Successfully completed 1 epoch")

src/brittle_star_project/utils/__init__.py

@@ -0,0 +1,3 @@
+from .logged_jit import logged_jit
+
+__all__ = ["logged_jit"]

src/brittle_star_project/utils/logged_jit.py

@@ -0,0 +1,17 @@
+import jax
+from experiment_logger import get_logger
+
+
+def logged_jit(fn, **jit_kwargs):
+    logger = get_logger()
+    name = getattr(fn, "__name__", getattr(fn, "__qualname__", repr(fn)))
+
+    def decorator(func):
+        def traced_func(*args, **kwargs):
+            logger.debug(f"[JIT] Compiling {name}...")
+            return func(*args, **kwargs)
+
+        jitted = jax.jit(traced_func, **jit_kwargs)
+        return jitted
+
+    return decorator(fn)

src/experiment_logger/config_logger.py

@@ -31,3 +31,6 @@ class LoggingConfig:
                 "Configuration Error: 'upload_checkpoints' is True, but it requires "
                 "both 'track' and 'save_checkpoints' to also be True."
             )
+
+        # NOTE: Checkpoint evaluation settings live under the project's
+        # `evaluation` config group (see brittle_star_project.configs).

src/experiment_logger/simple_logger.py

@@ -71,6 +71,10 @@ class SimpleLogger:
     def save_final_model(self, params: Any, metadata: Optional[Dict[str, Any]] = None):
         print("[SAVE] Final model would be saved (SimpleLogger: No-Op)")
 
+    def sync_file(self, path: Any):
+        """No-op for SimpleLogger."""
+        pass
+
     def finish(self):
         print(f"[FINISH] SimpleLogger finished for run: {self.run_name}")
 

src/experiment_logger/unified_logger.py

@@ -432,6 +432,21 @@ class UnifiedLogger:
         except Exception as e:
             self.error(f"Error saving final model: {e}")
 
+    def sync_file(self, path: Path) -> None:
+        """Upload a file to W&B if tracking is enabled.
+
+        Best-effort: logs a warning on failure, never raises.
+        """
+        if self.wandb_run is None:
+            return
+        try:
+            import wandb
+
+            # "Simple sync" behavior: wandb will copy this file into the run.
+            wandb.save(str(path), base_path=str(path.parent))
+        except Exception as e:
+            self.warning(f"Failed to sync file to W&B: {e}")
+
     def finish(self):
         """Finalize logging and cleanup."""
         # Flush remaining metrics

tests/test_adjacency.py

@@ -0,0 +1,98 @@
+import jax.numpy as jnp
+import numpy as np
+
+from brittle_star_project.MLPs import build_adjacency
+from brittle_star_project.environment.env_config import MorphMode
+
+
+def assert_symmetric(adj):
+    assert jnp.all(adj == adj.T)
+
+
+def test_centralized():
+    adj = build_adjacency([4, 4, 4, 4, 4], MorphMode.CENTRALIZED)
+
+    assert adj.shape == (1, 1)
+    assert adj[0, 0] == 1
+
+
+def test_fully_connected():
+    adj = build_adjacency([4, 4, 4, 4, 4], MorphMode.FULLY_CONNECTED)
+
+    assert adj.shape == (5, 5)
+    assert jnp.all(adj == 1)
+
+
+def test_ring():
+    adj = build_adjacency([4, 4, 4, 4, 4], MorphMode.RING)
+
+    assert adj.shape == (5, 5)
+    assert_symmetric(adj)
+
+    # each node should connect to itself + 2 neighbors
+    for node in range(5):
+        assert adj[node, node] == 1
+        assert jnp.sum(adj[node]) == 3
+        neighbor1 = (node - 1) % 5
+        neighbor2 = (node + 1) % 5
+        assert adj[neighbor1, node] == 1
+        assert adj[node, neighbor2] == 1  # Symmetrical
+
+
+def test_segment_structure():
+    segments = [4, 4, 4, 4, 4]
+    adj = build_adjacency(segments, MorphMode.SEGMENT)
+
+    num_arms = 5
+    num_segments = sum(segments)
+    num_nodes = num_arms + num_segments
+
+    assert adj.shape == (num_nodes, num_nodes)
+
+    # --- ring connectivity ---
+    for i in range(num_arms):
+        assert adj[i, i] == 1
+        assert adj[i, (i - 1) % num_arms] == 1
+        assert adj[i, (i + 1) % num_arms] == 1
+
+    # --- segment chain checks ---
+    offset = num_arms
+    for arm in range(5):
+        for i in range(4):
+            node = offset + arm * 4 + i
+
+            # self
+            assert adj[node, node] == 1
+
+            # chain neighbors
+            if i > 0:
+                assert adj[node, node - 1] == 1
+            if i < 3:
+                assert adj[node, node + 1] == 1
+
+    # --- ring ↔ segment connections ---
+    for arm in range(5):
+        first_seg = num_arms + arm * 4
+        assert adj[arm, first_seg] == 1
+        assert adj[first_seg, arm] == 1
+
+    save_adj(adj)
+
+
+def save_adj(adj, name="adjacency_debug.txt"):
+    a = np.array(adj)
+
+    with open(name, "w") as f:
+        f.write("\nAdjacency matrix:\n")
+        f.write("   " + " ".join([f"{i:2d}" for i in range(a.shape[0])]) + "\n")
+
+        for i, row in enumerate(a):
+            line = f"{i:2d}  " + "  ".join(["█" if x > 0 else "." for x in row])
+            f.write(line + "\n")
+
+
+def test_no_isolated_nodes():
+    adj = build_adjacency([4, 4, 4, 4, 4], MorphMode.SEGMENT)
+
+    # no node should be completely isolated
+    assert jnp.all(jnp.sum(adj, axis=0) > 0)

tests/test_evaluation.py

@@ -0,0 +1,218 @@
+import numpy as np
+import pytest
+import yaml
+from pathlib import Path
+
+from brittle_star_project.evaluation.checkpoint import (
+    metadata_to_configs,
+    TrainingConfig,
+    load_metadata,
+)
+from brittle_star_project.evaluation.rollout import _maybe_clip_action
+from brittle_star_project.environment.env_config import (
+    MorphologyConfig,
+    ArenaConfig,
+    EnvConfig,
+    ObservationBoundsConfig,
+    MorphMode,
+)
+from brittle_star_project.environment.env_types import Task
+
+
+def test_metadata_to_configs():
+    """Test that a raw metadata dictionary correctly instantiates the typed configs."""
+    mock_metadata = {
+        "morphology": {
+            "segments_per_arm": [4, 0, 4, 0, 0],
+            "use_p_control": False,
+        },
+        "arena": {"sand_ground_color": False, "size": [15.0, 10.0]},
+        "environment": {
+            "task": "LIGHT_ESCAPE",
+            "simulation_time": 5000.0,
+        },
+        "obs_bounds": {"joint_velocity": [-10.0, 10.0]},
+    }
+
+    config = metadata_to_configs(mock_metadata)
+
+    assert isinstance(config, TrainingConfig)
+
+    # Check MorphologyConfig
+    assert isinstance(config.morphology, MorphologyConfig)
+    assert config.morphology.segments_per_arm == [4, 0, 4, 0, 0]
+    assert config.morphology.use_p_control is False
+    assert config.morphology.use_torque_control is False  # default
+
+    # Check ArenaConfig
+    assert isinstance(config.arena, ArenaConfig)
+    assert config.arena.sand_ground_color is False
+    assert config.arena.size == [15.0, 10.0]
+    assert config.arena.wall_height == 1.5  # default
+
+    # Check EnvConfig
+    assert isinstance(config.environment, EnvConfig)
+    assert config.environment.task == Task.LIGHT_ESCAPE
+    assert config.environment.simulation_time == 5000.0
+    assert config.environment.time_scale == 2  # default
+
+    # Check ObservationBoundsConfig
+    assert isinstance(config.obs_bounds, ObservationBoundsConfig)
+    assert config.obs_bounds.joint_velocity == [-10.0, 10.0]
+    assert config.obs_bounds.segment_contact == [0.0, 1.0]  # default
+
+
+def test_maybe_clip_action():
+    """Test action clipping against boundaries."""
+    # Test valid clipping
+    action = np.array([1.5, -2.5, 0.0])
+    low = np.array([-1.0, -1.0, -1.0])
+    high = np.array([1.0, 1.0, 1.0])
+
+    clipped = _maybe_clip_action(action, low, high)
+    np.testing.assert_array_equal(clipped, np.array([1.0, -1.0, 0.0]))
+
+    # Test skipping when bounds are None
+    unclipped_1 = _maybe_clip_action(action, None, high)
+    np.testing.assert_array_equal(unclipped_1, action)
+
+    unclipped_2 = _maybe_clip_action(action, low, None)
+    np.testing.assert_array_equal(unclipped_2, action)
+
+    # Test skipping on shape mismatch
+    wrong_low = np.array([-1.0, -1.0])  # Shape mismatch
+    unclipped_3 = _maybe_clip_action(action, wrong_low, high)
+    np.testing.assert_array_equal(unclipped_3, action)
+
+
+def test_load_metadata_with_override(tmp_path: Path):
+    """Test that metadata can be loaded from both default and override paths."""
+    # 1. Setup
+    model_path = tmp_path / "model.flax"
+    model_path.write_bytes(b"dummy")
+
+    default_metadata_path = tmp_path / "model_metadata.yaml"
+    default_content = {"version": "default", "seed": 42}
+    with open(default_metadata_path, "w") as f:
+        yaml.dump(default_content, f)
+
+    override_path = tmp_path / "custom_metadata.yaml"
+    override_content = {"version": "override", "seed": 1337}
+    with open(override_path, "w") as f:
+        yaml.dump(override_content, f)
+
+    # 2. Test default behavior
+    loaded_default = load_metadata(model_path)
+    assert loaded_default == default_content
+
+    # 3. Test override behavior
+    loaded_override = load_metadata(model_path, metadata_override_path=override_path)
+    assert loaded_override == override_content
+
+    # 4. Test Error Case
+    non_existent = tmp_path / "missing.yaml"
+    with pytest.raises(FileNotFoundError, match="Could not find metadata YAML at"):
+        load_metadata(model_path, metadata_override_path=non_existent)
+
+
+@pytest.fixture
+def mock_training_config():
+    return TrainingConfig(
+        morphology=MorphologyConfig(
+            segments_per_arm=[1, 1, 1, 1, 1], morph_mode=MorphMode.CENTRALIZED
+        ),
+        arena=ArenaConfig(),
+        environment=EnvConfig(),
+        obs_bounds=ObservationBoundsConfig(),
+    )
+
+
+@pytest.fixture
+def mock_metadata():
+    return {"architecture": {"message_passing_steps": 2}}
+
+
+def test_build_eval_env_training_morphology(tmp_path, mock_training_config, mock_metadata):
+    from brittle_star_project.evaluation.eval_env_builder import build_eval_env
+    from unittest.mock import patch
+
+    model_path = tmp_path / "model.flax"
+    patch_target = "brittle_star_project.evaluation.eval_env_builder.PolicyAgent.from_checkpoint"
+    with patch(patch_target) as mock_agent:
+        mock_agent.return_value = "mock_policy"
+        bundle = build_eval_env(
+            model_path=model_path,
+            training=mock_training_config,
+            metadata=mock_metadata,
+            morphology_override_path=None,
+        )
+        assert bundle.segments_per_arm == [1, 1, 1, 1, 1]
+        assert bundle.num_active_arms == 5
+        assert bundle.architecture == "CENTRALIZED"
+        assert bundle.policy == "mock_policy"
+
+
+def test_build_eval_env_override_morphology(tmp_path, mock_training_config, mock_metadata):
+    from brittle_star_project.evaluation.eval_env_builder import build_eval_env
+    from unittest.mock import patch
+
+    model_path = tmp_path / "model.flax"
+    override_path = tmp_path / "override.yaml"
+    override_path.write_text(yaml.dump({"segments_per_arm": [1, 0, 1, 0, 1]}))
+
+    with patch("brittle_star_project.evaluation.eval_env_builder.PolicyAgent.from_checkpoint"):
+        bundle = build_eval_env(
+            model_path=model_path,
+            training=mock_training_config,
+            metadata=mock_metadata,
+            morphology_override_path=override_path,
+        )
+        assert bundle.segments_per_arm == [1, 0, 1, 0, 1]
+        assert bundle.num_active_arms == 3
+        # Should be smaller than 5*N
+        assert sum(bundle.action_mask) < len(bundle.action_mask)
+
+
+def test_build_eval_env_action_mask_shape(tmp_path, mock_training_config, mock_metadata):
+    from brittle_star_project.evaluation.eval_env_builder import build_eval_env
+    from unittest.mock import patch
+
+    model_path = tmp_path / "model.flax"
+    override_path = tmp_path / "override.yaml"
+    override_path.write_text(yaml.dump({"segments_per_arm": [1, 0, 1, 0, 0]}))
+
+    with patch("brittle_star_project.evaluation.eval_env_builder.PolicyAgent.from_checkpoint"):
+        bundle = build_eval_env(
+            model_path=model_path,
+            training=mock_training_config,
+            metadata=mock_metadata,
+            morphology_override_path=override_path,
+        )
+        # For each segment with P-control, there's 2 actions (pitch and yaw).
+        # Total segments = 5 -> 10 actions for training.
+        assert len(bundle.action_mask) == 10
+        # Active segments = 2 -> 4 actions active.
+        assert sum(bundle.action_mask) == 4
+
+
+def test_build_eval_env_morph_mode_inherited(tmp_path, mock_training_config, mock_metadata):
+    from brittle_star_project.evaluation.eval_env_builder import build_eval_env
+    from brittle_star_project.environment.env_config import MorphMode
+    from unittest.mock import patch
+
+    model_path = tmp_path / "model.flax"
+    override_path = tmp_path / "override.yaml"
+    # No morph_mode in the override YAML
+    override_path.write_text(yaml.dump({"segments_per_arm": [1, 0, 1, 0, 1]}))
+
+    # Change training config to be RING
+    mock_training_config.morphology.morph_mode = MorphMode.RING
+
+    with patch("brittle_star_project.evaluation.eval_env_builder.PolicyAgent.from_checkpoint"):
+        bundle = build_eval_env(
+            model_path=model_path,
+            training=mock_training_config,
+            metadata=mock_metadata,
+            morphology_override_path=override_path,
+        )
+        assert bundle.architecture == "RING"

tests/test_morphology_render.py

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

tests/test_network_shapes.py

@@ -1,9 +1,10 @@
 import jax
 import jax.numpy as jnp
+from brittle_star_project.environment.env_config import MorphMode
 from brittle_star_project.environment.padded_obs_wrapper import (
     compute_padding_masks,
-    pad_observations_batched,
 )
+from brittle_star_project.environment.obs_processing import create_obs_processor
 
 # We use Actor and OneDenseLayerMLP (as the critic) based on your mlps.py
 from brittle_star_project.MLPs.mlps import Actor, OneDenseLayerMLP
@@ -20,28 +21,29 @@ def test_centralized_forward_pass_with_padding():
         "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)
+    segments_per_arm = jnp.array((4, 0, 4, 2, 4))
+    num_arms = jnp.where(segments_per_arm > 0, 1, 0).sum().item()
 
-    # 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"
+    # 2. Process and Pad Observation
+    masks = compute_padding_masks(segments_per_arm=list(segments_per_arm))
+    obs_processor = create_obs_processor(
+        bounds_dict={},
+        needed_copies=1,
+        num_arms=num_arms,
+        padding_masks=masks,
+        morph_mode=MorphMode.CENTRALIZED,
+        segments_per_arm=segments_per_arm,
+    )
+    global_state = obs_processor(amputated_obs)
+
+    # joint_position: 5 arms × 8 joints (padded) = 40
+    # joint_velocity:  5 arms × 8 joints (padded) = 40
+    # segment_contact: 5 arms × 4 segs  (padded) = 20
+    # Total = 100 (no disk or direction keys supplied)
+    assert global_state.shape == (batch_size, 1, 100), (
+        f"Expected global state shape (2, 1, 100), got {global_state.shape}"
     )
 
-    # 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
 
@@ -56,9 +58,11 @@ def test_centralized_forward_pass_with_padding():
     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_mean.shape == (batch_size, 1, 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,), (
+    assert value.shape == (batch_size, 1, 1) or value.shape == (batch_size,), (
         f"Critic value shape mismatch: {value.shape}"
     )
 

tests/test_obs_processor.py

@@ -0,0 +1,124 @@
+import jax
+import jax.numpy as jnp
+
+from brittle_star_project.environment.obs_processing import create_obs_processor
+from brittle_star_project.environment.env_config import MorphMode, ObservationBoundsConfig
+
+
+obs_bounds = ObservationBoundsConfig().to_bounds_dict()
+
+# Features per decentralized agent (one arm's data):
+#   disk_z_tilt              → scalar                    → 1 feat
+#   joint_actuator_force     → 4 segs × 2 joints         → 8 feat
+#   joint_position           → 4 segs × 2 joints         → 8 feat
+#   joint_velocity           → 4 segs × 2 joints         → 8 feat
+#   robot_direction_to_target→ (x, y)                    → 2 feat
+#   segment_contact          → 4 segs                    → 4 feat
+# Total per agent: 1+8+8+8+2+4 = 31
+
+NUM_ARMS = 5
+SEGS_PER_ARM = 4  # healthy segments per arm
+JOINTS_PER_SEG = 2  # from _build_joint_indices: segs * 2
+
+SEGS_HEALTHY = [4, 4, 4, 4, 4]
+SEGS_DAMAGED = [4, 4, 4, 4, 0]  # arm 4 fully disabled
+SEGS_DAMAGED_2 = [4, 0, 4, 2, 4]  # arm 3 fully disabled
+AGENT_INDICES = [0, 1, 2, 3, 4]
+
+FEAT_PER_AGENT = 1 + 8 + 8 + 8 + 2 + 4  # = 31
+
+# Centralized flattening (needed_copies=1, one copy of global features):
+#   disk_z_tilt              → repeated once              →  1 feat
+#   joint_actuator_force     → 5 arms × 8 joints          → 40 feat
+#   joint_position           → 5 arms × 8 joints          → 40 feat
+#   joint_velocity           → 5 arms × 8 joints          → 40 feat
+#   robot_direction_to_target→ repeated once              →  2 feat
+#   segment_contact          → 5 arms × 4 segs            → 20 feat
+# Total: 1+40+40+40+2+20 = 143
+FEAT_CENTRALIZED = 1 + 40 + 40 + 40 + 2 + 20  # = 143
+
+
+def make_obs(segs_per_arm: list[int]) -> dict:
+    total_segs = sum(segs_per_arm)
+    total_joints = JOINTS_PER_SEG * total_segs
+
+    return {
+        "actuator_force": jnp.ones(total_joints),
+        "disk_angular_velocity": jnp.zeros(3),
+        "disk_linear_velocity": jnp.zeros(3),
+        "disk_position": jnp.zeros(3),
+        "disk_rotation": jnp.array([0.1, 0.1, 0.5]),  # (roll, pitch, yaw)
+        "joint_actuator_force": jnp.full(total_joints, 1.0),
+        "joint_position": jnp.full(total_joints, 0.5),
+        "joint_velocity": jnp.full(total_joints, 2.0),
+        "segment_contact": jnp.ones(total_segs),
+        "tendon_position": jnp.zeros(0),
+        "tendon_velocity": jnp.zeros(0),
+        "unit_xy_direction_to_target": jnp.array([1.0, 0.0]),
+        "xy_distance_to_target": jnp.array([3.5]),
+    }
+
+
+def batch_obs(obs: dict):
+    return jax.tree_util.tree_map(lambda x: x[None, :], obs)
+
+
+def make_processor(morph_mode: MorphMode, needed_copies: int, segments_per_arm: list[int]):
+    return create_obs_processor(
+        bounds_dict=obs_bounds,
+        num_arms=NUM_ARMS,
+        needed_copies=needed_copies,
+        morph_mode=morph_mode,
+        segments_per_arm=segments_per_arm,
+        agent_indices=AGENT_INDICES,
+    )
+
+
+def test_centralized_no_damage():
+    proc = make_processor(MorphMode.CENTRALIZED, 1, SEGS_HEALTHY)
+    obs = make_obs(SEGS_HEALTHY)
+    obs = batch_obs(obs)
+    global_state = proc(obs)
+
+    # Centralized: 5 agents flattened into 1 → shape (1, 1, 155)
+    assert global_state.shape == (1, 1, FEAT_CENTRALIZED)
+
+
+def test_centralized_damaged_1_arm():
+    proc = make_processor(MorphMode.CENTRALIZED, 1, SEGS_DAMAGED)
+    obs = make_obs(SEGS_DAMAGED)
+    obs = batch_obs(obs)
+    global_state = proc(obs)
+
+    # shape test
+    assert global_state.shape == (1, 1, FEAT_CENTRALIZED)
+
+
+def test_centralized_damaged_2_arms():
+    proc = make_processor(MorphMode.CENTRALIZED, 1, SEGS_DAMAGED_2)
+    obs = make_obs(SEGS_DAMAGED_2)
+    obs = batch_obs(obs)
+    global_state = proc(obs)
+
+    # shape test
+    assert global_state.shape == (1, 1, FEAT_CENTRALIZED)
+
+
+def test_decentralized_fully_connected_no_damage():
+    proc = make_processor(MorphMode.FULLY_CONNECTED, NUM_ARMS, SEGS_HEALTHY)
+    obs = make_obs(SEGS_HEALTHY)
+    obs = batch_obs(obs)
+    global_state = proc(obs)
+
+    # shape test
+    assert global_state.shape == (1, NUM_ARMS, FEAT_PER_AGENT)
+
+
+def test_decentralized_fully_connected_damaged_1_arm():
+    proc = make_processor(MorphMode.FULLY_CONNECTED, NUM_ARMS, SEGS_DAMAGED)
+    obs = make_obs(SEGS_DAMAGED)
+    obs = batch_obs(obs)
+    global_state = proc(obs)
+
+    # shape test
+    assert global_state.shape == (1, NUM_ARMS, FEAT_PER_AGENT)

tests/test_target_direction.py

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

uv.lock

@@ -42,6 +42,10 @@ analysis = [
 cuda = [
     { name = "jax", extra = ["cuda13"] },
 ]
+evaluation = [
+    { name = "imageio" },
+    { name = "imageio-ffmpeg" },
+]
 
 [package.dev-dependencies]
 dev = [
@@ -58,6 +62,8 @@ requires-dist = [
     { name = "flax", specifier = ">=0.12.2" },
     { name = "gymnasium", specifier = ">=1.2.3" },
     { name = "hydra-core", specifier = ">=1.3.2" },
+    { name = "imageio", marker = "extra == 'evaluation'", specifier = ">=2.35.0" },
+    { name = "imageio-ffmpeg", marker = "extra == 'evaluation'", specifier = ">=0.5.1" },
     { 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" },
@@ -75,7 +81,7 @@ requires-dist = [
     { name = "wandb", specifier = "==0.24.2" },
     { name = "warp-lang" },
 ]
-provides-extras = ["cuda", "analysis"]
+provides-extras = ["cuda", "analysis", "evaluation"]
 
 [package.metadata.requires-dev]
 dev = [
@@ -795,6 +801,20 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/49/fa/391e437a34e55095173dca5f24070d89cbc233ff85bf1c29c93248c6588d/imageio-2.37.3-py3-none-any.whl", hash = "sha256:46f5bb8522cd421c0f5ae104d8268f569d856b29eb1a13b92829d1970f32c9f0", size = 317646, upload-time = "2026-03-09T11:31:10.771Z" },
 ]
 
+[[package]]
+name = "imageio-ffmpeg"
+version = "0.6.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/44/bd/c3343c721f2a1b0c9fc71c1aebf1966a3b7f08c2eea8ed5437a2865611d6/imageio_ffmpeg-0.6.0.tar.gz", hash = "sha256:e2556bed8e005564a9f925bb7afa4002d82770d6b08825078b7697ab88ba1755", size = 25210, upload-time = "2025-01-16T21:34:32.747Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/da/58/87ef68ac83f4c7690961bce288fd8e382bc5f1513860fc7f90a9c1c1c6bf/imageio_ffmpeg-0.6.0-py3-none-macosx_10_9_intel.macosx_10_9_x86_64.whl", hash = "sha256:9d2baaf867088508d4a3458e61eeb30e945c4ad8016025545f66c4b5aaef0a61", size = 24932969, upload-time = "2025-01-16T21:34:20.464Z" },
+    { url = "https://files.pythonhosted.org/packages/40/5c/f3d8a657d362cc93b81aab8feda487317da5b5d31c0e1fdfd5e986e55d17/imageio_ffmpeg-0.6.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:b1ae3173414b5fc5f538a726c4e48ea97edc0d2cdc11f103afee655c463fa742", size = 21113891, upload-time = "2025-01-16T21:34:00.277Z" },
+    { url = "https://files.pythonhosted.org/packages/33/e7/1925bfbc563c39c1d2e82501d8372734a5c725e53ac3b31b4c2d081e895b/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:1d47bebd83d2c5fc770720d211855f208af8a596c82d17730aa51e815cdee6dc", size = 25632706, upload-time = "2025-01-16T21:33:53.475Z" },
+    { url = "https://files.pythonhosted.org/packages/a0/2d/43c8522a2038e9d0e7dbdf3a61195ecc31ca576fb1527a528c877e87d973/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:c7e46fcec401dd990405049d2e2f475e2b397779df2519b544b8aab515195282", size = 29498237, upload-time = "2025-01-16T21:34:13.726Z" },
+    { url = "https://files.pythonhosted.org/packages/a0/13/59da54728351883c3c1d9fca1710ab8eee82c7beba585df8f25ca925f08f/imageio_ffmpeg-0.6.0-py3-none-win32.whl", hash = "sha256:196faa79366b4a82f95c0f4053191d2013f4714a715780f0ad2a68ff37483cc2", size = 19652251, upload-time = "2025-01-16T21:34:06.812Z" },
+    { url = "https://files.pythonhosted.org/packages/2c/c6/fa760e12a2483469e2bf5058c5faff664acf66cadb4df2ad6205b016a73d/imageio_ffmpeg-0.6.0-py3-none-win_amd64.whl", hash = "sha256:02fa47c83703c37df6bfe4896aab339013f62bf02c5ebf2dce6da56af04ffc0a", size = 31246824, upload-time = "2025-01-16T21:34:28.6Z" },
+]
+
 [[package]]
 name = "iniconfig"
 version = "2.3.0"