refactor: lift apply_per_node

scripts/evaluate_checkpoints.py

@@ -26,7 +26,7 @@ from brittle_star_project.MLPs.mlps import (
 )
 from brittle_star_project.MLPs.adjancency_builder import build_adjacency
 from brittle_star_project.environment import MorphMode
-from brittle_star_project.trainers.PPOTrainer import apply_per_node
+from brittle_star_project.MLPs.routing import apply_per_node
 import logging
 import re
 from pathlib import Path
@@ -219,8 +219,8 @@ def main(dict_cfg: DictConfig) -> None:
     eval_fn = build_eval_rollout_fn(
         env=env,
         obs_processor=obs_processor,
-        sensor_apply=lambda p, x: apply_per_node(sensor, p, x),
+        sensor_apply=lambda p, x: apply_per_node(sensor.apply, p, x),
-        actor_apply=lambda p, x: apply_per_node(actor, p, x),
+        actor_apply=lambda p, x: apply_per_node(actor.apply, p, x),
         message_passer_apply=(None if message_passer is None else message_passer.apply),
         action_low=action_low,
         action_high=action_high,

src/brittle_star_project/MLPs/routing.py

@@ -0,0 +1,22 @@
+"""Shared JAX routing utilities for decentralized multi-agent models."""
+
+import jax
+
+
+def apply_per_node(apply_fn, params, x):
+    """Apply a Flax module independently to each node.
+
+    Args:
+        apply_fn: The module's ``apply`` method (e.g. ``sensor.apply``).
+        params: Per-node parameters with shape ``(num_nodes, ...)``.
+        x: Input tensor with shape ``(batch, num_nodes, features)``.
+
+    Returns:
+        Output tensor with shape ``(batch, num_nodes, out_features)``.
+    """
+
+    def apply_single_node(p, x_node):
+        # x_node: (batch, feat) — one node's input across the batch
+        return jax.vmap(lambda xi: apply_fn(p, xi))(x_node)
+
+    return jax.vmap(apply_single_node, in_axes=(0, 1), out_axes=1)(params, x)

src/brittle_star_project/evaluation/policy.py

@@ -7,6 +7,7 @@ import jax
 import jax.numpy as jnp
 import numpy as np
 
+from brittle_star_project.MLPs.routing import apply_per_node
 from brittle_star_project.evaluation.checkpoint import load_params
 
 
@@ -147,22 +148,12 @@ class PolicyAgent:
             obs_processor=obs_processor,
         )
 
-    def _apply_per_node(self, net, params, x):
-        # params: (nodes, ...)
-        # x: (batch, nodes, feat)
-
-        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 act(self, *, observations: dict[str, Any]) -> np.ndarray:
         """Return deterministic action (actor mean, no exploration noise)."""
         batched_obs = jax.tree.map(lambda x: jnp.asarray(x)[None, ...], observations)
         obs = self._obs_processor(batched_obs)
 
-        hidden = self._apply_per_node(self._sensor, self._params["sensor_params"], obs)
+        hidden = apply_per_node(self._sensor.apply, self._params["sensor_params"], obs)
 
         if self._message_passer is not None:
             mp_params = self._params.get("message_passer_params")
@@ -172,6 +163,6 @@ class PolicyAgent:
                 )
             hidden = jax.vmap(lambda x: self._message_passer.apply(mp_params, x))(hidden)
 
-        mean, _log_std = self._apply_per_node(self._actor, self._params["actor_params"], hidden)
+        mean, _log_std = apply_per_node(self._actor.apply, self._params["actor_params"], hidden)
 
         return np.asarray(mean, dtype=np.float32).ravel()

src/brittle_star_project/trainers/PPOTrainer.py

@@ -23,6 +23,7 @@ from brittle_star_project.evaluation.evaluate_mjx import (
     build_eval_rollout_fn,
     evaluate_checkpoint_mjx,
 )
+from brittle_star_project.MLPs.routing import apply_per_node
 from brittle_star_project.MLPs.mlps import (
     Actor,
     AgentParams,
@@ -71,7 +72,7 @@ def _get_action_and_value_noise(
     action_high,
 ):
     # (B, n_nodes, feat)
-    hidden = apply_per_node(sensor, agent_state.params["sensor_params"], next_obs)
+    hidden = apply_per_node(sensor.apply, agent_state.params["sensor_params"], next_obs)
 
     if message_passer is not None:
         params = agent_state.params["message_passer_params"]
@@ -82,7 +83,7 @@ def _get_action_and_value_noise(
         feature_extractor, agent_state.params["feature_extractor_params"], next_obs
     )
 
-    mean, log_std = apply_per_node(actor, agent_state.params["actor_params"], hidden)
+    mean, log_std = apply_per_node(actor.apply, agent_state.params["actor_params"], hidden)
     log_std = jnp.clip(log_std, -5, 2)
     key, subkey = jax.random.split(key)
     noise = jax.random.normal(subkey, shape=mean.shape)
@@ -272,17 +273,6 @@ def _step_env_wrapped(
     )
 
 
-def apply_per_node(net, params, x):
-    # params: (nodes, ...)
-    # x: (batch, nodes, feat)
-
-    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:
@@ -516,10 +506,10 @@ class PPOTrainer:
         )
 
         def apply_sensor(p, x):
-            return apply_per_node(self.sensor, p, x)
+            return apply_per_node(self.sensor.apply, p, x)
 
         def apply_actor(p, x):
-            return apply_per_node(self.actor, p, x)
+            return apply_per_node(self.actor.apply, p, x)
 
         def apply_critic(p, x):
             return apply_shared(self.critic, p, x)
@@ -903,8 +893,8 @@ class PPOTrainer:
                 self._eval_fn = build_eval_rollout_fn(
                     env=self.env,
                     obs_processor=self.obs_processor,
-                    sensor_apply=lambda p, x: apply_per_node(self.sensor, p, x),
+                    sensor_apply=lambda p, x: apply_per_node(self.sensor.apply, p, x),
-                    actor_apply=lambda p, x: apply_per_node(self.actor, p, x),
+                    actor_apply=lambda p, x: apply_per_node(self.actor.apply, p, x),
                     message_passer_apply=(
                         None if self.message_passer is None else self.message_passer.apply
                     ),