From 22c43f05b1a2388fcc13cabc5f1de65b5f62a329 Mon Sep 17 00:00:00 2001 From: JibrilExe Date: Fri, 1 May 2026 10:40:02 +0200 Subject: [PATCH] fix: critic init now fixed? --- scripts/hpc/train.pbs | 2 +- src/brittle_star_project/ppo.py | 10 ++--- .../trainers/PPOTrainer.py | 45 ++++++++++++++----- 3 files changed, 41 insertions(+), 16 deletions(-) diff --git a/scripts/hpc/train.pbs b/scripts/hpc/train.pbs index 0f95b18..a0e79cc 100644 --- a/scripts/hpc/train.pbs +++ b/scripts/hpc/train.pbs @@ -66,7 +66,7 @@ fi # Run training using Hydra overrides python scripts/train.py \ hydra.run.dir="$SCRATCH_RUNDIR" \ - ppo=stable \ + ppo=smoke_test \ logging=hpc echo ">>> Staging out results to $DATA_RUNDIR..." diff --git a/src/brittle_star_project/ppo.py b/src/brittle_star_project/ppo.py index 1ea37a8..0b1cd95 100644 --- a/src/brittle_star_project/ppo.py +++ b/src/brittle_star_project/ppo.py @@ -8,7 +8,7 @@ import jax.numpy as jnp # Chose to use a class as it seemed the easiest way to integrate the CleanRL code style # with our need to seperate concerns class PPO: - def __init__(self, args, sensor, actor, critic, feature_extractor, message_passer=None): + def __init__(self, args, sensor_apply, actor_apply, critic_apply, feature_extractor_apply, message_passer=None): self.args = args if not message_passer: @@ -18,10 +18,10 @@ class PPO: partial( ppo_loss, args=args, - sensor_apply=sensor.apply, - actor_apply=actor.apply, - critic_apply=critic.apply, - feature_extractor_apply=feature_extractor.apply, + sensor_apply=sensor_apply, + actor_apply=actor_apply, + critic_apply=critic_apply, + feature_extractor_apply=feature_extractor_apply, message_passer=message_passer, ), has_aux=True, diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index e027a6a..502dd6d 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -217,7 +217,7 @@ def _get_action_and_value_noise( adj_matrix: jnp.ndarray, ): hidden = apply_per_node(sensor, agent_state.params["sensor_params"], next_obs) - hidden_critic = apply_per_node( + hidden_critic = apply_shared( feature_extractor, agent_state.params["feature_extractor_params"], next_obs ) @@ -226,11 +226,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) + raw_action = mean + noise * std 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) + + logprob = -0.5 * (((raw_action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1) + value = apply_shared(critic, agent_state.params["critic_params"], hidden_critic) + + raw_action = raw_action.reshape(-1) + raw_action = raw_action.reshape(raw_action.shape[0], -1) # concat the per agent, keep the envs dim + clipped_action = _clip_action(raw_action, action_low, action_high) return clipped_action, raw_action, logprob, value.squeeze(-1), mean, std, key @@ -333,11 +339,22 @@ def _step_env_wrapped(episode_stats, env_state, action, env_step_fn, morph_mode, ), ) - def apply_per_node(net, params, x): + # params: (nodes, ...) # x: (batch, nodes, feat) - return jax.vmap(net.apply, in_axes=(0, 1), out_axes=1)(params, x) + def apply_single_node(p, x_node): + # x_node: (batch, feat) + return jax.vmap(lambda xi: net.apply(p, xi))(x_node) + + return jax.vmap(apply_single_node, in_axes=(0, 1), out_axes=1)(params, x) + +def apply_shared(net, params, x): + # x: (batch, nodes, feat) + # If the critic expects a single vector per environment: + batch_size = x.shape[0] + x_flattened = x.reshape(batch_size, -1) + return jax.vmap(lambda xi: net.apply(params, xi))(x_flattened) # TODO: update to work with extra dimension + message passing def _rollout_jit( @@ -399,9 +416,9 @@ def _compute_gae_jit( critic, adj_matrix: jnp.ndarray, ): - next_value = critic.apply( + next_value = apply_shared(critic, agent_state.params["critic_params"], - feature_extractor.apply(agent_state.params["feature_extractor_params"], next_obs), + apply_shared(feature_extractor,agent_state.params["feature_extractor_params"], next_obs), ).squeeze(-1) advantages = jnp.zeros((num_envs,)) @@ -511,7 +528,12 @@ class PPOTrainer: ) ) - self._ppo = PPO(self.ppo, self.sensor, self.actor, self.critic, self.feature_extractor) + apply_sensor = lambda p, x: apply_per_node(self.sensor, p, x) + apply_actor = lambda p, x: apply_per_node(self.actor, p, x) + apply_critic = lambda p, x: apply_shared(self.critic, p, x) + apply_feature = lambda p, x: apply_shared(self.feature_extractor, p, x) + + self._ppo = PPO(self.ppo, apply_sensor, apply_actor, apply_critic, apply_feature) self.agent_state = self._init_agent_state() @@ -583,9 +605,12 @@ class PPOTrainer: ) )(message_passer_keys) - feature_extractor_params = self.feature_extractor.init(feature_extractor_key, sample_obs) + flat_obs = sample_obs.reshape(-1) # BECAUSE 1 centralized critic + feature_extractor_params = self.feature_extractor.init(feature_extractor_key, flat_obs) + critic_params = self.critic.init( - critic_key, self.feature_extractor.apply(feature_extractor_params, sample_obs) + critic_key, + self.feature_extractor.apply(feature_extractor_params, flat_obs) ) return TrainState.create(