From e40a4b979f859c35fd181addac5156d43e4fdb8b Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 15 Apr 2026 18:34:28 +0200 Subject: [PATCH] fix: typos and nonexistent calls --- configs/main_config.yaml | 1 + configs/ppo/smoke_test.yaml | 19 +++++++ scripts/train.py | 18 +++---- .../configs/config_architecture.py | 52 ++++++------------- .../configs/register_configs.py | 4 +- .../trainers/PPOTrainer.py | 6 +-- src/experiment_logger/unified_logger.py | 2 +- 7 files changed, 52 insertions(+), 50 deletions(-) create mode 100644 configs/ppo/smoke_test.yaml diff --git a/configs/main_config.yaml b/configs/main_config.yaml index 0583327..7e3d642 100644 --- a/configs/main_config.yaml +++ b/configs/main_config.yaml @@ -3,6 +3,7 @@ # Sub-configs are loaded from the relative directories. defaults: + - brittle_star_config - experiment: base - logging: default - ppo: default diff --git a/configs/ppo/smoke_test.yaml b/configs/ppo/smoke_test.yaml new file mode 100644 index 0000000..40e4c07 --- /dev/null +++ b/configs/ppo/smoke_test.yaml @@ -0,0 +1,19 @@ +# Fast PPO Configuration +# Lower timestep count for quick iterations/testing. + +learning_rate: 0.0005 +total_timesteps: 65536 +num_envs: 512 +num_steps: 128 +anneal_lr: true +gamma: 0.99 +gae_lambda: 0.95 +num_minibatches: 4 +update_epochs: 4 +norm_adv: true +clip_coef: 0.2 +clip_vloss: true +ent_coef: 0.01 +vf_coef: 0.5 +max_grad_norm: 0.5 +target_kl: null diff --git a/scripts/train.py b/scripts/train.py index 4224eb9..409a196 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -23,7 +23,7 @@ def make_env(cfg: BrittleStarConfig) -> BrittleStarJaxEnvWrapper: @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) + config: BrittleStarConfig = OmegaConf.to_object(dict_cfg) # 2. Setup run metadata # Hydra changes CWD to the output directory by default. @@ -31,25 +31,25 @@ def main(dict_cfg: DictConfig): run_name = os.path.basename(run_dir) # 3. Initialize Logger - resolved_cfg_dict = OmegaConf.to_container(dict_cfg, resolve=True, throw_on_missing=True) + 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, + config=cfg_dict, + project_name=config.logging.wandb_project_name, + entity=config.logging.wandb_entity, base_dir=os.path.dirname(run_dir), - use_wandb=cfg.logging.track, + use_wandb=config.logging.track, ) logger = get_logger() logger.info(f"Hydra-initialized run: {run_name}") logger.info(f"Output directory: {run_dir}") # 4. Setup Environment and Torch - env = make_env(cfg) - torch.backends.cudnn.deterministic = cfg.experiment.torch_deterministic + env = make_env(config) + torch.backends.cudnn.deterministic = config.experiment.torch_deterministic # 5. Train - pass structured config directly - ppo_trainer = PPOTrainer(cfg, env, run_dir, run_name) + ppo_trainer = PPOTrainer(config, env, run_dir, run_name) ppo_trainer.train() diff --git a/src/brittle_star_project/configs/config_architecture.py b/src/brittle_star_project/configs/config_architecture.py index bc140dc..e8cf5a5 100644 --- a/src/brittle_star_project/configs/config_architecture.py +++ b/src/brittle_star_project/configs/config_architecture.py @@ -1,5 +1,5 @@ from dataclasses import dataclass, field -from typing import List +from typing import List, Optional @dataclass @@ -18,14 +18,20 @@ class ArchitectureConfig: See docs/design/actor-critic.md for the full design rationale. """ - # Input network: maps global observation to a hidden representation. - feature_extractor: LayerConfig = field( - default_factory=lambda: LayerConfig(hidden_dims=[64, 64], activation="tanh") - ) - # Output network: maps the hidden representation to a scalar value estimate. - critic: LayerConfig = field( - default_factory=lambda: LayerConfig(hidden_dims=[], activation="tanh") - ) + name: str = "base" + + # Actor pipeline + sensor: Optional[LayerConfig] = None + propagator: Optional[LayerConfig] = None + motor: Optional[LayerConfig] = None + + # Critic pipeline + feature_extractor: Optional[LayerConfig] = None + critic: Optional[LayerConfig] = None + + # Decentralized + message_passing_steps: Optional[int] = None + topology_type: Optional[str] = None # Supported values: "ring", "fully_connected" @dataclass @@ -39,14 +45,7 @@ class CentralizedConfig(ArchitectureConfig): See docs/design/actor-critic.md for the full design rationale. """ - # Input network: maps the global observation to a hidden state. - sensor: LayerConfig = field( - default_factory=lambda: LayerConfig(hidden_dims=[64, 64], activation="tanh") - ) - # Output network: projects hidden state to (mean, log_std) over all joints. - motor: LayerConfig = field( - default_factory=lambda: LayerConfig(hidden_dims=[], activation="tanh") - ) + name: str = "centralized" @dataclass @@ -64,21 +63,4 @@ class DecentralizedConfig(ArchitectureConfig): full design rationale. """ - # Input network: local observation -> initial hidden state per node. - sensor: LayerConfig = field( - default_factory=lambda: LayerConfig(hidden_dims=[64, 64], activation="tanh") - ) - # Message-passing network: aggregates neighbour messages and updates hidden state. - propagator: LayerConfig = field( - default_factory=lambda: LayerConfig(hidden_dims=[64, 64], activation="tanh") - ) - # Output network: final hidden state -> joint offset for this node only. - motor: LayerConfig = field( - default_factory=lambda: LayerConfig(hidden_dims=[], activation="tanh") - ) - - # Number of synchronous message-passing rounds per control step. - message_passing_steps: int = 1 - # Graph topology used for neighbour connections. - # Supported values: "ring", "fully_connected". - topology_type: str = "ring" + name: str = "decentralized" diff --git a/src/brittle_star_project/configs/register_configs.py b/src/brittle_star_project/configs/register_configs.py index 697eda9..74360a3 100644 --- a/src/brittle_star_project/configs/register_configs.py +++ b/src/brittle_star_project/configs/register_configs.py @@ -29,8 +29,8 @@ def register_configs() -> None: cs.store(group="ppo", name="base_ppo", node=PPOConfig) # Architecture variants — swap via CLI: architecture=decentralized - cs.store(group="architecture", name="centralized", node=CentralizedConfig) - cs.store(group="architecture", name="decentralized", node=DecentralizedConfig) + cs.store(group="architecture", name="centralized_schema", node=CentralizedConfig) + cs.store(group="architecture", name="decentralized_schema", node=DecentralizedConfig) # Environment configs cs.store(group="morphology", name="base_morphology", node=MorphologyConfig) diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 3dc9666..1e49cc2 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -527,11 +527,11 @@ class PPOTrainer: f"ETA {eta_str}" ) - if getattr(self.config.experiment, "debug_sanity", False): - self.logger.log( + if getattr(self.cfg.experiment, "debug_sanity", False): + self.logger.info( "\n[SANITY CHECK] Successfully completed 1 epoch of data collection and gradient updates." ) - self.logger.log("[SANITY CHECK] Gradients flowed without NaN. Exiting gracefully.") + self.logger.info("[SANITY CHECK] Gradients flowed without NaN. Exiting gracefully.") break if self.logging_cfg.save_model: diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py index 4eadd5b..a02bdb1 100644 --- a/src/experiment_logger/unified_logger.py +++ b/src/experiment_logger/unified_logger.py @@ -53,7 +53,7 @@ def init_logger(**kwargs) -> "UnifiedLogger": the output directories and set up all logging backends. """ global _active_logger - logger = UnifiedLogger(_set_as_global=False, **kwargs) + logger = UnifiedLogger(**kwargs) _active_logger = logger return logger