Merge branch 'dev' into docs/reward_and_mlp-design

.devcontainer/devcontainer.json

@@ -34,8 +34,7 @@
     },
     "remoteUser": "vscode",
     "runArgs": [
-        "--gpus",
+        "--device", "nvidia.com/gpu=all"
-        "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": {}
     }
-}
+}

.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>"

.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

.pre-commit-config.yaml

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

.vscode/settings.json

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

README.md

@@ -15,3 +15,7 @@ example command:
 ```bash
 uv run src/train.py --model_name my_model --epochs 50 --batch_size 32
 ```
+
+## HPC
+
+See **[docs/HPC.md](docs/HPC.md)** for the full guide, including environment setup, cluster selection, interactive debugging, and job submission.

configs/example.json

@@ -0,0 +1,8 @@
+{
+  "morphology": {
+    "num_arms": 2,
+    "num_segments_per_arm": 4,
+    "use_p_control": true,
+    "use_torque_control": false
+  }
+}

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

configs/production_training.yaml

@@ -0,0 +1,24 @@
+# Full PPO training config for Brittle Star (HPC Production)
+exp_name: "production_training"
+seed: 1
+track: true
+capture_video: true
+save_model: true
+checkpoint_frequency: 100  # not yet implemented in train.py but here for future use
+
+# Scaling for HPC (using A100 GPU slices)
+num_envs: 128
+total_timesteps: 10000000
+num_steps: 128
+num_minibatches: 4
+update_epochs: 4
+
+# Algorithm
+learning_rate: 2.5e-4
+anneal_lr: true
+gamma: 0.99
+gae_lambda: 0.95
+clip_coef: 0.1
+ent_coef: 0.01
+vf_coef: 0.5
+cuda: true

docs/CONTRIBUTING.md

@@ -32,4 +32,4 @@ 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.

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.

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
+```

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
-```

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

env/hpc/requirements.txt

@@ -0,0 +1,19 @@
+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
+tyro>=1.0.10
+wandb==0.24.2
+torch>=2.4.0

pyproject.toml

@@ -12,6 +12,10 @@ 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",
@@ -19,12 +23,16 @@ dependencies = [
     "pyopengl-accelerate>=3.1.10",
     "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 = [

ruff.toml

@@ -1,4 +1,5 @@
 line-length = 100
+exclude = ["wandb"]
 
 [lint]
 extend-select = [
@@ -394,4 +395,5 @@ extend-ignore = [
   # "PLW1404",  # implicit-str-concat
 ]
 
-
+[lint.per-file-ignores]
+"__init__.py" = ["F401"]

scripts/analysis/README.md

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

scripts/analysis/explore_tensorboard.py

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

scripts/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__).parent.parent
+
+
+def normalise(name: str) -> str:
+    """Normalise a PyPI package name for comparison."""
+    return re.sub(r"[-_.]+", "-", name).lower()
+
+
+def pkg_name(dep: str) -> str:
+    """Extract the bare package name from a PEP 508 dependency string."""
+    return re.split(r"[\[=><~!;]", dep)[0].strip()
+
+
+def main() -> None:
+    import tomllib
+
+    modules_path = ROOT / "env" / "hpc" / "modules.txt"
+    if not modules_path.exists():
+        print(f"Error: {modules_path} not found.", file=sys.stderr)
+        sys.exit(1)
+
+    # Read normalized module names from base modules only
+    # Library modules (like PyTorch) are kept in requirements for portability
+    module_names = [
+        normalise(line.split()[0].split("/")[0])
+        for line in modules_path.read_text().splitlines()
+        if line.strip() and not line.startswith("#")
+    ]
+
+    pyproject_path = ROOT / "pyproject.toml"
+    with pyproject_path.open("rb") as f:
+        data = tomllib.load(f)
+
+    # Collect all dependencies, merging 'cuda' extras into base dependencies
+    dep_dict: dict[str, str] = {}
+    for dep in data.get("project", {}).get("dependencies", []):
+        dep_dict[normalise(pkg_name(dep))] = dep
+
+    # Add cuda extras (takes precedence for HPC)
+    optional_deps = data.get("project", {}).get("optional-dependencies", {})
+    for group in ["cuda"]:
+        for dep in optional_deps.get(group, []):
+            dep_dict[normalise(pkg_name(dep))] = dep
+
+    deps = list(dep_dict.values())
+
+    final_deps: list[str] = []
+    print("Checking dependencies against HPC module list...", file=sys.stderr)
+    for dep in deps:
+        name = normalise(pkg_name(dep))
+        # Smart check: if the package name is a substring of any loaded module name
+        # (e.g. 'torch' in 'pytorch', 'scipy' in 'scipy-bundle')
+        if any(name in mod for mod in module_names):
+            print(f"  [skip – module provider found] {dep}", file=sys.stderr)
+            continue
+
+        final_deps.append(dep)
+        print(f"  [pip] {dep}", file=sys.stderr)
+
+    hpc_dir = ROOT / "env" / "hpc"
+    output_path = hpc_dir / "requirements.txt"
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    output_path.write_text("\n".join(final_deps) + "\n")
+    print(f"\nWrote {len(final_deps)} requirement(s) to {output_path}", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    main()

scripts/hpc/install.sh

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

scripts/hpc/train.pbs

@@ -0,0 +1,63 @@
+# 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"
+
+python src/train.py \
+    --env-config-path configs/hpc/smoke_test.yaml \
+    --run-dir "$SCRATCH_RUNDIR"
+
+echo ">>> Staging out results to $DATA_RUNDIR..."
+cp -r "$SCRATCH_RUNDIR/." "$DATA_RUNDIR/"
+
+echo ">>> Done"

scripts/simulate.py

@@ -4,17 +4,15 @@ import argparse
 from pathlib import Path
 
 from brittle_star_project import (
-    ArenaConfig,
     Backend,
     BrittleStarEnv,
     BrittleStarEnvFactory,
-    EnvConfig,
+    SimulationConfig,
-    MorphologyConfig,
+    simulate_policy,
-    Task,
 )
+from brittle_star_project.environment import from_json
 from brittle_star_project.rl import RLModel  # imports concrete models via rl.__init__
 from brittle_star_project.rl.base import get_rl_model_registry
-from brittle_star_project.renderer import SimulationConfig, simulate_policy
 
 MODEL_BY_NAME = get_rl_model_registry()
 MODEL_OPTIONS = sorted(MODEL_BY_NAME)
@@ -35,9 +33,9 @@ def parse_args() -> argparse.Namespace:
         help="Which model class to instantiate when --model is omitted.",
     )
     p.add_argument(
-        "--task",
+        "--backend",
-        choices=[t.value for t in Task],
+        choices=[b for b in Backend],
-        default=Task.DIRECTED_LOCOMOTION.value,
+        default=Backend.MJX,
     )
     p.add_argument("--seed", type=int, default=None)
     return p.parse_args()
@@ -46,14 +44,11 @@ def parse_args() -> argparse.Namespace:
 def main() -> None:
     args = parse_args()
 
+    morphology_cfg, arena_cfg, env_cfg = from_json("../configs/test.json")
+
     # ======= ENVIRONMENT SETUP =======
 
-    backend = Backend.MJC
+    backend = args.backend
-    task = Task(args.task)
-
-    morphology_cfg = MorphologyConfig()
-    arena_cfg = ArenaConfig(attach_target=(task == Task.DIRECTED_LOCOMOTION))
-    env_cfg = EnvConfig(task=task)
 
     factory = BrittleStarEnvFactory()
     raw_env = factory.create_environment(backend, morphology_cfg, arena_cfg, env_cfg)

scripts/train.py

@@ -0,0 +1,75 @@
+import subprocess
+import time
+
+import torch
+import tyro
+import yaml
+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
+
+
+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(log: bool = True) -> PPOArgs:
+    temp_args = tyro.cli(PPOArgs)
+
+    if temp_args.hyperparameter_config_path is not None:
+        if log:
+            print(f"Loading hyperparameter config from {temp_args.hyperparameter_config_path}")
+
+        with open(temp_args.hyperparameter_config_path, "r") as f:
+            config = yaml.safe_load(f)
+            if config:
+                # parse PPOArgs with defaults from yaml.
+                for key, value in config.items():
+                    if hasattr(temp_args, key):
+                        setattr(temp_args, key, value)
+
+                # Reparse CLI to ensure they OVERRIDE the yaml
+                args = tyro.cli(PPOArgs, default=temp_args)
+    else:
+        if log:
+            print("No hyperparameter config provided, using default config")
+
+        args = temp_args
+    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)
+
+    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()

src/MLPs/mlps.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):
+# semi generic so we can easily make a config for it in experiments
-    """
+class GenericDenseLayersWithActivation(nn.Module):
-    Dummy model only used for testing purposes
+    layer_sizes: Sequence[int] = field(default_factory=lambda: [64, 64])
-
+    activation: Callable = nn.tanh
-    inspired by: https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/ppo_atari_envpool_xla_jax_scan.py
-    """
-
-    hidden_dim: int = 195
 
     @nn.compact
     def __call__(self, x):
-        x = nn.Dense(self.hidden_dim, kernel_init=orthogonal(np.sqrt(2)), bias_init=constant(0.0))(
+        for size in self.layer_sizes:
-            x
+            x = nn.Dense(size, kernel_init=orthogonal(jnp.sqrt(2)))(x)
-        )
+            x = self.activation(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)
         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,9 +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
+    feature_extractor_params: flax.core.FrozenDict
 
 
 @jax.tree_util.register_dataclass

src/__init__.py

@@ -0,0 +1,25 @@
+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",
+]

src/brittle_star_project/__init__.py

@@ -2,6 +2,7 @@ from .environment.env_types import Backend, Task
 from .environment.env_config import ArenaConfig, EnvConfig, MorphologyConfig
 from .environment.factory import BrittleStarEnvFactory
 from .environment.env_wrapper import BrittleStarEnv
+from .render import simulate_policy, SimulationConfig, ControlPolicy
 
 __all__ = [
     "ArenaConfig",
@@ -11,4 +12,7 @@ __all__ = [
     "EnvConfig",
     "MorphologyConfig",
     "Task",
+    "simulate_policy",
+    "SimulationConfig",
+    "ControlPolicy",
 ]

src/brittle_star_project/dataclasses/PPOArgs.py

@@ -1,15 +1,30 @@
 from dataclasses import dataclass
 
+import jax
 
+
+@jax.tree_util.register_dataclass
 @dataclass
 class PPOArgs:
     """
     source: https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/ppo_atari_envpool_xla_jax_scan.py
     """
 
+    # path to environment config file, if None, use default config
+    env_config_path: str | None = None
+
+    # path to hyperparameter config file (yaml), if None, use default config
+    hyperparameter_config_path: str | None = None
+
     # the name of this experiment
     exp_name: str = "brittle_star_ppo"
 
+    # the directory to save the experiment results
+    run_dir: str | None = None
+
+    # how often to save checkpoints (0 to disable)
+    checkpoint_frequency: int = 0
+
     # seed of the experiment
     seed: int = 1
 
@@ -41,8 +56,6 @@ class PPOArgs:
     hf_entity: str = ""
 
     # ==== Algorithm specific dataclasses ====
-    # the id of the environment
-    env_id: str = ""  # todo
 
     # total timesteps of the experiments
     total_timesteps: int = 10000000
@@ -51,7 +64,7 @@ class PPOArgs:
     learning_rate: float = 2.5e-4
 
     # the number of parallel game environments
-    num_envs: int = 16
+    num_envs: int = 100
 
     # the number of steps to run in each environment per policy rollout
     num_steps: int = 128

src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py

@@ -8,6 +8,7 @@ from brittle_star_project import (
     ArenaConfig,
     Backend,
 )
+from brittle_star_project.environment import from_file
 
 
 class BrittleStarJaxEnvWrapper:
@@ -76,3 +77,21 @@ class BrittleStarJaxEnvWrapper:
         return BrittleStarJaxEnvWrapper(
             morphology, arena, env_config, num_envs=num_envs, backend=backend
         )
+
+    @staticmethod
+    def from_config(
+        config_path: str, num_envs: int, backend: Backend = Backend.MJX
+    ) -> "BrittleStarJaxEnvWrapper":
+        morphology_cfg, arena_cfg, env_cfg = from_file(config_path)
+        return BrittleStarJaxEnvWrapper(
+            morphology_cfg, arena_cfg, env_cfg, num_envs=num_envs, backend=backend
+        )
+
+    def __str__(self):
+        morphology_str = str(self._morphology)
+        arena_str = str(self._arena)
+        env_config_str = str(self._env_config)
+        return (
+            f"BrittleStarJaxEnvWrapper(backend={self._backend}, num_envs={self._num_envs}, "
+            + f"morphology={morphology_str}, arena={arena_str}, env_config={env_config_str})"
+        )

src/brittle_star_project/environment/__init__.py

@@ -1,4 +1,4 @@
-from .env_config import ArenaConfig, EnvConfig, MorphologyConfig
+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,4 +12,5 @@ __all__ = [
     "BrittleStarEnv",
     "StepResult",
     "BrittleStarEnvFactory",
+    "from_file",
 ]

src/brittle_star_project/environment/env_config.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
+import json
 
 from .env_types import Task
 
@@ -48,6 +49,18 @@ class EnvConfig:
     # Per docs in upstream env config: integer factors of 200.
     light_perlin_noise_scale: int = 0
 
-    @staticmethod
+
-    def from_json(path: str) -> EnvConfig:
+def from_file(path: str) -> tuple[MorphologyConfig, ArenaConfig, EnvConfig]:
-        pass
+    """Load configurations from a JSON or YAML file."""
+    with open(path, "r") as f:
+        if path.endswith(".yaml") or path.endswith(".yml"):
+            import yaml
+
+            config_dict = yaml.safe_load(f)
+        else:
+            config_dict = json.load(f)
+
+        morphology = MorphologyConfig(**config_dict.get("morphology", {}))
+        arena = ArenaConfig(**config_dict.get("arena", {}))
+        env = EnvConfig(**config_dict.get("env", {}))
+        return morphology, arena, env

src/brittle_star_project/render/__init__.py

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

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",
-]

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)

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

src/brittle_star_project/trainers/PPOTrainer.py

@@ -0,0 +1,536 @@
+import datetime
+import random
+import sys
+import time
+from dataclasses import asdict, dataclass
+from functools import partial
+from typing import Any
+
+import flax
+import jax
+import jax.numpy as jnp
+import numpy as np
+import optax
+import tqdm
+from flax.training.train_state import TrainState
+from torch.utils.tensorboard import SummaryWriter
+
+from brittle_star_project.dataclasses import EpisodeStatistics, PPOArgs
+from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
+from MLPs.mlps import (
+    Actor,
+    AgentParams,
+    GenericDenseLayersWithActivation,
+    OneDenseLayerMLP,
+    Storage,
+)
+from ppo import PPO
+
+
+@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
+    )
+
+
+# removed jit: used in _rollout_jit, so will be compiled with _rollout_jit
+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
+    )
+
+    # 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_critic)
+    return action, logprob, value.squeeze(-1), key
+
+
+# removed jit: used in _rollout_jit, so will be compiled with _rollout_jit
+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
+
+
+# removed jit: used in _rollout_jit, so will be compiled with _rollout_jit
+def _step_env_wrapped(episode_stats, env_state, action, env_step_fn):
+    next_env_state = env_step_fn(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),
+    )
+
+
+# jit applied in wrapper method self._rollout_jit using partial
+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
+
+
+# removed jit: used in _compute_gae_jit, so will be compiled with _compute_gae_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
+
+
+# jit applied on partial-wrapped wrapper method self._compute_gae_jit
+def _compute_gae_jit(
+    agent_state, storage, next_obs, next_done, gamma, gae_lambda, num_envs, sensor, critic
+):
+    next_value = critic.apply(
+        agent_state.params["critic_params"],
+        sensor.apply(agent_state.params["sensor_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 LossInfo:
+    # todo: better typing
+    loss: Any
+    pg_loss: Any
+    v_loss: Any
+    entropy_loss: Any
+    approx_kl: Any
+    avg_episodic_return: 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.writer = SummaryWriter(self.run_dir)
+
+        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,
+                sensor=self.sensor,
+                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, log: bool = True):
+        if log:
+            print(f"[RANDOM]: Setting random seed to {self.args.seed}")
+
+        random.seed(self.args.seed)
+        np.random.seed(self.args.seed)
+
+    def _init_agent(self, log: bool = True):
+        if log:
+            print("[AGENT]: Initializing agent...")
+
+        sensor = GenericDenseLayersWithActivation()
+        feature_extractor = GenericDenseLayersWithActivation()
+        actor = Actor(
+            action_dim=self.env.single_action_space.shape[0]
+        )  # continuous actions for MJX
+        critic = OneDenseLayerMLP()
+        # messenger = OneDenseLayerMLP()
+        return sensor, feature_extractor, actor, critic
+
+    def _init_agent_state(self, log: bool = True) -> TrainState:
+        if log:
+            print("[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, log: bool = True) -> EpisodeStatistics:
+        if log:
+            print("[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[Storage, ...]:
+        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,
+        loss_info,
+    ):
+
+        self.writer.add_scalar(
+            "charts/avg_episodic_return", loss_info.avg_episodic_return, global_step
+        )
+        self.writer.add_scalar(
+            "charts/avg_episodic_length",
+            np.mean(jax.device_get(episode_stats.returned_episode_lengths)),
+            global_step,
+        )
+        self.writer.add_scalar(
+            "charts/learning_rate",
+            self.agent_state.opt_state[1].hyperparams["learning_rate"].item(),
+            global_step,
+        )
+        self.writer.add_scalar("losses/value_loss", loss_info.v_loss[-1, -1].item(), global_step)
+        self.writer.add_scalar("losses/policy_loss", loss_info.pg_loss[-1, -1].item(), global_step)
+        self.writer.add_scalar("losses/entropy", loss_info.entropy_loss[-1, -1].item(), global_step)
+        self.writer.add_scalar("losses/approx_kl", loss_info.approx_kl[-1, -1].item(), global_step)
+        self.writer.add_scalar("losses/loss", loss_info.loss[-1, -1].item(), global_step)
+        self.writer.add_scalar(
+            "charts/SPS", int(global_step / (time.time() - start_time)), global_step
+        )
+        self.writer.add_scalar(
+            "charts/SPS_update",
+            int(self.args.num_envs * self.args.num_steps / (time.time() - iteration_time_start)),
+            global_step,
+        )
+
+    def _step(
+        self, env_state, next_obs, next_done, is_tty: bool, iteration: int, log: bool = True
+    ) -> tuple:
+        if log and not is_tty and iteration == 1:
+            print(f">>> [HPC] Starting first rollout (JIT): {time.ctime()}", flush=True)
+
+        (
+            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 log and not is_tty and iteration == 1:
+            print(f">>> [HPC] First rollout completed: {time.ctime()}", flush=True)
+
+        storage = self._compute_gae(storage, next_obs, next_done)
+
+        if log and not is_tty and iteration == 1:
+            print(f">>> [HPC] Starting first PPO update (JIT): {time.ctime()}", flush=True)
+
+        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 log and not is_tty and iteration == 1:
+            print(f">>> [HPC] First PPO update completed: {time.ctime()}", flush=True)
+
+        avg_episodic_return = float(
+            jnp.mean(jax.device_get(self.episode_stats.returned_episode_returns)).item()
+        )
+
+        return (
+            next_env_state,
+            next_obs,
+            next_done,
+            LossInfo(
+                loss=loss,
+                pg_loss=pg_loss,
+                v_loss=v_loss,
+                entropy_loss=entropy_loss,
+                approx_kl=approx_kl,
+                avg_episodic_return=avg_episodic_return,
+            ),
+        )
+
+    def _close(self):
+        self.env.close()
+        self.writer.close()
+
+    def _save_model(self, model_path: str, log: bool = True):
+        if log:
+            print(f"[SAVE]: Saving the model to: {model_path}...")
+
+        with open(model_path, "wb") as f:
+            f.write(
+                flax.serialization.to_bytes(
+                    [
+                        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"],
+                        ],
+                    ]
+                )
+            )
+
+    def train(self, log: bool = True):
+        """
+        Train the PPO agent for a specified number of iterations
+        (passed through PPOArgs in constructor).
+        Closes the environment at the end of training.
+        """
+        if log:
+            print(f"running name: {self.run_name}")
+
+        is_tty = sys.stdout.isatty()
+        if log:
+            print("[TRAIN]: Resetting environment...")
+
+            if not is_tty:
+                print(f">>> [HPC] Initial reset started: {time.ctime()}", flush=True)
+
+        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_)
+
+        if log and not is_tty:
+            print(f">>> [HPC] Initial reset completed: {time.ctime()}", flush=True)
+
+        global_step = 0
+        start_time = time.time()
+
+        if self.args.track:
+            import wandb
+
+            if log:
+                print("[TRAIN]: Initializing Weights and Biases...")
+
+            wandb.init(
+                project=self.args.wandb_project_name,
+                entity=self.args.wandb_entity,
+                sync_tensorboard=True,
+                config=vars(self.args),
+                name=self.run_name,
+                save_code=True,
+            )
+
+        if log:
+            print("[TRAIN]: Adding hyperparameters to TensorBoard...")
+
+        self.writer.add_text(
+            "hyperparameters",
+            "|param|value|\n|---|---|\n"
+            + "\n".join(f"|{k}|{v}|" for k, v in vars(self.args).items()),
+        )
+
+        iter_bar = tqdm.tqdm(
+            range(1, self.args.num_iterations + 1),
+            disable=not is_tty,
+        )
+        for iteration in iter_bar:
+            iteration_time_start = time.time()
+
+            env_state, next_obs, next_done, loss_info = self._step(
+                env_state, next_obs, next_done, is_tty=is_tty, iteration=iteration
+            )
+
+            global_step += self.args.num_steps * self.args.num_envs
+            self._log(global_step, self.episode_stats, start_time, iteration_time_start, loss_info)
+
+            if log and not is_tty:
+                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))
+
+                print(
+                    f"Iteration {iteration}/{self.args.num_iterations} | "
+                    f"Step {global_step}/{self.args.total_timesteps} | "
+                    f"SPS {sps} | "
+                    f"Return {loss_info.avg_episodic_return:.4f} | "
+                    f"ETA {eta_str}",
+                    flush=True,
+                )
+
+        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()

src/main.py

@@ -1,4 +0,0 @@
-import jax
-
-if __name__ == "__main__":
-    print(jax.devices())

src/ppo.py

@@ -0,0 +1,159 @@
+from functools import partial
+
+import flax
+import jax
+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, sensor, actor, critic, feature_extractor, message_passer=None):
+        self.args = args
+
+        if not message_passer:
+            message_passer = identity
+
+        self.ppo_loss_grad_fn = jax.value_and_grad(
+            partial(
+                ppo_loss,
+                args=args,
+                sensor_apply=sensor.apply,
+                actor_apply=actor.apply,
+                critic_apply=critic.apply,
+                feature_extractor_apply=feature_extractor.apply,
+                message_passer=message_passer,
+            ),
+            has_aux=True,
+        )
+
+    # This PPO class should be initialized only once,
+    # or this function will need to recompile
+    @partial(jax.jit, static_argnums=0)
+    def update_ppo(self, agent_state, storage, key):
+        args = self.args
+        ppo_loss_grad_fn = self.ppo_loss_grad_fn
+
+        def update_epoch(carry, _):
+            agent_state, key = carry
+            key, subkey = jax.random.split(key)
+
+            def flatten(x):
+                return x.reshape((-1,) + x.shape[2:])
+
+            def convert_data(x):
+                x = jax.random.permutation(subkey, x)
+                return jnp.reshape(x, (args.num_minibatches, -1) + x.shape[1:])
+
+            flatten_storage = jax.tree.map(flatten, storage)
+            shuffled_storage = jax.tree.map(convert_data, flatten_storage)
+
+            def update_minibatch(agent_state, minibatch):
+                (loss, (pg_loss, v_loss, entropy_loss, approx_kl)), grads = ppo_loss_grad_fn(
+                    agent_state.params,
+                    minibatch.obs,
+                    minibatch.actions,
+                    minibatch.logprobs,
+                    minibatch.advantages,
+                    minibatch.returns,
+                )
+                agent_state = agent_state.apply_gradients(grads=grads)
+                return agent_state, (
+                    loss,
+                    pg_loss,
+                    v_loss,
+                    entropy_loss,
+                    approx_kl,
+                    grads,
+                )
+
+            agent_state, metrics = jax.lax.scan(update_minibatch, agent_state, shuffled_storage)
+            return (agent_state, key), metrics
+
+        (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads) = jax.lax.scan(
+            update_epoch, (agent_state, key), (), length=args.update_epochs
+        )
+        return agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key
+
+
+"""
+Should be ok to use partial here, since the references to network,
+actor and critic should not change at runtime
+The cost of seperating concerns is to somehow pass these values
+that are now not in the same scope
+"""
+
+
+@partial(jax.jit, static_argnums=(0, 1, 2, 3, 4))
+def get_action_and_value(
+    sensor_apply,
+    actor_apply,
+    message_passer,
+    critic_apply,
+    feature_extractor_apply,
+    params: flax.core.FrozenDict,
+    x: jnp.ndarray,
+    action: jnp.ndarray,
+):
+    hidden_sensor = sensor_apply(params["sensor_params"], x)
+    hidden_critic = feature_extractor_apply(params["feature_extractor_params"], x)
+    hidden_sensor = message_passer(hidden_sensor)
+    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)
+    entropy = (0.5 + 0.5 * jnp.log(2 * jnp.pi) + log_std).sum(-1)
+    value = critic_apply(params["critic_params"], hidden_critic).squeeze(-1)
+
+    return logprob, entropy, value
+
+
+def ppo_loss(
+    params,
+    x,
+    a,
+    logp,
+    mb_advantages,
+    mb_returns,
+    args,
+    sensor_apply,
+    actor_apply,
+    message_passer,
+    critic_apply,
+    feature_extractor_apply,
+):
+    newlogprob, entropy, newvalue = get_action_and_value(
+        sensor_apply,
+        actor_apply,
+        message_passer,
+        critic_apply,
+        feature_extractor_apply,
+        params,
+        x,
+        a,
+    )
+    logratio = newlogprob - logp
+    ratio = jnp.exp(logratio)
+    approx_kl = ((ratio - 1) - logratio).mean()
+
+    if args.norm_adv:
+        mb_advantages = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8)
+
+    pg_loss1 = -mb_advantages * ratio
+    pg_loss2 = -mb_advantages * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef)
+    pg_loss = jnp.maximum(pg_loss1, pg_loss2).mean()
+    v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean()
+    entropy_loss = entropy.mean()
+    loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef
+    return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl))
+
+
+def identity(hidden):
+    """
+    Used for seamless jax integration,
+    avoids having branching inside jitted function,
+    used as message_passer in case it is not given,
+    (in case of centralized lvl)
+    """
+
+    return hidden

src/train.py

@@ -1,377 +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 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 Network, Actor, Critic, AgentParams, Storage
-
-
-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(num_envs: int) -> Callable:
-    def thunk():
-        return BrittleStarJaxEnvWrapper.default(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 = jax.random.split(key, 4)
-
-    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(num_envs=args.num_envs)()
-
-    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()
-    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)
-    actor_params = actor.init(actor_key, network.apply(network_params, sample_obs))
-    critic_params = critic.init(critic_key, network.apply(network_params, sample_obs))
-
-    agent_state = TrainState.create(
-        apply_fn=None,
-        params=asdict(AgentParams(network_params, actor_params, critic_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)
-    actor.apply = jax.jit(actor.apply)
-    critic.apply = jax.jit(critic.apply)
-
-    @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 get_action_and_value(
-        params: flax.core.FrozenDict,
-        x: jnp.ndarray,
-        action: np.ndarray,
-    ):
-        hidden = network.apply(params["network_params"], x)
-        mean, log_std = actor.apply(params["actor_params"], hidden)
-        std = jnp.exp(log_std)
-        logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
-        entropy = (0.5 + 0.5 * jnp.log(2 * jnp.pi) + log_std).sum(-1)
-        value = critic.apply(params["critic_params"], hidden).squeeze(-1)
-        return logprob, entropy, value
-
-    @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)
-
-    def ppo_loss(params, x, a, logp, mb_advantages, mb_returns):
-        newlogprob, entropy, newvalue = get_action_and_value(params, x, a)
-        logratio = newlogprob - logp
-        ratio = jnp.exp(logratio)
-        approx_kl = ((ratio - 1) - logratio).mean()
-
-        if args.norm_adv:
-            mb_advantages = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8)
-
-        pg_loss1 = -mb_advantages * ratio
-        pg_loss2 = -mb_advantages * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef)
-        pg_loss = jnp.maximum(pg_loss1, pg_loss2).mean()
-        v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean()
-        entropy_loss = entropy.mean()
-        loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef
-        return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl))
-
-    ppo_loss_grad_fn = jax.value_and_grad(ppo_loss, has_aux=True)
-
-    @jax.jit
-    def update_ppo(agent_state, storage, key):
-        def update_epoch(carry, _):
-            agent_state, key = carry
-            key, subkey = jax.random.split(key)
-
-            def flatten(x):
-                return x.reshape((-1,) + x.shape[2:])
-
-            def convert_data(x):
-                x = jax.random.permutation(subkey, x)
-                return jnp.reshape(x, (args.num_minibatches, -1) + x.shape[1:])
-
-            flatten_storage = jax.tree.map(flatten, storage)
-            shuffled_storage = jax.tree.map(convert_data, flatten_storage)
-
-            def update_minibatch(agent_state, minibatch):
-                (loss, (pg_loss, v_loss, entropy_loss, approx_kl)), grads = ppo_loss_grad_fn(
-                    agent_state.params,
-                    minibatch.obs,
-                    minibatch.actions,
-                    minibatch.logprobs,
-                    minibatch.advantages,
-                    minibatch.returns,
-                )
-                agent_state = agent_state.apply_gradients(grads=grads)
-                return agent_state, (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads)
-
-            agent_state, metrics = jax.lax.scan(update_minibatch, agent_state, shuffled_storage)
-            return (agent_state, key), metrics
-
-        (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads) = jax.lax.scan(
-            update_epoch, (agent_state, key), (), length=args.update_epochs
-        )
-        return agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key
-
-    # --- 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))
-    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 = update_ppo(
-            agent_state, storage, key
-        )
-
-        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()
-
-
-def main() -> None:
-    args = tyro.cli(PPOArgs)
-    train(args)
-
-
-if __name__ == "__main__":
-    main()

uv.lock

@@ -22,14 +22,22 @@ dependencies = [
     { name = "jax" },
     { name = "matplotlib" },
     { name = "mediapy" },
+    { name = "mujoco-warp" },
+    { name = "numpy" },
     { name = "optax" },
+    { name = "protobuf" },
     { name = "pyopengl" },
     { name = "pyopengl-accelerate" },
+    { name = "torch" },
     { name = "tyro" },
     { name = "wandb" },
+    { name = "warp-lang" },
 ]
 
 [package.optional-dependencies]
+analysis = [
+    { name = "tensorboard" },
+]
 cuda = [
     { name = "jax", extra = ["cuda13"] },
 ]
@@ -53,13 +61,19 @@ 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 = "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 +1276,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 +1660,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 +2458,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"