From ac352ef431d243f04aa4069ce8a030ea581e68a5 Mon Sep 17 00:00:00 2001 From: cedric Date: Fri, 10 Apr 2026 10:17:18 +0000 Subject: [PATCH 01/32] feat: start of debug setup, experiment description,.. --- configs/hpc/debug.yaml | 13 +++ .../used_variables.md | 82 +++++++++++++++++++ scripts/train.py | 14 +++- 3 files changed, 108 insertions(+), 1 deletion(-) create mode 100644 configs/hpc/debug.yaml create mode 100644 experiments/debug-experiment-10042026/used_variables.md diff --git a/configs/hpc/debug.yaml b/configs/hpc/debug.yaml new file mode 100644 index 0000000..d62df3e --- /dev/null +++ b/configs/hpc/debug.yaml @@ -0,0 +1,13 @@ +# Configuration for debug session +exp_name: "debug-experiment-10042026" # started on april 10 +seed: 42 +track: true +wandb_project_name: "Let's-find-that-bug" +wandb_entity: "SEL3-2026-Groep-4" + +num_envs: 32 +num_steps: 32 +total_timesteps: 102400 + +cuda: true + diff --git a/experiments/debug-experiment-10042026/used_variables.md b/experiments/debug-experiment-10042026/used_variables.md new file mode 100644 index 0000000..ed10279 --- /dev/null +++ b/experiments/debug-experiment-10042026/used_variables.md @@ -0,0 +1,82 @@ +## Default envconfig +task: Task = Task.DIRECTED_LOCOMOTION +simulation_time: float = 5.0 +num_physics_steps_per_control_step: int = 10 +time_scale: int = 2 +camera_ids: list[int] = field(default_factory=lambda: [0, 1]) +render_size: tuple[int, int] = (480, 640) +joint_randomization_noise_scale: float = 0.0 +target_distance: float = 3.0 +light_perlin_noise_scale: int = 0 + + +## Default ppoargs +seed: int = 1 +torch_deterministic: bool = True +cuda: bool = True +track: bool = False +checkpoint_frequency: int = 100 +learning_rate: float = 2.5e-4 +anneal_lr: bool = True +gamma: float = 0.99 +gae_lambda: float = 0.95 +num_minibatches: int = 4 +update_epochs: int = 4 +norm_adv: bool = True +clip_coef: float = 0.1 +clip_vloss: bool = True +ent_coef: float = 0.01 +vf_coef: float = 0.5 +max_grad_norm: float = 0.5 +target_kl: float | None = None +batch_size: int = 0 +minibatch_size: int = 0 +num_iterations: int = 0 + +## Used config file: +num_envs: 32 +num_steps: 32 +total_timesteps: 102400 + +## Arena config: +size: tuple[float, float] = (10.0, 5.0) +sand_ground_color: bool = True +attach_target: bool = True +wall_height: float = 1.5 +wall_thickness: float = 0.1 + +## Morphology: +num_arms: int = 5 +num_segments_per_arm: int = 4 +use_p_control: bool = True +use_torque_control: bool = False + +## MLPs: +### Sensor & Feature_extractor: +class GenericDenseLayersWithActivation(nn.Module): + layer_sizes: Sequence[int] = field(default_factory=lambda: [64, 64]) + activation: Callable = nn.tanh + + @nn.compact + def __call__(self, x): + for size in self.layer_sizes: + x = nn.Dense(size, kernel_init=orthogonal(jnp.sqrt(2)))(x) + x = self.activation(x) + return x + +### Actor: +class Actor(nn.Module): + action_dim: int + @nn.compact + def __call__(self, x): + mean = nn.Dense(self.action_dim, kernel_init=orthogonal(0.01), bias_init=constant(0.0))(x) + log_std = self.param("log_std", nn.initializers.zeros, (self.action_dim,)) + return mean, log_std + +### Critic: +class OneDenseLayerMLP(nn.Module): + @nn.compact + def __call__(self, x): + return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x) + +### Observations: diff --git a/scripts/train.py b/scripts/train.py index 492c887..5c646bf 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -68,7 +68,19 @@ if __name__ == "__main__": print_config(args, title="PPO Training Configuration") env = make_env(args.env_config_path, args.num_envs) - + raw_env = env.raw + print( + "\n\n\n Observation space \n", + raw_env.observation_space, + "Action space \n", + raw_env.action_space, + ) + print( + "\n\n\n Observation space \n", + raw_env.observation_space, + "Action space \n", + raw_env.action_space, + ) torch.backends.cudnn.deterministic = args.torch_deterministic ppo_trainer = PPOTrainer(args, env, run_dir, run_name) From 1ef257544a0e08942500690420dfd31bb706ab7e Mon Sep 17 00:00:00 2001 From: JibrilExe Date: Fri, 10 Apr 2026 15:18:33 +0200 Subject: [PATCH 02/32] feat: first approach to action clipping --- configs/hpc/debug.yaml | 3 ++- .../used_variables.md | 2 +- scripts/hpc/train.pbs | 4 +-- scripts/train.py | 14 +++------- .../trainers/PPOTrainer.py | 27 ++++++++++++++++--- 5 files changed, 32 insertions(+), 18 deletions(-) diff --git a/configs/hpc/debug.yaml b/configs/hpc/debug.yaml index d62df3e..20c65e8 100644 --- a/configs/hpc/debug.yaml +++ b/configs/hpc/debug.yaml @@ -8,6 +8,7 @@ wandb_entity: "SEL3-2026-Groep-4" num_envs: 32 num_steps: 32 total_timesteps: 102400 - +num_arms: 2 +num_segments_per_arm: 1 cuda: true diff --git a/experiments/debug-experiment-10042026/used_variables.md b/experiments/debug-experiment-10042026/used_variables.md index ed10279..0dcba35 100644 --- a/experiments/debug-experiment-10042026/used_variables.md +++ b/experiments/debug-experiment-10042026/used_variables.md @@ -46,7 +46,7 @@ wall_height: float = 1.5 wall_thickness: float = 0.1 ## Morphology: -num_arms: int = 5 +num_arms: int = 2 num_segments_per_arm: int = 4 use_p_control: bool = True use_torque_control: bool = False diff --git a/scripts/hpc/train.pbs b/scripts/hpc/train.pbs index e7d21f0..e4ce355 100644 --- a/scripts/hpc/train.pbs +++ b/scripts/hpc/train.pbs @@ -65,8 +65,8 @@ fi # TODO Once experiments get serious, change the config python scripts/train.py \ - --env-config-path configs/hpc/smoke_test.yaml \ - --hyperparameter-config-path configs/hpc/smoke_test.yaml \ + --env-config-path configs/hpc/debug.yaml \ + --hyperparameter-config-path configs/hpc/debug.yaml \ --run-dir "$SCRATCH_RUNDIR" echo ">>> Staging out results to $DATA_RUNDIR..." diff --git a/scripts/train.py b/scripts/train.py index 5c646bf..1f543f7 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -69,18 +69,10 @@ if __name__ == "__main__": env = make_env(args.env_config_path, args.num_envs) raw_env = env.raw - print( - "\n\n\n Observation space \n", - raw_env.observation_space, - "Action space \n", - raw_env.action_space, - ) - print( - "\n\n\n Observation space \n", - raw_env.observation_space, - "Action space \n", - raw_env.action_space, + logger.log( + {"run_dir": run_dir} ) + torch.backends.cudnn.deterministic = args.torch_deterministic ppo_trainer = PPOTrainer(args, env, run_dir, run_name) diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 8c238f1..432030a 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -24,6 +24,9 @@ from brittle_star_project.MLPs.mlps import ( ) from brittle_star_project.ppo import PPO +@jax.jit +def _clip_action(action: jnp.ndarray, low: jnp.ndarray, high: jnp.ndarray) -> jnp.ndarray: + return jnp.clip(action, low, high) @jax.jit def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, learning_rate): @@ -47,6 +50,8 @@ def _get_action_and_value_noise( agent_state: TrainState, next_obs: jnp.ndarray, key: jax.random.PRNGKey, + action_low, + action_high ): hidden = sensor.apply(agent_state.params["sensor_params"], next_obs) hidden_critic = feature_extractor.apply( @@ -59,9 +64,14 @@ def _get_action_and_value_noise( noise = jax.random.normal(subkey, shape=mean.shape) std = jnp.exp(log_std) action = mean + noise * std - logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1) + clipped_action = _clip_action( + action, + action_low, + action_high + ) + logprob = -0.5 * (((clipped_action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1) value = critic.apply(agent_state.params["critic_params"], hidden_critic) - return action, logprob, value.squeeze(-1), key + return clipped_action, logprob, value.squeeze(-1), key # removed jit: used in _rollout_jit, so will be compiled with _rollout_jit @@ -73,10 +83,12 @@ def _step_once( feature_extractor: GenericDenseLayersWithActivation, actor: Actor, critic: OneDenseLayerMLP, + action_low, + action_high ): agent_state, episode_stats, obs, done, key, env_state = carry action, logprob, value, key = _get_action_and_value_noise( - sensor, feature_extractor, actor, critic, agent_state, obs, key + sensor, feature_extractor, actor, critic, agent_state, obs, key, action_low, action_high ) episode_stats, env_state, (next_obs, reward, next_done) = env_step_fn( @@ -140,6 +152,8 @@ def _rollout_jit( feature_extractor: GenericDenseLayersWithActivation, actor: Actor, critic: OneDenseLayerMLP, + action_low, + action_high ): (agent_state, episode_stats, next_obs, next_done, key, env_state), storage = jax.lax.scan( partial( @@ -149,6 +163,8 @@ def _rollout_jit( actor=actor, critic=critic, env_step_fn=step_env_fn, + action_low=action_low, + action_high=action_high ), (agent_state, episode_stats, next_obs, next_done, key, env_state), (), @@ -215,6 +231,9 @@ class PPOTrainer: self.actor.apply = jax.jit(self.actor.apply) self.critic.apply = jax.jit(self.critic.apply) + action_low = jnp.asarray(self.env.single_action_space.low, dtype=jnp.float32) + action_high = jnp.asarray(self.env.single_action_space.high, dtype=jnp.float32) + self._rollout_jit = jax.jit( partial( _rollout_jit, @@ -224,6 +243,8 @@ class PPOTrainer: feature_extractor=self.feature_extractor, actor=self.actor, critic=self.critic, + action_low=action_low, + action_high=action_high ) ) self._compute_gae_jit = jax.jit( From c0e7569773c17c1ded9e70a7a2ec18b1cebad3e6 Mon Sep 17 00:00:00 2001 From: JibrilExe Date: Sat, 11 Apr 2026 13:42:52 +0200 Subject: [PATCH 03/32] feat: More low lvl logs, reward, advantage, returns --- configs/hpc/debug.yaml | 4 +- scripts/train.py | 5 +- src/brittle_star_project/MLPs/mlps.py | 4 ++ .../trainers/PPOTrainer.py | 57 +++++++++++++++---- 4 files changed, 54 insertions(+), 16 deletions(-) diff --git a/configs/hpc/debug.yaml b/configs/hpc/debug.yaml index 20c65e8..c387dc1 100644 --- a/configs/hpc/debug.yaml +++ b/configs/hpc/debug.yaml @@ -1,10 +1,10 @@ # Configuration for debug session -exp_name: "debug-experiment-10042026" # started on april 10 +exp_name: "debug-experiment" seed: 42 track: true wandb_project_name: "Let's-find-that-bug" wandb_entity: "SEL3-2026-Groep-4" - + num_envs: 32 num_steps: 32 total_timesteps: 102400 diff --git a/scripts/train.py b/scripts/train.py index 1f543f7..886bd54 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -50,7 +50,7 @@ if __name__ == "__main__": run_name = f"{args.exp_name}__seed_{args.seed}__{git_hash}__{int(time.time())}" if args.run_dir is None: - run_dir = f"runs/{run_name}" + run_dir = f"/data/gent/465/vsc46589/runs/{run_name}" else: run_dir = args.run_dir @@ -69,9 +69,6 @@ if __name__ == "__main__": env = make_env(args.env_config_path, args.num_envs) raw_env = env.raw - logger.log( - {"run_dir": run_dir} - ) torch.backends.cudnn.deterministic = args.torch_deterministic diff --git a/src/brittle_star_project/MLPs/mlps.py b/src/brittle_star_project/MLPs/mlps.py index 6abb540..9568a36 100644 --- a/src/brittle_star_project/MLPs/mlps.py +++ b/src/brittle_star_project/MLPs/mlps.py @@ -58,6 +58,10 @@ class Storage: returns: jnp.array rewards: jnp.array + raw_actions: jnp.ndarray = None # before clipping + means: jnp.ndarray = None # policy mean + stds: jnp.ndarray = None # policy std + def replace(self, **kwargs) -> "Storage": fs = fields(self) return Storage(**{f.name: kwargs.get(f.name, getattr(self, f.name)) for f in fs}) diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 1e554f8..8a9b341 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -11,8 +11,6 @@ import numpy as np import optax from flax.training.train_state import TrainState -from experiment_logger import get_logger - from brittle_star_project.dataclasses import EpisodeStatistics, PPOArgs from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper from brittle_star_project.MLPs.mlps import ( @@ -23,6 +21,7 @@ from brittle_star_project.MLPs.mlps import ( Storage, ) from brittle_star_project.ppo import PPO +from experiment_logger import get_logger @jax.jit def _clip_action(action: jnp.ndarray, low: jnp.ndarray, high: jnp.ndarray) -> jnp.ndarray: @@ -67,15 +66,17 @@ def _get_action_and_value_noise( key, subkey = jax.random.split(key) noise = jax.random.normal(subkey, shape=mean.shape) std = jnp.exp(log_std) - action = mean + noise * std + raw_action = mean + noise * std clipped_action = _clip_action( - action, + raw_action, action_low, action_high ) - logprob = -0.5 * (((clipped_action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1) + logprob = -0.5 * (((raw_action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1) value = critic.apply(agent_state.params["critic_params"], hidden_critic) - return clipped_action, logprob, value.squeeze(-1), key + + return clipped_action, raw_action, logprob, value.squeeze(-1), mean, std, key + def _step_once( @@ -90,21 +91,24 @@ def _step_once( action_high ): agent_state, episode_stats, obs, done, key, env_state = carry - action, logprob, value, key = _get_action_and_value_noise( + clipped_action, raw_action, logprob, value, mean, std, key = _get_action_and_value_noise( sensor, feature_extractor, actor, critic, agent_state, obs, key, action_low, action_high ) episode_stats, env_state, (next_obs, reward, next_done) = env_step_fn( - episode_stats, env_state, action + episode_stats, env_state, clipped_action ) storage = Storage( obs=obs, - actions=action, + actions=clipped_action, + raw_actions=raw_action, logprobs=logprob, dones=done, values=value, rewards=reward, + means=mean, + stds=std, returns=jnp.zeros_like(reward), advantages=jnp.zeros_like(reward), ) @@ -371,7 +375,37 @@ class PPOTrainer: start_time, iteration_time_start, training_measurements, + storage ): + data = jax.device_get({ + 'rewards': storage.rewards[0], # (num_steps,) + 'values': storage.values[0], + 'returns': storage.returns[0], # NEW + 'advantages': storage.advantages[0],# NEW + '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/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'])), + } + metrics = { "charts/avg_episodic_return": training_measurements.avg_episodic_return, "charts/avg_episodic_length": np.mean( @@ -394,6 +428,7 @@ class PPOTrainer: "charts/SPS_update": int( self.args.num_envs * self.args.num_steps / (time.time() - iteration_time_start) ), + **storage_metrics } self.logger.log(metrics, step=global_step) @@ -464,6 +499,7 @@ class PPOTrainer: avg_terminated_length=avg_terminated_length, avg_truncated_length=avg_truncated_length, ), + storage ) def _close(self): @@ -507,7 +543,7 @@ class PPOTrainer: for iteration in iter_bar: iteration_time_start = time.time() - env_state, next_obs, next_done, training_measurements = self._step( + env_state, next_obs, next_done, training_measurements, storage = self._step( env_state, next_obs, next_done, iteration=iteration ) @@ -518,6 +554,7 @@ class PPOTrainer: start_time, iteration_time_start, training_measurements, + storage ) sps = int(global_step / (time.time() - start_time)) From 1c328bdde1fc56e87db0ab5547a335dc89a9a927 Mon Sep 17 00:00:00 2001 From: JibrilExe Date: Sun, 12 Apr 2026 11:12:46 +0200 Subject: [PATCH 04/32] fix: observation filtering, action clipping, log_std clip, reward discount, bigger MLP --- scripts/train.py | 2 +- src/brittle_star_project/ppo.py | 1 + .../trainers/PPOTrainer.py | 53 +++++++++++++------ 3 files changed, 39 insertions(+), 17 deletions(-) diff --git a/scripts/train.py b/scripts/train.py index 886bd54..e4b5786 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -50,7 +50,7 @@ if __name__ == "__main__": run_name = f"{args.exp_name}__seed_{args.seed}__{git_hash}__{int(time.time())}" if args.run_dir is None: - run_dir = f"/data/gent/465/vsc46589/runs/{run_name}" + run_dir = f"runs/{run_name}" else: run_dir = args.run_dir diff --git a/src/brittle_star_project/ppo.py b/src/brittle_star_project/ppo.py index cf6c69e..60d9671 100644 --- a/src/brittle_star_project/ppo.py +++ b/src/brittle_star_project/ppo.py @@ -99,6 +99,7 @@ def get_action_and_value( hidden_critic = feature_extractor_apply(params["feature_extractor_params"], x) hidden_sensor = message_passer(hidden_sensor) mean, log_std = actor_apply(params["actor_params"], hidden_sensor) + log_std = jnp.clip(log_std, -5, 2) std = jnp.exp(log_std) logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1) diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 8a9b341..2dc02f9 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -23,6 +23,21 @@ from brittle_star_project.MLPs.mlps import ( from brittle_star_project.ppo import PPO from experiment_logger import get_logger +# TODO: move to config +_ALLOWED_OBS_KEYS = { + "joint_position", + "joint_velocity", + "joint_actuator_force", + "actuator_force", + "disk_position", + "disk_rotation", + "disk_linear_velocity", + "disk_angular_velocity", + "unit_xy_direction_to_target", + "xy_distance_to_target", +} +# TODO: clip scaled reward? + @jax.jit def _clip_action(action: jnp.ndarray, low: jnp.ndarray, high: jnp.ndarray) -> jnp.ndarray: return jnp.clip(action, low, high) @@ -41,9 +56,17 @@ def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, lear @jax.jit def _convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray: - return jax.vmap(lambda o: jnp.concatenate([v.flatten() for v in o.values() if v.size > 0]))( - obs_dict - ) + """Convert the raw observation dict → flat array, filtering unwanted keys.""" + def _filter_and_flatten(o: dict) -> jnp.ndarray: + values = [] + for key in sorted(o.keys()): + if key in _ALLOWED_OBS_KEYS: #TODO: NORMALIZATION or .. of observations?? + v = o[key] + if v.size > 0: + values.append(jnp.asarray(v).flatten()) + return jnp.concatenate(values) + + return jax.vmap(_filter_and_flatten)(obs_dict) def _get_action_and_value_noise( @@ -63,6 +86,7 @@ def _get_action_and_value_noise( ) mean, log_std = actor.apply(agent_state.params["actor_params"], hidden) + log_std = jnp.clip(log_std, -5, 2) key, subkey = jax.random.split(key) noise = jax.random.normal(subkey, shape=mean.shape) std = jnp.exp(log_std) @@ -75,7 +99,7 @@ def _get_action_and_value_noise( logprob = -0.5 * (((raw_action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1) value = critic.apply(agent_state.params["critic_params"], hidden_critic) - return clipped_action, raw_action, logprob, value.squeeze(-1), mean, std, key + return raw_action, raw_action, logprob, value.squeeze(-1), mean, std, key @@ -119,6 +143,7 @@ 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 *= 100 terminated = next_env_state.terminated truncated = next_env_state.truncated done = terminated | truncated @@ -211,7 +236,9 @@ def _compute_gae_jit( (dones[1:], values[1:], values[:-1], storage.rewards), reverse=True, ) - return storage.replace(advantages=advantages, returns=advantages + storage.values) + returns = advantages + storage.values + advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8) + return storage.replace(advantages=advantages, returns=returns) @dataclass @@ -289,8 +316,8 @@ class PPOTrainer: def _init_agent(self): self.logger.info("[AGENT]: Initializing agent...") - sensor = GenericDenseLayersWithActivation() - feature_extractor = GenericDenseLayersWithActivation() + sensor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300]) + feature_extractor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300]) actor = Actor(action_dim=self.env.single_action_space.shape[0]) critic = OneDenseLayerMLP() return sensor, feature_extractor, actor, critic @@ -302,15 +329,9 @@ class PPOTrainer: self.key, 5 ) - sample_obs = jnp.concatenate( - [ - v.flatten() - for v in self.env.single_observation_space.sample( - rng=jax.random.PRNGKey(0) - ).values() - if v.size > 0 - ] - ) + dummy_reset = self.env.reset(seed=0) + sample_obs = _convert_obs_dict_to_array(dummy_reset.observations)[0] # take first env + sensor_params = self.sensor.init(sensor_key, sample_obs) feature_extractor_params = self.feature_extractor.init(feature_extractor_key, sample_obs) actor_params = self.actor.init(actor_key, self.sensor.apply(sensor_params, sample_obs)) From c5a8dd2b3aeabd86cd214b0e086958fb50ea72b0 Mon Sep 17 00:00:00 2001 From: JibrilExe Date: Sun, 12 Apr 2026 12:54:48 +0200 Subject: [PATCH 05/32] feat: expanded used var doc, and added distance to target log --- configs/hpc/debug.yaml | 2 +- .../used_variables.md | 16 +++++++++++++- .../environment/env_config.py | 2 +- .../trainers/PPOTrainer.py | 21 ++++++++++++++++--- 4 files changed, 35 insertions(+), 6 deletions(-) diff --git a/configs/hpc/debug.yaml b/configs/hpc/debug.yaml index c387dc1..e2a9f9d 100644 --- a/configs/hpc/debug.yaml +++ b/configs/hpc/debug.yaml @@ -4,7 +4,7 @@ seed: 42 track: true wandb_project_name: "Let's-find-that-bug" wandb_entity: "SEL3-2026-Groep-4" - +run_dir: "/data/gent/465/vsc46589" num_envs: 32 num_steps: 32 total_timesteps: 102400 diff --git a/experiments/debug-experiment-10042026/used_variables.md b/experiments/debug-experiment-10042026/used_variables.md index 0dcba35..e49c9e4 100644 --- a/experiments/debug-experiment-10042026/used_variables.md +++ b/experiments/debug-experiment-10042026/used_variables.md @@ -1,6 +1,6 @@ ## Default envconfig task: Task = Task.DIRECTED_LOCOMOTION -simulation_time: float = 5.0 +simulation_time: float = 500.0 num_physics_steps_per_control_step: int = 10 time_scale: int = 2 camera_ids: list[int] = field(default_factory=lambda: [0, 1]) @@ -53,6 +53,8 @@ use_torque_control: bool = False ## MLPs: ### Sensor & Feature_extractor: +Both with 3 layers of 300 neurons per layer. + class GenericDenseLayersWithActivation(nn.Module): layer_sizes: Sequence[int] = field(default_factory=lambda: [64, 64]) activation: Callable = nn.tanh @@ -80,3 +82,15 @@ class OneDenseLayerMLP(nn.Module): return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x) ### Observations: +_ALLOWED_OBS_KEYS = { + "joint_position", + "joint_velocity", + "joint_actuator_force", + "actuator_force", + "disk_position", + "disk_rotation", + "disk_linear_velocity", + "disk_angular_velocity", + "unit_xy_direction_to_target", + "xy_distance_to_target", +} \ No newline at end of file diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index 78083e9..c2e8067 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/src/brittle_star_project/environment/env_config.py @@ -31,7 +31,7 @@ class EnvConfig: task: Task = Task.DIRECTED_LOCOMOTION - simulation_time: float = 5.0 + simulation_time: float = 500.0 num_physics_steps_per_control_step: int = 10 time_scale: int = 2 diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 2dc02f9..6997911 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -38,6 +38,12 @@ _ALLOWED_OBS_KEYS = { } # TODO: clip scaled reward? +@jax.jit +def _get_xy_distance_to_target(obs_dict: dict) -> jnp.ndarray: + """Extract xy_distance_to_target for all environments.""" + # obs_dict is a dict of arrays with leading batch dimension (num_envs, ...) + return obs_dict["xy_distance_to_target"].squeeze(-1) # shape: (num_envs,) + @jax.jit def _clip_action(action: jnp.ndarray, low: jnp.ndarray, high: jnp.ndarray) -> jnp.ndarray: return jnp.clip(action, low, high) @@ -143,7 +149,7 @@ 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 *= 100 + reward *= 1000 terminated = next_env_state.terminated truncated = next_env_state.truncated done = terminated | truncated @@ -396,7 +402,9 @@ class PPOTrainer: start_time, iteration_time_start, training_measurements, - storage + storage, + next_obs, + xy_distance ): data = jax.device_get({ 'rewards': storage.rewards[0], # (num_steps,) @@ -425,6 +433,9 @@ class PPOTrainer: "rollout/env0/action_mean": float(np.mean(data['actions'])), "rollout/env0/raw_action_mean": float(np.mean(data['raw_actions'])), + + "charts/env0_xy_distance_to_target": float(xy_distance[0]), + "charts/env1_xy_distance_to_target": float(xy_distance[1]), } metrics = { @@ -568,6 +579,8 @@ class PPOTrainer: env_state, next_obs, next_done, iteration=iteration ) + xy_distance = _get_xy_distance_to_target(env_state.observations) + global_step += self.args.num_steps * self.args.num_envs self._log( global_step, @@ -575,7 +588,9 @@ class PPOTrainer: start_time, iteration_time_start, training_measurements, - storage + storage, + next_obs, + xy_distance ) sps = int(global_step / (time.time() - start_time)) From 631d01ebf660156881f7d158c0e71be122e23f8d Mon Sep 17 00:00:00 2001 From: JibrilExe Date: Sun, 12 Apr 2026 13:09:53 +0200 Subject: [PATCH 06/32] feat: log distance to target for each env --- configs/hpc/debug.yaml | 3 +-- src/brittle_star_project/trainers/PPOTrainer.py | 12 ++++++------ 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/configs/hpc/debug.yaml b/configs/hpc/debug.yaml index e2a9f9d..9fb9170 100644 --- a/configs/hpc/debug.yaml +++ b/configs/hpc/debug.yaml @@ -7,8 +7,7 @@ wandb_entity: "SEL3-2026-Groep-4" run_dir: "/data/gent/465/vsc46589" num_envs: 32 num_steps: 32 -total_timesteps: 102400 +total_timesteps: 409600 num_arms: 2 -num_segments_per_arm: 1 cuda: true diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 6997911..c924519 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -407,10 +407,10 @@ class PPOTrainer: xy_distance ): data = jax.device_get({ - 'rewards': storage.rewards[0], # (num_steps,) + 'rewards': storage.rewards[0], 'values': storage.values[0], - 'returns': storage.returns[0], # NEW - 'advantages': storage.advantages[0],# NEW + 'returns': storage.returns[0], + 'advantages': storage.advantages[0], 'actions': storage.actions[0], 'raw_actions': storage.raw_actions[0], 'means': storage.means[0], @@ -433,11 +433,11 @@ class PPOTrainer: "rollout/env0/action_mean": float(np.mean(data['actions'])), "rollout/env0/raw_action_mean": float(np.mean(data['raw_actions'])), - - "charts/env0_xy_distance_to_target": float(xy_distance[0]), - "charts/env1_xy_distance_to_target": float(xy_distance[1]), } + 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( From 1f8fdbdc592609ac52dc40f499b8878c9298141e Mon Sep 17 00:00:00 2001 From: JibrilExe Date: Tue, 14 Apr 2026 19:07:15 +0200 Subject: [PATCH 07/32] feat: observation normalization --- .../trainers/PPOTrainer.py | 30 +++++++++++++++++-- 1 file changed, 28 insertions(+), 2 deletions(-) diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index c924519..ec19e26 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -59,6 +59,9 @@ def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, lear frac = 1.0 - (count // (minibatch_count * update_epochs)) / num_iterations return learning_rate * frac +@jax.jit +def _normalize_obs(obs, mean, var, eps=1e-8): + return jnp.clip((obs - mean) / jnp.sqrt(var + eps), -10.0, 10.0) @jax.jit def _convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray: @@ -105,7 +108,7 @@ def _get_action_and_value_noise( logprob = -0.5 * (((raw_action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1) value = critic.apply(agent_state.params["critic_params"], hidden_critic) - return raw_action, raw_action, logprob, value.squeeze(-1), mean, std, key + return clipped_action, raw_action, logprob, value.squeeze(-1), mean, std, key @@ -337,7 +340,9 @@ class PPOTrainer: dummy_reset = self.env.reset(seed=0) sample_obs = _convert_obs_dict_to_array(dummy_reset.observations)[0] # take first env - + self.obs_mean = jnp.zeros((len(sample_obs),)) + self.obs_var = jnp.ones((len(sample_obs),)) + self.obs_count = 1e-4 sensor_params = self.sensor.init(sensor_key, sample_obs) feature_extractor_params = self.feature_extractor.init(feature_extractor_key, sample_obs) actor_params = self.actor.init(actor_key, self.sensor.apply(sensor_params, sample_obs)) @@ -376,6 +381,25 @@ class PPOTrainer: returned_episode_returns=jnp.zeros(self.args.num_envs, jnp.float32), returned_episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32), ) + + def _update_obs_stats(self, obs: jnp.ndarray): + batch_mean = jnp.mean(obs, axis=0) + batch_var = jnp.var(obs, axis=0) + batch_count = obs.shape[0] + + delta = batch_mean - self.obs_mean + total_count = self.obs_count + batch_count + + new_mean = self.obs_mean + delta * batch_count / total_count + + m_a = self.obs_var * self.obs_count + m_b = batch_var * batch_count + M2 = m_a + m_b + delta**2 * self.obs_count * batch_count / total_count + new_var = M2 / total_count + + self.obs_mean = new_mean + self.obs_var = new_var + self.obs_count = total_count def _rollout(self, env_state, next_obs, next_done) -> tuple[Any, ...]: return self._rollout_jit( @@ -578,6 +602,8 @@ class PPOTrainer: env_state, next_obs, next_done, training_measurements, storage = self._step( env_state, next_obs, next_done, iteration=iteration ) + 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) From 7a87facc8aaa43936ddeb97b4b1dff9fe45c3841 Mon Sep 17 00:00:00 2001 From: JibrilExe Date: Tue, 14 Apr 2026 19:23:20 +0200 Subject: [PATCH 08/32] feat: value loss clipping --- src/brittle_star_project/ppo.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/src/brittle_star_project/ppo.py b/src/brittle_star_project/ppo.py index 60d9671..ffe2064 100644 --- a/src/brittle_star_project/ppo.py +++ b/src/brittle_star_project/ppo.py @@ -143,7 +143,19 @@ def ppo_loss( pg_loss1 = -mb_advantages * ratio pg_loss2 = -mb_advantages * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef) pg_loss = jnp.maximum(pg_loss1, pg_loss2).mean() - v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean() + old_value = mb_returns - mb_advantages + + v_clipped = old_value + jnp.clip( + newvalue - old_value, + -args.clip_coef, + args.clip_coef + ) + + v_loss_unclipped = (newvalue - mb_returns) ** 2 + v_loss_clipped = (v_clipped - mb_returns) ** 2 + + v_loss = 0.5 * jnp.maximum(v_loss_unclipped, v_loss_clipped).mean() + entropy_loss = entropy.mean() loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl)) From 0a9bf2e0a55f8d02c35e451cc2812a002c5e6d88 Mon Sep 17 00:00:00 2001 From: JibrilExe Date: Tue, 14 Apr 2026 20:30:21 +0200 Subject: [PATCH 09/32] fix: dont store clipped action, and upgraded reward scale --- src/brittle_star_project/ppo.py | 15 ++------------- src/brittle_star_project/trainers/PPOTrainer.py | 5 +++-- 2 files changed, 5 insertions(+), 15 deletions(-) diff --git a/src/brittle_star_project/ppo.py b/src/brittle_star_project/ppo.py index ffe2064..0156b82 100644 --- a/src/brittle_star_project/ppo.py +++ b/src/brittle_star_project/ppo.py @@ -143,19 +143,8 @@ def ppo_loss( pg_loss1 = -mb_advantages * ratio pg_loss2 = -mb_advantages * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef) pg_loss = jnp.maximum(pg_loss1, pg_loss2).mean() - old_value = mb_returns - mb_advantages - - v_clipped = old_value + jnp.clip( - newvalue - old_value, - -args.clip_coef, - args.clip_coef - ) - - v_loss_unclipped = (newvalue - mb_returns) ** 2 - v_loss_clipped = (v_clipped - mb_returns) ** 2 - - v_loss = 0.5 * jnp.maximum(v_loss_unclipped, v_loss_clipped).mean() - + + v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean() entropy_loss = entropy.mean() loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl)) diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index ec19e26..3dcf6db 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -134,7 +134,7 @@ def _step_once( storage = Storage( obs=obs, - actions=clipped_action, + actions=raw_action, raw_actions=raw_action, logprobs=logprob, dones=done, @@ -152,7 +152,8 @@ 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 *= 1000 + reward *= 5000 + reward = jnp.clip(reward, -1, 1) terminated = next_env_state.terminated truncated = next_env_state.truncated done = terminated | truncated From 20ba63a303c771d4bface9b35bb6217f017a719c Mon Sep 17 00:00:00 2001 From: JibrilExe Date: Tue, 14 Apr 2026 21:14:26 +0200 Subject: [PATCH 10/32] fix: ruff format --- configs/hpc/debug.yaml | 7 ++ src/brittle_star_project/MLPs/mlps.py | 6 +- .../environment/env_config.py | 2 +- src/brittle_star_project/ppo.py | 2 +- .../trainers/PPOTrainer.py | 92 +++++++++---------- 5 files changed, 57 insertions(+), 52 deletions(-) diff --git a/configs/hpc/debug.yaml b/configs/hpc/debug.yaml index 9fb9170..a0f5640 100644 --- a/configs/hpc/debug.yaml +++ b/configs/hpc/debug.yaml @@ -7,7 +7,14 @@ wandb_entity: "SEL3-2026-Groep-4" run_dir: "/data/gent/465/vsc46589" num_envs: 32 num_steps: 32 +num_minibatches: 32 total_timesteps: 409600 num_arms: 2 cuda: true +ent_coef: 0.005 +vf_coef: 1.0 +clip_coef: 0.2 + +anneal_lr: true +learning_rate: 0.0003 \ No newline at end of file diff --git a/src/brittle_star_project/MLPs/mlps.py b/src/brittle_star_project/MLPs/mlps.py index 9568a36..5e2deb5 100644 --- a/src/brittle_star_project/MLPs/mlps.py +++ b/src/brittle_star_project/MLPs/mlps.py @@ -58,9 +58,9 @@ class Storage: returns: jnp.array rewards: jnp.array - raw_actions: jnp.ndarray = None # before clipping - means: jnp.ndarray = None # policy mean - stds: jnp.ndarray = None # policy std + raw_actions: jnp.ndarray = None # before clipping + means: jnp.ndarray = None # policy mean + stds: jnp.ndarray = None # policy std def replace(self, **kwargs) -> "Storage": fs = fields(self) diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index c2e8067..65ec229 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/src/brittle_star_project/environment/env_config.py @@ -31,7 +31,7 @@ class EnvConfig: task: Task = Task.DIRECTED_LOCOMOTION - simulation_time: float = 500.0 + simulation_time: float = 10000.0 num_physics_steps_per_control_step: int = 10 time_scale: int = 2 diff --git a/src/brittle_star_project/ppo.py b/src/brittle_star_project/ppo.py index 0156b82..1ea37a8 100644 --- a/src/brittle_star_project/ppo.py +++ b/src/brittle_star_project/ppo.py @@ -143,7 +143,7 @@ def ppo_loss( pg_loss1 = -mb_advantages * ratio pg_loss2 = -mb_advantages * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef) pg_loss = jnp.maximum(pg_loss1, pg_loss2).mean() - + v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean() entropy_loss = entropy.mean() loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 3dcf6db..58a477e 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -38,16 +38,19 @@ _ALLOWED_OBS_KEYS = { } # TODO: clip scaled reward? + @jax.jit def _get_xy_distance_to_target(obs_dict: dict) -> jnp.ndarray: """Extract xy_distance_to_target for all environments.""" # obs_dict is a dict of arrays with leading batch dimension (num_envs, ...) return obs_dict["xy_distance_to_target"].squeeze(-1) # shape: (num_envs,) + @jax.jit def _clip_action(action: jnp.ndarray, low: jnp.ndarray, high: jnp.ndarray) -> jnp.ndarray: return jnp.clip(action, low, high) + def _compute_explained_variance(values: jnp.ndarray, returns: jnp.ndarray) -> float: var_returns = jnp.var(returns) explained_var = 1.0 - jnp.var(returns - values) / (var_returns + 1e-8) @@ -59,22 +62,25 @@ def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, lear frac = 1.0 - (count // (minibatch_count * update_epochs)) / num_iterations return learning_rate * frac + @jax.jit def _normalize_obs(obs, mean, var, eps=1e-8): return jnp.clip((obs - mean) / jnp.sqrt(var + eps), -10.0, 10.0) + @jax.jit def _convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray: """Convert the raw observation dict → flat array, filtering unwanted keys.""" + def _filter_and_flatten(o: dict) -> jnp.ndarray: values = [] for key in sorted(o.keys()): - if key in _ALLOWED_OBS_KEYS: #TODO: NORMALIZATION or .. of observations?? + if key in _ALLOWED_OBS_KEYS: # TODO: NORMALIZATION or .. of observations?? v = o[key] if v.size > 0: values.append(jnp.asarray(v).flatten()) return jnp.concatenate(values) - + return jax.vmap(_filter_and_flatten)(obs_dict) @@ -87,7 +93,7 @@ def _get_action_and_value_noise( next_obs: jnp.ndarray, key: jax.random.PRNGKey, action_low, - action_high + action_high, ): hidden = sensor.apply(agent_state.params["sensor_params"], next_obs) hidden_critic = feature_extractor.apply( @@ -100,16 +106,11 @@ def _get_action_and_value_noise( noise = jax.random.normal(subkey, shape=mean.shape) std = jnp.exp(log_std) raw_action = mean + noise * std - clipped_action = _clip_action( - raw_action, - action_low, - action_high - ) + clipped_action = _clip_action(raw_action, action_low, action_high) logprob = -0.5 * (((raw_action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1) value = critic.apply(agent_state.params["critic_params"], hidden_critic) - - return clipped_action, raw_action, logprob, value.squeeze(-1), mean, std, key + return clipped_action, raw_action, logprob, value.squeeze(-1), mean, std, key def _step_once( @@ -121,7 +122,7 @@ def _step_once( actor: Actor, critic: OneDenseLayerMLP, action_low, - action_high + action_high, ): agent_state, episode_stats, obs, done, key, env_state = carry clipped_action, raw_action, logprob, value, mean, std, key = _get_action_and_value_noise( @@ -152,8 +153,8 @@ 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 *= 5000 - reward = jnp.clip(reward, -1, 1) + reward *= 20000 + reward = jnp.clip(reward, -10, 10) terminated = next_env_state.terminated truncated = next_env_state.truncated done = terminated | truncated @@ -192,7 +193,7 @@ def _rollout_jit( actor: Actor, critic: OneDenseLayerMLP, action_low, - action_high + action_high, ): (agent_state, episode_stats, next_obs, next_done, key, env_state), storage = jax.lax.scan( partial( @@ -203,7 +204,7 @@ def _rollout_jit( critic=critic, env_step_fn=step_env_fn, action_low=action_low, - action_high=action_high + action_high=action_high, ), (agent_state, episode_stats, next_obs, next_done, key, env_state), (), @@ -295,7 +296,7 @@ class PPOTrainer: actor=self.actor, critic=self.critic, action_low=action_low, - action_high=action_high + action_high=action_high, ) ) self._compute_gae_jit = jax.jit( @@ -382,7 +383,7 @@ class PPOTrainer: returned_episode_returns=jnp.zeros(self.args.num_envs, jnp.float32), returned_episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32), ) - + def _update_obs_stats(self, obs: jnp.ndarray): batch_mean = jnp.mean(obs, axis=0) batch_var = jnp.var(obs, axis=0) @@ -429,39 +430,36 @@ class PPOTrainer: training_measurements, storage, next_obs, - xy_distance + 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], - }) + 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], + } + ) 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/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"])), } for i in range(len(xy_distance)): - storage_metrics[f"env_data/env{i}_xy_dist_target"] = float(xy_distance[i]) + storage_metrics[f"env_data/env{i}_xy_dist_target"] = float(xy_distance[i]) metrics = { "charts/avg_episodic_return": training_measurements.avg_episodic_return, @@ -485,7 +483,7 @@ class PPOTrainer: "charts/SPS_update": int( self.args.num_envs * self.args.num_steps / (time.time() - iteration_time_start) ), - **storage_metrics + **storage_metrics, } self.logger.log(metrics, step=global_step) @@ -556,7 +554,7 @@ class PPOTrainer: avg_terminated_length=avg_terminated_length, avg_truncated_length=avg_truncated_length, ), - storage + storage, ) def _close(self): @@ -617,7 +615,7 @@ class PPOTrainer: training_measurements, storage, next_obs, - xy_distance + xy_distance, ) sps = int(global_step / (time.time() - start_time)) From 67f620d5993fdfb26ac902337bda0e95c3eaba13 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 14 Apr 2026 21:16:33 +0200 Subject: [PATCH 11/32] chore: setup hydra dependencies --- .gitignore | 2 ++ env/hpc/modules.txt | 1 + pyproject.toml | 2 +- uv.lock | 72 ++++++++++++++++++++++----------------------- 4 files changed, 39 insertions(+), 38 deletions(-) diff --git a/.gitignore b/.gitignore index b0567ce..43d51e5 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,8 @@ artifacts/* runs/* wandb/ +outputs/ +multirun/ # Python-generated files __pycache__/ diff --git a/env/hpc/modules.txt b/env/hpc/modules.txt index 4c4a1d7..2ea2c6e 100644 --- a/env/hpc/modules.txt +++ b/env/hpc/modules.txt @@ -1,3 +1,4 @@ GCCcore/13.3.0 Python/3.12.3-GCCcore-13.3.0 FFmpeg/7.0.2-GCCcore-13.3.0 +Hydra/1.3.2-GCCcore-13.3.0 diff --git a/pyproject.toml b/pyproject.toml index 5687167..607e704 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ dependencies = [ "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", ] diff --git a/uv.lock b/uv.lock index 01caa5d..bca7772 100644 --- a/uv.lock +++ b/uv.lock @@ -18,6 +18,7 @@ dependencies = [ { name = "evosax" }, { name = "flax" }, { name = "gymnasium" }, + { name = "hydra-core" }, { name = "ipykernel" }, { name = "jax" }, { name = "matplotlib" }, @@ -30,7 +31,6 @@ dependencies = [ { name = "pyopengl-accelerate" }, { name = "pyyaml" }, { name = "torch" }, - { name = "tyro" }, { name = "wandb" }, { name = "warp-lang" }, ] @@ -57,6 +57,7 @@ requires-dist = [ { name = "evosax", specifier = "==0.2.0" }, { name = "flax", specifier = ">=0.12.2" }, { name = "gymnasium", specifier = ">=1.2.3" }, + { name = "hydra-core", specifier = ">=1.3.2" }, { name = "ipykernel", specifier = "==7.2.0" }, { name = "jax", specifier = "==0.9.0.1" }, { name = "jax", extras = ["cuda13"], marker = "extra == 'cuda'", specifier = "==0.9.0.1" }, @@ -71,7 +72,6 @@ requires-dist = [ { name = "pyyaml", specifier = ">=6.0" }, { name = "tensorboard", marker = "extra == 'analysis'" }, { name = "torch", specifier = ">=2.4.0" }, - { name = "tyro", specifier = ">=1.0.10" }, { name = "wandb", specifier = "==0.24.2" }, { name = "warp-lang" }, ] @@ -111,6 +111,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" }, ] +[[package]] +name = "antlr4-python3-runtime" +version = "4.9.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3e/38/7859ff46355f76f8d19459005ca000b6e7012f2f1ca597746cbcd1fbfe5e/antlr4-python3-runtime-4.9.3.tar.gz", hash = "sha256:f224469b4168294902bb1efa80a8bf7855f24c99aef99cbefc1bcd3cce77881b", size = 117034, upload-time = "2021-11-06T17:52:23.524Z" } + [[package]] name = "appnope" version = "0.1.4" @@ -492,15 +498,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ea/97/4f78412f73a9350bc8f934441bae5b68b102c8f4240a7f06b4114b51d6de/dm_tree-0.1.9-cp312-cp312-win_amd64.whl", hash = "sha256:9020a5ce256fcc83aa4bc190cc96dd66e87685db0a6e501b0c06aa492c2e38fc", size = 102022, upload-time = "2025-01-30T20:45:28.701Z" }, ] -[[package]] -name = "docstring-parser" -version = "0.17.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/b2/9d/c3b43da9515bd270df0f80548d9944e389870713cc1fe2b8fb35fe2bcefd/docstring_parser-0.17.0.tar.gz", hash = "sha256:583de4a309722b3315439bb31d64ba3eebada841f2e2cee23b99df001434c912", size = 27442, upload-time = "2025-07-21T07:35:01.868Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/55/e2/2537ebcff11c1ee1ff17d8d0b6f4db75873e3b0fb32c2d4a2ee31ecb310a/docstring_parser-0.17.0-py3-none-any.whl", hash = "sha256:cf2569abd23dce8099b300f9b4fa8191e9582dda731fd533daf54c4551658708", size = 36896, upload-time = "2025-07-21T07:35:00.684Z" }, -] - [[package]] name = "dotmap" version = "1.3.30" @@ -753,6 +750,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c5/7b/bca5613a0c3b542420cf92bd5e5fb8ebd5435ce1011a091f66bb7693285e/humanize-4.15.0-py3-none-any.whl", hash = "sha256:b1186eb9f5a9749cd9cb8565aee77919dd7c8d076161cf44d70e59e3301e1769", size = 132203, upload-time = "2025-12-20T20:16:11.67Z" }, ] +[[package]] +name = "hydra-core" +version = "1.3.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "antlr4-python3-runtime" }, + { name = "omegaconf" }, + { name = "packaging" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6d/8e/07e42bc434a847154083b315779b0a81d567154504624e181caf2c71cd98/hydra-core-1.3.2.tar.gz", hash = "sha256:8a878ed67216997c3e9d88a8e72e7b4767e81af37afb4ea3334b269a4390a824", size = 3263494, upload-time = "2023-02-23T18:33:43.03Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c6/50/e0edd38dcd63fb26a8547f13d28f7a008bc4a3fd4eb4ff030673f22ad41a/hydra_core-1.3.2-py3-none-any.whl", hash = "sha256:fa0238a9e31df3373b35b0bfb672c34cc92718d21f81311d8996a16de1141d8b", size = 154547, upload-time = "2023-02-23T18:33:40.801Z" }, +] + [[package]] name = "identify" version = "2.6.18" @@ -1531,6 +1542,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/11/3f1ee9dce24b41812dd572a037c4436d4d21f759fbe373cc271b0ce98805/nvidia_nvvm-13.2.51-py3-none-win_amd64.whl", hash = "sha256:a4809baaa5429eabe1878853761ce31f0ba15216e2348710b7898dc591f5fc14", size = 56751075, upload-time = "2026-03-09T10:11:09.994Z" }, ] +[[package]] +name = "omegaconf" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "antlr4-python3-runtime" }, + { name = "pyyaml" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/09/48/6388f1bb9da707110532cb70ec4d2822858ddfb44f1cdf1233c20a80ea4b/omegaconf-2.3.0.tar.gz", hash = "sha256:d5d4b6d29955cc50ad50c46dc269bcd92c6e00f5f90d23ab5fee7bfca4ba4cc7", size = 3298120, upload-time = "2022-12-08T20:59:22.753Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e3/94/1843518e420fa3ed6919835845df698c7e27e183cb997394e4a670973a65/omegaconf-2.3.0-py3-none-any.whl", hash = "sha256:7b4df175cdb08ba400f45cae3bdcae7ba8365db4d165fc65fd04b050ab63b46b", size = 79500, upload-time = "2022-12-08T20:59:19.686Z" }, +] + [[package]] name = "opencv-python" version = "4.13.0.92" @@ -2315,18 +2339,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74caf5e34b66d9f3a429af689c1c7128daba1d8208df60e81106b115c00d6fca", size = 188266850, upload-time = "2026-01-20T16:00:43.041Z" }, ] -[[package]] -name = "typeguard" -version = "4.5.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/2b/e8/66e25efcc18542d58706ce4e50415710593721aae26e794ab1dec34fb66f/typeguard-4.5.1.tar.gz", hash = "sha256:f6f8ecbbc819c9bc749983cc67c02391e16a9b43b8b27f15dc70ed7c4a007274", size = 80121, upload-time = "2026-02-19T16:09:03.392Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/91/88/b55b3117287a8540b76dbdd87733808d4d01c8067a3b339408c250bb3600/typeguard-4.5.1-py3-none-any.whl", hash = "sha256:44d2bf329d49a244110a090b55f5f91aa82d9a9834ebfd30bcc73651e4a8cc40", size = 36745, upload-time = "2026-02-19T16:09:01.6Z" }, -] - [[package]] name = "typing-extensions" version = "4.15.0" @@ -2361,20 +2373,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" }, ] -[[package]] -name = "tyro" -version = "1.0.10" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "docstring-parser" }, - { name = "typeguard" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/30/c1/0a5850badd3f18373d6a0366091638674cec6780b558c1c5b846adea938b/tyro-1.0.10.tar.gz", hash = "sha256:2822eacac963a4922bf7eafe3b156a1f0f7fe8e34148202987581224f25565c2", size = 481084, upload-time = "2026-03-18T08:24:17.307Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/05/be/a0b4c9fa64999a2e337cbefcdedd2e101e8dd88a84e4fa497bd0e4531dc1/tyro-1.0.10-py3-none-any.whl", hash = "sha256:8de87a3a40c8a91f10831f8f0638cd0eed00f0e4de9cd3d561e967f407477210", size = 183433, upload-time = "2026-03-18T08:24:16.012Z" }, -] - [[package]] name = "tzdata" version = "2025.3" From 157e061f86a0a1757fea539a7be54382fc75f95b Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 14 Apr 2026 21:53:52 +0200 Subject: [PATCH 12/32] chore: draft new configs --- .../configs/config_experiment.py | 9 +++++++ .../configs/config_networks.py | 23 ++++++++++++++++ .../configs/config_ppo.py | 22 +++++++++++++++ .../configs/main_config.py | 18 +++++++++++++ .../configs/register_configs.py | 27 +++++++++++++++++++ src/experiment_logger/config_logger.py | 14 ++++++++++ 6 files changed, 113 insertions(+) create mode 100644 src/brittle_star_project/configs/config_experiment.py create mode 100644 src/brittle_star_project/configs/config_networks.py create mode 100644 src/brittle_star_project/configs/config_ppo.py create mode 100644 src/brittle_star_project/configs/main_config.py create mode 100644 src/brittle_star_project/configs/register_configs.py create mode 100644 src/experiment_logger/config_logger.py diff --git a/src/brittle_star_project/configs/config_experiment.py b/src/brittle_star_project/configs/config_experiment.py new file mode 100644 index 0000000..69ceca4 --- /dev/null +++ b/src/brittle_star_project/configs/config_experiment.py @@ -0,0 +1,9 @@ +from dataclasses import dataclass + + +@dataclass +class ExperimentConfig: + exp_name: str = "brittle_star_ppo" + seed: int = 1 + torch_deterministic: bool = True + cuda: bool = True diff --git a/src/brittle_star_project/configs/config_networks.py b/src/brittle_star_project/configs/config_networks.py new file mode 100644 index 0000000..347936b --- /dev/null +++ b/src/brittle_star_project/configs/config_networks.py @@ -0,0 +1,23 @@ +from dataclasses import dataclass, field +from typing import List, Optional + + +@dataclass +class LayerConfig: + hidden_dims: List[int] = field(default_factory=lambda: [64, 64]) + activation: str = "relu" + + +@dataclass +class MLPConfig: + actor: LayerConfig = field(default_factory=LayerConfig) + critic: LayerConfig = field(default_factory=LayerConfig) + sensor: LayerConfig = field(default_factory=LayerConfig) + feature_extractor: LayerConfig = field(default_factory=LayerConfig) + + +@dataclass +class NetworksConfig: + mlp: MLPConfig = field(default_factory=MLPConfig) + # Future placeholder for message_passing + # message_passing: Optional[MessagePassingConfig] = None diff --git a/src/brittle_star_project/configs/config_ppo.py b/src/brittle_star_project/configs/config_ppo.py new file mode 100644 index 0000000..d0a29bf --- /dev/null +++ b/src/brittle_star_project/configs/config_ppo.py @@ -0,0 +1,22 @@ +from dataclasses import dataclass +from typing import Optional + + +@dataclass +class PPOConfig: + learning_rate: float = 2.5e-4 + total_timesteps: int = 10000000 + num_envs: int = 100 + num_steps: int = 128 + anneal_lr: bool = True + gamma: float = 0.99 + gae_lambda: float = 0.95 + num_minibatches: int = 4 + update_epochs: int = 4 + norm_adv: bool = True + clip_coef: float = 0.1 + clip_vloss: bool = True + ent_coef: float = 0.01 + vf_coef: float = 0.5 + max_grad_norm: float = 0.5 + target_kl: Optional[float] = None diff --git a/src/brittle_star_project/configs/main_config.py b/src/brittle_star_project/configs/main_config.py new file mode 100644 index 0000000..941a5f2 --- /dev/null +++ b/src/brittle_star_project/configs/main_config.py @@ -0,0 +1,18 @@ +from dataclasses import dataclass, field + +from experiment_logger.config_logger import LoggingConfig +from brittle_star_project.configs.config_experiment import ExperimentConfig +from brittle_star_project.configs.config_ppo import PPOConfig +from brittle_star_project.configs.config_networks import NetworksConfig +from brittle_star_project.environment.env_config import MorphologyConfig, ArenaConfig, EnvConfig + + +@dataclass +class BrittleStarConfig: + experiment: ExperimentConfig = field(default_factory=ExperimentConfig) + logging: LoggingConfig = field(default_factory=LoggingConfig) + ppo: PPOConfig = field(default_factory=PPOConfig) + networks: NetworksConfig = field(default_factory=NetworksConfig) + morphology: MorphologyConfig = field(default_factory=MorphologyConfig) + arena: ArenaConfig = field(default_factory=ArenaConfig) + environment: EnvConfig = field(default_factory=EnvConfig) diff --git a/src/brittle_star_project/configs/register_configs.py b/src/brittle_star_project/configs/register_configs.py new file mode 100644 index 0000000..22d93ca --- /dev/null +++ b/src/brittle_star_project/configs/register_configs.py @@ -0,0 +1,27 @@ +from hydra.core.config_store import ConfigStore + +from experiment_logger.config_logger import LoggingConfig +from brittle_star_project.configs.config_experiment import ExperimentConfig +from brittle_star_project.configs.config_ppo import PPOConfig +from brittle_star_project.configs.config_networks import NetworksConfig, MLPConfig, LayerConfig +from brittle_star_project.environment.env_config import MorphologyConfig, ArenaConfig, EnvConfig +from brittle_star_project.configs.main_config import BrittleStarConfig + + +def register_configs(): + """Register dataclasses with Hydra's ConfigStore.""" + cs = ConfigStore.instance() + + # Store the main config schema + cs.store(name="brittle_star_config", node=BrittleStarConfig) + + # Store individual structured configs for validation + cs.store(group="experiment", name="base_experiment", node=ExperimentConfig) + cs.store(group="logging", name="base_logging", node=LoggingConfig) + cs.store(group="ppo", name="base_ppo", node=PPOConfig) + cs.store(group="networks", name="base_networks", node=NetworksConfig) + + # Environment configs (reusing existing environment config objects) + cs.store(group="morphology", name="base_morphology", node=MorphologyConfig) + cs.store(group="arena", name="base_arena", node=ArenaConfig) + cs.store(group="environment", name="base_environment", node=EnvConfig) diff --git a/src/experiment_logger/config_logger.py b/src/experiment_logger/config_logger.py new file mode 100644 index 0000000..5d39b8e --- /dev/null +++ b/src/experiment_logger/config_logger.py @@ -0,0 +1,14 @@ +from dataclasses import dataclass +from typing import Optional + + +@dataclass +class LoggingConfig: + track: bool = False + wandb_project_name: str = "PPO-Modularity" + wandb_entity: Optional[str] = "SEL3-2026-Groep-4" + capture_video: bool = False + save_model: bool = True + checkpoint_frequency: int = 100 + upload_model: bool = False + hf_entity: str = "" From 82dfcabede87e54c41dcc7f7facaf38149c718c6 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 14 Apr 2026 22:06:15 +0200 Subject: [PATCH 13/32] feat(config): introduce polymorphic architecture configs centralized/decentralized --- configs/architecture/centralized.yaml | 21 +++++ configs/architecture/decentralized.yaml | 31 +++++++ .../configs/config_architecture.py | 84 +++++++++++++++++++ .../configs/config_networks.py | 23 ----- .../configs/main_config.py | 11 ++- .../configs/register_configs.py | 25 ++++-- 6 files changed, 163 insertions(+), 32 deletions(-) create mode 100644 configs/architecture/centralized.yaml create mode 100644 configs/architecture/decentralized.yaml create mode 100644 src/brittle_star_project/configs/config_architecture.py delete mode 100644 src/brittle_star_project/configs/config_networks.py diff --git a/configs/architecture/centralized.yaml b/configs/architecture/centralized.yaml new file mode 100644 index 0000000..62ff16d --- /dev/null +++ b/configs/architecture/centralized.yaml @@ -0,0 +1,21 @@ +# Centralized Actor-Critic Architecture +# Baseline configuration with a single global sensor and motor. + +# Default values are defined in CentralizedConfig dataclass. +# Use this configuration for standard PPO experiments. + +sensor: + hidden_dims: [64, 64] + activation: "tanh" + +motor: + hidden_dims: [] + activation: "tanh" + +feature_extractor: + hidden_dims: [64, 64] + activation: "tanh" + +critic: + hidden_dims: [] + activation: "tanh" diff --git a/configs/architecture/decentralized.yaml b/configs/architecture/decentralized.yaml new file mode 100644 index 0000000..d87e144 --- /dev/null +++ b/configs/architecture/decentralized.yaml @@ -0,0 +1,31 @@ +# Decentralized Actor Architecture (NerveNet-MLP variant) +# Multi-agent/distributed configuration using local sensors, propagators, and motors. + +# Default values are defined in DecentralizedConfig dataclass. +# Use this configuration for decentralized execution experiments. + +sensor: + hidden_dims: [64, 64] + activation: "tanh" + +propagator: + hidden_dims: [64, 64] + activation: "tanh" + +motor: + hidden_dims: [] + activation: "tanh" + +feature_extractor: + hidden_dims: [64, 64] + activation: "tanh" + +critic: + hidden_dims: [] + activation: "tanh" + +# Synchronous message-passing rounds per control step +message_passing_steps: 1 + +# Connectivity topology (e.g., ring, fully_connected) +topology_type: "ring" diff --git a/src/brittle_star_project/configs/config_architecture.py b/src/brittle_star_project/configs/config_architecture.py new file mode 100644 index 0000000..bc140dc --- /dev/null +++ b/src/brittle_star_project/configs/config_architecture.py @@ -0,0 +1,84 @@ +from dataclasses import dataclass, field +from typing import List + + +@dataclass +class LayerConfig: + hidden_dims: List[int] = field(default_factory=lambda: [64, 64]) + activation: str = "tanh" + + +@dataclass +class ArchitectureConfig: + """Base class for actor-critic network configurations. + + Both centralized and decentralized architectures share a centralized critic + composed of a feature extractor followed by a shallow output layer. + + See docs/design/actor-critic.md for the full design rationale. + """ + + # Input network: maps global observation to a hidden representation. + feature_extractor: LayerConfig = field( + default_factory=lambda: LayerConfig(hidden_dims=[64, 64], activation="tanh") + ) + # Output network: maps the hidden representation to a scalar value estimate. + critic: LayerConfig = field( + default_factory=lambda: LayerConfig(hidden_dims=[], activation="tanh") + ) + + +@dataclass +class CentralizedConfig(ArchitectureConfig): + """Centralized actor-critic architecture (baseline). + + The actor is a single global policy composed of a sensor (input network) + and a motor (output network). The sensor receives the full concatenated + global observation; the motor projects the hidden state to all joint actions. + + See docs/design/actor-critic.md for the full design rationale. + """ + + # Input network: maps the global observation to a hidden state. + sensor: LayerConfig = field( + default_factory=lambda: LayerConfig(hidden_dims=[64, 64], activation="tanh") + ) + # Output network: projects hidden state to (mean, log_std) over all joints. + motor: LayerConfig = field( + default_factory=lambda: LayerConfig(hidden_dims=[], activation="tanh") + ) + + +@dataclass +class DecentralizedConfig(ArchitectureConfig): + """Decentralized actor architecture (NerveNet-MLP variant). + + Each node runs a local sensor, exchanges messages with neighbours via a + propagator for a fixed number of steps, and then a local motor produces + the joint offset for that node only. + + The critic remains centralized (shared with the base class): it receives the + full concatenated global observation and outputs a single scalar. + + See docs/design/actor-critic.md and docs/design/communication.md for the + full design rationale. + """ + + # Input network: local observation -> initial hidden state per node. + sensor: LayerConfig = field( + default_factory=lambda: LayerConfig(hidden_dims=[64, 64], activation="tanh") + ) + # Message-passing network: aggregates neighbour messages and updates hidden state. + propagator: LayerConfig = field( + default_factory=lambda: LayerConfig(hidden_dims=[64, 64], activation="tanh") + ) + # Output network: final hidden state -> joint offset for this node only. + motor: LayerConfig = field( + default_factory=lambda: LayerConfig(hidden_dims=[], activation="tanh") + ) + + # Number of synchronous message-passing rounds per control step. + message_passing_steps: int = 1 + # Graph topology used for neighbour connections. + # Supported values: "ring", "fully_connected". + topology_type: str = "ring" diff --git a/src/brittle_star_project/configs/config_networks.py b/src/brittle_star_project/configs/config_networks.py deleted file mode 100644 index 347936b..0000000 --- a/src/brittle_star_project/configs/config_networks.py +++ /dev/null @@ -1,23 +0,0 @@ -from dataclasses import dataclass, field -from typing import List, Optional - - -@dataclass -class LayerConfig: - hidden_dims: List[int] = field(default_factory=lambda: [64, 64]) - activation: str = "relu" - - -@dataclass -class MLPConfig: - actor: LayerConfig = field(default_factory=LayerConfig) - critic: LayerConfig = field(default_factory=LayerConfig) - sensor: LayerConfig = field(default_factory=LayerConfig) - feature_extractor: LayerConfig = field(default_factory=LayerConfig) - - -@dataclass -class NetworksConfig: - mlp: MLPConfig = field(default_factory=MLPConfig) - # Future placeholder for message_passing - # message_passing: Optional[MessagePassingConfig] = None diff --git a/src/brittle_star_project/configs/main_config.py b/src/brittle_star_project/configs/main_config.py index 941a5f2..aab864b 100644 --- a/src/brittle_star_project/configs/main_config.py +++ b/src/brittle_star_project/configs/main_config.py @@ -3,16 +3,23 @@ from dataclasses import dataclass, field from experiment_logger.config_logger import LoggingConfig from brittle_star_project.configs.config_experiment import ExperimentConfig from brittle_star_project.configs.config_ppo import PPOConfig -from brittle_star_project.configs.config_networks import NetworksConfig +from brittle_star_project.configs.config_architecture import ArchitectureConfig, CentralizedConfig from brittle_star_project.environment.env_config import MorphologyConfig, ArenaConfig, EnvConfig @dataclass class BrittleStarConfig: + """Root configuration for a brittle star training run. + + Composed of strictly separated sub-configs. Each sub-config can be swapped + independently via CLI or a different YAML file. See configs/README.md. + """ + experiment: ExperimentConfig = field(default_factory=ExperimentConfig) logging: LoggingConfig = field(default_factory=LoggingConfig) ppo: PPOConfig = field(default_factory=PPOConfig) - networks: NetworksConfig = field(default_factory=NetworksConfig) + # Default to centralized; swap with architecture=decentralized on the CLI. + architecture: ArchitectureConfig = field(default_factory=CentralizedConfig) morphology: MorphologyConfig = field(default_factory=MorphologyConfig) arena: ArenaConfig = field(default_factory=ArenaConfig) environment: EnvConfig = field(default_factory=EnvConfig) diff --git a/src/brittle_star_project/configs/register_configs.py b/src/brittle_star_project/configs/register_configs.py index 22d93ca..697eda9 100644 --- a/src/brittle_star_project/configs/register_configs.py +++ b/src/brittle_star_project/configs/register_configs.py @@ -3,25 +3,36 @@ from hydra.core.config_store import ConfigStore from experiment_logger.config_logger import LoggingConfig from brittle_star_project.configs.config_experiment import ExperimentConfig from brittle_star_project.configs.config_ppo import PPOConfig -from brittle_star_project.configs.config_networks import NetworksConfig, MLPConfig, LayerConfig +from brittle_star_project.configs.config_architecture import ( + CentralizedConfig, + DecentralizedConfig, +) from brittle_star_project.environment.env_config import MorphologyConfig, ArenaConfig, EnvConfig from brittle_star_project.configs.main_config import BrittleStarConfig -def register_configs(): - """Register dataclasses with Hydra's ConfigStore.""" +def register_configs() -> None: + """Register all dataclasses with Hydra's ConfigStore. + + This must be called before hydra.main() processes the config, ensuring + every structured config is validated against its Python schema. Typos in + YAML keys will raise ConfigAttributeError at startup. + """ cs = ConfigStore.instance() - # Store the main config schema + # Root schema cs.store(name="brittle_star_config", node=BrittleStarConfig) - # Store individual structured configs for validation + # Sub-config groups — each group corresponds to a configs/ subdirectory. cs.store(group="experiment", name="base_experiment", node=ExperimentConfig) cs.store(group="logging", name="base_logging", node=LoggingConfig) cs.store(group="ppo", name="base_ppo", node=PPOConfig) - cs.store(group="networks", name="base_networks", node=NetworksConfig) - # Environment configs (reusing existing environment config objects) + # Architecture variants — swap via CLI: architecture=decentralized + cs.store(group="architecture", name="centralized", node=CentralizedConfig) + cs.store(group="architecture", name="decentralized", node=DecentralizedConfig) + + # Environment configs cs.store(group="morphology", name="base_morphology", node=MorphologyConfig) cs.store(group="arena", name="base_arena", node=ArenaConfig) cs.store(group="environment", name="base_environment", node=EnvConfig) From ba32fc90767d73ed166b639ba25a7b9796d212be Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 14 Apr 2026 22:22:17 +0200 Subject: [PATCH 14/32] feat(morphology): support per-arm segment counts for partial amputations --- configs/example.yaml | 3 +-- .../environment/env_config.py | 17 +++++++++++++++-- src/brittle_star_project/environment/factory.py | 2 +- 3 files changed, 17 insertions(+), 5 deletions(-) diff --git a/configs/example.yaml b/configs/example.yaml index 14bf3d7..baeb1b0 100644 --- a/configs/example.yaml +++ b/configs/example.yaml @@ -1,5 +1,4 @@ morphology: - num_arms: 2 - num_segments_per_arm: 4 + segments_per_arm: [4, 4] use_p_control: true use_torque_control: false diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index 78083e9..206e902 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/src/brittle_star_project/environment/env_config.py @@ -7,11 +7,24 @@ from .env_types import Task @dataclass(frozen=True, slots=True) class MorphologyConfig: - num_arms: int = 5 - num_segments_per_arm: int = 4 + """Brittle star morphology configuration. + + segments_per_arm defines the number of segments for each arm. The length of + this list implicitly sets the number of arms. Use 0 segments to represent + a fully amputated arm (e.g., [4, 0, 4, 2, 4] for a 5-arm morphology with + arm 1 removed and arm 3 shortened). + + The upstream biorobot library natively supports per-arm segment counts. + """ + + segments_per_arm: tuple[int, ...] = (4, 4, 4, 4, 4) use_p_control: bool = True use_torque_control: bool = False + @property + def num_arms(self) -> int: + return len(self.segments_per_arm) + @dataclass(frozen=True, slots=True) class ArenaConfig: diff --git a/src/brittle_star_project/environment/factory.py b/src/brittle_star_project/environment/factory.py index 1a891ea..523feb2 100644 --- a/src/brittle_star_project/environment/factory.py +++ b/src/brittle_star_project/environment/factory.py @@ -22,7 +22,7 @@ class BrittleStarEnvFactory: spec = default_brittle_star_morphology_specification( num_arms=config.num_arms, - num_segments_per_arm=config.num_segments_per_arm, + num_segments_per_arm=list(config.segments_per_arm), use_p_control=config.use_p_control, use_torque_control=config.use_torque_control, ) From 3b8c347715b6958f3f610de9874cda74caeffadd Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 14 Apr 2026 22:36:56 +0200 Subject: [PATCH 15/32] feat(env): add observation padding wrapper for amputated morphologies --- .../environment/padded_obs_wrapper.py | 91 +++++++++++++++++++ 1 file changed, 91 insertions(+) create mode 100644 src/brittle_star_project/environment/padded_obs_wrapper.py diff --git a/src/brittle_star_project/environment/padded_obs_wrapper.py b/src/brittle_star_project/environment/padded_obs_wrapper.py new file mode 100644 index 0000000..8c07248 --- /dev/null +++ b/src/brittle_star_project/environment/padded_obs_wrapper.py @@ -0,0 +1,91 @@ +"""Observation padding wrapper for amputated brittle star morphologies. + +When using a centralized controller, the global observation vector must remain +a constant size regardless of how many segments are amputated. This wrapper pads +the observation dictionary values with zeros using spatial insertion so that the +flattened observation maintains the correct physical mapping to the neural network. +""" +from __future__ import annotations + +from typing import Any +import jax.numpy as jnp + +# Observation keys whose size scales with the number of joints (2 per segment). +_JOINT_SCALED_KEYS = frozenset({ + "joint_position", + "joint_velocity", + "joint_actuator_force", + "actuator_force", +}) + +# Observation keys whose size scales with the number of segments (1 per segment). +_SEGMENT_SCALED_KEYS = frozenset({ + "segment_contact", +}) + + +def compute_padding_masks( + segments_per_arm: tuple[int, ...], + reference_segments_per_arm: tuple[int, ...] = (4, 4, 4, 4, 4), +) -> dict[str, Any]: + """Pre-compute boolean masks for spatial insertion of observations. + + Args: + segments_per_arm: The current (possibly amputated) morphology. + reference_segments_per_arm: The full morphology that defines the expected size. + + Returns: + A dict containing 1D boolean masks and target sizes. + """ + mask_1x = [] + mask_2x = [] + + for actual, ref in zip(segments_per_arm, reference_segments_per_arm): + # 1x scaling (e.g., contacts: 1 value per segment) + mask_1x.extend([True] * actual + [False] * (ref - actual)) + # 2x scaling (e.g., joints: 2 values per segment) + mask_2x.extend([True] * (actual * 2) + [False] * ((ref - actual) * 2)) + + return { + "mask_1x": jnp.array(mask_1x, dtype=bool), + "mask_2x": jnp.array(mask_2x, dtype=bool), + "target_size_1x": sum(reference_segments_per_arm), + "target_size_2x": sum(reference_segments_per_arm) * 2, + } + + +def pad_observation( + obs: dict[str, Any], + masks: dict[str, Any], +) -> dict[str, Any]: + """Pad an observation dict using spatial insertion.""" + padded = {} + for key, value in obs.items(): + if key in _JOINT_SCALED_KEYS: + out = jnp.zeros(masks["target_size_2x"], dtype=value.dtype) + padded[key] = out.at[masks["mask_2x"]].set(value) + elif key in _SEGMENT_SCALED_KEYS: + out = jnp.zeros(masks["target_size_1x"], dtype=value.dtype) + padded[key] = out.at[masks["mask_1x"]].set(value) + else: + padded[key] = value + return padded + + +def pad_observations_batched( + obs: dict[str, Any], + masks: dict[str, Any], +) -> dict[str, Any]: + """Pad a batched observation dict (leading batch dimension) using spatial insertion.""" + padded = {} + for key, value in obs.items(): + batch_size = value.shape[0] + if key in _JOINT_SCALED_KEYS: + out = jnp.zeros((batch_size, masks["target_size_2x"]), dtype=value.dtype) + padded[key] = out.at[:, masks["mask_2x"]].set(value) + elif key in _SEGMENT_SCALED_KEYS: + out = jnp.zeros((batch_size, masks["target_size_1x"]), dtype=value.dtype) + padded[key] = out.at[:, masks["mask_1x"]].set(value) + else: + padded[key] = value + return padded From 2a7537f51cbec072f3e96265da08578669ac2a21 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 14 Apr 2026 22:43:28 +0200 Subject: [PATCH 16/32] feat: add modular configuration files --- configs/arena/default.yaml | 8 ++++++++ configs/environment/directed_locomotion.yaml | 12 ++++++++++++ configs/environment/light_escape.yaml | 12 ++++++++++++ configs/experiment/base.yaml | 7 +++++++ configs/experiment/dev_test.yaml | 7 +++++++ configs/experiment/hpc_smoke_test.yaml | 7 +++++++ configs/logging/default.yaml | 11 +++++++++++ configs/logging/wandb_enabled.yaml | 11 +++++++++++ configs/morphology/3_arms.yaml | 6 ++++++ configs/morphology/5_arms_full.yaml | 6 ++++++ configs/morphology/partial_amputation.yaml | 6 ++++++ configs/ppo/default.yaml | 19 +++++++++++++++++++ configs/ppo/fast.yaml | 19 +++++++++++++++++++ configs/ppo/stable.yaml | 19 +++++++++++++++++++ 14 files changed, 150 insertions(+) create mode 100644 configs/arena/default.yaml create mode 100644 configs/environment/directed_locomotion.yaml create mode 100644 configs/environment/light_escape.yaml create mode 100644 configs/experiment/base.yaml create mode 100644 configs/experiment/dev_test.yaml create mode 100644 configs/experiment/hpc_smoke_test.yaml create mode 100644 configs/logging/default.yaml create mode 100644 configs/logging/wandb_enabled.yaml create mode 100644 configs/morphology/3_arms.yaml create mode 100644 configs/morphology/5_arms_full.yaml create mode 100644 configs/morphology/partial_amputation.yaml create mode 100644 configs/ppo/default.yaml create mode 100644 configs/ppo/fast.yaml create mode 100644 configs/ppo/stable.yaml diff --git a/configs/arena/default.yaml b/configs/arena/default.yaml new file mode 100644 index 0000000..c42ec3c --- /dev/null +++ b/configs/arena/default.yaml @@ -0,0 +1,8 @@ +# Default Arena Configuration +# Base aquarium environment settings. + +size: [10.0, 5.0] +sand_ground_color: true +attach_target: true +wall_height: 1.5 +wall_thickness: 0.1 diff --git a/configs/environment/directed_locomotion.yaml b/configs/environment/directed_locomotion.yaml new file mode 100644 index 0000000..63ad520 --- /dev/null +++ b/configs/environment/directed_locomotion.yaml @@ -0,0 +1,12 @@ +# Directed Locomotion Environment +# Baseline task setting. + +task: DIRECTED_LOCOMOTION +simulation_time: 5.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 +light_perlin_noise_scale: 0 diff --git a/configs/environment/light_escape.yaml b/configs/environment/light_escape.yaml new file mode 100644 index 0000000..c79cfb9 --- /dev/null +++ b/configs/environment/light_escape.yaml @@ -0,0 +1,12 @@ +# Light Escape Environment +# Advanced task requiring movement away from light source. + +task: LIGHT_ESCAPE +simulation_time: 5.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 +light_perlin_noise_scale: 200 # Must be integer factor of 200 diff --git a/configs/experiment/base.yaml b/configs/experiment/base.yaml new file mode 100644 index 0000000..9f31d27 --- /dev/null +++ b/configs/experiment/base.yaml @@ -0,0 +1,7 @@ +# Base Experiment Configuration +# Default values align with ExperimentConfig dataclass. + +exp_name: "brittle_star_ppo" +seed: 1 +torch_deterministic: true +cuda: true diff --git a/configs/experiment/dev_test.yaml b/configs/experiment/dev_test.yaml new file mode 100644 index 0000000..941a1e9 --- /dev/null +++ b/configs/experiment/dev_test.yaml @@ -0,0 +1,7 @@ +# Testing Experiment Configuration +# Quick experiment for local development/testing. + +exp_name: "dev_test_brittle_star" +seed: 42 +torch_deterministic: true +cuda: true diff --git a/configs/experiment/hpc_smoke_test.yaml b/configs/experiment/hpc_smoke_test.yaml new file mode 100644 index 0000000..4264801 --- /dev/null +++ b/configs/experiment/hpc_smoke_test.yaml @@ -0,0 +1,7 @@ +# HPC Smoke Test Configuration +# Uses minimal settings but simulates HPC environment. + +exp_name: "hpc_smoke_test" +seed: 123 +torch_deterministic: true +cuda: true diff --git a/configs/logging/default.yaml b/configs/logging/default.yaml new file mode 100644 index 0000000..4a0cb92 --- /dev/null +++ b/configs/logging/default.yaml @@ -0,0 +1,11 @@ +# Default Logging Configuration +# Offline local-only setup (WandB disabled). + +track: false +wandb_project_name: "PPO-Modularity" +wandb_entity: "SEL3-2026-Groep-4" +capture_video: false +save_model: true +checkpoint_frequency: 100 +upload_model: false +hf_entity: "" diff --git a/configs/logging/wandb_enabled.yaml b/configs/logging/wandb_enabled.yaml new file mode 100644 index 0000000..2e82781 --- /dev/null +++ b/configs/logging/wandb_enabled.yaml @@ -0,0 +1,11 @@ +# WandB Enabled Logging Configuration +# For production/cloud experiments with weights synced. + +track: true +wandb_project_name: "PPO-Modularity" +wandb_entity: "SEL3-2026-Groep-4" +capture_video: false +save_model: true +checkpoint_frequency: 100 +upload_model: false +hf_entity: "" diff --git a/configs/morphology/3_arms.yaml b/configs/morphology/3_arms.yaml new file mode 100644 index 0000000..ebc1665 --- /dev/null +++ b/configs/morphology/3_arms.yaml @@ -0,0 +1,6 @@ +# 3 Arms Morphology Configuration +# Symmetric amputation (arms 1 and 3 removed). + +segments_per_arm: [4, 0, 4, 0, 4] +use_p_control: true +use_torque_control: false diff --git a/configs/morphology/5_arms_full.yaml b/configs/morphology/5_arms_full.yaml new file mode 100644 index 0000000..408b7c9 --- /dev/null +++ b/configs/morphology/5_arms_full.yaml @@ -0,0 +1,6 @@ +# 5 Arms Full Morphology Configuration +# Baseline 5-arm brittle star. + +segments_per_arm: [4, 4, 4, 4, 4] +use_p_control: true +use_torque_control: false diff --git a/configs/morphology/partial_amputation.yaml b/configs/morphology/partial_amputation.yaml new file mode 100644 index 0000000..8003035 --- /dev/null +++ b/configs/morphology/partial_amputation.yaml @@ -0,0 +1,6 @@ +# Partial Amputation Configuration +# Random partial amputation for robustness testing. + +segments_per_arm: [4, 2, 4, 4, 4] +use_p_control: true +use_torque_control: false diff --git a/configs/ppo/default.yaml b/configs/ppo/default.yaml new file mode 100644 index 0000000..50107b4 --- /dev/null +++ b/configs/ppo/default.yaml @@ -0,0 +1,19 @@ +# Default PPO Configuration +# Standard hyperparams from original codebase. + +learning_rate: 0.00025 +total_timesteps: 10000000 +num_envs: 100 +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.1 +clip_vloss: true +ent_coef: 0.01 +vf_coef: 0.5 +max_grad_norm: 0.5 +target_kl: null diff --git a/configs/ppo/fast.yaml b/configs/ppo/fast.yaml new file mode 100644 index 0000000..5001ef2 --- /dev/null +++ b/configs/ppo/fast.yaml @@ -0,0 +1,19 @@ +# 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 diff --git a/configs/ppo/stable.yaml b/configs/ppo/stable.yaml new file mode 100644 index 0000000..35cade4 --- /dev/null +++ b/configs/ppo/stable.yaml @@ -0,0 +1,19 @@ +# Stable PPO Configuration +# Standard hyperparams with lower LR and larger batch. + +learning_rate: 0.0001 +total_timesteps: 10000000 +num_envs: 100 +num_steps: 256 +anneal_lr: true +gamma: 0.99 +gae_lambda: 0.95 +num_minibatches: 8 +update_epochs: 4 +norm_adv: true +clip_coef: 0.1 +clip_vloss: true +ent_coef: 0.01 +vf_coef: 0.5 +max_grad_norm: 0.5 +target_kl: null From 71aedc4853d666b944f54f2d706af6e09a13a73d Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 14 Apr 2026 22:47:09 +0200 Subject: [PATCH 17/32] docs: add main configuration entry point and usage guidelines --- configs/README.md | 53 +++++++++++++++++++++++++++++++--------- configs/main_config.yaml | 13 ++++++++++ 2 files changed, 54 insertions(+), 12 deletions(-) create mode 100644 configs/main_config.yaml diff --git a/configs/README.md b/configs/README.md index 58c6aed..21f95b1 100644 --- a/configs/README.md +++ b/configs/README.md @@ -1,23 +1,52 @@ -# Configuration Files +# Brittle Star Configuration System -This directory contains configuration files for training experiments. +This project uses **Hydra** for a modular, hierarchical, and strictly-typed configuration system. -## Usage +## Core Concepts -Use `--config` with `scripts/train.py` to run an experiment: +1. **Composition over Inheritance**: Instead of one giant config file, the configuration is composed of small, domain-specific modules (PPO settings, architecture, morphology, etc.). +2. **Strict Typing**: Every configuration is validated against a Python dataclass schema (`ConfigStore`). Misspelled keys throw a `ConfigAttributeError` immediately. +3. **CLI Swapping**: You can swap entire modules or override individual values from the command line without touching code. +## Directory Structure + +- `main_config.yaml`: The root entry point defining the default composition. +- `experiment/`: High-level experiment settings (seed, device). +- `logging/`: WandB and checkpointing configuration. +- `ppo/`: PPO training hyperparameters. +- `architecture/`: Polymorphic network architectures (centralized vs. decentralized). +- `morphology/`: Physical robot definitions (number of segments, amputations). +- `arena/`: Environment physics and visual settings. +- `environment/`: Task-specific settings (Directed Locomotion, Light Escape). + +## Common Commands + +### Local Debugging +Run a quick test with minimal iterations: ```bash -python scripts/train.py --config configs/default_ppo.yaml +python scripts/train.py experiment=dev_test ppo=fast ``` -You can overriding settings via CLI: +### Swapping Architectures or Morphologies +Test a decentralized controller on a 3-arm robot: ```bash -python scripts/train.py --config configs/default_ppo.yaml --learning-rate 0.001 +python scripts/train.py architecture=decentralized morphology=3_arms ``` -## Available Configurations +### HPC Production +Run stable PPO with WandB enabled (HPC submission scripts handle the `hydra.run.dir` redirection): +```bash +python scripts/train.py ppo=stable logging=wandb_enabled +``` -- `default_ppo.yaml`: Baseline config. -- `dev_test.yaml`: Fast iteration for development. -- `production_training.yaml`: Full-scale training. -- `personal_template.yaml`: Template for team members to customize. +### Dry-Run Validation +Check if your configuration is valid without starting the simulation: +```bash +python scripts/train.py --cfg job +``` + +## Developer Notes + +- **Adding a new group**: Create a subdirectory in `configs/` and register the new dataclass in `src/brittle_star_project/configs/register_configs.py`. +- **Typo Catching**: If you see a `ConfigAttributeError`, check for typos in your YAML keys or CLI overrides. +- **Output Redirection**: Hydra automatically creates `outputs/` directories. On HPC, ensure `hydra.run.dir` is set to a fast scratch storage. diff --git a/configs/main_config.yaml b/configs/main_config.yaml new file mode 100644 index 0000000..0583327 --- /dev/null +++ b/configs/main_config.yaml @@ -0,0 +1,13 @@ +# Brittle Star Project - Main Configuration +# This file defines the default composition of the hierarchical configuration. +# Sub-configs are loaded from the relative directories. + +defaults: + - experiment: base + - logging: default + - ppo: default + - architecture: centralized + - morphology: 5_arms_full + - arena: default + - environment: directed_locomotion + - _self_ From eabc64009abffadaa472aab7553d45d37a2f8bc4 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 14 Apr 2026 22:53:17 +0200 Subject: [PATCH 18/32] feat(log): implement LoggerProxy to defer initialization and prevent premature directory creation --- src/experiment_logger/__init__.py | 3 +- src/experiment_logger/config_utils.py | 24 +++---- src/experiment_logger/unified_logger.py | 90 +++++++++++++++---------- 3 files changed, 69 insertions(+), 48 deletions(-) diff --git a/src/experiment_logger/__init__.py b/src/experiment_logger/__init__.py index 64d4be4..c89af7f 100644 --- a/src/experiment_logger/__init__.py +++ b/src/experiment_logger/__init__.py @@ -5,7 +5,7 @@ This package provides a unified interface for logging to multiple backends """ from experiment_logger.config_utils import load_yaml_config, merge_config_with_cli -from experiment_logger.unified_logger import UnifiedLogger, get_logger +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 @@ -13,6 +13,7 @@ __all__ = [ "UnifiedLogger", "SimpleLogger", "get_logger", + "init_logger", "init_wandb", "finish_wandb", "load_yaml_config", diff --git a/src/experiment_logger/config_utils.py b/src/experiment_logger/config_utils.py index 4c5b79f..09cb4ce 100644 --- a/src/experiment_logger/config_utils.py +++ b/src/experiment_logger/config_utils.py @@ -8,8 +8,6 @@ from dataclasses import fields, is_dataclass from experiment_logger.unified_logger import get_logger -log = get_logger() - T = TypeVar("T") @@ -24,7 +22,7 @@ def load_yaml_config(config_path: str) -> Dict[str, Any]: if config is None: return {} - log.info(f"Loaded configuration from: {config_path}") + get_logger().info(f"Loaded configuration from: {config_path}") return config @@ -35,7 +33,7 @@ def save_yaml_config(config: Dict[str, Any], config_path: str): with open(config_path, "w") as f: yaml.dump(config, f, default_flow_style=False, indent=2, sort_keys=False) - log.info(f"Saved configuration to: {config_path}") + get_logger().info(f"Saved configuration to: {config_path}") def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T: @@ -66,10 +64,10 @@ def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T: else: filtered_config[key] = field.type(value) if value is not None else None # type: ignore except (ValueError, TypeError) as e: - log.warning(f"Could not convert {key}={value} to {field.type}: {e}") + get_logger().warning(f"Could not convert {key}={value} to {field.type}: {e}") filtered_config[key] = value else: - log.warning(f"Unknown configuration parameter: {key}") + get_logger().warning(f"Unknown configuration parameter: {key}") return cls(**filtered_config) @@ -101,9 +99,9 @@ def merge_config_with_cli(config_class: Type[T], config_file: str | None = None) yaml_config = {} if extracted_config_file and os.path.exists(extracted_config_file): yaml_config = load_yaml_config(extracted_config_file) - log.info(f"Merging YAML config from {extracted_config_file} with CLI args") + get_logger().info(f"Merging YAML config from {extracted_config_file} with CLI args") elif extracted_config_file: - log.warning(f"Config file not found: {extracted_config_file}, using CLI args only") + get_logger().warning(f"Config file not found: {extracted_config_file}, using CLI args only") # Create default instance to know what the defaults are default_instance = config_class() @@ -126,22 +124,22 @@ def merge_config_with_cli(config_class: Type[T], config_file: str | None = None) if cli_value != default_value: final_config[field_name] = cli_value if yaml_value != default_value and yaml_value != cli_value: - log.info(f"CLI override: {field_name}={cli_value} (YAML had {yaml_value})") + get_logger().info(f"CLI override: {field_name}={cli_value} (YAML had {yaml_value})") else: final_config[field_name] = yaml_value if yaml_value != default_value: - log.info(f"YAML config: {field_name}={yaml_value}") + get_logger().info(f"YAML config: {field_name}={yaml_value}") return config_class(**final_config) def print_config(config: Any, title: str = "Configuration"): """Pretty print configuration.""" - log.info(f"{title}:") + get_logger().info(f"{title}:") if is_dataclass(config): for field in fields(config): value = getattr(config, field.name) - log.info(f" {field.name}: {value}") + get_logger().info(f" {field.name}: {value}") else: for key, value in vars(config).items(): - log.info(f" {key}: {value}") + get_logger().info(f" {key}: {value}") diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py index 5cfb3f1..4eadd5b 100644 --- a/src/experiment_logger/unified_logger.py +++ b/src/experiment_logger/unified_logger.py @@ -21,38 +21,67 @@ import numpy as np from experiment_logger.wandb_utils import finish_wandb, init_wandb -# Global singleton storage -_global_logger = None +# Global storage for the active logger and the proxy singleton +_active_logger: Optional[Any] = None +_proxy_instance: Optional["LoggerProxy"] = None -def get_logger() -> "UnifiedLogger": - """Retrieve the global UnifiedLogger. If not initialized, fallback to auto-initialization.""" - global _global_logger - if _global_logger is None: - try: - commit_hash = ( - subprocess.check_output( - ["git", "rev-parse", "--short", "HEAD"], stderr=subprocess.STDOUT - ) - .decode("utf-8") - .strip() - ) - except Exception: - commit_hash = "unknown" +def get_logger() -> "LoggerProxy": + """Retrieve the global LoggerProxy. - timestamp_str = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") - generic_name = f"{timestamp_str}_{commit_hash}_brittle_star" + This should be used for all logging calls. It returns a proxy that + delegates to the active logger (defaulting to a SimpleLogger until + init_logger is called). + """ + global _proxy_instance, _active_logger + if _proxy_instance is None: + if _active_logger is None: + # Fallback to SimpleLogger to avoid premature directory creation + from experiment_logger.simple_logger import SimpleLogger - # Initialize generic fallback logger without WandB - _global_logger = UnifiedLogger( - run_name=generic_name, - config={"auto_initialized": True}, - use_wandb=False, - _set_as_global=False, # Prevent recursive call inside __init__ - ) - _global_logger.warning(f"UnifiedLogger auto-initialized with name: {generic_name}") + _active_logger = SimpleLogger(run_name="pre_init") - return _global_logger + _proxy_instance = LoggerProxy() + + return _proxy_instance + + +def init_logger(**kwargs) -> "UnifiedLogger": + """Initialize the full UnifiedLogger and set it as the active logger. + + This should be called once the configuration is ready. It will create + the output directories and set up all logging backends. + """ + global _active_logger + logger = UnifiedLogger(_set_as_global=False, **kwargs) + _active_logger = logger + return logger + + +class LoggerProxy: + """Proxy that delegates all method calls to the active logger instance. + + This allows the logger to be swapped out (e.g., from a SimpleLogger to + a UnifiedLogger) without any clients needing to update their references. + """ + + def _get_logger(self) -> Any: + global _active_logger + if _active_logger is None: + # This shouldn't normally happen since get_logger handles it + from experiment_logger.simple_logger import SimpleLogger + + _active_logger = SimpleLogger(run_name="pre_init_fallback") + return _active_logger + + def __getattr__(self, name: str) -> Any: + return getattr(self._get_logger(), name) + + def __enter__(self): + return self._get_logger().__enter__() + + def __exit__(self, exc_type, exc_val, exc_tb): + return self._get_logger().__exit__(exc_type, exc_val, exc_tb) class UnifiedLogger: @@ -68,7 +97,6 @@ class UnifiedLogger: use_wandb: bool = True, save_code: bool = True, log_level: int = logging.INFO, - _set_as_global: bool = True, ): """Initialize the unified logger. @@ -80,7 +108,6 @@ class UnifiedLogger: base_dir: Base directory for local storage use_wandb: Whether to use WandB logging save_code: Whether to save code to WandB - _set_as_global: Internal flag to override the global singleton """ self.run_name = run_name self.config = config @@ -119,11 +146,6 @@ class UnifiedLogger: self._text_logger.addHandler(fh) self._text_logger.addHandler(ch) - # Set as global singleton - global _global_logger - if _set_as_global: - _global_logger = self - # Save config to disk self._save_config() From ff83af8cef2cffabb2fdb07260136b7baa38697e Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 14 Apr 2026 23:03:43 +0200 Subject: [PATCH 19/32] migrate main training entry point to Hydra and wire structured configs --- scripts/train.py | 136 ++++++++++++++++++++++++++++++----------------- 1 file changed, 86 insertions(+), 50 deletions(-) diff --git a/scripts/train.py b/scripts/train.py index 492c887..5ef2258 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -1,75 +1,111 @@ -import subprocess -import time - -import torch import os +import time +import torch +import hydra +from omegaconf import DictConfig, OmegaConf +from brittle_star_project.configs.main_config import BrittleStarConfig +from brittle_star_project.configs.register_configs import register_configs from brittle_star_project.dataclasses import PPOArgs from brittle_star_project.trainers.PPOTrainer import PPOTrainer from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper - -from experiment_logger import UnifiedLogger -from experiment_logger.config_utils import merge_config_with_cli, print_config +from experiment_logger import init_logger, get_logger -def make_env(config_path: str | None, num_envs: int) -> BrittleStarJaxEnvWrapper: - if config_path is None: - return BrittleStarJaxEnvWrapper.default(num_envs=num_envs) - return BrittleStarJaxEnvWrapper.from_config(config_path, num_envs=num_envs) +def make_env(cfg: BrittleStarConfig) -> BrittleStarJaxEnvWrapper: + """Create the environment using the structured configuration.""" + return BrittleStarJaxEnvWrapper( + morphology=cfg.morphology, + arena=cfg.arena, + env_config=cfg.environment, + num_envs=cfg.ppo.num_envs, + ) -def parse_args() -> PPOArgs: - import argparse +def create_ppo_args_compat(cfg: BrittleStarConfig, run_dir: str) -> PPOArgs: + """Temporary adapter to bridge BrittleStarConfig to the legacy PPOArgs. - # Use argparse to reliably extract just the config path without swallowing --help - parser = argparse.ArgumentParser(add_help=False) - parser.add_argument("--hyperparameter-config-path", type=str, default=None) - known_args, _ = parser.parse_known_args() - - args = merge_config_with_cli(PPOArgs, config_file=known_args.hyperparameter_config_path) - return args - - -def get_git_hash() -> str: - try: - return ( - subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode("ascii").strip() - ) - except (subprocess.CalledProcessError, UnicodeDecodeError): - return "none" - - -if __name__ == "__main__": - args = parse_args() + This will be removed in Step 4.4 once PPOTrainer is refactored. + """ + # Flatten the hierarchical config into the expected PPOArgs format + args = PPOArgs( + exp_name=cfg.experiment.exp_name, + seed=cfg.experiment.seed, + torch_deterministic=cfg.experiment.torch_deterministic, + cuda=cfg.experiment.cuda, + track=cfg.logging.track, + wandb_project_name=cfg.logging.wandb_project_name, + wandb_entity=cfg.logging.wandb_entity, + capture_video=cfg.logging.capture_video, + save_model=cfg.logging.save_model, + checkpoint_frequency=cfg.logging.checkpoint_frequency, + upload_model=cfg.logging.upload_model, + hf_entity=cfg.logging.hf_entity, + total_timesteps=cfg.ppo.total_timesteps, + learning_rate=cfg.ppo.learning_rate, + num_envs=cfg.ppo.num_envs, + num_steps=cfg.ppo.num_steps, + anneal_lr=cfg.ppo.anneal_lr, + gamma=cfg.ppo.gamma, + gae_lambda=cfg.ppo.gae_lambda, + num_minibatches=cfg.ppo.num_minibatches, + update_epochs=cfg.ppo.update_epochs, + norm_adv=cfg.ppo.norm_adv, + clip_coef=cfg.ppo.clip_coef, + clip_vloss=cfg.ppo.clip_vloss, + ent_coef=cfg.ppo.ent_coef, + vf_coef=cfg.ppo.vf_coef, + max_grad_norm=cfg.ppo.max_grad_norm, + target_kl=cfg.ppo.target_kl, + run_dir=run_dir, + ) + # Compute runtime fields args.batch_size = args.num_envs * args.num_steps args.minibatch_size = args.batch_size // args.num_minibatches args.num_iterations = args.total_timesteps // args.batch_size - git_hash = get_git_hash() - run_name = f"{args.exp_name}__seed_{args.seed}__{git_hash}__{int(time.time())}" + return args - if args.run_dir is None: - run_dir = f"runs/{run_name}" - else: - run_dir = args.run_dir - os.makedirs(run_dir, exist_ok=True) +@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3") +def main(dict_cfg: DictConfig): + # 1. Convert DictConfig to structured dataclass + cfg: BrittleStarConfig = OmegaConf.to_object(dict_cfg) - # Initialize Global Logger - logger = UnifiedLogger( - config=vars(args), - project_name=args.wandb_project_name, # or default PPO-Modularity if missing + # 2. Setup run metadata + # Hydra changes CWD to the output directory by default. + # We use that as our run_dir. + run_dir = os.getcwd() + run_name = os.path.basename(run_dir) + + # 3. Initialize Logger + # We pass the resolved dictionary for WandB/YAML logging + resolved_cfg_dict = OmegaConf.to_container(dict_cfg, resolve=True, throw_on_missing=True) + init_logger( run_name=run_name, + config=resolved_cfg_dict, + project_name=cfg.logging.wandb_project_name, + entity=cfg.logging.wandb_entity, base_dir=os.path.dirname(run_dir), - use_wandb=args.track, + use_wandb=cfg.logging.track, ) + logger = get_logger() + logger.info(f"Hydra-initialized run: {run_name}") + logger.info(f"Output directory: {run_dir}") - print_config(args, title="PPO Training Configuration") + # 4. Prepare compatibility object for PPOTrainer + ppo_args = create_ppo_args_compat(cfg, run_dir) - env = make_env(args.env_config_path, args.num_envs) + # 5. Setup Environment and Torch + env = make_env(cfg) + torch.backends.cudnn.deterministic = cfg.experiment.torch_deterministic - torch.backends.cudnn.deterministic = args.torch_deterministic - - ppo_trainer = PPOTrainer(args, env, run_dir, run_name) + # 6. Train + ppo_trainer = PPOTrainer(ppo_args, env, run_dir, run_name) ppo_trainer.train() + + +if __name__ == "__main__": + register_configs() + main() From ae6beb174bb341ad8498ee032f8a11b9aee7b383 Mon Sep 17 00:00:00 2001 From: cedric Date: Wed, 15 Apr 2026 05:39:18 +0000 Subject: [PATCH 20/32] fix: updated used vars --- .../used_variables.md | 25 +++++++++++++------ 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/experiments/debug-experiment-10042026/used_variables.md b/experiments/debug-experiment-10042026/used_variables.md index e49c9e4..92f23a1 100644 --- a/experiments/debug-experiment-10042026/used_variables.md +++ b/experiments/debug-experiment-10042026/used_variables.md @@ -20,23 +20,35 @@ learning_rate: float = 2.5e-4 anneal_lr: bool = True gamma: float = 0.99 gae_lambda: float = 0.95 -num_minibatches: int = 4 update_epochs: int = 4 norm_adv: bool = True -clip_coef: float = 0.1 clip_vloss: bool = True -ent_coef: float = 0.01 -vf_coef: float = 0.5 max_grad_norm: float = 0.5 target_kl: float | None = None batch_size: int = 0 minibatch_size: int = 0 num_iterations: int = 0 -## Used config file: +## Used config file: (hpc/debug.yaml) +exp_name: "debug-experiment" +seed: 42 +track: true +wandb_project_name: "Let's-find-that-bug" +wandb_entity: "SEL3-2026-Groep-4" +run_dir: "/data/gent/465/vsc46589" num_envs: 32 num_steps: 32 -total_timesteps: 102400 +num_minibatches: 32 +total_timesteps: 409600 +num_arms: 2 +cuda: true + +ent_coef: 0.005 +vf_coef: 1.0 +clip_coef: 0.2 + +anneal_lr: true +learning_rate: 0.0003 ## Arena config: size: tuple[float, float] = (10.0, 5.0) @@ -46,7 +58,6 @@ wall_height: float = 1.5 wall_thickness: float = 0.1 ## Morphology: -num_arms: int = 2 num_segments_per_arm: int = 4 use_p_control: bool = True use_torque_control: bool = False From 93bff112088d5e5e4b8824e77f680cd550c5968a Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 15 Apr 2026 15:18:34 +0200 Subject: [PATCH 21/32] refactor: use hydra for configs in simulation and training scripts --- scripts/simulate.py | 82 +++++++---------- scripts/train.py | 59 +------------ .../environment/BrittleStarJaxEnvWrapper.py | 10 --- .../environment/__init__.py | 3 +- .../environment/env_config.py | 13 --- .../trainers/PPOTrainer.py | 87 +++++++++++-------- src/experiment_logger/__init__.py | 3 +- src/experiment_logger/config_utils.py | 61 ------------- tests/test_config.py | 5 -- 9 files changed, 87 insertions(+), 236 deletions(-) diff --git a/scripts/simulate.py b/scripts/simulate.py index 957f6fa..63b1cfe 100644 --- a/scripts/simulate.py +++ b/scripts/simulate.py @@ -1,6 +1,14 @@ +"""Simulate a trained policy in the MuJoCo viewer. + +Uses Hydra to load the same BrittleStarConfig that was used during training. +Override settings via CLI, e.g.: + python scripts/simulate.py morphology=3_arms +""" + from __future__ import annotations -import argparse +import hydra +from omegaconf import DictConfig, OmegaConf from pathlib import Path from brittle_star_project import ( @@ -10,86 +18,58 @@ from brittle_star_project import ( SimulationConfig, simulate_policy, ) -from brittle_star_project.environment import from_file -from brittle_star_project.rl import RLModel # imports concrete models via rl.__init__ +from brittle_star_project.configs.main_config import BrittleStarConfig +from brittle_star_project.configs.register_configs import register_configs +from brittle_star_project.rl import RLModel from brittle_star_project.rl.base import get_rl_model_registry MODEL_BY_NAME = get_rl_model_registry() MODEL_OPTIONS = sorted(MODEL_BY_NAME) -def parse_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description="Simulate a trained policy in the MuJoCo viewer.") - p.add_argument( - "--model", - type=str, - default=None, - help="Path to a saved model artifact. If omitted, a model is created from --model-type.", - ) - p.add_argument( - "--model-type", - choices=MODEL_OPTIONS, - default="random", - help="Which model class to instantiate when --model is omitted.", - ) - p.add_argument( - "--backend", - choices=[b for b in Backend], - default=Backend.MJX, - ) - p.add_argument("--seed", type=int, default=None) - return p.parse_args() +@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3") +def main(dict_cfg: DictConfig) -> None: + cfg: BrittleStarConfig = OmegaConf.to_object(dict_cfg) - -def main() -> None: - args = parse_args() - - morphology_cfg, arena_cfg, env_cfg = from_file("../configs/test.yaml") + # Allow overriding these via Hydra CLI or a dedicated simulate config group + # For now, defaults matching the old argparse behavior + backend = Backend.MJX + model_type = "random" + model_path = None + seed = cfg.experiment.seed # ======= ENVIRONMENT SETUP ======= - - backend = args.backend - factory = BrittleStarEnvFactory() - raw_env = factory.create_environment(backend, morphology_cfg, arena_cfg, env_cfg) - env = BrittleStarEnv(raw_env, backend=backend, config=env_cfg) + raw_env = factory.create_environment(backend, cfg.morphology, cfg.arena, cfg.environment) + env = BrittleStarEnv(raw_env, backend=backend, config=cfg.environment) - seed_for_env = int(args.seed) if args.seed is not None else 0 - state = env.reset(seed=seed_for_env) + state = env.reset(seed=seed) # ======= MODEL SETUP ======= - - # Extract the number of actuators (nu) from the environment's model, so we can pass it to the - # policy/model. nu = int(state.mj_model.nu) - if args.model is not None: - model_path = Path(args.model) - policy = RLModel.load(model_path) + if model_path is not None: + policy = RLModel.load(Path(model_path)) if hasattr(policy, "nu"): policy.nu = nu else: - model_cls = MODEL_BY_NAME[str(args.model_type)] - policy = model_cls(seed=seed_for_env) + model_cls = MODEL_BY_NAME[model_type] + policy = model_cls(seed=seed) if hasattr(policy, "nu"): policy.nu = nu - # If the policy/model has a `seed` attribute, use the provided seed (or default) to reset it. - default_seed = int(getattr(policy, "seed", seed_for_env)) - if args.seed is not None and hasattr(policy, "reset"): - policy.reset(int(args.seed)) + default_seed = int(getattr(policy, "seed", seed)) # ======= SIMULATION ======= - rollout_cfg = SimulationConfig( realtime=True, - seed=int(args.seed) if args.seed is not None else default_seed, + seed=default_seed, ) simulate_policy(policy, rollout_cfg, state) - env.close() if __name__ == "__main__": + register_configs() main() diff --git a/scripts/train.py b/scripts/train.py index 5ef2258..4224eb9 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -1,12 +1,10 @@ import os -import time import torch import hydra from omegaconf import DictConfig, OmegaConf from brittle_star_project.configs.main_config import BrittleStarConfig from brittle_star_project.configs.register_configs import register_configs -from brittle_star_project.dataclasses import PPOArgs from brittle_star_project.trainers.PPOTrainer import PPOTrainer from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper from experiment_logger import init_logger, get_logger @@ -22,52 +20,6 @@ def make_env(cfg: BrittleStarConfig) -> BrittleStarJaxEnvWrapper: ) -def create_ppo_args_compat(cfg: BrittleStarConfig, run_dir: str) -> PPOArgs: - """Temporary adapter to bridge BrittleStarConfig to the legacy PPOArgs. - - This will be removed in Step 4.4 once PPOTrainer is refactored. - """ - # Flatten the hierarchical config into the expected PPOArgs format - args = PPOArgs( - exp_name=cfg.experiment.exp_name, - seed=cfg.experiment.seed, - torch_deterministic=cfg.experiment.torch_deterministic, - cuda=cfg.experiment.cuda, - track=cfg.logging.track, - wandb_project_name=cfg.logging.wandb_project_name, - wandb_entity=cfg.logging.wandb_entity, - capture_video=cfg.logging.capture_video, - save_model=cfg.logging.save_model, - checkpoint_frequency=cfg.logging.checkpoint_frequency, - upload_model=cfg.logging.upload_model, - hf_entity=cfg.logging.hf_entity, - total_timesteps=cfg.ppo.total_timesteps, - learning_rate=cfg.ppo.learning_rate, - num_envs=cfg.ppo.num_envs, - num_steps=cfg.ppo.num_steps, - anneal_lr=cfg.ppo.anneal_lr, - gamma=cfg.ppo.gamma, - gae_lambda=cfg.ppo.gae_lambda, - num_minibatches=cfg.ppo.num_minibatches, - update_epochs=cfg.ppo.update_epochs, - norm_adv=cfg.ppo.norm_adv, - clip_coef=cfg.ppo.clip_coef, - clip_vloss=cfg.ppo.clip_vloss, - ent_coef=cfg.ppo.ent_coef, - vf_coef=cfg.ppo.vf_coef, - max_grad_norm=cfg.ppo.max_grad_norm, - target_kl=cfg.ppo.target_kl, - run_dir=run_dir, - ) - - # Compute runtime fields - args.batch_size = args.num_envs * args.num_steps - args.minibatch_size = args.batch_size // args.num_minibatches - args.num_iterations = args.total_timesteps // args.batch_size - - return args - - @hydra.main(config_path="../configs", config_name="main_config", version_base="1.3") def main(dict_cfg: DictConfig): # 1. Convert DictConfig to structured dataclass @@ -75,12 +27,10 @@ def main(dict_cfg: DictConfig): # 2. Setup run metadata # Hydra changes CWD to the output directory by default. - # We use that as our run_dir. run_dir = os.getcwd() run_name = os.path.basename(run_dir) # 3. Initialize Logger - # We pass the resolved dictionary for WandB/YAML logging resolved_cfg_dict = OmegaConf.to_container(dict_cfg, resolve=True, throw_on_missing=True) init_logger( run_name=run_name, @@ -94,15 +44,12 @@ def main(dict_cfg: DictConfig): logger.info(f"Hydra-initialized run: {run_name}") logger.info(f"Output directory: {run_dir}") - # 4. Prepare compatibility object for PPOTrainer - ppo_args = create_ppo_args_compat(cfg, run_dir) - - # 5. Setup Environment and Torch + # 4. Setup Environment and Torch env = make_env(cfg) torch.backends.cudnn.deterministic = cfg.experiment.torch_deterministic - # 6. Train - ppo_trainer = PPOTrainer(ppo_args, env, run_dir, run_name) + # 5. Train - pass structured config directly + ppo_trainer = PPOTrainer(cfg, env, run_dir, run_name) ppo_trainer.train() diff --git a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py index b143d9c..7a6d161 100644 --- a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py +++ b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py @@ -8,7 +8,6 @@ from brittle_star_project import ( ArenaConfig, Backend, ) -from brittle_star_project.environment import from_file class BrittleStarJaxEnvWrapper: @@ -86,15 +85,6 @@ class BrittleStarJaxEnvWrapper: morphology, arena, env_config, num_envs=num_envs, backend=backend ) - @staticmethod - def from_config( - config_path: str, num_envs: int, backend: Backend = Backend.MJX - ) -> "BrittleStarJaxEnvWrapper": - morphology_cfg, arena_cfg, env_cfg = from_file(config_path) - return BrittleStarJaxEnvWrapper( - morphology_cfg, arena_cfg, env_cfg, num_envs=num_envs, backend=backend - ) - def __str__(self): morphology_str = str(self._morphology) arena_str = str(self._arena) diff --git a/src/brittle_star_project/environment/__init__.py b/src/brittle_star_project/environment/__init__.py index aad8c35..78896cb 100644 --- a/src/brittle_star_project/environment/__init__.py +++ b/src/brittle_star_project/environment/__init__.py @@ -1,4 +1,4 @@ -from .env_config import ArenaConfig, EnvConfig, MorphologyConfig, from_file +from .env_config import ArenaConfig, EnvConfig, MorphologyConfig from .env_types import Backend, Task from .env_wrapper import BrittleStarEnv, StepResult from .factory import BrittleStarEnvFactory @@ -12,5 +12,4 @@ __all__ = [ "BrittleStarEnv", "StepResult", "BrittleStarEnvFactory", - "from_file", ] diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index 206e902..48f0e34 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/src/brittle_star_project/environment/env_config.py @@ -60,16 +60,3 @@ class EnvConfig: # Light escape # Per docs in upstream env config: integer factors of 200. light_perlin_noise_scale: int = 0 - - -def from_file(path: str) -> tuple[MorphologyConfig, ArenaConfig, EnvConfig]: - """Load configurations from a YAML file.""" - import yaml - - with open(path, "r") as f: - config_dict = yaml.safe_load(f) - - morphology = MorphologyConfig(**config_dict.get("morphology", {})) - arena = ArenaConfig(**config_dict.get("arena", {})) - env = EnvConfig(**config_dict.get("env", {})) - return morphology, arena, env diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 6bc94ef..b20c4dd 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -13,7 +13,8 @@ from flax.training.train_state import TrainState from experiment_logger import get_logger -from brittle_star_project.dataclasses import EpisodeStatistics, PPOArgs +from brittle_star_project.configs.main_config import BrittleStarConfig +from brittle_star_project.dataclasses import EpisodeStatistics from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper from brittle_star_project.MLPs.mlps import ( Actor, @@ -210,14 +211,23 @@ class TrainingMeasurements: class PPOTrainer: - def __init__(self, args: PPOArgs, env: BrittleStarJaxEnvWrapper, run_dir: str, run_name: str): - self.args = args + def __init__( + self, cfg: BrittleStarConfig, env: BrittleStarJaxEnvWrapper, run_dir: str, run_name: str + ): + self.cfg = cfg + self.ppo = cfg.ppo + self.experiment = cfg.experiment + self.logging_cfg = cfg.logging self.env = env self.run_dir = run_dir self.run_name = run_name self.logger = get_logger() - self.key = jax.random.PRNGKey(args.seed) + # Derived runtime fields + self.batch_size = self.ppo.num_envs * self.ppo.num_steps + self.num_iterations = self.ppo.total_timesteps // self.batch_size + + self.key = jax.random.PRNGKey(self.experiment.seed) self.sensor, self.feature_extractor, self.actor, self.critic = self._init_agent() self.sensor.apply = jax.jit(self.sensor.apply) @@ -228,7 +238,7 @@ class PPOTrainer: self._rollout_jit = jax.jit( partial( _rollout_jit, - max_steps=self.args.num_steps, + max_steps=self.ppo.num_steps, step_env_fn=partial(_step_env_wrapped, env_step_fn=self.env.step), sensor=self.sensor, feature_extractor=self.feature_extractor, @@ -239,15 +249,15 @@ class PPOTrainer: self._compute_gae_jit = jax.jit( partial( _compute_gae_jit, - num_envs=self.args.num_envs, - gamma=self.args.gamma, - gae_lambda=self.args.gae_lambda, + num_envs=self.ppo.num_envs, + gamma=self.ppo.gamma, + gae_lambda=self.ppo.gae_lambda, feature_extractor=self.feature_extractor, critic=self.critic, ) ) - self._ppo = PPO(self.args, self.sensor, self.actor, self.critic, self.feature_extractor) + self._ppo = PPO(self.ppo, self.sensor, self.actor, self.critic, self.feature_extractor) self.agent_state = self._init_agent_state() @@ -256,10 +266,10 @@ class PPOTrainer: self._init_random() def _init_random(self): - self.logger.info(f"[RANDOM]: Setting random seed to {self.args.seed}") + self.logger.info(f"[RANDOM]: Setting random seed to {self.experiment.seed}") - random.seed(self.args.seed) - np.random.seed(self.args.seed) + random.seed(self.experiment.seed) + np.random.seed(self.experiment.seed) def _init_agent(self): self.logger.info("[AGENT]: Initializing agent...") @@ -299,17 +309,17 @@ class PPOTrainer: AgentParams(sensor_params, actor_params, critic_params, feature_extractor_params) ), tx=optax.chain( - optax.clip_by_global_norm(self.args.max_grad_norm), + optax.clip_by_global_norm(self.ppo.max_grad_norm), optax.inject_hyperparams(optax.adam)( learning_rate=partial( _linear_schedule, - minibatch_count=self.args.num_minibatches, - update_epochs=self.args.update_epochs, - num_iterations=self.args.num_iterations, - learning_rate=self.args.learning_rate, + minibatch_count=self.ppo.num_minibatches, + update_epochs=self.ppo.update_epochs, + num_iterations=self.num_iterations, + learning_rate=self.ppo.learning_rate, ) - if self.args.anneal_lr - else self.args.learning_rate, + if self.ppo.anneal_lr + else self.ppo.learning_rate, eps=1e-5, ), ), @@ -319,10 +329,10 @@ class PPOTrainer: self.logger.info("[EPISODE STATS]: Initializing episode stats...") return EpisodeStatistics( - episode_returns=jnp.zeros(self.args.num_envs, dtype=jnp.float32), - episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32), - returned_episode_returns=jnp.zeros(self.args.num_envs, jnp.float32), - returned_episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32), + episode_returns=jnp.zeros(self.ppo.num_envs, dtype=jnp.float32), + episode_lengths=jnp.zeros(self.ppo.num_envs, dtype=jnp.int32), + returned_episode_returns=jnp.zeros(self.ppo.num_envs, jnp.float32), + returned_episode_lengths=jnp.zeros(self.ppo.num_envs, dtype=jnp.int32), ) def _rollout(self, env_state, next_obs, next_done) -> tuple[Any, ...]: @@ -371,7 +381,7 @@ class PPOTrainer: "losses/loss": training_measurements.loss[-1, -1].item(), "charts/SPS": int(global_step / (time.time() - start_time)), "charts/SPS_update": int( - self.args.num_envs * self.args.num_steps / (time.time() - iteration_time_start) + self.ppo.num_envs * self.ppo.num_steps / (time.time() - iteration_time_start) ), } self.logger.log(metrics, step=global_step) @@ -451,8 +461,14 @@ class PPOTrainer: def _save_model(self, model_path: str): self.logger.info("[SAVE]: Saving the final model...") + from dataclasses import asdict as _asdict + + config_dict = { + "experiment": _asdict(self.experiment), + "ppo": _asdict(self.ppo), + } params = [ - vars(self.args), + config_dict, [ self.agent_state.params["sensor_params"], self.agent_state.params["actor_params"], @@ -464,8 +480,7 @@ class PPOTrainer: def train(self): """ - Train the PPO agent for a specified number of iterations - (passed through PPOArgs in constructor). + Train the PPO agent for a specified number of iterations. Closes the environment at the end of training. """ self.logger.info(f"running name: {self.run_name}") @@ -473,16 +488,16 @@ class PPOTrainer: self.logger.info("[TRAIN]: Resetting environment...") self.logger.log_non_interactive(f"Initial reset started: {time.ctime()}") - env_state = self.env.reset(seed=self.args.seed) + env_state = self.env.reset(seed=self.experiment.seed) next_obs = _convert_obs_dict_to_array(env_state.observations) - next_done = jnp.zeros(self.args.num_envs, dtype=jnp.bool_) + next_done = jnp.zeros(self.ppo.num_envs, dtype=jnp.bool_) self.logger.log_non_interactive(f"Initial reset completed: {time.ctime()}") global_step = 0 start_time = time.time() - iter_bar = self.logger.progress_bar(range(1, self.args.num_iterations + 1)) + iter_bar = self.logger.progress_bar(range(1, self.num_iterations + 1)) for iteration in iter_bar: iteration_time_start = time.time() @@ -490,7 +505,7 @@ class PPOTrainer: env_state, next_obs, next_done, iteration=iteration ) - global_step += self.args.num_steps * self.args.num_envs + global_step += self.ppo.num_steps * self.ppo.num_envs self._log( global_step, self.episode_stats, @@ -500,20 +515,20 @@ class PPOTrainer: ) sps = int(global_step / (time.time() - start_time)) - remaining_steps = self.args.total_timesteps - global_step + remaining_steps = self.ppo.total_timesteps - global_step eta_seconds = int(remaining_steps / sps) if sps > 0 else 0 eta_str = str(datetime.timedelta(seconds=eta_seconds)) self.logger.log_non_interactive( - f"Iteration {iteration}/{self.args.num_iterations} | " - f"Step {global_step}/{self.args.total_timesteps} | " + f"Iteration {iteration}/{self.num_iterations} | " + f"Step {global_step}/{self.ppo.total_timesteps} | " f"SPS {sps} | " f"Return {training_measurements.avg_episodic_return:.4f} | " f"ETA {eta_str}" ) - if self.args.save_model: - model_path = f"{self.run_dir}/{self.args.exp_name}.cleanrl_model" + if self.logging_cfg.save_model: + model_path = f"{self.run_dir}/{self.experiment.exp_name}.cleanrl_model" self._save_model(model_path=model_path) self._close() diff --git a/src/experiment_logger/__init__.py b/src/experiment_logger/__init__.py index c89af7f..53e57c0 100644 --- a/src/experiment_logger/__init__.py +++ b/src/experiment_logger/__init__.py @@ -4,7 +4,7 @@ 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, merge_config_with_cli +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 @@ -17,6 +17,5 @@ __all__ = [ "init_wandb", "finish_wandb", "load_yaml_config", - "merge_config_with_cli", ] __version__ = "0.1.0" diff --git a/src/experiment_logger/config_utils.py b/src/experiment_logger/config_utils.py index 09cb4ce..6017041 100644 --- a/src/experiment_logger/config_utils.py +++ b/src/experiment_logger/config_utils.py @@ -72,67 +72,6 @@ def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T: return cls(**filtered_config) -def merge_config_with_cli(config_class: Type[T], config_file: str | None = None) -> T: - """Merge YAML config with CLI arguments, with CLI taking precedence. - - Args: - config_class: Dataclass type to create - config_file: Path to YAML config file (optional) - - Returns: - Instance of config_class with merged configuration - """ - # Parse CLI args first to get the default/CLI values - import tyro - - # Check if --config is in sys.argv and extract it - extracted_config_file = config_file - if "--config" in sys.argv: - config_idx = sys.argv.index("--config") - if config_idx + 1 < len(sys.argv): - extracted_config_file = sys.argv[config_idx + 1] - # Remove from sys.argv so tyro doesn't see it - sys.argv.pop(config_idx) # Remove --config - sys.argv.pop(config_idx) # Remove config file path - - # Load YAML config if available - yaml_config = {} - if extracted_config_file and os.path.exists(extracted_config_file): - yaml_config = load_yaml_config(extracted_config_file) - get_logger().info(f"Merging YAML config from {extracted_config_file} with CLI args") - elif extracted_config_file: - get_logger().warning(f"Config file not found: {extracted_config_file}, using CLI args only") - - # Create default instance to know what the defaults are - default_instance = config_class() - default_dict = {f.name: getattr(default_instance, f.name) for f in fields(config_class)} # type: ignore - - # Parse CLI args - cli_instance = tyro.cli(config_class) - cli_dict = {f.name: getattr(cli_instance, f.name) for f in fields(config_class)} # type: ignore - - # Merge configs: YAML as base, CLI overrides non-default values - final_config = {} - - for field in fields(config_class): # type: ignore - field_name = field.name - default_value = default_dict[field_name] - yaml_value = yaml_config.get(field_name, default_value) - cli_value = cli_dict[field_name] - - # Use CLI value if it's different from default, otherwise use YAML value - if cli_value != default_value: - final_config[field_name] = cli_value - if yaml_value != default_value and yaml_value != cli_value: - get_logger().info(f"CLI override: {field_name}={cli_value} (YAML had {yaml_value})") - else: - final_config[field_name] = yaml_value - if yaml_value != default_value: - get_logger().info(f"YAML config: {field_name}={yaml_value}") - - return config_class(**final_config) - - def print_config(config: Any, title: str = "Configuration"): """Pretty print configuration.""" get_logger().info(f"{title}:") diff --git a/tests/test_config.py b/tests/test_config.py index 212e4b4..e0ee0ef 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -31,8 +31,3 @@ class TestYamlConfig: with pytest.raises(FileNotFoundError): load_yaml_config("nonexistent.yaml") - - def test_merge_config_with_cli_is_callable(self): - from experiment_logger.config_utils import merge_config_with_cli - - assert callable(merge_config_with_cli) From b4f1e98f8cf60c5d34f87ff8769f74c6a0d10a8c Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 15 Apr 2026 15:29:32 +0200 Subject: [PATCH 22/32] chore: cleanup and HPC integration --- configs/default_ppo.yaml | 48 -------- configs/dev_test.yaml | 42 ------- configs/example.yaml | 4 - configs/hpc/smoke_test.yaml | 11 -- configs/hpc/wandb_expand.yaml | 10 -- configs/hpc/wandb_test.yaml | 11 -- configs/personal_template.yaml | 40 ------ configs/production_training.yaml | 42 ------- scripts/hpc/train.pbs | 8 +- .../dataclasses/PPOArgs.py | 116 ------------------ .../dataclasses/__init__.py | 2 - 11 files changed, 4 insertions(+), 330 deletions(-) delete mode 100644 configs/default_ppo.yaml delete mode 100644 configs/dev_test.yaml delete mode 100644 configs/example.yaml delete mode 100644 configs/hpc/smoke_test.yaml delete mode 100644 configs/hpc/wandb_expand.yaml delete mode 100644 configs/hpc/wandb_test.yaml delete mode 100644 configs/personal_template.yaml delete mode 100644 configs/production_training.yaml delete mode 100644 src/brittle_star_project/dataclasses/PPOArgs.py diff --git a/configs/default_ppo.yaml b/configs/default_ppo.yaml deleted file mode 100644 index 1b06c3d..0000000 --- a/configs/default_ppo.yaml +++ /dev/null @@ -1,48 +0,0 @@ -# PPO Training Configuration Template -# -# This file provides an example configuration for PPO training. -# Copy this file and modify it for your specific experiments. -# -# Usage: -# python src/train.py --config-path configs/my_config.yaml -# Or override specific parameters: -# python src/train.py --learning-rate 0.001 --num-envs 32 - -# Experiment settings -exp_name: "brittle_star_ppo" -seed: 1 - -# Tracking settings -track: false # Set to true to enable WandB logging -wandb_project_name: "PPO-Modularity" -wandb_entity: "SEL3-2026-Groep-4" # Set to your WandB username or team name - -# Model saving -save_model: true -checkpoint_frequency: 100 # Save checkpoint every N iterations (0 = no checkpoints) - -# Environment settings -num_envs: 16 - -# Training hyperparameters -total_timesteps: 10000000 -learning_rate: 0.00025 -num_steps: 128 -anneal_lr: true - -# PPO specific -gamma: 0.99 -gae_lambda: 0.95 -num_minibatches: 4 -update_epochs: 4 -norm_adv: true -clip_coef: 0.1 -clip_vloss: true -ent_coef: 0.01 -vf_coef: 0.5 -max_grad_norm: 0.5 -target_kl: null - -# Hardware -cuda: true -torch_deterministic: true diff --git a/configs/dev_test.yaml b/configs/dev_test.yaml deleted file mode 100644 index b64d330..0000000 --- a/configs/dev_test.yaml +++ /dev/null @@ -1,42 +0,0 @@ -# Quick Development/Testing Configuration -# -# Fast configuration for development and testing with short runs. - -# Experiment settings -exp_name: "brittle_star_dev_test" -seed: 123 - -# Tracking settings - IMPORTANT: Set your own wandb_entity! -track: true -wandb_project_name: "PPO-Modularity-Dev" -wandb_entity: "SEL3-2026-Groep-4" # ⚠️ SET THIS TO YOUR WANDB USERNAME OR TEAM - -# Model saving -save_model: true -checkpoint_frequency: 10 # More frequent checkpoints for testing - -# Environment settings -num_envs: 4 # Smaller for faster iteration - -# Training hyperparameters - Fast/testing -total_timesteps: 100000 # Short run for testing -learning_rate: 0.001 # Higher learning rate for faster learning -num_steps: 64 # Shorter rollouts -anneal_lr: true - -# PPO specific - Optimized for quick results -gamma: 0.99 -gae_lambda: 0.95 -num_minibatches: 2 -update_epochs: 2 # Fewer epochs for speed -norm_adv: true -clip_coef: 0.1 -clip_vloss: true -ent_coef: 0.02 # Higher entropy for exploration -vf_coef: 0.5 -max_grad_norm: 0.5 -target_kl: null - -# Hardware -cuda: true -torch_deterministic: true \ No newline at end of file diff --git a/configs/example.yaml b/configs/example.yaml deleted file mode 100644 index baeb1b0..0000000 --- a/configs/example.yaml +++ /dev/null @@ -1,4 +0,0 @@ -morphology: - segments_per_arm: [4, 4] - use_p_control: true - use_torque_control: false diff --git a/configs/hpc/smoke_test.yaml b/configs/hpc/smoke_test.yaml deleted file mode 100644 index 1dd2fc0..0000000 --- a/configs/hpc/smoke_test.yaml +++ /dev/null @@ -1,11 +0,0 @@ -# Minimal config to verify HPC setup is functional. -# Run with: python scripts/train.py --config-path configs/hpc/smoke_test.yaml -exp_name: "hpc_smoke_test" -seed: 0 -track: false # Test WandB integration -capture_video: false # No rendering for smoke test -save_model: true # Test the end-of-training save routine -num_envs: 512 -total_timesteps: 65536 -num_steps: 128 -cuda: true diff --git a/configs/hpc/wandb_expand.yaml b/configs/hpc/wandb_expand.yaml deleted file mode 100644 index 432160f..0000000 --- a/configs/hpc/wandb_expand.yaml +++ /dev/null @@ -1,10 +0,0 @@ -exp_name: "explained_var_fun_more_steps" -seed: 42 -track: true -wandb_project_name: "LET-THERE-BE-MORE-LOGGING" -wandb_entity: "SEL3-2026-Groep-4" - -num_envs: 16 -num_steps: 256 -total_timesteps: 50000 -cuda: true \ No newline at end of file diff --git a/configs/hpc/wandb_test.yaml b/configs/hpc/wandb_test.yaml deleted file mode 100644 index bf2229d..0000000 --- a/configs/hpc/wandb_test.yaml +++ /dev/null @@ -1,11 +0,0 @@ -# Configuration to verify WandB online tracking -exp_name: "hpc_wandb_verification" -seed: 42 -track: true # Enabled for testing WandB -wandb_project_name: "PPO-Modularity" -wandb_entity: "SEL3-2026-Groep-4" - -num_envs: 128 -total_timesteps: 50000 # Short run for quick verification -num_steps: 128 -cuda: true diff --git a/configs/personal_template.yaml b/configs/personal_template.yaml deleted file mode 100644 index 67caef7..0000000 --- a/configs/personal_template.yaml +++ /dev/null @@ -1,40 +0,0 @@ -# Personal Configuration Example for Team Member -# -# Copy this template and customize for your personal experiments - -# Experiment settings - PERSONALIZE THESE -exp_name: "YOUR_NAME_experiment_v1" # ⚠️ Change YOUR_NAME -seed: 42 - -# WandB settings - ⚠️ IMPORTANT: Set your credentials! -track: true # Enable WandB tracking -wandb_project_name: "PPO-Modularity" -wandb_entity: "SEL3-2026-Groep-4" # ⚠️ CHANGE THIS to your WandB username/team - -# Quick experiment settings (modify as needed) -total_timesteps: 500000 # 500K for quick results -num_envs: 8 -learning_rate: 0.0005 -num_steps: 128 - -# Model saving -save_model: true -checkpoint_frequency: 25 # Save checkpoints frequently - -# Standard PPO settings (usually don't need to change) -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 -anneal_lr: true - -# Hardware -cuda: true -torch_deterministic: true \ No newline at end of file diff --git a/configs/production_training.yaml b/configs/production_training.yaml deleted file mode 100644 index 6dce29b..0000000 --- a/configs/production_training.yaml +++ /dev/null @@ -1,42 +0,0 @@ -# Production Training Configuration -# -# Full-scale training configuration for production runs -# with wandb logging enabled. - -# Experiment settings -exp_name: "brittle_star_production_training" -seed: 42 - -# Tracking -track: true -capture_video: false -wandb_project_name: "PPO-Modularity" -wandb_entity: "SEL3-2026-Groep-4" - -# Model saving -save_model: true -checkpoint_frequency: 100 # Save checkpoint every 100 iterations - -# Environment settings -num_envs: 512 - -# Training hyperparameters -total_timesteps: 50000000 -num_steps: 256 -num_minibatches: 4 -update_epochs: 4 - -learning_rate: 2.5e-4 -anneal_lr: true -gamma: 0.99 -gae_lambda: 0.95 -clip_coef: 0.1 -clip_vloss: true -ent_coef: 0.01 -vf_coef: 0.5 -max_grad_norm: 0.5 -target_kl: null - -# Hardware -cuda: true -torch_deterministic: true \ No newline at end of file diff --git a/scripts/hpc/train.pbs b/scripts/hpc/train.pbs index b9c8519..60ba4d8 100644 --- a/scripts/hpc/train.pbs +++ b/scripts/hpc/train.pbs @@ -63,11 +63,11 @@ elif [ -f "$PBS_O_WORKDIR/.env" ]; then export $(grep -v '^#' "$PBS_O_WORKDIR/.env" | xargs) fi -# TODO Once experiments get serious, change the config +# Run training using Hydra overrides python scripts/train.py \ - --env-config-path configs/hpc/wandb_expand.yaml \ - --hyperparameter-config-path configs/hpc/wandb_expand.yaml \ - --run-dir "$SCRATCH_RUNDIR" + hydra.run.dir="$SCRATCH_RUNDIR" \ + ppo=stable \ + logging=wandb_enabled echo ">>> Staging out results to $DATA_RUNDIR..." cp -r "$SCRATCH_RUNDIR/." "$DATA_RUNDIR/" diff --git a/src/brittle_star_project/dataclasses/PPOArgs.py b/src/brittle_star_project/dataclasses/PPOArgs.py deleted file mode 100644 index 036b44f..0000000 --- a/src/brittle_star_project/dataclasses/PPOArgs.py +++ /dev/null @@ -1,116 +0,0 @@ -from dataclasses import dataclass - -import jax - - -@jax.tree_util.register_dataclass -@dataclass -class PPOArgs: - """ - source: https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/ppo_atari_envpool_xla_jax_scan.py - """ - - # path to environment config file, if None, use default config - env_config_path: str | None = None - - # path to hyperparameter config file (yaml), if None, use default config - hyperparameter_config_path: str | None = None - - # the name of this experiment - exp_name: str = "brittle_star_ppo" - - # the directory to save the experiment results - run_dir: str | None = None - - # seed of the experiment - seed: int = 1 - - # if toggled, `torch.backends.cudnn.deterministic=False` - torch_deterministic: bool = True - - # if toggled, cuda will be enabled by default - cuda: bool = True - - # if toggled, this experiment will be tracked with Weights and Biases - track: bool = False - - # the wandb's project name - wandb_project_name: str = "PPO-Modularity" - - # the entity (team) of wandb's project - wandb_entity: str | None = "SEL3-2026-Groep-4" - - # whether to capture videos of the agent performances (check out `videos` folder) - capture_video: bool = False - - # whether to save model into the `runs/{run_name}` folder - save_model: bool = True - - # checkpoint frequency (in iterations, 0 = no intermediate checkpoints) - checkpoint_frequency: int = 100 - - # whether to upload the saved model to huggingface - upload_model: bool = False - - # the user or org name of the model repository from the Hugging Face Hub - hf_entity: str = "" - - # ==== Algorithm specific dataclasses ==== - - # total timesteps of the experiments - total_timesteps: int = 10000000 - - # the learning rate of the optimizer - learning_rate: float = 2.5e-4 - - # the number of parallel game environments - num_envs: int = 100 - - # the number of steps to run in each environment per policy rollout - num_steps: int = 128 - - # Toggle learning rate annealing for policy and value networks - anneal_lr: bool = True - - # the discount factor gamma - gamma: float = 0.99 - - # the lambda for the general advantage estimation - gae_lambda: float = 0.95 - - # the number of mini-batches - num_minibatches: int = 4 - - # the K epochs to update the policy - update_epochs: int = 4 - - # Toggles advantages normalization - norm_adv: bool = True - - # the surrogate clipping coefficient - clip_coef: float = 0.1 - - # Toggles whether or not to use a clipped loss for the value function, as per the paper. - clip_vloss: bool = True - - # coefficient of the entropy - ent_coef: float = 0.01 - - # coefficient of the value function - vf_coef: float = 0.5 - - # the maximum norm for the gradient clipping - max_grad_norm: float = 0.5 - - # the target KL divergence threshold - target_kl: float | None = None - - # ==== to be filled in runtime ==== - # the batch size (computed in runtime) - batch_size: int = 0 - - # the mini-batch size (computed in runtime) - minibatch_size: int = 0 - - # the number of iterations (computed in runtime) - num_iterations: int = 0 diff --git a/src/brittle_star_project/dataclasses/__init__.py b/src/brittle_star_project/dataclasses/__init__.py index 3501b2b..f2d2ad1 100644 --- a/src/brittle_star_project/dataclasses/__init__.py +++ b/src/brittle_star_project/dataclasses/__init__.py @@ -1,8 +1,6 @@ -from .PPOArgs import PPOArgs from .EpisodeStatistics import EpisodeStatistics __all__ = [ - "PPOArgs", "EpisodeStatistics", ] From 70bd78833d1306bdb1aadb263f04790ad790f0d3 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 15 Apr 2026 15:34:18 +0200 Subject: [PATCH 23/32] chore: remove old .gitkeeps --- configs/.gitkeep | 0 tests/.gitkeep | 0 2 files changed, 0 insertions(+), 0 deletions(-) delete mode 100644 configs/.gitkeep delete mode 100644 tests/.gitkeep diff --git a/configs/.gitkeep b/configs/.gitkeep deleted file mode 100644 index e69de29..0000000 diff --git a/tests/.gitkeep b/tests/.gitkeep deleted file mode 100644 index e69de29..0000000 From eff0c7c1df16f4647e1d55e114022bc60034f9d6 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 15 Apr 2026 16:47:17 +0200 Subject: [PATCH 24/32] test: additional testing and visual confirmations --- .../configs/config_experiment.py | 1 + .../configs/main_config.py | 5 +- .../environment/env_config.py | 12 ++-- .../trainers/PPOTrainer.py | 7 ++ tests/test_config.py | 33 --------- tests/test_configs.py | 42 ++++++++++++ tests/test_morphology_render.py | 58 ++++++++++++++++ tests/test_network_shapes.py | 67 +++++++++++++++++++ 8 files changed, 184 insertions(+), 41 deletions(-) delete mode 100644 tests/test_config.py create mode 100644 tests/test_configs.py create mode 100644 tests/test_morphology_render.py create mode 100644 tests/test_network_shapes.py diff --git a/src/brittle_star_project/configs/config_experiment.py b/src/brittle_star_project/configs/config_experiment.py index 69ceca4..716d3b6 100644 --- a/src/brittle_star_project/configs/config_experiment.py +++ b/src/brittle_star_project/configs/config_experiment.py @@ -7,3 +7,4 @@ class ExperimentConfig: seed: int = 1 torch_deterministic: bool = True cuda: bool = True + debug_sanity: bool = False diff --git a/src/brittle_star_project/configs/main_config.py b/src/brittle_star_project/configs/main_config.py index aab864b..cd54473 100644 --- a/src/brittle_star_project/configs/main_config.py +++ b/src/brittle_star_project/configs/main_config.py @@ -18,8 +18,9 @@ class BrittleStarConfig: experiment: ExperimentConfig = field(default_factory=ExperimentConfig) logging: LoggingConfig = field(default_factory=LoggingConfig) ppo: PPOConfig = field(default_factory=PPOConfig) - # Default to centralized; swap with architecture=decentralized on the CLI. - architecture: ArchitectureConfig = field(default_factory=CentralizedConfig) + # This field is polymorphic; defaults to the base class to allow subclasses + # (CentralizedConfig, DecentralizedConfig) to be merged in via Hydra. + architecture: ArchitectureConfig = field(default_factory=ArchitectureConfig) morphology: MorphologyConfig = field(default_factory=MorphologyConfig) arena: ArenaConfig = field(default_factory=ArenaConfig) environment: EnvConfig = field(default_factory=EnvConfig) diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index 48f0e34..ee30e87 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/src/brittle_star_project/environment/env_config.py @@ -5,7 +5,7 @@ from dataclasses import dataclass, field from .env_types import Task -@dataclass(frozen=True, slots=True) +@dataclass class MorphologyConfig: """Brittle star morphology configuration. @@ -17,7 +17,7 @@ class MorphologyConfig: The upstream biorobot library natively supports per-arm segment counts. """ - segments_per_arm: tuple[int, ...] = (4, 4, 4, 4, 4) + segments_per_arm: list[int] = field(default_factory=lambda: [4, 4, 4, 4, 4]) use_p_control: bool = True use_torque_control: bool = False @@ -26,16 +26,16 @@ class MorphologyConfig: return len(self.segments_per_arm) -@dataclass(frozen=True, slots=True) +@dataclass class ArenaConfig: - size: tuple[float, float] = (10.0, 5.0) + size: list[float] = field(default_factory=lambda: [10.0, 5.0]) sand_ground_color: bool = True attach_target: bool = True wall_height: float = 1.5 wall_thickness: float = 0.1 -@dataclass(frozen=True, slots=True) +@dataclass class EnvConfig: """Shared environment settings. @@ -50,7 +50,7 @@ class EnvConfig: camera_ids: list[int] = field(default_factory=lambda: [0, 1]) # (height, width) - render_size: tuple[int, int] = (480, 640) + render_size: list[int] = field(default_factory=lambda: [480, 640]) joint_randomization_noise_scale: float = 0.0 diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index b20c4dd..3dc9666 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -527,6 +527,13 @@ class PPOTrainer: f"ETA {eta_str}" ) + if getattr(self.config.experiment, "debug_sanity", False): + self.logger.log( + "\n[SANITY CHECK] Successfully completed 1 epoch of data collection and gradient updates." + ) + self.logger.log("[SANITY CHECK] Gradients flowed without NaN. Exiting gracefully.") + break + if self.logging_cfg.save_model: model_path = f"{self.run_dir}/{self.experiment.exp_name}.cleanrl_model" self._save_model(model_path=model_path) diff --git a/tests/test_config.py b/tests/test_config.py deleted file mode 100644 index e0ee0ef..0000000 --- a/tests/test_config.py +++ /dev/null @@ -1,33 +0,0 @@ -"""Tests for YAML config loading.""" - -import sys -from pathlib import Path - -import pytest - -# Ensure src is on the path when running from the project root -sys.path.insert(0, str(Path(__file__).parent.parent / "src")) - -CONFIGS_DIR = Path(__file__).parent.parent / "configs" - - -class TestYamlConfig: - def test_load_yaml_config(self): - from experiment_logger.config_utils import load_yaml_config - - config = load_yaml_config(str(CONFIGS_DIR / "default_ppo.yaml")) - assert isinstance(config, dict) - assert "total_timesteps" in config - assert "learning_rate" in config - - def test_load_dev_test_config(self): - from experiment_logger.config_utils import load_yaml_config - - config = load_yaml_config(str(CONFIGS_DIR / "dev_test.yaml")) - assert config["total_timesteps"] == 100000 - - def test_missing_config_raises(self): - from experiment_logger.config_utils import load_yaml_config - - with pytest.raises(FileNotFoundError): - load_yaml_config("nonexistent.yaml") diff --git a/tests/test_configs.py b/tests/test_configs.py new file mode 100644 index 0000000..40ff35e --- /dev/null +++ b/tests/test_configs.py @@ -0,0 +1,42 @@ +import pytest +from hydra import compose, initialize +from omegaconf import OmegaConf +from brittle_star_project.configs.main_config import BrittleStarConfig +from brittle_star_project.configs.register_configs import register_configs + +# Registration must happen before composition to enable validation against schemas +register_configs() + + +def test_config_composition_centralized(): + """Test that the centralized configuration composes and validates correctly.""" + with initialize(version_base="1.3", config_path="../configs"): + # We compose the config; it follows main_config.yaml + cfg = compose(config_name="main_config", overrides=["architecture=centralized"]) + + # Merge with the dataclass class to get a structured DictConfig, + # then convert to a real dataclass instance to verify validation. + structured_cfg = OmegaConf.to_object(OmegaConf.merge(BrittleStarConfig, cfg)) + + # Basic assertions + assert "CentralizedConfig" in str(type(structured_cfg.architecture)) + assert isinstance(structured_cfg.ppo.learning_rate, float) + assert structured_cfg.ppo.learning_rate > 0 + + +def test_config_composition_decentralized(): + """Test that the decentralized configuration composes and validates correctly.""" + with initialize(version_base="1.3", config_path="../configs"): + cfg = compose(config_name="main_config", overrides=["architecture=decentralized"]) + + # Merge and convert to dataclass instance + structured_cfg = OmegaConf.to_object(OmegaConf.merge(BrittleStarConfig, cfg)) + + # Basic assertions + assert "DecentralizedConfig" in str(type(structured_cfg.architecture)) + assert isinstance(structured_cfg.ppo.learning_rate, float) + assert structured_cfg.ppo.learning_rate > 0 + + # Decentralized specifics + assert hasattr(structured_cfg.architecture, "message_passing_steps") + assert structured_cfg.architecture.message_passing_steps > 0 diff --git a/tests/test_morphology_render.py b/tests/test_morphology_render.py new file mode 100644 index 0000000..c127a01 --- /dev/null +++ b/tests/test_morphology_render.py @@ -0,0 +1,58 @@ +import jax +import os +import mujoco +from PIL import Image +from brittle_star_project.environment.env_config import EnvConfig, MorphologyConfig, ArenaConfig +from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper + + +def test_render_morphologies(): + base_dir = "runs/renders" + os.makedirs(base_dir, exist_ok=True) + + # --- 1. Full 5-Arm Morphology --- + morph_full = MorphologyConfig(segments_per_arm=[4, 4, 4, 4, 4]) + env_full = BrittleStarJaxEnvWrapper( + morphology=morph_full, arena=ArenaConfig(), env_config=EnvConfig(), num_envs=1 + ) + state_full = env_full.reset(seed=0) + + model_full = state_full.mj_model + data_full = state_full.mj_data + + # 1. Compute forward kinematics so geoms are correctly positioned + mujoco.mj_forward(model_full, data_full) + + # 2. Render using the environment's primary camera (camera=0) + renderer_full = mujoco.Renderer(model=model_full) + renderer_full.update_scene(data_full, camera=1) + pixels_full = renderer_full.render() + image_path = os.path.join(base_dir, "full_5_arm.png") + Image.fromarray(pixels_full).save(image_path) + print(f"Generated full morphology render: {image_path}") + + # --- 2. Partially Amputated Morphology --- + morph_amp = MorphologyConfig(segments_per_arm=[4, 0, 4, 2, 4]) + env_amp = BrittleStarJaxEnvWrapper( + morphology=morph_amp, arena=ArenaConfig(), env_config=EnvConfig(), num_envs=1 + ) + state_amp = env_amp.reset(seed=0) + + model_amp = state_amp.mj_model + data_amp = state_amp.mj_data + + # Compute forward kinematics + mujoco.mj_forward(model_amp, data_amp) + + renderer_amp = mujoco.Renderer(model=model_amp) + renderer_amp.update_scene(data_amp, camera=1) + pixels_amp = renderer_amp.render() + image_path = os.path.join(base_dir, "amputated_arm.png") + Image.fromarray(pixels_amp).save(image_path) + print(f"Generated amputated morphology render: {image_path}") + + print("Morphology render test successful!") + + +if __name__ == "__main__": + test_render_morphologies() diff --git a/tests/test_network_shapes.py b/tests/test_network_shapes.py new file mode 100644 index 0000000..32e9898 --- /dev/null +++ b/tests/test_network_shapes.py @@ -0,0 +1,67 @@ +import jax +import jax.numpy as jnp +from brittle_star_project.environment.padded_obs_wrapper import ( + compute_padding_masks, + pad_observations_batched, +) + +# We use Actor and OneDenseLayerMLP (as the critic) based on your mlps.py +from brittle_star_project.MLPs.mlps import Actor, OneDenseLayerMLP + + +def test_centralized_forward_pass_with_padding(): + batch_size = 2 + + # 1. Simulate Amputated Observation [4, 0, 4, 2, 4] -> 14 segments total + # 14 segments * 2 = 28 joints + amputated_obs = { + "joint_position": jnp.zeros((batch_size, 28)), + "joint_velocity": jnp.zeros((batch_size, 28)), + "segment_contact": jnp.zeros((batch_size, 14)), + } + + # 2. Pad Observation using the boolean scattering wrapper + masks = compute_padding_masks(segments_per_arm=(4, 0, 4, 2, 4)) + padded_obs = pad_observations_batched(amputated_obs, masks) + + # Assertions to ensure padding sizes are correct (40 joints, 20 segments) + assert padded_obs["joint_position"].shape == (batch_size, 40), "Padding failed for joint keys" + assert padded_obs["segment_contact"].shape == (batch_size, 20), ( + "Padding failed for segment keys" + ) + + # 3. Concatenate for Centralized MLP (simulating the global state vector) + global_state = jnp.concatenate( + [padded_obs["joint_position"], padded_obs["joint_velocity"], padded_obs["segment_contact"]], + axis=-1, + ) + + # 40 + 40 + 20 = 100 dimensions + assert global_state.shape == (batch_size, 100), ( + f"Expected global state shape (2, 100), got {global_state.shape}" + ) + + # 4. Initialize dummy networks (40 actuators for the max morphology output) + actor = Actor(action_dim=40) + critic = OneDenseLayerMLP() # Acts as the centralized critic + + rng = jax.random.PRNGKey(0) + rng_a, rng_c = jax.random.split(rng) + + # Initialize Flax variables + actor_params = actor.init(rng_a, global_state) + critic_params = critic.init(rng_c, global_state) + + # 5. Forward Pass Assertions + action_mean, action_log_std = actor.apply(actor_params, global_state) + value = critic.apply(critic_params, global_state) + + assert action_mean.shape == (batch_size, 40), f"Actor mean shape mismatch: {action_mean.shape}" + assert action_log_std.shape == (40,), f"Actor log_std shape mismatch: {action_log_std.shape}" + assert value.shape == (batch_size, 1) or value.shape == (batch_size,), ( + f"Critic value shape mismatch: {value.shape}" + ) + + +if __name__ == "__main__": + test_centralized_forward_pass_with_padding() From 5a358fadded334577b1d0395c6dbf62277ef6c34 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 15 Apr 2026 18:07:12 +0200 Subject: [PATCH 25/32] fix: skip test in ci that requires buffer --- tests/test_morphology_render.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_morphology_render.py b/tests/test_morphology_render.py index c127a01..d8648d8 100644 --- a/tests/test_morphology_render.py +++ b/tests/test_morphology_render.py @@ -1,4 +1,5 @@ import jax +import pytest import os import mujoco from PIL import Image @@ -6,6 +7,7 @@ from brittle_star_project.environment.env_config import EnvConfig, MorphologyCon from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper +@pytest.mark.skipif(os.getenv("CI") == "true", reason="No OpenGL display in CI") def test_render_morphologies(): base_dir = "runs/renders" os.makedirs(base_dir, exist_ok=True) From e40a4b979f859c35fd181addac5156d43e4fdb8b Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 15 Apr 2026 18:34:28 +0200 Subject: [PATCH 26/32] fix: typos and nonexistent calls --- configs/main_config.yaml | 1 + configs/ppo/smoke_test.yaml | 19 +++++++ scripts/train.py | 18 +++---- .../configs/config_architecture.py | 52 ++++++------------- .../configs/register_configs.py | 4 +- .../trainers/PPOTrainer.py | 6 +-- src/experiment_logger/unified_logger.py | 2 +- 7 files changed, 52 insertions(+), 50 deletions(-) create mode 100644 configs/ppo/smoke_test.yaml diff --git a/configs/main_config.yaml b/configs/main_config.yaml index 0583327..7e3d642 100644 --- a/configs/main_config.yaml +++ b/configs/main_config.yaml @@ -3,6 +3,7 @@ # Sub-configs are loaded from the relative directories. defaults: + - brittle_star_config - experiment: base - logging: default - ppo: default diff --git a/configs/ppo/smoke_test.yaml b/configs/ppo/smoke_test.yaml new file mode 100644 index 0000000..40e4c07 --- /dev/null +++ b/configs/ppo/smoke_test.yaml @@ -0,0 +1,19 @@ +# Fast PPO Configuration +# Lower timestep count for quick iterations/testing. + +learning_rate: 0.0005 +total_timesteps: 65536 +num_envs: 512 +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/scripts/train.py b/scripts/train.py index 4224eb9..409a196 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -23,7 +23,7 @@ def make_env(cfg: BrittleStarConfig) -> BrittleStarJaxEnvWrapper: @hydra.main(config_path="../configs", config_name="main_config", version_base="1.3") def main(dict_cfg: DictConfig): # 1. Convert DictConfig to structured dataclass - cfg: BrittleStarConfig = OmegaConf.to_object(dict_cfg) + config: BrittleStarConfig = OmegaConf.to_object(dict_cfg) # 2. Setup run metadata # Hydra changes CWD to the output directory by default. @@ -31,25 +31,25 @@ def main(dict_cfg: DictConfig): run_name = os.path.basename(run_dir) # 3. Initialize Logger - resolved_cfg_dict = OmegaConf.to_container(dict_cfg, resolve=True, throw_on_missing=True) + cfg_dict = OmegaConf.to_container(dict_cfg, resolve=True, throw_on_missing=True) init_logger( run_name=run_name, - config=resolved_cfg_dict, - project_name=cfg.logging.wandb_project_name, - entity=cfg.logging.wandb_entity, + config=cfg_dict, + project_name=config.logging.wandb_project_name, + entity=config.logging.wandb_entity, base_dir=os.path.dirname(run_dir), - use_wandb=cfg.logging.track, + use_wandb=config.logging.track, ) logger = get_logger() logger.info(f"Hydra-initialized run: {run_name}") logger.info(f"Output directory: {run_dir}") # 4. Setup Environment and Torch - env = make_env(cfg) - torch.backends.cudnn.deterministic = cfg.experiment.torch_deterministic + env = make_env(config) + torch.backends.cudnn.deterministic = config.experiment.torch_deterministic # 5. Train - pass structured config directly - ppo_trainer = PPOTrainer(cfg, env, run_dir, run_name) + ppo_trainer = PPOTrainer(config, env, run_dir, run_name) ppo_trainer.train() diff --git a/src/brittle_star_project/configs/config_architecture.py b/src/brittle_star_project/configs/config_architecture.py index bc140dc..e8cf5a5 100644 --- a/src/brittle_star_project/configs/config_architecture.py +++ b/src/brittle_star_project/configs/config_architecture.py @@ -1,5 +1,5 @@ from dataclasses import dataclass, field -from typing import List +from typing import List, Optional @dataclass @@ -18,14 +18,20 @@ class ArchitectureConfig: See docs/design/actor-critic.md for the full design rationale. """ - # Input network: maps global observation to a hidden representation. - feature_extractor: LayerConfig = field( - default_factory=lambda: LayerConfig(hidden_dims=[64, 64], activation="tanh") - ) - # Output network: maps the hidden representation to a scalar value estimate. - critic: LayerConfig = field( - default_factory=lambda: LayerConfig(hidden_dims=[], activation="tanh") - ) + name: str = "base" + + # Actor pipeline + sensor: Optional[LayerConfig] = None + propagator: Optional[LayerConfig] = None + motor: Optional[LayerConfig] = None + + # Critic pipeline + feature_extractor: Optional[LayerConfig] = None + critic: Optional[LayerConfig] = None + + # Decentralized + message_passing_steps: Optional[int] = None + topology_type: Optional[str] = None # Supported values: "ring", "fully_connected" @dataclass @@ -39,14 +45,7 @@ class CentralizedConfig(ArchitectureConfig): See docs/design/actor-critic.md for the full design rationale. """ - # Input network: maps the global observation to a hidden state. - sensor: LayerConfig = field( - default_factory=lambda: LayerConfig(hidden_dims=[64, 64], activation="tanh") - ) - # Output network: projects hidden state to (mean, log_std) over all joints. - motor: LayerConfig = field( - default_factory=lambda: LayerConfig(hidden_dims=[], activation="tanh") - ) + name: str = "centralized" @dataclass @@ -64,21 +63,4 @@ class DecentralizedConfig(ArchitectureConfig): full design rationale. """ - # Input network: local observation -> initial hidden state per node. - sensor: LayerConfig = field( - default_factory=lambda: LayerConfig(hidden_dims=[64, 64], activation="tanh") - ) - # Message-passing network: aggregates neighbour messages and updates hidden state. - propagator: LayerConfig = field( - default_factory=lambda: LayerConfig(hidden_dims=[64, 64], activation="tanh") - ) - # Output network: final hidden state -> joint offset for this node only. - motor: LayerConfig = field( - default_factory=lambda: LayerConfig(hidden_dims=[], activation="tanh") - ) - - # Number of synchronous message-passing rounds per control step. - message_passing_steps: int = 1 - # Graph topology used for neighbour connections. - # Supported values: "ring", "fully_connected". - topology_type: str = "ring" + name: str = "decentralized" diff --git a/src/brittle_star_project/configs/register_configs.py b/src/brittle_star_project/configs/register_configs.py index 697eda9..74360a3 100644 --- a/src/brittle_star_project/configs/register_configs.py +++ b/src/brittle_star_project/configs/register_configs.py @@ -29,8 +29,8 @@ def register_configs() -> None: cs.store(group="ppo", name="base_ppo", node=PPOConfig) # Architecture variants — swap via CLI: architecture=decentralized - cs.store(group="architecture", name="centralized", node=CentralizedConfig) - cs.store(group="architecture", name="decentralized", node=DecentralizedConfig) + cs.store(group="architecture", name="centralized_schema", node=CentralizedConfig) + cs.store(group="architecture", name="decentralized_schema", node=DecentralizedConfig) # Environment configs cs.store(group="morphology", name="base_morphology", node=MorphologyConfig) diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 3dc9666..1e49cc2 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -527,11 +527,11 @@ class PPOTrainer: f"ETA {eta_str}" ) - if getattr(self.config.experiment, "debug_sanity", False): - self.logger.log( + if getattr(self.cfg.experiment, "debug_sanity", False): + self.logger.info( "\n[SANITY CHECK] Successfully completed 1 epoch of data collection and gradient updates." ) - self.logger.log("[SANITY CHECK] Gradients flowed without NaN. Exiting gracefully.") + self.logger.info("[SANITY CHECK] Gradients flowed without NaN. Exiting gracefully.") break if self.logging_cfg.save_model: diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py index 4eadd5b..a02bdb1 100644 --- a/src/experiment_logger/unified_logger.py +++ b/src/experiment_logger/unified_logger.py @@ -53,7 +53,7 @@ def init_logger(**kwargs) -> "UnifiedLogger": the output directories and set up all logging backends. """ global _active_logger - logger = UnifiedLogger(_set_as_global=False, **kwargs) + logger = UnifiedLogger(**kwargs) _active_logger = logger return logger From 07ae66b8ffe998ab8e6fae86873fabb28631210b Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 15 Apr 2026 18:44:37 +0200 Subject: [PATCH 27/32] fix: rediect logs --- .gitignore | 1 + configs/README.md | 4 +++- configs/main_config.yaml | 6 ++++++ src/brittle_star_project/configs/config_experiment.py | 1 + 4 files changed, 11 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 43d51e5..83cfc3f 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,7 @@ runs/* wandb/ outputs/ multirun/ +metrics/ # Python-generated files __pycache__/ diff --git a/configs/README.md b/configs/README.md index 21f95b1..91208f2 100644 --- a/configs/README.md +++ b/configs/README.md @@ -49,4 +49,6 @@ python scripts/train.py --cfg job - **Adding a new group**: Create a subdirectory in `configs/` and register the new dataclass in `src/brittle_star_project/configs/register_configs.py`. - **Typo Catching**: If you see a `ConfigAttributeError`, check for typos in your YAML keys or CLI overrides. -- **Output Redirection**: Hydra automatically creates `outputs/` directories. On HPC, ensure `hydra.run.dir` is set to a fast scratch storage. +- **Output Redirection**: We use `experiment.base_run_dir` to configure where logs and models are stored (defaults to `runs/`). + - To change it locally: `python scripts/train.py experiment.base_run_dir=/path/to/custom/dir` + - On HPC, ensure this points to a fast scratch storage. diff --git a/configs/main_config.yaml b/configs/main_config.yaml index 7e3d642..a4efbb7 100644 --- a/configs/main_config.yaml +++ b/configs/main_config.yaml @@ -12,3 +12,9 @@ defaults: - arena: default - environment: directed_locomotion - _self_ + +hydra: + job: + chdir: True + run: + dir: ${experiment.base_run_dir}/${experiment.exp_name}/${now:%Y-%m-%d}/${now:%H-%M-%S} diff --git a/src/brittle_star_project/configs/config_experiment.py b/src/brittle_star_project/configs/config_experiment.py index 716d3b6..7409811 100644 --- a/src/brittle_star_project/configs/config_experiment.py +++ b/src/brittle_star_project/configs/config_experiment.py @@ -8,3 +8,4 @@ class ExperimentConfig: torch_deterministic: bool = True cuda: bool = True debug_sanity: bool = False + base_run_dir: str = "runs" From 6ed4ad80606ee43bfc9692f10e596f78771dc253 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 15 Apr 2026 18:48:05 +0200 Subject: [PATCH 28/32] style: ruff format & check --- .../configs/main_config.py | 2 +- .../environment/padded_obs_wrapper.py | 23 +++++++++++-------- .../trainers/PPOTrainer.py | 5 +--- src/experiment_logger/config_utils.py | 1 - src/experiment_logger/unified_logger.py | 2 -- tests/test_configs.py | 1 - tests/test_morphology_render.py | 1 - 7 files changed, 16 insertions(+), 19 deletions(-) diff --git a/src/brittle_star_project/configs/main_config.py b/src/brittle_star_project/configs/main_config.py index cd54473..2164a30 100644 --- a/src/brittle_star_project/configs/main_config.py +++ b/src/brittle_star_project/configs/main_config.py @@ -3,7 +3,7 @@ from dataclasses import dataclass, field from experiment_logger.config_logger import LoggingConfig from brittle_star_project.configs.config_experiment import ExperimentConfig from brittle_star_project.configs.config_ppo import PPOConfig -from brittle_star_project.configs.config_architecture import ArchitectureConfig, CentralizedConfig +from brittle_star_project.configs.config_architecture import ArchitectureConfig from brittle_star_project.environment.env_config import MorphologyConfig, ArenaConfig, EnvConfig diff --git a/src/brittle_star_project/environment/padded_obs_wrapper.py b/src/brittle_star_project/environment/padded_obs_wrapper.py index 8c07248..95f310b 100644 --- a/src/brittle_star_project/environment/padded_obs_wrapper.py +++ b/src/brittle_star_project/environment/padded_obs_wrapper.py @@ -5,23 +5,28 @@ a constant size regardless of how many segments are amputated. This wrapper pads the observation dictionary values with zeros using spatial insertion so that the flattened observation maintains the correct physical mapping to the neural network. """ + from __future__ import annotations from typing import Any import jax.numpy as jnp # Observation keys whose size scales with the number of joints (2 per segment). -_JOINT_SCALED_KEYS = frozenset({ - "joint_position", - "joint_velocity", - "joint_actuator_force", - "actuator_force", -}) +_JOINT_SCALED_KEYS = frozenset( + { + "joint_position", + "joint_velocity", + "joint_actuator_force", + "actuator_force", + } +) # Observation keys whose size scales with the number of segments (1 per segment). -_SEGMENT_SCALED_KEYS = frozenset({ - "segment_contact", -}) +_SEGMENT_SCALED_KEYS = frozenset( + { + "segment_contact", + } +) def compute_padding_masks( diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 1e49cc2..b3321b2 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -528,10 +528,7 @@ class PPOTrainer: ) if getattr(self.cfg.experiment, "debug_sanity", False): - self.logger.info( - "\n[SANITY CHECK] Successfully completed 1 epoch of data collection and gradient updates." - ) - self.logger.info("[SANITY CHECK] Gradients flowed without NaN. Exiting gracefully.") + self.logger.info("\n[SANITY CHECK] Successfully completed 1 epoch") break if self.logging_cfg.save_model: diff --git a/src/experiment_logger/config_utils.py b/src/experiment_logger/config_utils.py index 6017041..80d1748 100644 --- a/src/experiment_logger/config_utils.py +++ b/src/experiment_logger/config_utils.py @@ -1,7 +1,6 @@ """Configuration utilities for loading YAML configs and merging with CLI args.""" import os -import sys from typing import Dict, Any, Type, TypeVar import yaml from dataclasses import fields, is_dataclass diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py index a02bdb1..2f5d06f 100644 --- a/src/experiment_logger/unified_logger.py +++ b/src/experiment_logger/unified_logger.py @@ -6,9 +6,7 @@ This logger ensures all experimental data is preserved by writing to: 3. stdout (for real-time monitoring) """ -import datetime import logging -import subprocess import yaml import sys import time diff --git a/tests/test_configs.py b/tests/test_configs.py index 40ff35e..84a7cde 100644 --- a/tests/test_configs.py +++ b/tests/test_configs.py @@ -1,4 +1,3 @@ -import pytest from hydra import compose, initialize from omegaconf import OmegaConf from brittle_star_project.configs.main_config import BrittleStarConfig diff --git a/tests/test_morphology_render.py b/tests/test_morphology_render.py index d8648d8..423a575 100644 --- a/tests/test_morphology_render.py +++ b/tests/test_morphology_render.py @@ -1,4 +1,3 @@ -import jax import pytest import os import mujoco From 143dbbc710ba90496491ad9eaeb045f4ac1f9438 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 15 Apr 2026 19:26:33 +0200 Subject: [PATCH 29/32] fix: explicit arms length check --- configs/architecture/centralized.yaml | 1 + configs/architecture/decentralized.yaml | 1 + .../environment/padded_obs_wrapper.py | 14 +++++++++++++- tests/test_configs.py | 5 +++-- 4 files changed, 18 insertions(+), 3 deletions(-) diff --git a/configs/architecture/centralized.yaml b/configs/architecture/centralized.yaml index 62ff16d..3b7783d 100644 --- a/configs/architecture/centralized.yaml +++ b/configs/architecture/centralized.yaml @@ -4,6 +4,7 @@ # Default values are defined in CentralizedConfig dataclass. # Use this configuration for standard PPO experiments. +name: "centralized" sensor: hidden_dims: [64, 64] activation: "tanh" diff --git a/configs/architecture/decentralized.yaml b/configs/architecture/decentralized.yaml index d87e144..fbf07bc 100644 --- a/configs/architecture/decentralized.yaml +++ b/configs/architecture/decentralized.yaml @@ -4,6 +4,7 @@ # Default values are defined in DecentralizedConfig dataclass. # Use this configuration for decentralized execution experiments. +name: "decentralized" sensor: hidden_dims: [64, 64] activation: "tanh" diff --git a/src/brittle_star_project/environment/padded_obs_wrapper.py b/src/brittle_star_project/environment/padded_obs_wrapper.py index 95f310b..3f22038 100644 --- a/src/brittle_star_project/environment/padded_obs_wrapper.py +++ b/src/brittle_star_project/environment/padded_obs_wrapper.py @@ -42,10 +42,22 @@ def compute_padding_masks( Returns: A dict containing 1D boolean masks and target sizes. """ + if len(segments_per_arm) != len(reference_segments_per_arm): + raise ValueError( + f"Morphology mismatch: current has {len(segments_per_arm)} arms, " + f"but reference requires {len(reference_segments_per_arm)} arms." + ) + mask_1x = [] mask_2x = [] - for actual, ref in zip(segments_per_arm, reference_segments_per_arm): + for arm_idx, (actual, ref) in enumerate(zip(segments_per_arm, reference_segments_per_arm)): + if not (0 <= actual <= ref): + raise ValueError( + f"Invalid amputation at arm {arm_idx}: " + f"actual segments ({actual}) must be between 0 and reference ({ref})." + ) + # 1x scaling (e.g., contacts: 1 value per segment) # 1x scaling (e.g., contacts: 1 value per segment) mask_1x.extend([True] * actual + [False] * (ref - actual)) # 2x scaling (e.g., joints: 2 values per segment) diff --git a/tests/test_configs.py b/tests/test_configs.py index 84a7cde..c8cad63 100644 --- a/tests/test_configs.py +++ b/tests/test_configs.py @@ -18,7 +18,8 @@ def test_config_composition_centralized(): structured_cfg = OmegaConf.to_object(OmegaConf.merge(BrittleStarConfig, cfg)) # Basic assertions - assert "CentralizedConfig" in str(type(structured_cfg.architecture)) + assert structured_cfg.architecture.name == "centralized" + assert structured_cfg.architecture.propagator is None assert isinstance(structured_cfg.ppo.learning_rate, float) assert structured_cfg.ppo.learning_rate > 0 @@ -32,7 +33,7 @@ def test_config_composition_decentralized(): structured_cfg = OmegaConf.to_object(OmegaConf.merge(BrittleStarConfig, cfg)) # Basic assertions - assert "DecentralizedConfig" in str(type(structured_cfg.architecture)) + assert structured_cfg.architecture.name == "decentralized" assert isinstance(structured_cfg.ppo.learning_rate, float) assert structured_cfg.ppo.learning_rate > 0 From a72471611771af206abc69e8f55f42334119a2bb Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 15 Apr 2026 19:39:39 +0200 Subject: [PATCH 30/32] feat: simulation config --- configs/main_config.yaml | 1 + configs/simulation/default.yaml | 11 ++++++++++ scripts/simulate.py | 22 +++++++++++-------- scripts/train.py | 6 +++-- .../configs/config_simulation.py | 12 ++++++++++ .../configs/main_config.py | 2 ++ .../configs/register_configs.py | 2 ++ tests/test_configs.py | 11 ++++++---- 8 files changed, 52 insertions(+), 15 deletions(-) create mode 100644 configs/simulation/default.yaml create mode 100644 src/brittle_star_project/configs/config_simulation.py diff --git a/configs/main_config.yaml b/configs/main_config.yaml index a4efbb7..745755a 100644 --- a/configs/main_config.yaml +++ b/configs/main_config.yaml @@ -11,6 +11,7 @@ defaults: - morphology: 5_arms_full - arena: default - environment: directed_locomotion + - simulation: default - _self_ hydra: diff --git a/configs/simulation/default.yaml b/configs/simulation/default.yaml new file mode 100644 index 0000000..00c4e44 --- /dev/null +++ b/configs/simulation/default.yaml @@ -0,0 +1,11 @@ +# Default Simulation Settings +# These values are used by scripts/simulate.py + +# Path to the trained model (optional) +model_path: null + +# Type of model to use if no path is provided (e.g., random) +model_type: "random" + +# Execution backend (MJX or BRAX) +backend: "MJX" diff --git a/scripts/simulate.py b/scripts/simulate.py index 63b1cfe..8fb597d 100644 --- a/scripts/simulate.py +++ b/scripts/simulate.py @@ -29,19 +29,23 @@ MODEL_OPTIONS = sorted(MODEL_BY_NAME) @hydra.main(config_path="../configs", config_name="main_config", version_base="1.3") def main(dict_cfg: DictConfig) -> None: - cfg: BrittleStarConfig = OmegaConf.to_object(dict_cfg) + # 1. Convert DictConfig to structured dataclass, ensuring the root schema is applied correctly. + config: BrittleStarConfig = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg) + ) - # Allow overriding these via Hydra CLI or a dedicated simulate config group - # For now, defaults matching the old argparse behavior - backend = Backend.MJX - model_type = "random" - model_path = None - seed = cfg.experiment.seed + # Use the configurable settings from the simulation group + backend = config.simulation.backend + model_type = config.simulation.model_type + model_path = config.simulation.model_path + seed = config.experiment.seed # ======= ENVIRONMENT SETUP ======= factory = BrittleStarEnvFactory() - raw_env = factory.create_environment(backend, cfg.morphology, cfg.arena, cfg.environment) - env = BrittleStarEnv(raw_env, backend=backend, config=cfg.environment) + raw_env = factory.create_environment( + backend, config.morphology, config.arena, config.environment + ) + env = BrittleStarEnv(raw_env, backend=backend, config=config.environment) state = env.reset(seed=seed) diff --git a/scripts/train.py b/scripts/train.py index 409a196..409f050 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -22,8 +22,10 @@ def make_env(cfg: BrittleStarConfig) -> BrittleStarJaxEnvWrapper: @hydra.main(config_path="../configs", config_name="main_config", version_base="1.3") def main(dict_cfg: DictConfig): - # 1. Convert DictConfig to structured dataclass - config: BrittleStarConfig = OmegaConf.to_object(dict_cfg) + # 1. Convert DictConfig to structured dataclass, ensuring the root schema is applied correctly. + config: BrittleStarConfig = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg) + ) # 2. Setup run metadata # Hydra changes CWD to the output directory by default. diff --git a/src/brittle_star_project/configs/config_simulation.py b/src/brittle_star_project/configs/config_simulation.py new file mode 100644 index 0000000..56747ae --- /dev/null +++ b/src/brittle_star_project/configs/config_simulation.py @@ -0,0 +1,12 @@ +from dataclasses import dataclass +from typing import Optional +from brittle_star_project.environment.env_types import Backend + + +@dataclass +class SimulationSettings: + """Settings for the simulation script.""" + + model_path: Optional[str] = None + model_type: str = "random" + backend: Backend = Backend.MJX diff --git a/src/brittle_star_project/configs/main_config.py b/src/brittle_star_project/configs/main_config.py index 2164a30..5a937a7 100644 --- a/src/brittle_star_project/configs/main_config.py +++ b/src/brittle_star_project/configs/main_config.py @@ -4,6 +4,7 @@ from experiment_logger.config_logger import LoggingConfig from brittle_star_project.configs.config_experiment import ExperimentConfig from brittle_star_project.configs.config_ppo import PPOConfig from brittle_star_project.configs.config_architecture import ArchitectureConfig +from brittle_star_project.configs.config_simulation import SimulationSettings from brittle_star_project.environment.env_config import MorphologyConfig, ArenaConfig, EnvConfig @@ -24,3 +25,4 @@ class BrittleStarConfig: morphology: MorphologyConfig = field(default_factory=MorphologyConfig) arena: ArenaConfig = field(default_factory=ArenaConfig) environment: EnvConfig = field(default_factory=EnvConfig) + simulation: SimulationSettings = field(default_factory=SimulationSettings) diff --git a/src/brittle_star_project/configs/register_configs.py b/src/brittle_star_project/configs/register_configs.py index 74360a3..b677b31 100644 --- a/src/brittle_star_project/configs/register_configs.py +++ b/src/brittle_star_project/configs/register_configs.py @@ -7,6 +7,7 @@ from brittle_star_project.configs.config_architecture import ( CentralizedConfig, DecentralizedConfig, ) +from brittle_star_project.configs.config_simulation import SimulationSettings from brittle_star_project.environment.env_config import MorphologyConfig, ArenaConfig, EnvConfig from brittle_star_project.configs.main_config import BrittleStarConfig @@ -36,3 +37,4 @@ def register_configs() -> None: cs.store(group="morphology", name="base_morphology", node=MorphologyConfig) cs.store(group="arena", name="base_arena", node=ArenaConfig) cs.store(group="environment", name="base_environment", node=EnvConfig) + cs.store(group="simulation", name="base_simulation", node=SimulationSettings) diff --git a/tests/test_configs.py b/tests/test_configs.py index c8cad63..43ef95e 100644 --- a/tests/test_configs.py +++ b/tests/test_configs.py @@ -13,9 +13,10 @@ def test_config_composition_centralized(): # We compose the config; it follows main_config.yaml cfg = compose(config_name="main_config", overrides=["architecture=centralized"]) - # Merge with the dataclass class to get a structured DictConfig, - # then convert to a real dataclass instance to verify validation. - structured_cfg = OmegaConf.to_object(OmegaConf.merge(BrittleStarConfig, cfg)) + # Merge with the structured schema and convert to a real dataclass instance to verify validation. + structured_cfg = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), cfg) + ) # Basic assertions assert structured_cfg.architecture.name == "centralized" @@ -30,7 +31,9 @@ def test_config_composition_decentralized(): cfg = compose(config_name="main_config", overrides=["architecture=decentralized"]) # Merge and convert to dataclass instance - structured_cfg = OmegaConf.to_object(OmegaConf.merge(BrittleStarConfig, cfg)) + structured_cfg = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), cfg) + ) # Basic assertions assert structured_cfg.architecture.name == "decentralized" From a83a4bedbeee531adc91c79d34ae5c3c52dcdd88 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 15 Apr 2026 20:09:11 +0200 Subject: [PATCH 31/32] fix: cleanup --- scripts/simulate.py | 6 +++++- .../environment/BrittleStarJaxEnvWrapper.py | 20 +++++++++++++++++-- tests/test_configs.py | 11 ++++++---- tests/test_morphology_render.py | 6 ++++++ 4 files changed, 36 insertions(+), 7 deletions(-) diff --git a/scripts/simulate.py b/scripts/simulate.py index 8fb597d..666d048 100644 --- a/scripts/simulate.py +++ b/scripts/simulate.py @@ -12,7 +12,6 @@ from omegaconf import DictConfig, OmegaConf from pathlib import Path from brittle_star_project import ( - Backend, BrittleStarEnv, BrittleStarEnvFactory, SimulationConfig, @@ -37,7 +36,12 @@ def main(dict_cfg: DictConfig) -> None: # Use the configurable settings from the simulation group backend = config.simulation.backend model_type = config.simulation.model_type + + # Hydra chdir changes CWD; we map CLI relative paths relative to invocation originally. model_path = config.simulation.model_path + if model_path is not None: + model_path = hydra.utils.to_absolute_path(model_path) + seed = config.experiment.seed # ======= ENVIRONMENT SETUP ======= diff --git a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py index 7a6d161..538ba7a 100644 --- a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py +++ b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py @@ -28,6 +28,10 @@ class BrittleStarJaxEnvWrapper: self._backend, self._morphology, self._arena, self._env_config ) + from brittle_star_project.environment.padded_obs_wrapper import compute_padding_masks + + 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)) self._vectorized_action_sample = jax.jit(jax.vmap(self._env.action_space.sample)) @@ -61,7 +65,13 @@ class BrittleStarJaxEnvWrapper: self.logger.info(f"Resetting vectorized environment environments with seed {seed}") self._action_rng, env_rng = jax.random.split(jax.random.PRNGKey(seed), 2) env_rngs = jnp.array(jax.random.split(env_rng, self._num_envs)) - return self._vectorized_reset(rng=env_rngs) + state = self._vectorized_reset(rng=env_rngs) + from brittle_star_project.environment.padded_obs_wrapper import pad_observations_batched + + state = state.replace( + observations=pad_observations_batched(state.observations, self._padding_masks) + ) + return state def sample_actions(self): assert self._action_rng is not None, "Call reset() before sample_actions()" @@ -71,7 +81,13 @@ class BrittleStarJaxEnvWrapper: return self._vectorized_action_sample(rng=jnp.array(sub_rngs)) def step(self, state, action): - return self._vectorized_step(state=state, action=action) + next_state = self._vectorized_step(state=state, action=action) + from brittle_star_project.environment.padded_obs_wrapper import pad_observations_batched + + next_state = next_state.replace( + observations=pad_observations_batched(next_state.observations, self._padding_masks) + ) + return next_state def close(self): self._env.close() diff --git a/tests/test_configs.py b/tests/test_configs.py index 43ef95e..0a6a651 100644 --- a/tests/test_configs.py +++ b/tests/test_configs.py @@ -1,4 +1,5 @@ -from hydra import compose, initialize +from pathlib import Path +from hydra import compose, initialize_config_dir from omegaconf import OmegaConf from brittle_star_project.configs.main_config import BrittleStarConfig from brittle_star_project.configs.register_configs import register_configs @@ -9,11 +10,12 @@ register_configs() def test_config_composition_centralized(): """Test that the centralized configuration composes and validates correctly.""" - with initialize(version_base="1.3", config_path="../configs"): + config_dir = str(Path(__file__).parent.parent / "configs") + with initialize_config_dir(version_base="1.3", config_dir=config_dir): # We compose the config; it follows main_config.yaml cfg = compose(config_name="main_config", overrides=["architecture=centralized"]) - # Merge with the structured schema and convert to a real dataclass instance to verify validation. + # Merge with the structured schema and convert to a real dataclass instance structured_cfg = OmegaConf.to_object( OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), cfg) ) @@ -27,7 +29,8 @@ def test_config_composition_centralized(): def test_config_composition_decentralized(): """Test that the decentralized configuration composes and validates correctly.""" - with initialize(version_base="1.3", config_path="../configs"): + config_dir = str(Path(__file__).parent.parent / "configs") + with initialize_config_dir(version_base="1.3", config_dir=config_dir): cfg = compose(config_name="main_config", overrides=["architecture=decentralized"]) # Merge and convert to dataclass instance diff --git a/tests/test_morphology_render.py b/tests/test_morphology_render.py index 423a575..6b435cf 100644 --- a/tests/test_morphology_render.py +++ b/tests/test_morphology_render.py @@ -1,5 +1,11 @@ import pytest import os +import sys + +# CRITICAL for headless cross-platform testing (devcontainers etc) +if sys.platform == "linux" and "DISPLAY" not in os.environ and "WAYLAND_DISPLAY" not in os.environ: + os.environ.setdefault("MUJOCO_GL", "egl") + import mujoco from PIL import Image from brittle_star_project.environment.env_config import EnvConfig, MorphologyConfig, ArenaConfig From a33b8a2e8e5198c6af97bc410388481f1fe3a3b2 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Thu, 16 Apr 2026 10:26:26 +0200 Subject: [PATCH 32/32] refactor: top level imports Co-authored-by: Robin Meersman --- .../environment/BrittleStarJaxEnvWrapper.py | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) diff --git a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py index 538ba7a..a5175a7 100644 --- a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py +++ b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py @@ -1,13 +1,11 @@ import jax import jax.numpy as jnp -from brittle_star_project import ( - EnvConfig, - BrittleStarEnvFactory, - MorphologyConfig, - ArenaConfig, - Backend, -) +from experiment_logger import get_logger +from .env_config import EnvConfig, MorphologyConfig, ArenaConfig +from .env_types import Backend +from .factory import BrittleStarEnvFactory +from .padded_obs_wrapper import compute_padding_masks, pad_observations_batched class BrittleStarJaxEnvWrapper: @@ -28,8 +26,7 @@ class BrittleStarJaxEnvWrapper: self._backend, self._morphology, self._arena, self._env_config ) - from brittle_star_project.environment.padded_obs_wrapper import compute_padding_masks - + # Pre-compute masks for observation padding self._padding_masks = compute_padding_masks(self._morphology.segments_per_arm) self._vectorized_reset = jax.jit(jax.vmap(self._env.reset)) @@ -38,8 +35,6 @@ class BrittleStarJaxEnvWrapper: self._action_rng = None - from experiment_logger import get_logger - self.logger = get_logger() self.logger.info( f"Initialized BrittleStarJaxEnvWrapper with {num_envs} envs on {backend.value}" @@ -66,7 +61,6 @@ class BrittleStarJaxEnvWrapper: self._action_rng, env_rng = jax.random.split(jax.random.PRNGKey(seed), 2) env_rngs = jnp.array(jax.random.split(env_rng, self._num_envs)) state = self._vectorized_reset(rng=env_rngs) - from brittle_star_project.environment.padded_obs_wrapper import pad_observations_batched state = state.replace( observations=pad_observations_batched(state.observations, self._padding_masks) @@ -82,7 +76,6 @@ class BrittleStarJaxEnvWrapper: def step(self, state, action): next_state = self._vectorized_step(state=state, action=action) - from brittle_star_project.environment.padded_obs_wrapper import pad_observations_batched next_state = next_state.replace( observations=pad_observations_batched(next_state.observations, self._padding_masks)