restructure

scripts/hpc/export_requirements.py

@@ -5,9 +5,6 @@ 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

scripts/simulate.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from brittle_star_project import (
+    Backend,
+    BrittleStarEnv,
+    BrittleStarEnvFactory,
+    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
+
+MODEL_BY_NAME = get_rl_model_registry()
+MODEL_OPTIONS = sorted(MODEL_BY_NAME)
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Simulate a trained policy in the MuJoCo viewer.")
+    p.add_argument(
+        "--model",
+        type=str,
+        default=None,
+        help="Path to a saved model artifact. If omitted, a model is created from --model-type.",
+    )
+    p.add_argument(
+        "--model-type",
+        choices=MODEL_OPTIONS,
+        default="random",
+        help="Which model class to instantiate when --model is omitted.",
+    )
+    p.add_argument(
+        "--backend",
+        choices=[b for b in Backend],
+        default=Backend.MJX,
+    )
+    p.add_argument("--seed", type=int, default=None)
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    morphology_cfg, arena_cfg, env_cfg = from_json("../configs/test.json")
+
+    # ======= 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)
+
+    if args.model is not None:
+        model_path = Path(args.model)
+        policy = RLModel.load(model_path)
+        if hasattr(policy, "nu"):
+            policy.nu = nu
+    else:
+        model_cls = MODEL_BY_NAME[str(args.model_type)]
+        policy = model_cls(seed=seed_for_env)
+        if hasattr(policy, "nu"):
+            policy.nu = nu
+
+    # If the policy/model has a `seed` attribute, use the provided seed (or default) to reset it.
+    default_seed = int(getattr(policy, "seed", seed_for_env))
+    if args.seed is not None and hasattr(policy, "reset"):
+        policy.reset(int(args.seed))
+
+    # ======= SIMULATION =======
+
+    rollout_cfg = SimulationConfig(
+        realtime=True,
+        seed=int(args.seed) if args.seed is not None else default_seed,
+    )
+
+    simulate_policy(policy, rollout_cfg, state)
+
+    env.close()
+
+
+if __name__ == "__main__":
+    main()

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()