Merge branch 'dev' into feat/wandb-logging
This commit is contained in:
commit
512272d6ab
SSH key fingerprint:
SHA256:u/h/LVoqKF1Iz02uOyxe6hcjmoZASCGV2HM0TG9ZMoU
41 changed files with 1517 additions and 792 deletions
|
|
@ -45,4 +45,4 @@
|
|||
"features": {
|
||||
"ghcr.io/devcontainers/features/common-utils:1": {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
38
.github/workflows/update_hpc_requirements.yml
vendored
Normal file
38
.github/workflows/update_hpc_requirements.yml
vendored
Normal file
|
|
@ -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>"
|
||||
4
.gitignore
vendored
4
.gitignore
vendored
|
|
@ -1,6 +1,7 @@
|
|||
# Model files
|
||||
artifacts/*
|
||||
runs/*
|
||||
wandb/
|
||||
|
||||
# Experiment tracking
|
||||
wandb/
|
||||
|
|
@ -375,7 +376,6 @@ celerybeat.pid
|
|||
# Environments
|
||||
.env
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
|
|
@ -475,7 +475,6 @@ tags
|
|||
[Ll]ib
|
||||
[Ll]ib64
|
||||
[Ll]ocal
|
||||
[Ss]cripts
|
||||
pyvenv.cfg
|
||||
.venv
|
||||
pip-selfcheck.json
|
||||
|
|
@ -522,3 +521,4 @@ Icon
|
|||
Network Trash Folder
|
||||
Temporary Items
|
||||
.apdisk
|
||||
*.pdf
|
||||
|
|
|
|||
|
|
@ -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']
|
||||
7
.vscode/settings.json
vendored
Normal file
7
.vscode/settings.json
vendored
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
{
|
||||
"python.testing.pytestArgs": [
|
||||
"tests"
|
||||
],
|
||||
"python.testing.unittestEnabled": false,
|
||||
"python.testing.pytestEnabled": true
|
||||
}
|
||||
26
README.md
26
README.md
|
|
@ -1,12 +1,10 @@
|
|||
# Brittle Star
|
||||
|
||||
Reinforcement learning research on brittle star locomotion using PPO.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Installation
|
||||
|
||||
Set up the environment using UV:
|
||||
To set up the UV module, you can run the following command:
|
||||
|
||||
```bash
|
||||
uv sync --frozen
|
||||
|
|
@ -33,7 +31,7 @@ uv sync --frozen
|
|||
|
||||
### Training
|
||||
|
||||
Run training with your configuration:
|
||||
example command:
|
||||
|
||||
```bash
|
||||
uv run python src/train.py
|
||||
|
|
@ -60,22 +58,6 @@ The training script uses a unified logging framework that:
|
|||
|
||||
All experiment data is preserved locally, even if WandB is unavailable.
|
||||
|
||||
## Project Structure
|
||||
## HPC
|
||||
|
||||
```
|
||||
src/brittle_star_project/ # Core library (reusable components)
|
||||
├── logging/ # Unified logging framework
|
||||
├── environment/ # Environment wrappers
|
||||
├── rl/ # RL algorithms and models
|
||||
└── dataclasses/ # Configuration dataclasses
|
||||
|
||||
configs/ # Training configurations
|
||||
runs/ # Training outputs (checkpoints, metrics)
|
||||
```
|
||||
|
||||
## For Researchers
|
||||
|
||||
**Important:** Do not commit your personal WandB credentials to the repository.
|
||||
Instead, create your own config file (e.g., `configs/yourname.yaml`) and add it to `.gitignore` if needed.
|
||||
|
||||
See [configs/README.md](configs/README.md) for more details on configuration management.
|
||||
See **[docs/HPC.md](docs/HPC.md)** for the full guide, including environment setup, cluster selection, interactive debugging, and job submission.
|
||||
|
|
|
|||
8
configs/example.json
Normal file
8
configs/example.json
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
{
|
||||
"morphology": {
|
||||
"num_arms": 2,
|
||||
"num_segments_per_arm": 4,
|
||||
"use_p_control": true,
|
||||
"use_torque_control": false
|
||||
}
|
||||
}
|
||||
11
configs/hpc/smoke_test.yaml
Normal file
11
configs/hpc/smoke_test.yaml
Normal file
|
|
@ -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
|
||||
|
|
@ -4,34 +4,33 @@
|
|||
# with wandb logging enabled.
|
||||
|
||||
# Experiment settings
|
||||
exp_name: "brittle_star_production"
|
||||
exp_name: "brittle_star_production_training"
|
||||
seed: 42
|
||||
|
||||
# Tracking settings - IMPORTANT: Set your own wandb_entity!
|
||||
# Tracking
|
||||
track: true
|
||||
capture_video: false
|
||||
wandb_project_name: "PPO-Modularity"
|
||||
wandb_entity: "SEL3-2026-Groep-4" # ⚠️ SET THIS TO YOUR WANDB USERNAME OR TEAM
|
||||
wandb_entity: "SEL3-2026-Groep-4"
|
||||
|
||||
# Model saving
|
||||
save_model: true
|
||||
checkpoint_frequency: 100 # Save checkpoint every 100 iterations
|
||||
|
||||
# Environment settings
|
||||
num_envs: 32 # Increased for production
|
||||
num_envs: 512
|
||||
|
||||
# Training hyperparameters - Production scale
|
||||
total_timesteps: 50000000 # 50M timesteps for full training
|
||||
learning_rate: 0.00025
|
||||
num_steps: 256 # Longer rollouts
|
||||
# Training hyperparameters
|
||||
total_timesteps: 50000000
|
||||
num_steps: 256
|
||||
num_minibatches: 4
|
||||
update_epochs: 4
|
||||
|
||||
learning_rate: 2.5e-4
|
||||
anneal_lr: true
|
||||
|
||||
# PPO specific - Fine-tuned
|
||||
gamma: 0.99
|
||||
gae_lambda: 0.95
|
||||
num_minibatches: 8 # More minibatches for stability
|
||||
update_epochs: 4
|
||||
norm_adv: true
|
||||
clip_coef: 0.2
|
||||
clip_coef: 0.1
|
||||
clip_vloss: true
|
||||
ent_coef: 0.01
|
||||
vf_coef: 0.5
|
||||
|
|
|
|||
|
|
@ -32,7 +32,7 @@ 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
|
||||
|
||||
|
|
|
|||
85
docs/HPC.md
Normal file
85
docs/HPC.md
Normal file
|
|
@ -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.
|
||||
14
docs/api/simulate.md
Normal file
14
docs/api/simulate.md
Normal file
|
|
@ -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
|
||||
```
|
||||
|
|
@ -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
|
||||
```
|
||||
3
env/hpc/modules.txt
vendored
Normal file
3
env/hpc/modules.txt
vendored
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
GCCcore/13.3.0
|
||||
Python/3.12.3-GCCcore-13.3.0
|
||||
FFmpeg/7.0.2-GCCcore-13.3.0
|
||||
19
env/hpc/requirements.txt
vendored
Normal file
19
env/hpc/requirements.txt
vendored
Normal file
|
|
@ -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
|
||||
|
|
@ -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",
|
||||
|
|
@ -20,12 +24,16 @@ dependencies = [
|
|||
"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 = [
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
27
scripts/analysis/README.md
Normal file
27
scripts/analysis/README.md
Normal file
|
|
@ -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`)
|
||||
143
scripts/analysis/explore_tensorboard.py
Normal file
143
scripts/analysis/explore_tensorboard.py
Normal file
|
|
@ -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()
|
||||
83
scripts/hpc/export_requirements.py
Normal file
83
scripts/hpc/export_requirements.py
Normal file
|
|
@ -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()
|
||||
54
scripts/hpc/install.sh
Normal file
54
scripts/hpc/install.sh
Normal file
|
|
@ -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'
|
||||
63
scripts/hpc/train.pbs
Normal file
63
scripts/hpc/train.pbs
Normal file
|
|
@ -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"
|
||||
|
|
@ -4,17 +4,15 @@ import argparse
|
|||
from pathlib import Path
|
||||
|
||||
from brittle_star_project import (
|
||||
ArenaConfig,
|
||||
Backend,
|
||||
BrittleStarEnv,
|
||||
BrittleStarEnvFactory,
|
||||
EnvConfig,
|
||||
MorphologyConfig,
|
||||
Task,
|
||||
SimulationConfig,
|
||||
simulate_policy,
|
||||
)
|
||||
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",
|
||||
choices=[t.value for t in Task],
|
||||
default=Task.DIRECTED_LOCOMOTION.value,
|
||||
"--backend",
|
||||
choices=[b for b in Backend],
|
||||
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
|
||||
task = Task(args.task)
|
||||
|
||||
morphology_cfg = MorphologyConfig()
|
||||
arena_cfg = ArenaConfig(attach_target=(task == Task.DIRECTED_LOCOMOTION))
|
||||
env_cfg = EnvConfig(task=task)
|
||||
backend = args.backend
|
||||
|
||||
factory = BrittleStarEnvFactory()
|
||||
raw_env = factory.create_environment(backend, morphology_cfg, arena_cfg, env_cfg)
|
||||
75
scripts/train.py
Normal file
75
scripts/train.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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,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
|
||||
|
|
@ -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",
|
||||
]
|
||||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
@ -44,8 +59,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
|
||||
|
|
@ -54,7 +67,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
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from brittle_star_project import (
|
|||
ArenaConfig,
|
||||
Backend,
|
||||
)
|
||||
from brittle_star_project.environment import from_file
|
||||
|
||||
|
||||
class BrittleStarJaxEnvWrapper:
|
||||
|
|
@ -84,3 +85,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})"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
pass
|
||||
|
||||
def from_file(path: str) -> tuple[MorphologyConfig, ArenaConfig, EnvConfig]:
|
||||
"""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
|
||||
|
|
|
|||
|
|
@ -0,0 +1,3 @@
|
|||
from .renderer import simulate_policy, SimulationConfig, ControlPolicy
|
||||
|
||||
__all__ = ["simulate_policy", "SimulationConfig", "ControlPolicy"]
|
||||
|
|
@ -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",
|
||||
]
|
||||
|
|
@ -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)
|
||||
|
|
@ -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)),
|
||||
)
|
||||
536
src/brittle_star_project/trainers/PPOTrainer.py
Normal file
536
src/brittle_star_project/trainers/PPOTrainer.py
Normal file
|
|
@ -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()
|
||||
0
src/brittle_star_project/trainers/__init__.py
Normal file
0
src/brittle_star_project/trainers/__init__.py
Normal file
|
|
@ -1,6 +0,0 @@
|
|||
import jax
|
||||
from experiment_logger import get_logger
|
||||
|
||||
if __name__ == "__main__":
|
||||
logger = get_logger()
|
||||
logger.info(f"JAX devices: {jax.devices()}")
|
||||
159
src/ppo.py
Normal file
159
src/ppo.py
Normal file
|
|
@ -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
|
||||
426
src/train.py
426
src/train.py
|
|
@ -1,426 +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
|
||||
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 experiment_logger import UnifiedLogger, get_logger
|
||||
from experiment_logger.config_utils import merge_config_with_cli
|
||||
|
||||
|
||||
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())}"
|
||||
get_logger().info(f"Run name: {run_name}")
|
||||
|
||||
# Initialize unified logger (replaces wandb.init and tensorboard writer)
|
||||
logger = UnifiedLogger(
|
||||
run_name=run_name,
|
||||
config=vars(args),
|
||||
project_name=args.wandb_project_name,
|
||||
entity=args.wandb_entity,
|
||||
use_wandb=args.track,
|
||||
save_code=True,
|
||||
)
|
||||
|
||||
# Keep TensorBoard writer for backward compatibility
|
||||
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")
|
||||
device = "cpu" # Force CPU for JAX
|
||||
logger.info(f"Device: {device}")
|
||||
|
||||
logger.info("Creating 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
|
||||
|
||||
logger.info("Initializing 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
|
||||
logger.info("Resetting 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,
|
||||
)
|
||||
|
||||
logger.info("Starting training...")
|
||||
iters_bar = tqdm.tqdm(range(1, args.num_iterations + 1))
|
||||
for iteration 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))
|
||||
avg_episodic_length = np.mean(jax.device_get(episode_stats.returned_episode_lengths))
|
||||
learning_rate = agent_state.opt_state[1].hyperparams["learning_rate"].item()
|
||||
sps = int(global_step / (time.time() - start_time))
|
||||
sps_update = int(args.num_envs * args.num_steps / (time.time() - iteration_time_start))
|
||||
|
||||
iters_bar.set_postfix_str(
|
||||
f"global_step={global_step}, avg_episodic_return={avg_episodic_return}"
|
||||
)
|
||||
|
||||
# Log to unified logger
|
||||
logger.log(
|
||||
{
|
||||
"charts/avg_episodic_return": avg_episodic_return,
|
||||
"charts/avg_episodic_length": avg_episodic_length,
|
||||
"charts/learning_rate": learning_rate,
|
||||
"charts/SPS": sps,
|
||||
"charts/SPS_update": sps_update,
|
||||
"losses/value_loss": v_loss[-1, -1].item(),
|
||||
"losses/policy_loss": pg_loss[-1, -1].item(),
|
||||
"losses/entropy": entropy_loss[-1, -1].item(),
|
||||
"losses/approx_kl": approx_kl[-1, -1].item(),
|
||||
"losses/loss": loss[-1, -1].item(),
|
||||
},
|
||||
step=global_step,
|
||||
)
|
||||
|
||||
# Also log to TensorBoard for backward compatibility
|
||||
writer.add_scalar("charts/avg_episodic_return", avg_episodic_return, global_step)
|
||||
writer.add_scalar("charts/avg_episodic_length", avg_episodic_length, global_step)
|
||||
writer.add_scalar("charts/learning_rate", learning_rate, 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)
|
||||
writer.add_scalar("charts/SPS", sps, global_step)
|
||||
writer.add_scalar("charts/SPS_update", sps_update, global_step)
|
||||
|
||||
# Save periodic checkpoints
|
||||
if args.checkpoint_frequency > 0 and iteration % args.checkpoint_frequency == 0:
|
||||
logger.save_checkpoint(
|
||||
params={
|
||||
"network_params": agent_state.params["network_params"],
|
||||
"actor_params": agent_state.params["actor_params"],
|
||||
"critic_params": agent_state.params["critic_params"],
|
||||
},
|
||||
step=global_step,
|
||||
metadata={
|
||||
"iteration": iteration,
|
||||
"avg_episodic_return": float(avg_episodic_return),
|
||||
"avg_episodic_length": float(avg_episodic_length),
|
||||
},
|
||||
)
|
||||
|
||||
if args.save_model:
|
||||
# Save using unified logger (better organization and WandB integration)
|
||||
logger.save_final_model(
|
||||
params={
|
||||
"network_params": agent_state.params["network_params"],
|
||||
"actor_params": agent_state.params["actor_params"],
|
||||
"critic_params": agent_state.params["critic_params"],
|
||||
},
|
||||
metadata={
|
||||
"global_step": global_step,
|
||||
"avg_episodic_return": float(avg_episodic_return),
|
||||
"config": vars(args),
|
||||
},
|
||||
)
|
||||
|
||||
# Also save in old format for backward compatibility
|
||||
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"],
|
||||
],
|
||||
]
|
||||
)
|
||||
)
|
||||
logger.info(f"Legacy model saved to {model_path}")
|
||||
|
||||
# Finalize logging
|
||||
logger.finish()
|
||||
env.close()
|
||||
writer.close()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
# Enhanced argument parsing with YAML config support
|
||||
args = merge_config_with_cli(PPOArgs)
|
||||
|
||||
# Print final configuration
|
||||
from experiment_logger.config_utils import print_config
|
||||
|
||||
print_config(args, "Final Training Configuration")
|
||||
|
||||
train(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
46
uv.lock
generated
46
uv.lock
generated
|
|
@ -22,15 +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"] },
|
||||
]
|
||||
|
|
@ -54,14 +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 = [
|
||||
|
|
@ -1264,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"
|
||||
|
|
@ -2430,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"
|
||||
|
|
|
|||
Reference in a new issue