Merge pull request #62 from SELab-3-2026/dev

Release of project

.agents/rules/architecture.md

@@ -0,0 +1,20 @@
+# Architecture and Code Formatting Rules
+
+When writing or modifying code in this project, adhere strictly to the following rules:
+
+## 1. Core Frameworks & Tooling
+- **PPO & CleanRL**: Proximal Policy Optimization (PPO) is the baseline algorithm. Use CleanRL as the starting framework for PPO, ensuring adaptation for continuous action spaces.
+- **JAX / Flax**: All Artificial Neural Network (ANN) controller architectures must be implemented using Flax (neural networks in JAX). Ensure full compatibility with the JAX/Flax ecosystem.
+- **MuJoCo**: The simulation environment uses a MuJoCo brittle star. Ensure that XML structures (sensors, actuators, joints, morphology) respect realistic constraints and adhere to the project requirements.
+
+## 2. Clean Code Principles
+- **Naming Conventions**: Variables and functions must have consistent, intention-revealing names. A descriptive name is universally preferred over a short name with an explanatory comment.
+- **Single Responsibility Function Design**: Functions must be modular. Minimize arguments and completely avoid boolean flag arguments that control execution behavior.
+- **Commenting**: Code explains the "how". Comments are strictly reserved for explaining the "why". 
+- **No Commented-out Code**: The AI must **never** generate commented-out or dead code. Delete it using version control instead.
+- **YAGNI & Complexity Management**: You Aren't Gonna Need It. Avoid premature optimization or unnecessary abstraction. Only introduce complexity with documented justification. Break large functions into small, testable blocks.
+- **No Notebooks for Core Logic**: Jupyter Notebooks are explicitly forbidden for general software development as they discourage modularity. They should only be used for prototyping, tutorials, or post-processing analysis.
+
+## 3. Formatting & Linting
+- **Ruff**: Output perfectly formatted code adhering to the Google style standard. Always format and lint the code using `ruff` (see `pyproject.toml` and `ruff.toml`).
+- **Separation of Concerns**: Configuration code must be completely separated from implementation logic. Core logic must never be mixed with scripts or notebooks.

.agents/rules/general.md

@@ -0,0 +1,15 @@
+# General AI Agent Rules
+
+When assisting with this project, the AI Agent must strictly abide by these overarching operational rules:
+
+## 1. Scientific Integrity
+- **No Hallucinations**: You must never fabricate results, hallucinate citations, or generate false empirical claims. Do not guess what happened if a process fails; rely strictly on outputs and logs.
+
+## 2. Agent Operational Constraints
+- **Absolute Paths**: Always use absolute paths when making tool calls or reading/writing files.
+- **Refactoring Guardrails**: Do not commence massive files/directory refactors or major system migrations without explicitly communicating the plan and asking for user clarification or approval first.
+- **No Boilerplate Feedback**: The AI must not produce generic boilerplate summaries or overly generic advice. Ensure all outputs are completely contextual, robust, and well-reasoned.
+
+## 3. Communication
+- **Artifacts and UI**: Use artifacts (like `implementation_plan.md` or `task.md`) appropriately for tracking design phases and updates.
+- **Explicit Documenting**: If performing design choices, list them explicitly. Keep the output focused on the exact task.

.agents/rules/git-workflow.md

@@ -0,0 +1,25 @@
+# Git Workflow & Repository Structure Rules
+
+When performing Git operations and managing the repository layout, follow these rules:
+
+## 1. Committing Practices
+- **Frequent & Small**: Produce small, logical commits instead of massive monolithic ones.
+- **Conventional Commits**: Commit messages must adhere to the Conventional Commits specification (e.g., `feat: ...`, `fix: ...`, `refactor: ...`). 
+- **Single Functionality**: Each commit should relate to exactly one piece of functionality or distinct structural change.
+
+## 2. Branching & Merging
+- **Branch `dev`**: The `dev` branch is the primary integration branch for pushing and merging code.
+- **Branch `main`**: Only stable, finalized releases may be pushed to `main`.
+- **Feature Branches**: Organize distinct work into logical feature branches when pushing to the remote server, maintaining an organized Git history.
+
+## 3. Artifact Management & Exclusions
+- **LFS Only**: Data files, trained models, and large datasets must **never** be committed directly to Git. Ensure they are tracked with Git Large File Storage (LFS). **CRITICAL**: Every developer and AI agent must have `git-lfs` installed locally for the repository hooks to successfully pull these large files. Run `git lfs install` after cloning or setting up your environment.
+
+## 4. Repository Layout Strictness
+Ensure generated code is meticulously placed in the correct directories:
+- `src/` for algorithms, network designs, and core agent modules.
+- `env/` for MuJoCo wrappers and environment definitions.
+- `config/` for experiment configurations (using json, gin, or yaml).
+- `experiments/` for executable scripts.
+- `docs/` for ReadTheDocs or Doxygen documentation, and decision logs.
+- `tests/` for unit tests and verification scripts.

.agents/rules/method.md

@@ -0,0 +1,14 @@
+# Methodology and Process Rules
+
+The agent must adhere to the following scientific and operational practices, focused on robustness and reproducibility:
+
+## 1. Scientific Context & Methodology
+- **Research Focus**: Maintain focus on the project's objective: studying how controller modularity affects learning speed, coordination, and fault tolerance in brittle-star locomotion.
+- **Hypothesis-Driven Design**: Base execution on clear hypotheses. Document all design decisions prior to implementation (in `/docs/decisions/` or via Artifacts/Plans).
+- **Scaffolding Approach**: Start development with simple setups before scaling to complex environments and varying morphologies.
+- **Value of Negative Results**: Understand that a controller failing to learn locomotion, when coupled with a thorough analysis of the failure, holds strong scientific value. Do not artificially force a positive result.
+
+## 2. Reproducibility Protection
+- **Dependency Management (uv)**: This project strictly prefers `uv`. Do **not** manually modify the `uv.lock` file. Add dependencies via `uv add <package>` and sync environments via `uv sync --frozen` (or via devcontainers).
+- **Consistent Initialization**: The AI must avoid hidden randomness. Ensure that all runs use consistent seed initialization.
+- **Experiment Tracking**: Use Weights & Biases (wandb) for tracking and logging all experiments and parameters. Ensure every experiment-friendly parameter is properly logged.

.agents/skills/lint/SKILL.md

@@ -0,0 +1,24 @@
+---
+name: Lint and Format Code
+description: Instructions for checking code style and formatting using Ruff.
+---
+
+# Skill: Lint and Format Code
+
+The goal of this skill is to enforce the project's adherence to the Google style standard and to maintain high code quality across the python source files.
+
+## Instructions
+
+1. **Format Code**: To automatically format all Python files according to the `ruff.toml` specifications:
+   ```bash
+   uv run ruff format .
+   ```
+
+2. **Check for Lints / Auto-fix**: To check the repository for style violations and automatically fix safe corrections:
+   ```bash
+   uv run ruff check --fix .
+   ```
+
+## Important Considerations
+- The `ruff.toml` file at the root handles all lint and format configuration. Do not ignore configurations when applying fixes.
+- If Ruff points out complex errors that cannot be auto-fixed, analyze the code and manually address the violations, prioritizing descriptive naming and adherence to the single-responsibility principle.

.agents/skills/test/SKILL.md

@@ -0,0 +1,25 @@
+---
+name: Run Tests
+description: Instructions for executing the project test suite to verify code correctness.
+---
+
+# Skill: Run Tests
+
+The goal of this skill is to verify that the project is functioning correctly after development changes. 
+
+## Instructions
+
+1. **Verify Environment**: The project operates dynamically with hardware acceleration (GPU via JAX) and uses `uv` for dependency/execution management.
+2. **Execute Tests**:
+   - To run basic verification (e.g., checking JAX initialization and hardware detection), run the test script:
+     ```bash
+     uv run python -m tests.test_jax_init
+     ```
+   - If a broader suite of modular tests is added (e.g., `pytest`), execute tests in the `tests/` directory with:
+     ```bash
+     uv run pytest tests/
+     ```
+
+## Important Considerations
+- If you are running tests inside the local environment without a Devcontainer, verify whether `uv sync --frozen` (for CPU) or `uv sync --frozen --extra cuda` (for GPU) has been executed to avoid import errors.
+- Do not run bare `python ...` without `uv run` locally, unless you are strictly operating inside a pre-activated `.venv` inside a Devcontainer.

.commitlintrc.json

@@ -0,0 +1,3 @@
+{
+  "extends": ["@commitlint/config-conventional"]
+}

.devcontainer/Dockerfile

@@ -0,0 +1,30 @@
+FROM mcr.microsoft.com/devcontainers/python:3.12
+
+# Install uv for dependency management
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
+
+# Install system dependencies for JAX, OpenGL, and Mujoco
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    libgl1 \
+    libosmesa6-dev \
+    libglew-dev \
+    libglfw3 \
+    patchelf \
+    && apt-get clean && rm -rf /var/lib/apt/lists/*
+
+# Set environment variables for hardware acceleration and rendering
+ENV LD_LIBRARY_PATH=/usr/lib/nvidia
+ENV NVIDIA_VISIBLE_DEVICES=all
+ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility,graphics
+
+# Set the workspace path (defaults to /workspaces/project but can be overridden)
+ARG WORKSPACE_PATH=/workspaces/project
+ENV WORKSPACE_PATH=$WORKSPACE_PATH
+ENV PATH="$WORKSPACE_PATH/.venv/bin:$PATH"
+
+# Ensure the workspace and .venv directories exist and are owned by the vscode user
+# This prevents permission errors when the postCreateCommand runs uv sync
+RUN mkdir -p $WORKSPACE_PATH/.venv && chown -R vscode:vscode $WORKSPACE_PATH
+
+USER vscode
+WORKDIR $WORKSPACE_PATH

.devcontainer/devcontainer.json

@@ -0,0 +1,48 @@
+{
+    "name": "Brittle Star JAX/CUDA",
+    "build": {
+        "dockerfile": "Dockerfile",
+        "context": "..",
+        "args": {
+            "WORKSPACE_PATH": "/workspaces/${localWorkspaceFolderBasename}"
+        }
+    },
+    "workspaceFolder": "/workspaces/${localWorkspaceFolderBasename}",
+    "customizations": {
+        "vscode": {
+            "extensions": [
+                "ms-python.python",
+                "charliermarsh.ruff",
+                "ms-python.vscode-pylance",
+                "ms-toolsai.jupyter"
+            ],
+            "settings": {
+                "python.defaultInterpreterPath": "${containerWorkspaceFolder}/.venv/bin/python",
+                "python.analysis.localRoot": "${containerWorkspaceFolder}",
+                "python.analysis.extraPaths": [
+                    "${containerWorkspaceFolder}/.venv/lib/python3.12/site-packages"
+                ]
+            }
+        },
+        "jetbrains": {
+            "plugins": [
+                "Pythonid",
+                "fleet.python",
+                "com.koxudaxi.ruff"
+            ]
+        }
+    },
+    "remoteUser": "vscode",
+    "runArgs": [
+        "--device", "nvidia.com/gpu=all"
+    ],
+    // Ensure the .venv persists using a named volume for performance and parity
+    "mounts": [
+        "source=${localWorkspaceFolderBasename}-venv,target=${containerWorkspaceFolder}/.venv,type=volume"
+    ],
+    // Invoke the hardware-aware sync script
+    "postCreateCommand": "bash ./.devcontainer/post-create.sh",
+    "features": {
+        "ghcr.io/devcontainers/features/common-utils:1": {}
+    }
+}

.devcontainer/post-create.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+set -e
+
+# Detect if a GPU is available via nvidia-smi.
+# This works for Linux hosts and Windows (WSL2) with NVIDIA Container Toolkit.
+# On Mac (Apple Silicon) or systems without NVIDIA GPUs, this will skip the 'cuda' extra.
+if command -v nvidia-smi &> /dev/null && nvidia-smi &> /dev/null; then
+    echo "GPU detected. Syncing with 'cuda' extra..."
+    uv sync --frozen --extra cuda
+else
+    echo "No GPU detected or nvidia-smi failed. Syncing without 'cuda' extra..."
+    uv sync --frozen
+fi
+
+echo "Environment synced successfully."
+
+# Install pre-commit hooks so they are active in the devcontainer
+uv run pre-commit install

.env.example

@@ -0,0 +1,12 @@
+# Brittle Star Project Environment Variables
+# Copy this file to .env and fill in your values.
+# IMPORTANT: Never commit the actual .env file, it is in .gitignore
+
+# ---------------------------- #
+# Weights and Biases API Key   #
+# ---------------------------- #
+# To find your API key:
+# 1. Log in to wandb.ai
+# 2. Go to User Settings (https://wandb.ai/settings)
+# 3. Scroll down to the "API keys" section
+WANDB_API_KEY=your_api_key_here

.gitattributes

@@ -0,0 +1,14 @@
+# Model weights
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+
+# Media
+*.mp4 filter=lfs diff=lfs merge=lfs -text
+*.avi filter=lfs diff=lfs merge=lfs -text
+
+# Datasets
+*.csv filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text

.github/scripts/prepare_docs.py

@@ -0,0 +1,28 @@
+import os
+import glob
+import re
+import shutil
+
+folders_to_copy = ["src", "scripts", "configs"]
+for folder in folders_to_copy:
+    if os.path.exists(folder):
+        shutil.copytree(folder, f"docs/{folder}", dirs_exist_ok=True)
+
+for filepath in glob.glob("docs/**/*.md", recursive=True):
+    with open(filepath, "r", encoding="utf-8") as f:
+        content = f.read()
+
+    # RULE A: Fix links pointing OUT to src/, scripts/, or configs/
+    # Logic: Because the folders were moved one level deeper, we remove exactly ONE '../'
+    content = re.sub(
+        r"\]\(\.\./((?:\.\./)*)(src|scripts|configs)/([^)]*)\)", r"](\1\2/\3)", content
+    )
+
+    # RULE B: Fix links pointing FROM the copied files back TO the original docs/ folder
+    # Logic: Since these files are now inside docs/, the 'docs/' segment in the path is redundant.
+    content = re.sub(r"\]\(((?:\.\./)+)docs/([^)]*)\)", r"](\1\2)", content)
+
+    with open(filepath, "w", encoding="utf-8") as f:
+        f.write(content)
+
+print("Successfully imported external files and adjusted markdown links.")

.github/workflows/lint.yml

@@ -0,0 +1,23 @@
+name: Lint
+
+on:
+  pull_request:
+    types: [opened, synchronize, ready_for_review]
+    branches:
+      - main
+      - dev
+
+jobs:
+  ruff:
+    if: ${{ !github.event.pull_request.draft }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Run Ruff check
+        uses: astral-sh/ruff-action@v1
+        with:
+          args: "check"
+      - name: Run Ruff format check
+        uses: astral-sh/ruff-action@v1
+        with:
+          args: "format --check"

.github/workflows/publish-docs.yml

@@ -0,0 +1,38 @@
+name: Publish docs via GitHub Pages
+
+on:
+  push:
+    branches:
+      - main
+      - dev
+      - docs/*
+
+jobs:
+  build:
+    name: Deploy docs
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+          ssh-key: ${{ secrets.DEPLOY_KEY }}
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+
+      - name: Install dependencies
+        run: pip install mkdocs-material
+
+      - name: Prepare external docs
+        run: python .github/scripts/prepare_docs.py
+
+      - name: Configure Git identity
+        run: |
+          git config --global user.name "github-actions[bot]"
+          git config --global user.email "github-actions[bot]@users.noreply.github.com"
+
+      - name: Deploy to GitHub Pages
+        run: mkdocs gh-deploy --force

.github/workflows/test.yml

@@ -0,0 +1,31 @@
+name: Test
+
+on:
+  pull_request:
+    types: [opened, synchronize, ready_for_review]
+    branches:
+      - main
+      - dev
+
+jobs:
+  pytest:
+    if: ${{ !github.event.pull_request.draft }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+        with:
+          version: "latest"
+          
+      - name: Setup python
+        uses: actions/setup-python@v5
+        with:
+          python-version-file: ".python-version"
+          
+      - name: Install dependencies
+        run: uv sync --frozen
+        
+      - name: Run tests with pytest
+        run: uv run pytest tests/

.github/workflows/update_hpc_requirements.yml

@@ -0,0 +1,43 @@
+name: Update HPC requirements
+
+on:
+  push:
+    paths:
+      - pyproject.toml
+    branches:
+      - main
+      - dev
+      - "ci/**"
+
+jobs:
+  update-hpc-requirements:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+      pull-requests: write
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.head_ref || github.ref_name }}
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+
+      - name: Regenerate env/hpc/requirements.txt
+        run: uv run scripts/hpc/export_requirements.py
+
+      - name: Create Pull Request with updated requirements
+        uses: peter-evans/create-pull-request@v6
+        with:
+          token: ${{ secrets.GITHUB_TOKEN }}
+          commit-message: "chore(hpc): update env/hpc/requirements.txt from pyproject.toml"
+          title: "chore(hpc): update HPC requirements"
+          body: "Automatically generated pull request to update `env/hpc/requirements.txt` based on recent changes to `pyproject.toml`."
+          branch: chore/auto-update-hpc-requirements
+          base: ${{ github.ref_name }}
+          author: "github-actions[bot] <github-actions[bot]@users.noreply.github.com>"

.gitignore

@@ -1,3 +1,13 @@
+# Model files
+artifacts/*
+runs/*
+wandb/
+outputs/
+multirun/
+metrics/
+adjacency_debug.txt
+vids/
+
 # Python-generated files
 __pycache__/
 *.py[oc]
@@ -368,7 +378,6 @@ celerybeat.pid
 # Environments
 .env
 .venv
-env/
 venv/
 ENV/
 env.bak/
@@ -468,7 +477,6 @@ tags
 [Ll]ib
 [Ll]ib64
 [Ll]ocal
-[Ss]cripts
 pyvenv.cfg
 .venv
 pip-selfcheck.json
@@ -515,4 +523,7 @@ Icon
 Network Trash Folder
 Temporary Items
 .apdisk
+*.pdf
 
+# plot directory
+poster_plots/

.pre-commit-config.yaml

@@ -0,0 +1,21 @@
+repos:
+  - repo: local
+    hooks:
+      - id: ruff-format
+        name: ruff format (uv)
+        entry: uv run ruff format
+        language: system
+        types: [python]
+
+  - repo: https://github.com/alessandrojcm/commitlint-pre-commit-hook
+    rev: v9.16.0
+    hooks:
+      - id: commitlint
+        stages: [commit-msg]
+        additional_dependencies: ["@commitlint/config-conventional"]
+
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v5.0.0
+    hooks:
+      - id: no-commit-to-branch
+        args: ['--branch', 'main', '--branch', 'dev']

.vscode/settings.json

@@ -0,0 +1,7 @@
+{
+    "python.testing.pytestArgs": [
+        "tests"
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true
+}

README.md

@@ -0,0 +1,62 @@
+# Brittle Star
+
+> What is the impact of different levels of controller-modularity on the learning-speed, coordination and tolerance for
+defects (e.g. amputations) in brittle-star-like robots trained with Reinforcement Learning?
+
+## Quick start
+
+### Local setup
+
+To set up the UV module, you can run the following command:
+
+```bash
+uv sync --frozen
+```
+
+## Repository Structure
+
+```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)
+│   └── experiment_logger/  # Standalone logging package
+└── 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
+   uv run python scripts/train.py ppo.learning_rate=0.001 logging.track=true
+   ```
+
+2. **Monitor progress:**
+   See [Tracking & Monitoring](docs/api/tracking.md).
+
+3. **Simulate a trained model:**
+   See [Simulation & Evaluation](docs/api/simulation.md).
+
+4. **Compare fault tolerance of models:**
+   See [Checkpoint & Model Evaluation](docs/api/evaluation.md)
+
+## Results & Reproduction
+
+See **[docs/api/reproduction.md](docs/api/reproduction.md)** to learn how to access our public [Weights & Biases (WandB) project](https://wandb.ai/SEL3-2026-Groep-4/final-models-v2?nw=96mloffsyq), retrieve specific run parameters, and run the training/evaluation reproduction workflow.
+
+## HPC
+
+See **[docs/HPC.md](docs/HPC.md)** for the full guide, including environment setup, cluster selection, interactive debugging, and job submission.
+
+## Documentation
+
+Please find all documentation and a starting point for more information in [corresponding README](./docs/README.md).

configs/README.md

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

configs/architecture/centralized.yaml

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

configs/architecture/decentralized.yaml

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

configs/arena/default.yaml

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

configs/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-v2/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: 10
+  wandb_project_name: "final-models-v2"
+
+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

@@ -0,0 +1,12 @@
+# Directed Locomotion Environment
+# Baseline task setting.
+
+task: DIRECTED_LOCOMOTION
+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: 3.0
+light_perlin_noise_scale: 0

configs/environment/light_escape.yaml

@@ -0,0 +1,12 @@
+# Light Escape Environment
+# Advanced task requiring movement away from light source.
+
+task: LIGHT_ESCAPE
+simulation_time: 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: 3.0
+light_perlin_noise_scale: 200 # Must be integer factor of 200

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,25 @@
+# @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: 10
+comparison_output_csv: "runs/evaluation/comparison.csv"
+
+# Paths to the .cleanrl_model files to be compared (relative to workspace root).
+comparison_models:
+  - "runs/final-v2-centralized/artifacts/12-19-01_checkpoint_v22/checkpoint_step_230.flax"
+  - "runs/final-v2-fully-conn/artifacts/14-02-00_checkpoint_v17/checkpoint_step_180.flax"
+  - "runs/final-v2-ring/artifacts/15-27-03_checkpoint_v21/checkpoint_step_220.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/base.yaml

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

configs/experiment/dev_test.yaml

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

configs/experiment/hpc_smoke_test.yaml

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

configs/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-v2/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: 10
+  wandb_project_name: "final-models-v2"
+
+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

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

configs/logging/hpc.yaml

@@ -0,0 +1,9 @@
+track: true
+wandb_project_name: "hpc-default"
+wandb_entity: "SEL3-2026-Groep-4"
+save_model: true
+save_checkpoints: true
+upload_final_model: true
+upload_checkpoints: true
+checkpoint_frequency: 100
+hf_entity: ""

configs/logging/wandb_enabled.yaml

@@ -0,0 +1,13 @@
+# WandB Enabled Logging Configuration
+# For production/cloud experiments with weights synced.
+
+track: true
+wandb_project_name: "default-project"
+wandb_entity: "SEL3-2026-Groep-4"
+capture_video: false
+save_model: true
+save_checkpoints: true
+checkpoint_frequency: 100
+upload_final_model: true
+upload_checkpoints: false
+hf_entity: ""

configs/main_config.yaml

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

configs/morphology/2_arms.yaml

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

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/3_arms.yaml

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

configs/morphology/5_arms_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.yaml

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

configs/morphology/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/morphology/partial_amputation.yaml

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

configs/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/debug.yaml

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

configs/ppo/default.yaml

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

configs/ppo/dev_larger_timesteps_larger_rolloutsteps.yaml

@@ -0,0 +1,16 @@
+learning_rate: 0.0003
+total_timesteps: 1228800
+num_envs: 32
+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.005
+vf_coef: 1.0
+max_grad_norm: 0.5
+target_kl: null

configs/ppo/smoke_test.yaml

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

configs/ppo/stable.yaml

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

configs/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-v2/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: 10
+  wandb_project_name: "final-models-v2"
+
+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

@@ -0,0 +1,30 @@
+# Default Simulation Settings
+# These values are used by scripts/simulate.py
+
+# Path to the trained model (optional)
+model_path: null
+
+# Script behavior
+headless: false
+# In headless mode this is required; in viewer mode null means "infinite".
+max_steps: 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
+
+video_width: 640
+video_height: 80
+video_fps: 60
+
+# 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/CONTRIBUTING.md

@@ -0,0 +1,44 @@
+# Contribution Guidelines
+
+This document outlines the contribution protocols for the scientific software engineering project focusing on bio-inspired control architectures for brittle-star-like robots. The primary objective of this project is to produce scientific insight, rather than a commercial product.
+
+## 1. Scientific Context & Methodology
+
+* **Research Focus:** The goal is to study how controller modularity affects learning speed, coordination, and fault tolerance in brittle-star locomotion.
+* **Hypothesis-Driven Design:** Clear hypotheses must dictate a structured methodology and rigorous evaluation. All design decisions must be formally documented prior to implementation.
+* **Scaffolding Approach:** Development must start with simple setups before progressively increasing the complexity of environments and morphologies.
+* **Evaluation of Results:** Negative results possess scientific validity when thoroughly analyzed. If a controller fails to learn locomotion, providing a comprehensive analysis of the failure is considered a strong scientific contribution.
+* **Reproducibility:** Contributors must utilize fixed library versions. Configuration systems (such as json, gin, or yaml) must be employed to ensure reproducible runs.
+
+## 2. Clean Code & Code Quality
+
+Code readability is paramount, as code is read far more frequently than it is written.
+
+* **Naming Conventions:** Variables and functions must utilize consistent, intention-revealing names. A long, descriptive name is strictly preferred over a short name accompanied by a comment.
+* **Function Design:** Functions must be modular and adhere to the single responsibility principle. Arguments must be minimized, and boolean flag arguments controlling behavior should be avoided.
+* **Commenting:** Code must document the "how," while comments are strictly reserved for documenting the "why". Commented-out code is prohibited and must be deleted via version control.
+* **YAGNI:** Contributors must adhere to the "You Aren't Gonna Need It" (YAGNI) principle and actively avoid premature optimization.
+* **Notebooks:** Jupyter Notebooks are strictly limited to quick prototyping, tutorials, demonstrations, or post-processing analysis. They are explicitly forbidden for general software development because they discourage modularity.
+
+## 3. Version Control & Repository Structure
+
+* **Git Practices:** Commits must be frequent and small. Each commit should relate to exactly one piece of functionality.
+* **Branching Strategy:** The `dev` branch serves as the integration branch for pushing and merging code. Only stable releases may be pushed to the `main` branch.
+* **Artifact Management:** Data files, trained models, and large datasets must never be committed directly to Git. Git Large File Storage (LFS) must be used for tracking large files. **All developers must have `git-lfs` installed locally** (see [DEVELOPMENT.md](./DEVELOPMENT.md) for setup).
+* **Repository Layout:** The repository must maintain the following core directories: `src/` for algorithms, `env/` for MuJoCo wrappers, `config/` for experiment configurations, `experiments/` for scripts, `docs/` for Doxygen or ReadTheDocs documentation, and `tests/` for unit tests.
+
+## 4. Architecture & Tooling
+
+* **Algorithms & Frameworks:** Proximal Policy Optimization (PPO) is the recommended baseline algorithm. CleanRL should be used as a starting point and adapted for continuous action spaces. All Artificial Neural Network (ANN) controller architectures must be implemented using Flax.
+* **Simulation:** The simulation environment utilizes a MuJoCo brittle star. XML MuJoCo structures must remain realistic and respect morphological constraints.
+* **Experiment Tracking:** Weights & Biases (wandb) must be utilized for tracking and logging all experiments.
+* **Code Styling:** All code must conform to the chosen style guide (Google standard). This is enforced via `uv` using **ruff** and pre-commit hooks.
+
+## 5. AI-Assisted Development & Code Review
+
+This project supports AI-assisted development to enhance productivity, but contributors must take full responsibility for all AI-generated outputs.
+
+* **Self-Review Requirement:** Contributors must thoroughly self-review all AI-assisted code, documentation, and configurations before requesting peer review. This includes verifying correctness, adherence to project standards, scientific validity, and integration with existing code.
+* **Quality Standards:** AI-generated content must meet the same rigorous standards as manually written code, including proper testing, documentation, and alignment with the scientific methodology outlined in Section 1.
+* **Available Skills:** This project provides specific AI skills for common tasks (located in `.agents/skills/`), including linting and testing workflows. Contributors should leverage these skills to maintain consistency and quality.
+* **Transparency:** When using AI assistance for complex algorithmic decisions or scientific design choices, contributors should document the rationale in commit messages or code comments where appropriate.

docs/DEVELOPMENT.md

@@ -0,0 +1,72 @@
+# Development Guide
+
+This guide outlines how to set up the development environment for this project, prioritizing **reproducible builds**, **environment parity**, and **cross-hardware compatibility**.
+
+## Reproducibility & uv
+
+This project uses [uv](https://github.com/astral-sh/uv) to manage dependencies and virtual environments. The `uv.lock` file is the absolute source of truth for package versions and must always be committed.
+
+### Source of Truth
+
+- **Never modify `uv.lock` manually.**
+- To add a dependency, run `uv add <package>`.
+- To update dependencies, run `uv lock --upgrade`.
+- To sync your environment with the lockfile, run `uv sync --frozen`.
+
+## Git LFS (Critical)
+
+**All developers must have Git LFS installed locally.** This repository tracks model weights (`.pt`, `.safetensors`, etc.), recordings (`.mp4`), and datasets using Git LFS.
+
+- **Setup:** Run `git lfs install` after cloning this repository. If you are using the `.devcontainer` or `flake.nix`, LFS is typically available automatically.
+- If you clone without LFS installed, run `git lfs pull` after installation to fetch the actual data files instead of the small pointer files.
+
+## Devcontainer Setup (Recommended)
+
+The devcontainer provides an identical experience to local development but with all system dependencies pre-configured. It automatically detects your hardware (GPU vs CPU) and syncs the appropriate dependencies.
+
+### Prerequisites
+
+- Docker Desktop or Docker Engine.
+- [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) (for GPU support).
+
+### Setup for VS Code
+
+1. Install the [Dev Containers](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-containers) extension.
+2. Open the project and click **Reopen in Container**.
+3. On first launch, the `post-create.sh` script will:
+   - Detect if an NVIDIA GPU is available via `nvidia-smi`.
+   - Run `uv sync --frozen --extra cuda` if a GPU is found.
+   - Run `uv sync --frozen` otherwise.
+4. The environment is stored in a **named volume** for `.venv` to ensure persistence and performance.
+
+### Setup for JetBrains IDEs
+
+1. The IDE will detect the `.devcontainer/devcontainer.json` file.
+2. The environment is pre-configured to point to `/workspaces/project/.venv`.
+3. The hardware-aware sync will run automatically during container creation.
+
+## Local Development (Alternative)
+
+If you prefer not to use Docker:
+
+1. Install [uv](https://docs.astral.sh/uv/getting-started/installation/).
+2. Run `uv sync --frozen` (CPU) or `uv sync --frozen --extra cuda` (GPU).
+
+## Hardware Acceleration (JAX)
+
+Verify your setup by running the JAX initialization test:
+
+```bash
+uv run pytest tests/test_jax_init.py
+```
+
+In the devcontainer, this will succeed on both CPU and GPU. A `GpuDevice` is expected if a GPU is detected and the `cuda` extra was installed.
+
+## Logging & Monitoring
+
+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/HPC.md

@@ -0,0 +1,89 @@
+# HPC Guide
+
+Full documentation: [https://docs.hpc.ugent.be/](https://docs.hpc.ugent.be/)
+
+## Storage Overview
+
+- **Run Outputs**: Written to `$VSC_SCRATCH` during the job (fast I/O) and copied to `$VSC_DATA` at the end for persistence.
+- **Virtual Environments**: Managed on **`$VSC_DATA`** by mirroring configuration files. This avoids the 3GB home quota without requiring symlinks in the project root.
+
+## Initial Environment Setup
+
+Run **once** after cloning the repository. This script handles all modules, mirroring, and environment synchronization.
+
+```bash
+# Option A: Interactive (on a compute node)
+module swap cluster/donphan  # Debug cluster (CPU only)
+# OR for GPU clusters:
+# module swap cluster/joltik
+# module swap cluster/accelgor
+# module swap cluster/litleo
+
+qsub -I -l nodes=1:gpus=1  # Only for GPU clusters
+cd "${PBS_O_WORKDIR}"
+bash scripts/hpc/install.sh
+
+# Option B: Batch (Run in background)
+# NOTE: GPU clusters (joltik/accelgor/litleo) require -l gpus=1 at runtime
+qsub -l gpus=1 scripts/hpc/install.sh
+```
+
+## Production vs. Debug Clusters
+
+Our scripts are cluster-agnostic and do **not** have hardcoded GPU requirements. Instead, you must request GPUs at runtime using the `-l gpus=1` flag when submitting to a production GPU cluster.
+
+### Debugging (Donphan)
+
+The `donphan` cluster does not support GPUs. Simply run the scripts without extra resource flags:
+```bash
+module swap cluster/donphan
+qsub scripts/hpc/train.pbs
+```
+
+### Production (Joltik, Accelgor, Litleo)
+
+These clusters provide GPU acceleration and **require** a GPU request at runtime:
+```bash
+module swap cluster/joltik  # or accelgor/litleo
+qsub -l gpus=1 scripts/hpc/train.pbs
+```
+
+## Interactive Debugging
+
+To activate your environment for interactive work, simply run the same `install.sh` script.
+
+```bash
+qsub -I -l nodes=1:ppn=4 -l walltime=1:00:00
+cd "$PBS_O_WORKDIR"
+bash scripts/hpc/install.sh
+```
+
+### Verification Commands
+
+After installation, run these commands to ensure your environment is set up correctly:
+
+1. **Verify Quota Safety**:
+   ```bash
+   ls -d venvs 2>/dev/null && echo "FAIL" || echo ">>> PASS: Project root is clean."
+   ```
+
+2. **Verify Library Versions (NumPy Fix)**:
+   ```bash
+   python -c "import numpy; print(f'NumPy: {numpy.__version__}')"
+   # Expected: 2.x.x (Venv version), not 1.2x (System version)
+   ```
+
+3. **Verify GPU Access**:
+   ```bash
+   python -c "import torch, jax; print(f'GPU: {torch.cuda.is_available()}'); print(f'JAX: {jax.devices()}')"
+   ```
+
+## Managing Dependencies
+
+`env/hpc/requirements.txt` is auto-generated from `pyproject.toml`. To regenerate:
+
+```bash
+uv run scripts/hpc/export_requirements.py
+```
+
+Modules listed in `env/hpc/modules.txt` are automatically excluded from the pip requirements to save space and use HPC-optimized binaries.

docs/README.md

@@ -0,0 +1,52 @@
+# Documentation
+
+Welcome to the Brittle Star project documentation. This codebase contains the implementations and research for the scientific evaluation of controller modularity in brittle-star-like robots trained using Reinforcement Learning.
+
+For the core codebase, scripts, and contribution history, visit our [GitHub Repository](https://github.com/SELab-3-2026/SEL3-2026-Groep-4).
+
+## Core Requirements & Guides
+
+- **[Installation Instructions](./DEVELOPMENT.md)**: Steps to set up your development environment locally or in a devcontainer using `uv`, including GPU configuration. For High-Performance Computing (HPC) setup details, see the **[HPC Guide](./HPC.md)**.
+- **[How to Run Experiments](./api/training.md)**: A complete guide on running training jobs, setting custom hyperparameters, and overriding config options using Hydra.
+- **[Results & Reproduction](./api/reproduction.md)**: Guide on how to access our public WandB training runs table and reproduce our training and evaluation phases (determining the best checkpoint vs. comparing architectures).
+- **[Contribution Guidelines](./CONTRIBUTING.md)**: Standards, rules, and best practices for developing and adding code to the repository.
+- **[Repository Structure](#repository-structure)**: Overview of the directories and files within the codebase.
+
+## Repository Structure
+
+```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)
+│   └── experiment_logger/  # Standalone logging package
+└── tests/                  # Unit and integration tests
+```
+
+## 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 Scheme](./design/communication.md): Message propagation, Nerve-Net style.
+- [Modularity & Topology](./design/controllers.md): Macroscopic brain topology, centralized, arm-level, segment-level.
+- [Input & Action Spaces](./design/input_action_spaces.md): Description of the model's input and output.
+- [Reinforcement Learning Algorithm](./design/learning_algorithm.md): RL techniques, i.e. PPO.
+- [Reward Function & Observation Space](./design/reward_function.md): Goals, fitness tracking, and reward structures.
+
+## API reference (`/api`)
+
+If you are interested in the "how do I use it?"
+
+- [Brittle Star Environment](./api/environment.md): MuJoCo environment interaction and configuration.
+- [Training Models](./api/training.md): How to configure and run experiments.
+- [Tracking & Monitoring](./api/tracking.md): Setting up WandB and TensorBoard to monitor runs.
+- [Checkpoint & Model Evaluation](./api/evaluation.md): Evaluating checkpoints and comparing fault tolerance.
+- [Interactive Simulation & Visualization](./api/simulation.md): Visualizing models in the MuJoCo viewer or rendering simulation videos.
+- [Analysis & Plotting Tools](./api/analysis.md): Comparing checkpoints and generating plots.
+- [Results & Reproduction](./api/reproduction.md): Accessing WandB results and running reproduction pipelines.

docs/api/analysis.md

@@ -0,0 +1,89 @@
+# 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/environment.md

@@ -0,0 +1,30 @@
+# Brittle star environment
+
+## Creation
+
+The environment package contains a factory class `BrittleStarEnvFactory` 
+that creates instances of the environment/morphologies/... It uses the
+configuration classes defined in `env_config.py` to create the instances.
+
+## Configuration
+
+The data classes in `env_config` have default values as stated in the tutorials.
+
+* MorphologyConfig: configuration for the morphology of the brittle star. Contains
+number of arms, number of segments per arm, and control mode.
+* ArenaConfig: configuration for the arena. Sets the size of the arena, whether to 
+set the ground floor to sand, attach a target and sizes of the walls.
+* EnvConfig: configuration for the environment. These set shared settings
+such as camera locations, simulation time and the task.
+
+## Backend and Task enums
+
+The Backend enum specifies either an MJC or MJX backend.
+
+* MJC: runs on CPU
+* MJX: uses jax on the gpu
+
+The Task enum specifies which task to use. 2 items are present:
+
+* DIRECTED_LOCOMOTION: move to a target location
+* LIGHT_ESCAPE: situation where the robot must move to a darker location

docs/api/evaluation.md

@@ -0,0 +1,60 @@
+# 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 & Fault 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.
+
+---
+
+For a step-by-step walkthrough on using these evaluation phases to reproduce our project results, see the **[Results & Reproduction Guide](./reproduction.md)**.

docs/api/reproduction.md

@@ -0,0 +1,108 @@
+# Results & Reproduction
+
+This guide explains how to access our official training logs and reproduce our results.
+
+Our official training runs, model configurations, and metrics are publicly hosted on Weights & Biases (WandB).
+
+---
+
+## Weights & Biases (WandB) Project
+
+All experiments, final models, and training logs are tracked in our public WandB project:
+
+* **Official Runs Table**: [WandB final-models-v2 Table](https://wandb.ai/SEL3-2026-Groep-4/final-models-v2/table?nw=96mloffsyq)
+
+This page lists the verified runs with their architecture types, morphology definitions, evaluation metrics, and final model performance.
+
+### How to Reproduce a Run from WandB
+
+Weights & Biases provides a built-in feature to extract the exact parameters and commands used for any given run:
+
+1. Open the [WandB final-models-v2 Table](https://wandb.ai/SEL3-2026-Groep-4/final-models-v2/table?nw=96mloffsyq).
+2. Click on the name of the run you wish to reproduce to open its detail page.
+3. In the top-right corner of the run header (next to the run name, not the main workspace header), click the **three dots (`...`)** menu.
+4. Select **"Reproduce run"**. This will display the exact command-line arguments and configuration settings used to execute that run.
+
+---
+
+## Local & HPC Reproduction Workflow
+
+To reproduce our training and evaluation phases locally or on an HPC cluster, follow the procedures below.
+
+### 1. Environment Setup
+
+To ensure identical package versions (including JAX, Flax, and MuJoCo), sync your environment using the lockfile:
+
+```bash
+uv sync --frozen
+```
+
+### 2. Training Phase
+
+Run the training script using the exact parameters retrieved from WandB's "Reproduce run" page or from a downloaded `_metadata.yaml` file:
+
+```bash
+uv run python scripts/train.py experiment=my_experiment ppo.learning_rate=0.001 experiment.seed=42
+```
+
+---
+
+## Evaluation Phases
+
+Reproducing our evaluation results is divided into two distinct phases:
+
+### Phase 1: Determining the Best Checkpoint
+
+During training, checkpoints are saved at regular intervals. To determine which of these checkpoints performed the best:
+
+1. **Evaluate Checkpoints Post-Training**:
+   If checkpoint evaluation was not run during training, scan the completed run's checkpoints folder by pointing to the final model path:
+
+   ```bash
+   uv run python scripts/evaluate_checkpoints.py simulation.model_path=runs/your_run_dir/final_model.flax
+   ```
+
+   This script runs deterministic rollouts for every checkpoint in `runs/your_run_dir/checkpoints/`.
+
+2. **Locate the Results**:
+   The evaluations are saved to:
+
+   ```text
+   runs/your_run_dir/metrics/checkpoint_evaluation.csv
+   ```
+
+   Analyze this CSV to find the checkpoint iteration with the highest average return or target success rate. This checkpoint will be used for cross-architecture comparisons.
+
+### Phase 2: Comparing Checkpoints Between Architectures
+
+Once the best checkpoints for each architecture are identified, they are compared under shared, standardized environments (including fault tolerance checks such as leg amputations).
+
+1. **Configure the Comparison Models**:
+   Open or create an evaluation config file (e.g., `configs/evaluation/poster.yaml`) and add the paths to the best checkpoints:
+
+   ```yaml
+   # configs/evaluation/poster.yaml
+   evaluation:
+     comparison_models:
+       - runs/run_arch_centralized/checkpoints/checkpoint_best.flax
+       - runs/run_arch_decentralized/checkpoints/checkpoint_best.flax
+   ```
+
+2. **Execute the Comparison Script**:
+   Run the comparison script using your config:
+
+   ```bash
+   uv run python scripts/compare_models.py evaluation=poster
+   ```
+
+   This script runs multiple sequential evaluation episodes (defined by `comparison_num_episodes` starting at `comparison_base_seed`) for every model across the selected morphologies.
+
+3. **Analyze Comparison Metrics**:
+   The script writes a consolidated CSV file to `metrics/model_comparison.csv` containing:
+
+   * **`eval_return`**: The cumulative return.
+   * **`approx_max_velocity`**: The distance covered per step.
+   * **`reached_target`**: Navigational success rates.
+   * **`arm_0` to `arm_4`**: Active segments per arm (indicating damage/amputations).
+
+This CSV can then be passed to the plotting scripts (e.g., `scripts/plots/analyze_comparisons.py`) to generate visualization plots. For details on configuration and outputs, see the **[Analysis & Plotting Guide](./analysis.md)**.

docs/api/simulation.md

@@ -0,0 +1,54 @@
+# Interactive Simulation & Visualization
+
+The simulation pipeline allows you to visualize trained models and observe their behavior 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`
+
+### Top-Down and Follow Cameras
+
+Using the following script, you can render a top-down and follow camera view for multiple models at once:
+
+```bash
+uv run scripts/poster_visualisations/render_poster_videos.py \
+  runs/final-models/centralized/.../final_model.flax \
+  runs/final-models/fully-connected/.../final_model.flax \
+  runs/final-models/ring/.../final_model.flax \
+  --max-steps 10000 --width 640 --height 480 --fps 60 \
+  --output-root vids/poster/
+```
+For batch evaluation, checkpoint analysis, and cross-model architecture comparisons, see the **[Checkpoint & Model Evaluation Guide](./evaluation.md)**.
+

docs/api/tracking.md

@@ -0,0 +1,64 @@
+# 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](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).
+
+## Developer Logging API
+
+For details on the developer API of our internal logging library (how backend routing, checkpoint synchronization, and singleton initialization works), see the **[Experiment Logger API Guide](../../src/experiment_logger/README.md)**.

docs/api/training.md

@@ -0,0 +1,64 @@
+# 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.
+
+For a detailed guide on the structure, validation, and usage of our Hydra configuration files, see the **[Brittle Star Configuration System Guide](../../configs/README.md)**.
+
+### 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
+```
+
+## Reproducing Experiments
+
+For detailed steps on how to reproduce training runs, locate run configuration metadata, or reproduce our experiments using Weights & Biases (WandB), see the **[Results & Reproduction Guide](./reproduction.md)**.
+
+---
+
+For more details on evaluation metrics and comparison tools, see [Checkpoint & Model Evaluation](./evaluation.md).
+
+For more details on tracking your experiments, see [Tracking & Monitoring](./tracking.md).

docs/design/actor-critic.md

@@ -0,0 +1,121 @@
+# Actor-Critic Architecture
+
+To process observations into actions, our controllers utilize an Actor-Critic architecture. Because we use Proximal
+Policy Optimization (PPO), the pipeline fundamentally requires separate networks for the policy (Actor) and the value
+estimation (Critic).
+
+**Centralized Architecture (Baseline)**
+
+This pipeline treats the agent as a single entity and uses standard Proximal Policy Optimization (PPO).
+
+- Centralized Actor: Composed of two chained MLPs (Sensor $\rightarrow$ Motor) passing a hidden state between them. The
+  centralized sensor receives the concatenated global state vector of all limbs at once and processes it into a hidden
+  state. The centralized motor receives this hidden state and outputs the joint offsets for all actuators
+  simultaneously. This is mathematically equivalent to using one large MLP with hidden layers, but splitting makes the
+  implementation easier by allowing us to reuse the same components for the decentralized modules.
+- Centralized Critic: Composed of two sequential MLPs (Feature Extractor $\rightarrow$ Critic). Because PPO evaluates
+  the state-value function, this network only receives the concatenated global state vector (no actions). It outputs a
+  single scalar estimating the expected future reward for the entire agent.
+
+Our policy and value networks use separate input networks/feature extractors as advised by the SEL3 course assistants and the blog. For continuous actions this should allow better learning at a small cost.
+
+```mermaid
+graph TD
+    Obs([Global Observation])
+
+    Sens[Sensor]
+    Act[Motor]
+    OutAct([Action Distribution<br/>mean, log_std])
+
+    Feat[Feature extractor]
+    Crit[Critic]
+    OutCrit([Value Estimate<br/>scalar])
+
+    Obs --> Sens
+    Obs --> Feat
+
+    Sens -->|"Hidden state"| Act
+    Feat -->|"Hidden state"| Crit
+
+    Act --> OutAct
+    Crit --> OutCrit
+```
+
+**Decentralized Architecture**
+
+This pipeline utilizes the "Centralized Training with Decentralized Execution" principle, specifically the NerveNet-MLP
+variant.
+
+- Decentralized Actor, split into three distinct models:
+  - Sensor: A local model at each node. It receives its local state plus the goal vector directly, processing them into
+    an initial hidden state.
+  - Propagator: Nodes synchronously compute and exchange messages with connected neighbors for $N$ steps to update
+    their hidden states. See [communication.md](./communication.md) for details.
+  - Motor: A local model uses its final updated hidden state to output the joint offset strictly for its own actuator.
+- Centralized Critic: Composed of two sequential MLPs (Feature Extractor $\rightarrow$ Critic). During training, it
+  acts globally by taking the concatenated state vectors from all sensors to output a single, global state-value scalar
+  evaluating the entire agent's pose.
+
+To keep the implementation simple, we should use one critic per node in our architecture, but only a single, global
+critic for all nodes at once, for the following reasons:
+
+1. Credit Assignment Problem (Ha, 2017): The MuJoCo simulator provides an overall reward based on the brittle star
+  movement progression, e.g. total distance travelled. Using an isolated critic for each node in the network would not
+  allow to determine which local action contributed to the global success. A global critic solves this by evaluating
+  the combined state of the agent at once.
+2. Implementation simplicity: Building a second decentralized message-passing graph for the critic (NerveNet-2) would
+  require more coding. Using a standard MLP that concatenates all raw input vectors is much easier to program while
+  mathematically equivalent.
+
+```mermaid
+graph TD
+    Obs([Local Observation])
+
+    Sens[Sensor]
+    Prop[Propagator]
+    Feat[Feature extractor]
+
+    Mot[Motor]
+    Crit[Critic]
+
+    OutMot([Action Distribution<br/>mean, log_std])
+    OutCrit([Value Estimate<br/>scalar])
+
+    Obs --> Sens
+    Sens -->|"Hidden state"| Prop
+    Obs --> Feat
+
+    Prop -->|"Hidden state"| Mot
+
+
+    Feat -->|"Hidden state"| Crit
+
+    Mot --> OutMot
+    Crit --> OutCrit
+
+    Prop -.->|"message passing"|Prop
+```
+
+## Implementation Details (Network Depth)
+
+Inspired by: [PPO Implementation Details](https://iclr-blog-track.github.io/2022/03/25/ppo-implementation-details/)
+
+The MLPs used in both pipelines are defined with specific hidden layer configurations to balance learning capability
+and computational cost. As of right now, though this might change as we make progress in our experiments, we use:
+
+- Input Networks (Sensors & Feature Extractors): These networks map the raw state inputs to internal hidden states.
+  They are configured as standard dense networks with 3 hidden layers of 300 nodes each (`[300, 300, 300]`) and utilize `tanh`
+  activation functions.
+- Output Networks (Motors, Actors & Critics): The final output models are intentionally kept shallow. The Actor
+  directly projects the hidden state to a continuous action distribution (`mean` and `log_std`) using a single dense
+  output layer (zero hidden layers) initialized orthogonally. The Critic functions similarly, mapping the hidden
+  representation to a single scalar value.
+
+Note: For the continuous action distributions outputted by the Motor, we explicitly use `mean` and `log_std` as advised
+by previous research to maintain learning stability.
+
+**References**
+
+- Ha, D. (2017, October 29). A Visual Guide to Evolution Strategies. 大トロ ・ Machine Learning. [https://blog.otoro.net/2017/10/29/visual-evolution-strategies/](https://blog.otoro.net/2017/10/29/visual-evolution-strategies/)
+- Schulman, John, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. ‘Proximal Policy Optimization Algorithms’. arXiv:1707.06347. Preprint, arXiv, 28 August 2017. [https://doi.org/10.48550/arXiv.1707.06347](https://doi.org/10.48550/arXiv.1707.06347).
+- Wang, Tingwu, Renjie Liao, Jimmy Ba, and S. Fidler. ‘NerveNet: Learning Structured Policy with Graph Neural Networks’. Conference paper presented at International Conference on Learning Representations. 15 February 2018. [https://www.semanticscholar.org/paper/NerveNet:-Learning-Structured-Policy-with-Graph-Wang-Liao/249408527106d7595d45dd761dd53c83e5a02613](https://www.semanticscholar.org/paper/NerveNet:-Learning-Structured-Policy-with-Graph-Wang-Liao/249408527106d7595d45dd761dd53c83e5a02613).

docs/design/communication.md

@@ -0,0 +1,44 @@
+# Communication scheme (Message Passing)
+
+Remember our research question:
+
+> "What is the impact of different levels of controller modularity on learning speed, coordination, and fault tolerance
+> (e.g. amputations) in brittle-star-like robots trained with Reinforcement Learning?"
+
+To test decentralized modularity (such as arm-level or segment-level controllers), the various modules *must* be able to
+communicate with each other to achieve coordinated locomotion. This is accomplished through *message passing* in a Graph
+Neural Network (GNN)-like architecture. Two prominent communication styles from the literature are N-step NerveNet (Wang
+et al., 2018) and bottom-up top-down Shared Modular Policies (Huang et al., 2020).
+
+We have chosen to apply **one uniform communication style** across all modular architectures, specifically opting for
+**N-step NerveNet**.
+
+## Rationale
+
+Initially, our idea was to equip arm-level controllers with NerveNet message passing and segment-level controllers with
+SMP. However, we evaluated that this introduces a threat to the validity of our research question. If we observe
+differences in performance, it would be impossible to determine whether the variance is caused by the *level of
+modularity*, or by the difference in the message passing scheme. To purely compare modularity, the communication scheme
+style must remain constant.
+
+Second, we decided that NerveNet is a better fit for our research. The morphology of our brittle star contains cycles at
+the decentralized level (e.g., a ring of segments or arms around the body). NerveNet has proven to be robust for
+arbitrary structures, including graphs with cycles. SMP inherently expects a tree structure for its bottom-up and
+top-down pass. Applying SMP to a ring structure requires a workaround to break that cycle.
+
+## Limitations and alternatives
+
+Choosing NerveNet introduces a scalability issue as the morphology grows. In NerveNet, a message advances only one
+segment or node per propagation step. When dealing with long arms (e.g., > 5 segments), this requires a large number of
+propagation steps to transmit information from one tip of an arm to another.
+
+If we were to use SMP instead - which is possible - the inner states of nodes are shared across the entire graph in just
+two passes. For very large or long morphologies, this would be much more scalable.
+
+By rejecting SMP, we accept that our model might learn slower or require more computational power for highly segmented,
+extended morphologies.
+
+**References**
+
+- Wang, Tingwu, Renjie Liao, Jimmy Ba, and S. Fidler. ‘NerveNet: Learning Structured Policy with Graph Neural Networks’. Conference paper presented at International Conference on Learning Representations. 15 February 2018. [https://www.semanticscholar.org/paper/NerveNet:-Learning-Structured-Policy-with-Graph-Wang-Liao/249408527106d7595d45dd761dd53c83e5a02613](https://www.semanticscholar.org/paper/NerveNet:-Learning-Structured-Policy-with-Graph-Wang-Liao/249408527106d7595d45dd761dd53c83e5a02613).
+- Huang, Wenlong, Igor Mordatch, and Deepak Pathak. ‘One Policy to Control Them All: Shared Modular Policies for Agent-Agnostic Control’. arXiv:2007.04976. Preprint, arXiv, 9 July 2020. [https://doi.org/10.48550/arXiv.2007.04976](https://doi.org/10.48550/arXiv.2007.04976).

docs/design/controllers.md

@@ -0,0 +1,28 @@
+# Levels of modularity and topology
+
+The brittle star can be controlled at different levels. A monolithic controller processes all inputs and outputs at
+once, whereas modular controllers divide the brains across the body, inspired by the biology of brittle stars.
+
+We define four architectures to compare:
+
+1. **Centralized, monolithic**: A single Multi Layer Perceptron per robot that receives all observations and outputs
+   all actions.
+2. **Fully connected arm-level**: Each arm contains an MLP that processes the inputs for that arm, an MLP that processes
+   the communicated inner-states, and an MLP that outputs the actions for that arm. One policy for these MLPs is shared
+   across the arms. The controllers in each arm are connected to each other and form a fully connected graph. There is
+   no central disk, but the controllers are fully connected.
+3. **Ring arm-level**: Identical setup to the fully connected arm-level, but the controllers are connected in a ring
+   structure. This setup is considered less centralized than the fully connected graph.
+4. **Segment-level**: Each segment contains the three MLPs discussed above. The base segments, attached to the body,
+   form a ring structure, with the remaining segments attached as extended "strings". Segments can only communicate with
+   segments that are physically connected to it.
+
+## Rationale
+
+To fairly compare decentralized modularity against centralized control, the decentralized models should not be allowed
+to contain a central organ acting as a bottleneck or coordinator. By removing the central disk in the decentralized
+models and replacing it with a ring topology, we closely approximate the biological reality of the brittle star and test
+a decentralized morphology.
+
+The fully connected graph functions as an intermediate step in between a fully centralized and a decentralized ring. We
+use it to test whether our models scale to more complex structures.

docs/design/input_action_spaces.md

@@ -0,0 +1,164 @@
+# Input (state) and output (action) spaces
+
+To effectively learn locomotion and navigation, the agent requires a well-defined observation space (inputs) and action
+space (outputs). The control models map these observations directly to physical movements.
+
+**Inputs (state space)**
+
+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.
+
+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)**
+
+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?*
+
+- Inclusion of Joint Velocities: Because our control models do not inherently possess memory of previous timesteps,
+  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.
+- 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
+
+Alternative state and action formulations include:
+
+- Torque-based continuous control: In many continuous control tasks (like standard MuJoCo benchmarks), actions
+  represent continuous torques applied to joints. While this provides more granular, low-level physical control, it
+  heavily complicates training and does not align well with the physical reality of servo-driven hardware.
+- Recurrent Neural Networks (RNNs) / Frame Stacking: Instead of explicitly passing velocities in the state space, the
+  network could infer momentum by observing a history of past states. Using RNNs or frame stacking allows the agent to
+  build an internal memory of movement. However, this significantly increases architectural complexity and training
+  time compared to explicitly providing the velocity data.
+- 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
+```

docs/design/learning_algorithm.md

@@ -0,0 +1,30 @@
+# Reinforcement Learning Algorithm
+
+To control the *continuous* action space (the joints of the robot) based on sensor data, we require a reliable
+Reinforcement Learning (RL) algorithm or optimization strategy.
+
+We have chosen **Proximal Policy Optimization (PPO)** (Schulman et al., 2017).
+
+## Rationale
+
+PPO is an on-policy algorithm known for its stability and robustness (safe training without excessive variance). More
+importantly, it requires relatively little hyperparameter tuning compared to other algorithms. Since NerveNet was
+successfully trained using PPO (Wang et al., 2018), selecting PPO significantly reduces the risk of convergence
+failures.
+
+## Limitations and alternatives
+
+Alternative learning algorithms include:
+
+- **Twin Delayed DDPG (Fujimoto et al., 2018)**: TD3 is a strong off-policy alternative used in the SMP paper (Huang et
+  al., 2020). It is highly sample-efficient and reportedly excels at zero-shot adaptations. However, this approach would
+  be more complex and error-prone than with PPO.
+- **Evolution strategies (ES)**: Evolution strategies are useful for optimizing Central Pattern Generators (CPGs), e.g.
+  CMA-ES, OpenAI-ES. While this method is easier to distribute and parallelize, ES typically scales worse with
+  exceptionally large observation spaces compared to gradient-based RL methods like PPO.
+
+**References**
+
+- Fujimoto, Scott, Herke Hoof, and David Meger. ‘Addressing Function Approximation Error in Actor-Critic Methods’. Proceedings of the 35th International Conference on Machine Learning, 3 July 2018, 1587-96. [https://proceedings.mlr.press/v80/fujimoto18a.html](https://proceedings.mlr.press/v80/fujimoto18a.html).
+- Huang, Wenlong, Igor Mordatch, and Deepak Pathak. ‘One Policy to Control Them All: Shared Modular Policies for Agent-Agnostic Control’. arXiv:2007.04976. Preprint, arXiv, 9 July 2020. [https://doi.org/10.48550/arXiv.2007.04976](https://doi.org/10.48550/arXiv.2007.04976).
+- Schulman, John, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. ‘Proximal Policy Optimization Algorithms’. arXiv:1707.06347. Preprint, arXiv, 28 August 2017. [https://doi.org/10.48550/arXiv.1707.06347](https://doi.org/10.48550/arXiv.1707.06347).

docs/design/reward_function.md

@@ -0,0 +1,32 @@
+# Reward function and observation space
+
+The robot needs to know whether its movements contribute to the ultimate goal of locomotion towards a target. Sensor
+inputs must be distributed fairly to guarantee an objective comparison between different architectures.
+
+- The distance from the robot to the target and/or the light intensity are treated as global inputs.
+- Positions and joints, which are normalized to floating-point values between 0 and 1, are considered local inputs.
+- The reward function is centered around minimizing the distance to the goal or maximizing the movement towards the goal
+  within a finite number of timesteps $T$.
+- To motivate efficient movement, the amount of timesteps taken to reach the goal will be used as penalty.
+- An extra penalty based on movement relative to the current step and the previous is used to penalize a movement away from the target.
+
+## From reward to PPO
+
+The resulting reward is passed to our PPO library. Our critic network (value function) predicts how good our eventual reward will be for the current state, this value is combined with the reward from the reward function to get advantages. These advantages are then used to calculate the losses to update both our critic and actor pipeline.
+
+## Rationale
+
+Using a light source (or a gradient) is biologically plausible for many simple organisms. By normalizing all signals
+between 0 and 1, PPO training is highly stabilized. The timesteps must be finite to reset the environment in a timely
+manner if the policy gets stuck in a local minimum.
+
+## Limitations and alternatives
+
+Providing global information to all individual decentralized segments can be considered biologically cheating or
+practically infeasible once the robot would be physically built. Some sensory input cannot be put in each joint, for
+example.
+
+The alternative is to provide the global input to the outermost segments of the arms, or a specific set of segments
+assigned with this functionality. The network would then have to learn to propagate this signal throughout the body via
+message passing. While biologically more accurate, this drastically complicates the learning process. We have written
+this down as potential future research.

docs/javascripts/katex.js

@@ -0,0 +1,16 @@
+const renderMath = (el) => {
+  renderMathInElement(el, {
+    delimiters: [
+      { left: "$$",  right: "$$",  display: true },
+      { left: "$",   right: "$",   display: false },
+      { left: "\\(", right: "\\)", display: false },
+      { left: "\\[", right: "\\]", display: true }
+    ],
+  });
+};
+
+if (typeof document$ !== "undefined") {
+  document$.subscribe(({ body }) => renderMath(body));
+} else {
+  document.addEventListener("DOMContentLoaded", () => renderMath(document.body));
+}

env/hpc/modules.txt

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

env/hpc/requirements.txt

@@ -0,0 +1,20 @@
+biorobot==0.4.2
+cleanrl>=0.4.8
+evosax==0.2.0
+flax>=0.12.2
+gymnasium>=1.2.3
+ipykernel==7.2.0
+jax[cuda13]==0.9.0.1
+numpy>=2.0.0
+protobuf>=5.0.0
+warp-lang
+mujoco-warp
+matplotlib==3.10.8
+mediapy==1.2.6
+optax>=0.2.6
+pyopengl>=3.1.10
+pyopengl-accelerate>=3.1.10
+pyyaml>=6.0
+hydra-core>=1.3.2
+wandb==0.24.2
+torch>=2.4.0

flake.nix

@@ -29,6 +29,10 @@
         # Editor of your choice
         (nix-jetbrains-plugins.lib.buildIdeWithPlugins pkgs "pycharm" pluginList)
       ];
+
+      shellHook = ''
+        uv run pre-commit install
+      '';
     };
   });
 }

mkdocs.yml

@@ -0,0 +1,41 @@
+site_name: Brittle Star Project
+theme:
+  name: material
+
+nav:
+  - Home: README.md
+  - Design & Architecture:
+      - Actor-Critic Architecture: design/actor-critic.md
+      - Communication Scheme: design/communication.md
+      - Modularity & Topology: design/controllers.md
+      - Input & Action Spaces: design/input_action_spaces.md
+      - Reinforcement Learning Algorithm: design/learning_algorithm.md
+      - Reward Function & Observation Space: design/reward_function.md
+  - API Reference:
+      - Brittle Star Environment: api/environment.md
+      - Training Models: api/training.md
+      - Tracking & Monitoring: api/tracking.md
+      - Checkpoint & Model Evaluation: api/evaluation.md
+      - Interactive Simulation & Visualization: api/simulation.md
+      - Analysis & Plotting Tools: api/analysis.md
+      - Results & Reproduction: api/reproduction.md
+  - HPC Guide: HPC.md
+  - Contribution Guidelines: CONTRIBUTING.md
+  - Development Guide: DEVELOPMENT.md
+
+markdown_extensions:
+  - pymdownx.superfences:
+      custom_fences:
+        - name: mermaid
+          class: mermaid
+          format: !!python/name:pymdownx.superfences.fence_code_format
+  - pymdownx.arithmatex:
+      generic: true
+
+extra_css:
+  - https://unpkg.com/katex@0/dist/katex.min.css
+
+extra_javascript:
+  - javascripts/katex.js
+  - https://unpkg.com/katex@0/dist/katex.min.js
+  - https://unpkg.com/katex@0/dist/contrib/auto-render.min.js

pyproject.toml

@@ -1,22 +1,84 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
 [project]
 name = "2026sel3-project"
 version = "0.1.0"
 description = "Add your description here"
 readme = "README.md"
-requires-python = ">=3.12, <3.13"
+requires-python = ">= 3.12, < 3.13"
 dependencies = [
     "biorobot==0.4.2",
+    "cleanrl>=0.4.8",
     "evosax==0.2.0",
+    "flax>=0.12.2",
+    "gymnasium>=1.2.3",
     "ipykernel==7.2.0",
-    "jax[cuda13]==0.9.0.1",
+    "jax==0.9.0.1",
+    "numpy>=2.0.0",
+    "protobuf>=5.0.0",
+    "warp-lang",
+    "mujoco-warp",
     "matplotlib==3.10.8",
     "mediapy==1.2.6",
+    "optax>=0.2.6",
     "pyopengl>=3.1.10",
     "pyopengl-accelerate>=3.1.10",
+    "pyyaml>=6.0",
+    "hydra-core>=1.3.2",
     "wandb==0.24.2",
+    "torch>=2.4.0",
+]
+
+[project.optional-dependencies]
+cuda = [
+    "jax[cuda13]==0.9.0.1",
+]
+analysis = [
+    "tensorboard",
+]
+evaluation = [
+    "imageio>=2.35.0",
+    "imageio-ffmpeg>=0.5.1",
 ]
 
 [dependency-groups]
 dev = [
+    "pre-commit>=4.0.0",
+    "pytest>=8.0.0",
     "ruff>=0.15.2",
 ]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/brittle_star_project", "src/experiment_logger"]
+
+[tool.mypy]
+mypy_path = "src"
+check_untyped_defs = false
+warn_return_any = false
+
+[[tool.mypy.overrides]]
+module = [
+    "jax.*",
+    "flax.*",
+    "wandb.*",
+    "torch.*",
+    "mujoco.*",
+    "mujoco_warp.*",
+    "optax.*",
+    "tyro.*",
+    "biorobot.*",
+    "gymnasium.*",
+    "matplotlib.*",
+    "mediapy.*",
+    "matplotlib.*",
+    "mediapy.*",
+    "pytest.*",
+    "tensorboard.*",
+    "tqdm.*",
+    "numpy.*",
+    "yaml.*",
+    "moojoco.*"
+]
+ignore_missing_imports = true

ruff.toml

@@ -1,3 +1,7 @@
+line-length = 100
+exclude = ["wandb"]
+
+[lint]
 extend-select = [
   # "PLC0103",  # invalid-name
   # "PLC0104",  # disallowed-name
@@ -40,10 +44,10 @@ extend-select = [
   "PLC2401",  # non-ascii-name
   # "PLC2403",  # non-ascii-module-import
   # "PLC2503",  # bad-file-encoding
-  "PLC2801",  # unnecessary-dunder-call
+  # "PLC2801",  # unnecessary-dunder-call (requires preview)
   # "PLC3001",  # unnecessary-lambda-assignment
   "PLC3002",  # unnecessary-direct-lambda-call
-  "E999",  # syntax-error
+  # "E999",  # syntax-error (removed from ruff)
   # "PLE0011",  # unrecognized-inline-option
   # "PLE0013",  # bad-plugin-value
   # "PLE0014",  # bad-configuration-section
@@ -130,7 +134,7 @@ extend-select = [
   # "PLE1137",  # unsupported-assignment-operation
   # "PLE1138",  # unsupported-delete-operation
   # "PLE1139",  # invalid-metaclass
-  "PLE1141",  # dict-iter-missing-items
+  # "PLE1141",  # dict-iter-missing-items (requires preview)
   "PLE1142",  # await-outside-async
   # "PLE1143",  # unhashable-member
   # "PLE1144",  # invalid-slice-step
@@ -168,7 +172,7 @@ extend-select = [
   # "PLE3102",  # positional-only-arguments-expected
   # "PLE3701",  # invalid-field-call
   # "PLE4702",  # modified-iterating-dict
-  "PLE4703",  # modified-iterating-set
+  # "PLE4703",  # modified-iterating-set (requires preview)
   # "PLF0001",  # fatal
   # "PLF0002",  # astroid-error
   # "PLF0010",  # parse-error
@@ -210,7 +214,7 @@ extend-select = [
   # "PLW0238",  # unused-private-member
   # "PLW0239",  # overridden-final-method
   # "PLW0240",  # subclassed-final-class
-  "PLW0244",  # redefined-slots-in-subclass
+  # "PLW0244",  # redefined-slots-in-subclass (requires preview)
   "PLW0245",  # super-without-brackets
   # "PLW0246",  # useless-parent-delegation
   # "PLW0301",  # unnecessary-semicolon
@@ -274,7 +278,7 @@ extend-select = [
   "PLW1508",  # invalid-envvar-default
   "PLW1509",  # subprocess-popen-preexec-fn
   # "PLW1510",  # subprocess-run-check
-  "PLW1514",  # unspecified-encoding
+  # "PLW1514",  # unspecified-encoding (requires preview)
   # "PLW1515",  # forgotten-debug-statement
   # "PLW1518",  # method-cache-max-size-none
   "PLW2101",  # useless-with-lock
@@ -298,7 +302,7 @@ extend-select = [
   # "PLW4906",  # deprecated-attribute
 ]
 
-ignore = [
+extend-ignore = [
   # "PLC0116",  # missing-function-docstring
   # "PLC0200",  # consider-using-enumerate
   # "PLC0305",  # trailing-newlines
@@ -346,7 +350,7 @@ ignore = [
   # "PLR1705",  # no-else-return
   # "PLR1706",  # consider-using-ternary
   # "PLR1707",  # trailing-comma-tuple
-  "PLR1708",  # stop-iteration-return
+  # "PLR1708",  # stop-iteration-return (deprecated)
   # "PLR1709",  # simplify-boolean-expression
   # "PLR1710",  # inconsistent-return-statements
   "PLR1711",  # useless-return
@@ -391,4 +395,5 @@ ignore = [
   # "PLW1404",  # implicit-str-concat
 ]
 
-
+[lint.per-file-ignores]
+"__init__.py" = ["F401"]

scripts/analysis/README.md

@@ -0,0 +1,27 @@
+# Experiment Analysis Tools
+
+This directory contains scripts for post-processing and analyzing experiment results, including TensorBoard logs and saved model weights.
+
+## Scripts
+
+### 1. `explore_tensorboard.py`
+A CLI tool to summarize TensorBoard `tfevents` files without a GUI.
+
+**Key Features:**
+- Displays last values, min, max, and step counts for all scalar metrics.
+- Calculates total run duration and estimated completion percentage.
+- Exports granular scalar data to CSV for analysis in Excel/Pandas.
+
+**Usage:**
+```bash
+# General usage
+python explore_tensorboard.py <run_directory>
+
+# Exporting data
+python explore_tensorboard.py <run_directory> --csv data.csv
+```
+
+**Requirements:**
+- `pandas`
+- `tensorboard`
+- `tensorflow-cpu` (or `tensorflow`)

scripts/analysis/explore_tensorboard.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python3
+"""
+Reproducible CLI tool to explore TensorBoard logs.
+Designed for both local development and HPC diagnostics.
+
+Requirements:
+    pip install tensorboard
+
+Usage:
+    python explore_tensorboard.py <path_to_run_directory> [--csv output.csv]
+"""
+
+import argparse
+import os
+import sys
+import csv
+
+try:
+    from tensorboard.backend.event_processing import event_accumulator
+except ImportError:
+    print("Error: Missing dependency. Please run: pip install tensorboard")
+    sys.exit(1)
+
+
+def explore_run(log_dir):
+    """
+    Extracts and displays a summary of scalar metrics from a TensorBoard log directory.
+    """
+    print(f"\n{'=' * 20} Exploring Run {'=' * 20}")
+    print(f"Directory: {log_dir}")
+    print(f"{'=' * 55}\n")
+
+    if not os.path.exists(log_dir):
+        print(f"Error: Directory '{log_dir}' does not exist.")
+        return None
+
+    # Initialize EventAccumulator
+    # size_guidance=0 loads all data points for each tag.
+    ea = event_accumulator.EventAccumulator(
+        log_dir,
+        size_guidance={
+            event_accumulator.SCALARS: 0,
+            event_accumulator.TENSORS: 0,
+        },
+    )
+
+    print("Loading event files (this may take a moment for large runs)...")
+    ea.Reload()
+
+    tags = ea.Tags()
+    scalar_tags = tags.get("scalars", [])
+
+    if not scalar_tags:
+        print("No scalar metrics found in this directory.")
+        return None
+
+    print(f"Found {len(scalar_tags)} scalar metrics.\n")
+
+    data = {}
+    summary = []
+
+    # Process scalar values
+    for tag in scalar_tags:
+        events = ea.Scalars(tag)
+        if not events:
+            continue
+
+        values = [e.value for e in events]
+        last_event = events[-1]
+        data[tag] = values
+
+        summary.append(
+            {
+                "Metric": tag,
+                "Steps": len(events),
+                "Last Value": f"{last_event.value:.4f}",
+                "Max": f"{max(values):.4f}",
+                "Min": f"{min(values):.4f}",
+            }
+        )
+
+    # Display summary table formatted manually
+    summary = sorted(summary, key=lambda x: x["Metric"])
+    print(f"{'Metric':<30} {'Steps':>10} {'Last':>12} {'Max':>12} {'Min':>12}")
+    print("-" * 80)
+    for row in summary:
+        print(
+            f"{row['Metric']:<30} {row['Steps']:>10} {row['Last Value']:>12} "
+            f"{row['Max']:>12} {row['Min']:>12}"
+        )
+
+    # Calculate and display global metadata
+    if "charts/SPS" in data:
+        sps_events = ea.Scalars("charts/SPS")
+        if len(sps_events) > 1:
+            total_duration_hours = (sps_events[-1].wall_time - sps_events[0].wall_time) / 3600
+            print(f"\nTotal Recorded Duration: {total_duration_hours:.2f} hours")
+
+            # Estimate completion if total_timesteps is available in hyperparameters
+            try:
+                hp_tags = [t for t in tags.get("tensors", []) if "hyperparameters" in t]
+                if hp_tags:
+                    hp_event = ea.Tensors(hp_tags[0])[0]
+                    hp_text = hp_event.tensor_proto.string_val[0].decode("utf-8")
+                    if "total_timesteps" in hp_text:
+                        for line in hp_text.split("\n"):
+                            if "total_timesteps" in line:
+                                target = int(line.split("|")[2].strip())
+                                current = ea.Scalars(scalar_tags[0])[-1].step
+                                percent = (current / target) * 100
+                                print(f"Progress: {current:,} / {target:,} steps ({percent:.1f}%)")
+            except Exception:
+                pass
+
+    return data
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Reproducible TensorBoard exploration tool.")
+    parser.add_argument("log_dir", help="Path to the TensorBoard run directory.")
+    parser.add_argument("--csv", help="Optional: Path to export scalar data to CSV.", default=None)
+
+    args = parser.parse_args()
+
+    scalar_data = explore_run(args.log_dir)
+
+    if args.csv and scalar_data:
+        # Reloading for wall_time and steps
+        ea = event_accumulator.EventAccumulator(args.log_dir).Reload()
+        with open(args.csv, mode="w", newline="") as f:
+            writer = csv.DictWriter(f, fieldnames=["tag", "step", "value", "wall_time"])
+            writer.writeheader()
+            for tag in scalar_data.keys():
+                for e in ea.Scalars(tag):
+                    writer.writerow(
+                        {"tag": tag, "step": e.step, "value": e.value, "wall_time": e.wall_time}
+                    )
+
+        print(f"\nData exported to: {args.csv}")
+
+
+if __name__ == "__main__":
+    main()

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/hpc/export_requirements.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+"""Export HPC pip requirements from pyproject.toml.
+
+This is a LOCAL DEVELOPER UTILITY — run it on your own machine before pushing
+code whenever pyproject.toml dependencies change. It reads the modules from
+env/hpc/modules.txt and the full dependency list from pyproject.toml, then
+writes the remainder to env/hpc/requirements.txt.
+"""
+
+from __future__ import annotations
+
+import re
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[2]
+
+
+def normalise(name: str) -> str:
+    """Normalise a PyPI package name for comparison."""
+    return re.sub(r"[-_.]+", "-", name).lower()
+
+
+def pkg_name(dep: str) -> str:
+    """Extract the bare package name from a PEP 508 dependency string."""
+    return re.split(r"[\[=><~!;]", dep)[0].strip()
+
+
+def main() -> None:
+    import tomllib
+
+    modules_path = ROOT / "env" / "hpc" / "modules.txt"
+    if not modules_path.exists():
+        print(f"Error: {modules_path} not found.", file=sys.stderr)
+        sys.exit(1)
+
+    # Read normalized module names from base modules only
+    # Library modules (like PyTorch) are kept in requirements for portability
+    module_names = [
+        normalise(line.split()[0].split("/")[0])
+        for line in modules_path.read_text().splitlines()
+        if line.strip() and not line.startswith("#")
+    ]
+
+    pyproject_path = ROOT / "pyproject.toml"
+    with pyproject_path.open("rb") as f:
+        data = tomllib.load(f)
+
+    # Collect all dependencies, merging 'cuda' extras into base dependencies
+    dep_dict: dict[str, str] = {}
+    for dep in data.get("project", {}).get("dependencies", []):
+        dep_dict[normalise(pkg_name(dep))] = dep
+
+    # Add cuda extras (takes precedence for HPC)
+    optional_deps = data.get("project", {}).get("optional-dependencies", {})
+    for group in ["cuda"]:
+        for dep in optional_deps.get(group, []):
+            dep_dict[normalise(pkg_name(dep))] = dep
+
+    deps = list(dep_dict.values())
+
+    final_deps: list[str] = []
+    print("Checking dependencies against HPC module list...", file=sys.stderr)
+    for dep in deps:
+        name = normalise(pkg_name(dep))
+        # Smart check: if the package name is a substring of any loaded module name
+        # (e.g. 'torch' in 'pytorch', 'scipy' in 'scipy-bundle')
+        if any(name in mod for mod in module_names):
+            print(f"  [skip – module provider found] {dep}", file=sys.stderr)
+            continue
+
+        final_deps.append(dep)
+        print(f"  [pip] {dep}", file=sys.stderr)
+
+    hpc_dir = ROOT / "env" / "hpc"
+    output_path = hpc_dir / "requirements.txt"
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    output_path.write_text("\n".join(final_deps) + "\n")
+    print(f"\nWrote {len(final_deps)} requirement(s) to {output_path}", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    main()

scripts/hpc/install.sh

@@ -0,0 +1,54 @@
+#!/bin/bash -l
+# scripts/hpc/install.sh
+#
+# Usage (on any compute node):
+#   bash scripts/hpc/install.sh
+#
+# Batch usage:
+#   qsub scripts/hpc/install.sh
+
+#PBS -N brittlestar-install
+#PBS -l walltime=00:15:00
+
+set -euo pipefail
+
+# Preliminary status echo
+echo ">>> Starting installation job $PBS_JOBID on $(hostname)..."
+
+if [ -n "$PBS_O_WORKDIR" ]; then
+    cd "$PBS_O_WORKDIR"
+fi
+
+mkdir -p "${PBS_O_WORKDIR}/runs"
+
+# Mirror configs to $VSC_DATA to avoid home quota limits (3GB)
+# vsc-venv manages environments relative to the requirements file
+PROJ_NAME=$(basename "$PWD")
+HPC_CONFIG_DIR="$VSC_DATA/$PROJ_NAME/env/hpc"
+mkdir -p "$HPC_CONFIG_DIR"
+cp env/hpc/*.txt "$HPC_CONFIG_DIR/"
+
+# Keep caches off $VSC_HOME (quota ~3 GB).
+export PIP_CACHE_DIR="$VSC_SCRATCH/.cache/pip"
+export UV_CACHE_DIR="$VSC_SCRATCH/.cache/uv"
+mkdir -p "$PIP_CACHE_DIR" "$UV_CACHE_DIR"
+
+module load vsc-venv
+
+echo ">>> Synchronizing and activating environment (vsc-venv)..."
+# cd to $VSC_DATA so vsc-venv creates its venvs/ directory there, not in $HOME.
+mkdir -p "$VSC_DATA/$PROJ_NAME"
+cd "$VSC_DATA/$PROJ_NAME"
+set +euo pipefail
+source vsc-venv --activate \
+    --modules "$HPC_CONFIG_DIR/modules.txt" \
+    --requirements "$HPC_CONFIG_DIR/requirements.txt"
+set -euo pipefail
+cd "$PBS_O_WORKDIR"
+
+echo '>>> Installing ipykernel...'
+CLUSTER_ID="${VSC_INSTITUTE_CLUSTER:-generic}"
+python -m ipykernel install --user --name="sel3_${CLUSTER_ID}" \
+    --display-name "SEL3 (${CLUSTER_ID})"
+
+echo '>>> Done'

scripts/hpc/train.pbs

@@ -0,0 +1,75 @@
+# Production training (requires GPU at runtime):
+#   qsub -l gpus=1 scripts/hpc/train.pbs
+# Debug/CPU training:
+#   qsub scripts/hpc/train.pbs
+
+#PBS -N brittlestar-ppo
+#PBS -l nodes=1:ppn=8
+#PBS -l walltime=24:00:00
+#PBS -o runs/brittlestar-ppo.o$PBS_JOBID
+#PBS -e runs/brittlestar-ppo.e$PBS_JOBID
+
+set -euo pipefail
+
+# Preliminary status echo
+echo ">>> Starting training job $PBS_JOBID on $(hostname)..."
+
+if [ -n "$PBS_O_WORKDIR" ]; then
+    cd "$PBS_O_WORKDIR"
+fi
+
+# Set up storage paths dynamically
+PROJ_NAME=$(basename "$PWD")
+RUN_ID="brittlestar_${PBS_JOBID}"
+SCRATCH_RUNDIR="$VSC_SCRATCH/runs/$RUN_ID"
+DATA_RUNDIR="$VSC_DATA/runs/$RUN_ID"
+mkdir -p "$SCRATCH_RUNDIR" "$DATA_RUNDIR" runs/
+
+# Keep caches off $VSC_HOME (quota ~3 GB).
+export PIP_CACHE_DIR="$VSC_SCRATCH/.cache/pip"
+export UV_CACHE_DIR="$VSC_SCRATCH/.cache/uv"
+mkdir -p "$PIP_CACHE_DIR" "$UV_CACHE_DIR"
+
+module load vsc-venv
+
+echo ">>> Synchronizing and activating environment (vsc-venv)..."
+HPC_CONFIG_DIR="$VSC_DATA/$PROJ_NAME/env/hpc"
+if [ ! -d "$HPC_CONFIG_DIR" ]; then
+    echo "ERROR: HPC_CONFIG_DIR ($HPC_CONFIG_DIR) does not exist. Run install.sh first."
+    exit 1
+fi
+
+# cd to $VSC_DATA so vsc-venv finds its venvs/ directory there, not in $HOME.
+cd "$VSC_DATA/$PROJ_NAME"
+set +euo pipefail
+source vsc-venv --activate \
+    --modules "$HPC_CONFIG_DIR/modules.txt" \
+    --requirements "$HPC_CONFIG_DIR/requirements.txt"
+set -euo pipefail
+cd "$PBS_O_WORKDIR"
+
+
+echo ">>> Starting BrittleStar training..."
+export MUJOCO_GL=egl
+export WANDB_DIR="$SCRATCH_RUNDIR"
+
+export PYTHONPATH="$PBS_O_WORKDIR/src:${PYTHONPATH:-}"
+
+if [ -f "$VSC_DATA/$PROJ_NAME/.env" ]; then
+    echo ">>> Sourcing API keys from .env..."
+    export $(grep -v '^#' "$VSC_DATA/$PROJ_NAME/.env" | xargs)
+elif [ -f "$PBS_O_WORKDIR/.env" ]; then
+    echo ">>> Sourcing API keys from .env..."
+    export $(grep -v '^#' "$PBS_O_WORKDIR/.env" | xargs)
+fi
+
+# Run training using Hydra overrides
+python scripts/train.py \
+    hydra.run.dir="$SCRATCH_RUNDIR" \
+    ppo=stable \
+    logging=hpc
+
+echo ">>> Staging out results to $DATA_RUNDIR..."
+cp -r "$SCRATCH_RUNDIR/." "$DATA_RUNDIR/"
+
+echo ">>> Done"

scripts/plots/analyze_comparisons.py

@@ -0,0 +1,418 @@
+"""
+Poster Comparison Visualizations
+
+This script generates a Forward Velocity plot and three secondary plots (Accumulated Reward, Success
+Rate, Distance Remaining).
+"""
+
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+from plot_config import (
+    BEST_PERFORMER_COLOR,
+    BEST_PERFORMER_TEXT,
+    COLORS,
+    LEGEND_KWARGS,
+    apply_style,
+    create_common_parser,
+)
+
+
+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
+    group_spacing = 1.3
+    x_indices = np.arange(len(morphologies)) * group_spacing
+    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, _ 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  # center the label in the 3 bars
+        + (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")  # 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}")
+        )
+
+    ax.legend(**LEGEND_KWARGS, ncol=len(architectures))
+    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}")
+        )
+
+    ax.legend(**LEGEND_KWARGS, ncol=len(morphologies))
+    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="",
+            title="Maximal forward velocity (in cm/s)",
+            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,385 @@
+"""
+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."""
+
+    CHECKPOINT = "checkpoint"
+    ARCH = "architecture"
+    TIMESTEPS = "trained_timesteps"
+    REWARD = "eval_return"
+    VELOCITY = "velocity"
+    EVAL_STEPS = "eval_steps"
+    FINAL_XY_DIST = "final_xy_dist"
+    INITIAL_XY_DIST = "initial_xy_dist"
+    REACHED_TARGET = "reached_target"
+
+
+# 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/final-v2-centralized/checkpoint_evaluation.csv",
+    "decentralized fully connected": "runs/final-v2-fully-conn/checkpoint_evaluation.csv",
+    "decentralized ring-level": "runs/final-v2-ring/checkpoint_evaluation.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.CHECKPOINT,
+        Columns.TIMESTEPS,
+        Columns.REWARD,
+        Columns.INITIAL_XY_DIST,
+        Columns.FINAL_XY_DIST,
+        Columns.EVAL_STEPS,
+    ]
+    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.VELOCITY] = (df[Columns.INITIAL_XY_DIST] - df[Columns.FINAL_XY_DIST]) / df[
+            Columns.EVAL_STEPS
+        ]
+        df[Columns.ARCH] = arch_name
+        df[Columns.VELOCITY] = (df[Columns.INITIAL_XY_DIST] - df[Columns.FINAL_XY_DIST]) / df[
+            Columns.EVAL_STEPS
+        ]
+
+        dfs.append(df)
+
+    return pd.concat(dfs, ignore_index=True) if dfs else pd.DataFrame()
+
+
+def _convergence_timestep(
+    series: pd.Series, timesteps: pd.Series, checkpoints: pd.Series
+) -> tuple[float, int, int]:
+    """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
+
+    mask = smoothed >= threshold
+    first_idx = mask.idxmax()
+
+    return timesteps.loc[first_idx], first_idx, checkpoints.loc[first_idx]
+
+
+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 = []
+
+    centralized_base = 0
+
+    for arch in df[Columns.ARCH].unique():
+        arch_data = df[df[Columns.ARCH] == arch].sort_values(Columns.TIMESTEPS)
+
+        reward_timestep, reward_checkpoint_idx, reward_checkpoint = _convergence_timestep(
+            arch_data[Columns.REWARD],
+            arch_data[Columns.TIMESTEPS],
+            arch_data[Columns.CHECKPOINT],
+        )
+
+        velocity_timestep, velocity_checkpoint_idx, velocity_checkpoint = _convergence_timestep(
+            arch_data[Columns.VELOCITY],
+            arch_data[Columns.TIMESTEPS],
+            arch_data[Columns.CHECKPOINT],
+        )
+
+        results.append(
+            {
+                "Architecture": arch,
+                "Reward_Convergence_Timestep": reward_timestep,
+                "Reward_Convergence_Checkpoint_Idx": reward_checkpoint_idx,
+                "Reward_Convergence_Checkpoint": reward_checkpoint,
+                "Velocity_Convergence_Timestep": velocity_timestep,
+                "Velocity_Convergence_Checkpoint_Idx": velocity_checkpoint_idx,
+                "Velocity_Convergence_Checkpoint": velocity_checkpoint,
+            }
+        )
+
+        if arch == "centralized 5 arms":
+            centralized_base = reward_checkpoint
+        else:
+            print(arch, "speedup:", 1 - reward_checkpoint / centralized_base)
+
+    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)
+    print(
+        results[
+            ["Architecture", "Reward_Convergence_Checkpoint_Idx", "Reward_Convergence_Checkpoint"]
+        ]
+    )
+
+    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,75 @@
+import argparse
+import matplotlib.pyplot as plt
+
+# Shared Color Palette (Colorblind friendly, high contrast)
+# Matches poster design
+COLORS = {
+    "CENTRALIZED": "#0D567C",  # Blue
+    "FULLY_CONNECTED": "#8C0E0F",  # Reddish
+    "RING": "#FCB305",  # Pale Yellow
+}
+
+
+def apply_style(font_size=36):
+    """
+    Applies the shared typography and aesthetic settings to Matplotlib.
+    """
+    plt.rcParams.update(
+        {
+            "font.size": font_size,
+            "axes.labelsize": font_size,
+            "axes.titlesize": font_size,
+            "xtick.labelsize": font_size,
+            "ytick.labelsize": font_size,
+            "legend.fontsize": font_size,
+            "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.12),
+    "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/poster_visualisations/render_poster_videos.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+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.video import record_episode_multi_camera
+
+_ARCH_DIR_MAP = {
+    "CENTRALIZED": "centralized",
+    "FULLY_CONNECTED": "fully-connected",
+    "RING": "ring",
+    "SEGMENT": "segment",
+}
+
+ROBOT_COLOR_MAP = {
+    "CENTRALIZED": "#0D567C",  # Blue
+    "FULLY_CONNECTED": "#8C0E0F",  # Reddish
+    "RING": "#FCB304",  # Pale Yellow
+}
+
+
+def _arch_dir(name: str) -> str:
+    return _ARCH_DIR_MAP.get(name, name.lower())
+
+
+def _resolve_overrides(overrides: list[str], count: int) -> list[str | None]:
+    if not overrides:
+        return [None] * count
+    if len(overrides) == 1 and count > 1:
+        return overrides * count
+    if len(overrides) != count:
+        raise ValueError("morphology overrides must match the number of models")
+    return overrides
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Render top-down and follow videos for poster.")
+    parser.add_argument("models", nargs="+", help="Paths to .flax checkpoints")
+    parser.add_argument(
+        "--morphology-override",
+        action="append",
+        default=[],
+        help="Override morphology YAML path (repeat to match models)",
+    )
+    parser.add_argument("--output-root", default="vids/poster")
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--max-steps", type=int, default=5000)
+    parser.add_argument("--topdown-camera", type=int, default=0)
+    parser.add_argument("--follow-camera", type=int, default=1)
+    parser.add_argument("--topdown-camera-x", type=float, default=-3.0)
+    parser.add_argument("--topdown-camera-y", type=float, default=0.0)
+    parser.add_argument("--topdown-camera-z", type=float, default=4.5)
+    parser.add_argument("--topdown-camera-fovy", type=float, default=None)
+    parser.add_argument("--target-x", type=float, default=-6.0)
+    parser.add_argument("--target-y", type=float, default=0.0)
+    parser.add_argument("--width", type=int, default=2160)
+    parser.add_argument("--height", type=int, default=960)
+    parser.add_argument("--fps", type=int, default=60)
+    parser.add_argument(
+        "--robot-color",
+        default="#2B4162",
+        help="Hex color for the brittle star robot",
+    )
+    args = parser.parse_args()
+
+    if (args.target_x is None) != (args.target_y is None):
+        raise ValueError("target-x and target-y must be provided together")
+
+    target_xy = None
+    if args.target_x is not None:
+        target_xy = (float(args.target_x), float(args.target_y))
+
+    camera_fovy = None
+    if args.topdown_camera_fovy is not None:
+        camera_fovy = {args.topdown_camera: float(args.topdown_camera_fovy)}
+
+    camera_x = None
+    if args.topdown_camera_x is not None:
+        camera_x = {args.topdown_camera: float(args.topdown_camera_x)}
+
+    camera_y = None
+    if args.topdown_camera_y is not None:
+        camera_y = {args.topdown_camera: float(args.topdown_camera_y)}
+
+    camera_z = None
+    if args.topdown_camera_z is not None:
+        camera_z = {args.topdown_camera: float(args.topdown_camera_z)}
+
+    camera_xyz = (camera_x, camera_y, camera_z)
+
+    overrides = _resolve_overrides(args.morphology_override, len(args.models))
+    output_root = Path(args.output_root)
+
+    for model_path_str, override in zip(args.models, overrides):
+        model_path = Path(model_path_str)
+        metadata = load_metadata(model_path, None)
+        training = metadata_to_configs(metadata)
+
+        bundle = build_eval_env(
+            model_path=model_path,
+            training=training,
+            metadata=metadata,
+            morphology_override_path=override,
+        )
+
+        arch_dir = _arch_dir(bundle.architecture)
+        arms_dir = f"{bundle.num_active_arms}arms"
+        out_dir = output_root / arms_dir / arch_dir
+        out_dir.mkdir(parents=True, exist_ok=True)
+
+        output_paths = {
+            args.topdown_camera: out_dir / "topdown.mp4",
+            args.follow_camera: out_dir / "follow.mp4",
+        }
+
+        print(bundle.architecture)
+        color = ROBOT_COLOR_MAP.get(bundle.architecture, args.robot_color)
+
+        result = record_episode_multi_camera(
+            env=bundle.env,
+            policy=bundle.policy,
+            seed=args.seed,
+            max_steps=args.max_steps,
+            action_low=bundle.action_low,
+            action_high=bundle.action_high,
+            action_mask=bundle.action_mask,
+            output_paths=output_paths,
+            camera_ids=[args.topdown_camera, args.follow_camera],
+            camera_fovy=camera_fovy,
+            camera_xyz=camera_xyz,
+            target_xy=target_xy,
+            robot_color=color,
+            width=args.width,
+            height=args.height,
+            fps=args.fps,
+        )
+
+        final_dist = "n/a" if result.final_xy_dist is None else f"{result.final_xy_dist:.3f}"
+        print(
+            f"{arms_dir}/{arch_dir}: return={result.return_:.6f}, len={result.length}, "
+            f"target_reached={result.reached_target}, final_xy_dist={final_dist}"
+        )
+
+        bundle.env.close()
+
+
+if __name__ == "__main__":
+    main()

scripts/poster_visualisations/render_poster_videos.sh

@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+
+# Multi-camera renders per model -> vids/poster/{arms}arms/{arch}/topdown.mp4 + follow.mp4
+
+path=$1
+
+uv run scripts/poster_visualisations/render_poster_videos.py \
+  "$path"/centralized.flax \
+  "$path"/fully-connected.flax \
+  "$path"/ring.flax \

scripts/poster_visualisations/render_static_path_image.py

@@ -0,0 +1,241 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+
+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 (
+    _get_observations,
+    _maybe_clip_action,
+    _target_reached,
+)
+from brittle_star_project.evaluation.video import (
+    _apply_camera_overrides,
+    _ensure_offscreen_size,
+    hex_to_rgba,
+)
+
+ROBOT_COLOR_MAP = {
+    "CENTRALIZED": "#0D567C",  # Blue
+    "FULLY_CONNECTED": "#8C0E0F",  # Reddish
+    "RING": "#FCB304",  # Pale Yellow
+}
+
+
+def _enum_value(enum_obj, *names: str) -> int:
+    for name in names:
+        if hasattr(enum_obj, name):
+            return int(getattr(enum_obj, name))
+    raise AttributeError(f"Could not find any of {names!r} on {enum_obj!r}")
+
+
+def _append_sphere(scene, mujoco, center: np.ndarray, radius: float, rgba: np.ndarray) -> None:
+    geom = scene.geoms[scene.ngeom]
+    mujoco.mjv_initGeom(
+        geom,
+        mujoco.mjtGeom.mjGEOM_SPHERE,
+        np.asarray([radius, 0.0, 0.0], dtype=np.float32),
+        center,
+        np.eye(3, dtype=np.float32).reshape(-1),
+        rgba,
+    )
+    scene.ngeom += 1
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Render a static path image from a rollout.")
+    parser.add_argument("model", help="Path to .flax checkpoint")
+    parser.add_argument("--morphology-override", default=None)
+    parser.add_argument("--output-path", required=True)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--max-steps", type=int, default=5000)
+    parser.add_argument("--body-name", default="BrittleStarMorphology/central_disk")
+    parser.add_argument("--camera-id", type=int, default=0)
+    parser.add_argument("--camera-x", type=float, default=-3.0)
+    parser.add_argument("--camera-y", type=float, default=0.0)
+    parser.add_argument("--camera-z", type=float, default=4.5)
+    parser.add_argument("--camera-fovy", type=float, default=None)
+    parser.add_argument("--target-x", type=float, default=-6.0)
+    parser.add_argument("--target-y", type=float, default=0.0)
+    parser.add_argument("--width", type=int, default=2160)
+    parser.add_argument("--height", type=int, default=960)
+    parser.add_argument("--frame-stride", type=int, default=15)
+    parser.add_argument(
+        "--path-color", default="#FA9F42"
+    )  # ring = #888888, centralized = #2B4162, fully connected = FA9F42
+    parser.add_argument(
+        "--robot-color",
+        default="#FA9F42",
+        help="Hex color for brittle star robot (e.g. #ff0000)",
+    )
+    args = parser.parse_args()
+
+    model_path = Path(args.model)
+    metadata = load_metadata(model_path, None)
+    training = metadata_to_configs(metadata)
+
+    bundle = build_eval_env(
+        model_path=model_path,
+        training=training,
+        metadata=metadata,
+        morphology_override_path=args.morphology_override,
+    )
+
+    if (args.target_x is None) != (args.target_y is None):
+        raise ValueError("target-x and target-y must be provided together")
+
+    target_xy = None
+    if args.target_x is not None:
+        target_xy = (float(args.target_x), float(args.target_y))
+
+    try:
+        import imageio
+        import mujoco
+    except ImportError as e:
+        raise ImportError(
+            "Static image rendering requires 'mujoco' and 'imageio'. "
+            "Please install the evaluation dependencies: `uv pip install .[evaluation]`"
+        ) from e
+
+    reset_kwargs = {}
+    if target_xy is not None:
+        reset_kwargs["target_position"] = (target_xy[0], target_xy[1], 0.0)
+
+    state = bundle.env.reset(seed=args.seed, **reset_kwargs)
+    model = state.mj_model
+    data = state.mj_data
+
+    _apply_camera_overrides(
+        model,
+        camera_fovy={args.camera_id: float(args.camera_fovy)}
+        if args.camera_fovy is not None
+        else None,
+        camera_xyz=(
+            {args.camera_id: float(args.camera_x)} if args.camera_x is not None else None,
+            {args.camera_id: float(args.camera_y)} if args.camera_y is not None else None,
+            {args.camera_id: float(args.camera_z)} if args.camera_z is not None else None,
+        ),
+    )
+    _ensure_offscreen_size(model, args.width, args.height)
+
+    body_id = mujoco.mj_name2id(
+        model, mujoco.mjtObj.mjOBJ_BODY, "BrittleStarMorphology/central_disk"
+    )
+
+    robot_rgba = hex_to_rgba(ROBOT_COLOR_MAP.get(bundle.architecture, args.robot_color), 1.0)
+
+    # Optionally override robot color by recoloring geoms belonging to the robot's body subtree.
+    if args.robot_color is not None:
+        # Collect body IDs in the subtree rooted at `body_id` by walking parent links.
+        nbody = int(model.nbody)
+        body_parent = model.body_parentid
+        robot_body_ids = set([int(body_id)])
+        for i in range(1, nbody):
+            cur = int(i)
+            # walk up until root (0) or until we hit the robot root
+            while cur not in (-1, 0, int(body_id)):
+                cur = int(body_parent[cur])
+            if cur == int(body_id):
+                robot_body_ids.add(i)
+
+        # Recolor geoms whose body id is in the robot subtree
+        for g in range(int(model.ngeom)):
+            if int(model.geom_bodyid[g]) in robot_body_ids:
+                model.geom_rgba[g][:] = robot_rgba
+
+    # Optionally override robot color by recoloring geoms belonging to the robot's body subtree.
+    if args.robot_color is not None:
+        # Collect body IDs in the subtree rooted at `body_id` by walking parent links.
+        nbody = int(model.nbody)
+        body_parent = model.body_parentid
+        robot_body_ids = set([int(body_id)])
+        for i in range(1, nbody):
+            cur = int(i)
+            # walk up until root (0) or until we hit the robot root
+            while cur not in (-1, 0, int(body_id)):
+                cur = int(body_parent[cur])
+            if cur == int(body_id):
+                robot_body_ids.add(i)
+
+        # Recolor geoms whose body id is in the robot subtree
+        for g in range(int(model.ngeom)):
+            if int(model.geom_bodyid[g]) in robot_body_ids:
+                model.geom_rgba[g][:] = robot_rgba
+
+    positions = []
+    observations = _get_observations(state)
+
+    for _ in range(int(args.max_steps)):
+        positions.append(np.asarray(data.xpos[body_id], dtype=np.float32))
+
+        obs_dict = observations or {}
+        action = bundle.policy.act(observations=obs_dict)
+        if bundle.action_mask is not None:
+            action = action[bundle.action_mask]
+        action = _maybe_clip_action(action, bundle.action_low, bundle.action_high)
+
+        state = bundle.env.step(state=state, action=action)
+        data = state.mj_data
+        observations = _get_observations(state)
+
+        if _target_reached(state=state):
+            break
+
+    positions_arr = np.vstack(positions)
+    if len(positions_arr) < 2:
+        raise ValueError("Need at least two rollout positions to render a path")
+
+    path_points = positions_arr.copy()
+    path_points[:, 2] -= 0.02
+
+    path_step = max(1, int(args.frame_stride))
+    path_points_visible = path_points[::path_step]
+    path_rgba = hex_to_rgba(ROBOT_COLOR_MAP.get(bundle.architecture, args.path_color), 0.92)
+
+    ctx = mujoco.GLContext(args.width, args.height)
+    ctx.make_current()
+    try:
+        catmask = _enum_value(mujoco.mjtCatBit, "mjCAT_ALL")
+        camera_type = _enum_value(mujoco.mjtCamera, "mjCAMERA_FIXED")
+        font_scale = _enum_value(mujoco.mjtFontScale, "mjFONTSCALE_100")
+
+        maxgeom = int(model.ngeom + len(path_points_visible) + 8)
+        scene = mujoco.MjvScene(model, maxgeom=maxgeom)
+        option = mujoco.MjvOption()
+        perturb = mujoco.MjvPerturb()
+        camera = mujoco.MjvCamera()
+        mujoco.mjv_defaultOption(option)
+        mujoco.mjv_defaultPerturb(perturb)
+        mujoco.mjv_defaultCamera(camera)
+        camera.type = camera_type
+        camera.fixedcamid = int(args.camera_id)
+        if hasattr(camera, "trackbodyid"):
+            camera.trackbodyid = -1
+
+        context = mujoco.MjrContext(model, font_scale)
+        viewport = mujoco.MjrRect(0, 0, args.width, args.height)
+
+        mujoco.mjv_updateScene(model, data, option, perturb, camera, catmask, scene)
+
+        for idx, path_point in enumerate(path_points_visible):
+            path_rgba[3] = 0.10 + 0.70 * (idx / max(len(path_points_visible) - 1, 1))
+            _append_sphere(scene, mujoco, path_point, 0.03, path_rgba)
+
+        rgb = np.empty((args.height, args.width, 3), dtype=np.uint8)
+        depth = np.empty((args.height, args.width), dtype=np.float32)
+        mujoco.mjr_render(viewport, scene, context)
+        mujoco.mjr_readPixels(rgb, depth, viewport, context)
+        imageio.imwrite(args.output_path, np.flipud(rgb))
+
+        context.free()
+    finally:
+        ctx.free()
+
+    bundle.env.close()
+
+
+if __name__ == "__main__":
+    main()

scripts/poster_visualisations/render_static_path_image.sh

@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+
+# Static path image + optional ghost render
+
+path=$1 # path to .flax model with metadata.yaml alongside it
+
+uv run scripts/poster_visualisations/render_static_path_image.py \
+  "$path" \
+  --output-path vids/poster/5arms/centralized/path.png \
+  --ghost-overlay

scripts/simulate.py

@@ -0,0 +1,180 @@
+"""Simulate a trained policy in the MuJoCo viewer.
+
+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
+
+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.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,
+)
+
+
+@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
+def main(dict_cfg: DictConfig) -> None:
+    # 1. Hydra composes ONLY SimulationSettings
+    cfg = OmegaConf.to_object(OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg))
+    sim_cfg = cfg.simulation
+
+    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)"
+        )
+
+    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}'.")
+
+    # 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)
+
+    # 8. Run simulation
+    headless = bool(sim_cfg.headless)
+    max_steps = sim_cfg.max_steps
+
+    if sim_cfg.record_video:
+        if max_steps is None:
+            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")
+
+        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,
+            action_mask=action_mask,
+            output_path=output_path,
+            camera_id=sim_cfg.camera_id,
+            width=sim_cfg.video_width,
+            height=sim_cfg.video_height,
+            fps=sim_cfg.video_fps,
+        )
+
+        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={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 = max_steps_i
+
+        model_dt = float(state0.mj_model.opt.timestep)
+        control_dt = model_dt * float(training.environment.num_physics_steps_per_control_step)
+
+        rollout_viewer(
+            env=env,
+            policy=policy,
+            seed=seed,
+            state=state0,
+            control_dt=control_dt,
+            max_steps=max_steps_val,
+            action_low=action_low,
+            action_high=action_high,
+            action_mask=action_mask,
+        )
+
+    env.close()
+
+
+if __name__ == "__main__":
+    register_configs()
+    main()

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/download_wandb_project.py

@@ -0,0 +1,142 @@
+from pathlib import Path
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import threading
+import wandb
+import argparse
+
+# tune these depending on network / W&B limits
+MAX_RUN_WORKERS = 8
+MAX_FILE_WORKERS = 16
+MAX_ARTIFACT_WORKERS = 8
+
+api = wandb.Api()
+
+print_lock = threading.Lock()
+
+
+def safe_print(*args, **kwargs):
+    with print_lock:
+        print(*args, **kwargs)
+
+
+def download_file(file, run_dir):
+    target = run_dir / file.name
+
+    try:
+        # skip existing files
+        if target.exists():
+            return f"SKIP FILE {target}"
+
+        target.parent.mkdir(parents=True, exist_ok=True)
+
+        file.download(root=run_dir, replace=False)
+
+        return f"DONE FILE {target}"
+
+    except Exception as e:
+        return f"FAIL FILE {target}: {e}"
+
+
+def sanitize_artifact_name(name: str):
+    return name.replace(":", "_")
+
+
+def download_artifact(artifact, artifact_root):
+    try:
+        artifact_name = sanitize_artifact_name(artifact.name)
+        artifact_dir = artifact_root / artifact_name
+
+        if artifact_dir.exists() and any(artifact_dir.iterdir()):
+            return f"SKIP ARTIFACT {artifact.name}"
+
+        artifact_dir.mkdir(parents=True, exist_ok=True)
+
+        artifact.download(root=artifact_dir)
+
+        return f"DONE ARTIFACT {artifact.name}"
+
+    except Exception as e:
+        return f"FAIL ARTIFACT {artifact.name}: {e}"
+
+
+def download_run(run, root):
+    run_dir = root / f"{run.name}"
+    run_dir.mkdir(parents=True, exist_ok=True)
+
+    safe_print(f"\n=== {run.name} ({run.id}) ===")
+
+    # -------------------------
+    # Download regular run files
+    # -------------------------
+    files = list(run.files())
+
+    with ThreadPoolExecutor(max_workers=MAX_FILE_WORKERS) as executor:
+        futures = [executor.submit(download_file, file, run_dir) for file in files]
+
+        for future in as_completed(futures):
+            safe_print(future.result())
+
+    # -------------------------
+    # Download logged artifacts
+    # -------------------------
+    artifact_root = run_dir / "artifacts"
+
+    try:
+        artifacts = list(run.logged_artifacts())
+        safe_print(f"Found {len(artifacts)} artifacts for {run.name}")
+
+        with ThreadPoolExecutor(max_workers=MAX_ARTIFACT_WORKERS) as executor:
+            futures = [
+                executor.submit(download_artifact, artifact, artifact_root)
+                for artifact in artifacts
+            ]
+
+            for future in as_completed(futures):
+                safe_print(future.result())
+
+    except Exception as e:
+        safe_print(f"Artifact download failed for {run.name}: {e}")
+
+    # -------------------------
+    # OPTIONAL: download used/input artifacts
+    # -------------------------
+    # try:
+    #     used_artifacts = list(run.used_artifacts())
+    #     used_root = run_dir / "used_artifacts"
+    #
+    #     for artifact in used_artifacts:
+    #         download_artifact(artifact, used_root)
+    # except Exception as e:
+    #     safe_print(f"Used artifact download failed: {e}")
+
+    safe_print(f"Finished {run.name}")
+
+
+def main(entity: str, project: str, root: Path):
+    root.mkdir(exist_ok=True)
+
+    runs = list(api.runs(f"{entity}/{project}"))
+
+    safe_print(f"Found {len(runs)} runs")
+
+    with ThreadPoolExecutor(max_workers=MAX_RUN_WORKERS) as executor:
+        futures = [executor.submit(download_run, run, root) for run in runs]
+
+        for future in as_completed(futures):
+            try:
+                future.result()
+            except Exception as e:
+                safe_print("RUN FAILED:", e)
+
+    safe_print("\nAll downloads complete.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--entity", type=str, default="SEL3-2026-Groep-4")
+    parser.add_argument("--project", type=str, required=True)
+    parser.add_argument("--root", type=str, default="runs")
+    args = parser.parse_args()
+
+    root = Path(args.root)
+    main(entity=args.entity, project=args.project, root=root)

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()