Merge branch 'dev' into reward-fix

.github/workflows/update_hpc_requirements.yml

@@ -14,6 +14,7 @@ jobs:
     runs-on: ubuntu-latest
     permissions:
       contents: write
+      pull-requests: write
     steps:
       - uses: actions/checkout@v4
         with:
@@ -30,9 +31,13 @@ jobs:
       - name: Regenerate env/hpc/requirements.txt
         run: uv run scripts/hpc/export_requirements.py
 
-      - name: Commit updated requirements if changed
+      - name: Create Pull Request with updated requirements
-        uses: stefanzweifel/git-auto-commit-action@v5
+        uses: peter-evans/create-pull-request@v6
         with:
-          commit_message: "chore(hpc): update env/hpc/requirements.txt from pyproject.toml [skip ci]"
+          token: ${{ secrets.GITHUB_TOKEN }}
-          file_pattern: env/hpc/requirements.txt
+          commit-message: "chore(hpc): update env/hpc/requirements.txt from pyproject.toml"
-          commit_author: "github-actions[bot] <github-actions[bot]@users.noreply.github.com>"
+          title: "chore(hpc): update HPC requirements"
+          body: "Automatically generated pull request to update `env/hpc/requirements.txt` based on recent changes to `pyproject.toml`."
+          branch: chore/auto-update-hpc-requirements
+          base: ${{ github.ref_name }}
+          author: "github-actions[bot] <github-actions[bot]@users.noreply.github.com>"

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

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

configs/logging/wandb_enabled.yaml

@@ -2,10 +2,12 @@
 # For production/cloud experiments with weights synced.
 
 track: true
-wandb_project_name: "PPO-Modularity - reward engineering"
+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: ""

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

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.
+        # 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 to support this raw dictionary natively.
         policy = RLModel.load(Path(model_path))
         if hasattr(policy, "nu"):
             policy.nu = nu

scripts/train.py

@@ -36,11 +36,9 @@ def main(dict_cfg: DictConfig):
     cfg_dict = OmegaConf.to_container(dict_cfg, resolve=True, throw_on_missing=True)
     init_logger(
         run_name=run_name,
-        config=cfg_dict,
+        full_config=cfg_dict,
-        project_name=config.logging.wandb_project_name,
+        logging_cfg=config.logging,
-        entity=config.logging.wandb_entity,
         base_dir=os.path.dirname(run_dir),
-        use_wandb=config.logging.track,
     )
     logger = get_logger()
     logger.info(f"Hydra-initialized run: {run_name}")

src/brittle_star_project/trainers/PPOTrainer.py

@@ -575,23 +575,13 @@ class PPOTrainer:
 
     def _save_model(self, model_path: str):
         self.logger.info("[SAVE]: Saving the final model...")
+        self.logger.save_final_model(params=self.agent_state.params, metadata=asdict(self.cfg))
 
-        from dataclasses import asdict as _asdict
+    def _save_checkpoint(self, iteration: int):
-
+        self.logger.info(f"[SAVE]: Saving checkpoint at iteration {iteration}...")
-        config_dict = {
+        self.logger.save_checkpoint(
-            "experiment": _asdict(self.experiment),
+            params=self.agent_state.params, step=iteration, metadata=asdict(self.cfg)
-            "ppo": _asdict(self.ppo),
+        )
-        }
-        params = [
-            config_dict,
-            [
-                self.agent_state.params["sensor_params"],
-                self.agent_state.params["actor_params"],
-                self.agent_state.params["critic_params"],
-                self.agent_state.params["feature_extractor_params"],
-            ],
-        ]
-        self.logger.save_final_model(params=params)
 
     def train(self):
         """
@@ -647,6 +637,10 @@ class PPOTrainer:
                 f"ETA {eta_str}"
             )
 
+            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)
+
             if getattr(self.cfg.experiment, "debug_sanity", False):
                 self.logger.info("\n[SANITY CHECK] Successfully completed 1 epoch")
                 break

src/experiment_logger/__init__.py

@@ -4,7 +4,6 @@ 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
 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
@@ -16,6 +15,5 @@ __all__ = [
     "init_logger",
     "init_wandb",
     "finish_wandb",
-    "load_yaml_config",
 ]
 __version__ = "0.1.0"

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/config_utils.py

@@ -1,83 +0,0 @@
-"""Configuration utilities for loading YAML configs and merging with CLI args."""
-
-import os
-from typing import Dict, Any, Type, TypeVar
-import yaml
-from dataclasses import fields, is_dataclass
-
-from experiment_logger.unified_logger import get_logger
-
-T = TypeVar("T")
-
-
-def load_yaml_config(config_path: str) -> Dict[str, Any]:
-    """Load configuration from YAML file."""
-    if not os.path.exists(config_path):
-        raise FileNotFoundError(f"Config file not found: {config_path}")
-
-    with open(config_path, "r") as f:
-        config = yaml.safe_load(f)
-
-    if config is None:
-        return {}
-
-    get_logger().info(f"Loaded configuration from: {config_path}")
-    return config
-
-
-def save_yaml_config(config: Dict[str, Any], config_path: str):
-    """Save configuration to YAML file."""
-    os.makedirs(os.path.dirname(config_path), exist_ok=True)
-
-    with open(config_path, "w") as f:
-        yaml.dump(config, f, default_flow_style=False, indent=2, sort_keys=False)
-
-    get_logger().info(f"Saved configuration to: {config_path}")
-
-
-def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T:
-    """Create dataclass instance from dictionary, handling type conversions."""
-    if not is_dataclass(cls):
-        raise ValueError(f"{cls} is not a dataclass")
-
-    # Get field names and types
-    field_map = {f.name: f for f in fields(cls)}  # type: ignore
-
-    # Filter config to only include valid fields
-    filtered_config: Dict[str, Any] = {}
-    for key, value in config_dict.items():
-        if key in field_map:
-            field = field_map[key]
-            # Handle type conversion if needed
-            try:
-                # Handle None values and optional types
-                if value is None:
-                    filtered_config[key] = None
-                elif hasattr(field.type, "__origin__") and field.type.__origin__ is type(None):
-                    # Optional type (Union[X, None])
-                    filtered_config[key] = value
-                else:
-                    # Try to convert to the expected type
-                    if field.type is bool and isinstance(value, str):
-                        filtered_config[key] = value.lower() in ("true", "1", "yes", "on")
-                    else:
-                        filtered_config[key] = field.type(value) if value is not None else None  # type: ignore
-            except (ValueError, TypeError) as e:
-                get_logger().warning(f"Could not convert {key}={value} to {field.type}: {e}")
-                filtered_config[key] = value
-        else:
-            get_logger().warning(f"Unknown configuration parameter: {key}")
-
-    return cls(**filtered_config)
-
-
-def print_config(config: Any, title: str = "Configuration"):
-    """Pretty print configuration."""
-    get_logger().info(f"{title}:")
-    if is_dataclass(config):
-        for field in fields(config):
-            value = getattr(config, field.name)
-            get_logger().info(f"  {field.name}: {value}")
-    else:
-        for key, value in vars(config).items():
-            get_logger().info(f"  {key}: {value}")

src/experiment_logger/simple_logger.py

@@ -14,18 +14,16 @@ class SimpleLogger:
     def __init__(
         self,
         run_name: str = "simple_run",
-        config: Optional[Dict[str, Any]] = None,
+        full_config: Optional[Dict[str, Any]] = None,
-        project_name: str = "none",
+        logging_cfg: Optional[Any] = None,
-        entity: Optional[str] = None,
         base_dir: str = "runs",
-        use_wandb: bool = False,
         save_code: bool = False,
         log_level: int = logging.INFO,
         _set_as_global: bool = False,
     ):
         self.is_interactive = True
         self.run_name = run_name
-        self.config = config or {}
+        self.full_config = full_config or {}
         print(f"[INIT] SimpleLogger initialized for run: {run_name}")
 
     def set_level(self, level: int):

src/experiment_logger/unified_logger.py

@@ -18,6 +18,7 @@ import jax.numpy as jnp
 import numpy as np
 
 from experiment_logger.wandb_utils import finish_wandb, init_wandb
+from experiment_logger.config_logger import LoggingConfig
 
 # Global storage for the active logger and the proxy singleton
 _active_logger: Optional[Any] = None
@@ -88,11 +89,9 @@ class UnifiedLogger:
     def __init__(
         self,
         run_name: str,
-        config: Dict[str, Any],
+        full_config: Dict[str, Any],
-        project_name: str = "PPO-Modularity",
+        logging_cfg: LoggingConfig,
-        entity: Optional[str] = None,
         base_dir: str = "runs",
-        use_wandb: bool = True,
         save_code: bool = True,
         log_level: int = logging.INFO,
     ):
@@ -100,16 +99,16 @@ class UnifiedLogger:
 
         Args:
             run_name: Unique name for this run
-            config: Configuration dictionary with hyperparameters
+            full_config: Full configuration dictionary with hyperparameters to be saved
-            project_name: WandB project name
+            logging_cfg: Structured logging configuration dataclass
-            entity: WandB entity (team/user name)
             base_dir: Base directory for local storage
-            use_wandb: Whether to use WandB logging
             save_code: Whether to save code to WandB
         """
         self.run_name = run_name
-        self.config = config
+        self.full_config = full_config
-        self.use_wandb = use_wandb
+        self.use_wandb = logging_cfg.track
+        self.upload_final_model = logging_cfg.upload_final_model
+        self.upload_checkpoints = logging_cfg.upload_checkpoints
         self.wandb_available = False
         self.wandb_run = None
         self.is_interactive = sys.stdout.isatty()
@@ -159,7 +158,7 @@ class UnifiedLogger:
 
         # Initialize WandB if requested
         if self.use_wandb:
-            self._init_wandb(project_name, entity, save_code)
+            self._init_wandb(logging_cfg.wandb_project_name, logging_cfg.wandb_entity, save_code)
 
         # Initialize metrics storage
         self.metrics_buffer: List[Dict[str, Any]] = []
@@ -207,7 +206,7 @@ class UnifiedLogger:
             project=project_name,
             entity=entity,
             name=self.run_name,
-            config=self.config,
+            config=self.full_config,
             save_code=save_code,
             resume="allow",
         )
@@ -217,7 +216,7 @@ class UnifiedLogger:
         """Save configuration to disk."""
         try:
             with open(self.config_file, "w") as f:
-                yaml.dump(self.config, f, default_flow_style=False, indent=2, sort_keys=False)
+                yaml.dump(self.full_config, f, default_flow_style=False, indent=2, sort_keys=False)
             self.info(f"Config saved to {self.config_file}")
         except Exception as e:
             self.error(f"Error saving config: {e}")
@@ -327,7 +326,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 +362,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