Merge branch 'dev' into simulate-results

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

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

configs/morphology/2_arms.yaml

@@ -0,0 +1,5 @@
+# 2 Arms Morphology Configuration
+
+segments_per_arm: [4, 0, 4, 0, 0]
+use_p_control: true
+use_torque_control: false

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

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

configs/ppo/fast.yaml

@@ -1,19 +0,0 @@
-# Fast PPO Configuration
-# Lower timestep count for quick iterations/testing.
-
-learning_rate: 0.0005
-total_timesteps: 500000
-num_envs: 8
-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

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
 

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

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

@@ -367,6 +367,7 @@ def main(dict_cfg: DictConfig) -> None:
     observations0 = _get_observations(state0)
     env_obs_dim = int(_transform_obs_dict(observations0 or {}).shape[0])
     ckpt_obs_dim = _infer_checkpoint_obs_dim(policy)
+    
     if ckpt_obs_dim is not None and ckpt_obs_dim != env_obs_dim:
         raise ValueError(
             "Checkpoint/env mismatch: "
@@ -374,7 +375,7 @@ def main(dict_cfg: DictConfig) -> None:
             "Use the same Hydra config (morphology/arena/environment) "
             "that was used during training."
         )
-
+    
     # ======= SIMULATION =======
     headless = bool(config.simulation.headless)
     max_steps = config.simulation.max_steps

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/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, ...],
+    segments_per_arm: Sequence[int],
-    reference_segments_per_arm: tuple[int, ...] = (4, 4, 4, 4, 4),
+    reference_segments_per_arm: Sequence[int] = (4, 4, 4, 4, 4),
 ) -> dict[str, Any]:
     """Pre-compute boolean masks for spatial insertion of observations.
 

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 = _reward_fn(env_state, next_env_state)
-    reward *= 20000
-    reward = jnp.clip(reward, -10, 10)
     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],
+                "rewards": storage.rewards,
-                "values": storage.values[0],
+                "values": storage.values,
-                "returns": storage.returns[0],
+                "returns": storage.returns,
-                "advantages": storage.advantages[0],
+                "advantages": storage.advantages,
-                "actions": storage.actions[0],
-                "raw_actions": storage.raw_actions[0],
-                "means": storage.means[0],
-                "stds": storage.stds[0],
-                "logprobs": storage.logprobs[0],
             }
         )
 
-        storage_metrics = {
+        rollout_metrics = {
-            "rollout/env0/return_mean": float(np.mean(data["returns"])),
+            "rollout/reward_mean": float(np.mean(data["rewards"])),
-            "rollout/env0/advantage_mean": float(np.mean(data["advantages"])),
+            "rollout/return_mean": float(np.mean(data["returns"])),
-            "rollout/env0/value_mean": float(np.mean(data["values"])),
+            "rollout/value_mean": float(np.mean(data["values"])),
-            "rollout/env0/value_vs_return_diff": float(np.mean(data["values"] - data["returns"])),
+            "rollout/advantage_mean": float(np.mean(data["advantages"])),
-            "rollout/env0/reward_mean": float(np.mean(data["rewards"])),
+            "rollout/advantage_std": float(np.std(data["advantages"])),
-            "rollout/env0/mean_mean": float(np.mean(data["means"])),
+            "rollout/value_vs_return_mse": float(np.mean((data["values"] - data["returns"]) ** 2)),
-            "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"])),
         }
 
-        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/episodic_return": training_measurements.avg_episodic_return,
-            "charts/avg_episodic_length": np.mean(
+            "charts/episodic_length": float(
-                jax.device_get(episode_stats.returned_episode_lengths)
+                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:
@@ -573,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):
         """
@@ -620,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,
@@ -630,8 +620,6 @@ class PPOTrainer:
                 iteration_time_start,
                 training_measurements,
                 storage,
-                next_obs,
-                xy_distance,
             )
 
             sps = int(global_step / (time.time() - start_time))
@@ -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

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