restructure

experiments/simulate.py

@@ -1,95 +0,0 @@
-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()