feat: message passing step added to architecture to support decentralized network architectures

adds support for:
* fully connected
* ring
architectures

configs/architecture/decentralized.yaml

@@ -26,7 +26,7 @@ critic:
   activation: "tanh"
 
 # Synchronous message-passing rounds per control step
-message_passing_steps: 1
+message_passing_steps: 4
 
 # Connectivity topology (e.g., ring, fully_connected)
-topology_type: "ring"
+topology_type: "fully_connected"

configs/environment/dir_loc_further.yaml

@@ -0,0 +1,2 @@
+simulation_time: 50000.0
+target_distance: 3.0

configs/experiment/long_2arm.yaml

@@ -0,0 +1,6 @@
+# Testing chicken dinner 4 but further distance.
+
+exp_name: "long2arm"
+seed: 123
+torch_deterministic: true
+cuda: true

configs/morphology/2_arms_decentralized.yaml

@@ -0,0 +1,6 @@
+# 2 Arms Morphology Configuration
+
+segments_per_arm: [4, 0, 4, 0, 0]
+use_p_control: true
+use_torque_control: false
+morph_mode: FULLY_CONNECTED

configs/morphology/5_arms_damaged.yaml

@@ -0,0 +1,6 @@
+# 5 Arms Full Morphology Configuration
+# Baseline 5-arm brittle star.
+
+segments_per_arm: [4, 4, 0, 4, 4]
+use_p_control: true
+use_torque_control: false

configs/morphology/5_arms_full_fullconnected.yaml

@@ -0,0 +1,7 @@
+# 5 Arms Full Morphology Configuration
+# Baseline 5-arm brittle star.
+
+segments_per_arm: [4, 4, 4, 4, 4]
+use_p_control: true
+use_torque_control: false
+morph_mode: FULLY_CONNECTED

configs/ppo/chickendinnerwinner.yaml

@@ -0,0 +1,16 @@
+anneal_lr: true
+clip_coef: 0.2
+clip_vloss: true
+ent_coef: 0.001
+gae_lambda: 0.95
+gamma: 0.99
+learning_rate: 0.0001
+max_grad_norm: 0.5
+norm_adv: true
+num_envs: 32
+num_minibatches: 32
+num_steps: 64
+target_kl: 0.02
+total_timesteps: 12288000
+update_epochs: 4
+vf_coef: 1.0

configs/ppo/smoke_test.yaml

@@ -2,9 +2,9 @@
 # Lower timestep count for quick iterations/testing.
 
 learning_rate: 0.0005
-total_timesteps: 65536
+total_timesteps: 1024
-num_envs: 512
+num_envs: 32
-num_steps: 128
+num_steps: 32
 anneal_lr: true
 gamma: 0.99
 gae_lambda: 0.95

experiments/debug-experiment-10042026/used_variables.md

@@ -1,107 +0,0 @@
-## Default envconfig
-task: Task = Task.DIRECTED_LOCOMOTION
-simulation_time: float = 500.0
-num_physics_steps_per_control_step: int = 10
-time_scale: int = 2
-camera_ids: list[int] = field(default_factory=lambda: [0, 1])
-render_size: tuple[int, int] = (480, 640)
-joint_randomization_noise_scale: float = 0.0
-target_distance: float = 3.0
-light_perlin_noise_scale: int = 0
-
-
-## Default ppoargs
-seed: int = 1
-torch_deterministic: bool = True
-cuda: bool = True
-track: bool = False
-checkpoint_frequency: int = 100
-learning_rate: float = 2.5e-4
-anneal_lr: bool = True
-gamma: float = 0.99
-gae_lambda: float = 0.95
-update_epochs: int = 4
-norm_adv: bool = True
-clip_vloss: bool = True
-max_grad_norm: float = 0.5
-target_kl: float | None = None
-batch_size: int = 0
-minibatch_size: int = 0
-num_iterations: int = 0
-
-## Used config file: (hpc/debug.yaml)
-exp_name: "debug-experiment"
-seed: 42
-track: true
-wandb_project_name: "Let's-find-that-bug"
-wandb_entity: "SEL3-2026-Groep-4"
-run_dir: "/data/gent/465/vsc46589"
-num_envs: 32
-num_steps: 32
-num_minibatches: 32
-total_timesteps: 409600 
-num_arms: 2
-cuda: true
-
-ent_coef: 0.005
-vf_coef: 1.0
-clip_coef: 0.2
-
-anneal_lr: true
-learning_rate: 0.0003
-
-## Arena config:
-size: tuple[float, float] = (10.0, 5.0)
-sand_ground_color: bool = True
-attach_target: bool = True
-wall_height: float = 1.5
-wall_thickness: float = 0.1
-
-## Morphology:
-num_segments_per_arm: int = 4
-use_p_control: bool = True
-use_torque_control: bool = False
-
-## MLPs:
-### Sensor & Feature_extractor:
-Both with 3 layers of 300 neurons per layer.
-
-class GenericDenseLayersWithActivation(nn.Module):
-    layer_sizes: Sequence[int] = field(default_factory=lambda: [64, 64])
-    activation: Callable = nn.tanh
-
-    @nn.compact
-    def __call__(self, x):
-        for size in self.layer_sizes:
-            x = nn.Dense(size, kernel_init=orthogonal(jnp.sqrt(2)))(x)
-            x = self.activation(x)
-        return x
-
-### Actor:
-class Actor(nn.Module):
-    action_dim: int
-    @nn.compact
-    def __call__(self, x):
-        mean = nn.Dense(self.action_dim, kernel_init=orthogonal(0.01), bias_init=constant(0.0))(x)
-        log_std = self.param("log_std", nn.initializers.zeros, (self.action_dim,))
-        return mean, log_std
-
-### Critic:
-class OneDenseLayerMLP(nn.Module):
-    @nn.compact
-    def __call__(self, x):
-        return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x)
-
-### Observations:
-_ALLOWED_OBS_KEYS = {
-    "joint_position",
-    "joint_velocity",
-    "joint_actuator_force",
-    "actuator_force",
-    "disk_position",
-    "disk_rotation",
-    "disk_linear_velocity",
-    "disk_angular_velocity",
-    "unit_xy_direction_to_target",
-    "xy_distance_to_target",
-}

scripts/simulate.py

@@ -17,12 +17,14 @@ import numpy as np
 from omegaconf import DictConfig, OmegaConf
 import yaml
 
+import jax.numpy as jnp
+
 from brittle_star_project import Backend, BrittleStarEnv, BrittleStarEnvFactory
 from brittle_star_project.configs.main_config import BrittleStarConfig
 from brittle_star_project.configs.register_configs import register_configs
 from brittle_star_project.environment.padded_obs_wrapper import compute_padding_masks
 from brittle_star_project.environment.obs_processing import create_obs_processor
-from brittle_star_project.environment.env_config import MorphologyConfig
+from brittle_star_project.environment.env_config import MorphMode, MorphologyConfig
 
 from brittle_star_project.evaluation.checkpoint import load_metadata, metadata_to_configs
 from brittle_star_project.evaluation.policy import PolicyAgent
@@ -32,6 +34,7 @@ from brittle_star_project.evaluation.video import (
     create_evaluation_dir,
     save_evaluation_metadata,
 )
+from brittle_star_project.MLPs.adjancency_builder import build_adjacency
 
 
 @hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
@@ -78,9 +81,34 @@ def main(dict_cfg: DictConfig) -> None:
         segments_per_arm=env_morphology.segments_per_arm,
         reference_segments_per_arm=training.morphology.segments_per_arm,
     )
+
+    segs_per_arm = jnp.array(env_morphology.segments_per_arm)
+
+    needed_copies = 0
+    agent_indices = [0, 1, 2, 3, 4]
+    match env_morphology.morph_mode:
+        case MorphMode.CENTRALIZED:
+            needed_copies = 1
+        case MorphMode.FULLY_CONNECTED | MorphMode.RING:
+            agent_mask = segs_per_arm > 0
+            agent_indices = jnp.where(agent_mask)[0]
+            needed_copies = jnp.where(segs_per_arm > 0, 1, 0).sum().item()
+        case MorphMode.SEGMENT:
+            agent_mask = segs_per_arm > 0
+            agent_indices = jnp.where(agent_mask)[0]
+            needed_copies = jnp.where(segs_per_arm > 0, 1, 0).sum().item()
+            needed_copies = (segs_per_arm.sum() + jnp.where(segs_per_arm > 0, 1, 0).sum()).item()
+
+    num_arms = jnp.where(segs_per_arm > 0, 1, 0).sum().item()
+
     obs_processor = create_obs_processor(
         bounds_dict=training.obs_bounds.to_bounds_dict(),
         padding_masks=padding_masks,
+        needed_copies=needed_copies,
+        num_arms=num_arms,
+        morph_mode=env_morphology.morph_mode,
+        segments_per_arm=env_morphology.segments_per_arm,
+        agent_indices=agent_indices,
     )
 
     # 6. Build environment
@@ -104,11 +132,24 @@ def main(dict_cfg: DictConfig) -> None:
     state0 = env.reset(seed=seed)
 
     # Calculate the action dimension the model was trained with
-    trained_action_dim = sum(training.morphology.segments_per_arm) * 2
+    trained_action_dim = raw_env.action_space.shape[0] // needed_copies
 
     # 7. Load policy
+    message_passing_steps = (metadata.get("architecture", {}) or {}).get("message_passing_steps")
+    if message_passing_steps is None:
+        message_passing_steps = 4
+    message_passing_steps = int(message_passing_steps)
+
+    adj_matrix = None
+    if env_morphology.morph_mode != MorphMode.CENTRALIZED:
+        adj_matrix = build_adjacency(env_morphology.segments_per_arm, env_morphology.morph_mode)
+
     policy = PolicyAgent.from_checkpoint(
-        model_path, action_dim=trained_action_dim, obs_processor=obs_processor
+        model_path,
+        action_dim=trained_action_dim,
+        obs_processor=obs_processor,
+        message_passing_steps=message_passing_steps,
+        adj_matrix=adj_matrix,
     )
 
     # Convert the JAX boolean mask to a numpy array for easy indexing

scripts/simulate.sh

@@ -0,0 +1,9 @@
+#!/usr/bin/env bash
+
+path=$1
+
+uv run simulate.py \
+    simulation.model_path="$path"/final_model.flax \
+    simulation.record_video=True \
+    simulation.video_output_path=../vids/simulation.mp4 \
+    simulation.max_steps=10000

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
+    sensor_params: FrozenDict | dict
-    actor_params: flax.core.FrozenDict
+    actor_params: FrozenDict | dict
-    critic_params: flax.core.FrozenDict
+    critic_params: FrozenDict | dict
-    feature_extractor_params: flax.core.FrozenDict
+    feature_extractor_params: FrozenDict | dict
+    message_passer_params: FrozenDict | dict
 
 
 @jax.tree_util.register_dataclass
 @dataclass
 class Storage:
-    obs: jnp.array
+    obs: jnp.ndarray
-    actions: jnp.array
+    actions: jnp.ndarray
-    logprobs: jnp.array
+    logprobs: jnp.ndarray
-    dones: jnp.array
+    dones: jnp.ndarray
-    values: jnp.array
+    values: jnp.ndarray
-    advantages: jnp.array
+    advantages: jnp.ndarray
-    returns: jnp.array
+    returns: jnp.ndarray
-    rewards: jnp.array
+    rewards: jnp.ndarray
 
-    raw_actions: jnp.ndarray = None  # before clipping
+    raw_actions: jnp.ndarray | None = None  # before clipping
-    means: jnp.ndarray = None  # policy mean
+    means: jnp.ndarray | None = None  # policy mean
-    stds: jnp.ndarray = None  # policy std
+    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_returns: jnp.ndarray
-    episode_lengths: jnp.array
+    episode_lengths: jnp.ndarray
-    returned_episode_returns: jnp.array
+    returned_episode_returns: jnp.ndarray
-    returned_episode_lengths: jnp.array
+    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:
+    def _split_to_agents(obs: dict, morph_mode) -> dict:
-        padded = {}
+        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:
+            if arr.size == 0:
-                padded_arr = jnp.zeros(padding_masks["target_size_2x"], dtype=arr.dtype)
+                continue
-                padded[key] = padded_arr.at[padding_masks["mask_2x"]].set(arr)
+
-            elif key in _SEGMENT_SCALED_KEYS:
+            if arr.ndim == 0:
-                padded_arr = jnp.zeros(padding_masks["target_size_1x"], dtype=arr.dtype)
+                arr = arr.reshape(1)
-                padded[key] = padded_arr.at[padding_masks["mask_1x"]].set(arr)
+
+            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
+                arr = jnp.repeat(arr[None, :], num_agents, axis=0)
-        return padded
+
+            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",
+        Input:
-            "joint_actuator_force",
+            key -> (num_arms, feat_per_key)
-            "joint_position",
+
-            "joint_velocity",
+        Output:
-            "robot_direction_to_target",
+            (num_arms, total_features)
-            "segment_contact",
+        """
-        ]
         values = []
+
         for key in ordered_keys:
-            if key in obs:
+            if key not in obs:
-                arr = jnp.asarray(obs[key]).flatten()
+                continue
-                if arr.size > 0:
+
+            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)
+
+        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 = _split_to_agents(processed, morph_mode)
-            processed = _pad_features(processed)
+        flat = _flatten_features(processed)  # (num_arms, total_feat)
-        return _flatten_features(processed)
+
+        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,6 +68,9 @@ 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
@@ -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]
+        obs = self._obs_processor(batched_obs)
-        hidden = self._sensor_apply(self._params["sensor_params"], obs)
+
-        mean, _log_std = self._actor_apply(self._params["actor_params"], hidden)
+        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

@@ -114,18 +114,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,
+                sensor_apply=sensor_apply,
-                actor_apply=actor.apply,
+                actor_apply=actor_apply,
-                critic_apply=critic.apply,
+                critic_apply=critic_apply,
-                feature_extractor_apply=feature_extractor.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, (
+                return agent_state, (loss, pg_loss, v_loss, entropy_loss, approx_kl)
-                    loss,
-                    pg_loss,
-                    v_loss,
-                    entropy_loss,
-                    approx_kl,
-                    grads,
-                )
 
             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)
+    logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi))
-    entropy = (0.5 + 0.5 * jnp.log(2 * jnp.pi) + log_std).sum(-1)
+    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
+@logged_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
 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,
+    sensor: nn.Module,
-    feature_extractor: GenericDenseLayersWithActivation,
+    feature_extractor: nn.Module,
-    actor: Actor,
+    actor: nn.Module,
-    critic: OneDenseLayerMLP,
+    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)
+    # (B, n_nodes, feat)
-    hidden_critic = feature_extractor.apply(
+    hidden = apply_per_node(sensor, agent_state.params["sensor_params"], next_obs)
-        agent_state.params["feature_extractor_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,
+    sensor: nn.Module,
-    feature_extractor: GenericDenseLayersWithActivation,
+    feature_extractor: nn.Module,
-    actor: Actor,
+    actor: nn.Module,
-    critic: OneDenseLayerMLP,
+    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(
+    flat_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
+        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,
+    sensor: nn.Module,
-    feature_extractor: GenericDenseLayersWithActivation,
+    feature_extractor: nn.Module,
-    actor: Actor,
+    actor: nn.Module,
-    critic: OneDenseLayerMLP,
+    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.logger.debug(f"needed copies = {self.needed_copies}")
-        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)
 
         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, sensor_key, actor_key, critic_key, feature_extractor_key, message_passer_key = (
-            self.key, 5
+            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)
+        self.logger.debug(f"[_init_agent_state] sample_obs: {sample_obs.shape}")
-        actor_params = self.actor.init(actor_key, self.sensor.apply(sensor_params, sample_obs))
+        self.obs_mean = jnp.zeros((sample_obs.shape[-1],))
-        critic_params = self.critic.init(
+        self.obs_var = jnp.ones((sample_obs.shape[-1],))
-            critic_key, self.feature_extractor.apply(feature_extractor_params, sample_obs)
+        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,
+                    sensor_apply=lambda p, x: apply_per_node(self.sensor, p, x),
-                    actor_apply=self.actor.apply,
+                    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)

tests/test_adjacency.py

@@ -0,0 +1,98 @@
+import jax.numpy as jnp
+import numpy as np
+
+from brittle_star_project.MLPs import build_adjacency
+from brittle_star_project.environment.env_config import MorphMode
+
+
+def assert_symmetric(adj):
+    assert jnp.all(adj == adj.T)
+
+
+def test_centralized():
+    adj = build_adjacency([4, 4, 4, 4, 4], MorphMode.CENTRALIZED)
+
+    assert adj.shape == (1, 1)
+    assert adj[0, 0] == 1
+
+
+def test_fully_connected():
+    adj = build_adjacency([4, 4, 4, 4, 4], MorphMode.FULLY_CONNECTED)
+
+    assert adj.shape == (5, 5)
+    assert jnp.all(adj == 1)
+
+
+def test_ring():
+    adj = build_adjacency([4, 4, 4, 4, 4], MorphMode.RING)
+
+    assert adj.shape == (5, 5)
+    assert_symmetric(adj)
+
+    # each node should connect to itself + 2 neighbors
+    for node in range(5):
+        assert adj[node, node] == 1
+        assert jnp.sum(adj[node]) == 3
+        neighbor1 = (node - 1) % 5
+        neighbor2 = (node + 1) % 5
+        assert adj[neighbor1, node] == 1
+        assert adj[node, neighbor2] == 1  # Symmetrical
+
+
+def test_segment_structure():
+    segments = [4, 4, 4, 4, 4]
+    adj = build_adjacency(segments, MorphMode.SEGMENT)
+
+    num_arms = 5
+    num_segments = sum(segments)
+    num_nodes = num_arms + num_segments
+
+    assert adj.shape == (num_nodes, num_nodes)
+
+    # --- ring connectivity ---
+    for i in range(num_arms):
+        assert adj[i, i] == 1
+        assert adj[i, (i - 1) % num_arms] == 1
+        assert adj[i, (i + 1) % num_arms] == 1
+
+    # --- segment chain checks ---
+    offset = num_arms
+    for arm in range(5):
+        for i in range(4):
+            node = offset + arm * 4 + i
+
+            # self
+            assert adj[node, node] == 1
+
+            # chain neighbors
+            if i > 0:
+                assert adj[node, node - 1] == 1
+            if i < 3:
+                assert adj[node, node + 1] == 1
+
+    # --- ring ↔ segment connections ---
+    for arm in range(5):
+        first_seg = num_arms + arm * 4
+        assert adj[arm, first_seg] == 1
+        assert adj[first_seg, arm] == 1
+
+    save_adj(adj)
+
+
+def save_adj(adj, name="adjacency_debug.txt"):
+    a = np.array(adj)
+
+    with open(name, "w") as f:
+        f.write("\nAdjacency matrix:\n")
+        f.write("   " + " ".join([f"{i:2d}" for i in range(a.shape[0])]) + "\n")
+
+        for i, row in enumerate(a):
+            line = f"{i:2d}  " + "  ".join(["█" if x > 0 else "." for x in row])
+            f.write(line + "\n")
+
+
+def test_no_isolated_nodes():
+    adj = build_adjacency([4, 4, 4, 4, 4], MorphMode.SEGMENT)
+
+    # no node should be completely isolated
+    assert jnp.all(jnp.sum(adj, axis=0) > 0)

tests/test_network_shapes.py

@@ -1,5 +1,6 @@
 import jax
 import jax.numpy as jnp
+from brittle_star_project.environment.env_config import MorphMode
 from brittle_star_project.environment.padded_obs_wrapper import (
     compute_padding_masks,
 )
@@ -20,14 +21,24 @@ def test_centralized_forward_pass_with_padding():
         "segment_contact": jnp.zeros((batch_size, 14)),
     }
 
+    segments_per_arm = jnp.array((4, 0, 4, 2, 4))
+    num_arms = jnp.where(segments_per_arm > 0, 1, 0).sum().item()
+
     # 2. Process and Pad Observation
-    masks = compute_padding_masks(segments_per_arm=(4, 0, 4, 2, 4))
+    masks = compute_padding_masks(segments_per_arm=list(segments_per_arm))
-    obs_processor = create_obs_processor(bounds_dict={}, padding_masks=masks)
+    obs_processor = create_obs_processor(
+        bounds_dict={},
+        needed_copies=1,
+        num_arms=num_arms,
+        padding_masks=masks,
+        morph_mode=MorphMode.CENTRALIZED,
+        segments_per_arm=segments_per_arm,
+    )
     global_state = obs_processor(amputated_obs)
 
-    # 40 + 40 + 20 = 100 dimensions
+    # 40 + 40 + 20 + padding = 145 dimensions
-    assert global_state.shape == (batch_size, 100), (
+    assert global_state.shape == (batch_size, 1, 145), (
-        f"Expected global state shape (2, 100), got {global_state.shape}"
+        f"Expected global state shape (2, 1, 145), got {global_state.shape}"
     )
 
     # 4. Initialize dummy networks (40 actuators for the max morphology output)
@@ -45,9 +56,11 @@ def test_centralized_forward_pass_with_padding():
     action_mean, action_log_std = actor.apply(actor_params, global_state)
     value = critic.apply(critic_params, global_state)
 
-    assert action_mean.shape == (batch_size, 40), f"Actor mean shape mismatch: {action_mean.shape}"
+    assert action_mean.shape == (batch_size, 1, 40), (
+        f"Actor mean shape mismatch: {action_mean.shape}"
+    )
     assert action_log_std.shape == (40,), f"Actor log_std shape mismatch: {action_log_std.shape}"
-    assert value.shape == (batch_size, 1) or value.shape == (batch_size,), (
+    assert value.shape == (batch_size, 1, 1) or value.shape == (batch_size,), (
         f"Critic value shape mismatch: {value.shape}"
     )
 

tests/test_obs_processor.py

@@ -0,0 +1,119 @@
+import jax
+import jax.numpy as jnp
+
+from brittle_star_project.environment.obs_processing import create_obs_processor
+from brittle_star_project.environment.env_config import MorphMode, ObservationBoundsConfig
+
+
+obs_bounds = ObservationBoundsConfig().to_bounds_dict()
+
+"""
+Test for obs_processor.
+
+Centralized: 40 features per agent:
+    disk_z_tilt              → scalar → reshaped to (1,)  → 1 feat
+    joint_actuator_force     → 8 joints padded to 8       → 8 feat
+    joint_position           → 8 joints padded to 8       → 8 feat
+    joint_velocity           → 8 joints padded to 8       → 8 feat
+    robot_direction_to_target→ (x, y)                     → 2 feat
+    segment_contact          → 4 segs, pre-padded by 9    → 13 feat (9 leading + 4)
+"""
+
+NUM_ARMS = 5
+SEGS_PER_ARM = 4  # healthy segments per arm
+JOINTS_PER_SEG = 2  # from _build_joint_indices: segs * 2
+
+SEGS_HEALTHY = [4, 4, 4, 4, 4]
+SEGS_DAMAGED = [4, 4, 4, 4, 0]  # arm 4 fully disabled
+SEGS_DAMAGED_2 = [4, 0, 4, 2, 4]  # arm 3 fully disabled
+AGENT_INDICES = [0, 1, 2, 3, 4]
+
+FEAT_PER_AGENT = 1 + 8 + 8 + 8 + 2 + 13  # = 40
+
+
+def make_obs(segs_per_arm: list[int]) -> dict:
+    total_segs = sum(segs_per_arm)
+    total_joints = JOINTS_PER_SEG * total_segs
+
+    return {
+        "actuator_force": jnp.ones(total_joints),
+        "disk_angular_velocity": jnp.zeros(3),
+        "disk_linear_velocity": jnp.zeros(3),
+        "disk_position": jnp.zeros(3),
+        "disk_rotation": jnp.array([0.1, 0.1, 0.5]),  # (roll, pitch, yaw)
+        "joint_actuator_force": jnp.full(total_joints, 1.0),
+        "joint_position": jnp.full(total_joints, 0.5),
+        "joint_velocity": jnp.full(total_joints, 2.0),
+        "segment_contact": jnp.ones(total_segs),
+        "tendon_position": jnp.zeros(0),
+        "tendon_velocity": jnp.zeros(0),
+        "unit_xy_direction_to_target": jnp.array([1.0, 0.0]),
+        "xy_distance_to_target": jnp.array([3.5]),
+    }
+
+
+def batch_obs(obs: dict):
+    return jax.tree_util.tree_map(lambda x: x[None, :], obs)
+
+
+def make_processor(morph_mode: MorphMode, needed_copies: int, segments_per_arm: list[int]):
+    return create_obs_processor(
+        bounds_dict=obs_bounds,
+        num_arms=NUM_ARMS,
+        needed_copies=needed_copies,
+        morph_mode=morph_mode,
+        segments_per_arm=segments_per_arm,
+        agent_indices=AGENT_INDICES,
+    )
+
+
+def test_centralized_no_damage():
+    proc = make_processor(MorphMode.CENTRALIZED, 1, SEGS_HEALTHY)
+    obs = make_obs(SEGS_HEALTHY)
+    obs = batch_obs(obs)
+    global_state = proc(obs)
+
+    # shape test
+    assert global_state.shape == (1, 1, 188)
+
+    # TODO: more?
+
+
+def test_centralized_damaged_1_arm():
+    proc = make_processor(MorphMode.CENTRALIZED, 1, SEGS_DAMAGED)
+    obs = make_obs(SEGS_DAMAGED)
+    obs = batch_obs(obs)
+    global_state = proc(obs)
+
+    # shape test
+    assert global_state.shape == (1, 1, 188)
+
+
+def test_centralized_damaged_2_arms():
+    proc = make_processor(MorphMode.CENTRALIZED, 1, SEGS_DAMAGED_2)
+    obs = make_obs(SEGS_DAMAGED_2)
+    obs = batch_obs(obs)
+    global_state = proc(obs)
+
+    # shape test
+    assert global_state.shape == (1, 1, 188)
+
+
+def test_decentralized_fully_connected_no_damage():
+    proc = make_processor(MorphMode.FULLY_CONNECTED, NUM_ARMS, SEGS_HEALTHY)
+    obs = make_obs(SEGS_HEALTHY)
+    obs = batch_obs(obs)
+    global_state = proc(obs)
+
+    # shape test
+    assert global_state.shape == (1, NUM_ARMS, FEAT_PER_AGENT)
+
+
+def test_decentralized_fully_connected_damaged_1_arm():
+    proc = make_processor(MorphMode.FULLY_CONNECTED, NUM_ARMS, SEGS_DAMAGED)
+    obs = make_obs(SEGS_DAMAGED)
+    obs = batch_obs(obs)
+    global_state = proc(obs)
+
+    # shape test
+    assert global_state.shape == (1, NUM_ARMS, FEAT_PER_AGENT)

tests/test_target_direction.py

@@ -1,6 +1,7 @@
 import jax.numpy as jnp
 
 from brittle_star_project.configs.main_config import BrittleStarConfig
+from brittle_star_project.environment.env_config import MorphMode
 from brittle_star_project.environment.env_types import Backend
 from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
 from brittle_star_project.environment.obs_processing import create_obs_processor
@@ -43,8 +44,16 @@ def test_processor_converts_to_egocentric_direction():
     cfg = BrittleStarConfig()
     env = BrittleStarJaxEnvWrapper.default(num_envs=1, backend=Backend.MJX)
 
+    segments_per_arm = jnp.array((4, 4, 4, 4, 4))
+    num_arms = jnp.where(segments_per_arm > 0, 1, 0).sum().item()
+
     obs_processor = create_obs_processor(
-        bounds_dict=cfg.obs_bounds.to_bounds_dict(), padding_masks=env.padding_masks
+        bounds_dict=cfg.obs_bounds.to_bounds_dict(),
+        needed_copies=1,
+        num_arms=num_arms,
+        padding_masks=env.padding_masks,
+        morph_mode=MorphMode.CENTRALIZED,
+        segments_per_arm=segments_per_arm,
     )
 
     env_state = env.reset(seed=42)
@@ -66,10 +75,13 @@ def test_processor_converts_to_egocentric_direction():
     processed_2 = obs_processor(dummy_obs_2)
 
     # Find the indices of the elements that changed
-    diff_array = jnp.abs(processed_1[0] - processed_2[0])
+    diff_array = jnp.abs(processed_1[0, 0] - processed_2[0, 0])
     changed_indices = jnp.where(diff_array > 1e-4)[0]
 
-    local_target = processed_1[0, changed_indices]
+    # (143,)
+    local_target = processed_1[0, 0, changed_indices]
+
+    # (2,)
     expected_local_target = jnp.array([0.0, -1.0])
 
     assert jnp.sum(jnp.abs(local_target - expected_local_target)) < 1e-4, (