fix(ppo): added message_passer to ppo

src/brittle_star_project/MLPs/mlps.py

@@ -40,16 +40,17 @@ class Actor(nn.Module):
 class MessagePasser(nn.Module):
     hidden_dim: int
     num_propagation_steps: int
+    adj_matrix: jnp.ndarray
 
     @nn.compact
-    def __call__(self, x: jnp.ndarray, adj_matrix: jnp.ndarray):
+    def __call__(self, x: jnp.ndarray):
         for _ in range(self.num_propagation_steps):
             # (n_nodes, feat)
             messages = nn.Dense(self.hidden_dim)(x)
             messages = nn.tanh(messages)
 
             # note: if mean is wanted: adj_matrix / (adj.sum(axis=-1, keepdims=True) + 1e-8)
-            agg = adj_matrix
+            agg = self.adj_matrix
             aggregated = agg @ messages
 
             x_concat = jnp.concatenate([x, aggregated], axis=-1)

src/brittle_star_project/environment/obs_processing.py

@@ -72,6 +72,10 @@ def create_obs_processor(
                 padded[key] = arr
         return padded
 
+    # TODO
+    def _split_to_agents() -> dict:
+        return {}
+
     def _flatten_features(obs: dict) -> jnp.ndarray:
         ordered_keys = [
             "disk_z_tilt",
@@ -93,7 +97,7 @@ def create_obs_processor(
             if v.size == 0:
                 continue
 
-            # reshape scalars
+            # reshape scalars to 1D array
             if v.ndim == 0:
                 v = v.reshape(1)
 
@@ -102,19 +106,8 @@ def create_obs_processor(
                 values.append(v.reshape(1, -1))  # (1, feat)
                 continue
 
-            # -------- SPLIT TO SEGMENTS --------
+            # -------- SCALE WITH SEGMENTS --------
             if key in _JOINT_SCALED_KEYS:
-                if morph_mode == MorphMode.SEGMENT:
-                    joint_count = 3
-                    axis_per_joint = 2
-
-                    center_size = num_arms * joint_count * axis_per_joint
-                    v_center = v[:center_size].reshape(
-                        num_arms, joint_count * axis_per_joint
-                    )  # (arms, 6)
-                    v_segs = v[center_size:].reshape(-1, 2)  # (segs, 2)
-                    values.append(jnp.concatenate([v_center, v_segs], axis=0))  # (arms+segs, ?)
-                    continue
                 v = v.reshape(num_arms, -1)  # (n_arms, 2)
 
             elif key in _SEGMENT_SCALED_KEYS:
@@ -122,25 +115,20 @@ def create_obs_processor(
 
             else:
                 # global key, broadcast to all nodes
-                n_nodes = (num_segments + num_arms) if morph_mode == MorphMode.SEGMENT else num_arms
+                v = jnp.repeat(v[None, :], num_arms, axis=0)  # (num_arms, feat)
-                v = jnp.repeat(v[None, :], n_nodes, axis=0)  # (n_nodes, feat)
 
-            # -------- SEGMENT MODE --------
+            # RING + FULLY CONNECTED
-            if morph_mode == MorphMode.SEGMENT:
-                values.append(v)  # (n_nodes, feat)
-                continue
-
-            # -------- ARM MODE --------
             v = v.reshape(num_arms, -1)
             values.append(v)  # (n_arms, feat)
 
-        return jnp.concatenate(values, axis=-1)
+        return jnp.concatenate(values, axis=-1)  # (agent, feat)
 
     def _process_single(obs_dict: dict) -> jnp.ndarray:
-        processed = _add_derived_features(obs_dict)
+        processed = _add_derived_features(obs_dict)  # key |--> (feat-count, feat-lengths, )
-        processed = _normalize_features(processed)
+        processed = _normalize_features(processed)  # key |--> (feat-count, feat-lengths, )
+        # TODO: split to agents # key |--> (agents, feat-count, feat-lengths)
         if padding_masks is not None:
-            processed = _pad_features(processed)
+            processed = _pad_features(processed)  # key |--> (agents, feat-count, feat-lengths')
-        return _flatten_features(processed)
+        return _flatten_features(processed)  # (agents, feat)
 
     return jax.jit(jax.vmap(_process_single))

src/brittle_star_project/ppo.py

@@ -119,7 +119,7 @@ def get_action_and_value(
 ):
     hidden_sensor = sensor_apply(params["sensor_params"], x)
     hidden_critic = feature_extractor_apply(params["feature_extractor_params"], x)
-    hidden_sensor = message_passer(hidden_sensor)
+    hidden_sensor = message_passer(params["message_passer_params"], hidden_sensor)
 
     debug.callback(logger.debug, f"[SHAPE] hidden_sensor: {hidden_sensor.shape}")
     debug.callback(logger.debug, f"[SHAPE] hidden_critic: {hidden_critic.shape}")

src/brittle_star_project/trainers/PPOTrainer.py

@@ -62,7 +62,6 @@ def _get_action_and_value_noise(
     key,
     action_low,
     action_high,
-    adj_matrix: jnp.ndarray,
 ):
     # (B, n_nodes, feat)
     hidden = apply_per_node(sensor, agent_state.params["sensor_params"], next_obs)
@@ -70,7 +69,7 @@ def _get_action_and_value_noise(
     if message_passer is not None:
         params = agent_state.params["message_passer_params"]
         # (n_nodes, feat) --> let each node talk with its neighbours ==> vmap over B dimension
-        hidden = jax.vmap(lambda x: message_passer.apply(params, x, adj_matrix))(hidden)
+        hidden = jax.vmap(lambda x: message_passer.apply(params, x))(hidden)
 
     hidden_critic = apply_shared(
         feature_extractor, agent_state.params["feature_extractor_params"], next_obs
@@ -100,7 +99,6 @@ def _step_once(
     carry,
     _,
     env_step_fn,
-    adj_matrix,
     sensor: nn.Module,
     feature_extractor: nn.Module,
     actor: nn.Module,
@@ -121,7 +119,6 @@ def _step_once(
         key,
         action_low,
         action_high,
-        adj_matrix,
     )
     logger11.debug(f"[_step_once] raw_action: {raw_action.shape}")
     logger11.debug(f"[_step_once] clipped_action: {flat_clipped_action.shape}")
@@ -242,7 +239,6 @@ def _rollout_jit(
     message_passer: Optional[nn.Module],
     action_low,
     action_high,
-    adj_matrix,
 ):
     (agent_state, episode_stats, next_obs, next_done, key, env_state), storage = jax.lax.scan(
         partial(
@@ -255,7 +251,6 @@ def _rollout_jit(
             env_step_fn=step_env_fn,
             action_low=action_low,
             action_high=action_high,
-            adj_matrix=adj_matrix,
         ),
         (agent_state, episode_stats, next_obs, next_done, key, env_state),
         (),
@@ -393,7 +388,6 @@ class PPOTrainer:
                 message_passer=self.message_passer,
                 action_low=action_low,
                 action_high=action_high,
-                adj_matrix=self.adj,
             )
         )
         self._compute_gae_jit = logged_jit(
@@ -419,7 +413,13 @@ class PPOTrainer:
         def apply_feature(p, x):
             return apply_shared(self.feature_extractor, p, x)
 
-        self._ppo = PPO(self.ppo, apply_sensor, apply_actor, apply_critic, apply_feature)
+        def apply_message_passer(p, x):
+            assert self.message_passer is not None
+            return jax.vmap(lambda x_in: self.message_passer.apply(p, x_in))(x)
+
+        self._ppo = PPO(
+            self.ppo, apply_sensor, apply_actor, apply_critic, apply_feature, apply_message_passer
+        )
 
         self.agent_state = self._init_agent_state()
 
@@ -453,6 +453,7 @@ class PPOTrainer:
             MessagePasser(
                 hidden_dim=300,
                 num_propagation_steps=self.cfg.architecture.message_passing_steps or 4,
+                adj_matrix=self.adj,
             )
             if self.morph_mode != MorphMode.CENTRALIZED
             else None
@@ -517,7 +518,6 @@ class PPOTrainer:
             message_passer_params = self.message_passer.init(
                 message_passer_key,
                 self.sensor.apply(single_sensor_param, sample_obs),
-                self.adj,
             )
         self.logger.debug(
             f"[_init_agent_state] message_passer_params: {