From bb14507fc7206f103404bf65671499ed0215d1da Mon Sep 17 00:00:00 2001 From: Robin Meersman Date: Sat, 2 May 2026 23:34:52 +0200 Subject: [PATCH] fix(message passing): forgot vmapping over params :thumbs_up: --- src/brittle_star_project/trainers/PPOTrainer.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index a297a07..7e11b29 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -220,16 +220,18 @@ def _get_action_and_value_noise( action_high, adj_matrix: jnp.ndarray, ): - def apply_message_passing(h): + def apply_message_passing(p, h): assert message_passer is not None, ( "message passing shouldn't be performed if morphology mode is Centralized" ) - return message_passer.apply(agent_state.params["message_passer_params"], h, adj_matrix) + return jax.vmap(lambda x: message_passer.apply(p, x, adj_matrix))(h) hidden = apply_per_node(sensor, agent_state.params["sensor_params"], next_obs) if message_passer is not None: - hidden = jax.vmap(apply_message_passing)(hidden) + hidden = jax.vmap(apply_message_passing)( + agent_state.params["message_passer_params"], hidden + ) hidden_critic = apply_shared( feature_extractor, agent_state.params["feature_extractor_params"], next_obs @@ -334,7 +336,6 @@ def _reward_fn(env_state, next_env_state): return jnp.where(next_env_state.terminated, 50.0, clipped_env_reward - penalty) -# TODO: message passing def _step_env_wrapped(episode_stats, env_state, action, env_step_fn, morph_mode, segments_per_arm): next_env_state = env_step_fn(env_state, action)