From 8b6dcbae7c899a02919f288341a224f2813f6d2a Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Thu, 16 Apr 2026 12:27:56 +0200 Subject: [PATCH 01/12] feat: checkpoints --- .../trainers/PPOTrainer.py | 27 ++++++++----------- src/brittle_star_project/trainers/utils.py | 21 +++++++++++++++ 2 files changed, 32 insertions(+), 16 deletions(-) create mode 100644 src/brittle_star_project/trainers/utils.py diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index e5ecca6..ea77f12 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -24,6 +24,7 @@ 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 = { @@ -573,23 +574,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) - from dataclasses import asdict as _asdict - - config_dict = { - "experiment": _asdict(self.experiment), - "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 _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) def train(self): """ @@ -647,6 +638,10 @@ class PPOTrainer: f"ETA {eta_str}" ) + if self.logging_cfg.save_model 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 diff --git a/src/brittle_star_project/trainers/utils.py b/src/brittle_star_project/trainers/utils.py new file mode 100644 index 0000000..e271366 --- /dev/null +++ b/src/brittle_star_project/trainers/utils.py @@ -0,0 +1,21 @@ +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 From 06aa1c54c6426b45a3a3cfaf4112ffb00a8426b2 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Thu, 16 Apr 2026 12:37:09 +0200 Subject: [PATCH 02/12] ci: create pull request --- .github/workflows/update_hpc_requirements.yml | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/.github/workflows/update_hpc_requirements.yml b/.github/workflows/update_hpc_requirements.yml index 38dbc42..969b433 100644 --- a/.github/workflows/update_hpc_requirements.yml +++ b/.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 - uses: stefanzweifel/git-auto-commit-action@v5 + - name: Create Pull Request with updated requirements + uses: peter-evans/create-pull-request@v6 with: - commit_message: "chore(hpc): update env/hpc/requirements.txt from pyproject.toml [skip ci]" - file_pattern: env/hpc/requirements.txt - commit_author: "github-actions[bot] " + token: ${{ secrets.GITHUB_TOKEN }} + commit-message: "chore(hpc): update env/hpc/requirements.txt from pyproject.toml" + 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] " From 4d0f729aee33d88193465cad412ad8d3fa3a6e88 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Thu, 16 Apr 2026 14:35:23 +0200 Subject: [PATCH 03/12] chore(log): remove dead code --- src/experiment_logger/__init__.py | 2 - src/experiment_logger/config_utils.py | 83 --------------------------- 2 files changed, 85 deletions(-) delete mode 100644 src/experiment_logger/config_utils.py diff --git a/src/experiment_logger/__init__.py b/src/experiment_logger/__init__.py index 53e57c0..e1b2d09 100644 --- a/src/experiment_logger/__init__.py +++ b/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" diff --git a/src/experiment_logger/config_utils.py b/src/experiment_logger/config_utils.py deleted file mode 100644 index 80d1748..0000000 --- a/src/experiment_logger/config_utils.py +++ /dev/null @@ -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}") From 37a4b59e04c8a43cd1a25dd6981481949c1d1896 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Thu, 16 Apr 2026 15:10:21 +0200 Subject: [PATCH 04/12] 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 From c5b08e817e82bf352d124e13a1abd4b8fdfbcbb7 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Thu, 16 Apr 2026 15:18:40 +0200 Subject: [PATCH 05/12] refactor(log): use dataclass for config --- scripts/train.py | 8 ++----- src/experiment_logger/simple_logger.py | 8 +++---- src/experiment_logger/unified_logger.py | 29 ++++++++++--------------- 3 files changed, 17 insertions(+), 28 deletions(-) diff --git a/scripts/train.py b/scripts/train.py index a338484..c367000 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -36,13 +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, - project_name=config.logging.wandb_project_name, - entity=config.logging.wandb_entity, + full_config=cfg_dict, + logging_cfg=config.logging, 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/experiment_logger/simple_logger.py b/src/experiment_logger/simple_logger.py index 0ed1ab5..7e4a816 100644 --- a/src/experiment_logger/simple_logger.py +++ b/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, - project_name: str = "none", - entity: Optional[str] = None, + full_config: Optional[Dict[str, Any]] = None, + logging_cfg: Optional[Any] = 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): diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py index c98da37..634638d 100644 --- a/src/experiment_logger/unified_logger.py +++ b/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,13 +89,9 @@ class UnifiedLogger: def __init__( self, run_name: str, - config: Dict[str, Any], - project_name: str = "PPO-Modularity", - entity: Optional[str] = None, + full_config: Dict[str, Any], + logging_cfg: LoggingConfig, 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, ): @@ -102,18 +99,16 @@ class UnifiedLogger: Args: run_name: Unique name for this run - config: Configuration dictionary with hyperparameters - project_name: WandB project name - entity: WandB entity (team/user name) + full_config: Full configuration dictionary with hyperparameters to be saved + logging_cfg: Structured logging configuration dataclass 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.use_wandb = use_wandb - self.upload_final_model = upload_final_model - self.upload_checkpoints = upload_checkpoints + self.full_config = full_config + 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() @@ -163,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]] = [] @@ -211,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", ) @@ -221,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}") From 4581694fa6e06592c9560d041968d33030b14291 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Thu, 16 Apr 2026 15:25:48 +0200 Subject: [PATCH 06/12] style: ruff format + check --- scripts/simulate.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/scripts/simulate.py b/scripts/simulate.py index b49e50f..3c4aa42 100644 --- a/scripts/simulate.py +++ b/scripts/simulate.py @@ -59,7 +59,7 @@ def main(dict_cfg: DictConfig) -> None: 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. + # 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(): @@ -69,7 +69,7 @@ def main(dict_cfg: DictConfig) -> None: # 'critic_params': FrozenDict({...}), # ... # } - # Update the RLModel.load function or subsequent destructuring to support this raw dictionary natively. + # Update to support this raw dictionary natively. policy = RLModel.load(Path(model_path)) if hasattr(policy, "nu"): policy.nu = nu From ad17401c4ce82b4d651246e08b62633c53afd35b Mon Sep 17 00:00:00 2001 From: Robin Meersman Date: Fri, 17 Apr 2026 14:23:57 +0200 Subject: [PATCH 07/12] fix: custom reward function dependent on env reward + extensions --- configs/environment/directed_locomotion.yaml | 4 +- configs/logging/wandb_enabled.yaml | 2 +- configs/morphology/2_arms.yaml | 5 ++ configs/ppo/debug.yaml | 2 +- configs/ppo/fast.yaml | 2 +- configs/ppo/fast2.yaml | 16 ++++ .../environment/BrittleStarJaxEnvWrapper.py | 2 +- .../environment/padded_obs_wrapper.py | 6 +- .../trainers/PPOTrainer.py | 82 +++++++++---------- 9 files changed, 71 insertions(+), 50 deletions(-) create mode 100644 configs/morphology/2_arms.yaml create mode 100644 configs/ppo/fast2.yaml diff --git a/configs/environment/directed_locomotion.yaml b/configs/environment/directed_locomotion.yaml index 63ad520..465d80e 100644 --- a/configs/environment/directed_locomotion.yaml +++ b/configs/environment/directed_locomotion.yaml @@ -2,11 +2,11 @@ # Baseline task setting. task: DIRECTED_LOCOMOTION -simulation_time: 5.0 +simulation_time: 5000.0 num_physics_steps_per_control_step: 10 time_scale: 2 camera_ids: [0, 1] render_size: [480, 640] joint_randomization_noise_scale: 0.0 -target_distance: 3.0 +target_distance: 0.6 light_perlin_noise_scale: 0 diff --git a/configs/logging/wandb_enabled.yaml b/configs/logging/wandb_enabled.yaml index 2e82781..40c52bf 100644 --- a/configs/logging/wandb_enabled.yaml +++ b/configs/logging/wandb_enabled.yaml @@ -2,7 +2,7 @@ # For production/cloud experiments with weights synced. track: true -wandb_project_name: "PPO-Modularity" +wandb_project_name: "PPO-Modularity - reward engineering" wandb_entity: "SEL3-2026-Groep-4" capture_video: false save_model: true diff --git a/configs/morphology/2_arms.yaml b/configs/morphology/2_arms.yaml new file mode 100644 index 0000000..dc6f159 --- /dev/null +++ b/configs/morphology/2_arms.yaml @@ -0,0 +1,5 @@ +# 2 Arms Morphology Configuration + +segments_per_arm: [4, 4] +use_p_control: true +use_torque_control: false \ No newline at end of file diff --git a/configs/ppo/debug.yaml b/configs/ppo/debug.yaml index 58dbd3a..7732fd3 100644 --- a/configs/ppo/debug.yaml +++ b/configs/ppo/debug.yaml @@ -1,5 +1,5 @@ learning_rate: 0.0003 -total_timesteps: 409600 +total_timesteps: 409600 num_envs: 32 num_steps: 32 anneal_lr: true diff --git a/configs/ppo/fast.yaml b/configs/ppo/fast.yaml index 5001ef2..ecd65fa 100644 --- a/configs/ppo/fast.yaml +++ b/configs/ppo/fast.yaml @@ -3,7 +3,7 @@ learning_rate: 0.0005 total_timesteps: 500000 -num_envs: 8 +num_envs: 32 num_steps: 128 anneal_lr: true gamma: 0.99 diff --git a/configs/ppo/fast2.yaml b/configs/ppo/fast2.yaml new file mode 100644 index 0000000..b2f629a --- /dev/null +++ b/configs/ppo/fast2.yaml @@ -0,0 +1,16 @@ +learning_rate: 0.0003 +total_timesteps: 1228800 +num_envs: 32 +num_steps: 64 +anneal_lr: true +gamma: 0.99 +gae_lambda: 0.95 +num_minibatches: 32 +update_epochs: 4 +norm_adv: true +clip_coef: 0.2 +clip_vloss: true +ent_coef: 0.005 +vf_coef: 1.0 +max_grad_norm: 0.5 +target_kl: null \ No newline at end of file diff --git a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py index a5175a7..3122696 100644 --- a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py +++ b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py @@ -27,7 +27,7 @@ class BrittleStarJaxEnvWrapper: ) # Pre-compute masks for observation padding - self._padding_masks = compute_padding_masks(self._morphology.segments_per_arm) + self._padding_masks = compute_padding_masks(self._morphology.segments_per_arm, (4, 4)) self._vectorized_reset = jax.jit(jax.vmap(self._env.reset)) self._vectorized_step = jax.jit(jax.vmap(self._env.step)) diff --git a/src/brittle_star_project/environment/padded_obs_wrapper.py b/src/brittle_star_project/environment/padded_obs_wrapper.py index 3f22038..4886284 100644 --- a/src/brittle_star_project/environment/padded_obs_wrapper.py +++ b/src/brittle_star_project/environment/padded_obs_wrapper.py @@ -8,7 +8,7 @@ flattened observation maintains the correct physical mapping to the neural netwo from __future__ import annotations -from typing import Any +from typing import Any, Sequence import jax.numpy as jnp # Observation keys whose size scales with the number of joints (2 per segment). @@ -30,8 +30,8 @@ _SEGMENT_SCALED_KEYS = frozenset( def compute_padding_masks( - segments_per_arm: tuple[int, ...], - reference_segments_per_arm: tuple[int, ...] = (4, 4, 4, 4, 4), + segments_per_arm: Sequence[int], + reference_segments_per_arm: Sequence[int] = (4, 4, 4, 4, 4), ) -> dict[str, Any]: """Pre-compute boolean masks for spatial insertion of observations. diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index e5ecca6..3b102d4 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -151,12 +151,27 @@ def _step_once( return (agent_state, episode_stats, next_obs, next_done, key, env_state), storage +def _reward_fn(env_state, next_env_state): + # if delta distance positive ==> brittle star walking away from target + delta_distance = ( + next_env_state.observations["xy_distance_to_target"] + - env_state.observations["xy_distance_to_target"] + ).squeeze(-1) + + env_reward = next_env_state.reward + clipped_env_reward = jnp.clip(100 * env_reward, -10, 10) + + time_penalty = 0.1 + distance_penalty = jnp.clip(0.5 * delta_distance, -0.5, 0.5) + penalty = time_penalty + distance_penalty + + return jnp.where(next_env_state.terminated, 50.0, clipped_env_reward - penalty) + + def _step_env_wrapped(episode_stats, env_state, action, env_step_fn): next_env_state = env_step_fn(env_state, action) - reward = next_env_state.reward - reward *= 20000 - reward = jnp.clip(reward, -10, 10) + reward = _reward_fn(env_state, next_env_state) terminated = next_env_state.terminated truncated = next_env_state.truncated done = terminated | truncated @@ -440,62 +455,49 @@ class PPOTrainer: iteration_time_start, training_measurements, storage, - next_obs, - xy_distance, ): data = jax.device_get( { - "rewards": storage.rewards[0], - "values": storage.values[0], - "returns": storage.returns[0], - "advantages": storage.advantages[0], - "actions": storage.actions[0], - "raw_actions": storage.raw_actions[0], - "means": storage.means[0], - "stds": storage.stds[0], - "logprobs": storage.logprobs[0], + "rewards": storage.rewards, + "values": storage.values, + "returns": storage.returns, + "advantages": storage.advantages, } ) - storage_metrics = { - "rollout/env0/return_mean": float(np.mean(data["returns"])), - "rollout/env0/advantage_mean": float(np.mean(data["advantages"])), - "rollout/env0/value_mean": float(np.mean(data["values"])), - "rollout/env0/value_vs_return_diff": float(np.mean(data["values"] - data["returns"])), - "rollout/env0/reward_mean": float(np.mean(data["rewards"])), - "rollout/env0/mean_mean": float(np.mean(data["means"])), - "rollout/env0/logprob_mean": float(np.mean(data["logprobs"])), - "rollout/env0/action_mean": float(np.mean(data["actions"])), - "rollout/env0/raw_action_mean": float(np.mean(data["raw_actions"])), + rollout_metrics = { + "rollout/reward_mean": float(np.mean(data["rewards"])), + "rollout/return_mean": float(np.mean(data["returns"])), + "rollout/value_mean": float(np.mean(data["values"])), + "rollout/advantage_mean": float(np.mean(data["advantages"])), + "rollout/advantage_std": float(np.std(data["advantages"])), + "rollout/value_vs_return_mse": float(np.mean((data["values"] - data["returns"]) ** 2)), } - for i in range(len(xy_distance)): - storage_metrics[f"env_data/env{i}_xy_dist_target"] = float(xy_distance[i]) - metrics = { - "charts/avg_episodic_return": training_measurements.avg_episodic_return, - "charts/avg_episodic_length": np.mean( - jax.device_get(episode_stats.returned_episode_lengths) + "charts/episodic_return": training_measurements.avg_episodic_return, + "charts/episodic_length": float( + np.mean(jax.device_get(episode_stats.returned_episode_lengths)) ), - "charts/learning_rate": self.agent_state.opt_state[1] - .hyperparams["learning_rate"] - .item(), "charts/explained_variance": training_measurements.explained_variance, - "charts/num_terminated": training_measurements.num_terminated, - "charts/num_truncated": training_measurements.num_truncated, - "charts/avg_terminated_ep_length": training_measurements.avg_terminated_length, - "charts/avg_truncated_ep_length": training_measurements.avg_truncated_length, "losses/value_loss": training_measurements.v_loss[-1, -1].item(), "losses/policy_loss": training_measurements.pg_loss[-1, -1].item(), "losses/entropy": training_measurements.entropy_loss[-1, -1].item(), "losses/approx_kl": training_measurements.approx_kl[-1, -1].item(), - "losses/loss": training_measurements.loss[-1, -1].item(), + "charts/learning_rate": self.agent_state.opt_state[1] + .hyperparams["learning_rate"] + .item(), "charts/SPS": int(global_step / (time.time() - start_time)), "charts/SPS_update": int( self.ppo.num_envs * self.ppo.num_steps / (time.time() - iteration_time_start) ), - **storage_metrics, + "termi_trunci/num_terminated": training_measurements.num_terminated, + "termi_trunci/num_truncated": training_measurements.num_truncated, + "termi_trunci/avg_terminated_ep_length": training_measurements.avg_terminated_length, + "termi_trunci/avg_truncated_ep_length": training_measurements.avg_truncated_length, + **rollout_metrics, } + self.logger.log(metrics, step=global_step) def _step(self, env_state, next_obs, next_done, iteration: int) -> tuple: @@ -630,8 +632,6 @@ class PPOTrainer: iteration_time_start, training_measurements, storage, - next_obs, - xy_distance, ) sps = int(global_step / (time.time() - start_time)) From c880285dc293779b75e2f79242367915e5a54509 Mon Sep 17 00:00:00 2001 From: Robin Meersman Date: Fri, 17 Apr 2026 14:28:11 +0200 Subject: [PATCH 08/12] feat(docs): added delta_distance extension to documentation for reward function --- docs/design/reward_function.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/design/reward_function.md b/docs/design/reward_function.md index a72fb33..9cf178f 100644 --- a/docs/design/reward_function.md +++ b/docs/design/reward_function.md @@ -8,6 +8,8 @@ inputs must be distributed fairly to guarantee an objective comparison between d - The reward function is centered around minimizing the distance to the goal or maximizing the movement towards the goal within a finite number of timesteps $T$. - To motivate efficient movement, the amount of timesteps taken to reach the goal will be used as penalty. +- An extra penalty based on movement relative to the current step and +the previous is used to penalize a movement away from the target. ## From reward to PPO From 2ef2597430cb04c3b56a890d4b01513a06f1cd3d Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Fri, 17 Apr 2026 22:09:20 +0200 Subject: [PATCH 09/12] fix: remove unused var --- src/brittle_star_project/trainers/PPOTrainer.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index f78821b..7cdc1ff 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -612,8 +612,6 @@ class PPOTrainer: self._update_obs_stats(next_obs) next_obs = _normalize_obs(next_obs, self.obs_mean, self.obs_var) - xy_distance = _get_xy_distance_to_target(env_state.observations) - global_step += self.ppo.num_steps * self.ppo.num_envs self._log( global_step, From 6c47883bc90b25b7c71a82470308550f8445dfed Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Fri, 17 Apr 2026 22:23:23 +0200 Subject: [PATCH 10/12] fix(hpc): outdated requirements --- env/hpc/requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/env/hpc/requirements.txt b/env/hpc/requirements.txt index 84515f9..c85ce2d 100644 --- a/env/hpc/requirements.txt +++ b/env/hpc/requirements.txt @@ -15,6 +15,6 @@ optax>=0.2.6 pyopengl>=3.1.10 pyopengl-accelerate>=3.1.10 pyyaml>=6.0 -tyro>=1.0.10 +hydra-core>=1.3.2 wandb==0.24.2 torch>=2.4.0 From 34e8ba91c21a46f06cea9b91dd3863dfe48b1a8b Mon Sep 17 00:00:00 2001 From: Robin Meersman Date: Sat, 18 Apr 2026 11:07:01 +0200 Subject: [PATCH 11/12] fix(PR#41): applied comments for PR #41 review --- configs/morphology/2_arms.yaml | 2 +- configs/ppo/fast.yaml | 19 ------------------- configs/ppo/fast2.yaml | 16 ---------------- .../environment/BrittleStarJaxEnvWrapper.py | 2 +- 4 files changed, 2 insertions(+), 37 deletions(-) delete mode 100644 configs/ppo/fast.yaml delete mode 100644 configs/ppo/fast2.yaml diff --git a/configs/morphology/2_arms.yaml b/configs/morphology/2_arms.yaml index dc6f159..c301ec4 100644 --- a/configs/morphology/2_arms.yaml +++ b/configs/morphology/2_arms.yaml @@ -1,5 +1,5 @@ # 2 Arms Morphology Configuration -segments_per_arm: [4, 4] +segments_per_arm: [4, 0, 4, 0, 0] use_p_control: true use_torque_control: false \ No newline at end of file diff --git a/configs/ppo/fast.yaml b/configs/ppo/fast.yaml deleted file mode 100644 index ecd65fa..0000000 --- a/configs/ppo/fast.yaml +++ /dev/null @@ -1,19 +0,0 @@ -# Fast PPO Configuration -# Lower timestep count for quick iterations/testing. - -learning_rate: 0.0005 -total_timesteps: 500000 -num_envs: 32 -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/configs/ppo/fast2.yaml b/configs/ppo/fast2.yaml deleted file mode 100644 index b2f629a..0000000 --- a/configs/ppo/fast2.yaml +++ /dev/null @@ -1,16 +0,0 @@ -learning_rate: 0.0003 -total_timesteps: 1228800 -num_envs: 32 -num_steps: 64 -anneal_lr: true -gamma: 0.99 -gae_lambda: 0.95 -num_minibatches: 32 -update_epochs: 4 -norm_adv: true -clip_coef: 0.2 -clip_vloss: true -ent_coef: 0.005 -vf_coef: 1.0 -max_grad_norm: 0.5 -target_kl: null \ No newline at end of file diff --git a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py index 3122696..a5175a7 100644 --- a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py +++ b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py @@ -27,7 +27,7 @@ class BrittleStarJaxEnvWrapper: ) # Pre-compute masks for observation padding - self._padding_masks = compute_padding_masks(self._morphology.segments_per_arm, (4, 4)) + self._padding_masks = compute_padding_masks(self._morphology.segments_per_arm) self._vectorized_reset = jax.jit(jax.vmap(self._env.reset)) self._vectorized_step = jax.jit(jax.vmap(self._env.step)) From dfcdfdad2b4d759383c4781bd580409d3c3219e8 Mon Sep 17 00:00:00 2001 From: Robin Meersman Date: Sat, 18 Apr 2026 11:07:26 +0200 Subject: [PATCH 12/12] fix: renamed fast2 to better describing name --- ...dev_larger_timesteps_larger_rolloutsteps.yaml | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 configs/ppo/dev_larger_timesteps_larger_rolloutsteps.yaml diff --git a/configs/ppo/dev_larger_timesteps_larger_rolloutsteps.yaml b/configs/ppo/dev_larger_timesteps_larger_rolloutsteps.yaml new file mode 100644 index 0000000..b2f629a --- /dev/null +++ b/configs/ppo/dev_larger_timesteps_larger_rolloutsteps.yaml @@ -0,0 +1,16 @@ +learning_rate: 0.0003 +total_timesteps: 1228800 +num_envs: 32 +num_steps: 64 +anneal_lr: true +gamma: 0.99 +gae_lambda: 0.95 +num_minibatches: 32 +update_epochs: 4 +norm_adv: true +clip_coef: 0.2 +clip_vloss: true +ent_coef: 0.005 +vf_coef: 1.0 +max_grad_norm: 0.5 +target_kl: null \ No newline at end of file