From 079480324d1c9c4cfccabea560b30d81a8c90e26 Mon Sep 17 00:00:00 2001 From: Robin Meersman Date: Tue, 5 May 2026 12:11:37 +0200 Subject: [PATCH] fix(testing): forward pass test updated to new abstraction --- src/brittle_star_project/environment/obs_processing.py | 10 +++++++--- .../environment/padded_obs_wrapper.py | 6 ++++++ tests/test_network_shapes.py | 10 +++++++--- tests/test_target_direction.py | 3 --- 4 files changed, 20 insertions(+), 9 deletions(-) diff --git a/src/brittle_star_project/environment/obs_processing.py b/src/brittle_star_project/environment/obs_processing.py index 73c2e94..87ebe10 100644 --- a/src/brittle_star_project/environment/obs_processing.py +++ b/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 diff --git a/src/brittle_star_project/environment/padded_obs_wrapper.py b/src/brittle_star_project/environment/padded_obs_wrapper.py index ea00713..3216dfe 100644 --- a/src/brittle_star_project/environment/padded_obs_wrapper.py +++ b/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}: " diff --git a/tests/test_network_shapes.py b/tests/test_network_shapes.py index aad79fc..1183377 100644 --- a/tests/test_network_shapes.py +++ b/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}" ) diff --git a/tests/test_target_direction.py b/tests/test_target_direction.py index 825c50c..64a6361 100644 --- a/tests/test_target_direction.py +++ b/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()