From 3cb3c765cc6d3a06996ffa4a1150d9cdfa1d4581 Mon Sep 17 00:00:00 2001 From: Robin Meersman Date: Sat, 2 May 2026 23:27:31 +0200 Subject: [PATCH] feat(message passing): fixed tracing bug, todo: dense layer input dimension incorrect --- src/brittle_star_project/trainers/PPOTrainer.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 72c4734..a297a07 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -134,8 +134,8 @@ def _normalize_obs(obs, mean, var, eps=1e-8): def _convert_obs_dict_to_array_morphology(obs_dict, morph_mode, segments_per_arm: jnp.ndarray): - num_segments = segments_per_arm.sum() - num_arms = jnp.where(segments_per_arm > 0, 1, 0).sum() + num_segments = int(segments_per_arm.sum()) + num_arms = int(jnp.where(segments_per_arm > 0, 1, 0).sum()) @logged_jit def _filter_and_flatten(o) -> jnp.ndarray: @@ -668,9 +668,11 @@ class PPOTrainer: message_passer_params = {} if self.morph_mode != MorphMode.CENTRALIZED: + assert self.message_passer is not None, "MessagePasser is None" + message_passer_params = jax.vmap( lambda k: self.message_passer.init( - k, self.sensor.apply(single_sensor_param, sample_obs) + k, self.sensor.apply(single_sensor_param, sample_obs), self.adj ) )(message_passer_keys) self.logger.info(