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

fix(message passing): forgot vmapping over params :thumbs_up:

Browse source
This commit is contained in:
Robin Meersman 2026-05-02 23:34:52 +02:00
parent 3cb3c765cc
commit bb14507fc7
1 changed files with 5 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -220,16 +220,18 @@ def _get_action_and_value_noise(
action_high, action_high,
adj_matrix: jnp.ndarray, adj_matrix: jnp.ndarray,
): ):
def apply_message_passing(h): def apply_message_passing(p, h):
assert message_passer is not None, ( assert message_passer is not None, (
"message passing shouldn't be performed if morphology mode is Centralized" "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) hidden = apply_per_node(sensor, agent_state.params["sensor_params"], next_obs)
if message_passer is not None: 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( hidden_critic = apply_shared(
feature_extractor, agent_state.params["feature_extractor_params"], next_obs 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) 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): 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) next_env_state = env_step_fn(env_state, action)