Merge branch 'dev' into simulate-results

scripts/analysis/README.md

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

scripts/analysis/explore_tensorboard.py

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

scripts/hpc/export_requirements.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+"""Export HPC pip requirements from pyproject.toml.
+
+This is a LOCAL DEVELOPER UTILITY — run it on your own machine before pushing
+code whenever pyproject.toml dependencies change. It reads the modules from
+env/hpc/modules.txt and the full dependency list from pyproject.toml, then
+writes the remainder to env/hpc/requirements.txt.
+"""
+
+from __future__ import annotations
+
+import re
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[2]
+
+
+def normalise(name: str) -> str:
+    """Normalise a PyPI package name for comparison."""
+    return re.sub(r"[-_.]+", "-", name).lower()
+
+
+def pkg_name(dep: str) -> str:
+    """Extract the bare package name from a PEP 508 dependency string."""
+    return re.split(r"[\[=><~!;]", dep)[0].strip()
+
+
+def main() -> None:
+    import tomllib
+
+    modules_path = ROOT / "env" / "hpc" / "modules.txt"
+    if not modules_path.exists():
+        print(f"Error: {modules_path} not found.", file=sys.stderr)
+        sys.exit(1)
+
+    # Read normalized module names from base modules only
+    # Library modules (like PyTorch) are kept in requirements for portability
+    module_names = [
+        normalise(line.split()[0].split("/")[0])
+        for line in modules_path.read_text().splitlines()
+        if line.strip() and not line.startswith("#")
+    ]
+
+    pyproject_path = ROOT / "pyproject.toml"
+    with pyproject_path.open("rb") as f:
+        data = tomllib.load(f)
+
+    # Collect all dependencies, merging 'cuda' extras into base dependencies
+    dep_dict: dict[str, str] = {}
+    for dep in data.get("project", {}).get("dependencies", []):
+        dep_dict[normalise(pkg_name(dep))] = dep
+
+    # Add cuda extras (takes precedence for HPC)
+    optional_deps = data.get("project", {}).get("optional-dependencies", {})
+    for group in ["cuda"]:
+        for dep in optional_deps.get(group, []):
+            dep_dict[normalise(pkg_name(dep))] = dep
+
+    deps = list(dep_dict.values())
+
+    final_deps: list[str] = []
+    print("Checking dependencies against HPC module list...", file=sys.stderr)
+    for dep in deps:
+        name = normalise(pkg_name(dep))
+        # Smart check: if the package name is a substring of any loaded module name
+        # (e.g. 'torch' in 'pytorch', 'scipy' in 'scipy-bundle')
+        if any(name in mod for mod in module_names):
+            print(f"  [skip – module provider found] {dep}", file=sys.stderr)
+            continue
+
+        final_deps.append(dep)
+        print(f"  [pip] {dep}", file=sys.stderr)
+
+    hpc_dir = ROOT / "env" / "hpc"
+    output_path = hpc_dir / "requirements.txt"
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    output_path.write_text("\n".join(final_deps) + "\n")
+    print(f"\nWrote {len(final_deps)} requirement(s) to {output_path}", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    main()

scripts/hpc/install.sh

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

scripts/hpc/train.pbs

@@ -0,0 +1,75 @@
+# 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"
+
+export PYTHONPATH="$PBS_O_WORKDIR/src:${PYTHONPATH:-}"
+
+if [ -f "$VSC_DATA/$PROJ_NAME/.env" ]; then
+    echo ">>> Sourcing API keys from .env..."
+    export $(grep -v '^#' "$VSC_DATA/$PROJ_NAME/.env" | xargs)
+elif [ -f "$PBS_O_WORKDIR/.env" ]; then
+    echo ">>> Sourcing API keys from .env..."
+    export $(grep -v '^#' "$PBS_O_WORKDIR/.env" | xargs)
+fi
+
+# TODO Once experiments get serious, change the config
+python scripts/train.py \
+    --env-config-path configs/hpc/wandb_expand.yaml \
+    --hyperparameter-config-path configs/hpc/wandb_expand.yaml \
+    --run-dir "$SCRATCH_RUNDIR"
+
+echo ">>> Staging out results to $DATA_RUNDIR..."
+cp -r "$SCRATCH_RUNDIR/." "$DATA_RUNDIR/"
+
+echo ">>> Done"

scripts/simulate.py

@@ -0,0 +1,390 @@
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+from typing import Any
+import flax
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from brittle_star_project import (
+    Backend,
+)
+from brittle_star_project.environment import from_file
+
+def _flatten_obs_dict(obs_dict: dict[str, Any]) -> jnp.ndarray:
+    """Flatten the env's observation dict into a 1D vector.
+
+    concatenates values in the dict's iteration order and skips empty arrays.
+    """
+
+    parts: list[jnp.ndarray] = []
+    for v in obs_dict.values():
+        arr = jnp.asarray(v)
+        if arr.size == 0:
+            continue
+        parts.append(arr.reshape((-1,)))
+
+    if not parts:
+        return jnp.zeros((0,), dtype=jnp.float32)
+    return jnp.concatenate(parts, axis=0)
+
+
+# A minimal policy class to load a CleanRL/Flax checkpoint and run inference.
+class CleanRLPPOPolicy:
+
+    def __init__(
+        self,
+        *,
+        network_params: Any,
+        actor_params: Any,
+        action_dim: int,
+    ) -> None:
+        from brittle_star_project.rl import Actor, Network
+
+        self._network = Network()
+        self._actor = Actor(action_dim=action_dim)
+        self._network_apply = jax.jit(self._network.apply)
+        self._actor_apply = jax.jit(self._actor.apply)
+        self._params = {
+            "network_params": network_params,
+            "actor_params": actor_params,
+        }
+
+    @staticmethod
+    def load(
+        path: Path,
+        *,
+        action_dim: int,
+    ) -> "CleanRLPPOPolicy":
+        def _get_index(container: Any, idx: int) -> Any:
+            if isinstance(container, (list, tuple)):
+                return container[idx]
+            if isinstance(container, dict):
+                return container.get(idx, container.get(str(idx)))
+            raise KeyError(idx)
+
+        def _looks_like_indexed_dict(container: Any) -> bool:
+            return (
+                isinstance(container, dict)
+                and container
+                and all(str(k).isdigit() for k in container.keys())
+            )
+
+        def _parse_checkpoint(restored_obj: Any) -> tuple[Any, Any, Any, Any]:
+            """Extract (args_dict, network_params, actor_params, critic_params).
+
+            `src/train.py` saves:
+              flax.serialization.to_bytes([vars(args), [net, actor, critic]])
+
+            `msgpack_restore()` occasionally restores lists as dicts keyed by
+            string indices ("0", "1", ...), so we accept both shapes.
+            """
+
+            args_part: Any | None = None
+            params_part: Any = restored_obj
+
+            if isinstance(restored_obj, (list, tuple)) and len(restored_obj) >= 2:
+                args_part = restored_obj[0]
+                params_part = restored_obj[1]
+            elif _looks_like_indexed_dict(restored_obj) and (
+                "0" in restored_obj or "1" in restored_obj
+            ):
+                args_part = restored_obj.get("0", restored_obj.get(0))
+                params_part = restored_obj.get("1", restored_obj.get(1))
+
+            if _looks_like_indexed_dict(params_part):
+                network_params = _get_index(params_part, 0)
+                actor_params = _get_index(params_part, 1)
+                critic_params = _get_index(params_part, 2)
+                if network_params is None or actor_params is None:
+                    raise ValueError("Missing required params in checkpoint")
+                return args_part, network_params, actor_params, critic_params
+
+            if isinstance(params_part, (list, tuple)) and len(params_part) >= 2:
+                network_params = params_part[0]
+                actor_params = params_part[1]
+                critic_params = params_part[2] if len(params_part) >= 3 else None
+                return args_part, network_params, actor_params, critic_params
+
+            raise ValueError(
+                f"Unexpected .cleanrl_model structure in {path}. "
+                "Expected [args_dict, [network_params, actor_params, critic_params]] "
+                "or an equivalent dict-indexed variant."
+            )
+
+        payload = path.read_bytes()
+        restored = flax.serialization.msgpack_restore(payload)
+        _args_dict, network_params, actor_params, _critic_params = _parse_checkpoint(restored)
+
+        return CleanRLPPOPolicy(
+            network_params=network_params,
+            actor_params=actor_params,
+            action_dim=action_dim,
+        )
+
+    def act(self, *, observations: dict[str, Any]) -> np.ndarray:
+        obs = _flatten_obs_dict(observations)
+        hidden = self._network_apply(self._params["network_params"], obs)
+        mean, _log_std = self._actor_apply(self._params["actor_params"], hidden)
+
+        # Always evaluate with the actor mean.
+        # (Sampling adds exploration noise, which is useful for training but not for evaluation.)
+        return np.asarray(mean, dtype=np.float32).ravel()
+
+
+def _get_observations(state: Any) -> dict[str, Any]:
+    return getattr(state, "observations", None)
+
+
+def _get_xy_distance_to_target(observations: dict[str, Any]) -> float | None:
+    return float(np.asarray(observations["xy_distance_to_target"]).reshape(-1)[0])
+
+
+def _target_reached(*, state: Any) -> bool:
+    return bool(getattr(state, "terminated", False))
+
+
+def _rollout_one_episode_headless(
+    *,
+    env: Any,
+    policy: CleanRLPPOPolicy,
+    seed: int,
+    max_steps: int,
+) -> tuple[float, int, bool, float | None]:
+    """Run one rollout up to `max_steps`.
+
+    Returns (return, length, reached_target, final_xy_dist).
+    """
+    state = env.reset(seed=seed)
+
+    ep_return = 0.0
+
+    observations = _get_observations(state)
+    prev_dist = _get_xy_distance_to_target(observations)
+    reached_target = _target_reached(state=state)
+
+    # NOTE: In the MJC backend, `state.reward` is always 0.0.
+    # To get a meaningful return, we compute a simple progress reward:
+    #   r_t = d_{t-1} - d_t
+    # where d is `xy_distance_to_target`.
+    steps = 0
+    for _ in range(int(max_steps)):
+        action = policy.act(observations=observations)
+
+        nu = int(state.mj_model.nu)
+        if nu > 0 and action.shape != (nu,):
+            raise ValueError(f"Policy returned action shape {action.shape}, expected ({nu},)")
+
+        state = env.step(state=state, action=action)
+        steps += 1
+
+        observations = _get_observations(state)
+        cur_dist = _get_xy_distance_to_target(observations)
+        if prev_dist is not None and cur_dist is not None:
+            ep_return += prev_dist - cur_dist
+        prev_dist = cur_dist
+
+        reached_target = _target_reached(state=state)
+        if reached_target:
+            break
+
+    final_dist = _get_xy_distance_to_target(observations)
+    return ep_return, steps, reached_target, final_dist
+
+
+def _run_one_episode_viewer(
+    *,
+    env: Any,
+    policy: CleanRLPPOPolicy,
+    seed: int,
+    state: Any,
+    control_dt: float,
+    max_steps: int,
+) -> None:
+    import mujoco.viewer
+
+    model = state.mj_model
+    data = state.mj_data
+
+    seed = int(seed)
+    episode_return = 0.0
+    observations = _get_observations(state)
+    prev_dist = _get_xy_distance_to_target(observations)
+    reached_target = _target_reached(state=state)
+
+    viewer = mujoco.viewer.launch_passive(model, data)
+    try:
+        steps = 0
+        for _step_idx in range(int(max_steps)):
+            if not viewer.is_running():
+                break
+            step_start = time.time()
+
+            # One control step. We do the env step under the viewer lock.
+            action = policy.act(observations=observations)
+            if model.nu > 0 and action.shape != (int(model.nu),):
+                raise ValueError(
+                    f"Policy returned action shape {action.shape}, expected ({int(model.nu)},)"
+                )
+            # The passive viewer runs a GUI thread; protect MuJoCo state mutation.
+            with viewer.lock():
+                state = env.step(state=state, action=action)
+            if not viewer.is_running():
+                break
+            viewer.sync()
+
+            steps += 1
+
+            observations = _get_observations(state)
+            cur_dist = _get_xy_distance_to_target(observations)
+            if prev_dist is not None and cur_dist is not None:
+                episode_return += prev_dist - cur_dist
+            prev_dist = cur_dist
+
+            reached_target = _target_reached(state=state)
+            if reached_target:
+                break
+
+            # Real-time pacing so the viewer doesn't run as fast as possible.
+            remaining = control_dt - (time.time() - step_start)
+            if remaining > 0:
+                time.sleep(remaining)
+
+        # Done: target reached, fixed horizon reached, or window closed.
+        if viewer.is_running():
+            dist = _get_xy_distance_to_target(observations)
+            dist_str = "n/a" if dist is None else f"{dist:.3f}"
+            print(
+                "episode done: "
+                f"return={episode_return:.6f}, len={steps}, "
+                f"target_reached={reached_target}, final_xy_dist={dist_str}"
+            )
+            viewer.close()
+    finally:
+        # Ensure the GUI thread stops before the env/model/data are torn down.
+        try:
+            viewer.close()
+        except Exception:
+            pass
+        for _ in range(200):
+            if not viewer.is_running():
+                break
+            time.sleep(0.01)
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description="Run a trained policy for exactly one episode (viewer or headless)."
+    )
+    p.add_argument(
+        "--model",
+        type=str,
+        required=True,
+        help=("Path to a CleanRL/Flax '.cleanrl_model' checkpoint (saved by src/train.py)."),
+    )
+    p.add_argument(
+        "--headless",
+        action="store_true",
+        help="Run without the MuJoCo viewer (still exactly one episode).",
+    )
+    p.add_argument(
+        "--max-steps",
+        type=int,
+        required=True,
+        help=(
+            "Number of control steps to run (fixed horizon). "
+            "This script stops when this many steps are reached, or earlier if "
+            "the target is reached (directed locomotion)."
+        ),
+    )
+    p.add_argument(
+        "--backend",
+        choices=[b for b in Backend],
+        default=Backend.MJC,
+    )
+    p.add_argument("--seed", type=int, default=0)
+    return p.parse_args()
+
+
+def main() -> None:
+    from brittle_star_project.environment import (
+        BrittleStarEnv,
+        BrittleStarEnvFactory,
+    )
+
+    args = parse_args()
+
+    morphology_cfg, arena_cfg, env_cfg = from_file("../configs/test.yaml")
+
+    # ======= ENVIRONMENT SETUP =======
+
+    backend = args.backend
+
+    factory = BrittleStarEnvFactory()
+    raw_env = factory.create_environment(backend, morphology_cfg, arena_cfg, env_cfg)
+    env = BrittleStarEnv(raw_env, backend=backend, config=env_cfg)
+
+    seed_for_env = int(args.seed) if args.seed is not None else 0
+    state = env.reset(seed=seed_for_env)
+
+    # ======= MODEL SETUP =======
+
+    # Extract the number of actuators (nu) from the environment's model, so we can pass it to the
+    # policy/model.
+    nu = int(state.mj_model.nu)
+
+    model_path = Path(args.model)
+    if model_path.suffix != ".cleanrl_model":
+        raise ValueError(f"Expected a '.cleanrl_model' checkpoint, got '{model_path.name}'.")
+
+    policy = CleanRLPPOPolicy.load(
+        model_path,
+        action_dim=nu,
+    )
+
+    default_seed = seed_for_env
+
+    # ======= SIMULATION =======
+
+    if args.headless:
+        max_steps = int(args.max_steps)
+        if max_steps <= 0:
+            raise ValueError("--max-steps must be > 0")
+
+        ep_seed = int(args.seed) if args.seed is not None else default_seed
+        ep_return, ep_len, reached_target, final_dist = _rollout_one_episode_headless(
+            env=env,
+            policy=policy,
+            seed=ep_seed,
+            max_steps=max_steps,
+        )
+        final_dist_str = "n/a" if final_dist is None else f"{final_dist:.3f}"
+        print(
+            "episode done: "
+            f"return={ep_return:.6f}, len={ep_len}, "
+            f"target_reached={reached_target}, final_xy_dist={final_dist_str}"
+        )
+    else:
+        max_steps = int(args.max_steps)
+        if max_steps <= 0:
+            raise ValueError("--max-steps must be > 0")
+
+        model_dt = float(state.mj_model.opt.timestep)
+        control_dt = model_dt * float(env_cfg.num_physics_steps_per_control_step)
+        _run_one_episode_viewer(
+            env=env,
+            policy=policy,
+            seed=int(args.seed) if args.seed is not None else default_seed,
+            state=state,
+            control_dt=control_dt,
+            max_steps=max_steps,
+        )
+
+    env.close()
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,75 @@
+import subprocess
+import time
+
+import torch
+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
+
+from experiment_logger import UnifiedLogger
+from experiment_logger.config_utils import merge_config_with_cli, print_config
+
+
+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() -> PPOArgs:
+    import argparse
+
+    # Use argparse to reliably extract just the config path without swallowing --help
+    parser = argparse.ArgumentParser(add_help=False)
+    parser.add_argument("--hyperparameter-config-path", type=str, default=None)
+    known_args, _ = parser.parse_known_args()
+
+    args = merge_config_with_cli(PPOArgs, config_file=known_args.hyperparameter_config_path)
+    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)
+
+    # Initialize Global Logger
+    logger = UnifiedLogger(
+        config=vars(args),
+        project_name=args.wandb_project_name,  # or default PPO-Modularity if missing
+        run_name=run_name,
+        base_dir=os.path.dirname(run_dir),
+        use_wandb=args.track,
+    )
+
+    print_config(args, title="PPO Training Configuration")
+
+    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()