feat: configure checkpoints saving

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)
 

src/brittle_star_project/trainers/utils.py

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

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."
+            )

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