migrate main training entry point to Hydra and wire structured configs

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