feat(obs_processor): rewrote split method, todo: testing + code cleanup

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

scripts/train.py

@@ -42,7 +42,7 @@ def main(dict_cfg: DictConfig):
         base_dir=os.path.dirname(run_dir),
     )
     logger = get_logger()
-    logger.set_level(logging.INFO)
+    logger.set_level(logging.DEBUG)
     logger.info(f"Hydra-initialized run: {run_name}")
     logger.info(f"Output directory: {run_dir}")
 

simulate.sh

@@ -0,0 +1,9 @@
+#!/usr/bin/env bash
+
+path=$1
+
+uv run scripts/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/environment/obs_processing.py

@@ -102,61 +102,47 @@ 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
         for key, arr in obs.items():
-            if key not in ordered_keys or arr.size == 0:
+            if arr.size == 0:
                 continue
 
-            logger.debug(f"[INPUT] {key}: {arr.shape}")
-
             if arr.ndim == 0:
                 arr = arr.reshape(1)
 
-            # -------- CENTRALIZED --------
-            if morph_mode == MorphMode.CENTRALIZED:
-                output[key] = arr.reshape(1, -1)
-                # TODO: padding for centralized
-                continue
-
-            # -------- SEGMENTS --------
             if key in _SEGMENT_SCALED_KEYS:
                 per_agent = []
-
+                for i, _ in enumerate(agent_indices):
-                for i, agent_id in enumerate(agent_indices):
                     idx = segment_indices[i]
-                    taken = jnp.take(arr, idx, axis=0)  # (segs, ...)
+                    taken = jnp.take(arr, idx, axis=0)
-                    logger.debug(f"WHY {taken.shape}")
-
-                    # pad to 4 (segments per arm?)
                     pad_len = 4 - taken.shape[0]
-                    padded = jnp.pad(taken, [(0, pad_len)] + [(0, 0)] * (taken.ndim - 1))
+                    padded = jnp.pad(taken, [(9, pad_len)] + [(0, 0)] * (taken.ndim - 1))
 
                     per_agent.append(padded.reshape(-1))
-
+                arr = jnp.stack(per_agent)
-                out = jnp.stack(per_agent)
-
-            # -------- JOINTS --------
             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)  # (joint_n, ...)
+                    taken = jnp.take(arr, idx, axis=0)
-                    # 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))
-
+                arr = jnp.stack(per_agent)
-                out = jnp.stack(per_agent)
-
-            # -------- GLOBAL --------
             else:
-                out = jnp.repeat(arr[None, :], num_agents, axis=0)
+                arr = jnp.repeat(arr[None, :], num_agents, axis=0)
 
-            logger.debug(f"[OUTPUT] {key}: {out.shape}")
+            if morph_mode == MorphMode.CENTRALIZED:
-            output[key] = out
+                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
 

src/brittle_star_project/trainers/PPOTrainer.py

@@ -277,7 +277,6 @@ def apply_shared(net, params, x):
     return jax.vmap(lambda xi: net.apply(params, xi))(x_flattened)
 
 
-# TODO: update to work with extra dimension + message passing
 def _rollout_jit(
     agent_state,
     episode_stats,