tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
Archived
1
Fork 0
Code Issues Releases Activity

refactor: use hydra for configs in simulation and training scripts

Browse source
This commit is contained in:
Tibo De Peuter 2026-04-15 15:18:34 +02:00
parent ff83af8cef
commit 93bff11208
Signed by: tdpeuter
SSH key fingerprint: SHA256:u/h/LVoqKF1Iz02uOyxe6hcjmoZASCGV2HM0TG9ZMoU
9 changed files with 87 additions and 236 deletions
Download patch file Download diff file Expand all files Collapse all files

82
scripts/simulate.py
View file

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

59
scripts/train.py
View file

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

10
src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py
View file

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

3
src/brittle_star_project/environment/__init__.py
View file

@ -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",
]

13
src/brittle_star_project/environment/env_config.py
View file

@ -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

87
src/brittle_star_project/trainers/PPOTrainer.py
View file

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

3
src/experiment_logger/__init__.py
View file

@ -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"

61
src/experiment_logger/config_utils.py
View file

@ -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}:")

5
tests/test_config.py
View file

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