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.configs.main_config import BrittleStarConfig
-from brittle_star_project.rl import RLModel  # imports concrete models via rl.__init__
+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:
+@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
-    p = argparse.ArgumentParser(description="Simulate a trained policy in the MuJoCo viewer.")
+def main(dict_cfg: DictConfig) -> None:
-    p.add_argument(
+    cfg: BrittleStarConfig = OmegaConf.to_object(dict_cfg)
-        "--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()
 
-
+    # Allow overriding these via Hydra CLI or a dedicated simulate config group
-def main() -> None:
+    # For now, defaults matching the old argparse behavior
-    args = parse_args()
+    backend = Backend.MJX
-
+    model_type = "random"
-    morphology_cfg, arena_cfg, env_cfg = from_file("../configs/test.yaml")
+    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)
+    raw_env = factory.create_environment(backend, cfg.morphology, cfg.arena, cfg.environment)
-    env = BrittleStarEnv(raw_env, backend=backend, config=env_cfg)
+    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)
-    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:
+    if model_path is not None:
-        model_path = Path(args.model)
+        policy = RLModel.load(Path(model_path))
-        policy = RLModel.load(model_path)
         if hasattr(policy, "nu"):
             policy.nu = nu
     else:
-        model_cls = MODEL_BY_NAME[str(args.model_type)]
+        model_cls = MODEL_BY_NAME[model_type]
-        policy = model_cls(seed=seed_for_env)
+        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))
-    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,
+        seed=default_seed,
     )
 
     simulate_policy(policy, rollout_cfg, state)
-
     env.close()
 
 
 if __name__ == "__main__":
+    register_configs()
     main()

scripts/train.py

@@ -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
+    # 4. Setup Environment and Torch
-    ppo_args = create_ppo_args_compat(cfg, run_dir)
-
-    # 5. Setup Environment and Torch
     env = make_env(cfg)
     torch.backends.cudnn.deterministic = cfg.experiment.torch_deterministic
 
-    # 6. Train
+    # 5. Train - pass structured config directly
-    ppo_trainer = PPOTrainer(ppo_args, env, run_dir, run_name)
+    ppo_trainer = PPOTrainer(cfg, env, run_dir, run_name)
     ppo_trainer.train()
 
 

src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py

@@ -8,7 +8,6 @@ from brittle_star_project import (
     ArenaConfig,
     Backend,
 )
-from brittle_star_project.environment import from_file
 
 
 class BrittleStarJaxEnvWrapper:
@@ -86,15 +85,6 @@ class BrittleStarJaxEnvWrapper:
             morphology, arena, env_config, num_envs=num_envs, backend=backend
         )
 
-    @staticmethod
-    def from_config(
-        config_path: str, num_envs: int, backend: Backend = Backend.MJX
-    ) -> "BrittleStarJaxEnvWrapper":
-        morphology_cfg, arena_cfg, env_cfg = from_file(config_path)
-        return BrittleStarJaxEnvWrapper(
-            morphology_cfg, arena_cfg, env_cfg, num_envs=num_envs, backend=backend
-        )
-
     def __str__(self):
         morphology_str = str(self._morphology)
         arena_str = str(self._arena)

src/brittle_star_project/environment/__init__.py

@@ -1,4 +1,4 @@
-from .env_config import ArenaConfig, EnvConfig, MorphologyConfig, from_file
+from .env_config import ArenaConfig, EnvConfig, MorphologyConfig
 from .env_types import Backend, Task
 from .env_wrapper import BrittleStarEnv, StepResult
 from .factory import BrittleStarEnvFactory
@@ -12,5 +12,4 @@ __all__ = [
     "BrittleStarEnv",
     "StepResult",
     "BrittleStarEnvFactory",
-    "from_file",
 ]

src/brittle_star_project/environment/env_config.py

@@ -60,16 +60,3 @@ class EnvConfig:
     # Light escape
     # Per docs in upstream env config: integer factors of 200.
     light_perlin_noise_scale: int = 0
-
-
-def from_file(path: str) -> tuple[MorphologyConfig, ArenaConfig, EnvConfig]:
-    """Load configurations from a YAML file."""
-    import yaml
-
-    with open(path, "r") as f:
-        config_dict = yaml.safe_load(f)
-
-        morphology = MorphologyConfig(**config_dict.get("morphology", {}))
-        arena = ArenaConfig(**config_dict.get("arena", {}))
-        env = EnvConfig(**config_dict.get("env", {}))
-        return morphology, arena, env

src/brittle_star_project/trainers/PPOTrainer.py

@@ -13,7 +13,8 @@ from flax.training.train_state import TrainState
 
 from experiment_logger import get_logger
 
-from brittle_star_project.dataclasses import EpisodeStatistics, PPOArgs
+from brittle_star_project.configs.main_config import BrittleStarConfig
+from brittle_star_project.dataclasses import EpisodeStatistics
 from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
 from brittle_star_project.MLPs.mlps import (
     Actor,
@@ -210,14 +211,23 @@ class TrainingMeasurements:
 
 
 class PPOTrainer:
-    def __init__(self, args: PPOArgs, env: BrittleStarJaxEnvWrapper, run_dir: str, run_name: str):
+    def __init__(
-        self.args = args
+        self, cfg: BrittleStarConfig, env: BrittleStarJaxEnvWrapper, run_dir: str, run_name: str
+    ):
+        self.cfg = cfg
+        self.ppo = cfg.ppo
+        self.experiment = cfg.experiment
+        self.logging_cfg = cfg.logging
         self.env = env
         self.run_dir = run_dir
         self.run_name = run_name
         self.logger = get_logger()
 
-        self.key = jax.random.PRNGKey(args.seed)
+        # Derived runtime fields
+        self.batch_size = self.ppo.num_envs * self.ppo.num_steps
+        self.num_iterations = self.ppo.total_timesteps // self.batch_size
+
+        self.key = jax.random.PRNGKey(self.experiment.seed)
 
         self.sensor, self.feature_extractor, self.actor, self.critic = self._init_agent()
         self.sensor.apply = jax.jit(self.sensor.apply)
@@ -228,7 +238,7 @@ class PPOTrainer:
         self._rollout_jit = jax.jit(
             partial(
                 _rollout_jit,
-                max_steps=self.args.num_steps,
+                max_steps=self.ppo.num_steps,
                 step_env_fn=partial(_step_env_wrapped, env_step_fn=self.env.step),
                 sensor=self.sensor,
                 feature_extractor=self.feature_extractor,
@@ -239,15 +249,15 @@ class PPOTrainer:
         self._compute_gae_jit = jax.jit(
             partial(
                 _compute_gae_jit,
-                num_envs=self.args.num_envs,
+                num_envs=self.ppo.num_envs,
-                gamma=self.args.gamma,
+                gamma=self.ppo.gamma,
-                gae_lambda=self.args.gae_lambda,
+                gae_lambda=self.ppo.gae_lambda,
                 feature_extractor=self.feature_extractor,
                 critic=self.critic,
             )
         )
 
-        self._ppo = PPO(self.args, self.sensor, self.actor, self.critic, self.feature_extractor)
+        self._ppo = PPO(self.ppo, self.sensor, self.actor, self.critic, self.feature_extractor)
 
         self.agent_state = self._init_agent_state()
 
@@ -256,10 +266,10 @@ class PPOTrainer:
         self._init_random()
 
     def _init_random(self):
-        self.logger.info(f"[RANDOM]: Setting random seed to {self.args.seed}")
+        self.logger.info(f"[RANDOM]: Setting random seed to {self.experiment.seed}")
 
-        random.seed(self.args.seed)
+        random.seed(self.experiment.seed)
-        np.random.seed(self.args.seed)
+        np.random.seed(self.experiment.seed)
 
     def _init_agent(self):
         self.logger.info("[AGENT]: Initializing agent...")
@@ -299,17 +309,17 @@ class PPOTrainer:
                 AgentParams(sensor_params, actor_params, critic_params, feature_extractor_params)
             ),
             tx=optax.chain(
-                optax.clip_by_global_norm(self.args.max_grad_norm),
+                optax.clip_by_global_norm(self.ppo.max_grad_norm),
                 optax.inject_hyperparams(optax.adam)(
                     learning_rate=partial(
                         _linear_schedule,
-                        minibatch_count=self.args.num_minibatches,
+                        minibatch_count=self.ppo.num_minibatches,
-                        update_epochs=self.args.update_epochs,
+                        update_epochs=self.ppo.update_epochs,
-                        num_iterations=self.args.num_iterations,
+                        num_iterations=self.num_iterations,
-                        learning_rate=self.args.learning_rate,
+                        learning_rate=self.ppo.learning_rate,
                     )
-                    if self.args.anneal_lr
+                    if self.ppo.anneal_lr
-                    else self.args.learning_rate,
+                    else self.ppo.learning_rate,
                     eps=1e-5,
                 ),
             ),
@@ -319,10 +329,10 @@ class PPOTrainer:
         self.logger.info("[EPISODE STATS]: Initializing episode stats...")
 
         return EpisodeStatistics(
-            episode_returns=jnp.zeros(self.args.num_envs, dtype=jnp.float32),
+            episode_returns=jnp.zeros(self.ppo.num_envs, dtype=jnp.float32),
-            episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32),
+            episode_lengths=jnp.zeros(self.ppo.num_envs, dtype=jnp.int32),
-            returned_episode_returns=jnp.zeros(self.args.num_envs, jnp.float32),
+            returned_episode_returns=jnp.zeros(self.ppo.num_envs, jnp.float32),
-            returned_episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32),
+            returned_episode_lengths=jnp.zeros(self.ppo.num_envs, dtype=jnp.int32),
         )
 
     def _rollout(self, env_state, next_obs, next_done) -> tuple[Any, ...]:
@@ -371,7 +381,7 @@ class PPOTrainer:
             "losses/loss": training_measurements.loss[-1, -1].item(),
             "charts/SPS": int(global_step / (time.time() - start_time)),
             "charts/SPS_update": int(
-                self.args.num_envs * self.args.num_steps / (time.time() - iteration_time_start)
+                self.ppo.num_envs * self.ppo.num_steps / (time.time() - iteration_time_start)
             ),
         }
         self.logger.log(metrics, step=global_step)
@@ -451,8 +461,14 @@ class PPOTrainer:
     def _save_model(self, model_path: str):
         self.logger.info("[SAVE]: Saving the final model...")
 
+        from dataclasses import asdict as _asdict
+
+        config_dict = {
+            "experiment": _asdict(self.experiment),
+            "ppo": _asdict(self.ppo),
+        }
         params = [
-            vars(self.args),
+            config_dict,
             [
                 self.agent_state.params["sensor_params"],
                 self.agent_state.params["actor_params"],
@@ -464,8 +480,7 @@ class PPOTrainer:
 
     def train(self):
         """
-        Train the PPO agent for a specified number of iterations
+        Train the PPO agent for a specified number of iterations.
-        (passed through PPOArgs in constructor).
         Closes the environment at the end of training.
         """
         self.logger.info(f"running name: {self.run_name}")
@@ -473,16 +488,16 @@ class PPOTrainer:
         self.logger.info("[TRAIN]: Resetting environment...")
         self.logger.log_non_interactive(f"Initial reset started: {time.ctime()}")
 
-        env_state = self.env.reset(seed=self.args.seed)
+        env_state = self.env.reset(seed=self.experiment.seed)
         next_obs = _convert_obs_dict_to_array(env_state.observations)
-        next_done = jnp.zeros(self.args.num_envs, dtype=jnp.bool_)
+        next_done = jnp.zeros(self.ppo.num_envs, dtype=jnp.bool_)
 
         self.logger.log_non_interactive(f"Initial reset completed: {time.ctime()}")
 
         global_step = 0
         start_time = time.time()
 
-        iter_bar = self.logger.progress_bar(range(1, self.args.num_iterations + 1))
+        iter_bar = self.logger.progress_bar(range(1, self.num_iterations + 1))
         for iteration in iter_bar:
             iteration_time_start = time.time()
 
@@ -490,7 +505,7 @@ class PPOTrainer:
                 env_state, next_obs, next_done, iteration=iteration
             )
 
-            global_step += self.args.num_steps * self.args.num_envs
+            global_step += self.ppo.num_steps * self.ppo.num_envs
             self._log(
                 global_step,
                 self.episode_stats,
@@ -500,20 +515,20 @@ class PPOTrainer:
             )
 
             sps = int(global_step / (time.time() - start_time))
-            remaining_steps = self.args.total_timesteps - global_step
+            remaining_steps = self.ppo.total_timesteps - global_step
             eta_seconds = int(remaining_steps / sps) if sps > 0 else 0
             eta_str = str(datetime.timedelta(seconds=eta_seconds))
 
             self.logger.log_non_interactive(
-                f"Iteration {iteration}/{self.args.num_iterations} | "
+                f"Iteration {iteration}/{self.num_iterations} | "
-                f"Step {global_step}/{self.args.total_timesteps} | "
+                f"Step {global_step}/{self.ppo.total_timesteps} | "
                 f"SPS {sps} | "
                 f"Return {training_measurements.avg_episodic_return:.4f} | "
                 f"ETA {eta_str}"
             )
 
-        if self.args.save_model:
+        if self.logging_cfg.save_model:
-            model_path = f"{self.run_dir}/{self.args.exp_name}.cleanrl_model"
+            model_path = f"{self.run_dir}/{self.experiment.exp_name}.cleanrl_model"
             self._save_model(model_path=model_path)
 
         self._close()

src/experiment_logger/__init__.py

@@ -4,7 +4,7 @@ This package provides a unified interface for logging to multiple backends
 (WandB, disk, stdout) simultaneously, ensuring no data loss.
 """
 
-from experiment_logger.config_utils import load_yaml_config, merge_config_with_cli
+from experiment_logger.config_utils import load_yaml_config
 from experiment_logger.unified_logger import UnifiedLogger, get_logger, init_logger
 from experiment_logger.simple_logger import SimpleLogger
 from experiment_logger.wandb_utils import finish_wandb, init_wandb
@@ -17,6 +17,5 @@ __all__ = [
     "init_wandb",
     "finish_wandb",
     "load_yaml_config",
-    "merge_config_with_cli",
 ]
 __version__ = "0.1.0"

src/experiment_logger/config_utils.py

@@ -72,67 +72,6 @@ def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T:
     return cls(**filtered_config)
 
 
-def merge_config_with_cli(config_class: Type[T], config_file: str | None = None) -> T:
-    """Merge YAML config with CLI arguments, with CLI taking precedence.
-
-    Args:
-        config_class: Dataclass type to create
-        config_file: Path to YAML config file (optional)
-
-    Returns:
-        Instance of config_class with merged configuration
-    """
-    # Parse CLI args first to get the default/CLI values
-    import tyro
-
-    # Check if --config is in sys.argv and extract it
-    extracted_config_file = config_file
-    if "--config" in sys.argv:
-        config_idx = sys.argv.index("--config")
-        if config_idx + 1 < len(sys.argv):
-            extracted_config_file = sys.argv[config_idx + 1]
-            # Remove from sys.argv so tyro doesn't see it
-            sys.argv.pop(config_idx)  # Remove --config
-            sys.argv.pop(config_idx)  # Remove config file path
-
-    # Load YAML config if available
-    yaml_config = {}
-    if extracted_config_file and os.path.exists(extracted_config_file):
-        yaml_config = load_yaml_config(extracted_config_file)
-        get_logger().info(f"Merging YAML config from {extracted_config_file} with CLI args")
-    elif extracted_config_file:
-        get_logger().warning(f"Config file not found: {extracted_config_file}, using CLI args only")
-
-    # Create default instance to know what the defaults are
-    default_instance = config_class()
-    default_dict = {f.name: getattr(default_instance, f.name) for f in fields(config_class)}  # type: ignore
-
-    # Parse CLI args
-    cli_instance = tyro.cli(config_class)
-    cli_dict = {f.name: getattr(cli_instance, f.name) for f in fields(config_class)}  # type: ignore
-
-    # Merge configs: YAML as base, CLI overrides non-default values
-    final_config = {}
-
-    for field in fields(config_class):  # type: ignore
-        field_name = field.name
-        default_value = default_dict[field_name]
-        yaml_value = yaml_config.get(field_name, default_value)
-        cli_value = cli_dict[field_name]
-
-        # Use CLI value if it's different from default, otherwise use YAML value
-        if cli_value != default_value:
-            final_config[field_name] = cli_value
-            if yaml_value != default_value and yaml_value != cli_value:
-                get_logger().info(f"CLI override: {field_name}={cli_value} (YAML had {yaml_value})")
-        else:
-            final_config[field_name] = yaml_value
-            if yaml_value != default_value:
-                get_logger().info(f"YAML config: {field_name}={yaml_value}")
-
-    return config_class(**final_config)
-
-
 def print_config(config: Any, title: str = "Configuration"):
     """Pretty print configuration."""
     get_logger().info(f"{title}:")

tests/test_config.py

@@ -31,8 +31,3 @@ class TestYamlConfig:
 
         with pytest.raises(FileNotFoundError):
             load_yaml_config("nonexistent.yaml")
-
-    def test_merge_config_with_cli_is_callable(self):
-        from experiment_logger.config_utils import merge_config_with_cli
-
-        assert callable(merge_config_with_cli)