tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
Archived
1
Fork 0
Code Issues Releases Activity

Merge branch 'reward-fix' of github.com:SELab-3-2026/SEL3-2026-Groep-4 into reward-fix

Browse source
This commit is contained in:
Robin Meersman 2026-04-18 11:07:36 +02:00
commit 7e7837d2d1
14 changed files with 92 additions and 140 deletions
Download patch file Download diff file Expand all files Collapse all files

15
.github/workflows/update_hpc_requirements.yml vendored
View file

@ -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] <github-actions[bot]@users.noreply.github.com>"
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] <github-actions[bot]@users.noreply.github.com>"

4
configs/logging/default.yaml
View file

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

9
configs/logging/hpc.yaml Normal file
View file

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

6
configs/logging/wandb_enabled.yaml
View file

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

2
env/hpc/requirements.txt vendored
View file

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

2
scripts/hpc/train.pbs
View file

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

13
scripts/simulate.py
View file

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

6
scripts/train.py
View file

@ -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,
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,
)
logger = get_logger()
logger.info(f"Hydra-initialized run: {run_name}")

28
src/brittle_star_project/trainers/PPOTrainer.py
View file

@ -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
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}...")
self.logger.save_checkpoint(
params=self.agent_state.params, step=iteration, metadata=asdict(self.cfg)
)
def train(self):
"""
@ -622,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,
@ -647,6 +635,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

2
src/experiment_logger/__init__.py
View file

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

25
src/experiment_logger/config_logger.py
View file

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

83
src/experiment_logger/config_utils.py
View file

@ -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}")

8
src/experiment_logger/simple_logger.py
View file

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

29
src/experiment_logger/unified_logger.py
View file

@ -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],
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,
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
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.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()
@ -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