refactor: use hydra for configs in simulation and training scripts
This commit is contained in:
parent
ff83af8cef
commit
93bff11208
SSH key fingerprint:
SHA256:u/h/LVoqKF1Iz02uOyxe6hcjmoZASCGV2HM0TG9ZMoU
9 changed files with 87 additions and 236 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -1,12 +1,10 @@
|
|||
import os
|
||||
import time
|
||||
import torch
|
||||
import hydra
|
||||
from omegaconf import DictConfig, OmegaConf
|
||||
|
||||
from brittle_star_project.configs.main_config import BrittleStarConfig
|
||||
from brittle_star_project.configs.register_configs import register_configs
|
||||
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 init_logger, get_logger
|
||||
|
|
@ -22,52 +20,6 @@ def make_env(cfg: BrittleStarConfig) -> BrittleStarJaxEnvWrapper:
|
|||
)
|
||||
|
||||
|
||||
def create_ppo_args_compat(cfg: BrittleStarConfig, run_dir: str) -> PPOArgs:
|
||||
"""Temporary adapter to bridge BrittleStarConfig to the legacy PPOArgs.
|
||||
|
||||
This will be removed in Step 4.4 once PPOTrainer is refactored.
|
||||
"""
|
||||
# Flatten the hierarchical config into the expected PPOArgs format
|
||||
args = PPOArgs(
|
||||
exp_name=cfg.experiment.exp_name,
|
||||
seed=cfg.experiment.seed,
|
||||
torch_deterministic=cfg.experiment.torch_deterministic,
|
||||
cuda=cfg.experiment.cuda,
|
||||
track=cfg.logging.track,
|
||||
wandb_project_name=cfg.logging.wandb_project_name,
|
||||
wandb_entity=cfg.logging.wandb_entity,
|
||||
capture_video=cfg.logging.capture_video,
|
||||
save_model=cfg.logging.save_model,
|
||||
checkpoint_frequency=cfg.logging.checkpoint_frequency,
|
||||
upload_model=cfg.logging.upload_model,
|
||||
hf_entity=cfg.logging.hf_entity,
|
||||
total_timesteps=cfg.ppo.total_timesteps,
|
||||
learning_rate=cfg.ppo.learning_rate,
|
||||
num_envs=cfg.ppo.num_envs,
|
||||
num_steps=cfg.ppo.num_steps,
|
||||
anneal_lr=cfg.ppo.anneal_lr,
|
||||
gamma=cfg.ppo.gamma,
|
||||
gae_lambda=cfg.ppo.gae_lambda,
|
||||
num_minibatches=cfg.ppo.num_minibatches,
|
||||
update_epochs=cfg.ppo.update_epochs,
|
||||
norm_adv=cfg.ppo.norm_adv,
|
||||
clip_coef=cfg.ppo.clip_coef,
|
||||
clip_vloss=cfg.ppo.clip_vloss,
|
||||
ent_coef=cfg.ppo.ent_coef,
|
||||
vf_coef=cfg.ppo.vf_coef,
|
||||
max_grad_norm=cfg.ppo.max_grad_norm,
|
||||
target_kl=cfg.ppo.target_kl,
|
||||
run_dir=run_dir,
|
||||
)
|
||||
|
||||
# Compute runtime fields
|
||||
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
|
||||
|
||||
return args
|
||||
|
||||
|
||||
@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
|
||||
def main(dict_cfg: DictConfig):
|
||||
# 1. Convert DictConfig to structured dataclass
|
||||
|
|
@ -75,12 +27,10 @@ def main(dict_cfg: DictConfig):
|
|||
|
||||
# 2. Setup run metadata
|
||||
# Hydra changes CWD to the output directory by default.
|
||||
# We use that as our run_dir.
|
||||
run_dir = os.getcwd()
|
||||
run_name = os.path.basename(run_dir)
|
||||
|
||||
# 3. Initialize Logger
|
||||
# We pass the resolved dictionary for WandB/YAML logging
|
||||
resolved_cfg_dict = OmegaConf.to_container(dict_cfg, resolve=True, throw_on_missing=True)
|
||||
init_logger(
|
||||
run_name=run_name,
|
||||
|
|
@ -94,15 +44,12 @@ def main(dict_cfg: DictConfig):
|
|||
logger.info(f"Hydra-initialized run: {run_name}")
|
||||
logger.info(f"Output directory: {run_dir}")
|
||||
|
||||
# 4. Prepare compatibility object for PPOTrainer
|
||||
ppo_args = create_ppo_args_compat(cfg, run_dir)
|
||||
|
||||
# 5. Setup Environment and Torch
|
||||
# 4. Setup Environment and Torch
|
||||
env = make_env(cfg)
|
||||
torch.backends.cudnn.deterministic = cfg.experiment.torch_deterministic
|
||||
|
||||
# 6. Train
|
||||
ppo_trainer = PPOTrainer(ppo_args, env, run_dir, run_name)
|
||||
# 5. Train - pass structured config directly
|
||||
ppo_trainer = PPOTrainer(cfg, env, run_dir, run_name)
|
||||
ppo_trainer.train()
|
||||
|
||||
|
||||
|
|
|
|||
Reference in a new issue