diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index be1fa12..623526f 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -34,8 +34,7 @@ }, "remoteUser": "vscode", "runArgs": [ - "--gpus", - "all" + "--device", "nvidia.com/gpu=all" ], // Ensure the .venv persists using a named volume for performance and parity "mounts": [ @@ -46,4 +45,4 @@ "features": { "ghcr.io/devcontainers/features/common-utils:1": {} } -} \ 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..7e09b14 --- /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/export_hpc_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 8204ce9..b0567ce 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ # Model files artifacts/* runs/* +wandb/ # Python-generated files __pycache__/ @@ -372,7 +373,6 @@ celerybeat.pid # Environments .env .venv -env/ venv/ ENV/ env.bak/ @@ -472,7 +472,6 @@ tags [Ll]ib [Ll]ib64 [Ll]ocal -[Ss]cripts pyvenv.cfg .venv pip-selfcheck.json @@ -519,3 +518,4 @@ Icon Network Trash Folder Temporary Items .apdisk +*.pdf 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 47b0552..085f3d5 100644 --- a/README.md +++ b/README.md @@ -15,3 +15,7 @@ example command: ```bash uv run src/train.py --model_name my_model --epochs 50 --batch_size 32 ``` + +## HPC + +See **[docs/HPC.md](docs/HPC.md)** for the full guide, including environment setup, cluster selection, interactive debugging, and job submission. diff --git a/configs/hpc/smoke_test.yaml b/configs/hpc/smoke_test.yaml new file mode 100644 index 0000000..720f829 --- /dev/null +++ b/configs/hpc/smoke_test.yaml @@ -0,0 +1,11 @@ +# Minimal config to verify HPC setup is functional. +# Run with: python src/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 new file mode 100644 index 0000000..215b2be --- /dev/null +++ b/configs/production_training.yaml @@ -0,0 +1,24 @@ +# Full PPO training config for Brittle Star (HPC Production) +exp_name: "production_training" +seed: 1 +track: true +capture_video: true +save_model: true +checkpoint_frequency: 100 # not yet implemented in train.py but here for future use + +# Scaling for HPC (using A100 GPU slices) +num_envs: 128 +total_timesteps: 10000000 +num_steps: 128 +num_minibatches: 4 +update_epochs: 4 + +# Algorithm +learning_rate: 2.5e-4 +anneal_lr: true +gamma: 0.99 +gae_lambda: 0.95 +clip_coef: 0.1 +ent_coef: 0.01 +vf_coef: 0.5 +cuda: true diff --git a/docs/CONTRIBUTING.md b/docs/CONTRIBUTING.md index efd1399..e694bb3 100644 --- a/docs/CONTRIBUTING.md +++ b/docs/CONTRIBUTING.md @@ -32,4 +32,4 @@ Code readability is paramount, as code is read far more frequently than it is wr * **Algorithms & Frameworks:** Proximal Policy Optimization (PPO) is the recommended baseline algorithm. CleanRL should be used as a starting point and adapted for continuous action spaces. All Artificial Neural Network (ANN) controller architectures must be implemented using Flax. * **Simulation:** The simulation environment utilizes a MuJoCo brittle star. XML MuJoCo structures must remain realistic and respect morphological constraints. * **Experiment Tracking:** Weights & Biases (wandb) must be utilized for tracking and logging all experiments. -* **Code Styling:** All code must conform to the chosen style guide (i.e. Google standard). This is enforced using build tools and pre-commit hooks such as flake8, black, or isort. +* **Code Styling:** All code must conform to the chosen style guide (Google standard). This is enforced via `uv` using **ruff** and pre-commit hooks. diff --git a/docs/HPC.md b/docs/HPC.md new file mode 100644 index 0000000..ff18a5f --- /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/export_hpc_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/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/experiments/_train_backup.py b/experiments/_train_backup.py index 6d6a181..2cc38ab 100644 --- a/experiments/_train_backup.py +++ b/experiments/_train_backup.py @@ -1,4 +1,8 @@ +import datetime import random +import yaml +import subprocess +import sys import time from dataclasses import asdict from functools import partial @@ -7,6 +11,7 @@ from typing import Callable import flax import jax import jax.numpy as jnp +import matplotlib.pyplot as plt import numpy as np import optax import torch @@ -35,11 +40,11 @@ def convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray: ) -def make_env(config_path: str | None, num_envs: int) -> Callable: +def make_env(env_config_path: str | None, num_envs: int) -> Callable: def thunk(): - if config_path is None: + if env_config_path is None: return BrittleStarJaxEnvWrapper.default(num_envs=num_envs) - return BrittleStarJaxEnvWrapper.from_config(config_path, num_envs=num_envs) + return BrittleStarJaxEnvWrapper.from_config(env_config_path, num_envs=num_envs) return thunk @@ -63,11 +68,29 @@ def save_model(model_path: str, agent_state: TrainState, args: PPOArgs): 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 + + # Try to get git short hash + try: + git_hash = ( + subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode("ascii").strip() + ) + except Exception: + git_hash = "none" + + run_name = f"{args.exp_name}__seed_{args.seed}__{git_hash}__{int(time.time())}" # args.num_iterations = args.total_timesteps // args.batch_size args.num_iterations = 5 run_name = f"{args.exp_name}__seed_{args.seed}__{int(time.time())}" print(f"running name: {run_name}") + if args.run_dir is None: + args.run_dir = f"runs/{run_name}" + + import os + + os.makedirs(args.run_dir, exist_ok=True) + if args.track: import wandb @@ -80,7 +103,7 @@ def train(args: PPOArgs): save_code=True, ) - writer = SummaryWriter(f"runs/{run_name}") + writer = SummaryWriter(args.run_dir) writer.add_text( "hyperparameters", "|param|value|\n|---|---|\n" + "\n".join(f"|{k}|{v}|" for k, v in vars(args).items()), @@ -96,7 +119,7 @@ def train(args: PPOArgs): print(f"Running on device: {device}") print("Creating the environment...") - env = make_env(config_path=args.config_path, num_envs=args.num_envs)() + env = make_env(env_config_path=args.env_config_path, num_envs=args.num_envs)() print(f"Environment: {env}") episode_stats = EpisodeStatistics( @@ -235,9 +258,13 @@ def train(args: PPOArgs): # Reset once to get initial state print("Resetting the environment...") + if not sys.stdout.isatty(): + print(f">>> [HPC] Initial reset started: {time.ctime()}", flush=True) 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_) + if not sys.stdout.isatty(): + print(f">>> [HPC] Initial reset completed: {time.ctime()}", flush=True) def step_once(carry, _, env_step_fn): agent_state, episode_stats, obs, done, key, env_state = carry @@ -277,28 +304,45 @@ def train(args: PPOArgs): ) print("Starting training...") - iters_bar = tqdm.tqdm(range(1, args.num_iterations + 1)) + iters_bar = tqdm.tqdm( + range(1, args.num_iterations + 1), + disable=not sys.stdout.isatty(), + ) returns = [] - for _ in iters_bar: + is_tty = sys.stdout.isatty() + for iteration in iters_bar: iteration_time_start = time.time() + if not is_tty and iteration == 1: + print(f">>> [HPC] Starting first rollout (JIT): {time.ctime()}", flush=True) + 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 ) + if not is_tty and iteration == 1: + print(f">>> [HPC] First rollout completed: {time.ctime()}", flush=True) + global_step += args.num_steps * args.num_envs storage = compute_gae(agent_state, next_obs, next_done, storage) + + if not is_tty and iteration == 1: + print(f">>> [HPC] Starting first PPO update (JIT): {time.ctime()}", flush=True) + agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key = ppo_instance.update_ppo( agent_state, storage, key ) + if not is_tty and iteration == 1: + print(f">>> [HPC] First PPO update completed: {time.ctime()}", flush=True) + + losses.append(jnp.mean(loss)) + avg_episodic_return = np.mean(jax.device_get(episode_stats.returned_episode_returns)) iters_bar.set_postfix_str( f"global_step={global_step}, avg_episodic_return={avg_episodic_return}" ) - returns.append(avg_episodic_return) - writer.add_scalar("charts/avg_episodic_return", avg_episodic_return, global_step) writer.add_scalar( "charts/avg_episodic_length", @@ -325,9 +369,39 @@ def train(args: PPOArgs): global_step, ) + if not is_tty: + sps = int(global_step / (time.time() - start_time)) + remaining_steps = 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}/{args.num_iterations} | " + f"Step {global_step}/{args.total_timesteps} | " + f"SPS {sps} | " + f"Return {avg_episodic_return:.4f} | " + f"ETA {eta_str}", + flush=True, + ) + if args.save_model: model_path = f"runs/{run_name}/{args.exp_name}.cleanrl_model" save_model(model_path, agent_state, args) + model_path = f"{args.run_dir}/{args.exp_name}.cleanrl_model" + with open(model_path, "wb") as f: + f.write( + flax.serialization.to_bytes( + [ + vars(args), + [ + agent_state.params["sensor_params"], + agent_state.params["actor_params"], + agent_state.params["critic_params"], + agent_state.params["feature_extractor_params"], + ], + ] + ) + ) print(f"model saved to {model_path}") env.close() @@ -340,10 +414,29 @@ def train(args: PPOArgs): show_window=True, filename=f"runs/{run_name}/{args.exp_name}_losses.png", ) + plt.plot(losses) + plt.title("PPO Loss, mean over minibatches") + plt.savefig(f"{args.run_dir}/{args.exp_name}_losses.png") + plt.close() def main() -> None: - args = tyro.cli(PPOArgs) + temp_args = tyro.cli(PPOArgs) + + if temp_args.env_config_path is not None: + with open(temp_args.env_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) + + # Re-parse CLI to ensure they OVERRIDE the yaml + args = tyro.cli(PPOArgs, default=temp_args) + else: + args = temp_args + train(args) diff --git a/pyproject.toml b/pyproject.toml index 43c977b..171bae1 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", @@ -19,12 +23,16 @@ dependencies = [ "pyopengl-accelerate>=3.1.10", "tyro>=1.0.10", "wandb==0.24.2", + "torch>=2.4.0", ] [project.optional-dependencies] cuda = [ "jax[cuda13]==0.9.0.1", ] +analysis = [ + "tensorboard", +] [dependency-groups] dev = [ diff --git a/ruff.toml b/ruff.toml index 3149b75..d665ebb 100644 --- a/ruff.toml +++ b/ruff.toml @@ -1,4 +1,5 @@ line-length = 100 +exclude = ["wandb"] [lint] extend-select = [ 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/export_hpc_requirements.py b/scripts/export_hpc_requirements.py new file mode 100644 index 0000000..9ff3de7 --- /dev/null +++ b/scripts/export_hpc_requirements.py @@ -0,0 +1,86 @@ +#!/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. + +Usage: + uv run scripts/export_hpc_requirements.py +""" + +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/src/brittle_star_project/dataclasses/PPOArgs.py b/src/brittle_star_project/dataclasses/PPOArgs.py index 8d13148..f6b20bf 100644 --- a/src/brittle_star_project/dataclasses/PPOArgs.py +++ b/src/brittle_star_project/dataclasses/PPOArgs.py @@ -11,11 +11,17 @@ class PPOArgs: """ # path to environment config file, if None, use default config - config_path: str | None = None + env_config_path: str | None = None # 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 diff --git a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py index 259f1a5..7c86d51 100644 --- a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py +++ b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py @@ -8,7 +8,7 @@ from brittle_star_project import ( ArenaConfig, Backend, ) -from brittle_star_project.environment import from_json +from brittle_star_project.environment import from_file class BrittleStarJaxEnvWrapper: @@ -82,7 +82,7 @@ class BrittleStarJaxEnvWrapper: def from_config( config_path: str, num_envs: int, backend: Backend = Backend.MJX ) -> "BrittleStarJaxEnvWrapper": - morphology_cfg, arena_cfg, env_cfg = from_json(config_path) + morphology_cfg, arena_cfg, env_cfg = from_file(config_path) return BrittleStarJaxEnvWrapper( morphology_cfg, arena_cfg, env_cfg, num_envs=num_envs, backend=backend ) diff --git a/src/brittle_star_project/environment/__init__.py b/src/brittle_star_project/environment/__init__.py index f87eae5..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_json +from .env_config import ArenaConfig, EnvConfig, MorphologyConfig, from_file from .env_types import Backend, Task from .env_wrapper import BrittleStarEnv, StepResult from .factory import BrittleStarEnvFactory @@ -12,5 +12,5 @@ __all__ = [ "BrittleStarEnv", "StepResult", "BrittleStarEnvFactory", - "from_json", + "from_file", ] diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index 6f55f35..a693286 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/src/brittle_star_project/environment/env_config.py @@ -50,10 +50,17 @@ class EnvConfig: light_perlin_noise_scale: int = 0 -def from_json(path: str) -> tuple[MorphologyConfig, ArenaConfig, EnvConfig]: +def from_file(path: str) -> tuple[MorphologyConfig, ArenaConfig, EnvConfig]: + """Load configurations from a JSON or YAML file.""" with open(path, "r") as f: - config_json = json.load(f) - morphology = MorphologyConfig(**config_json.get("morphology", {})) - arena = ArenaConfig(**config_json.get("arena", {})) - env = EnvConfig(**config_json.get("env", {})) + 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/uv.lock b/uv.lock index e59c767..7fbaacb 100644 --- a/uv.lock +++ b/uv.lock @@ -22,11 +22,16 @@ dependencies = [ { name = "jax" }, { name = "matplotlib" }, { name = "mediapy" }, + { name = "mujoco-warp" }, + { name = "numpy" }, { name = "optax" }, + { name = "protobuf" }, { name = "pyopengl" }, { name = "pyopengl-accelerate" }, + { name = "torch" }, { name = "tyro" }, { name = "wandb" }, + { name = "warp-lang" }, ] [package.optional-dependencies] @@ -53,11 +58,16 @@ 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 = "torch", specifier = ">=2.4.0" }, { name = "tyro", specifier = ">=1.0.10" }, { name = "wandb", specifier = "==0.24.2" }, + { name = "warp-lang" }, ] provides-extras = ["cuda"] @@ -1262,6 +1272,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c2/7c/ad82beb7c4c9186d9fbef4799109d799692d70276bb1b3ee18a0674170d8/mujoco_mjx-3.6.0-py3-none-any.whl", hash = "sha256:c81000af0653f162b76009f48c153e9e6d19bfa8febe851e12466c81cbb7336a", size = 7013366, upload-time = "2026-03-11T01:46:19.148Z" }, ] +[[package]] +name = "mujoco-warp" +version = "3.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "etils", extra = ["epath"] }, + { name = "mujoco" }, + { name = "numpy" }, + { name = "warp-lang" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0e/de/b853418268e9777cad2792ee3a145c8397e3d4517d136499645847ffd7f2/mujoco_warp-3.6.0.tar.gz", hash = "sha256:3c4111a4e13dc61268ddac52593ac5032c05a7d80f0c5e3c98bf5881e32b5d06", size = 1887269, upload-time = "2026-03-11T01:11:44.343Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a3/b5/06c1e23c0cc4a06da268aa5f0fe05348d89c9ecbd3ecfb4c6d2b14ea23b2/mujoco_warp-3.6.0-py3-none-any.whl", hash = "sha256:371a405b186332cbfaa9630aabf35967405ef847cb7ea7c018ec67426a2ea160", size = 1965960, upload-time = "2026-03-11T01:11:42.527Z" }, +] + [[package]] name = "mypy-extensions" version = "1.1.0" @@ -1630,7 +1656,7 @@ name = "pexpect" version = "4.9.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "ptyprocess" }, + { name = "ptyprocess", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" } wheels = [ @@ -2428,6 +2454,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"