refactor: use hydra for configs in simulation and training scripts

scripts/simulate.py

@@ -1,6 +1,14 @@
+"""Simulate a trained policy in the MuJoCo viewer.
+
+Uses Hydra to load the same BrittleStarConfig that was used during training.
+Override settings via CLI, e.g.:
+    python scripts/simulate.py morphology=3_arms
+"""
+
 from __future__ import annotations
 
-import argparse
+import hydra
+from omegaconf import DictConfig, OmegaConf
 from pathlib import Path
 
 from brittle_star_project import (
@@ -10,86 +18,58 @@ from brittle_star_project import (
     SimulationConfig,
     simulate_policy,
 )
-from brittle_star_project.environment import from_file
-from brittle_star_project.rl import RLModel  # imports concrete models via rl.__init__
+from brittle_star_project.configs.main_config import BrittleStarConfig
+from brittle_star_project.configs.register_configs import register_configs
+from brittle_star_project.rl import RLModel
 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()
+@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
+def main(dict_cfg: DictConfig) -> None:
+    cfg: BrittleStarConfig = OmegaConf.to_object(dict_cfg)
 
-
-def main() -> None:
-    args = parse_args()
-
-    morphology_cfg, arena_cfg, env_cfg = from_file("../configs/test.yaml")
+    # Allow overriding these via Hydra CLI or a dedicated simulate config group
+    # For now, defaults matching the old argparse behavior
+    backend = Backend.MJX
+    model_type = "random"
+    model_path = None
+    seed = cfg.experiment.seed
 
     # ======= 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)
+    raw_env = factory.create_environment(backend, cfg.morphology, cfg.arena, cfg.environment)
+    env = BrittleStarEnv(raw_env, backend=backend, config=cfg.environment)
 
-    seed_for_env = int(args.seed) if args.seed is not None else 0
-    state = env.reset(seed=seed_for_env)
+    state = env.reset(seed=seed)
 
     # ======= 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 model_path is not None:
+        policy = RLModel.load(Path(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)
+        model_cls = MODEL_BY_NAME[model_type]
+        policy = model_cls(seed=seed)
         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))
+    default_seed = int(getattr(policy, "seed", seed))
 
     # ======= SIMULATION =======
-
     rollout_cfg = SimulationConfig(
         realtime=True,
-        seed=int(args.seed) if args.seed is not None else default_seed,
+        seed=default_seed,
     )
 
     simulate_policy(policy, rollout_cfg, state)
-
     env.close()
 
 
 if __name__ == "__main__":
+    register_configs()
     main()