backup

src/brittle_star_project/environment/obs_processing.py

@@ -22,11 +22,23 @@ _SEGMENT_SCALED_KEYS = frozenset(
 
 def create_obs_processor(
     bounds_dict: Dict[str, Tuple[float, float]],
-    num_segments: int,
     num_arms: int,
+    needed_copies: int,
     padding_masks: Optional[Dict] = None,
     morph_mode: MorphMode = MorphMode.CENTRALIZED,
 ):
+    # 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",
+        ]
+    )
+
     def _add_derived_features(obs: dict) -> dict:
         new_obs = dict(obs)
         if "disk_rotation" in new_obs:
@@ -57,49 +69,61 @@ def create_obs_processor(
                 normalized[key] = arr
         return normalized
 
-    def _pad_features(obs: dict) -> dict:
+    # ndarray (agents, features)
+    def _pad_features(obs: dict, agent_count: int) -> dict:
         assert padding_masks is not None
 
         padded = {}
+
+        # arr shape = (agent_count, ...) TODO
         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_arr = jnp.zeros(
-                padded[key] = padded_arr.at[padding_masks["mask_2x"]].set(arr)
+                    (agent_count, 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_arr = jnp.zeros(
-                padded[key] = padded_arr.at[padding_masks["mask_1x"]].set(arr)
+                    (agent_count, padding_masks["target_size_1x"]), dtype=arr.dtype
+                )
+                padded[key] = padded_arr.at[:, padding_masks["mask_1x"]].set(arr)
             else:
                 padded[key] = arr
         return padded
 
-    # TODO
+    def _split_to_agents(obs: dict, agent_count: int) -> dict:
-    def _split_to_agents() -> dict:
+        """
-        return {}
+        Each observation type is now updated to (agent_count, feature_shape),
+        thus duplicating the observation for each agent
+        """
 
+        output = {}
+
+        for key, arr in obs.items():
+            if key not in ordered_keys or arr.size == 0:
+                continue
+
+            if arr.ndim == 0:
+                arr = arr.reshape(1)
+
+            output[key] = jnp.repeat(arr[None, :], agent_count, axis=0)
+
+        return output
+
+    # TODO: update
     def _flatten_features(obs: dict) -> jnp.ndarray:
-        ordered_keys = [
+        """
-            "disk_z_tilt",
+        Collapse all observations into a single array
-            "joint_actuator_force",
+        """
-            "joint_position",
-            "joint_velocity",
-            "robot_direction_to_target",
-            "segment_contact",
-        ]
 
         values = []
 
         for key in sorted(obs.keys()):
             if key not in ordered_keys:
                 continue
-            v = obs[key]
 
-            # skip empty arrays and scalars
+            # empty arrays handled in split_to_agents
-            if v.size == 0:
+            v = obs[key]  # (agent_count, feat)
-                continue
-
-            # reshape scalars to 1D array
-            if v.ndim == 0:
-                v = v.reshape(1)
 
             # -------- CENTRALIZED --------
             if morph_mode == MorphMode.CENTRALIZED:
@@ -124,11 +148,18 @@ def create_obs_processor(
         return jnp.concatenate(values, axis=-1)  # (agent, feat)
 
     def _process_single(obs_dict: dict) -> jnp.ndarray:
-        processed = _add_derived_features(obs_dict)  # key |--> (feat-count, feat-lengths, )
+        processed = _add_derived_features(obs_dict)  # key |--> (feat-count,)
-        processed = _normalize_features(processed)  # key |--> (feat-count, feat-lengths, )
+        processed = _normalize_features(processed)  # key |--> (feat-count,)
-        # TODO: split to agents # key |--> (agents, feat-count, feat-lengths)
+        # TODO: split to agents
+        processed = _split_to_agents(processed, needed_copies)  # key |--> (agents, feat-count)
         if padding_masks is not None:
-            processed = _pad_features(processed)  # key |--> (agents, feat-count, feat-lengths')
+            processed = _pad_features(
+                processed, agent_count=needed_copies
+            )  # key |--> (agents, feat-count')
+        for k, v in processed.items():
+            print(k, v.shape)
+
+        exit(1)
         return _flatten_features(processed)  # (agents, feat)
 
     return jax.jit(jax.vmap(_process_single))

src/brittle_star_project/trainers/PPOTrainer.py

@@ -363,7 +363,7 @@ class PPOTrainer:
         # Build the centralized observation processor: derive -> normalize -> pad -> flatten.
         self.obs_processor = create_obs_processor(
             bounds_dict=self.cfg.obs_bounds.to_bounds_dict(),
-            num_segments=self.num_segments,
+            needed_copies=self.needed_copies,
             num_arms=self.num_arms,
             morph_mode=self.morph_mode,
             padding_masks=self.env.padding_masks,

tests/test_network_shapes.py

@@ -22,14 +22,13 @@ def test_centralized_forward_pass_with_padding():
     }
 
     segments_per_arm = jnp.array((4, 0, 4, 2, 4))
-    num_segments = segments_per_arm.sum().item()
     num_arms = jnp.where(segments_per_arm > 0, 1, 0).sum().item()
 
     # 2. Process and Pad Observation
     masks = compute_padding_masks(segments_per_arm=list(segments_per_arm))
     obs_processor = create_obs_processor(
         bounds_dict={},
-        num_segments=num_segments,
+        needed_copies=1,
         num_arms=num_arms,
         padding_masks=masks,
         morph_mode=MorphMode.CENTRALIZED,

tests/test_target_direction.py

@@ -45,12 +45,11 @@ def test_processor_converts_to_egocentric_direction():
     env = BrittleStarJaxEnvWrapper.default(num_envs=1, backend=Backend.MJX)
 
     segments_per_arm = jnp.array((4, 4, 4, 4, 4))
-    num_segments = segments_per_arm.sum().item()
     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(),
-        num_segments=num_segments,
+        needed_copies=1,
         num_arms=num_arms,
         padding_masks=env.padding_masks,
         morph_mode=MorphMode.CENTRALIZED,