Merge branch 'dev' into simulate-results

2026-04-15 14:27:30 +02:00 · 2026-04-15 14:27:30 +02:00 · 8876b3f2c1
commit 8876b3f2c1
parent 5329ae18ba a9ac7e7c0d
61 changed files with 2677 additions and 704 deletions
--- a/.devcontainer/devcontainer.json
+++ b/.devcontainer/devcontainer.json
@ -34,8 +34,7 @@
    },
    "remoteUser": "vscode",
    "runArgs": [
-        "--gpus",
-        "all"
+        "--device", "nvidia.com/gpu=all"
    ],
    // Ensure the .venv persists using a named volume for performance and parity
    "mounts": [
@ -46,4 +45,4 @@
    "features": {
        "ghcr.io/devcontainers/features/common-utils:1": {}
    }
-}
+}
--- a/.env.example
+++ b/.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
--- a/.github/workflows/update_hpc_requirements.yml
+++ b/.github/workflows/update_hpc_requirements.yml
@ -0,0 +1,38 @@
+name: Update HPC requirements
+
+on:
+  push:
+    paths:
+      - pyproject.toml
+    branches:
+      - main
+      - dev
+      - "ci/**"
+
+jobs:
+  update-hpc-requirements:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: 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: Commit updated requirements if changed
+        uses: stefanzweifel/git-auto-commit-action@v5
+        with:
+          commit_message: "chore(hpc): update env/hpc/requirements.txt from pyproject.toml [skip ci]"
+          file_pattern: env/hpc/requirements.txt
+          commit_author: "github-actions[bot] <github-actions[bot]@users.noreply.github.com>"
--- a/.gitignore
+++ b/.gitignore
@ -1,6 +1,7 @@
 # Model files
 artifacts/*
 runs/*
+wandb/

 # Python-generated files
 __pycache__/
@ -372,7 +373,6 @@ celerybeat.pid
 # Environments
 .env
 .venv
-env/
 venv/
 ENV/
 env.bak/
@ -472,7 +472,6 @@ tags
 [Ll]ib
 [Ll]ib64
 [Ll]ocal
-[Ss]cripts
 pyvenv.cfg
 .venv
 pip-selfcheck.json
@ -519,3 +518,4 @@ Icon
 Network Trash Folder
 Temporary Items
 .apdisk
+*.pdf
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@ -0,0 +1,7 @@
+{
+    "python.testing.pytestArgs": [
+        "tests"
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true
+}
--- a/README.md
+++ b/README.md
@ -1,8 +1,11 @@
 # Brittle Star

-## Usage
+> 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?

-### UV
+## Quick start
+
+### Local setup

 To set up the UV module, you can run the following command:

@ -10,8 +13,49 @@ To set up the UV module, you can run the following command:
 uv sync --frozen
 ```

+### Configuration
+
+1. **Copy the default configuration:**
+   ```bash
+   cp configs/default_ppo.yaml configs/my_experiment.yaml
+   ```
+
+2. **Edit `configs/my_experiment.yaml`** to set your WandB credentials:
+   ```yaml
+   track: true  # Enable WandB logging
+   wandb_entity: "your-wandb-username"  # Replace with your username/team
+   wandb_project_name: "PPO-Modularity"
+   ```
+
+3. **(Optional) Login to WandB:**
+   ```bash
+   uv run wandb login
+   ```
+
+### Training
+
 example command:

 ```bash
-uv run src/train.py --model_name my_model --epochs 50 --batch_size 32
+uv run python scripts/train.py
 ```
+
+Or use a custom config file:
+
+```bash
+uv run python scripts/train.py --config configs/my_experiment.yaml
+```
+
+Override specific parameters:
+
+```bash
+uv run python scripts/train.py --learning-rate 0.001 --num-envs 32 --track
+```
+
+## 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).
--- a/configs/README.md
+++ b/configs/README.md
@ -0,0 +1,23 @@
+# Configuration Files
+
+This directory contains configuration files for training experiments.
+
+## Usage
+
+Use `--config` with `scripts/train.py` to run an experiment:
+
+```bash
+python scripts/train.py --config configs/default_ppo.yaml
+```
+
+You can overriding settings via CLI:
+```bash
+python scripts/train.py --config configs/default_ppo.yaml --learning-rate 0.001
+```
+
+## Available Configurations
+
+- `default_ppo.yaml`: Baseline config.
+- `dev_test.yaml`: Fast iteration for development.
+- `production_training.yaml`: Full-scale training.
+- `personal_template.yaml`: Template for team members to customize.
--- a/configs/default_ppo.yaml
+++ b/configs/default_ppo.yaml
@ -0,0 +1,48 @@
+# PPO Training Configuration Template
+#
+# This file provides an example configuration for PPO training.
+# Copy this file and modify it for your specific experiments.
+#
+# Usage:
+#   python src/train.py --config-path configs/my_config.yaml
+# Or override specific parameters:
+#   python src/train.py --learning-rate 0.001 --num-envs 32
+
+# Experiment settings
+exp_name: "brittle_star_ppo"
+seed: 1
+
+# Tracking settings
+track: false  # Set to true to enable WandB logging
+wandb_project_name: "PPO-Modularity"
+wandb_entity: "SEL3-2026-Groep-4"  # Set to your WandB username or team name
+
+# Model saving
+save_model: true
+checkpoint_frequency: 100  # Save checkpoint every N iterations (0 = no checkpoints)
+
+# Environment settings
+num_envs: 16
+
+# Training hyperparameters
+total_timesteps: 10000000
+learning_rate: 0.00025
+num_steps: 128
+anneal_lr: true
+
+# PPO specific
+gamma: 0.99
+gae_lambda: 0.95
+num_minibatches: 4
+update_epochs: 4
+norm_adv: true
+clip_coef: 0.1
+clip_vloss: true
+ent_coef: 0.01
+vf_coef: 0.5
+max_grad_norm: 0.5
+target_kl: null
+
+# Hardware
+cuda: true
+torch_deterministic: true
--- a/configs/dev_test.yaml
+++ b/configs/dev_test.yaml
@ -0,0 +1,42 @@
+# Quick Development/Testing Configuration
+#
+# Fast configuration for development and testing with short runs.
+
+# Experiment settings
+exp_name: "brittle_star_dev_test"
+seed: 123
+
+# Tracking settings - IMPORTANT: Set your own wandb_entity!
+track: true
+wandb_project_name: "PPO-Modularity-Dev"
+wandb_entity: "SEL3-2026-Groep-4"  # ⚠️  SET THIS TO YOUR WANDB USERNAME OR TEAM
+
+# Model saving
+save_model: true
+checkpoint_frequency: 10  # More frequent checkpoints for testing
+
+# Environment settings
+num_envs: 4  # Smaller for faster iteration
+
+# Training hyperparameters - Fast/testing
+total_timesteps: 100000  # Short run for testing
+learning_rate: 0.001  # Higher learning rate for faster learning
+num_steps: 64  # Shorter rollouts
+anneal_lr: true
+
+# PPO specific - Optimized for quick results
+gamma: 0.99
+gae_lambda: 0.95
+num_minibatches: 2
+update_epochs: 2  # Fewer epochs for speed
+norm_adv: true
+clip_coef: 0.1
+clip_vloss: true
+ent_coef: 0.02  # Higher entropy for exploration
+vf_coef: 0.5
+max_grad_norm: 0.5
+target_kl: null
+
+# Hardware
+cuda: true
+torch_deterministic: true
--- a/configs/example.json
+++ b/configs/example.json
@ -1,8 +0,0 @@
-{
-  "morphology": {
-    "num_arms": 2,
-    "num_segments_per_arm": 4,
-    "use_p_control": true,
-    "use_torque_control": false
-  }
-}
--- a/configs/example.yaml
+++ b/configs/example.yaml
@ -0,0 +1,5 @@
+morphology:
+  num_arms: 2
+  num_segments_per_arm: 4
+  use_p_control: true
+  use_torque_control: false
--- a/configs/hpc/smoke_test.yaml
+++ b/configs/hpc/smoke_test.yaml
@ -0,0 +1,11 @@
+# Minimal config to verify HPC setup is functional.
+# Run with: python scripts/train.py --config-path configs/hpc/smoke_test.yaml
+exp_name: "hpc_smoke_test"
+seed: 0
+track: false           # Test WandB integration
+capture_video: false  # No rendering for smoke test
+save_model: true      # Test the end-of-training save routine
+num_envs: 512
+total_timesteps: 65536
+num_steps: 128
+cuda: true
--- a/configs/hpc/wandb_expand.yaml
+++ b/configs/hpc/wandb_expand.yaml
@ -0,0 +1,10 @@
+exp_name: "explained_var_fun_more_steps"
+seed: 42
+track: true
+wandb_project_name: "LET-THERE-BE-MORE-LOGGING"
+wandb_entity: "SEL3-2026-Groep-4"
+
+num_envs: 16
+num_steps: 256
+total_timesteps: 50000
+cuda: true
--- a/configs/hpc/wandb_test.yaml
+++ b/configs/hpc/wandb_test.yaml
@ -0,0 +1,11 @@
+# Configuration to verify WandB online tracking
+exp_name: "hpc_wandb_verification"
+seed: 42
+track: true               # Enabled for testing WandB
+wandb_project_name: "PPO-Modularity"
+wandb_entity: "SEL3-2026-Groep-4"
+
+num_envs: 128
+total_timesteps: 50000    # Short run for quick verification
+num_steps: 128
+cuda: true
--- a/configs/personal_template.yaml
+++ b/configs/personal_template.yaml
@ -0,0 +1,40 @@
+# Personal Configuration Example for Team Member
+#
+# Copy this template and customize for your personal experiments
+
+# Experiment settings - PERSONALIZE THESE
+exp_name: "YOUR_NAME_experiment_v1"  # ⚠️ Change YOUR_NAME 
+seed: 42
+
+# WandB settings - ⚠️ IMPORTANT: Set your credentials!
+track: true  # Enable WandB tracking
+wandb_project_name: "PPO-Modularity"
+wandb_entity: "SEL3-2026-Groep-4"  # ⚠️ CHANGE THIS to your WandB username/team
+
+# Quick experiment settings (modify as needed)
+total_timesteps: 500000  # 500K for quick results
+num_envs: 8
+learning_rate: 0.0005
+num_steps: 128
+
+# Model saving
+save_model: true  
+checkpoint_frequency: 25  # Save checkpoints frequently
+
+# Standard PPO settings (usually don't need to change)
+gamma: 0.99
+gae_lambda: 0.95
+num_minibatches: 4
+update_epochs: 4
+norm_adv: true
+clip_coef: 0.2
+clip_vloss: true
+ent_coef: 0.01
+vf_coef: 0.5
+max_grad_norm: 0.5
+target_kl: null
+anneal_lr: true
+
+# Hardware
+cuda: true
+torch_deterministic: true
--- a/configs/production_training.yaml
+++ b/configs/production_training.yaml
@ -0,0 +1,42 @@
+# Production Training Configuration
+#
+# Full-scale training configuration for production runs
+# with wandb logging enabled.
+
+# Experiment settings
+exp_name: "brittle_star_production_training"
+seed: 42
+
+# Tracking
+track: true
+capture_video: false
+wandb_project_name: "PPO-Modularity"
+wandb_entity: "SEL3-2026-Groep-4"
+
+# Model saving
+save_model: true
+checkpoint_frequency: 100  # Save checkpoint every 100 iterations
+
+# Environment settings
+num_envs: 512
+
+# Training hyperparameters
+total_timesteps: 50000000
+num_steps: 256
+num_minibatches: 4
+update_epochs: 4
+
+learning_rate: 2.5e-4
+anneal_lr: true
+gamma: 0.99
+gae_lambda: 0.95
+clip_coef: 0.1
+clip_vloss: true
+ent_coef: 0.01
+vf_coef: 0.5
+max_grad_norm: 0.5
+target_kl: null
+
+# Hardware
+cuda: true
+torch_deterministic: true
--- a/docs/CONTRIBUTING.md
+++ b/docs/CONTRIBUTING.md
@ -32,4 +32,13 @@ Code readability is paramount, as code is read far more frequently than it is wr
 * **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 (i.e. Google standard). This is enforced using build tools and pre-commit hooks such as flake8, black, or isort.
+* **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.
--- a/docs/DEVELOPMENT.md
+++ b/docs/DEVELOPMENT.md
@ -59,3 +59,25 @@ Verify your setup by running the JAX initialization test:
 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. For a full API reference, see the [package README](../src/experiment_logger/README.md).
+
+### Quick Setup
+
+1.  **Authorization**: Export your API key in your terminal to enable WandB synchronization:
+    ```bash
+    export WANDB_API_KEY=your_copied_api_key_here
+    ```
+2.  **Toggle Tracking**: Use the `--track` flag in `scripts/train.py` to enable online sync.
+3.  **Local Monitoring**: All runs are recorded in the `runs/` directory. View scalars with TensorBoard:
+    ```bash
+    tensorboard --logdir runs/
+    ```
+
+### Environment Awareness
+
+The logger automatically detects if it is running in an interactive terminal or a non-interactive environment (like an HPC Slurm job). It will automatically disable progress bars and switch to robust fallback modes (offline logging) to ensure your experiments never hang.
+
+
--- a/docs/HPC.md
+++ b/docs/HPC.md
@ -0,0 +1,85 @@
+# HPC Guide
+
+Full documentation: <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.
--- a/docs/README.md
+++ b/docs/README.md
@ -0,0 +1,15 @@
+# Documentation
+
+## Design & architecture ([`/design`](./design/))
+
+- [Actor/critic architecture](./design/actor-critic.md): Description of the actor-critic pipeline.
+- [Communication](./design/communication.md): Message propagation, Nerve-Net style.
+- [Controllers](./design/controllers.md): Macroscopig brain toplogy, centralized, arm-level, segment-level.
+- [Input/output](./design/input_action_spaces.md): Description of the model's input and output.
+- [Learning algorithm](./design/learning_algorithm.md): RL techniques, i.e. PPO.
+- [Reward function](./design/learning_algorithm.md): Goals, fitness tracking, and reward structures.
+
+## API reference ([`/api`](./api/))
+
+- [Environment](./api/environment.md): MuJoCo environment interaction, state retrieval, and configuration.
+- [Simulate](./api/simulate.md): Simulation rendering.
--- a/docs/api/simulate.md
+++ b/docs/api/simulate.md
@ -0,0 +1,14 @@
+# Training and Simulation for Brittle Star Models
+
+## Simulating a model
+
+In order to simulate and view the behavior of a trained model, you can use the `simulate.py` script. This script allows you to specify the path to a trained model and will launch a simulation using that model. This script has the following parameters:
+
+- `--model`: The path to the trained model artifact to simulate.
+- `--model-type`: The type of model to simulate (e.g., `random`, ...)
+- `--task`: The task to simulate (e.g., `directed_locomotion`, ...)
+- `--seed`: The random seed for reproducibility.
+
+```bash
+python simulate.py --model artifacts/my_model --model-type random --task directed_locomotion --seed 0
+```
--- a/docs/api/train_simulate.md
+++ b/docs/api/train_simulate.md
@ -1,30 +0,0 @@
-# Training and Simulation for Brittle Star Models
-
-## Training a model
-
-To train a model, you can use the `train.py` script. This script allows to pass some parameters to customize the training process:
-
- `--out`: The output path where the trained model will be saved.
- `--model_type`: The type of model to train (e.g., `random`, ...)
- `--task`: The task to train on (e.g., `directed_locomotion`, ...)
- `--seed`: The random seed for reproducibility.
- `--epochs`: The number of epochs to train for.
-
-This will then train the specified model on the specified task for the given number of epochs and save the trained model to the specified output path.
-
-```bash
-python train.py --out artifacts/my_model --model-type random --task directed_locomotion --seed 0 --epochs 50 
-```
-
-## Simulating a model
-
-In order to simulate and view the behavior of a trained model, you can use the `simulate.py` script. This script allows you to specify the path to a trained model and will launch a simulation using that model. This script has the following parameters:
-
- `--model`: The path to the trained model artifact to simulate.
- `--model-type`: The type of model to simulate (e.g., `random`, ...)
- `--task`: The task to simulate (e.g., `directed_locomotion`, ...)
- `--seed`: The random seed for reproducibility.
-
-```bash
-python simulate.py --model artifacts/my_model --model-type random --task directed_locomotion --seed 0
-```
--- a/docs/design/actor-critic.md
+++ b/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: 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 2 hidden layers of 64 nodes each (`[64, 64]`) 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/
+- 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.
+- 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.
--- a/docs/design/input_action_spaces.md
+++ b/docs/design/input_action_spaces.md
@ -0,0 +1,52 @@
+# 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 objective.
+
+- Joint positions: the current angles of all joints in the morphology.
+- Joint velocities: the current moving speed of the joints.
+- Goal vector: instad of just a scalar distance, the goal is represented asa a vector (distance and ange/direction) to
+  the target.
+
+**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.
+
+## 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.
+- Goal Vector (Distance + Angle): Providing only the scalar "distance to the goal" as an input is akin to blindfolding
+  the robot and asking it to find a target by playing "hot or cold." By providing a full vector, the agent knows
+  exactly where the target is relative to its current orientation, allowing for directed and efficient locomotion.
+- Absolute Joint Offsets: The physical Brittle Star robot relies on servo motors (if we were to build this simulated
+  robot), which are inherently position-controlled devices. (Continuous rotation servos exist, but they are less
+  commonly used for joints.) If our network outputted continuous torques (forces), a significant portion of the
+  reinforcement learning process would be wasted on learning low-level PID control dynamics (i.e., how much force to
+  apply to hold a position). Abstracting this away forces the learning algorithm to focus entirely on higher-level gait
+  generation and locomotion.
+
+## 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.
--- a/docs/design/reward_function.md
+++ b/docs/design/reward_function.md
@ -7,6 +7,11 @@ inputs must be distributed fairly to guarantee an objective comparison between d
 - 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.
+
+## 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

--- a/docs/design_decisions/.keep
+++ b/docs/design_decisions/.keep
--- a/env/hpc/modules.txt
+++ b/env/hpc/modules.txt
@ -0,0 +1,3 @@
+GCCcore/13.3.0
+Python/3.12.3-GCCcore-13.3.0
+FFmpeg/7.0.2-GCCcore-13.3.0
--- a/env/hpc/requirements.txt
+++ b/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
+tyro>=1.0.10
+wandb==0.24.2
+torch>=2.4.0
--- a/pyproject.toml
+++ b/pyproject.toml
@ -1,3 +1,7 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
 [project]
 name = "2026sel3-project"
 version = "0.1.0"
@ -12,19 +16,28 @@ dependencies = [
    "gymnasium>=1.2.3",
    "ipykernel==7.2.0",
    "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",
    "tyro>=1.0.10",
    "wandb==0.24.2",
+    "torch>=2.4.0",
 ]

 [project.optional-dependencies]
 cuda = [
    "jax[cuda13]==0.9.0.1",
 ]
+analysis = [
+    "tensorboard",
+]

 [dependency-groups]
 dev = [
@ -32,3 +45,36 @@ dev = [
    "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
--- a/ruff.toml
+++ b/ruff.toml
@ -1,4 +1,5 @@
 line-length = 100
+exclude = ["wandb"]

 [lint]
 extend-select = [
@ -349,7 +350,7 @@ extend-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
@ -393,3 +394,6 @@ extend-ignore = [
  "PLW0603",  # global-statement
  # "PLW1404",  # implicit-str-concat
 ]
+
+[lint.per-file-ignores]
+"__init__.py" = ["F401"]
--- a/scripts/analysis/README.md
+++ b/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`)
--- a/scripts/analysis/explore_tensorboard.py
+++ b/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()
--- a/scripts/hpc/export_requirements.py
+++ b/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()
--- a/scripts/hpc/install.sh
+++ b/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'
--- a/scripts/hpc/train.pbs
+++ b/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
+
+# TODO Once experiments get serious, change the config
+python scripts/train.py \
+    --env-config-path configs/hpc/wandb_expand.yaml \
+    --hyperparameter-config-path configs/hpc/wandb_expand.yaml \
+    --run-dir "$SCRATCH_RUNDIR"
+
+echo ">>> Staging out results to $DATA_RUNDIR..."
+cp -r "$SCRATCH_RUNDIR/." "$DATA_RUNDIR/"
+
+echo ">>> Done"
--- a/experiments/simulate.py
+++ b/experiments/simulate.py
@ -12,7 +12,7 @@ import numpy as np
 from brittle_star_project import (
    Backend,
 )
-from brittle_star_project.environment import from_json
+from brittle_star_project.environment import from_file

 def _flatten_obs_dict(obs_dict: dict[str, Any]) -> jnp.ndarray:
    """Flatten the env's observation dict into a 1D vector.
@ -317,7 +317,7 @@ def main() -> None:

    args = parse_args()

-    morphology_cfg, arena_cfg, env_cfg = from_json("../configs/test.json")
+    morphology_cfg, arena_cfg, env_cfg = from_file("../configs/test.yaml")

    # ======= ENVIRONMENT SETUP =======

--- a/scripts/train.py
+++ b/scripts/train.py
@ -0,0 +1,75 @@
+import subprocess
+import time
+
+import torch
+import os
+
+from brittle_star_project.dataclasses import PPOArgs
+from brittle_star_project.trainers.PPOTrainer import PPOTrainer
+from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
+
+from experiment_logger import UnifiedLogger
+from experiment_logger.config_utils import merge_config_with_cli, print_config
+
+
+def make_env(config_path: str | None, num_envs: int) -> BrittleStarJaxEnvWrapper:
+    if config_path is None:
+        return BrittleStarJaxEnvWrapper.default(num_envs=num_envs)
+    return BrittleStarJaxEnvWrapper.from_config(config_path, num_envs=num_envs)
+
+
+def parse_args() -> PPOArgs:
+    import argparse
+
+    # Use argparse to reliably extract just the config path without swallowing --help
+    parser = argparse.ArgumentParser(add_help=False)
+    parser.add_argument("--hyperparameter-config-path", type=str, default=None)
+    known_args, _ = parser.parse_known_args()
+
+    args = merge_config_with_cli(PPOArgs, config_file=known_args.hyperparameter_config_path)
+    return args
+
+
+def get_git_hash() -> str:
+    try:
+        return (
+            subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode("ascii").strip()
+        )
+    except (subprocess.CalledProcessError, UnicodeDecodeError):
+        return "none"
+
+
+if __name__ == "__main__":
+    args = parse_args()
+
+    args.batch_size = args.num_envs * args.num_steps
+    args.minibatch_size = args.batch_size // args.num_minibatches
+    args.num_iterations = args.total_timesteps // args.batch_size
+
+    git_hash = get_git_hash()
+    run_name = f"{args.exp_name}__seed_{args.seed}__{git_hash}__{int(time.time())}"
+
+    if args.run_dir is None:
+        run_dir = f"runs/{run_name}"
+    else:
+        run_dir = args.run_dir
+
+    os.makedirs(run_dir, exist_ok=True)
+
+    # Initialize Global Logger
+    logger = UnifiedLogger(
+        config=vars(args),
+        project_name=args.wandb_project_name,  # or default PPO-Modularity if missing
+        run_name=run_name,
+        base_dir=os.path.dirname(run_dir),
+        use_wandb=args.track,
+    )
+
+    print_config(args, title="PPO Training Configuration")
+
+    env = make_env(args.env_config_path, args.num_envs)
+
+    torch.backends.cudnn.deterministic = args.torch_deterministic
+
+    ppo_trainer = PPOTrainer(args, env, run_dir, run_name)
+    ppo_trainer.train()
--- a/src/init.py
+++ b/src/init.py
@ -1,25 +0,0 @@
-from .brittle_star_project import (
-    ArenaConfig,
-    Backend,
-    BrittleStarEnv,
-    BrittleStarEnvFactory,
-    EnvConfig,
-    MorphologyConfig,
-    Task,
-    simulate_policy,
-    SimulationConfig,
-    ControlPolicy,
-)
-
-__all__ = [
-    "ArenaConfig",
-    "Backend",
-    "BrittleStarEnv",
-    "BrittleStarEnvFactory",
-    "EnvConfig",
-    "MorphologyConfig",
-    "Task",
-    "simulate_policy",
-    "SimulationConfig",
-    "ControlPolicy",
-]
--- a/src/brittle_star_project/rl/DummyAgent.py
+++ b/src/brittle_star_project/rl/DummyAgent.py
@ -1,36 +1,27 @@
-from dataclasses import dataclass, fields
+from dataclasses import dataclass, fields, field

 import flax
 import flax.linen as nn
 import jax.numpy as jnp
 import jax.tree_util
-import numpy as np
+from typing import Sequence, Callable
 from flax.linen.initializers import constant, orthogonal


-class Network(nn.Module):
-    """
-    Dummy model only used for testing purposes
-
-    inspired by: https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/ppo_atari_envpool_xla_jax_scan.py
-    """
-
-    hidden_dim: int = 195
+# semi generic so we can easily make a config for it in experiments
+class GenericDenseLayersWithActivation(nn.Module):
+    layer_sizes: Sequence[int] = field(default_factory=lambda: [64, 64])
+    activation: Callable = nn.tanh

    @nn.compact
    def __call__(self, x):
-        x = nn.Dense(self.hidden_dim, kernel_init=orthogonal(np.sqrt(2)), bias_init=constant(0.0))(
-            x
-        )
-        x = nn.relu(x)
-        x = nn.Dense(self.hidden_dim, kernel_init=orthogonal(np.sqrt(2)), bias_init=constant(0.0))(
-            x
-        )
-        x = nn.relu(x)
+        for size in self.layer_sizes:
+            x = nn.Dense(size, kernel_init=orthogonal(jnp.sqrt(2)))(x)
+            x = self.activation(x)
        return x


-class Critic(nn.Module):
+class OneDenseLayerMLP(nn.Module):
    @nn.compact
    def __call__(self, x):
        return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x)
@ -49,10 +40,10 @@ class Actor(nn.Module):
@jax.tree_util.register_dataclass
@dataclass
 class AgentParams:
-    network_params: flax.core.FrozenDict
+    sensor_params: flax.core.FrozenDict
    actor_params: flax.core.FrozenDict
    critic_params: flax.core.FrozenDict
-    critic_network_params: flax.core.FrozenDict
+    feature_extractor_params: flax.core.FrozenDict


@jax.tree_util.register_dataclass
--- a/src/brittle_star_project/dataclasses/PPOArgs.py
+++ b/src/brittle_star_project/dataclasses/PPOArgs.py
@ -1,6 +1,9 @@
 from dataclasses import dataclass

+import jax

+
+@jax.tree_util.register_dataclass
@dataclass
 class PPOArgs:
    """
@ -8,11 +11,17 @@ class PPOArgs:
    """

    # path to environment config file, if None, use default config
-    config_path: str | None = None
+    env_config_path: str | None = None
+
+    # path to hyperparameter config file (yaml), if None, use default config
+    hyperparameter_config_path: str | None = None

    # the name of this experiment
    exp_name: str = "brittle_star_ppo"

+    # the directory to save the experiment results
+    run_dir: str | None = None
+
    # seed of the experiment
    seed: int = 1

@ -29,7 +38,7 @@ class PPOArgs:
    wandb_project_name: str = "PPO-Modularity"

    # the entity (team) of wandb's project
-    wandb_entity: str | None = None
+    wandb_entity: str | None = "SEL3-2026-Groep-4"

    # whether to capture videos of the agent performances (check out `videos` folder)
    capture_video: bool = False
@ -37,6 +46,9 @@ class PPOArgs:
    # whether to save model into the `runs/{run_name}` folder
    save_model: bool = True

+    # checkpoint frequency (in iterations, 0 = no intermediate checkpoints)
+    checkpoint_frequency: int = 100
+
    # whether to upload the saved model to huggingface
    upload_model: bool = False

--- a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py
+++ b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py
@ -8,7 +8,7 @@ from brittle_star_project import (
    ArenaConfig,
    Backend,
 )
-from brittle_star_project.environment import from_json
+from brittle_star_project.environment import from_file


 class BrittleStarJaxEnvWrapper:
@ -35,6 +35,13 @@ class BrittleStarJaxEnvWrapper:

        self._action_rng = None

+        from experiment_logger import get_logger
+
+        self.logger = get_logger()
+        self.logger.info(
+            f"Initialized BrittleStarJaxEnvWrapper with {num_envs} envs on {backend.value}"
+        )
+
    @property
    def backend(self):
        return self._backend
@ -52,6 +59,7 @@ class BrittleStarJaxEnvWrapper:
        return self._env.observation_space

    def reset(self, seed: int = 0):
+        self.logger.info(f"Resetting vectorized environment environments with seed {seed}")
        self._action_rng, env_rng = jax.random.split(jax.random.PRNGKey(seed), 2)
        env_rngs = jnp.array(jax.random.split(env_rng, self._num_envs))
        return self._vectorized_reset(rng=env_rngs)
@ -82,7 +90,7 @@ class BrittleStarJaxEnvWrapper:
    def from_config(
        config_path: str, num_envs: int, backend: Backend = Backend.MJX
    ) -> "BrittleStarJaxEnvWrapper":
-        morphology_cfg, arena_cfg, env_cfg = from_json(config_path)
+        morphology_cfg, arena_cfg, env_cfg = from_file(config_path)
        return BrittleStarJaxEnvWrapper(
            morphology_cfg, arena_cfg, env_cfg, num_envs=num_envs, backend=backend
        )
--- a/src/brittle_star_project/environment/init.py
+++ b/src/brittle_star_project/environment/init.py
@ -1,4 +1,4 @@
-from .env_config import ArenaConfig, EnvConfig, MorphologyConfig, from_json
+from .env_config import ArenaConfig, EnvConfig, MorphologyConfig, from_file
 from .env_types import Backend, Task
 from .env_wrapper import BrittleStarEnv, StepResult
 from .factory import BrittleStarEnvFactory
@ -12,5 +12,5 @@ __all__ = [
    "BrittleStarEnv",
    "StepResult",
    "BrittleStarEnvFactory",
-    "from_json",
+    "from_file",
 ]
--- a/src/brittle_star_project/environment/env_config.py
+++ b/src/brittle_star_project/environment/env_config.py
@ -1,7 +1,6 @@
 from __future__ import annotations

 from dataclasses import dataclass, field
-import json

 from .env_types import Task

@ -50,10 +49,14 @@ class EnvConfig:
    light_perlin_noise_scale: int = 0


-def from_json(path: str) -> tuple[MorphologyConfig, ArenaConfig, EnvConfig]:
+def from_file(path: str) -> tuple[MorphologyConfig, ArenaConfig, EnvConfig]:
+    """Load configurations from a YAML file."""
+    import yaml
+
    with open(path, "r") as f:
-        config_json = json.load(f)
-        morphology = MorphologyConfig(**config_json.get("morphology", {}))
-        arena = ArenaConfig(**config_json.get("arena", {}))
-        env = EnvConfig(**config_json.get("env", {}))
+        config_dict = yaml.safe_load(f)
+
+        morphology = MorphologyConfig(**config_dict.get("morphology", {}))
+        arena = ArenaConfig(**config_dict.get("arena", {}))
+        env = EnvConfig(**config_dict.get("env", {}))
        return morphology, arena, env
--- a/src/brittle_star_project/environment/factory.py
+++ b/src/brittle_star_project/environment/factory.py
@ -98,9 +98,15 @@ class BrittleStarEnvFactory:
            case _:
                raise ValueError(f"Unsupported task: {env_config.task}")

-        return env_class.from_morphology_and_arena(
+        env = env_class.from_morphology_and_arena(
            morphology=morphology,
            arena=arena,
            configuration=env_configuration,
            backend=backend.value,
        )
+
+        from experiment_logger import get_logger
+
+        get_logger().info(f"Created {env_config.task.value} env on backend {backend.value}")
+
+        return env
--- a/src/brittle_star_project/ppo.py
+++ b/src/brittle_star_project/ppo.py
@ -8,9 +8,7 @@ import jax.numpy as jnp
 # Chose to use a class as it seemed the easiest way to integrate the CleanRL code style
 # with our need to seperate concerns
 class PPO:
-    def __init__(
-        self, args, input_network, action_network, critic, critic_network, message_passer=None
-    ):
+    def __init__(self, args, sensor, actor, critic, feature_extractor, message_passer=None):
        self.args = args

        if not message_passer:
@ -20,10 +18,10 @@ class PPO:
            partial(
                ppo_loss,
                args=args,
-                input_network_apply=input_network.apply,
-                action_network_apply=action_network.apply,
+                sensor_apply=sensor.apply,
+                actor_apply=actor.apply,
                critic_apply=critic.apply,
-                critic_network_apply=critic_network.apply,
+                feature_extractor_apply=feature_extractor.apply,
                message_passer=message_passer,
            ),
            has_aux=True,
@ -87,20 +85,20 @@ that are now not in the same scope


@partial(jax.jit, static_argnums=(0, 1, 2, 3, 4))
-def get_action_and_value2(
-    input_apply,
-    action_apply,
+def get_action_and_value(
+    sensor_apply,
+    actor_apply,
    message_passer,
    critic_apply,
-    critic_network_apply,
+    feature_extractor_apply,
    params: flax.core.FrozenDict,
    x: jnp.ndarray,
    action: jnp.ndarray,
 ):
-    hidden_network = input_apply(params["network_params"], x)
-    hidden_critic = critic_network_apply(params["critic_network_params"], x)
-    hidden_network = message_passer(hidden_network)
-    mean, log_std = action_apply(params["actor_params"], hidden_network)
+    hidden_sensor = sensor_apply(params["sensor_params"], x)
+    hidden_critic = feature_extractor_apply(params["feature_extractor_params"], x)
+    hidden_sensor = message_passer(hidden_sensor)
+    mean, log_std = actor_apply(params["actor_params"], hidden_sensor)
    std = jnp.exp(log_std)

    logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
@ -118,18 +116,18 @@ def ppo_loss(
    mb_advantages,
    mb_returns,
    args,
-    input_network_apply,
-    action_network_apply,
+    sensor_apply,
+    actor_apply,
    message_passer,
    critic_apply,
-    critic_network_apply,
+    feature_extractor_apply,
 ):
-    newlogprob, entropy, newvalue = get_action_and_value2(
-        input_network_apply,
-        action_network_apply,
+    newlogprob, entropy, newvalue = get_action_and_value(
+        sensor_apply,
+        actor_apply,
        message_passer,
        critic_apply,
-        critic_network_apply,
+        feature_extractor_apply,
        params,
        x,
        a,
--- a/src/brittle_star_project/render/renderer.py
+++ b/src/brittle_star_project/render/renderer.py
@ -36,6 +36,9 @@ def simulate_policy(

    import mujoco.viewer

+    if state is None:
+        raise ValueError("A valid environment state must be provided.")
+
    model = state.mj_model
    data = state.mj_data

--- a/src/brittle_star_project/rl/init.py
+++ b/src/brittle_star_project/rl/init.py
@ -1,25 +0,0 @@
-from .DummyAgent import Network, Critic, Actor, AgentParams, Storage
-from .base import (
-    RLAlgorithm,
-    RLModel,
-    Transition,
-    create_model,
-    register_rl_model,
-    registered_model_types,
-)
-from .random_policy_model import RandomPolicyModel
-
-__all__ = [
-    "RLAlgorithm",
-    "RLModel",
-    "RandomPolicyModel",
-    "Transition",
-    "create_model",
-    "register_rl_model",
-    "registered_model_types",
-    "Network",
-    "Critic",
-    "Actor",
-    "AgentParams",
-    "Storage",
-]
--- a/src/brittle_star_project/rl/base.py
+++ b/src/brittle_star_project/rl/base.py
@ -1,162 +0,0 @@
-from __future__ import annotations
-
-from abc import ABC, abstractmethod
-from dataclasses import dataclass
-import json
-from pathlib import Path
-from typing import Any
-
-
-@dataclass(frozen=True, slots=True)
-class Transition:
-    """A minimal transition container for RL.
-
-    This is intentionally generic because the underlying env state type may be a
-    JAX pytree, a numpy struct, or something library-specific.
-    """
-
-    obs: Any
-    action: Any
-    reward: float
-    next_obs: Any
-    terminated: bool
-    truncated: bool
-    info: dict[str, Any] | None = None
-
-
-class RLAlgorithm(ABC):
-    """Insertable RL algorithm interface."""
-
-    @abstractmethod
-    def select_action(self, *, obs: Any, rng: Any | None = None) -> Any:
-        raise NotImplementedError
-
-    def observe(self, transition: Transition) -> None:
-        """Optional hook to store transitions."""
-
-    def update(self, *, rng: Any | None = None) -> dict[str, float]:
-        """Optional hook to run one training update."""
-
-        return {}
-
-    def save(self, path: str) -> None:
-        raise NotImplementedError("Save not implemented")
-
-    def load(self, path: str) -> None:
-        raise NotImplementedError("Load not implemented")
-
-
-_RL_MODEL_REGISTRY: dict[str, type["RLModel"]] = {}
-
-
-def registered_model_types() -> list[str]:
-    return sorted(_RL_MODEL_REGISTRY)
-
-
-def create_model(type_name: str, *, payload: dict[str, Any]) -> "RLModel":
-    model_cls = _RL_MODEL_REGISTRY.get(type_name)
-    if model_cls is None:
-        known = ", ".join(sorted(_RL_MODEL_REGISTRY)) or "<none>"
-        raise ValueError(f"Unknown RLModel type '{type_name}'. Known: {known}")
-    return model_cls.from_payload(payload)
-
-
-def get_rl_model_registry() -> dict[str, type["RLModel"]]:
-    """Return a copy of the current RLModel registry.
-
-    The registry is populated by importing concrete model modules that use the
-    `@register_rl_model(...)` decorator.
-    """
-
-    return dict(_RL_MODEL_REGISTRY)
-
-
-def register_rl_model(*type_names: str):
-    """Decorator to register an `RLModel` for generic loading.
-
-    Concrete model modules should apply this decorator, so `base.py` never needs
-    to import concrete models (avoids circular imports).
-    """
-
-    if not type_names:
-        raise TypeError("register_rl_model() requires at least one type name")
-
-    primary = type_names[0]
-
-    def _decorator(cls: type[RLModel]):
-        for name in type_names:
-            _RL_MODEL_REGISTRY[name] = cls
-        cls.type_name = primary
-        return cls
-
-    return _decorator
-
-
-class RLModel(ABC):
-    """Serializable policy/model interface.
-
-    This is the artifact that `train.py` writes and `simulate.py` loads.
-    """
-
-    # Overwritten by the `@register_rl_model(...)` decorator.
-    type_name: str = "RLModel"
-
-    def reset(self, seed: int | None = None) -> None:
-        """Optional hook for RNG/stateful models."""
-
-    @abstractmethod
-    def act(self, *, obs: Any | None = None, t: float = 0.0) -> Any:
-        raise NotImplementedError
-
-    def train(self, *, env: Any, num_epochs: int = 1) -> None:
-        """Optional training hook.
-
-        Many models won't learn; for those this can be a no-op.
-        """
-
-        _ = (env, num_epochs)
-
-    def to_payload(self) -> dict[str, Any]:
-        """Return JSON-serializable model parameters."""
-
-        return {}
-
-    @classmethod
-    def from_payload(cls, payload: dict[str, Any]) -> "RLModel":
-        """Reconstruct a model from `to_payload()` output."""
-
-        return cls(**payload)  # type: ignore[arg-type]
-
-    def save(self, path: str | Path) -> Path:
-        out = Path(path)
-        out.parent.mkdir(parents=True, exist_ok=True)
-        doc = {
-            "type": self.type_name,
-            "version": 1,
-            "payload": self.to_payload(),
-        }
-        out.write_text(json.dumps(doc, indent=2, sort_keys=True) + "\n")
-        return out
-
-    @classmethod
-    def load(cls, path: str | Path) -> "RLModel":
-        p = Path(path)
-        doc = json.loads(p.read_text())
-
-        type_name = doc.get("type")
-        if not isinstance(type_name, str):
-            raise ValueError("Model artifact missing string field 'type'")
-
-        model_cls = _RL_MODEL_REGISTRY.get(type_name)
-        if model_cls is None:
-            known = ", ".join(sorted(_RL_MODEL_REGISTRY)) or "<none>"
-            raise ValueError(f"Unknown RLModel type '{type_name}'. Known: {known}")
-
-        payload = doc.get("payload")
-        # Backward compatibility: older artifacts stored fields at top-level.
-        if payload is None:
-            payload = {k: v for k, v in doc.items() if k not in ("type", "version")}
-        if not isinstance(payload, dict):
-            raise ValueError("Model artifact field 'payload' must be an object")
-
-        return model_cls.from_payload(payload)
--- a/src/brittle_star_project/rl/random_policy_model.py
+++ b/src/brittle_star_project/rl/random_policy_model.py
@ -1,52 +0,0 @@
-from __future__ import annotations
-
-from dataclasses import dataclass, field
-
-import numpy as np
-
-from .base import RLModel, register_rl_model
-
-
-@register_rl_model("random")
-@dataclass(slots=True)
-class RandomPolicyModel(RLModel):
-    """A minimal, serializable policy model that outputs random controls.
-
-    This is intentionally *not* a learning algorithm yet. It exists so we can:
-    - produce a stable model artifact from `train.py`
-    - load that artifact in `simulate.py`
-    - drive the MuJoCo viewer with the model's actions
-    """
-
-    nu: int = 0
-    seed: int = 0
-    ctrl_noise_scale: float = 0.5
-
-    _rng: np.random.RandomState = field(init=False, repr=False)
-
-    def __post_init__(self) -> None:
-        self.reset(self.seed)
-
-    def reset(self, seed: int | None = None) -> None:
-        if seed is not None:
-            self.seed = int(seed)
-        self._rng = np.random.RandomState(self.seed)
-
-    def act(self, *, obs: np.ndarray | None = None, t: float = 0.0) -> np.ndarray:
-        if self.nu <= 0:
-            return np.zeros((0,), dtype=np.float32)
-        ctrl = self.ctrl_noise_scale * self._rng.randn(self.nu)
-        return ctrl.astype(np.float32)
-
-    def to_payload(self) -> dict[str, object]:
-        return {
-            "seed": int(self.seed),
-            "ctrl_noise_scale": float(self.ctrl_noise_scale),
-        }
-
-    @classmethod
-    def from_payload(cls, payload: dict[str, object]) -> RandomPolicyModel:
-        return cls(
-            seed=int(payload.get("seed", 0)),
-            ctrl_noise_scale=float(payload.get("ctrl_noise_scale", 0.5)),
-        )
--- a/src/brittle_star_project/trainers/PPOTrainer.py
+++ b/src/brittle_star_project/trainers/PPOTrainer.py
@ -0,0 +1,519 @@
+import datetime
+import random
+import time
+from dataclasses import asdict, dataclass
+from functools import partial
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import optax
+from flax.training.train_state import TrainState
+
+from experiment_logger import get_logger
+
+from brittle_star_project.dataclasses import EpisodeStatistics, PPOArgs
+from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
+from brittle_star_project.MLPs.mlps import (
+    Actor,
+    AgentParams,
+    GenericDenseLayersWithActivation,
+    OneDenseLayerMLP,
+    Storage,
+)
+from brittle_star_project.ppo import PPO
+
+
+def _compute_explained_variance(values: jnp.ndarray, returns: jnp.ndarray) -> float:
+    var_returns = jnp.var(returns)
+    explained_var = 1.0 - jnp.var(returns - values) / (var_returns + 1e-8)
+    return float(explained_var)
+
+
+@jax.jit
+def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, learning_rate):
+    frac = 1.0 - (count // (minibatch_count * update_epochs)) / num_iterations
+    return learning_rate * frac
+
+
+@jax.jit
+def _convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray:
+    return jax.vmap(lambda o: jnp.concatenate([v.flatten() for v in o.values() if v.size > 0]))(
+        obs_dict
+    )
+
+
+def _get_action_and_value_noise(
+    sensor: GenericDenseLayersWithActivation,
+    feature_extractor: GenericDenseLayersWithActivation,
+    actor: Actor,
+    critic: OneDenseLayerMLP,
+    agent_state: TrainState,
+    next_obs: jnp.ndarray,
+    key: jax.random.PRNGKey,
+):
+    hidden = sensor.apply(agent_state.params["sensor_params"], next_obs)
+    hidden_critic = feature_extractor.apply(
+        agent_state.params["feature_extractor_params"], next_obs
+    )
+
+    mean, log_std = actor.apply(agent_state.params["actor_params"], hidden)
+    key, subkey = jax.random.split(key)
+    noise = jax.random.normal(subkey, shape=mean.shape)
+    std = jnp.exp(log_std)
+    action = mean + noise * std
+    logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
+    value = critic.apply(agent_state.params["critic_params"], hidden_critic)
+    return action, logprob, value.squeeze(-1), key
+
+
+def _step_once(
+    carry,
+    _,
+    env_step_fn,
+    sensor: GenericDenseLayersWithActivation,
+    feature_extractor: GenericDenseLayersWithActivation,
+    actor: Actor,
+    critic: OneDenseLayerMLP,
+):
+    agent_state, episode_stats, obs, done, key, env_state = carry
+    action, logprob, value, key = _get_action_and_value_noise(
+        sensor, feature_extractor, actor, critic, agent_state, obs, key
+    )
+
+    episode_stats, env_state, (next_obs, reward, next_done) = env_step_fn(
+        episode_stats, env_state, action
+    )
+
+    storage = Storage(
+        obs=obs,
+        actions=action,
+        logprobs=logprob,
+        dones=done,
+        values=value,
+        rewards=reward,
+        returns=jnp.zeros_like(reward),
+        advantages=jnp.zeros_like(reward),
+    )
+    return (agent_state, episode_stats, next_obs, next_done, key, env_state), storage
+
+
+def _step_env_wrapped(episode_stats, env_state, action, env_step_fn):
+    next_env_state = env_step_fn(env_state, action)
+
+    reward = next_env_state.reward
+    terminated = next_env_state.terminated
+    truncated = next_env_state.truncated
+    done = terminated | truncated
+
+    new_episode_return = episode_stats.episode_returns + reward
+    new_episode_length = episode_stats.episode_lengths + 1
+
+    episode_stats = episode_stats.replace(
+        episode_returns=new_episode_return * (1 - done),
+        episode_lengths=new_episode_length * (1 - done),
+        returned_episode_returns=jnp.where(
+            done, new_episode_return, episode_stats.returned_episode_returns
+        ),
+        returned_episode_lengths=jnp.where(
+            done, new_episode_length, episode_stats.returned_episode_lengths
+        ),
+    )
+    return (
+        episode_stats,
+        next_env_state,
+        (_convert_obs_dict_to_array(next_env_state.observations), reward, done),
+    )
+
+
+def _rollout_jit(
+    agent_state,
+    episode_stats,
+    env_state,
+    next_obs,
+    next_done,
+    key,
+    max_steps,
+    step_env_fn,
+    sensor: GenericDenseLayersWithActivation,
+    feature_extractor: GenericDenseLayersWithActivation,
+    actor: Actor,
+    critic: OneDenseLayerMLP,
+):
+    (agent_state, episode_stats, next_obs, next_done, key, env_state), storage = jax.lax.scan(
+        partial(
+            _step_once,
+            sensor=sensor,
+            feature_extractor=feature_extractor,
+            actor=actor,
+            critic=critic,
+            env_step_fn=step_env_fn,
+        ),
+        (agent_state, episode_stats, next_obs, next_done, key, env_state),
+        (),
+        max_steps,
+    )
+    return agent_state, episode_stats, next_obs, next_done, storage, key, env_state
+
+
+def _compute_gae_once(carry, inp, gamma, gae_lambda):
+    advantages = carry
+    nextdone, nextvalues, curvalues, reward = inp
+    nextnonterminal = 1.0 - nextdone
+    delta = reward + gamma * nextvalues * nextnonterminal - curvalues
+    advantages = delta + gamma * gae_lambda * nextnonterminal * advantages
+    return advantages, advantages
+
+
+def _compute_gae_jit(
+    agent_state,
+    storage,
+    next_obs,
+    next_done,
+    gamma,
+    gae_lambda,
+    num_envs,
+    feature_extractor,
+    critic,
+):
+    next_value = critic.apply(
+        agent_state.params["critic_params"],
+        feature_extractor.apply(agent_state.params["feature_extractor_params"], next_obs),
+    ).squeeze(-1)
+
+    advantages = jnp.zeros((num_envs,))
+    dones = jnp.concatenate([storage.dones, next_done[None, :]], axis=0)
+    values = jnp.concatenate([storage.values, next_value[None, :]], axis=0)
+    _, advantages = jax.lax.scan(
+        partial(_compute_gae_once, gamma=gamma, gae_lambda=gae_lambda),
+        advantages,
+        (dones[1:], values[1:], values[:-1], storage.rewards),
+        reverse=True,
+    )
+    return storage.replace(advantages=advantages, returns=advantages + storage.values)
+
+
+@dataclass
+class TrainingMeasurements:
+    loss: jnp.ndarray
+    pg_loss: jnp.ndarray
+    v_loss: jnp.ndarray
+    entropy_loss: jnp.ndarray
+    approx_kl: jnp.ndarray
+    avg_episodic_return: float
+    explained_variance: float
+    num_terminated: int
+    num_truncated: int
+    avg_terminated_length: Any
+    avg_truncated_length: Any
+
+
+class PPOTrainer:
+    def __init__(self, args: PPOArgs, env: BrittleStarJaxEnvWrapper, run_dir: str, run_name: str):
+        self.args = args
+        self.env = env
+        self.run_dir = run_dir
+        self.run_name = run_name
+        self.logger = get_logger()
+
+        self.key = jax.random.PRNGKey(args.seed)
+
+        self.sensor, self.feature_extractor, self.actor, self.critic = self._init_agent()
+        self.sensor.apply = jax.jit(self.sensor.apply)
+        self.feature_extractor.apply = jax.jit(self.feature_extractor.apply)
+        self.actor.apply = jax.jit(self.actor.apply)
+        self.critic.apply = jax.jit(self.critic.apply)
+
+        self._rollout_jit = jax.jit(
+            partial(
+                _rollout_jit,
+                max_steps=self.args.num_steps,
+                step_env_fn=partial(_step_env_wrapped, env_step_fn=self.env.step),
+                sensor=self.sensor,
+                feature_extractor=self.feature_extractor,
+                actor=self.actor,
+                critic=self.critic,
+            )
+        )
+        self._compute_gae_jit = jax.jit(
+            partial(
+                _compute_gae_jit,
+                num_envs=self.args.num_envs,
+                gamma=self.args.gamma,
+                gae_lambda=self.args.gae_lambda,
+                feature_extractor=self.feature_extractor,
+                critic=self.critic,
+            )
+        )
+
+        self._ppo = PPO(self.args, self.sensor, self.actor, self.critic, self.feature_extractor)
+
+        self.agent_state = self._init_agent_state()
+
+        self.episode_stats = self._init_episode_stats()
+
+        self._init_random()
+
+    def _init_random(self):
+        self.logger.info(f"[RANDOM]: Setting random seed to {self.args.seed}")
+
+        random.seed(self.args.seed)
+        np.random.seed(self.args.seed)
+
+    def _init_agent(self):
+        self.logger.info("[AGENT]: Initializing agent...")
+
+        sensor = GenericDenseLayersWithActivation()
+        feature_extractor = GenericDenseLayersWithActivation()
+        actor = Actor(action_dim=self.env.single_action_space.shape[0])
+        critic = OneDenseLayerMLP()
+        return sensor, feature_extractor, actor, critic
+
+    def _init_agent_state(self) -> TrainState:
+        self.logger.info("[AGENT STATE]: Initializing agent state...")
+
+        self.key, sensor_key, actor_key, critic_key, feature_extractor_key = jax.random.split(
+            self.key, 5
+        )
+
+        sample_obs = jnp.concatenate(
+            [
+                v.flatten()
+                for v in self.env.single_observation_space.sample(
+                    rng=jax.random.PRNGKey(0)
+                ).values()
+                if v.size > 0
+            ]
+        )
+        sensor_params = self.sensor.init(sensor_key, sample_obs)
+        feature_extractor_params = self.feature_extractor.init(feature_extractor_key, sample_obs)
+        actor_params = self.actor.init(actor_key, self.sensor.apply(sensor_params, sample_obs))
+        critic_params = self.critic.init(
+            critic_key, self.feature_extractor.apply(feature_extractor_params, sample_obs)
+        )
+
+        return TrainState.create(
+            apply_fn=None,
+            params=asdict(
+                AgentParams(sensor_params, actor_params, critic_params, feature_extractor_params)
+            ),
+            tx=optax.chain(
+                optax.clip_by_global_norm(self.args.max_grad_norm),
+                optax.inject_hyperparams(optax.adam)(
+                    learning_rate=partial(
+                        _linear_schedule,
+                        minibatch_count=self.args.num_minibatches,
+                        update_epochs=self.args.update_epochs,
+                        num_iterations=self.args.num_iterations,
+                        learning_rate=self.args.learning_rate,
+                    )
+                    if self.args.anneal_lr
+                    else self.args.learning_rate,
+                    eps=1e-5,
+                ),
+            ),
+        )
+
+    def _init_episode_stats(self) -> EpisodeStatistics:
+        self.logger.info("[EPISODE STATS]: Initializing episode stats...")
+
+        return EpisodeStatistics(
+            episode_returns=jnp.zeros(self.args.num_envs, dtype=jnp.float32),
+            episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32),
+            returned_episode_returns=jnp.zeros(self.args.num_envs, jnp.float32),
+            returned_episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32),
+        )
+
+    def _rollout(self, env_state, next_obs, next_done) -> tuple[Any, ...]:
+        return self._rollout_jit(
+            self.agent_state,
+            self.episode_stats,
+            env_state,
+            next_obs,
+            next_done,
+            self.key,
+        )
+
+    def _compute_gae(self, storage, next_obs, next_done) -> Storage:
+        return self._compute_gae_jit(
+            self.agent_state,
+            storage,
+            next_obs,
+            next_done,
+        )
+
+    def _log(
+        self,
+        global_step,
+        episode_stats,
+        start_time,
+        iteration_time_start,
+        training_measurements,
+    ):
+        metrics = {
+            "charts/avg_episodic_return": training_measurements.avg_episodic_return,
+            "charts/avg_episodic_length": np.mean(
+                jax.device_get(episode_stats.returned_episode_lengths)
+            ),
+            "charts/learning_rate": self.agent_state.opt_state[1]
+            .hyperparams["learning_rate"]
+            .item(),
+            "charts/explained_variance": training_measurements.explained_variance,
+            "charts/num_terminated": training_measurements.num_terminated,
+            "charts/num_truncated": training_measurements.num_truncated,
+            "charts/avg_terminated_ep_length": training_measurements.avg_terminated_length,
+            "charts/avg_truncated_ep_length": training_measurements.avg_truncated_length,
+            "losses/value_loss": training_measurements.v_loss[-1, -1].item(),
+            "losses/policy_loss": training_measurements.pg_loss[-1, -1].item(),
+            "losses/entropy": training_measurements.entropy_loss[-1, -1].item(),
+            "losses/approx_kl": training_measurements.approx_kl[-1, -1].item(),
+            "losses/loss": training_measurements.loss[-1, -1].item(),
+            "charts/SPS": int(global_step / (time.time() - start_time)),
+            "charts/SPS_update": int(
+                self.args.num_envs * self.args.num_steps / (time.time() - iteration_time_start)
+            ),
+        }
+        self.logger.log(metrics, step=global_step)
+
+    def _step(self, env_state, next_obs, next_done, iteration: int) -> tuple:
+        if iteration == 1:
+            self.logger.log_non_interactive(f"Starting first rollout (JIT): {time.ctime()}")
+
+        (
+            self.agent_state,
+            self.episode_stats,
+            next_obs,
+            next_done,
+            storage,
+            self.key,
+            next_env_state,
+        ) = self._rollout(env_state, next_obs, next_done)
+
+        if iteration == 1:
+            self.logger.log_non_interactive(f"First rollout completed: {time.ctime()}")
+
+        storage = self._compute_gae(storage, next_obs, next_done)
+
+        if iteration == 1:
+            self.logger.log_non_interactive(f"Starting first PPO update (JIT): {time.ctime()}")
+
+        self.agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, self.key = (
+            self._ppo.update_ppo(self.agent_state, storage, self.key)
+        )
+
+        if iteration == 1:
+            self.logger.log_non_interactive(f"First PPO update completed: {time.ctime()}")
+
+        avg_episodic_return = float(
+            jnp.mean(jax.device_get(self.episode_stats.returned_episode_returns)).item()
+        )
+
+        explained_var = _compute_explained_variance(storage.values, storage.returns)
+
+        terminated = next_env_state.terminated
+        truncated = next_env_state.truncated
+        episode_lengths = self.episode_stats.returned_episode_lengths
+
+        num_terminated = int(jnp.sum(terminated).item())
+        num_truncated = int(jnp.sum(truncated).item())
+
+        avg_terminated_length = jnp.sum(episode_lengths * terminated) / jnp.maximum(
+            jnp.sum(terminated), 1
+        )
+
+        avg_truncated_length = jnp.sum(episode_lengths * truncated) / jnp.maximum(
+            jnp.sum(truncated), 1
+        )
+
+        return (
+            next_env_state,
+            next_obs,
+            next_done,
+            TrainingMeasurements(
+                loss=loss,
+                pg_loss=pg_loss,
+                v_loss=v_loss,
+                entropy_loss=entropy_loss,
+                approx_kl=approx_kl,
+                avg_episodic_return=avg_episodic_return,
+                explained_variance=explained_var,
+                num_terminated=num_terminated,
+                num_truncated=num_truncated,
+                avg_terminated_length=avg_terminated_length,
+                avg_truncated_length=avg_truncated_length,
+            ),
+        )
+
+    def _close(self):
+        self.env.close()
+
+    def _save_model(self, model_path: str):
+        self.logger.info("[SAVE]: Saving the final model...")
+
+        params = [
+            vars(self.args),
+            [
+                self.agent_state.params["sensor_params"],
+                self.agent_state.params["actor_params"],
+                self.agent_state.params["critic_params"],
+                self.agent_state.params["feature_extractor_params"],
+            ],
+        ]
+        self.logger.save_final_model(params=params)
+
+    def train(self):
+        """
+        Train the PPO agent for a specified number of iterations
+        (passed through PPOArgs in constructor).
+        Closes the environment at the end of training.
+        """
+        self.logger.info(f"running name: {self.run_name}")
+
+        self.logger.info("[TRAIN]: Resetting environment...")
+        self.logger.log_non_interactive(f"Initial reset started: {time.ctime()}")
+
+        env_state = self.env.reset(seed=self.args.seed)
+        next_obs = _convert_obs_dict_to_array(env_state.observations)
+        next_done = jnp.zeros(self.args.num_envs, dtype=jnp.bool_)
+
+        self.logger.log_non_interactive(f"Initial reset completed: {time.ctime()}")
+
+        global_step = 0
+        start_time = time.time()
+
+        iter_bar = self.logger.progress_bar(range(1, self.args.num_iterations + 1))
+        for iteration in iter_bar:
+            iteration_time_start = time.time()
+
+            env_state, next_obs, next_done, training_measurements = self._step(
+                env_state, next_obs, next_done, iteration=iteration
+            )
+
+            global_step += self.args.num_steps * self.args.num_envs
+            self._log(
+                global_step,
+                self.episode_stats,
+                start_time,
+                iteration_time_start,
+                training_measurements,
+            )
+
+            sps = int(global_step / (time.time() - start_time))
+            remaining_steps = self.args.total_timesteps - global_step
+            eta_seconds = int(remaining_steps / sps) if sps > 0 else 0
+            eta_str = str(datetime.timedelta(seconds=eta_seconds))
+
+            self.logger.log_non_interactive(
+                f"Iteration {iteration}/{self.args.num_iterations} | "
+                f"Step {global_step}/{self.args.total_timesteps} | "
+                f"SPS {sps} | "
+                f"Return {training_measurements.avg_episodic_return:.4f} | "
+                f"ETA {eta_str}"
+            )
+
+        if self.args.save_model:
+            model_path = f"{self.run_dir}/{self.args.exp_name}.cleanrl_model"
+            self._save_model(model_path=model_path)
+
+        self._close()
--- a/src/brittle_star_project/trainers/init.py
+++ b/src/brittle_star_project/trainers/init.py
--- a/src/experiment_logger/README.md
+++ b/src/experiment_logger/README.md
@ -0,0 +1,71 @@
+# Experiment Logger
+
+A standardized, unified interface for logging experiments across multiple backends (WandB, TensorBoard, and Local Disk).
+
+This library is designed to be a standalone package that decouples the logging logic from the core training routines in the `brittle_star_project`.
+
+## Quick Start
+
+The recommended way to use the logger is through the `get_logger()` singleton:
+
+```python
+from experiment_logger import UnifiedLogger, get_logger
+
+# Initialize at the start of your script (e.g., in train.py)
+logger = UnifiedLogger(
+    run_name="my_experiment_run",
+    config={"learning_rate": 3e-4},
+    project_name="MyProject",
+    base_dir="runs",
+    use_wandb=True
+)
+
+# In other files, retrieve the initialized singleton:
+# logger = get_logger()
+
+# Log metrics (Scalar values, numpy scalars, or JAX types)
+logger.log({"loss": 0.5, "accuracy": 0.98}, step=100)
+
+# Standard logging (Mirrored to disk and stdout)
+logger.info("Training started")
+logger.warning("Learning rate is very high")
+
+# Save checkpoints (Automatically synced to WandB as artifacts)
+logger.save_checkpoint(params, step=5000)
+```
+
+## Logger Classes
+
+### `UnifiedLogger`
+
+The full suite for production training. It manages:
+- **WandB**: Syncs metrics and uploads model checkpoints as artifacts.
+- **TensorBoard**: Writes events for local visualization.
+- **Local Disk**: Stores metrics in `metrics.yaml` and textual logs in `run.log`.
+
+### `SimpleLogger`
+
+A zero-dependency fallback that uses standard Python `print()` statements. Use this for standalone testing or minimal environments where you don't need persistent monitoring.
+
+```python
+from experiment_logger import SimpleLogger
+logger = SimpleLogger(run_name="test_run")
+```
+
+## API Features
+
+### `logger.progress_bar(iterable, **kwargs)`
+
+A smart wrapper around `tqdm` that automatically detects its environment. 
+- **Interactive Terminal**: Displays a normal progress bar.
+- **Non-Interactive (HPC)**: Automatically disables the bar to prevent log file bloat in `slurm.out`.
+
+### `logger.log_non_interactive(msg: str)`
+
+Prints a message *only* when running in non-interactive environments. Useful for high-level progress tracking (e.g., "Epoch 5 Complete") without interactive noise.
+
+### `logger.save_checkpoint(params, step, prefix="checkpoint")`
+
+Saves model parameters using Flax serialization.
+- **Local Location**: `runs/<run_name>/checkpoints/`
+- **WandB Logic**: Automatically uploads the `.flax` file as a model artifact for lineage tracking.
--- a/src/experiment_logger/init.py
+++ b/src/experiment_logger/init.py
@ -0,0 +1,21 @@
+"""Unified logging framework for machine learning experiments.
+
+This package provides a unified interface for logging to multiple backends
+(WandB, disk, stdout) simultaneously, ensuring no data loss.
+"""
+
+from experiment_logger.config_utils import load_yaml_config, merge_config_with_cli
+from experiment_logger.unified_logger import UnifiedLogger, get_logger
+from experiment_logger.simple_logger import SimpleLogger
+from experiment_logger.wandb_utils import finish_wandb, init_wandb
+
+__all__ = [
+    "UnifiedLogger",
+    "SimpleLogger",
+    "get_logger",
+    "init_wandb",
+    "finish_wandb",
+    "load_yaml_config",
+    "merge_config_with_cli",
+]
+__version__ = "0.1.0"
--- a/src/experiment_logger/config_utils.py
+++ b/src/experiment_logger/config_utils.py
@ -0,0 +1,147 @@
+"""Configuration utilities for loading YAML configs and merging with CLI args."""
+
+import os
+import sys
+from typing import Dict, Any, Type, TypeVar
+import yaml
+from dataclasses import fields, is_dataclass
+
+from experiment_logger.unified_logger import get_logger
+
+log = get_logger()
+
+T = TypeVar("T")
+
+
+def load_yaml_config(config_path: str) -> Dict[str, Any]:
+    """Load configuration from YAML file."""
+    if not os.path.exists(config_path):
+        raise FileNotFoundError(f"Config file not found: {config_path}")
+
+    with open(config_path, "r") as f:
+        config = yaml.safe_load(f)
+
+    if config is None:
+        return {}
+
+    log.info(f"Loaded configuration from: {config_path}")
+    return config
+
+
+def save_yaml_config(config: Dict[str, Any], config_path: str):
+    """Save configuration to YAML file."""
+    os.makedirs(os.path.dirname(config_path), exist_ok=True)
+
+    with open(config_path, "w") as f:
+        yaml.dump(config, f, default_flow_style=False, indent=2, sort_keys=False)
+
+    log.info(f"Saved configuration to: {config_path}")
+
+
+def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T:
+    """Create dataclass instance from dictionary, handling type conversions."""
+    if not is_dataclass(cls):
+        raise ValueError(f"{cls} is not a dataclass")
+
+    # Get field names and types
+    field_map = {f.name: f for f in fields(cls)}  # type: ignore
+
+    # Filter config to only include valid fields
+    filtered_config: Dict[str, Any] = {}
+    for key, value in config_dict.items():
+        if key in field_map:
+            field = field_map[key]
+            # Handle type conversion if needed
+            try:
+                # Handle None values and optional types
+                if value is None:
+                    filtered_config[key] = None
+                elif hasattr(field.type, "__origin__") and field.type.__origin__ is type(None):
+                    # Optional type (Union[X, None])
+                    filtered_config[key] = value
+                else:
+                    # Try to convert to the expected type
+                    if field.type is bool and isinstance(value, str):
+                        filtered_config[key] = value.lower() in ("true", "1", "yes", "on")
+                    else:
+                        filtered_config[key] = field.type(value) if value is not None else None  # type: ignore
+            except (ValueError, TypeError) as e:
+                log.warning(f"Could not convert {key}={value} to {field.type}: {e}")
+                filtered_config[key] = value
+        else:
+            log.warning(f"Unknown configuration parameter: {key}")
+
+    return cls(**filtered_config)
+
+
+def merge_config_with_cli(config_class: Type[T], config_file: str | None = None) -> T:
+    """Merge YAML config with CLI arguments, with CLI taking precedence.
+
+    Args:
+        config_class: Dataclass type to create
+        config_file: Path to YAML config file (optional)
+
+    Returns:
+        Instance of config_class with merged configuration
+    """
+    # Parse CLI args first to get the default/CLI values
+    import tyro
+
+    # Check if --config is in sys.argv and extract it
+    extracted_config_file = config_file
+    if "--config" in sys.argv:
+        config_idx = sys.argv.index("--config")
+        if config_idx + 1 < len(sys.argv):
+            extracted_config_file = sys.argv[config_idx + 1]
+            # Remove from sys.argv so tyro doesn't see it
+            sys.argv.pop(config_idx)  # Remove --config
+            sys.argv.pop(config_idx)  # Remove config file path
+
+    # Load YAML config if available
+    yaml_config = {}
+    if extracted_config_file and os.path.exists(extracted_config_file):
+        yaml_config = load_yaml_config(extracted_config_file)
+        log.info(f"Merging YAML config from {extracted_config_file} with CLI args")
+    elif extracted_config_file:
+        log.warning(f"Config file not found: {extracted_config_file}, using CLI args only")
+
+    # Create default instance to know what the defaults are
+    default_instance = config_class()
+    default_dict = {f.name: getattr(default_instance, f.name) for f in fields(config_class)}  # type: ignore
+
+    # Parse CLI args
+    cli_instance = tyro.cli(config_class)
+    cli_dict = {f.name: getattr(cli_instance, f.name) for f in fields(config_class)}  # type: ignore
+
+    # Merge configs: YAML as base, CLI overrides non-default values
+    final_config = {}
+
+    for field in fields(config_class):  # type: ignore
+        field_name = field.name
+        default_value = default_dict[field_name]
+        yaml_value = yaml_config.get(field_name, default_value)
+        cli_value = cli_dict[field_name]
+
+        # Use CLI value if it's different from default, otherwise use YAML value
+        if cli_value != default_value:
+            final_config[field_name] = cli_value
+            if yaml_value != default_value and yaml_value != cli_value:
+                log.info(f"CLI override: {field_name}={cli_value} (YAML had {yaml_value})")
+        else:
+            final_config[field_name] = yaml_value
+            if yaml_value != default_value:
+                log.info(f"YAML config: {field_name}={yaml_value}")
+
+    return config_class(**final_config)
+
+
+def print_config(config: Any, title: str = "Configuration"):
+    """Pretty print configuration."""
+    log.info(f"{title}:")
+    if is_dataclass(config):
+        for field in fields(config):
+            value = getattr(config, field.name)
+            log.info(f"  {field.name}: {value}")
+    else:
+        for key, value in vars(config).items():
+            log.info(f"  {key}: {value}")
--- a/src/experiment_logger/simple_logger.py
+++ b/src/experiment_logger/simple_logger.py
@ -0,0 +1,83 @@
+"""Simple terminal logger for running without external backends.
+
+This is used for standalone package usage where WandB or TensorBoard are not desired.
+It preserves the same API as UnifiedLogger but simply prints to stdout.
+"""
+
+import logging
+from typing import Any, Dict, Optional
+
+
+class SimpleLogger:
+    """Simple logger that implements the UnifiedLogger interface via print statements."""
+
+    def __init__(
+        self,
+        run_name: str = "simple_run",
+        config: Optional[Dict[str, Any]] = None,
+        project_name: str = "none",
+        entity: Optional[str] = None,
+        base_dir: str = "runs",
+        use_wandb: bool = False,
+        save_code: bool = False,
+        log_level: int = logging.INFO,
+        _set_as_global: bool = False,
+    ):
+        self.is_interactive = True
+        self.run_name = run_name
+        self.config = config or {}
+        print(f"[INIT] SimpleLogger initialized for run: {run_name}")
+
+    def set_level(self, level: int):
+        pass
+
+    def log_non_interactive(self, msg: str, *args, **kwargs):
+        """In SimpleLogger, we just print everything as we assume interactive use."""
+        self.info(msg, *args, **kwargs)
+
+    def progress_bar(self, iterable=None, *args, **kwargs):
+        """Standard tqdm wrapper that falls back to range if tqdm is missing."""
+        try:
+            import tqdm
+
+            return tqdm.tqdm(iterable, *args, **kwargs)
+        except ImportError:
+            return iterable
+
+    def info(self, msg: str, *args, **kwargs):
+        print(f"[INFO] {msg}")
+
+    def warning(self, msg: str, *args, **kwargs):
+        print(f"[WARNING] {msg}")
+
+    def error(self, msg: str, *args, **kwargs):
+        print(f"[ERROR] {msg}")
+
+    def debug(self, msg: str, *args, **kwargs):
+        print(f"[DEBUG] {msg}")
+
+    def log(self, metrics: Dict[str, Any], step: Optional[int] = None, commit: bool = True):
+        step_str = f"Step {step}" if step is not None else "Log"
+        metric_str = ", ".join(f"{k}: {v}" for k, v in metrics.items())
+        print(f"[{step_str}] {metric_str}")
+
+    def save_checkpoint(
+        self,
+        params: Any,
+        step: int,
+        prefix: str = "checkpoint",
+        metadata: Optional[Dict[str, Any]] = None,
+    ):
+        print(f"[SAVE] Checkpoint '{prefix}' would be saved at step {step} (SimpleLogger: No-Op)")
+
+    def save_final_model(self, params: Any, metadata: Optional[Dict[str, Any]] = None):
+        print("[SAVE] Final model would be saved (SimpleLogger: No-Op)")
+
+    def finish(self):
+        print(f"[FINISH] SimpleLogger finished for run: {self.run_name}")
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.finish()
--- a/src/experiment_logger/unified_logger.py
+++ b/src/experiment_logger/unified_logger.py
@ -0,0 +1,385 @@
+"""Unified logger that writes to multiple backends simultaneously.
+
+This logger ensures all experimental data is preserved by writing to:
+1. Weights & Biases (when available)
+2. Local disk (JSON files, model checkpoints, run.log)
+3. stdout (for real-time monitoring)
+"""
+
+import datetime
+import logging
+import subprocess
+import yaml
+import sys
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import flax
+import jax.numpy as jnp
+import numpy as np
+
+from experiment_logger.wandb_utils import finish_wandb, init_wandb
+
+# Global singleton storage
+_global_logger = None
+
+
+def get_logger() -> "UnifiedLogger":
+    """Retrieve the global UnifiedLogger. If not initialized, fallback to auto-initialization."""
+    global _global_logger
+    if _global_logger is None:
+        try:
+            commit_hash = (
+                subprocess.check_output(
+                    ["git", "rev-parse", "--short", "HEAD"], stderr=subprocess.STDOUT
+                )
+                .decode("utf-8")
+                .strip()
+            )
+        except Exception:
+            commit_hash = "unknown"
+
+        timestamp_str = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        generic_name = f"{timestamp_str}_{commit_hash}_brittle_star"
+
+        # Initialize generic fallback logger without WandB
+        _global_logger = UnifiedLogger(
+            run_name=generic_name,
+            config={"auto_initialized": True},
+            use_wandb=False,
+            _set_as_global=False,  # Prevent recursive call inside __init__
+        )
+        _global_logger.warning(f"UnifiedLogger auto-initialized with name: {generic_name}")
+
+    return _global_logger
+
+
+class UnifiedLogger:
+    """Unified logger for scientific experiments with redundant backup."""
+
+    def __init__(
+        self,
+        run_name: str,
+        config: Dict[str, Any],
+        project_name: str = "PPO-Modularity",
+        entity: Optional[str] = None,
+        base_dir: str = "runs",
+        use_wandb: bool = True,
+        save_code: bool = True,
+        log_level: int = logging.INFO,
+        _set_as_global: bool = True,
+    ):
+        """Initialize the unified logger.
+
+        Args:
+            run_name: Unique name for this run
+            config: Configuration dictionary with hyperparameters
+            project_name: WandB project name
+            entity: WandB entity (team/user name)
+            base_dir: Base directory for local storage
+            use_wandb: Whether to use WandB logging
+            save_code: Whether to save code to WandB
+            _set_as_global: Internal flag to override the global singleton
+        """
+        self.run_name = run_name
+        self.config = config
+        self.use_wandb = use_wandb
+        self.wandb_available = False
+        self.wandb_run = None
+        self.is_interactive = sys.stdout.isatty()
+
+        # Setup local storage
+        self.run_dir = Path(base_dir) / run_name
+        self.run_dir.mkdir(parents=True, exist_ok=True)
+
+        self.checkpoints_dir = self.run_dir / "checkpoints"
+        self.checkpoints_dir.mkdir(exist_ok=True)
+
+        self.metrics_dir = self.run_dir / "metrics"
+        self.metrics_dir.mkdir(exist_ok=True)
+
+        self.config_file = self.run_dir / "config.yaml"
+
+        # Setup standard Python logging mirror
+        self.text_log_file = self.run_dir / "run.log"
+        self._text_logger = logging.getLogger(f"UnifiedLogger_{self.run_name}")
+        self._text_logger.setLevel(log_level)
+        self._text_logger.propagate = False
+
+        # Avoid duplicate handlers if re-instantiated
+        if not self._text_logger.handlers:
+            fh = logging.FileHandler(self.text_log_file)
+            ch = logging.StreamHandler()
+
+            formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
+            fh.setFormatter(formatter)
+            ch.setFormatter(formatter)
+
+            self._text_logger.addHandler(fh)
+            self._text_logger.addHandler(ch)
+
+        # Set as global singleton
+        global _global_logger
+        if _set_as_global:
+            _global_logger = self
+
+        # Save config to disk
+        self._save_config()
+
+        # Setup TensorBoard
+        self.writer = None
+        try:
+            from torch.utils.tensorboard import SummaryWriter
+
+            self.writer = SummaryWriter(self.run_dir)
+            self.info("TensorBoard SummaryWriter initialized.")
+        except ImportError:
+            self.warning("tensorboard not installed. Skipping SummaryWriter.")
+
+        # Initialize WandB if requested
+        if self.use_wandb:
+            self._init_wandb(project_name, entity, save_code)
+
+        # Initialize metrics storage
+        self.metrics_buffer: List[Dict[str, Any]] = []
+        self.step_counter = 0
+
+        self.info(f"Initialized UnifiedLogger for run: {run_name}")
+        self.info(f"Local storage: {self.run_dir.absolute()}")
+        self.info(f"WandB logging: {self.wandb_available}")
+
+    def set_level(self, level: int):
+        """Dynamically update the verbosity of the stdout/text logger."""
+        self._text_logger.setLevel(level)
+
+    def log_non_interactive(self, msg: str, *args, **kwargs):
+        """Log an info message only if running in a non-interactive environment."""
+        if not self.is_interactive:
+            self.info(msg, *args, **kwargs)
+
+    def progress_bar(self, iterable=None, *args, **kwargs):
+        """Wrapper around tqdm that automatically disables in non-interactive environments."""
+        import tqdm
+
+        kwargs.setdefault("disable", not self.is_interactive)
+        return tqdm.tqdm(iterable, *args, **kwargs)
+
+    def info(self, msg: str, *args, **kwargs):
+        """Log an info message to stdout and disk."""
+        self._text_logger.info(msg, *args, **kwargs)
+
+    def warning(self, msg: str, *args, **kwargs):
+        """Log a warning message to stdout and disk."""
+        self._text_logger.warning(msg, *args, **kwargs)
+
+    def error(self, msg: str, *args, **kwargs):
+        """Log an error message to stdout and disk."""
+        self._text_logger.error(msg, *args, **kwargs)
+
+    def debug(self, msg: str, *args, **kwargs):
+        """Log a debug message to stdout and disk."""
+        self._text_logger.debug(msg, *args, **kwargs)
+
+    def _init_wandb(self, project_name: str, entity: Optional[str], save_code: bool):
+        """Initialize Weights & Biases logging."""
+        self.wandb_run = init_wandb(
+            project=project_name,
+            entity=entity,
+            name=self.run_name,
+            config=self.config,
+            save_code=save_code,
+            resume="allow",
+        )
+        self.wandb_available = self.wandb_run is not None
+
+    def _save_config(self):
+        """Save configuration to disk."""
+        try:
+            with open(self.config_file, "w") as f:
+                yaml.dump(self.config, f, default_flow_style=False, indent=2, sort_keys=False)
+            self.info(f"Config saved to {self.config_file}")
+        except Exception as e:
+            self.error(f"Error saving config: {e}")
+
+    def log(self, metrics: Dict[str, Any], step: Optional[int] = None, commit: bool = True):
+        """Log metrics to all backends.
+
+        Args:
+            metrics: Dictionary of metric name -> value
+            step: Global step counter (auto-incremented if None)
+            commit: Whether to commit to WandB immediately
+        """
+        if step is None:
+            step = self.step_counter
+            self.step_counter += 1
+
+        # Add timestamp
+        metrics_with_metadata = {
+            "step": step,
+            "timestamp": time.time(),
+            **metrics,
+        }
+
+        # Log to stdout
+        self._log_to_stdout(metrics_with_metadata)
+
+        # Log to WandB
+        if self.wandb_run is not None:
+            try:
+                self.wandb_run.log(metrics, step=step, commit=commit)
+            except Exception as e:
+                self.warning(f"WandB logging failed: {e}")
+
+        # Log to TensorBoard
+        if self.writer is not None:
+            for k, v in metrics.items():
+                if isinstance(v, (int, float, np.floating, np.integer)):
+                    self.writer.add_scalar(k, v, step)
+                elif hasattr(v, "item"):
+                    self.writer.add_scalar(k, v.item(), step)
+                elif isinstance(v, (np.ndarray, jnp.ndarray)) and v.size == 1:
+                    self.writer.add_scalar(k, v.item(), step)
+
+        # Buffer for disk storage
+        self.metrics_buffer.append(metrics_with_metadata)
+
+        # Periodically flush to disk
+        if len(self.metrics_buffer) >= 100:
+            self._flush_metrics()
+
+    def _log_to_stdout(self, metrics: Dict[str, Any]):
+        """Log metrics to stdout for real-time monitoring."""
+        step = metrics.get("step", "?")
+        metric_str = ", ".join(
+            f"{k}={v:.6f}" if isinstance(v, (float, np.floating)) else f"{k}={v}"
+            for k, v in metrics.items()
+            if k not in ["step", "timestamp"]
+        )
+        self.info(f"[Step {step}] {metric_str}")
+
+    def _flush_metrics(self):
+        """Flush buffered metrics to disk."""
+        if not self.metrics_buffer:
+            return
+
+        try:
+            metrics_file = self.metrics_dir / "metrics.yaml"
+            with open(metrics_file, "a") as f:
+                for metric in self.metrics_buffer:
+                    # Convert numpy/jax types to native Python types for YAML serialization
+                    serializable_metric = {}
+                    for k, v in metric.items():
+                        if hasattr(v, "item"):  # numpy/jax scalar
+                            serializable_metric[k] = v.item()
+                        elif isinstance(v, (np.ndarray, jnp.ndarray)):
+                            serializable_metric[k] = v.tolist()
+                        else:
+                            serializable_metric[k] = v
+                    f.write("---\n")
+                    yaml.dump(serializable_metric, f, default_flow_style=False)
+            self.metrics_buffer.clear()
+        except Exception as e:
+            self.error(f"Error flushing metrics: {e}")
+
+    def save_checkpoint(
+        self,
+        params: Any,
+        step: int,
+        prefix: str = "checkpoint",
+        metadata: Optional[Dict[str, Any]] = None,
+    ):
+        """Save model checkpoint to disk and optionally to WandB."""
+        checkpoint_name = f"{prefix}_step_{step}.flax"
+        checkpoint_path = self.checkpoints_dir / checkpoint_name
+
+        try:
+            # Save to disk using Flax serialization
+            with open(checkpoint_path, "wb") as f:
+                f.write(flax.serialization.to_bytes(params))
+
+            # Save metadata if provided
+            if metadata:
+                metadata_path = self.checkpoints_dir / f"{prefix}_step_{step}_metadata.yaml"
+                with open(metadata_path, "w") as f:
+                    yaml.dump(metadata, f, default_flow_style=False)
+
+            self.info(f"Checkpoint saved: {checkpoint_path}")
+
+            # Log to WandB as artifact
+            if self.wandb_run is not None:
+                try:
+                    import wandb
+
+                    artifact = wandb.Artifact(
+                        name=f"{self.run_name}_{prefix}",
+                        type="model",
+                        metadata=metadata or {},
+                    )
+                    artifact.add_file(str(checkpoint_path))
+                    if metadata:
+                        artifact.add_file(str(metadata_path))
+                    self.wandb_run.log_artifact(artifact)
+                    self.info("Checkpoint uploaded to WandB")
+                except Exception as e:
+                    self.warning(f"Could not upload checkpoint to WandB: {e}")
+
+        except Exception as e:
+            self.error(f"Error saving checkpoint: {e}")
+
+    def save_final_model(self, params: Any, metadata: Optional[Dict[str, Any]] = None):
+        """Save the final trained model."""
+        final_model_path = self.run_dir / "final_model.flax"
+
+        try:
+            with open(final_model_path, "wb") as f:
+                f.write(flax.serialization.to_bytes(params))
+
+            if metadata:
+                metadata_path = self.run_dir / "final_model_metadata.yaml"
+                with open(metadata_path, "w") as f:
+                    yaml.dump(metadata, f, default_flow_style=False)
+
+            self.info(f"Final model saved: {final_model_path}")
+
+            # Log to WandB
+            if self.wandb_run is not None:
+                try:
+                    import wandb
+
+                    artifact = wandb.Artifact(
+                        name=f"{self.run_name}_final_model",
+                        type="model",
+                        metadata=metadata or {},
+                    )
+                    artifact.add_file(str(final_model_path))
+                    if metadata:
+                        artifact.add_file(str(metadata_path))
+                    self.wandb_run.log_artifact(artifact)
+                except Exception as e:
+                    self.warning(f"Could not upload final model to WandB: {e}")
+
+        except Exception as e:
+            self.error(f"Error saving final model: {e}")
+
+    def finish(self):
+        """Finalize logging and cleanup."""
+        # Flush remaining metrics
+        self._flush_metrics()
+
+        if self.writer is not None:
+            self.writer.close()
+
+        self.info(f"Run complete. Results saved to: {self.run_dir.absolute()}")
+
+        # Finish WandB run
+        if self.wandb_available:
+            finish_wandb()
+
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit."""
+        self.finish()
--- a/src/experiment_logger/wandb_utils.py
+++ b/src/experiment_logger/wandb_utils.py
@ -0,0 +1,91 @@
+"""Centralized WandB initialization utilities."""
+
+import logging
+from typing import Any, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+
+def init_wandb(
+    project: str,
+    config: Dict[str, Any],
+    name: Optional[str] = None,
+    entity: Optional[str] = None,
+    sync_tensorboard: bool = False,
+    save_code: bool = True,
+    resume: str = "allow",
+    **kwargs,
+):
+    """Initialize WandB with standardized settings.
+
+    This function provides a centralized way to initialize WandB across different
+    scripts, ensuring consistent configuration and error handling.
+
+    Args:
+        project: WandB project name
+        config: Configuration dictionary to log
+        name: Run name (auto-generated if None)
+        entity: WandB entity (team/user name)
+        sync_tensorboard: Whether to sync tensorboard logs
+        save_code: Whether to save code snapshots
+        resume: Resume strategy ("allow", "must", "never", "auto")
+        **kwargs: Additional arguments to pass to wandb.init()
+
+    Returns:
+        wandb.Run object if successful, None otherwise
+    """
+    try:
+        import wandb
+        import os
+        import sys
+
+        # Robust HPC checking: check for API key
+        has_key = os.environ.get("WANDB_API_KEY") is not None
+        if not has_key:
+            try:
+                # Check if logged in locally via settings/netrc
+                has_key = wandb.setup().settings.api_key is not None
+            except Exception:
+                pass
+
+        is_interactive = sys.stdout.isatty()
+
+        if not has_key and not is_interactive and os.environ.get("WANDB_MODE") != "offline":
+            logger.warning(
+                "WANDB_API_KEY not found and environment is non-interactive. "
+                "Switching to offline mode."
+            )
+            sync_path = f"runs/{name}" if name else "runs"
+            logger.warning(f"WandB is offline. Use 'wandb sync {sync_path}' to upload logs later.")
+            os.environ["WANDB_MODE"] = "offline"
+
+        run = wandb.init(
+            project=project,
+            entity=entity,
+            name=name,
+            config=config,
+            sync_tensorboard=sync_tensorboard,
+            save_code=save_code,
+            resume=resume,
+            **kwargs,
+        )
+        logger.info(f"WandB initialized successfully for project '{project}', run '{run.name}'")
+        return run
+    except ImportError:
+        logger.warning("WandB not installed. Skipping WandB initialization.")
+        return None
+    except Exception as e:
+        logger.error(f"Failed to initialize WandB: {e}")
+        return None
+
+
+def finish_wandb():
+    """Safely finish the current WandB run."""
+    try:
+        import wandb
+
+        if wandb.run is not None:
+            wandb.finish()
+            logger.info("WandB run finished successfully")
+    except Exception as e:
+        logger.warning(f"Error finishing WandB run: {e}")
--- a/src/main.py
+++ b/src/main.py
@ -1,4 +0,0 @@
-import jax
-
-if __name__ == "__main__":
-    print(jax.devices())
--- a/src/train.py
+++ b/src/train.py
@ -1,328 +0,0 @@
-import random
-import time
-from dataclasses import asdict
-from functools import partial
-from typing import Callable
-
-import flax
-import jax
-import jax.numpy as jnp
-import matplotlib.pyplot as plt
-import numpy as np
-import optax
-import torch
-import tqdm
-import tyro
-from flax.training.train_state import TrainState
-from torch.utils.tensorboard import SummaryWriter
-
-from brittle_star_project.dataclasses import PPOArgs
-from brittle_star_project.dataclasses.EpisodeStatistics import EpisodeStatistics
-from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
-from brittle_star_project.rl import Actor, AgentParams, Critic, Network, Storage
-from ppo import PPO
-
-
-def convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray:
-    return jax.vmap(lambda o: jnp.concatenate([v.flatten() for v in o.values() if v.size > 0]))(
-        obs_dict
-    )
-
-
-def make_env(config_path: str | None, num_envs: int) -> Callable:
-    def thunk():
-        if config_path is None:
-            return BrittleStarJaxEnvWrapper.default(num_envs=num_envs)
-        return BrittleStarJaxEnvWrapper.from_config(config_path, num_envs=num_envs)
-
-    return thunk
-
-
-def train(args: PPOArgs):
-    args.batch_size = args.num_envs * args.num_steps
-    args.minibatch_size = args.batch_size // args.num_minibatches
-    args.num_iterations = args.total_timesteps // args.batch_size
-    run_name = f"{args.exp_name}__seed_{args.seed}__{int(time.time())}"
-    print(f"running name: {run_name}")
-
-    if args.track:
-        import wandb
-
-        wandb.init(
-            project=args.wandb_project_name,
-            entity=args.wandb_entity,
-            sync_tensorboard=True,
-            config=vars(args),
-            name=run_name,
-            save_code=True,
-        )
-
-    writer = SummaryWriter(f"runs/{run_name}")
-    writer.add_text(
-        "hyperparameters",
-        "|param|value|\n|---|---|\n" + "\n".join(f"|{k}|{v}|" for k, v in vars(args).items()),
-    )
-
-    random.seed(args.seed)
-    np.random.seed(args.seed)
-    key = jax.random.PRNGKey(args.seed)
-    key, network_key, actor_key, critic_key, critic_network_key = jax.random.split(key, 5)
-
-    torch.backends.cudnn.deterministic = args.torch_deterministic
-    device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu")
-    print(f"Running on device: {device}")
-
-    print("Creating the environment...")
-    env = make_env(config_path=args.config_path, num_envs=args.num_envs)()
-    print(f"Environment: {env}")
-
-    episode_stats = EpisodeStatistics(
-        episode_returns=jnp.zeros(args.num_envs, dtype=jnp.float32),
-        episode_lengths=jnp.zeros(args.num_envs, dtype=jnp.int32),
-        returned_episode_returns=jnp.zeros(args.num_envs, jnp.float32),
-        returned_episode_lengths=jnp.zeros(args.num_envs, dtype=jnp.int32),
-    )
-
-    def step_env_wrapped(episode_stats: EpisodeStatistics, env_state, action):
-        next_env_state = env.step(env_state, action)
-
-        # Extract per-environment signals from the state object
-        reward = next_env_state.reward  # (num_envs,)
-        terminated = next_env_state.terminated  # (num_envs,)
-        truncated = next_env_state.truncated  # (num_envs,)
-        done = terminated | truncated  # (num_envs,)
-
-        new_episode_return = episode_stats.episode_returns + reward
-        new_episode_length = episode_stats.episode_lengths + 1
-
-        episode_stats = episode_stats.replace(
-            episode_returns=new_episode_return * (1 - done),
-            episode_lengths=new_episode_length * (1 - done),
-            returned_episode_returns=jnp.where(
-                done, new_episode_return, episode_stats.returned_episode_returns
-            ),
-            returned_episode_lengths=jnp.where(
-                done, new_episode_length, episode_stats.returned_episode_lengths
-            ),
-        )
-        return (
-            episode_stats,
-            next_env_state,
-            (convert_obs_dict_to_array(next_env_state.observations), reward, done),
-        )
-
-    def linear_schedule(count):
-        frac = 1.0 - (count // (args.num_minibatches * args.update_epochs)) / args.num_iterations
-        return args.learning_rate * frac
-
-    print("Initializing the models...")
-    network = Network()
-    critic_network = Network()
-    actor = Actor(action_dim=env.single_action_space.shape[0])  # continuous actions for MJX
-    critic = Critic()
-
-    sample_obs = jnp.concatenate(
-        [
-            v.flatten()
-            for v in env.single_observation_space.sample(rng=jax.random.PRNGKey(0)).values()
-            if v.size > 0
-        ]
-    )
-    network_params = network.init(network_key, sample_obs)
-    critic_network_params = critic_network.init(critic_network_key, sample_obs)
-    actor_params = actor.init(actor_key, network.apply(network_params, sample_obs))
-    critic_params = critic.init(critic_key, critic_network.apply(critic_network_params, sample_obs))
-
-    agent_state = TrainState.create(
-        apply_fn=None,
-        params=asdict(
-            AgentParams(network_params, actor_params, critic_params, critic_network_params)
-        ),
-        tx=optax.chain(
-            optax.clip_by_global_norm(args.max_grad_norm),
-            optax.inject_hyperparams(optax.adam)(
-                learning_rate=linear_schedule if args.anneal_lr else args.learning_rate, eps=1e-5
-            ),
-        ),
-    )
-
-    network.apply = jax.jit(network.apply)
-    critic_network.apply = jax.jit(critic_network.apply)
-    actor.apply = jax.jit(actor.apply)
-    critic.apply = jax.jit(critic.apply)
-    ppo_instance = PPO(args, network, actor, critic, critic_network)
-
-    @jax.jit
-    def get_action_and_value_noise(
-        agent_state: TrainState,
-        next_obs: jnp.ndarray,
-        key: jax.random.PRNGKey,
-    ):
-        hidden = network.apply(agent_state.params["network_params"], next_obs)
-        # Continuous actions: sample from a Gaussian parameterized by the actor
-        mean, log_std = actor.apply(agent_state.params["actor_params"], hidden)
-        key, subkey = jax.random.split(key)
-        noise = jax.random.normal(subkey, shape=mean.shape)
-        std = jnp.exp(log_std)
-        action = mean + noise * std
-        logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
-        value = critic.apply(agent_state.params["critic_params"], hidden)
-        return action, logprob, value.squeeze(-1), key
-
-    @jax.jit
-    def compute_gae_once(carry, inp, gamma, gae_lambda):
-        advantages = carry
-        nextdone, nextvalues, curvalues, reward = inp
-        nextnonterminal = 1.0 - nextdone
-        delta = reward + gamma * nextvalues * nextnonterminal - curvalues
-        advantages = delta + gamma * gae_lambda * nextnonterminal * advantages
-        return advantages, advantages
-
-    @jax.jit
-    def compute_gae(agent_state, next_obs, next_done, storage):
-        next_value = critic.apply(
-            agent_state.params["critic_params"],
-            network.apply(agent_state.params["network_params"], next_obs),
-        ).squeeze(-1)
-
-        advantages = jnp.zeros((args.num_envs,))
-        dones = jnp.concatenate([storage.dones, next_done[None, :]], axis=0)
-        values = jnp.concatenate([storage.values, next_value[None, :]], axis=0)
-        _, advantages = jax.lax.scan(
-            partial(compute_gae_once, gamma=args.gamma, gae_lambda=args.gae_lambda),
-            advantages,
-            (dones[1:], values[1:], values[:-1], storage.rewards),
-            reverse=True,
-        )
-        return storage.replace(advantages=advantages, returns=advantages + storage.values)
-
-    # --- Main training loop ---
-    global_step = 0
-    start_time = time.time()
-
-    # Reset once to get initial state
-    print("Resetting the environment...")
-    next_env_state = env.reset(seed=args.seed)
-    next_obs = convert_obs_dict_to_array(next_env_state.observations)
-    next_done = jnp.zeros(args.num_envs, dtype=jnp.bool_)
-
-    def step_once(carry, _, env_step_fn):
-        agent_state, episode_stats, obs, done, key, env_state = carry
-        action, logprob, value, key = get_action_and_value_noise(agent_state, obs, key)
-
-        episode_stats, env_state, (next_obs, reward, next_done) = env_step_fn(
-            episode_stats, env_state, action
-        )
-
-        storage = Storage(
-            obs=obs,
-            actions=action,
-            logprobs=logprob,
-            dones=done,
-            values=value,
-            rewards=reward,
-            returns=jnp.zeros_like(reward),
-            advantages=jnp.zeros_like(reward),
-        )
-        return (agent_state, episode_stats, next_obs, next_done, key, env_state), storage
-
-    def rollout(
-        agent_state, episode_stats, next_obs, next_done, key, env_state, step_once_fn, max_steps
-    ):
-        (agent_state, episode_stats, next_obs, next_done, key, env_state), storage = jax.lax.scan(
-            step_once_fn,
-            (agent_state, episode_stats, next_obs, next_done, key, env_state),
-            (),
-            max_steps,
-        )
-        return agent_state, episode_stats, next_obs, next_done, storage, key, env_state
-
-    rollout = partial(
-        rollout,
-        step_once_fn=partial(step_once, env_step_fn=step_env_wrapped),
-        max_steps=args.num_steps,
-    )
-
-    print("Starting training...")
-    iters_bar = tqdm.tqdm(range(1, args.num_iterations + 1))
-    losses = []
-    for _ in iters_bar:
-        iteration_time_start = time.time()
-
-        agent_state, episode_stats, next_obs, next_done, storage, key, next_env_state = rollout(
-            agent_state, episode_stats, next_obs, next_done, key, next_env_state
-        )
-
-        global_step += args.num_steps * args.num_envs
-        storage = compute_gae(agent_state, next_obs, next_done, storage)
-        agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key = ppo_instance.update_ppo(
-            agent_state, storage, key
-        )
-
-        losses.append(jnp.mean(loss))
-
-        avg_episodic_return = np.mean(jax.device_get(episode_stats.returned_episode_returns))
-        iters_bar.set_postfix_str(
-            f"global_step={global_step}, avg_episodic_return={avg_episodic_return}"
-        )
-
-        writer.add_scalar("charts/avg_episodic_return", avg_episodic_return, global_step)
-        writer.add_scalar(
-            "charts/avg_episodic_length",
-            np.mean(jax.device_get(episode_stats.returned_episode_lengths)),
-            global_step,
-        )
-        writer.add_scalar(
-            "charts/learning_rate",
-            agent_state.opt_state[1].hyperparams["learning_rate"].item(),
-            global_step,
-        )
-        writer.add_scalar("losses/value_loss", v_loss[-1, -1].item(), global_step)
-        writer.add_scalar("losses/policy_loss", pg_loss[-1, -1].item(), global_step)
-        writer.add_scalar("losses/entropy", entropy_loss[-1, -1].item(), global_step)
-        writer.add_scalar("losses/approx_kl", approx_kl[-1, -1].item(), global_step)
-        writer.add_scalar("losses/loss", loss[-1, -1].item(), global_step)
-
-        # iters_bar.set_postfix_str(f"SPS: {int(global_step / (time.time() - start_time))}")
-
-        writer.add_scalar("charts/SPS", int(global_step / (time.time() - start_time)), global_step)
-        writer.add_scalar(
-            "charts/SPS_update",
-            int(args.num_envs * args.num_steps / (time.time() - iteration_time_start)),
-            global_step,
-        )
-
-    if args.save_model:
-        model_path = f"runs/{run_name}/{args.exp_name}.cleanrl_model"
-        with open(model_path, "wb") as f:
-            f.write(
-                flax.serialization.to_bytes(
-                    [
-                        vars(args),
-                        [
-                            agent_state.params["network_params"],
-                            agent_state.params["actor_params"],
-                            agent_state.params["critic_params"],
-                        ],
-                    ]
-                )
-            )
-        print(f"model saved to {model_path}")
-
-    env.close()
-    writer.close()
-
-    print("Saving loss plot...")
-    plt.plot(losses)
-    plt.title("PPO Loss, mean over minibatches")
-    plt.savefig(f"runs/{run_name}/{args.exp_name}_losses.png")
-    plt.close()
-
-
-def main() -> None:
-    args = tyro.cli(PPOArgs)
-    train(args)
-
-
-if __name__ == "__main__":
-    main()
--- a/tests/test_config.py
+++ b/tests/test_config.py
@ -0,0 +1,38 @@
+"""Tests for YAML config loading."""
+
+import sys
+from pathlib import Path
+
+import pytest
+
+# Ensure src is on the path when running from the project root
+sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
+
+CONFIGS_DIR = Path(__file__).parent.parent / "configs"
+
+
+class TestYamlConfig:
+    def test_load_yaml_config(self):
+        from experiment_logger.config_utils import load_yaml_config
+
+        config = load_yaml_config(str(CONFIGS_DIR / "default_ppo.yaml"))
+        assert isinstance(config, dict)
+        assert "total_timesteps" in config
+        assert "learning_rate" in config
+
+    def test_load_dev_test_config(self):
+        from experiment_logger.config_utils import load_yaml_config
+
+        config = load_yaml_config(str(CONFIGS_DIR / "dev_test.yaml"))
+        assert config["total_timesteps"] == 100000
+
+    def test_missing_config_raises(self):
+        from experiment_logger.config_utils import load_yaml_config
+
+        with pytest.raises(FileNotFoundError):
+            load_yaml_config("nonexistent.yaml")
+
+    def test_merge_config_with_cli_is_callable(self):
+        from experiment_logger.config_utils import merge_config_with_cli
+
+        assert callable(merge_config_with_cli)
--- a/uv.lock
+++ b/uv.lock
@ -11,7 +11,7 @@ resolution-markers = [
 [[package]]
 name = "2026sel3-project"
 version = "0.1.0"
-source = { virtual = "." }
+source = { editable = "." }
 dependencies = [
    { name = "biorobot" },
    { name = "cleanrl" },
@ -22,14 +22,23 @@ dependencies = [
    { name = "jax" },
    { name = "matplotlib" },
    { name = "mediapy" },
+    { name = "mujoco-warp" },
+    { name = "numpy" },
    { name = "optax" },
+    { name = "protobuf" },
    { name = "pyopengl" },
    { name = "pyopengl-accelerate" },
+    { name = "pyyaml" },
+    { name = "torch" },
    { name = "tyro" },
    { name = "wandb" },
+    { name = "warp-lang" },
 ]

 [package.optional-dependencies]
+analysis = [
+    { name = "tensorboard" },
+]
 cuda = [
    { name = "jax", extra = ["cuda13"] },
 ]
@ -53,13 +62,20 @@ requires-dist = [
    { name = "jax", extras = ["cuda13"], marker = "extra == 'cuda'", specifier = "==0.9.0.1" },
    { name = "matplotlib", specifier = "==3.10.8" },
    { name = "mediapy", specifier = "==1.2.6" },
+    { name = "mujoco-warp" },
+    { name = "numpy", specifier = ">=2.0.0" },
    { name = "optax", specifier = ">=0.2.6" },
+    { name = "protobuf", specifier = ">=5.0.0" },
    { name = "pyopengl", specifier = ">=3.1.10" },
    { name = "pyopengl-accelerate", specifier = ">=3.1.10" },
+    { name = "pyyaml", specifier = ">=6.0" },
+    { name = "tensorboard", marker = "extra == 'analysis'" },
+    { name = "torch", specifier = ">=2.4.0" },
    { name = "tyro", specifier = ">=1.0.10" },
    { name = "wandb", specifier = "==0.24.2" },
+    { name = "warp-lang" },
 ]
-provides-extras = ["cuda"]
+provides-extras = ["cuda", "analysis"]

 [package.metadata.requires-dev]
 dev = [
@ -1262,6 +1278,22 @@ wheels = [
    { url = "https://files.pythonhosted.org/packages/c2/7c/ad82beb7c4c9186d9fbef4799109d799692d70276bb1b3ee18a0674170d8/mujoco_mjx-3.6.0-py3-none-any.whl", hash = "sha256:c81000af0653f162b76009f48c153e9e6d19bfa8febe851e12466c81cbb7336a", size = 7013366, upload-time = "2026-03-11T01:46:19.148Z" },
 ]

+[[package]]
+name = "mujoco-warp"
+version = "3.6.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "absl-py" },
+    { name = "etils", extra = ["epath"] },
+    { name = "mujoco" },
+    { name = "numpy" },
+    { name = "warp-lang" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/0e/de/b853418268e9777cad2792ee3a145c8397e3d4517d136499645847ffd7f2/mujoco_warp-3.6.0.tar.gz", hash = "sha256:3c4111a4e13dc61268ddac52593ac5032c05a7d80f0c5e3c98bf5881e32b5d06", size = 1887269, upload-time = "2026-03-11T01:11:44.343Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/a3/b5/06c1e23c0cc4a06da268aa5f0fe05348d89c9ecbd3ecfb4c6d2b14ea23b2/mujoco_warp-3.6.0-py3-none-any.whl", hash = "sha256:371a405b186332cbfaa9630aabf35967405ef847cb7ea7c018ec67426a2ea160", size = 1965960, upload-time = "2026-03-11T01:11:42.527Z" },
+]
+
 [[package]]
 name = "mypy-extensions"
 version = "1.1.0"
@ -1630,7 +1662,7 @@ name = "pexpect"
 version = "4.9.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "ptyprocess" },
+    { name = "ptyprocess", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" }
 wheels = [
@ -2428,6 +2460,20 @@ wheels = [
    { url = "https://files.pythonhosted.org/packages/3a/9a/f3919d7ee7ba99dabf0aac7e299c6c328f5eae94f9f6b28c76005f882d5d/wandb-0.24.2-py3-none-win_arm64.whl", hash = "sha256:b42614b99f8b9af69f88c15a84283a973c8cd5750e9c4752aa3ce21f13dbac9a", size = 20268261, upload-time = "2026-02-05T00:12:14.353Z" },
 ]

+[[package]]
+name = "warp-lang"
+version = "1.12.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "numpy" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/7d/15/fadf3e3ba5c1c907530c20c98402aaef792da74bbbe382c848cef6e5affe/warp_lang-1.12.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:c78c3701d5cad86c30ef5017410d294ec46a396bb0d502ee1c98743494f3a62f", size = 24168341, upload-time = "2026-03-06T19:42:16.333Z" },
+    { url = "https://files.pythonhosted.org/packages/98/13/deab9dbae5c6aa753ac8ea1d3b1f85d20c5bab7bdebd8916ce242fbe1f0b/warp_lang-1.12.0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:a1436f60a1881cd94f787e751a83fc0987626be2d3e2b4e74c64a6947c6d1266", size = 136485344, upload-time = "2026-03-06T19:43:02.427Z" },
+    { url = "https://files.pythonhosted.org/packages/45/ce/9f5c57cac849edaba2f3335cb649b7019b09195b3af02221258482254559/warp_lang-1.12.0-py3-none-manylinux_2_34_aarch64.whl", hash = "sha256:a2d6decba693aba5b828573c4414fd6a3f4c4a934db9c322736ef2b3fa99fe76", size = 137735580, upload-time = "2026-03-06T19:44:22.279Z" },
+    { url = "https://files.pythonhosted.org/packages/7e/3f/1ddc888fe769447ae33915a9567a9dd7467e1fc7fc8010d39e01b339667f/warp_lang-1.12.0-py3-none-win_amd64.whl", hash = "sha256:697248edd2f1e2952f50e3db33b214af76173641a8894aacc467bed6dc247f8a", size = 119793582, upload-time = "2026-03-06T19:45:37.288Z" },
+]
+
 [[package]]
 name = "wcwidth"
 version = "0.6.0"