test: updated tests to allow padding

src/brittle_star_project/environment/obs_processing.py

@@ -102,9 +102,12 @@ def create_obs_processor(
         return normalized
 
     def _split_to_agents(obs: dict, morph_mode) -> dict:
-        # TODO: cleanup + MORE testing (works for centralized 5 arms + damaged arms)
         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 arr.size == 0:
                 continue
@@ -117,7 +120,7 @@ def create_obs_processor(
                 for i, _ in enumerate(agent_indices):
                     idx = segment_indices[i]
                     taken = jnp.take(arr, idx, axis=0)
-                    pad_len = 4 - taken.shape[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))
@@ -127,7 +130,7 @@ def create_obs_processor(
                 for i, _ in enumerate(agent_indices):
                     idx = joint_indices[i]
                     taken = jnp.take(arr, idx, axis=0)
-                    pad_len = 8 - taken.shape[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))

src/brittle_star_project/trainers/PPOTrainer.py

@@ -477,6 +477,8 @@ class PPOTrainer:
             agent_indices=self.agent_indices,
         )
 
+        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

tests/test_network_shapes.py

@@ -36,9 +36,9 @@ def test_centralized_forward_pass_with_padding():
     )
     global_state = obs_processor(amputated_obs)
 
-    # 40 + 40 + 20 = 100 dimensions
+    # 40 + 40 + 20 + padding = 145 dimensions
-    assert global_state.shape == (batch_size, 1, 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)

tests/test_obs_processor.py

@@ -25,6 +25,7 @@ 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
@@ -73,7 +74,7 @@ def test_centralized_no_damage():
     global_state = proc(obs)
 
     # shape test
-    assert global_state.shape == (1, 1, FEAT_PER_AGENT * NUM_ARMS)
+    assert global_state.shape == (1, 1, 188)
 
     # TODO: more?
 
@@ -85,7 +86,17 @@ def test_centralized_damaged_1_arm():
     global_state = proc(obs)
 
     # shape test
-    assert global_state.shape == (1, 1, FEAT_PER_AGENT * NUM_ARMS)
+    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():