feat(hpc): support run_dir staging and update docs

- Add run_dir and checkpoint_frequency to PPOArgs
- Update train.py to use run_dir for SummaryBoard, model saving, and loss plots
- Create configs/production_training.yaml for HPC production runs
- Update HPC.md with run_dir staging strategy details

configs/production_training.yaml

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

docs/HPC.md

@@ -6,39 +6,15 @@ Full documentation: <https://docs.hpc.ugent.be/>
 
 Choose the appropriate cluster before submitting a job with `module swap cluster/<name>`. The default login cluster is **doduo**.
 
-| Cluster   | Type                    | Use case                                      |
-|-----------|-------------------------|-----------------------------------------------|
-| `donphan` | Interactive / debug GPU | First-time setup, interactive debugging       |
-| `doduo`   | CPU (default login)     | Rapid iteration, CPU-only smoke tests         |
-| `joltik`  | GPU (A100 ¼-slice)      | Standard training runs                        |
-| `accelgor`| GPU (A100 full)         | Large-scale / long experiments                |
-| `litleo`  | GPU                     | Alternative GPU option                        |
-
-> **Rule:** use at most **1 GPU per group at a time** on shared GPU clusters.  
 > Check the current queue load at <https://shieldon.ugent.be:8083/pbsmon-web-users/>.
 
----
-
 ## Storage Overview
 
-The HPC provides three filesystems for different purposes. Understanding this is critical to avoid filling up your home directory.
-
-| Variable        | Typical size | Purpose                                              |
-|-----------------|-------------|------------------------------------------------------|
-| `$VSC_HOME`     | ~3 GB       | Config files, SSH keys, project source code only     |
-| `$VSC_DATA`     | ~25 GB      | Persistent outputs: trained models, final logs       |
-| `$VSC_SCRATCH`  | Large       | Fast I/O during jobs: caches, intermediate files     |
-
-**Important:**
-- Clone the repository into `$VSC_HOME` — it is small in size and accessible from all clusters.
-- All caches (pip, uv, matplotlib) must be redirected to `$VSC_SCRATCH` to avoid filling `$VSC_HOME`.
 - Run outputs are written to `$VSC_SCRATCH` during the job (fast I/O) and copied to `$VSC_DATA` at the end for persistence.
 - `$VSC_SCRATCH` may be purged periodically — do not use it as long-term storage.
 
 Check your quota: <https://account.vscentrum.be> (Usage section).
 
----
-
 ## Initial Environment Setup
 
 Run **once** from a login shell on `donphan` after cloning the repository:
@@ -56,21 +32,10 @@ cd 2026SEL3-project-BrittleStar
 bash scripts/hpc/install.sh
 ```
 
-This uses the official [`vsc-venv`](https://docs.hpc.ugent.be/Linux/setting_up_python_virtual_environments/#vsc-venv-python-virtual-environment-wrapper-script) wrapper to:
-- Redirect caches to `$VSC_SCRATCH` (to preserve your `$VSC_HOME` quota)
-- Load the EasyBuild modules listed in `env/hpc/modules.txt` (JAX, Flax, WandB, …)
-- Create a per-cluster virtual environment in `$VSC_DATA`
-- Pip-install the remaining packages from `env/hpc/requirements.txt`
-- Register a Jupyter kernel named `SEL3 (<cluster>)`
-
 > **Note:** virtual environments are cluster-specific. Re-run the script when switching to a new cluster.
 
----
-
 ## Interactive Debugging on donphan
 
-The `donphan` cluster provides quick access and is ideal for verifying your environment before submitting batch jobs.
-
 ### Option A — Interactive shell session
 
 ```bash
@@ -121,8 +86,6 @@ python src/train.py --config configs/hpc/smoke_test.yaml
 
 > **Warning:** JAX can only be loaded by one kernel at a time. Shut down other kernels before switching notebooks.
 
----
-
 ## Submitting Batch Training Jobs
 
 ```bash
@@ -135,10 +98,9 @@ qsub scripts/hpc/train.pbs
 ```
 
 The job script automatically:
-- Redirects all caches to `$VSC_SCRATCH`
-- Writes run outputs to `$VSC_SCRATCH/runs/<job_id>` during the run
-- Copies final results to `$VSC_DATA/runs/<job_id>` on completion
-- Writes PBS stdout/stderr to `$VSC_DATA/runs/<job_id>/job.out` / `job.err`
+- Writes run outputs to `$VSC_SCRATCH/runs/<job_id>` during the run using the `--run_dir` argument. This ensures that frequent I/O (like tensorboard logs and checkpoints) happens on the fastest available filesystem.
+- Copies the final results to `$VSC_DATA/runs/<job_id>` on completion for long-term persistence.
+- Writes PBS stdout/stderr to `runs/brittlestar-ppo.o<job_id>` / `.e<job_id>` (standard PBS convention, relative to the project root).
 
 Monitor your jobs:
 
@@ -148,8 +110,6 @@ qstat -f <id>  # detailed info for a specific job
 qdel <id>      # cancel a job
 ```
 
----
-
 ## Managing Dependencies
 
 `env/hpc/requirements.txt` is auto-generated by CI whenever `pyproject.toml` changes. To regenerate locally:

src/brittle_star_project/dataclasses/PPOArgs.py

@@ -13,6 +13,12 @@ class PPOArgs:
     # 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
 

src/main.py

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

src/train.py

@@ -46,6 +46,12 @@ def train(args: PPOArgs):
     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
 
@@ -58,7 +64,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()),
@@ -297,7 +303,7 @@ def train(args: PPOArgs):
         )
 
     if args.save_model:
-        model_path = f"runs/{run_name}/{args.exp_name}.cleanrl_model"
+        model_path = f"{args.run_dir}/{args.exp_name}.cleanrl_model"
         with open(model_path, "wb") as f:
             f.write(
                 flax.serialization.to_bytes(
@@ -319,7 +325,7 @@ def train(args: PPOArgs):
     print("Saving loss plot...")
     plt.plot(losses)
     plt.title("PPO Loss, mean over minibatches")
-    plt.savefig(f"runs/{run_name}/{args.exp_name}_losses.png")
+    plt.savefig(f"{args.run_dir}/{args.exp_name}_losses.png")
     plt.close()