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}: "