diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index 3e5c8d4..623526f 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -45,4 +45,4 @@ "features": { "ghcr.io/devcontainers/features/common-utils:1": {} } -} \ No newline at end of file +} diff --git a/.github/workflows/update_hpc_requirements.yml b/.github/workflows/update_hpc_requirements.yml new file mode 100644 index 0000000..38dbc42 --- /dev/null +++ b/.github/workflows/update_hpc_requirements.yml @@ -0,0 +1,38 @@ +name: Update HPC requirements + +on: + push: + paths: + - pyproject.toml + branches: + - main + - dev + - "ci/**" + +jobs: + update-hpc-requirements: + runs-on: ubuntu-latest + permissions: + contents: write + steps: + - uses: actions/checkout@v4 + with: + ref: ${{ github.head_ref || github.ref_name }} + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install uv + uses: astral-sh/setup-uv@v5 + + - name: Regenerate env/hpc/requirements.txt + run: uv run scripts/hpc/export_requirements.py + + - name: Commit updated requirements if changed + uses: stefanzweifel/git-auto-commit-action@v5 + with: + commit_message: "chore(hpc): update env/hpc/requirements.txt from pyproject.toml [skip ci]" + file_pattern: env/hpc/requirements.txt + commit_author: "github-actions[bot] " diff --git a/.gitignore b/.gitignore index 6cc1e08..79b1594 100644 --- a/.gitignore +++ b/.gitignore @@ -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 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c86e484..c289bb6 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,4 +1,12 @@ repos: + - repo: local + hooks: + - id: ruff-format + name: ruff format (uv) + entry: uv run ruff format + language: system + types: [python] + - repo: https://github.com/alessandrojcm/commitlint-pre-commit-hook rev: v9.16.0 hooks: @@ -6,15 +14,8 @@ repos: stages: [commit-msg] additional_dependencies: ["@commitlint/config-conventional"] - - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.9.9 - hooks: - - id: ruff - args: [ --fix ] - - id: ruff-format - - repo: https://github.com/pre-commit/pre-commit-hooks rev: v5.0.0 hooks: - id: no-commit-to-branch - args: ['--branch', 'main', '--branch', 'dev'] + args: ['--branch', 'main', '--branch', 'dev'] \ No newline at end of file diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..9b38853 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,7 @@ +{ + "python.testing.pytestArgs": [ + "tests" + ], + "python.testing.unittestEnabled": false, + "python.testing.pytestEnabled": true +} \ No newline at end of file diff --git a/README.md b/README.md index 74b001e..19c994d 100644 --- a/README.md +++ b/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. diff --git a/configs/example.json b/configs/example.json new file mode 100644 index 0000000..8b646db --- /dev/null +++ b/configs/example.json @@ -0,0 +1,8 @@ +{ + "morphology": { + "num_arms": 2, + "num_segments_per_arm": 4, + "use_p_control": true, + "use_torque_control": false + } +} \ No newline at end of file diff --git a/configs/hpc/smoke_test.yaml b/configs/hpc/smoke_test.yaml new file mode 100644 index 0000000..1dd2fc0 --- /dev/null +++ b/configs/hpc/smoke_test.yaml @@ -0,0 +1,11 @@ +# Minimal config to verify HPC setup is functional. +# Run with: python scripts/train.py --config-path configs/hpc/smoke_test.yaml +exp_name: "hpc_smoke_test" +seed: 0 +track: false # Test WandB integration +capture_video: false # No rendering for smoke test +save_model: true # Test the end-of-training save routine +num_envs: 512 +total_timesteps: 65536 +num_steps: 128 +cuda: true diff --git a/configs/production_training.yaml b/configs/production_training.yaml index 351ee83..6dce29b 100644 --- a/configs/production_training.yaml +++ b/configs/production_training.yaml @@ -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 diff --git a/docs/CONTRIBUTING.md b/docs/CONTRIBUTING.md index ee9b4dc..bd58706 100644 --- a/docs/CONTRIBUTING.md +++ b/docs/CONTRIBUTING.md @@ -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 diff --git a/docs/HPC.md b/docs/HPC.md new file mode 100644 index 0000000..3e0e8d1 --- /dev/null +++ b/docs/HPC.md @@ -0,0 +1,85 @@ +# HPC Guide + +Full documentation: + +## 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. diff --git a/docs/api/simulate.md b/docs/api/simulate.md new file mode 100644 index 0000000..9c4b21b --- /dev/null +++ b/docs/api/simulate.md @@ -0,0 +1,14 @@ +# Training and Simulation for Brittle Star Models + +## Simulating a model + +In order to simulate and view the behavior of a trained model, you can use the `simulate.py` script. This script allows you to specify the path to a trained model and will launch a simulation using that model. This script has the following parameters: + +- `--model`: The path to the trained model artifact to simulate. +- `--model-type`: The type of model to simulate (e.g., `random`, ...) +- `--task`: The task to simulate (e.g., `directed_locomotion`, ...) +- `--seed`: The random seed for reproducibility. + +```bash +python simulate.py --model artifacts/my_model --model-type random --task directed_locomotion --seed 0 +``` \ No newline at end of file diff --git a/docs/api/train_simulate.md b/docs/api/train_simulate.md deleted file mode 100644 index 9b299a9..0000000 --- a/docs/api/train_simulate.md +++ /dev/null @@ -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 -``` \ No newline at end of file diff --git a/env/hpc/modules.txt b/env/hpc/modules.txt new file mode 100644 index 0000000..4c4a1d7 --- /dev/null +++ b/env/hpc/modules.txt @@ -0,0 +1,3 @@ +GCCcore/13.3.0 +Python/3.12.3-GCCcore-13.3.0 +FFmpeg/7.0.2-GCCcore-13.3.0 diff --git a/env/hpc/requirements.txt b/env/hpc/requirements.txt new file mode 100644 index 0000000..819ea32 --- /dev/null +++ b/env/hpc/requirements.txt @@ -0,0 +1,19 @@ +biorobot==0.4.2 +cleanrl>=0.4.8 +evosax==0.2.0 +flax>=0.12.2 +gymnasium>=1.2.3 +ipykernel==7.2.0 +jax[cuda13]==0.9.0.1 +numpy>=2.0.0 +protobuf>=5.0.0 +warp-lang +mujoco-warp +matplotlib==3.10.8 +mediapy==1.2.6 +optax>=0.2.6 +pyopengl>=3.1.10 +pyopengl-accelerate>=3.1.10 +tyro>=1.0.10 +wandb==0.24.2 +torch>=2.4.0 diff --git a/pyproject.toml b/pyproject.toml index 9ba5a9a..b85ce0f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,10 @@ dependencies = [ "gymnasium>=1.2.3", "ipykernel==7.2.0", "jax==0.9.0.1", + "numpy>=2.0.0", + "protobuf>=5.0.0", + "warp-lang", + "mujoco-warp", "matplotlib==3.10.8", "mediapy==1.2.6", "optax>=0.2.6", @@ -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 = [ diff --git a/ruff.toml b/ruff.toml index 16e880c..ad24132 100644 --- a/ruff.toml +++ b/ruff.toml @@ -1,4 +1,5 @@ line-length = 100 +exclude = ["wandb"] [lint] extend-select = [ @@ -394,4 +395,5 @@ extend-ignore = [ # "PLW1404", # implicit-str-concat ] - +[lint.per-file-ignores] +"__init__.py" = ["F401"] diff --git a/scripts/analysis/README.md b/scripts/analysis/README.md new file mode 100644 index 0000000..a99f952 --- /dev/null +++ b/scripts/analysis/README.md @@ -0,0 +1,27 @@ +# Experiment Analysis Tools + +This directory contains scripts for post-processing and analyzing experiment results, including TensorBoard logs and saved model weights. + +## Scripts + +### 1. `explore_tensorboard.py` +A CLI tool to summarize TensorBoard `tfevents` files without a GUI. + +**Key Features:** +- Displays last values, min, max, and step counts for all scalar metrics. +- Calculates total run duration and estimated completion percentage. +- Exports granular scalar data to CSV for analysis in Excel/Pandas. + +**Usage:** +```bash +# General usage +python explore_tensorboard.py + +# Exporting data +python explore_tensorboard.py --csv data.csv +``` + +**Requirements:** +- `pandas` +- `tensorboard` +- `tensorflow-cpu` (or `tensorflow`) diff --git a/scripts/analysis/explore_tensorboard.py b/scripts/analysis/explore_tensorboard.py new file mode 100644 index 0000000..0ed1027 --- /dev/null +++ b/scripts/analysis/explore_tensorboard.py @@ -0,0 +1,143 @@ +#!/usr/bin/env python3 +""" +Reproducible CLI tool to explore TensorBoard logs. +Designed for both local development and HPC diagnostics. + +Requirements: + pip install tensorboard + +Usage: + python explore_tensorboard.py [--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() diff --git a/scripts/hpc/export_requirements.py b/scripts/hpc/export_requirements.py new file mode 100644 index 0000000..5908ad3 --- /dev/null +++ b/scripts/hpc/export_requirements.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +"""Export HPC pip requirements from pyproject.toml. + +This is a LOCAL DEVELOPER UTILITY — run it on your own machine before pushing +code whenever pyproject.toml dependencies change. It reads the modules from +env/hpc/modules.txt and the full dependency list from pyproject.toml, then +writes the remainder to env/hpc/requirements.txt. +""" + +from __future__ import annotations + +import re +import sys +from pathlib import Path + +ROOT = Path(__file__).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() diff --git a/scripts/hpc/install.sh b/scripts/hpc/install.sh new file mode 100644 index 0000000..d1b69db --- /dev/null +++ b/scripts/hpc/install.sh @@ -0,0 +1,54 @@ +#!/bin/bash -l +# scripts/hpc/install.sh +# +# Usage (on any compute node): +# bash scripts/hpc/install.sh +# +# Batch usage: +# qsub scripts/hpc/install.sh + +#PBS -N brittlestar-install +#PBS -l walltime=00:15:00 + +set -euo pipefail + +# Preliminary status echo +echo ">>> Starting installation job $PBS_JOBID on $(hostname)..." + +if [ -n "$PBS_O_WORKDIR" ]; then + cd "$PBS_O_WORKDIR" +fi + +mkdir "${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' diff --git a/scripts/hpc/train.pbs b/scripts/hpc/train.pbs new file mode 100644 index 0000000..3d53193 --- /dev/null +++ b/scripts/hpc/train.pbs @@ -0,0 +1,63 @@ +# Production training (requires GPU at runtime): +# qsub -l gpus=1 scripts/hpc/train.pbs +# Debug/CPU training: +# qsub scripts/hpc/train.pbs + +#PBS -N brittlestar-ppo +#PBS -l nodes=1:ppn=8 +#PBS -l walltime=24:00:00 +#PBS -o runs/brittlestar-ppo.o$PBS_JOBID +#PBS -e runs/brittlestar-ppo.e$PBS_JOBID + +set -euo pipefail + +# Preliminary status echo +echo ">>> Starting training job $PBS_JOBID on $(hostname)..." + +if [ -n "$PBS_O_WORKDIR" ]; then + cd "$PBS_O_WORKDIR" +fi + +# Set up storage paths dynamically +PROJ_NAME=$(basename "$PWD") +RUN_ID="brittlestar_${PBS_JOBID}" +SCRATCH_RUNDIR="$VSC_SCRATCH/runs/$RUN_ID" +DATA_RUNDIR="$VSC_DATA/runs/$RUN_ID" +mkdir -p "$SCRATCH_RUNDIR" "$DATA_RUNDIR" runs/ + +# Keep caches off $VSC_HOME (quota ~3 GB). +export PIP_CACHE_DIR="$VSC_SCRATCH/.cache/pip" +export UV_CACHE_DIR="$VSC_SCRATCH/.cache/uv" +mkdir -p "$PIP_CACHE_DIR" "$UV_CACHE_DIR" + +module load vsc-venv + +echo ">>> Synchronizing and activating environment (vsc-venv)..." +HPC_CONFIG_DIR="$VSC_DATA/$PROJ_NAME/env/hpc" +if [ ! -d "$HPC_CONFIG_DIR" ]; then + echo "ERROR: HPC_CONFIG_DIR ($HPC_CONFIG_DIR) does not exist. Run install.sh first." + exit 1 +fi + +# cd to $VSC_DATA so vsc-venv finds its venvs/ directory there, not in $HOME. +cd "$VSC_DATA/$PROJ_NAME" +set +euo pipefail +source vsc-venv --activate \ + --modules "$HPC_CONFIG_DIR/modules.txt" \ + --requirements "$HPC_CONFIG_DIR/requirements.txt" +set -euo pipefail +cd "$PBS_O_WORKDIR" + + +echo ">>> Starting BrittleStar training..." +export MUJOCO_GL=egl +export WANDB_DIR="$SCRATCH_RUNDIR" + +python src/train.py \ + --env-config-path configs/hpc/smoke_test.yaml \ + --run-dir "$SCRATCH_RUNDIR" + +echo ">>> Staging out results to $DATA_RUNDIR..." +cp -r "$SCRATCH_RUNDIR/." "$DATA_RUNDIR/" + +echo ">>> Done" diff --git a/experiments/simulate.py b/scripts/simulate.py similarity index 84% rename from experiments/simulate.py rename to scripts/simulate.py index 346bcda..28062fa 100644 --- a/experiments/simulate.py +++ b/scripts/simulate.py @@ -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) diff --git a/scripts/train.py b/scripts/train.py new file mode 100644 index 0000000..9ee6da6 --- /dev/null +++ b/scripts/train.py @@ -0,0 +1,75 @@ +import subprocess +import time + +import torch +import tyro +import yaml +import os + +from brittle_star_project.dataclasses import PPOArgs +from brittle_star_project.trainers.PPOTrainer import PPOTrainer +from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper + + +def make_env(config_path: str | None, num_envs: int) -> BrittleStarJaxEnvWrapper: + if config_path is None: + return BrittleStarJaxEnvWrapper.default(num_envs=num_envs) + return BrittleStarJaxEnvWrapper.from_config(config_path, num_envs=num_envs) + + +def parse_args(log: bool = True) -> PPOArgs: + temp_args = tyro.cli(PPOArgs) + + if temp_args.hyperparameter_config_path is not None: + if log: + print(f"Loading hyperparameter config from {temp_args.hyperparameter_config_path}") + + with open(temp_args.hyperparameter_config_path, "r") as f: + config = yaml.safe_load(f) + if config: + # parse PPOArgs with defaults from yaml. + for key, value in config.items(): + if hasattr(temp_args, key): + setattr(temp_args, key, value) + + # Reparse CLI to ensure they OVERRIDE the yaml + args = tyro.cli(PPOArgs, default=temp_args) + else: + if log: + print("No hyperparameter config provided, using default config") + + args = temp_args + return args + + +def get_git_hash() -> str: + try: + return ( + subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode("ascii").strip() + ) + except subprocess.CalledProcessError | UnicodeDecodeError: + return "none" + + +if __name__ == "__main__": + args = parse_args() + + args.batch_size = args.num_envs * args.num_steps + args.minibatch_size = args.batch_size // args.num_minibatches + args.num_iterations = args.total_timesteps // args.batch_size + + git_hash = get_git_hash() + run_name = f"{args.exp_name}__seed_{args.seed}__{git_hash}__{int(time.time())}" + if args.run_dir is None: + run_dir = f"runs/{run_name}" + else: + run_dir = args.run_dir + + os.makedirs(run_dir, exist_ok=True) + + env = make_env(args.env_config_path, args.num_envs) + + torch.backends.cudnn.deterministic = args.torch_deterministic + + ppo_trainer = PPOTrainer(args, env, run_dir, run_name) + ppo_trainer.train() diff --git a/src/brittle_star_project/rl/DummyAgent.py b/src/MLPs/mlps.py similarity index 64% rename from src/brittle_star_project/rl/DummyAgent.py rename to src/MLPs/mlps.py index fff4068..6abb540 100644 --- a/src/brittle_star_project/rl/DummyAgent.py +++ b/src/MLPs/mlps.py @@ -1,36 +1,27 @@ -from dataclasses import dataclass, fields +from dataclasses import dataclass, fields, field import flax import flax.linen as nn import jax.numpy as jnp import jax.tree_util -import numpy as np +from typing import Sequence, Callable from flax.linen.initializers import constant, orthogonal -class Network(nn.Module): - """ - 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 diff --git a/src/__init__.py b/src/__init__.py index e69de29..35ecd10 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -0,0 +1,25 @@ +from .brittle_star_project import ( + ArenaConfig, + Backend, + BrittleStarEnv, + BrittleStarEnvFactory, + EnvConfig, + MorphologyConfig, + Task, + simulate_policy, + SimulationConfig, + ControlPolicy, +) + +__all__ = [ + "ArenaConfig", + "Backend", + "BrittleStarEnv", + "BrittleStarEnvFactory", + "EnvConfig", + "MorphologyConfig", + "Task", + "simulate_policy", + "SimulationConfig", + "ControlPolicy", +] diff --git a/src/brittle_star_project/__init__.py b/src/brittle_star_project/__init__.py index b6b123a..3902cbb 100644 --- a/src/brittle_star_project/__init__.py +++ b/src/brittle_star_project/__init__.py @@ -2,6 +2,7 @@ from .environment.env_types import Backend, Task from .environment.env_config import ArenaConfig, EnvConfig, MorphologyConfig from .environment.factory import BrittleStarEnvFactory from .environment.env_wrapper import BrittleStarEnv +from .render import simulate_policy, SimulationConfig, ControlPolicy __all__ = [ "ArenaConfig", @@ -11,4 +12,7 @@ __all__ = [ "EnvConfig", "MorphologyConfig", "Task", + "simulate_policy", + "SimulationConfig", + "ControlPolicy", ] diff --git a/src/brittle_star_project/dataclasses/PPOArgs.py b/src/brittle_star_project/dataclasses/PPOArgs.py index d59defa..68265fa 100644 --- a/src/brittle_star_project/dataclasses/PPOArgs.py +++ b/src/brittle_star_project/dataclasses/PPOArgs.py @@ -1,15 +1,30 @@ from dataclasses import dataclass +import jax + +@jax.tree_util.register_dataclass @dataclass class PPOArgs: """ source: https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/ppo_atari_envpool_xla_jax_scan.py """ + # path to environment config file, if None, use default config + env_config_path: str | None = None + + # path to hyperparameter config file (yaml), if None, use default config + hyperparameter_config_path: str | None = None + # the name of this experiment exp_name: str = "brittle_star_ppo" + # the directory to save the experiment results + run_dir: str | None = None + + # how often to save checkpoints (0 to disable) + checkpoint_frequency: int = 0 + # seed of the experiment seed: int = 1 @@ -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 diff --git a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py index c7cb3c6..b143d9c 100644 --- a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py +++ b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py @@ -8,6 +8,7 @@ from brittle_star_project import ( ArenaConfig, Backend, ) +from brittle_star_project.environment import from_file class BrittleStarJaxEnvWrapper: @@ -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})" + ) diff --git a/src/brittle_star_project/environment/__init__.py b/src/brittle_star_project/environment/__init__.py index 78896cb..aad8c35 100644 --- a/src/brittle_star_project/environment/__init__.py +++ b/src/brittle_star_project/environment/__init__.py @@ -1,4 +1,4 @@ -from .env_config import ArenaConfig, EnvConfig, MorphologyConfig +from .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", ] diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index f50f4d8..a693286 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/src/brittle_star_project/environment/env_config.py @@ -1,6 +1,7 @@ from __future__ import annotations from dataclasses import dataclass, field +import json from .env_types import Task @@ -48,6 +49,18 @@ class EnvConfig: # Per docs in upstream env config: integer factors of 200. light_perlin_noise_scale: int = 0 - @staticmethod - def from_json(path: str) -> EnvConfig: - 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 diff --git a/src/brittle_star_project/render/__init__.py b/src/brittle_star_project/render/__init__.py index e69de29..51e0aa1 100644 --- a/src/brittle_star_project/render/__init__.py +++ b/src/brittle_star_project/render/__init__.py @@ -0,0 +1,3 @@ +from .renderer import simulate_policy, SimulationConfig, ControlPolicy + +__all__ = ["simulate_policy", "SimulationConfig", "ControlPolicy"] diff --git a/src/brittle_star_project/rl/__init__.py b/src/brittle_star_project/rl/__init__.py deleted file mode 100644 index e58e1c6..0000000 --- a/src/brittle_star_project/rl/__init__.py +++ /dev/null @@ -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", -] diff --git a/src/brittle_star_project/rl/base.py b/src/brittle_star_project/rl/base.py deleted file mode 100644 index 8e3c439..0000000 --- a/src/brittle_star_project/rl/base.py +++ /dev/null @@ -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 "" - 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 "" - 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) diff --git a/src/brittle_star_project/rl/random_policy_model.py b/src/brittle_star_project/rl/random_policy_model.py deleted file mode 100644 index 9372d49..0000000 --- a/src/brittle_star_project/rl/random_policy_model.py +++ /dev/null @@ -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)), - ) diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py new file mode 100644 index 0000000..7df57f0 --- /dev/null +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -0,0 +1,536 @@ +import datetime +import random +import sys +import time +from dataclasses import asdict, dataclass +from functools import partial +from typing import Any + +import flax +import jax +import jax.numpy as jnp +import numpy as np +import optax +import tqdm +from flax.training.train_state import TrainState +from torch.utils.tensorboard import SummaryWriter + +from brittle_star_project.dataclasses import EpisodeStatistics, PPOArgs +from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper +from MLPs.mlps import ( + Actor, + AgentParams, + GenericDenseLayersWithActivation, + OneDenseLayerMLP, + Storage, +) +from ppo import PPO + + +@jax.jit +def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, learning_rate): + frac = 1.0 - (count // (minibatch_count * update_epochs)) / num_iterations + return learning_rate * frac + + +@jax.jit +def _convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray: + return jax.vmap(lambda o: jnp.concatenate([v.flatten() for v in o.values() if v.size > 0]))( + obs_dict + ) + + +# removed jit: used in _rollout_jit, so will be compiled with _rollout_jit +def _get_action_and_value_noise( + sensor: GenericDenseLayersWithActivation, + feature_extractor: GenericDenseLayersWithActivation, + actor: Actor, + critic: OneDenseLayerMLP, + agent_state: TrainState, + next_obs: jnp.ndarray, + key: jax.random.PRNGKey, +): + hidden = sensor.apply(agent_state.params["sensor_params"], next_obs) + hidden_critic = feature_extractor.apply( + agent_state.params["feature_extractor_params"], next_obs + ) + + # Continuous actions: sample from a Gaussian parameterized by the actor + mean, log_std = actor.apply(agent_state.params["actor_params"], hidden) + key, subkey = jax.random.split(key) + noise = jax.random.normal(subkey, shape=mean.shape) + std = jnp.exp(log_std) + action = mean + noise * std + logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1) + value = critic.apply(agent_state.params["critic_params"], hidden_critic) + return action, logprob, value.squeeze(-1), key + + +# removed jit: used in _rollout_jit, so will be compiled with _rollout_jit +def _step_once( + carry, + _, + env_step_fn, + sensor: GenericDenseLayersWithActivation, + feature_extractor: GenericDenseLayersWithActivation, + actor: Actor, + critic: OneDenseLayerMLP, +): + agent_state, episode_stats, obs, done, key, env_state = carry + action, logprob, value, key = _get_action_and_value_noise( + sensor, feature_extractor, actor, critic, agent_state, obs, key + ) + + episode_stats, env_state, (next_obs, reward, next_done) = env_step_fn( + episode_stats, env_state, action + ) + + storage = Storage( + obs=obs, + actions=action, + logprobs=logprob, + dones=done, + values=value, + rewards=reward, + returns=jnp.zeros_like(reward), + advantages=jnp.zeros_like(reward), + ) + return (agent_state, episode_stats, next_obs, next_done, key, env_state), storage + + +# removed jit: used in _rollout_jit, so will be compiled with _rollout_jit +def _step_env_wrapped(episode_stats, env_state, action, env_step_fn): + next_env_state = env_step_fn(env_state, action) + + # Extract per-environment signals from the state object + reward = next_env_state.reward # (num_envs,) + terminated = next_env_state.terminated # (num_envs,) + truncated = next_env_state.truncated # (num_envs,) + done = terminated | truncated # (num_envs,) + + new_episode_return = episode_stats.episode_returns + reward + new_episode_length = episode_stats.episode_lengths + 1 + + episode_stats = episode_stats.replace( + episode_returns=new_episode_return * (1 - done), + episode_lengths=new_episode_length * (1 - done), + returned_episode_returns=jnp.where( + done, new_episode_return, episode_stats.returned_episode_returns + ), + returned_episode_lengths=jnp.where( + done, new_episode_length, episode_stats.returned_episode_lengths + ), + ) + return ( + episode_stats, + next_env_state, + (_convert_obs_dict_to_array(next_env_state.observations), reward, done), + ) + + +# jit applied in wrapper method self._rollout_jit using partial +def _rollout_jit( + agent_state, + episode_stats, + env_state, + next_obs, + next_done, + key, + max_steps, + step_env_fn, + sensor: GenericDenseLayersWithActivation, + feature_extractor: GenericDenseLayersWithActivation, + actor: Actor, + critic: OneDenseLayerMLP, +): + (agent_state, episode_stats, next_obs, next_done, key, env_state), storage = jax.lax.scan( + partial( + _step_once, + sensor=sensor, + feature_extractor=feature_extractor, + actor=actor, + critic=critic, + env_step_fn=step_env_fn, + ), + (agent_state, episode_stats, next_obs, next_done, key, env_state), + (), + max_steps, + ) + return agent_state, episode_stats, next_obs, next_done, storage, key, env_state + + +# removed jit: used in _compute_gae_jit, so will be compiled with _compute_gae_jit +def _compute_gae_once(carry, inp, gamma, gae_lambda): + advantages = carry + nextdone, nextvalues, curvalues, reward = inp + nextnonterminal = 1.0 - nextdone + delta = reward + gamma * nextvalues * nextnonterminal - curvalues + advantages = delta + gamma * gae_lambda * nextnonterminal * advantages + return advantages, advantages + + +# jit applied on partial-wrapped wrapper method self._compute_gae_jit +def _compute_gae_jit( + agent_state, storage, next_obs, next_done, gamma, gae_lambda, num_envs, sensor, critic +): + next_value = critic.apply( + agent_state.params["critic_params"], + sensor.apply(agent_state.params["sensor_params"], next_obs), + ).squeeze(-1) + + advantages = jnp.zeros((num_envs,)) + dones = jnp.concatenate([storage.dones, next_done[None, :]], axis=0) + values = jnp.concatenate([storage.values, next_value[None, :]], axis=0) + _, advantages = jax.lax.scan( + partial(_compute_gae_once, gamma=gamma, gae_lambda=gae_lambda), + advantages, + (dones[1:], values[1:], values[:-1], storage.rewards), + reverse=True, + ) + return storage.replace(advantages=advantages, returns=advantages + storage.values) + + +@dataclass +class LossInfo: + # todo: better typing + loss: Any + pg_loss: Any + v_loss: Any + entropy_loss: Any + approx_kl: Any + avg_episodic_return: Any + + +class PPOTrainer: + def __init__(self, args: PPOArgs, env: BrittleStarJaxEnvWrapper, run_dir: str, run_name: str): + self.args = args + self.env = env + self.run_dir = run_dir + self.run_name = run_name + self.writer = SummaryWriter(self.run_dir) + + self.key = jax.random.PRNGKey(args.seed) + + self.sensor, self.feature_extractor, self.actor, self.critic = self._init_agent() + self.sensor.apply = jax.jit(self.sensor.apply) + self.feature_extractor.apply = jax.jit(self.feature_extractor.apply) + self.actor.apply = jax.jit(self.actor.apply) + self.critic.apply = jax.jit(self.critic.apply) + + self._rollout_jit = jax.jit( + partial( + _rollout_jit, + max_steps=self.args.num_steps, + step_env_fn=partial(_step_env_wrapped, env_step_fn=self.env.step), + sensor=self.sensor, + feature_extractor=self.feature_extractor, + actor=self.actor, + critic=self.critic, + ) + ) + self._compute_gae_jit = jax.jit( + partial( + _compute_gae_jit, + num_envs=self.args.num_envs, + gamma=self.args.gamma, + gae_lambda=self.args.gae_lambda, + sensor=self.sensor, + critic=self.critic, + ) + ) + + self._ppo = PPO(self.args, self.sensor, self.actor, self.critic, self.feature_extractor) + + self.agent_state = self._init_agent_state() + + self.episode_stats = self._init_episode_stats() + + self._init_random() + + def _init_random(self, log: bool = True): + if log: + print(f"[RANDOM]: Setting random seed to {self.args.seed}") + + random.seed(self.args.seed) + np.random.seed(self.args.seed) + + def _init_agent(self, log: bool = True): + if log: + print("[AGENT]: Initializing agent...") + + sensor = GenericDenseLayersWithActivation() + feature_extractor = GenericDenseLayersWithActivation() + actor = Actor( + action_dim=self.env.single_action_space.shape[0] + ) # continuous actions for MJX + critic = OneDenseLayerMLP() + # messenger = OneDenseLayerMLP() + return sensor, feature_extractor, actor, critic + + def _init_agent_state(self, log: bool = True) -> TrainState: + if log: + print("[AGENT STATE]: Initializing agent state...") + + self.key, sensor_key, actor_key, critic_key, feature_extractor_key = jax.random.split( + self.key, 5 + ) + + sample_obs = jnp.concatenate( + [ + v.flatten() + for v in self.env.single_observation_space.sample( + rng=jax.random.PRNGKey(0) + ).values() + if v.size > 0 + ] + ) + sensor_params = self.sensor.init(sensor_key, sample_obs) + feature_extractor_params = self.feature_extractor.init(feature_extractor_key, sample_obs) + actor_params = self.actor.init(actor_key, self.sensor.apply(sensor_params, sample_obs)) + critic_params = self.critic.init( + critic_key, self.feature_extractor.apply(feature_extractor_params, sample_obs) + ) + + return TrainState.create( + apply_fn=None, + params=asdict( + AgentParams(sensor_params, actor_params, critic_params, feature_extractor_params) + ), + tx=optax.chain( + optax.clip_by_global_norm(self.args.max_grad_norm), + optax.inject_hyperparams(optax.adam)( + learning_rate=partial( + _linear_schedule, + minibatch_count=self.args.num_minibatches, + update_epochs=self.args.update_epochs, + num_iterations=self.args.num_iterations, + learning_rate=self.args.learning_rate, + ) + if self.args.anneal_lr + else self.args.learning_rate, + eps=1e-5, + ), + ), + ) + + def _init_episode_stats(self, log: bool = True) -> EpisodeStatistics: + if log: + print("[EPISODE STATS]: Initializing episode stats...") + + return EpisodeStatistics( + episode_returns=jnp.zeros(self.args.num_envs, dtype=jnp.float32), + episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32), + returned_episode_returns=jnp.zeros(self.args.num_envs, jnp.float32), + returned_episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32), + ) + + def _rollout(self, env_state, next_obs, next_done) -> tuple[Storage, ...]: + return self._rollout_jit( + self.agent_state, + self.episode_stats, + env_state, + next_obs, + next_done, + self.key, + ) + + def _compute_gae(self, storage, next_obs, next_done) -> Storage: + return self._compute_gae_jit( + self.agent_state, + storage, + next_obs, + next_done, + ) + + def _log( + self, + global_step, + episode_stats, + start_time, + iteration_time_start, + loss_info, + ): + + self.writer.add_scalar( + "charts/avg_episodic_return", loss_info.avg_episodic_return, global_step + ) + self.writer.add_scalar( + "charts/avg_episodic_length", + np.mean(jax.device_get(episode_stats.returned_episode_lengths)), + global_step, + ) + self.writer.add_scalar( + "charts/learning_rate", + self.agent_state.opt_state[1].hyperparams["learning_rate"].item(), + global_step, + ) + self.writer.add_scalar("losses/value_loss", loss_info.v_loss[-1, -1].item(), global_step) + self.writer.add_scalar("losses/policy_loss", loss_info.pg_loss[-1, -1].item(), global_step) + self.writer.add_scalar("losses/entropy", loss_info.entropy_loss[-1, -1].item(), global_step) + self.writer.add_scalar("losses/approx_kl", loss_info.approx_kl[-1, -1].item(), global_step) + self.writer.add_scalar("losses/loss", loss_info.loss[-1, -1].item(), global_step) + self.writer.add_scalar( + "charts/SPS", int(global_step / (time.time() - start_time)), global_step + ) + self.writer.add_scalar( + "charts/SPS_update", + int(self.args.num_envs * self.args.num_steps / (time.time() - iteration_time_start)), + global_step, + ) + + def _step( + self, env_state, next_obs, next_done, is_tty: bool, iteration: int, log: bool = True + ) -> tuple: + if log and not is_tty and iteration == 1: + print(f">>> [HPC] Starting first rollout (JIT): {time.ctime()}", flush=True) + + ( + self.agent_state, + self.episode_stats, + next_obs, + next_done, + storage, + self.key, + next_env_state, + ) = self._rollout(env_state, next_obs, next_done) + + if log and not is_tty and iteration == 1: + print(f">>> [HPC] First rollout completed: {time.ctime()}", flush=True) + + storage = self._compute_gae(storage, next_obs, next_done) + + if log and not is_tty and iteration == 1: + print(f">>> [HPC] Starting first PPO update (JIT): {time.ctime()}", flush=True) + + self.agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, self.key = ( + self._ppo.update_ppo(self.agent_state, storage, self.key) + ) + + if log and not is_tty and iteration == 1: + print(f">>> [HPC] First PPO update completed: {time.ctime()}", flush=True) + + avg_episodic_return = float( + jnp.mean(jax.device_get(self.episode_stats.returned_episode_returns)).item() + ) + + return ( + next_env_state, + next_obs, + next_done, + LossInfo( + loss=loss, + pg_loss=pg_loss, + v_loss=v_loss, + entropy_loss=entropy_loss, + approx_kl=approx_kl, + avg_episodic_return=avg_episodic_return, + ), + ) + + def _close(self): + self.env.close() + self.writer.close() + + def _save_model(self, model_path: str, log: bool = True): + if log: + print(f"[SAVE]: Saving the model to: {model_path}...") + + with open(model_path, "wb") as f: + f.write( + flax.serialization.to_bytes( + [ + vars(self.args), + [ + self.agent_state.params["sensor_params"], + self.agent_state.params["actor_params"], + self.agent_state.params["critic_params"], + self.agent_state.params["feature_extractor_params"], + ], + ] + ) + ) + + def train(self, log: bool = True): + """ + Train the PPO agent for a specified number of iterations + (passed through PPOArgs in constructor). + Closes the environment at the end of training. + """ + if log: + print(f"running name: {self.run_name}") + + is_tty = sys.stdout.isatty() + if log: + print("[TRAIN]: Resetting environment...") + + if not is_tty: + print(f">>> [HPC] Initial reset started: {time.ctime()}", flush=True) + + env_state = self.env.reset(seed=self.args.seed) + next_obs = _convert_obs_dict_to_array(env_state.observations) + next_done = jnp.zeros(self.args.num_envs, dtype=jnp.bool_) + + if log and not is_tty: + print(f">>> [HPC] Initial reset completed: {time.ctime()}", flush=True) + + global_step = 0 + start_time = time.time() + + if self.args.track: + import wandb + + if log: + print("[TRAIN]: Initializing Weights and Biases...") + + wandb.init( + project=self.args.wandb_project_name, + entity=self.args.wandb_entity, + sync_tensorboard=True, + config=vars(self.args), + name=self.run_name, + save_code=True, + ) + + if log: + print("[TRAIN]: Adding hyperparameters to TensorBoard...") + + self.writer.add_text( + "hyperparameters", + "|param|value|\n|---|---|\n" + + "\n".join(f"|{k}|{v}|" for k, v in vars(self.args).items()), + ) + + iter_bar = tqdm.tqdm( + range(1, self.args.num_iterations + 1), + disable=not is_tty, + ) + for iteration in iter_bar: + iteration_time_start = time.time() + + env_state, next_obs, next_done, loss_info = self._step( + env_state, next_obs, next_done, is_tty=is_tty, iteration=iteration + ) + + global_step += self.args.num_steps * self.args.num_envs + self._log(global_step, self.episode_stats, start_time, iteration_time_start, loss_info) + + if log and not is_tty: + sps = int(global_step / (time.time() - start_time)) + remaining_steps = self.args.total_timesteps - global_step + eta_seconds = int(remaining_steps / sps) if sps > 0 else 0 + eta_str = str(datetime.timedelta(seconds=eta_seconds)) + + print( + f"Iteration {iteration}/{self.args.num_iterations} | " + f"Step {global_step}/{self.args.total_timesteps} | " + f"SPS {sps} | " + f"Return {loss_info.avg_episodic_return:.4f} | " + f"ETA {eta_str}", + flush=True, + ) + + if self.args.save_model: + model_path = f"{self.run_dir}/{self.args.exp_name}.cleanrl_model" + self._save_model(model_path=model_path) + + self._close() diff --git a/src/brittle_star_project/trainers/__init__.py b/src/brittle_star_project/trainers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/main.py b/src/main.py deleted file mode 100644 index 61712ce..0000000 --- a/src/main.py +++ /dev/null @@ -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()}") diff --git a/src/ppo.py b/src/ppo.py new file mode 100644 index 0000000..cf6c69e --- /dev/null +++ b/src/ppo.py @@ -0,0 +1,159 @@ +from functools import partial + +import flax +import jax +import jax.numpy as jnp + + +# Chose to use a class as it seemed the easiest way to integrate the CleanRL code style +# with our need to seperate concerns +class PPO: + def __init__(self, args, sensor, actor, critic, feature_extractor, message_passer=None): + self.args = args + + if not message_passer: + message_passer = identity + + self.ppo_loss_grad_fn = jax.value_and_grad( + partial( + ppo_loss, + args=args, + sensor_apply=sensor.apply, + actor_apply=actor.apply, + critic_apply=critic.apply, + feature_extractor_apply=feature_extractor.apply, + message_passer=message_passer, + ), + has_aux=True, + ) + + # This PPO class should be initialized only once, + # or this function will need to recompile + @partial(jax.jit, static_argnums=0) + def update_ppo(self, agent_state, storage, key): + args = self.args + ppo_loss_grad_fn = self.ppo_loss_grad_fn + + def update_epoch(carry, _): + agent_state, key = carry + key, subkey = jax.random.split(key) + + def flatten(x): + return x.reshape((-1,) + x.shape[2:]) + + def convert_data(x): + x = jax.random.permutation(subkey, x) + return jnp.reshape(x, (args.num_minibatches, -1) + x.shape[1:]) + + flatten_storage = jax.tree.map(flatten, storage) + shuffled_storage = jax.tree.map(convert_data, flatten_storage) + + def update_minibatch(agent_state, minibatch): + (loss, (pg_loss, v_loss, entropy_loss, approx_kl)), grads = ppo_loss_grad_fn( + agent_state.params, + minibatch.obs, + minibatch.actions, + minibatch.logprobs, + minibatch.advantages, + minibatch.returns, + ) + agent_state = agent_state.apply_gradients(grads=grads) + return agent_state, ( + loss, + pg_loss, + v_loss, + entropy_loss, + approx_kl, + grads, + ) + + agent_state, metrics = jax.lax.scan(update_minibatch, agent_state, shuffled_storage) + return (agent_state, key), metrics + + (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads) = jax.lax.scan( + update_epoch, (agent_state, key), (), length=args.update_epochs + ) + return agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key + + +""" +Should be ok to use partial here, since the references to network, +actor and critic should not change at runtime +The cost of seperating concerns is to somehow pass these values +that are now not in the same scope +""" + + +@partial(jax.jit, static_argnums=(0, 1, 2, 3, 4)) +def get_action_and_value( + sensor_apply, + actor_apply, + message_passer, + critic_apply, + feature_extractor_apply, + params: flax.core.FrozenDict, + x: jnp.ndarray, + action: jnp.ndarray, +): + hidden_sensor = sensor_apply(params["sensor_params"], x) + hidden_critic = feature_extractor_apply(params["feature_extractor_params"], x) + hidden_sensor = message_passer(hidden_sensor) + mean, log_std = actor_apply(params["actor_params"], hidden_sensor) + std = jnp.exp(log_std) + + logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1) + entropy = (0.5 + 0.5 * jnp.log(2 * jnp.pi) + log_std).sum(-1) + value = critic_apply(params["critic_params"], hidden_critic).squeeze(-1) + + return logprob, entropy, value + + +def ppo_loss( + params, + x, + a, + logp, + mb_advantages, + mb_returns, + args, + sensor_apply, + actor_apply, + message_passer, + critic_apply, + feature_extractor_apply, +): + newlogprob, entropy, newvalue = get_action_and_value( + sensor_apply, + actor_apply, + message_passer, + critic_apply, + feature_extractor_apply, + params, + x, + a, + ) + logratio = newlogprob - logp + ratio = jnp.exp(logratio) + approx_kl = ((ratio - 1) - logratio).mean() + + if args.norm_adv: + mb_advantages = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8) + + pg_loss1 = -mb_advantages * ratio + pg_loss2 = -mb_advantages * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef) + pg_loss = jnp.maximum(pg_loss1, pg_loss2).mean() + v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean() + entropy_loss = entropy.mean() + loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef + return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl)) + + +def identity(hidden): + """ + Used for seamless jax integration, + avoids having branching inside jitted function, + used as message_passer in case it is not given, + (in case of centralized lvl) + """ + + return hidden diff --git a/src/train.py b/src/train.py deleted file mode 100644 index e40e751..0000000 --- a/src/train.py +++ /dev/null @@ -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() diff --git a/uv.lock b/uv.lock index 679acc3..42321e1 100644 --- a/uv.lock +++ b/uv.lock @@ -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"