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

Merge branch 'dev' into simulate-results

Browse source
This commit is contained in:
Jona Reynaert 2026-04-19 16:54:47 +02:00
commit c4447976ab
21 changed files with 151 additions and 207 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/environment/directed_locomotion.yaml
View file

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

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

5
configs/morphology/2_arms.yaml Normal file
View file

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

2
configs/ppo/debug.yaml
View file

@ -1,5 +1,5 @@
learning_rate: 0.0003
total_timesteps: 409600
total_timesteps: 409600
num_envs: 32
num_steps: 32
anneal_lr: true

16
configs/ppo/dev_larger_timesteps_larger_rolloutsteps.yaml Normal file
View file

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

19
configs/ppo/fast.yaml
View file

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

2
docs/design/reward_function.md
View file

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

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

3
scripts/simulate.py
View file

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

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

6
src/brittle_star_project/environment/padded_obs_wrapper.py
View file

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

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

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

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