tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
Archived
1
Fork 0
Code Issues Releases Activity

test: fix outdated imports etc.

Browse source
This commit is contained in:
Tibo De Peuter 2026-04-28 18:03:53 +02:00
parent 2a34b84f83
commit ba8184b034
Signed by: tdpeuter
SSH key fingerprint: SHA256:u/h/LVoqKF1Iz02uOyxe6hcjmoZASCGV2HM0TG9ZMoU
1 changed files with 4 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

19
tests/test_network_shapes.py
View file

@ -2,8 +2,8 @@ import jax
import jax.numpy as jnp import jax.numpy as jnp
from brittle_star_project.environment.padded_obs_wrapper import ( from brittle_star_project.environment.padded_obs_wrapper import (
compute_padding_masks, compute_padding_masks,
pad_observations_batched,
) )
from brittle_star_project.environment.obs_processing import create_obs_processor
# We use Actor and OneDenseLayerMLP (as the critic) based on your mlps.py # We use Actor and OneDenseLayerMLP (as the critic) based on your mlps.py
from brittle_star_project.MLPs.mlps import Actor, OneDenseLayerMLP from brittle_star_project.MLPs.mlps import Actor, OneDenseLayerMLP
@ -20,21 +20,10 @@ def test_centralized_forward_pass_with_padding():
"segment_contact": jnp.zeros((batch_size, 14)), "segment_contact": jnp.zeros((batch_size, 14)),
} }
# 2. Pad Observation using the boolean scattering wrapper # 2. Process and Pad Observation
masks = compute_padding_masks(segments_per_arm=(4, 0, 4, 2, 4)) masks = compute_padding_masks(segments_per_arm=(4, 0, 4, 2, 4))
padded_obs = pad_observations_batched(amputated_obs, masks) obs_processor = create_obs_processor(bounds_dict={}, padding_masks=masks)
global_state = obs_processor(amputated_obs)
# Assertions to ensure padding sizes are correct (40 joints, 20 segments)
assert padded_obs["joint_position"].shape == (batch_size, 40), "Padding failed for joint keys"
assert padded_obs["segment_contact"].shape == (batch_size, 20), (
"Padding failed for segment keys"
)
# 3. Concatenate for Centralized MLP (simulating the global state vector)
global_state = jnp.concatenate(
[padded_obs["joint_position"], padded_obs["joint_velocity"], padded_obs["segment_contact"]],
axis=-1,
)
# 40 + 40 + 20 = 100 dimensions # 40 + 40 + 20 = 100 dimensions
assert global_state.shape == (batch_size, 100), ( assert global_state.shape == (batch_size, 100), (