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