From 37a4b59e04c8a43cd1a25dd6981481949c1d1896 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Thu, 16 Apr 2026 15:10:21 +0200 Subject: [PATCH] feat: configure checkpoints saving --- configs/logging/default.yaml | 4 ++- configs/logging/hpc.yaml | 9 +++++++ configs/logging/wandb_enabled.yaml | 6 +++-- scripts/hpc/train.pbs | 2 +- scripts/simulate.py | 13 ++++++++++ scripts/train.py | 2 ++ .../trainers/PPOTrainer.py | 11 ++++---- src/brittle_star_project/trainers/utils.py | 21 ---------------- src/experiment_logger/config_logger.py | 25 ++++++++++++++++--- src/experiment_logger/unified_logger.py | 8 ++++-- 10 files changed, 65 insertions(+), 36 deletions(-) create mode 100644 configs/logging/hpc.yaml delete mode 100644 src/brittle_star_project/trainers/utils.py diff --git a/configs/logging/default.yaml b/configs/logging/default.yaml index 4a0cb92..1ecf7fd 100644 --- a/configs/logging/default.yaml +++ b/configs/logging/default.yaml @@ -6,6 +6,8 @@ wandb_project_name: "PPO-Modularity" wandb_entity: "SEL3-2026-Groep-4" capture_video: false save_model: true +save_checkpoints: true checkpoint_frequency: 100 -upload_model: false +upload_final_model: false +upload_checkpoints: false hf_entity: "" diff --git a/configs/logging/hpc.yaml b/configs/logging/hpc.yaml new file mode 100644 index 0000000..3c03918 --- /dev/null +++ b/configs/logging/hpc.yaml @@ -0,0 +1,9 @@ +track: true +wandb_project_name: "hpc-default" +wandb_entity: "SEL3-2026-Groep-4" +save_model: true +save_checkpoints: true +upload_final_model: true +upload_checkpoints: true +checkpoint_frequency: 100 +hf_entity: "" diff --git a/configs/logging/wandb_enabled.yaml b/configs/logging/wandb_enabled.yaml index 2e82781..7ddf95f 100644 --- a/configs/logging/wandb_enabled.yaml +++ b/configs/logging/wandb_enabled.yaml @@ -2,10 +2,12 @@ # For production/cloud experiments with weights synced. track: true -wandb_project_name: "PPO-Modularity" +wandb_project_name: "default-project" wandb_entity: "SEL3-2026-Groep-4" capture_video: false save_model: true +save_checkpoints: true checkpoint_frequency: 100 -upload_model: false +upload_final_model: true +upload_checkpoints: false hf_entity: "" diff --git a/scripts/hpc/train.pbs b/scripts/hpc/train.pbs index 60ba4d8..0f95b18 100644 --- a/scripts/hpc/train.pbs +++ b/scripts/hpc/train.pbs @@ -67,7 +67,7 @@ fi python scripts/train.py \ hydra.run.dir="$SCRATCH_RUNDIR" \ ppo=stable \ - logging=wandb_enabled + logging=hpc echo ">>> Staging out results to $DATA_RUNDIR..." cp -r "$SCRATCH_RUNDIR/." "$DATA_RUNDIR/" diff --git a/scripts/simulate.py b/scripts/simulate.py index 666d048..b49e50f 100644 --- a/scripts/simulate.py +++ b/scripts/simulate.py @@ -57,6 +57,19 @@ def main(dict_cfg: DictConfig) -> None: nu = int(state.mj_model.nu) if model_path is not None: + # TODO: Refactoring Notice - The .flax checkpoint payload no longer encapsulates the config + # and no longer wraps parameters into a hardcoded list. + # The file now natively contains solely the pure raw Jax 'agent_state.params' FrozenDict mapping. + # The entire BrittleStarConfig is safely exported alongside it down at '..._metadata.yaml'. + # + # Example parsed layout from flax.serialization.from_bytes(): + # { + # 'sensor_params': FrozenDict({...}), + # 'actor_params': FrozenDict({...}), + # 'critic_params': FrozenDict({...}), + # ... + # } + # Update the RLModel.load function or subsequent destructuring to support this raw dictionary natively. policy = RLModel.load(Path(model_path)) if hasattr(policy, "nu"): policy.nu = nu diff --git a/scripts/train.py b/scripts/train.py index 409f050..a338484 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -41,6 +41,8 @@ def main(dict_cfg: DictConfig): entity=config.logging.wandb_entity, base_dir=os.path.dirname(run_dir), use_wandb=config.logging.track, + upload_final_model=config.logging.upload_final_model, + upload_checkpoints=config.logging.upload_checkpoints, ) logger = get_logger() logger.info(f"Hydra-initialized run: {run_name}") diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index ea77f12..f1429a6 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -24,7 +24,6 @@ from brittle_star_project.MLPs.mlps import ( Storage, ) from brittle_star_project.ppo import PPO -from brittle_star_project.trainers.utils import serialize_training_state # TODO: move to config _ALLOWED_OBS_KEYS = { @@ -574,13 +573,13 @@ class PPOTrainer: def _save_model(self, model_path: str): self.logger.info("[SAVE]: Saving the final model...") - config_dict, params = serialize_training_state(self.cfg, self.agent_state) - self.logger.save_final_model(params=params, metadata=config_dict) + self.logger.save_final_model(params=self.agent_state.params, metadata=asdict(self.cfg)) def _save_checkpoint(self, iteration: int): self.logger.info(f"[SAVE]: Saving checkpoint at iteration {iteration}...") - config_dict, params = serialize_training_state(self.cfg, self.agent_state) - self.logger.save_checkpoint(params=params, step=iteration, metadata=config_dict) + self.logger.save_checkpoint( + params=self.agent_state.params, step=iteration, metadata=asdict(self.cfg) + ) def train(self): """ @@ -638,7 +637,7 @@ class PPOTrainer: f"ETA {eta_str}" ) - if self.logging_cfg.save_model and self.logging_cfg.checkpoint_frequency > 0: + if self.logging_cfg.save_checkpoints and self.logging_cfg.checkpoint_frequency > 0: if iteration % self.logging_cfg.checkpoint_frequency == 0: self._save_checkpoint(iteration) diff --git a/src/brittle_star_project/trainers/utils.py b/src/brittle_star_project/trainers/utils.py deleted file mode 100644 index e271366..0000000 --- a/src/brittle_star_project/trainers/utils.py +++ /dev/null @@ -1,21 +0,0 @@ -from flax.training.train_state import TrainState -from brittle_star_project.configs.main_config import BrittleStarConfig - - -def serialize_training_state(cfg: BrittleStarConfig, agent_state: TrainState): - from dataclasses import asdict as _asdict - - config_dict = { - "experiment": _asdict(cfg.experiment), - "ppo": _asdict(cfg.ppo), - } - params = [ - config_dict, - [ - agent_state.params["sensor_params"], - agent_state.params["actor_params"], - agent_state.params["critic_params"], - agent_state.params["feature_extractor_params"], - ], - ] - return config_dict, params diff --git a/src/experiment_logger/config_logger.py b/src/experiment_logger/config_logger.py index 5d39b8e..c34c162 100644 --- a/src/experiment_logger/config_logger.py +++ b/src/experiment_logger/config_logger.py @@ -5,10 +5,29 @@ from typing import Optional @dataclass class LoggingConfig: track: bool = False - wandb_project_name: str = "PPO-Modularity" + wandb_project_name: str = "default-project" wandb_entity: Optional[str] = "SEL3-2026-Groep-4" capture_video: bool = False - save_model: bool = True + + # Local Saving + save_model: bool = True # Final model + save_checkpoints: bool = True # Intermediate checkpoints checkpoint_frequency: int = 100 - upload_model: bool = False + + # Remote Uploading (WandB Artifacts) + upload_final_model: bool = False + upload_checkpoints: bool = False + hf_entity: str = "" + + def __post_init__(self): + if self.upload_final_model and not (self.track and self.save_model): + raise ValueError( + "Configuration Error: 'upload_final_model' is True, but it requires " + "both 'track' and 'save_model' to also be True." + ) + if self.upload_checkpoints and not (self.track and self.save_checkpoints): + raise ValueError( + "Configuration Error: 'upload_checkpoints' is True, but it requires " + "both 'track' and 'save_checkpoints' to also be True." + ) diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py index 2f5d06f..c98da37 100644 --- a/src/experiment_logger/unified_logger.py +++ b/src/experiment_logger/unified_logger.py @@ -93,6 +93,8 @@ class UnifiedLogger: entity: Optional[str] = None, base_dir: str = "runs", use_wandb: bool = True, + upload_final_model: bool = False, + upload_checkpoints: bool = False, save_code: bool = True, log_level: int = logging.INFO, ): @@ -110,6 +112,8 @@ class UnifiedLogger: self.run_name = run_name self.config = config self.use_wandb = use_wandb + self.upload_final_model = upload_final_model + self.upload_checkpoints = upload_checkpoints self.wandb_available = False self.wandb_run = None self.is_interactive = sys.stdout.isatty() @@ -327,7 +331,7 @@ class UnifiedLogger: self.info(f"Checkpoint saved: {checkpoint_path}") # Log to WandB as artifact - if self.wandb_run is not None: + if self.wandb_run is not None and self.upload_checkpoints: try: import wandb @@ -363,7 +367,7 @@ class UnifiedLogger: self.info(f"Final model saved: {final_model_path}") # Log to WandB - if self.wandb_run is not None: + if self.wandb_run is not None and self.upload_final_model: try: import wandb