Merge branch 'dev' into feat/training-evaluation

src/brittle_star_project/MLPs/__init__.py

@@ -0,0 +1,19 @@
+from .mlps import (
+    GenericDenseLayersWithActivation,
+    OneDenseLayerMLP,
+    Actor,
+    MessagePasser,
+    AgentParams,
+    Storage,
+)
+from .adjancency_builder import build_adjacency
+
+__all__ = [
+    "GenericDenseLayersWithActivation",
+    "OneDenseLayerMLP",
+    "Actor",
+    "MessagePasser",
+    "AgentParams",
+    "Storage",
+    "build_adjacency",
+]

src/brittle_star_project/MLPs/adjancency_builder.py

@@ -0,0 +1,67 @@
+from brittle_star_project.environment.env_config import MorphMode
+import jax.numpy as jnp
+
+
+def build_adjacency(segments_per_arm, mode: MorphMode):
+    num_arms = sum(1 for s in segments_per_arm if s > 0)
+    num_segments = sum(segments_per_arm)
+
+    # FOR NOW SEMI HARDCODE:
+    # CENTRALIZED: 1 agent, no stress, adja = 1,1 = [[1]]
+    # FULLY CONNECTED: 5 agents: adj = alle 1
+    # CENTRAL DISK:#arms= 5 agents, only neighbor as adjacent so diagonal kinda..
+    # ARM = #segments agents: diago kinda, but extra, center ring too, put center mlps first or..
+
+    if mode == MorphMode.CENTRALIZED:
+        return jnp.ones((1, 1))
+
+    if mode == MorphMode.FULLY_CONNECTED:
+        adj = jnp.ones((num_arms, num_arms))  # everybody adjacent everybody
+        return adj
+
+    if mode == MorphMode.RING:  # ring
+        adj = jnp.zeros((num_arms, num_arms))
+        for i in range(num_arms):
+            adj = adj.at[i, i].set(1)  # self
+            adj = adj.at[i, (i - 1) % num_arms].set(1)
+            adj = adj.at[i, (i + 1) % num_arms].set(1)  # left and right..
+        return adj
+
+    if mode == MorphMode.SEGMENT:
+        num_nodes = num_arms + num_segments
+        adj = jnp.zeros((num_nodes, num_nodes))
+
+        # first ring
+        for i in range(num_arms):
+            # self
+            adj = adj.at[i, i].set(1)
+
+            # ring neighbors
+            adj = adj.at[i, (i - 1) % num_arms].set(1)
+            adj = adj.at[i, (i + 1) % num_arms].set(1)
+
+        # then segment chains
+        idx = 0
+        for arm_idx, seg_count in enumerate(segments_per_arm):
+            for i in range(seg_count):
+                seg_node = num_arms + idx + i
+
+                adj = adj.at[seg_node, seg_node].set(1)
+                if i > 0:
+                    adj = adj.at[seg_node, seg_node - 1].set(1)
+                if i < seg_count - 1:
+                    adj = adj.at[seg_node, seg_node + 1].set(1)
+
+            idx += seg_count
+
+        idx = 0
+        for arm_idx, seg_count in enumerate(segments_per_arm):
+            first_seg = num_arms + idx  # first segment of this arm
+
+            # connect ring node first segment
+            adj = adj.at[arm_idx, first_seg].set(1)
+            adj = adj.at[first_seg, arm_idx].set(1)
+
+            idx += seg_count
+
+        return adj

src/brittle_star_project/MLPs/mlps.py

@@ -1,11 +1,11 @@
 from dataclasses import dataclass, fields, field
 
-import flax
 import flax.linen as nn
 import jax.numpy as jnp
 import jax.tree_util
 from typing import Sequence, Callable
 from flax.linen.initializers import constant, orthogonal
+from flax.core import FrozenDict
 
 
 # semi generic so we can easily make a config for it in experiments
@@ -37,30 +37,56 @@ class Actor(nn.Module):
         return mean, log_std
 
 
+class MessagePasser(nn.Module):
+    hidden_dim: int
+    num_propagation_steps: int
+    adj_matrix: jnp.ndarray
+
+    @nn.compact
+    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 = self.adj_matrix
+            aggregated = agg @ messages
+
+            x_concat = jnp.concatenate([x, aggregated], axis=-1)
+
+            gate = nn.sigmoid(nn.Dense(self.hidden_dim)(x_concat))
+            candidate = nn.tanh(nn.Dense(self.hidden_dim)(x_concat))
+            x = gate * x + (1 - gate) * candidate
+
+        return x
+
+
 @jax.tree_util.register_dataclass
 @dataclass
 class AgentParams:
-    sensor_params: flax.core.FrozenDict
-    actor_params: flax.core.FrozenDict
-    critic_params: flax.core.FrozenDict
-    feature_extractor_params: flax.core.FrozenDict
+    sensor_params: FrozenDict | dict
+    actor_params: FrozenDict | dict
+    critic_params: FrozenDict | dict
+    feature_extractor_params: FrozenDict | dict
+    message_passer_params: FrozenDict | dict
 
 
 @jax.tree_util.register_dataclass
 @dataclass
 class Storage:
-    obs: jnp.array
-    actions: jnp.array
-    logprobs: jnp.array
-    dones: jnp.array
-    values: jnp.array
-    advantages: jnp.array
-    returns: jnp.array
-    rewards: jnp.array
+    obs: jnp.ndarray
+    actions: jnp.ndarray
+    logprobs: jnp.ndarray
+    dones: jnp.ndarray
+    values: jnp.ndarray
+    advantages: jnp.ndarray
+    returns: jnp.ndarray
+    rewards: jnp.ndarray
 
-    raw_actions: jnp.ndarray = None  # before clipping
-    means: jnp.ndarray = None  # policy mean
-    stds: jnp.ndarray = None  # policy std
+    raw_actions: jnp.ndarray | None = None  # before clipping
+    means: jnp.ndarray | None = None  # policy mean
+    stds: jnp.ndarray | None = None  # policy std
 
     def replace(self, **kwargs) -> "Storage":
         fs = fields(self)

src/brittle_star_project/dataclasses/EpisodeStatistics.py

@@ -4,7 +4,7 @@ import jax.numpy as jnp
 
 @flax.struct.dataclass
 class EpisodeStatistics:
-    episode_returns: jnp.array
-    episode_lengths: jnp.array
-    returned_episode_returns: jnp.array
-    returned_episode_lengths: jnp.array
+    episode_returns: jnp.ndarray
+    episode_lengths: jnp.ndarray
+    returned_episode_returns: jnp.ndarray
+    returned_episode_lengths: jnp.ndarray

src/brittle_star_project/environment/__init__.py

@@ -1,4 +1,4 @@
-from .env_config import ArenaConfig, EnvConfig, MorphologyConfig
+from .env_config import ArenaConfig, EnvConfig, MorphologyConfig, MorphMode
 from .env_types import Backend, Task
 from .env_wrapper import BrittleStarEnv
 from .factory import BrittleStarEnvFactory
@@ -13,6 +13,7 @@ __all__ = [
     "Task",
     "BrittleStarEnv",
     "BrittleStarEnvFactory",
+    "MorphMode",
     "create_obs_processor",
     "compute_padding_masks",
 ]

src/brittle_star_project/environment/env_config.py

@@ -1,10 +1,18 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
+from enum import Enum
 
 from .env_types import Task
 
 
+class MorphMode(Enum):
+    CENTRALIZED = 0
+    FULLY_CONNECTED = 1
+    RING = 2
+    SEGMENT = 3
+
+
 @dataclass
 class MorphologyConfig:
     """Brittle star morphology configuration.
@@ -20,6 +28,7 @@ class MorphologyConfig:
     segments_per_arm: list[int] = field(default_factory=lambda: [4, 4, 4, 4, 4])
     use_p_control: bool = True
     use_torque_control: bool = False
+    morph_mode: MorphMode = MorphMode.CENTRALIZED
 
     @property
     def num_arms(self) -> int:

src/brittle_star_project/environment/obs_processing.py

@@ -2,6 +2,12 @@ import jax
 import jax.numpy as jnp
 from typing import Dict, Tuple, Optional
 
+from brittle_star_project.environment.env_config import MorphMode
+
+from experiment_logger import get_logger
+
+logger = get_logger()
+
 _JOINT_SCALED_KEYS = frozenset(
     {
         "joint_position",
@@ -18,9 +24,53 @@ _SEGMENT_SCALED_KEYS = frozenset(
 )
 
 
+def _build_joint_indices(segments_per_arm, indices_mlp):
+    indices = []
+    start = 0
+    for i, segs in enumerate(segments_per_arm):
+        # 2 joints per segment
+        if i in indices_mlp:
+            count = segs * 2
+            idx = jnp.arange(start, start + count)
+            indices.append(idx)
+            start += count
+    return indices
+
+
+def _build_segment_indices(segments_per_arm, indices_mlp):
+    indices = []
+    start = 0
+    for i, segs in enumerate(segments_per_arm):
+        if i in indices_mlp:
+            idx = jnp.arange(start, start + segs)
+            indices.append(idx)
+            start += segs
+    return indices
+
+
 def create_obs_processor(
-    bounds_dict: Dict[str, Tuple[float, float]], padding_masks: Optional[Dict] = None
+    bounds_dict: Dict[str, Tuple[float, float]],
+    num_arms: int,
+    needed_copies: int,
+    padding_masks: Optional[Dict] = None,
+    morph_mode: MorphMode = MorphMode.CENTRALIZED,
+    segments_per_arm=[4, 4, 4, 4, 4],
+    agent_indices=[0, 1, 2, 3, 4],
 ):
+    # made a set to allow O(1) search
+    ordered_keys = frozenset(
+        [
+            "disk_z_tilt",
+            "joint_actuator_force",
+            "joint_position",
+            "joint_velocity",
+            "robot_direction_to_target",
+            "segment_contact",
+        ]
+    )
+    segment_indices = _build_segment_indices(segments_per_arm, agent_indices)
+    joint_indices = _build_joint_indices(segments_per_arm, agent_indices)
+
     def _add_derived_features(obs: dict) -> dict:
         new_obs = dict(obs)
         if "disk_rotation" in new_obs:
@@ -51,41 +101,91 @@ def create_obs_processor(
                 normalized[key] = arr
         return normalized
 
-    def _pad_features(obs: dict) -> dict:
-        padded = {}
+    def _split_to_agents(obs: dict, morph_mode) -> dict:
+        output = {}
+        num_agents = needed_copies  # IMPORTANT: number of MLPs
+
+        segs_per_arm = 4
+        joints_per_segment = 2
+        joints_per_arm = segs_per_arm * joints_per_segment
         for key, arr in obs.items():
-            if key in _JOINT_SCALED_KEYS:
-                padded_arr = jnp.zeros(padding_masks["target_size_2x"], dtype=arr.dtype)
-                padded[key] = padded_arr.at[padding_masks["mask_2x"]].set(arr)
-            elif key in _SEGMENT_SCALED_KEYS:
-                padded_arr = jnp.zeros(padding_masks["target_size_1x"], dtype=arr.dtype)
-                padded[key] = padded_arr.at[padding_masks["mask_1x"]].set(arr)
+            if arr.size == 0:
+                continue
+
+            if arr.ndim == 0:
+                arr = arr.reshape(1)
+
+            if key in _SEGMENT_SCALED_KEYS:
+                per_agent = []
+                for i, _ in enumerate(agent_indices):
+                    idx = segment_indices[i]
+                    taken = jnp.take(arr, idx, axis=0)
+                    pad_len = segs_per_arm - taken.shape[0]
+                    padded = jnp.pad(taken, [(9, pad_len)] + [(0, 0)] * (taken.ndim - 1))
+
+                    per_agent.append(padded.reshape(-1))
+                arr = jnp.stack(per_agent)
+            elif key in _JOINT_SCALED_KEYS:
+                per_agent = []
+                for i, _ in enumerate(agent_indices):
+                    idx = joint_indices[i]
+                    taken = jnp.take(arr, idx, axis=0)
+                    pad_len = joints_per_arm - taken.shape[0]
+                    padded = jnp.pad(taken, [(0, pad_len)] + [(0, 0)] * (taken.ndim - 1))
+
+                    per_agent.append(padded.reshape(-1))
+                arr = jnp.stack(per_agent)
             else:
-                padded[key] = arr
-        return padded
+                arr = jnp.repeat(arr[None, :], num_agents, axis=0)
+
+            if morph_mode == MorphMode.CENTRALIZED:
+                output[key] = arr.reshape(1, -1)
+            elif key in _JOINT_SCALED_KEYS:
+                output[key] = arr.reshape(num_agents, -1)
+            elif key in _SEGMENT_SCALED_KEYS:
+                output[key] = arr[:, None]
+            else:
+                output[key] = arr
+
+        return output
 
     def _flatten_features(obs: dict) -> jnp.ndarray:
-        ordered_keys = [
-            "disk_z_tilt",
-            "joint_actuator_force",
-            "joint_position",
-            "joint_velocity",
-            "robot_direction_to_target",
-            "segment_contact",
-        ]
+        """
+        Input:
+            key -> (num_arms, feat_per_key)
+
+        Output:
+            (num_arms, total_features)
+        """
         values = []
+
         for key in ordered_keys:
-            if key in obs:
-                arr = jnp.asarray(obs[key]).flatten()
-                if arr.size > 0:
-                    values.append(arr)
-        return jnp.concatenate(values)
+            if key not in obs:
+                continue
+
+            arr = jnp.asarray(obs[key])  # (num_arms, feat)
+
+            if arr.size == 0:
+                continue
+
+            if arr.ndim == 1:
+                arr = arr[:, None]
+
+            arr = arr.reshape(arr.shape[0], -1)
+
+            values.append(arr)
+
+        return jnp.concatenate(values, axis=-1)  # (num_arms, total_feat)
 
     def _process_single(obs_dict: dict) -> jnp.ndarray:
         processed = _add_derived_features(obs_dict)
         processed = _normalize_features(processed)
-        if padding_masks is not None:
-            processed = _pad_features(processed)
-        return _flatten_features(processed)
+        processed = _split_to_agents(processed, morph_mode)
+        flat = _flatten_features(processed)  # (num_arms, total_feat)
+
+        logger.debug(f"[FLATTENED FINAL] shape: {flat.shape}")
+        logger.debug(f"[PER AGENT] example row 0 shape: {flat[0].shape}")
+
+        return flat  # (agents, feat)
 
     return jax.jit(jax.vmap(_process_single))

src/brittle_star_project/environment/padded_obs_wrapper.py

@@ -30,6 +30,12 @@ def compute_padding_masks(
     mask_2x = []
 
     for arm_idx, (actual, ref) in enumerate(zip(segments_per_arm, reference_segments_per_arm)):
+        if not isinstance(actual, int):
+            actual = actual.item()
+
+        if not isinstance(ref, int):
+            ref = ref.item()
+
         if not (0 <= actual <= ref):
             raise ValueError(
                 f"Invalid amputation at arm {arm_idx}: "

src/brittle_star_project/evaluation/checkpoint.py

@@ -4,6 +4,8 @@ import yaml
 from dataclasses import dataclass
 from pathlib import Path
 
+from collections.abc import Mapping
+
 import flax
 from omegaconf import OmegaConf
 
@@ -32,15 +34,19 @@ def load_params(path: Path) -> dict:
 
     sensor_params = None
     actor_params = None
+    message_passer_params = None
 
     # Extract params from restored checkpoint
-    if isinstance(restored, dict):
+    if isinstance(restored, Mapping):
         params_sub = restored.get("params", {})
         sensor_params = restored.get("sensor_params") or params_sub.get("sensor_params")
         actor_params = restored.get("actor_params") or params_sub.get("actor_params")
+        message_passer_params = restored.get("message_passer_params") or params_sub.get(
+            "message_passer_params"
+        )
     elif isinstance(restored, (list, tuple)) and len(restored) >= 2:
         params_part = restored[1]
-        if isinstance(params_part, dict):
+        if isinstance(params_part, Mapping):
             sensor_params = params_part.get("0", params_part.get(0))
             actor_params = params_part.get("1", params_part.get(1))
         elif isinstance(params_part, (list, tuple)) and len(params_part) >= 2:
@@ -53,6 +59,7 @@ def load_params(path: Path) -> dict:
     return {
         "sensor_params": sensor_params,
         "actor_params": actor_params,
+        "message_passer_params": message_passer_params,
     }
 
 

src/brittle_star_project/evaluation/evaluate_mjx.py

@@ -52,6 +52,7 @@ def build_eval_rollout_fn(
     obs_processor: Callable,
     sensor_apply: Callable,
     actor_apply: Callable,
+    message_passer_apply: Callable | None = None,
     action_low: jnp.ndarray,
     action_high: jnp.ndarray,
     reward_fn: Callable,
@@ -67,11 +68,14 @@ def build_eval_rollout_fn(
             returned by `create_obs_processor`.
         sensor_apply: The sensor network's `apply` method (JIT-compiled).
         actor_apply: The actor network's `apply` method (JIT-compiled).
+        message_passer_apply: Optional message-passing module apply method.
+            When provided, it is applied between the sensor and actor, using
+            `params["message_passer_params"]`.
         action_low: Per-joint action lower bound (JAX array, shape `(action_dim,)`).
         action_high: Per-joint action upper bound (JAX array, shape `(action_dim,)`).
         reward_fn: Shaped reward function with signature
             `reward_fn(env_state, next_env_state) -> jnp.ndarray`.
-            Typically the module-level `reward_fn` from `PPOTrainer`.
+            Typically, the module-level `reward_fn` from `PPOTrainer`.
 
     Returns:
         A JIT-compiled callable that runs one deterministic evaluation episode.
@@ -101,10 +105,14 @@ def build_eval_rollout_fn(
 
             obs = obs_processor(state.observations)
             hidden = sensor_apply(params["sensor_params"], obs)
+            if message_passer_apply is not None:
+                mp_params = params["message_passer_params"]
+                hidden = jax.vmap(lambda x: message_passer_apply(mp_params, x))(hidden)
             mean, _log_std = actor_apply(params["actor_params"], hidden)
 
             # Deterministic action: use the actor mean, no exploration noise.
-            action = jnp.clip(mean, action_low, action_high)
+            flat_mean = mean.reshape(mean.shape[0], -1)
+            action = jnp.clip(flat_mean, action_low, action_high)
             next_state = step_1(state=state, action=action)
 
             shaped_reward = reward_fn(state, next_state)

src/brittle_star_project/evaluation/policy.py

@@ -24,10 +24,17 @@ class PolicyAgent:
         *,
         sensor_params: Any,
         actor_params: Any,
+        message_passer_params: Any | None = None,
+        message_passing_steps: int | None = None,
+        adj_matrix: Any | None = None,
         action_dim: int,
         obs_processor: Any,
     ) -> None:
-        from brittle_star_project.MLPs.mlps import Actor, GenericDenseLayersWithActivation
+        from brittle_star_project.MLPs.mlps import (
+            Actor,
+            GenericDenseLayersWithActivation,
+            MessagePasser,
+        )
 
         # Infer layer sizes from params
         try:
@@ -45,7 +52,8 @@ class PolicyAgent:
             key = f"Dense_{idx}"
             if key not in dense_params:
                 break
-            layer_sizes.append(int(np.asarray(dense_params[key]["kernel"]).shape[1]))
+
+            layer_sizes.append(int(np.asarray(dense_params[key]["kernel"]).shape[-1]))
             idx += 1
 
         if not layer_sizes:
@@ -53,11 +61,31 @@ class PolicyAgent:
 
         self._sensor = GenericDenseLayersWithActivation(layer_sizes=layer_sizes)
         self._actor = Actor(action_dim=action_dim)
-        self._sensor_apply = jax.jit(self._sensor.apply)
-        self._actor_apply = jax.jit(self._actor.apply)
+
+        self._message_passer = None
+        if message_passer_params is not None and not (
+            isinstance(message_passer_params, dict) and len(message_passer_params) == 0
+        ):
+            if message_passing_steps is None or adj_matrix is None:
+                raise ValueError(
+                    "Checkpoint contains message_passer_params but PolicyAgent was not given "
+                    "message_passing_steps and adj_matrix. Pass these when constructing the agent "
+                    "so decentralized evaluation matches training."
+                )
+
+            hidden_dim = int(layer_sizes[-1])
+            self._message_passer = MessagePasser(
+                hidden_dim=hidden_dim,
+                num_propagation_steps=int(message_passing_steps),
+                adj_matrix=jnp.asarray(adj_matrix),
+            )
+            self._message_passer.apply = jax.jit(self._message_passer.apply)
+        self._sensor.apply = jax.jit(self._sensor.apply)
+        self._actor.apply = jax.jit(self._actor.apply)
         self._params = {
             "sensor_params": sensor_params,
             "actor_params": actor_params,
+            "message_passer_params": message_passer_params,
         }
         self._obs_processor = obs_processor
 
@@ -67,6 +95,9 @@ class PolicyAgent:
         *,
         sensor_params: Any,
         actor_params: Any,
+        message_passer_params: Any | None = None,
+        message_passing_steps: int | None = None,
+        adj_matrix: Any | None = None,
         action_dim: int,
         obs_processor: Any,
     ) -> "PolicyAgent":
@@ -74,14 +105,24 @@ class PolicyAgent:
         return cls(
             sensor_params=sensor_params,
             actor_params=actor_params,
+            message_passer_params=message_passer_params,
+            message_passing_steps=message_passing_steps,
+            adj_matrix=adj_matrix,
             action_dim=action_dim,
             obs_processor=obs_processor,
         )
 
-    def set_params(self, *, sensor_params: Any, actor_params: Any) -> None:
+    def set_params(
+        self,
+        *,
+        sensor_params: Any,
+        actor_params: Any,
+        message_passer_params: Any | None = None,
+    ) -> None:
         """Update parameters for evaluation without rebuilding the model."""
         self._params["sensor_params"] = sensor_params
         self._params["actor_params"] = actor_params
+        self._params["message_passer_params"] = message_passer_params
 
     @classmethod
     def from_checkpoint(
@@ -90,6 +131,8 @@ class PolicyAgent:
         *,
         action_dim: int,
         obs_processor: Any,
+        message_passing_steps: int | None = None,
+        adj_matrix: Any | None = None,
     ) -> "PolicyAgent":
         """Load params from .flax and construct the agent."""
         params = load_params(model_path)
@@ -97,15 +140,38 @@ class PolicyAgent:
         return cls(
             sensor_params=params["sensor_params"],
             actor_params=params["actor_params"],
+            message_passer_params=params.get("message_passer_params"),
+            message_passing_steps=message_passing_steps,
+            adj_matrix=adj_matrix,
             action_dim=action_dim,
             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)[0]
-        hidden = self._sensor_apply(self._params["sensor_params"], obs)
-        mean, _log_std = self._actor_apply(self._params["actor_params"], hidden)
+        obs = self._obs_processor(batched_obs)
+
+        hidden = self._apply_per_node(self._sensor, self._params["sensor_params"], obs)
+
+        if self._message_passer is not None:
+            mp_params = self._params.get("message_passer_params")
+            if mp_params is None or (isinstance(mp_params, dict) and len(mp_params) == 0):
+                raise ValueError(
+                    "PolicyAgent has a message passer but message_passer_params are missing/empty."
+                )
+            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)
 
         return np.asarray(mean, dtype=np.float32).ravel()

src/brittle_star_project/evaluation/rollout.py

@@ -117,18 +117,20 @@ def rollout_viewer(
 
     episode_return = 0.0
     observations = _get_observations(state)
+
     prev_dist = _get_xy_distance_to_target(observations) if observations else None
     reached_target = _target_reached(state=state)
 
     steps = 0
     with mujoco.viewer.launch_passive(model, data) as viewer:
         step_iter = range(int(max_steps)) if max_steps is not None else itertools.count()
-        for _step_idx in step_iter:
+        for _ in step_iter:
             if not viewer.is_running():
                 break
             step_start = time.time()
 
             obs_dict = observations or {}
+
             action = policy.act(observations=obs_dict)
             if action_mask is not None:
                 action = action[action_mask]

src/brittle_star_project/ppo.py

@@ -1,14 +1,27 @@
 from functools import partial
 
-import flax
 import jax
 import jax.numpy as jnp
+from jax import debug
+from flax.core import FrozenDict
+from experiment_logger import get_logger
+from brittle_star_project.utils import logged_jit
+
+logger = get_logger()
 
 
 # Chose to use a class as it seemed the easiest way to integrate the CleanRL code style
 # with our need to seperate concerns
 class PPO:
-    def __init__(self, args, sensor, actor, critic, feature_extractor, message_passer=None):
+    def __init__(
+        self,
+        args,
+        sensor_apply,
+        actor_apply,
+        critic_apply,
+        feature_extractor_apply,
+        message_passer=None,
+    ):
         self.args = args
 
         if not message_passer:
@@ -18,10 +31,10 @@ class PPO:
             partial(
                 ppo_loss,
                 args=args,
-                sensor_apply=sensor.apply,
-                actor_apply=actor.apply,
-                critic_apply=critic.apply,
-                feature_extractor_apply=feature_extractor.apply,
+                sensor_apply=sensor_apply,
+                actor_apply=actor_apply,
+                critic_apply=critic_apply,
+                feature_extractor_apply=feature_extractor_apply,
                 message_passer=message_passer,
             ),
             has_aux=True,
@@ -29,8 +42,14 @@ class PPO:
 
     # This PPO class should be initialized only once,
     # or this function will need to recompile
-    @partial(jax.jit, static_argnums=0)
+    @partial(logged_jit, static_argnums=0)
     def update_ppo(self, agent_state, storage, key):
+        debug.callback(logger.debug, f"[PPO] storage.obs shape: {storage.obs.shape}")
+        debug.callback(logger.debug, f"[PPO] storage.actions shape: {storage.actions.shape}")
+        debug.callback(logger.debug, f"[PPO] storage.logprobs shape: {storage.logprobs.shape}")
+        debug.callback(logger.debug, f"[PPO] storage.advantages shape: {storage.advantages.shape}")
+        debug.callback(logger.debug, f"[PPO] storage.returns shape: {storage.returns.shape}")
+
         args = self.args
         ppo_loss_grad_fn = self.ppo_loss_grad_fn
 
@@ -49,6 +68,16 @@ class PPO:
             shuffled_storage = jax.tree.map(convert_data, flatten_storage)
 
             def update_minibatch(agent_state, minibatch):
+                debug.callback(logger.debug, f"[PPO] minibatch.obs: {minibatch.obs.shape}")
+                debug.callback(logger.debug, f"[PPO] minibatch.actions: {minibatch.actions.shape}")
+                debug.callback(
+                    logger.debug, f"[PPO] minibatch.logprobs: {minibatch.logprobs.shape}"
+                )
+                debug.callback(
+                    logger.debug, f"[PPO] minibatch.advantages: {minibatch.advantages.shape}"
+                )
+                debug.callback(logger.debug, f"[PPO] minibatch.returns: {minibatch.returns.shape}")
+
                 (loss, (pg_loss, v_loss, entropy_loss, approx_kl)), grads = ppo_loss_grad_fn(
                     agent_state.params,
                     minibatch.obs,
@@ -58,19 +87,12 @@ class PPO:
                     minibatch.returns,
                 )
                 agent_state = agent_state.apply_gradients(grads=grads)
-                return agent_state, (
-                    loss,
-                    pg_loss,
-                    v_loss,
-                    entropy_loss,
-                    approx_kl,
-                    grads,
-                )
+                return agent_state, (loss, pg_loss, v_loss, entropy_loss, approx_kl)
 
             agent_state, metrics = jax.lax.scan(update_minibatch, agent_state, shuffled_storage)
             return (agent_state, key), metrics
 
-        (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads) = jax.lax.scan(
+        (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl) = jax.lax.scan(
             update_epoch, (agent_state, key), (), length=args.update_epochs
         )
         return agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key
@@ -84,27 +106,45 @@ that are now not in the same scope
 """
 
 
-@partial(jax.jit, static_argnums=(0, 1, 2, 3, 4))
+@partial(logged_jit, static_argnums=(0, 1, 2, 3, 4))
 def get_action_and_value(
     sensor_apply,
     actor_apply,
     message_passer,
     critic_apply,
     feature_extractor_apply,
-    params: flax.core.FrozenDict,
+    params: FrozenDict,
     x: jnp.ndarray,
     action: jnp.ndarray,
 ):
     hidden_sensor = sensor_apply(params["sensor_params"], x)
     hidden_critic = feature_extractor_apply(params["feature_extractor_params"], x)
-    hidden_sensor = message_passer(hidden_sensor)
+
+    # only apply message passing in decentralized context
+    if message_passer is not None:
+        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}")
+
     mean, log_std = actor_apply(params["actor_params"], hidden_sensor)
+
+    debug.callback(logger.debug, f"[SHAPE] mean: {mean.shape}")
+    debug.callback(logger.debug, f"[SHAPE] log_std: {log_std.shape}")
+    debug.callback(logger.debug, f"[SHAPE] action: {action.shape}")
+
     log_std = jnp.clip(log_std, -5, 2)
     std = jnp.exp(log_std)
 
-    logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
-    entropy = (0.5 + 0.5 * jnp.log(2 * jnp.pi) + log_std).sum(-1)
+    logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi))
+    debug.callback(logger.debug, f"[SHAPE] logprob pre-sum: {logprob.shape}")
+
+    logprob = logprob.sum(axis=(-2, -1))
+    debug.callback(logger.debug, f"[SHAPE] logprob final: {logprob.shape}")
+
+    entropy = (0.5 + 0.5 * jnp.log(2 * jnp.pi) + log_std).sum(axis=(-2, -1))
     value = critic_apply(params["critic_params"], hidden_critic).squeeze(-1)
+    debug.callback(logger.debug, f"[SHAPE] value: {value.shape}")
 
     return logprob, entropy, value
 
@@ -150,7 +190,7 @@ def ppo_loss(
     return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl))
 
 
-def identity(hidden):
+def identity(_, hidden):
     """
     Used for seamless jax integration,
     avoids having branching inside jitted function,

src/brittle_star_project/trainers/PPOTrainer.py

@@ -3,13 +3,14 @@ import random
 import time
 from dataclasses import asdict, dataclass
 from functools import partial
+from typing import Any, Optional
 
 import jax
 import jax.numpy as jnp
 import numpy as np
 import optax
+import flax.linen as nn
 from flax.training.train_state import TrainState
-from typing import Any
 
 from experiment_logger import get_logger
 
@@ -26,24 +27,21 @@ from brittle_star_project.MLPs.mlps import (
     Actor,
     AgentParams,
     GenericDenseLayersWithActivation,
+    MessagePasser,
     OneDenseLayerMLP,
     Storage,
 )
+from brittle_star_project.MLPs.adjancency_builder import build_adjacency
 from brittle_star_project.ppo import PPO
+from brittle_star_project.environment import MorphMode
+from brittle_star_project.utils import logged_jit
 
 from brittle_star_project.environment.env_types import Backend
 
 # TODO: clip scaled reward?
 
 
-@jax.jit
-def _get_xy_distance_to_target(obs_dict: dict) -> jnp.ndarray:
-    """Extract xy_distance_to_target for all environments."""
-    # obs_dict is a dict of arrays with leading batch dimension (num_envs, ...)
-    return obs_dict["xy_distance_to_target"].squeeze(-1)  # shape: (num_envs,)
-
-
-@jax.jit
+@logged_jit
 def _clip_action(action: jnp.ndarray, low: jnp.ndarray, high: jnp.ndarray) -> jnp.ndarray:
     return jnp.clip(action, low, high)
 
@@ -54,39 +52,54 @@ def _compute_explained_variance(values: jnp.ndarray, returns: jnp.ndarray) -> fl
     return float(explained_var)
 
 
-@jax.jit
+@logged_jit
 def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, learning_rate):
     frac = 1.0 - (count // (minibatch_count * update_epochs)) / num_iterations
     return learning_rate * frac
 
 
 def _get_action_and_value_noise(
-    sensor: GenericDenseLayersWithActivation,
-    feature_extractor: GenericDenseLayersWithActivation,
-    actor: Actor,
-    critic: OneDenseLayerMLP,
+    sensor: nn.Module,
+    feature_extractor: nn.Module,
+    actor: nn.Module,
+    critic: nn.Module,
+    message_passer: Optional[nn.Module],
     agent_state: TrainState,
     next_obs: jnp.ndarray,
-    key: jax.random.PRNGKey,
+    key,
     action_low,
     action_high,
 ):
-    hidden = sensor.apply(agent_state.params["sensor_params"], next_obs)
-    hidden_critic = feature_extractor.apply(
-        agent_state.params["feature_extractor_params"], next_obs
+    # (B, n_nodes, feat)
+    hidden = apply_per_node(sensor, agent_state.params["sensor_params"], next_obs)
+
+    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))(hidden)
+
+    hidden_critic = apply_shared(
+        feature_extractor, agent_state.params["feature_extractor_params"], next_obs
     )
 
-    mean, log_std = actor.apply(agent_state.params["actor_params"], hidden)
+    mean, log_std = apply_per_node(actor, 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)
     std = jnp.exp(log_std)
-    raw_action = mean + noise * std
-    clipped_action = _clip_action(raw_action, action_low, action_high)
-    logprob = -0.5 * (((raw_action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
-    value = critic.apply(agent_state.params["critic_params"], hidden_critic)
 
-    return clipped_action, raw_action, logprob, value.squeeze(-1), mean, std, key
+    raw_action = mean + noise * std
+    flat_action = raw_action.reshape(
+        raw_action.shape[0], -1
+    )  # concat the per agent, keep the envs dim (batch, agent * action)
+    flat_clipped_action = _clip_action(flat_action, action_low, action_high)
+
+    logprob = -0.5 * (((raw_action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(
+        axis=(-2, -1)
+    )
+    value = apply_shared(critic, agent_state.params["critic_params"], hidden_critic)
+
+    return flat_clipped_action, raw_action, logprob, value.squeeze(-1), mean, std, key
 
 
 def _step_once(
@@ -94,31 +107,59 @@ def _step_once(
     _,
     env_step_fn,
     num_envs: int,
-    sensor: GenericDenseLayersWithActivation,
-    feature_extractor: GenericDenseLayersWithActivation,
-    actor: Actor,
-    critic: OneDenseLayerMLP,
+    sensor: nn.Module,
+    feature_extractor: nn.Module,
+    actor: nn.Module,
+    critic: nn.Module,
+    message_passer: Optional[nn.Module],
     action_low,
     action_high,
 ):
     agent_state, episode_stats, obs, done, key, env_state, terminated_any, truncated_any = carry
-    clipped_action, raw_action, logprob, value, mean, std, key = _get_action_and_value_noise(
-        sensor, feature_extractor, actor, critic, agent_state, obs, key, action_low, action_high
+    flat_clipped_action, raw_action, logprob, value, mean, std, key = _get_action_and_value_noise(
+        sensor,
+        feature_extractor,
+        actor,
+        critic,
+        message_passer,
+        agent_state,
+        obs,
+        key,
+        action_low,
+        action_high,
     )
+    logger = get_logger()
 
+    logger.debug(f"[_step_once] raw_action: {raw_action.shape}")
+    logger.debug(f"[_step_once] clipped_action: {flat_clipped_action.shape}")
+
+    # Supporting signals (often where mismatch originates)
+    logger.debug(f"[_step_once] logprob: {logprob.shape}")
+    logger.debug(f"[_step_once] value: {value.shape}")
+    logger.debug(f"[_step_once] mean: {mean.shape}")
+    logger.debug(f"[_step_once] std: {std.shape}")
+
+    key, reset_key = jax.random.split(key)
+    reset_rngs = jax.random.split(reset_key, num_envs)
+
+    # ---- ENV STEP ----
     key, reset_key = jax.random.split(key)
     reset_rngs = jax.random.split(reset_key, num_envs)
 
     episode_stats, env_state, (next_obs, reward, next_done, terminated, truncated) = env_step_fn(
         episode_stats,
         env_state,
-        clipped_action,
+        flat_clipped_action,
         reset_rngs,
     )
 
     terminated_any = terminated_any | terminated
     truncated_any = truncated_any | truncated
 
+    logger.debug(f"[_step_once] next_obs: {next_obs.shape}")
+    logger.debug(f"[_step_once] reward: {reward.shape}")
+    logger.debug(f"[_step_once] next_done: {next_done.shape}")
+
     storage = Storage(
         obs=obs,
         actions=raw_action,
@@ -231,6 +272,25 @@ 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:
+    batch_size = x.shape[0]
+    x_flattened = x.reshape(batch_size, -1)
+    return jax.vmap(lambda xi: net.apply(params, xi))(x_flattened)
+
+
 def _rollout_jit(
     agent_state,
     episode_stats,
@@ -241,10 +301,11 @@ def _rollout_jit(
     max_steps,
     step_env_fn,
     num_envs: int,
-    sensor: GenericDenseLayersWithActivation,
-    feature_extractor: GenericDenseLayersWithActivation,
-    actor: Actor,
-    critic: OneDenseLayerMLP,
+    sensor: nn.Module,
+    feature_extractor: nn.Module,
+    actor: nn.Module,
+    critic: nn.Module,
+    message_passer: Optional[nn.Module],
     action_low,
     action_high,
 ):
@@ -270,6 +331,7 @@ def _rollout_jit(
             feature_extractor=feature_extractor,
             actor=actor,
             critic=critic,
+            message_passer=message_passer,
             env_step_fn=step_env_fn,
             num_envs=num_envs,
             action_low=action_low,
@@ -321,9 +383,10 @@ def _compute_gae_jit(
     feature_extractor,
     critic,
 ):
-    next_value = critic.apply(
+    next_value = apply_shared(
+        critic,
         agent_state.params["critic_params"],
-        feature_extractor.apply(agent_state.params["feature_extractor_params"], next_obs),
+        apply_shared(feature_extractor, agent_state.params["feature_extractor_params"], next_obs),
     ).squeeze(-1)
 
     advantages = jnp.zeros((num_envs,))
@@ -357,7 +420,11 @@ class TrainingMeasurements:
 
 class PPOTrainer:
     def __init__(
-        self, cfg: BrittleStarConfig, env: BrittleStarJaxEnvWrapper, run_dir: str, run_name: str
+        self,
+        cfg: BrittleStarConfig,
+        env: BrittleStarJaxEnvWrapper,
+        run_dir: str,
+        run_name: str,
     ):
         self.cfg = cfg
         self.ppo = cfg.ppo
@@ -375,24 +442,49 @@ class PPOTrainer:
 
         self.key = jax.random.PRNGKey(self.experiment.seed)
 
+        self.morph_mode = self.cfg.morphology.morph_mode
+
+        self.segments_per_arm = jnp.asarray(self.cfg.morphology.segments_per_arm, dtype=jnp.int32)
+        self.num_segments = self.segments_per_arm.sum().item()
+        self.num_arms = jnp.where(self.segments_per_arm > 0, 1, 0).sum().item()
+
+        self.logger.info(f"[INIT]: Used morphology mode {self.morph_mode}")
+        self.adj = build_adjacency(cfg.morphology.segments_per_arm, self.morph_mode)
+
+        (
+            self.sensor,
+            self.message_passer,
+            self.actor,
+            self.feature_extractor,
+            self.critic,
+            self.needed_copies,
+            self.agent_indices,
+        ) = self._init_agent()
+
+        self.sensor.apply = logged_jit(self.sensor.apply)
+        self.feature_extractor.apply = logged_jit(self.feature_extractor.apply)
+        self.actor.apply = logged_jit(self.actor.apply)
+        self.critic.apply = logged_jit(self.critic.apply)
+
         # Build the centralized observation processor: derive -> normalize -> pad -> flatten.
         self.obs_processor = create_obs_processor(
             bounds_dict=self.cfg.obs_bounds.to_bounds_dict(),
+            needed_copies=self.needed_copies,
+            num_arms=self.num_arms,
+            morph_mode=self.morph_mode,
             padding_masks=self.env.padding_masks,
+            segments_per_arm=self.segments_per_arm,
+            agent_indices=self.agent_indices,
         )
 
-        self.sensor, self.feature_extractor, self.actor, self.critic = self._init_agent()
-        self.sensor.apply = jax.jit(self.sensor.apply)
-        self.feature_extractor.apply = jax.jit(self.feature_extractor.apply)
-        self.actor.apply = jax.jit(self.actor.apply)
-        self.critic.apply = jax.jit(self.critic.apply)
+        self.logger.debug(f"needed copies = {self.needed_copies}")
 
         action_low = jnp.asarray(self.env.single_action_space.low, dtype=jnp.float32)
         action_high = jnp.asarray(self.env.single_action_space.high, dtype=jnp.float32)
         self._action_low = action_low
         self._action_high = action_high
 
-        self._rollout_jit = jax.jit(
+        self._rollout_jit = logged_jit(
             partial(
                 _rollout_jit,
                 max_steps=self.ppo.num_steps,
@@ -407,11 +499,12 @@ class PPOTrainer:
                 feature_extractor=self.feature_extractor,
                 actor=self.actor,
                 critic=self.critic,
+                message_passer=self.message_passer,
                 action_low=action_low,
                 action_high=action_high,
             )
         )
-        self._compute_gae_jit = jax.jit(
+        self._compute_gae_jit = logged_jit(
             partial(
                 _compute_gae_jit,
                 num_envs=self.ppo.num_envs,
@@ -422,7 +515,30 @@ class PPOTrainer:
             )
         )
 
-        self._ppo = PPO(self.ppo, self.sensor, self.actor, self.critic, self.feature_extractor)
+        def apply_sensor(p, x):
+            return apply_per_node(self.sensor, p, x)
+
+        def apply_actor(p, x):
+            return apply_per_node(self.actor, p, x)
+
+        def apply_critic(p, x):
+            return apply_shared(self.critic, p, x)
+
+        def apply_feature(p, x):
+            return apply_shared(self.feature_extractor, p, x)
+
+        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 if self.message_passer is not None else None,
+        )
 
         self.agent_state = self._init_agent_state()
 
@@ -440,33 +556,136 @@ class PPOTrainer:
 
     def _init_agent(self):
         self.logger.info("[AGENT]: Initializing agent...")
+        agent_indices = [0, 1, 2, 3, 4]
+        match self.morph_mode:
+            case MorphMode.CENTRALIZED:
+                needed_copies = 1
+            case MorphMode.FULLY_CONNECTED | MorphMode.RING:
+                agent_mask = self.segments_per_arm > 0
+                agent_indices = jnp.where(agent_mask)[0]
+                needed_copies = jnp.where(self.segments_per_arm > 0, 1, 0).sum().item()
+            case MorphMode.SEGMENT:
+                agent_mask = self.segments_per_arm > 0
+                agent_indices = jnp.where(agent_mask)[0]
+                needed_copies = (
+                    self.segments_per_arm.sum() + jnp.where(self.segments_per_arm > 0, 1, 0).sum()
+                ).item()
 
+        # scale actor output with size of model --> more models ==> less actions needed per model
+        actor = Actor(action_dim=self.env.single_action_space.shape[0] // needed_copies)
         sensor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300])
+        message_passer: Optional[nn.Module] = (
+            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
+        )
+
         feature_extractor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300])
-        actor = Actor(action_dim=self.env.single_action_space.shape[0])
         critic = OneDenseLayerMLP()
-        return sensor, feature_extractor, actor, critic
+        return (
+            sensor,
+            message_passer,
+            actor,
+            feature_extractor,
+            critic,
+            needed_copies,
+            agent_indices,
+        )
 
     def _init_agent_state(self) -> TrainState:
         self.logger.info("[AGENT STATE]: Initializing agent state...")
 
-        self.key, sensor_key, actor_key, critic_key, feature_extractor_key = jax.random.split(
-            self.key, 5
+        self.key, sensor_key, actor_key, critic_key, feature_extractor_key, message_passer_key = (
+            jax.random.split(self.key, 6)
         )
 
         dummy_reset = self.env.reset(seed=0)
+
+        for k, v in dummy_reset.observations.items():
+            self.logger.debug(k, v.shape)
+
         sample_obs = self.obs_processor(dummy_reset.observations)[0]  # take first env
-        sensor_params = self.sensor.init(sensor_key, sample_obs)
-        feature_extractor_params = self.feature_extractor.init(feature_extractor_key, sample_obs)
-        actor_params = self.actor.init(actor_key, self.sensor.apply(sensor_params, sample_obs))
-        critic_params = self.critic.init(
-            critic_key, self.feature_extractor.apply(feature_extractor_params, sample_obs)
+
+        self.logger.debug(f"[_init_agent_state] sample_obs: {sample_obs.shape}")
+        self.obs_mean = jnp.zeros((sample_obs.shape[-1],))
+        self.obs_var = jnp.ones((sample_obs.shape[-1],))
+        self.obs_count = 1e-4
+        self.logger.debug(f"[_init_agent_state] obs_mean: {self.obs_mean.shape}")
+        self.logger.debug(f"[_init_agent_state] obs_var: {self.obs_var.shape}")
+
+        self.logger.debug(f"[_init_agent_state]: Needed copies: {self.needed_copies}")
+        sensor_keys = jax.random.split(sensor_key, self.needed_copies)
+        actor_keys = jax.random.split(actor_key, self.needed_copies)
+
+        # (needed_copies, X)
+        sensor_params = jax.vmap(lambda k: self.sensor.init(k, sample_obs))(sensor_keys)
+        self.logger.debug(
+            f"[_init_agent_state] sensor_params: {jax.tree.map(lambda x: x.shape, sensor_params)}"
+        )
+
+        single_sensor_param = jax.tree.map(lambda x: x[0], sensor_params)
+        self.logger.debug(
+            f"[_init_agent_state] single_sensor_param: {
+                jax.tree.map(lambda x: x.shape, single_sensor_param)
+            }"
+        )
+
+        sensor_params_sample = self.sensor.apply(single_sensor_param, sample_obs)
+        self.logger.debug(
+            f"[_init_agent_state] sensor_params_sample shape: {sensor_params_sample.shape}"
+        )
+
+        actor_params = jax.vmap(lambda k: self.actor.init(k, sensor_params_sample))(actor_keys)
+        self.logger.debug(
+            f"[_init_agent_state] actor_params: {jax.tree.map(lambda x: x.shape, actor_params)}"
+        )
+
+        message_passer_params = {}
+        if self.morph_mode != MorphMode.CENTRALIZED:
+            assert self.message_passer is not None, "decentralized modes require a message passer"
+
+            message_passer_params = self.message_passer.init(
+                message_passer_key,
+                self.sensor.apply(single_sensor_param, sample_obs),
+            )
+            self.logger.debug(
+                f"[_init_agent_state] message_passer_params: {
+                    jax.tree.map(lambda x: x.shape, message_passer_params)
+                }"
+            )
+
+        flat_obs = sample_obs.reshape(-1)  # BECAUSE 1 centralized critic
+        self.logger.debug(f"[_init_agent_state] flat_obs: {flat_obs.shape}")
+
+        feature_extractor_params = self.feature_extractor.init(feature_extractor_key, flat_obs)
+        self.logger.debug(
+            f"[_init_agent_state] feature_extractor_params: {
+                jax.tree.map(lambda x: x.shape, feature_extractor_params)
+            }"
+        )
+
+        critic_input = self.feature_extractor.apply(feature_extractor_params, flat_obs)
+        self.logger.debug(f"[_init_agent_state] critic_input: {critic_input.shape}")
+
+        critic_params = self.critic.init(critic_key, critic_input)
+        self.logger.debug(
+            f"[_init_agent_state] critic_params: {jax.tree.map(lambda x: x.shape, critic_params)}"
         )
 
         return TrainState.create(
             apply_fn=None,
             params=asdict(
-                AgentParams(sensor_params, actor_params, critic_params, feature_extractor_params)
+                AgentParams(
+                    sensor_params,
+                    actor_params,
+                    critic_params,
+                    feature_extractor_params,
+                    message_passer_params,
+                )
             ),
             tx=optax.chain(
                 optax.clip_by_global_norm(self.ppo.max_grad_norm),
@@ -569,7 +788,7 @@ class PPOTrainer:
     def _step(self, env_state, next_obs, next_done, iteration: int) -> tuple:
         if iteration == 1:
             self.logger.log_non_interactive(f"Starting first rollout (JIT): {time.ctime()}")
-
+        self.logger.debug(f"[_step] next_obs (in): {next_obs.shape}")
         (
             self.agent_state,
             self.episode_stats,
@@ -581,12 +800,12 @@ class PPOTrainer:
             terminated_any,
             truncated_any,
         ) = self._rollout(env_state, next_obs, next_done)
-
+        self.logger.debug(f"[_step] next_obs (post-rollout): {next_obs.shape}")
         if iteration == 1:
             self.logger.log_non_interactive(f"First rollout completed: {time.ctime()}")
 
         storage = self._compute_gae(storage, next_obs, next_done)
-
+        self.logger.debug(f"[_step] storage.obs (post-gae): {storage.obs.shape}")
         if iteration == 1:
             self.logger.log_non_interactive(f"Starting first PPO update (JIT): {time.ctime()}")
 
@@ -684,8 +903,11 @@ class PPOTrainer:
                 self._eval_fn = build_eval_rollout_fn(
                     env=self.env,
                     obs_processor=self.obs_processor,
-                    sensor_apply=self.sensor.apply,
-                    actor_apply=self.actor.apply,
+                    sensor_apply=lambda p, x: apply_per_node(self.sensor, p, x),
+                    actor_apply=lambda p, x: apply_per_node(self.actor, p, x),
+                    message_passer_apply=(
+                        None if self.message_passer is None else self.message_passer.apply
+                    ),
                     action_low=self._action_low,
                     action_high=self._action_high,
                     reward_fn=reward_fn,
@@ -718,7 +940,10 @@ class PPOTrainer:
         self.logger.log_non_interactive(f"Initial reset started: {time.ctime()}")
 
         env_state = self.env.reset(seed=self.experiment.seed)
+
         next_obs = self.obs_processor(env_state.observations)
+        self.logger.debug(f"[train] next_obs: {next_obs.shape}")
+
         next_done = jnp.zeros(self.ppo.num_envs, dtype=jnp.bool_)
 
         self.logger.log_non_interactive(f"Initial reset completed: {time.ctime()}")

src/brittle_star_project/utils/__init__.py

@@ -0,0 +1,3 @@
+from .logged_jit import logged_jit
+
+__all__ = ["logged_jit"]

src/brittle_star_project/utils/logged_jit.py

@@ -0,0 +1,17 @@
+import jax
+from experiment_logger import get_logger
+
+
+def logged_jit(fn, **jit_kwargs):
+    logger = get_logger()
+    name = getattr(fn, "__name__", getattr(fn, "__qualname__", repr(fn)))
+
+    def decorator(func):
+        def traced_func(*args, **kwargs):
+            logger.debug(f"[JIT] Compiling {name}...")
+            return func(*args, **kwargs)
+
+        jitted = jax.jit(traced_func, **jit_kwargs)
+        return jitted
+
+    return decorator(fn)