diff --git a/scripts/train.py b/scripts/train.py index 492c887..5ef2258 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -1,75 +1,111 @@ -import subprocess -import time - -import torch 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 UnifiedLogger -from experiment_logger.config_utils import merge_config_with_cli, print_config +from experiment_logger import init_logger, get_logger -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 make_env(cfg: BrittleStarConfig) -> BrittleStarJaxEnvWrapper: + """Create the environment using the structured configuration.""" + return BrittleStarJaxEnvWrapper( + morphology=cfg.morphology, + arena=cfg.arena, + env_config=cfg.environment, + num_envs=cfg.ppo.num_envs, + ) -def parse_args() -> PPOArgs: - import argparse +def create_ppo_args_compat(cfg: BrittleStarConfig, run_dir: str) -> PPOArgs: + """Temporary adapter to bridge BrittleStarConfig to the legacy PPOArgs. - # 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() + 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 - git_hash = get_git_hash() - run_name = f"{args.exp_name}__seed_{args.seed}__{git_hash}__{int(time.time())}" + return args - 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) +@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3") +def main(dict_cfg: DictConfig): + # 1. Convert DictConfig to structured dataclass + cfg: BrittleStarConfig = OmegaConf.to_object(dict_cfg) - # Initialize Global Logger - logger = UnifiedLogger( - config=vars(args), - project_name=args.wandb_project_name, # or default PPO-Modularity if missing + # 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, + config=resolved_cfg_dict, + project_name=cfg.logging.wandb_project_name, + entity=cfg.logging.wandb_entity, base_dir=os.path.dirname(run_dir), - use_wandb=args.track, + use_wandb=cfg.logging.track, ) + logger = get_logger() + logger.info(f"Hydra-initialized run: {run_name}") + logger.info(f"Output directory: {run_dir}") - print_config(args, title="PPO Training Configuration") + # 4. Prepare compatibility object for PPOTrainer + ppo_args = create_ppo_args_compat(cfg, run_dir) - env = make_env(args.env_config_path, args.num_envs) + # 5. Setup Environment and Torch + env = make_env(cfg) + torch.backends.cudnn.deterministic = cfg.experiment.torch_deterministic - torch.backends.cudnn.deterministic = args.torch_deterministic - - ppo_trainer = PPOTrainer(args, env, run_dir, run_name) + # 6. Train + ppo_trainer = PPOTrainer(ppo_args, env, run_dir, run_name) ppo_trainer.train() + + +if __name__ == "__main__": + register_configs() + main()