fix: split pad dada

src/brittle_star_project/environment/obs_processing.py

@@ -24,25 +24,27 @@ _SEGMENT_SCALED_KEYS = frozenset(
 )
 
 
-def _build_joint_indices(segments_per_arm):
+def _build_joint_indices(segments_per_arm, indices_mlp):
     indices = []
     start = 0
-    for segs in segments_per_arm:
+    for i, segs in enumerate(segments_per_arm):
         # 2 joints per segment
-        count = segs * 2
+        if i in indices_mlp:
-        idx = jnp.arange(start, start + count)
+            count = segs * 2
-        indices.append(idx)
+            idx = jnp.arange(start, start + count)
-        start += count
+            indices.append(idx)
+            start += count
     return indices
 
 
-def _build_segment_indices(segments_per_arm):
+def _build_segment_indices(segments_per_arm, indices_mlp):
     indices = []
     start = 0
-    for segs in segments_per_arm:
+    for i, segs in enumerate(segments_per_arm):
-        idx = jnp.arange(start, start + segs)
+        if i in indices_mlp:
-        indices.append(idx)
+            idx = jnp.arange(start, start + segs)
-        start += segs
+            indices.append(idx)
+            start += segs
     return indices
 
 
@@ -53,6 +55,7 @@ def create_obs_processor(
     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(
@@ -65,8 +68,8 @@ def create_obs_processor(
             "segment_contact",
         ]
     )
-    segment_indices = _build_segment_indices(segments_per_arm)
+    segment_indices = _build_segment_indices(segments_per_arm, agent_indices)
-    joint_indices = _build_joint_indices(segments_per_arm)
+    joint_indices = _build_joint_indices(segments_per_arm, agent_indices)
 
     def _add_derived_features(obs: dict) -> dict:
         new_obs = dict(obs)
@@ -120,34 +123,64 @@ def create_obs_processor(
                 padded[key] = arr
         return padded
 
-    def _split_to_agents(obs: dict, morph_mode, segments_per_arm) -> dict:
+    def _split_to_agents(obs: dict, morph_mode) -> dict:
-        output = {}
+        total = 0
-        num_arms = len(segments_per_arm)
+        for k, v in obs.items():
+            if hasattr(v, "shape"):
+                size = v.size
+                logger11.info(f"[RAW] {k}: shape={v.shape}, size={size}")
+                total += size
+            else:
+                logger11.info(f"[RAW] {k}: non-array")
 
+        logger11.info(f"[RAW TOTAL FEATURES]: {total}")
+        output = {}
+        num_agents = needed_copies  # IMPORTANT: number of MLPs
         for key, arr in obs.items():
             if key not in ordered_keys or arr.size == 0:
                 continue
-
             logger11.info(f"[INPUT] {key}: {arr.shape}")
-
             if arr.ndim == 0:
                 arr = arr.reshape(1)
-
+            # -------- CENTRALIZED --------
             if morph_mode == MorphMode.CENTRALIZED:
-                out = arr.reshape(1, -1)
+                output[key] = arr.reshape(1, -1)
-                output[key] = out
                 continue
 
+            # -------- SEGMENTS --------
             if key in _SEGMENT_SCALED_KEYS:
-                per_agent = [jnp.take(arr, idx, axis=0) for idx in segment_indices]
+                per_agent = []
+
+                for i, agent_id in enumerate(agent_indices):
+                    idx = segment_indices[i]
+                    taken = jnp.take(arr, idx, axis=0)  # (segs, ...)
+                    logger11.info(f"WHY {taken.shape}")
+                    # pad to 4
+                    pad_len = 4 - taken.shape[0]
+                    padded = jnp.pad(taken, [(0, pad_len)] + [(0, 0)] * (taken.ndim - 1))
+
+                    per_agent.append(padded.reshape(-1))
+
                 out = jnp.stack(per_agent)
 
+            # -------- JOINTS --------
             elif key in _JOINT_SCALED_KEYS:
-                per_agent = [jnp.take(arr, idx, axis=0) for idx in joint_indices]
+                per_agent = []
+
+                for i, agent_id in enumerate(agent_indices):
+                    idx = joint_indices[i]
+                    taken = jnp.take(arr, idx, axis=0)  # (joint_n, ...)
+                    # pad to 8
+                    pad_len = 8 - taken.shape[0]
+
+                    padded = jnp.pad(taken, [(0, pad_len)] + [(0, 0)] * (taken.ndim - 1))
+                    per_agent.append(padded.reshape(-1))
+
                 out = jnp.stack(per_agent)
 
+            # -------- GLOBAL --------
             else:
-                out = jnp.repeat(arr[None, :], num_arms, axis=0)
+                out = jnp.repeat(arr[None, :], num_agents, axis=0)
 
             logger11.info(f"[OUTPUT] {key}: {out.shape}")
             output[key] = out
@@ -185,15 +218,10 @@ def create_obs_processor(
     def _process_single(obs_dict: dict) -> jnp.ndarray:
         processed = _add_derived_features(obs_dict)
         processed = _normalize_features(processed)
-        # morph_mode = MorphMode.FULLY_CONNECTED
+        processed = _split_to_agents(processed, morph_mode)
-        processed = _split_to_agents(processed, morph_mode, segments_per_arm)
-        # needed_copies = 5
-        if padding_masks is not None:
-            processed = _pad_features(processed, agent_count=needed_copies)
         flat = _flatten_features(processed)  # (num_arms, total_feat)
         logger11.info(f"[FLATTENED FINAL] shape: {flat.shape}")
         logger11.info(f"[PER AGENT] example row 0 shape: {flat[0].shape}")
-        exit(1)
         return _flatten_features(processed)  # (agents, feat)
 
     return jax.jit(jax.vmap(_process_single))

src/brittle_star_project/trainers/PPOTrainer.py

@@ -353,6 +353,7 @@ class PPOTrainer:
             self.feature_extractor,
             self.critic,
             self.needed_copies,
+            self.agent_indices,
         ) = self._init_agent()
 
         self.sensor.apply = logged_jit(self.sensor.apply)
@@ -368,6 +369,7 @@ class PPOTrainer:
             morph_mode=self.morph_mode,
             padding_masks=self.env.padding_masks,
             segments_per_arm=self.segments_per_arm,
+            agent_indices=self.agent_indices,
         )
 
         action_low = jnp.asarray(self.env.single_action_space.low, dtype=jnp.float32)
@@ -436,13 +438,18 @@ 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 = jnp.where(self.segments_per_arm > 0, 1, 0).sum().item()
                 needed_copies = (
                     self.segments_per_arm.sum() + jnp.where(self.segments_per_arm > 0, 1, 0).sum()
                 ).item()
@@ -462,7 +469,15 @@ class PPOTrainer:
 
         feature_extractor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300])
         critic = OneDenseLayerMLP()
-        return sensor, message_passer, actor, feature_extractor, critic, needed_copies
+        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...")