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/environment/obs_processing_tmp.py

@@ -0,0 +1,17 @@
+obs = {}
+segs = set()
+joints = set()
+
+for key, arr in obs.items():
+    if key in segs:
+        # padding using mask 1x
+        pass
+    elif key in joints:
+        # padding using mask 2x
+        pass
+    else:
+        # no padding
+        pass
+
+    # reshape according to segs, joints, or global
+    pass

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,

tmp.txt

@@ -0,0 +1,27 @@
+-- 5 arms, 1 disabled
+2026-05-08 14:21:57,035 - DEBUG - [INPUT] joint_actuator_force: (32,)
+2026-05-08 14:21:57,035 - DEBUG - [OUTPUT] joint_actuator_force: (32,)
+2026-05-08 14:21:57,035 - DEBUG - [INPUT] joint_position: (32,)
+2026-05-08 14:21:57,035 - DEBUG - [OUTPUT] joint_position: (32,)
+2026-05-08 14:21:57,036 - DEBUG - [INPUT] joint_velocity: (32,)
+2026-05-08 14:21:57,036 - DEBUG - [OUTPUT] joint_velocity: (32,)
+2026-05-08 14:21:57,036 - DEBUG - [INPUT] segment_contact: (16,)
+2026-05-08 14:21:57,036 - DEBUG - [OUTPUT] segment_contact: (16,)
+2026-05-08 14:21:57,036 - DEBUG - [INPUT] disk_z_tilt: ()
+2026-05-08 14:21:57,036 - DEBUG - [OUTPUT] disk_z_tilt: (1,)
+2026-05-08 14:21:57,036 - DEBUG - [INPUT] robot_direction_to_target: (2,)
+2026-05-08 14:21:57,036 - DEBUG - [OUTPUT] robot_direction_to_target: (2,)
+
+-- 5 arms, none disabled
+2026-05-08 14:27:04,472 - DEBUG - [INPUT] joint_actuator_force: (40,)
+2026-05-08 14:27:04,472 - DEBUG - [OUTPUT] joint_actuator_force: (40,)
+2026-05-08 14:27:04,472 - DEBUG - [INPUT] joint_position: (40,)
+2026-05-08 14:27:04,472 - DEBUG - [OUTPUT] joint_position: (40,)
+2026-05-08 14:27:04,473 - DEBUG - [INPUT] joint_velocity: (40,)
+2026-05-08 14:27:04,473 - DEBUG - [OUTPUT] joint_velocity: (40,)
+2026-05-08 14:27:04,473 - DEBUG - [INPUT] segment_contact: (20,)
+2026-05-08 14:27:04,473 - DEBUG - [OUTPUT] segment_contact: (20,)
+2026-05-08 14:27:04,473 - DEBUG - [INPUT] disk_z_tilt: ()
+2026-05-08 14:27:04,473 - DEBUG - [OUTPUT] disk_z_tilt: (1,)
+2026-05-08 14:27:04,474 - DEBUG - [INPUT] robot_direction_to_target: (2,)
+2026-05-08 14:27:04,474 - DEBUG - [OUTPUT] robot_direction_to_target: (2,)

vids/evaluation_metadata.yaml

@@ -0,0 +1,9 @@
+timestamp: '2026-05-07T20:43:29.612567'
+morphology_override: null
+seed: 1
+max_steps: 10000
+result:
+  return: 0.40030128204863225
+  length: 6705
+  reached_target: true
+  final_xy_dist: 0.19969871795136748