From 93bff112088d5e5e4b8824e77f680cd550c5968a Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 15 Apr 2026 15:18:34 +0200 Subject: [PATCH] refactor: use hydra for configs in simulation and training scripts --- scripts/simulate.py | 82 +++++++---------- scripts/train.py | 59 +------------ .../environment/BrittleStarJaxEnvWrapper.py | 10 --- .../environment/__init__.py | 3 +- .../environment/env_config.py | 13 --- .../trainers/PPOTrainer.py | 87 +++++++++++-------- src/experiment_logger/__init__.py | 3 +- src/experiment_logger/config_utils.py | 61 ------------- tests/test_config.py | 5 -- 9 files changed, 87 insertions(+), 236 deletions(-) diff --git a/scripts/simulate.py b/scripts/simulate.py index 957f6fa..63b1cfe 100644 --- a/scripts/simulate.py +++ b/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.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() diff --git a/scripts/train.py b/scripts/train.py index 5ef2258..4224eb9 100644 --- a/scripts/train.py +++ b/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 - 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() diff --git a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py index b143d9c..7a6d161 100644 --- a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py +++ b/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) diff --git a/src/brittle_star_project/environment/__init__.py b/src/brittle_star_project/environment/__init__.py index aad8c35..78896cb 100644 --- a/src/brittle_star_project/environment/__init__.py +++ b/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", ] diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index 206e902..48f0e34 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/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 diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 6bc94ef..b20c4dd 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/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): - self.args = args + def __init__( + 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, - gamma=self.args.gamma, - gae_lambda=self.args.gae_lambda, + num_envs=self.ppo.num_envs, + gamma=self.ppo.gamma, + 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) - np.random.seed(self.args.seed) + random.seed(self.experiment.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, - update_epochs=self.args.update_epochs, - num_iterations=self.args.num_iterations, - learning_rate=self.args.learning_rate, + minibatch_count=self.ppo.num_minibatches, + update_epochs=self.ppo.update_epochs, + num_iterations=self.num_iterations, + learning_rate=self.ppo.learning_rate, ) - if self.args.anneal_lr - else self.args.learning_rate, + if self.ppo.anneal_lr + 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_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32), - returned_episode_returns=jnp.zeros(self.args.num_envs, jnp.float32), - returned_episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32), + episode_returns=jnp.zeros(self.ppo.num_envs, dtype=jnp.float32), + episode_lengths=jnp.zeros(self.ppo.num_envs, dtype=jnp.int32), + returned_episode_returns=jnp.zeros(self.ppo.num_envs, jnp.float32), + 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 - (passed through PPOArgs in constructor). + Train the PPO agent for a specified number of iterations. 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"Step {global_step}/{self.args.total_timesteps} | " + f"Iteration {iteration}/{self.num_iterations} | " + 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: - model_path = f"{self.run_dir}/{self.args.exp_name}.cleanrl_model" + if self.logging_cfg.save_model: + model_path = f"{self.run_dir}/{self.experiment.exp_name}.cleanrl_model" self._save_model(model_path=model_path) self._close() diff --git a/src/experiment_logger/__init__.py b/src/experiment_logger/__init__.py index c89af7f..53e57c0 100644 --- a/src/experiment_logger/__init__.py +++ b/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" diff --git a/src/experiment_logger/config_utils.py b/src/experiment_logger/config_utils.py index 09cb4ce..6017041 100644 --- a/src/experiment_logger/config_utils.py +++ b/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}:") diff --git a/tests/test_config.py b/tests/test_config.py index 212e4b4..e0ee0ef 100644 --- a/tests/test_config.py +++ b/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)