fix(testing): forward pass test updated to new abstraction

src/brittle_star_project/environment/obs_processing.py

@@ -105,9 +105,13 @@ def create_obs_processor(
             # -------- SPLIT TO SEGMENTS --------
             if key in _JOINT_SCALED_KEYS:
                 if morph_mode == MorphMode.SEGMENT:
-                    # TODO: remove magic constants
-                    center_size = num_arms * 3 * 2
-                    v_center = v[:center_size].reshape(num_arms, 3 * 2)  # (arms, 6)
+                    joint_count = 3
+                    axis_per_joint = 2
+
+                    center_size = num_arms * joint_count * axis_per_joint
+                    v_center = v[:center_size].reshape(
+                        num_arms, joint_count * axis_per_joint
+                    )  # (arms, 6)
                     v_segs = v[center_size:].reshape(-1, 2)  # (segs, 2)
                     values.append(jnp.concatenate([v_center, v_segs], axis=0))  # (arms+segs, ?)
                     continue

src/brittle_star_project/environment/padded_obs_wrapper.py

@@ -30,6 +30,12 @@ def compute_padding_masks(
     mask_2x = []
 
     for arm_idx, (actual, ref) in enumerate(zip(segments_per_arm, reference_segments_per_arm)):
+        if not isinstance(actual, int):
+            actual = actual.item()
+
+        if not isinstance(ref, int):
+            ref = ref.item()
+
         if not (0 <= actual <= ref):
             raise ValueError(
                 f"Invalid amputation at arm {arm_idx}: "

tests/test_network_shapes.py

@@ -1,5 +1,6 @@
 import jax
 import jax.numpy as jnp
+from brittle_star_project.environment.env_config import MorphMode
 from brittle_star_project.environment.padded_obs_wrapper import (
     compute_padding_masks,
 )
@@ -31,11 +32,12 @@ def test_centralized_forward_pass_with_padding():
         num_segments=num_segments,
         num_arms=num_arms,
         padding_masks=masks,
+        morph_mode=MorphMode.CENTRALIZED,
     )
     global_state = obs_processor(amputated_obs)
 
     # 40 + 40 + 20 = 100 dimensions
-    assert global_state.shape == (batch_size, 100), (
+    assert global_state.shape == (batch_size, 1, 100), (
         f"Expected global state shape (2, 100), got {global_state.shape}"
     )
 
@@ -54,9 +56,11 @@ def test_centralized_forward_pass_with_padding():
     action_mean, action_log_std = actor.apply(actor_params, global_state)
     value = critic.apply(critic_params, global_state)
 
-    assert action_mean.shape == (batch_size, 40), f"Actor mean shape mismatch: {action_mean.shape}"
+    assert action_mean.shape == (batch_size, 1, 40), (
+        f"Actor mean shape mismatch: {action_mean.shape}"
+    )
     assert action_log_std.shape == (40,), f"Actor log_std shape mismatch: {action_log_std.shape}"
-    assert value.shape == (batch_size, 1) or value.shape == (batch_size,), (
+    assert value.shape == (batch_size, 1, 1) or value.shape == (batch_size,), (
         f"Critic value shape mismatch: {value.shape}"
     )
 

tests/test_target_direction.py

@@ -88,6 +88,3 @@ def test_processor_converts_to_egocentric_direction():
         f"The obs_processor did not correctly rotate the vector to egocentric. "
         f"Expected {expected_local_target}, but got {local_target}."
     )
-
-
-test_processor_converts_to_egocentric_direction()