1
Fork 0

fix: tests still assumed badly broadcasted segment_contacts

This commit is contained in:
Tibo De Peuter 2026-05-12 21:30:45 +02:00
parent da06fa2f2e
commit e21e5b1c3a
Signed by: tdpeuter
SSH key fingerprint: SHA256:u/h/LVoqKF1Iz02uOyxe6hcjmoZASCGV2HM0TG9ZMoU
2 changed files with 29 additions and 22 deletions

View file

@ -36,12 +36,14 @@ def test_centralized_forward_pass_with_padding():
)
global_state = obs_processor(amputated_obs)
# 40 + 40 + 20 + padding = 145 dimensions
assert global_state.shape == (batch_size, 1, 145), (
f"Expected global state shape (2, 1, 145), got {global_state.shape}"
# joint_position: 5 arms × 8 joints (padded) = 40
# joint_velocity: 5 arms × 8 joints (padded) = 40
# segment_contact: 5 arms × 4 segs (padded) = 20
# Total = 100 (no disk or direction keys supplied)
assert global_state.shape == (batch_size, 1, 100), (
f"Expected global state shape (2, 1, 100), got {global_state.shape}"
)
# 4. Initialize dummy networks (40 actuators for the max morphology output)
actor = Actor(action_dim=40)
critic = OneDenseLayerMLP() # Acts as the centralized critic

View file

@ -7,17 +7,14 @@ from brittle_star_project.environment.env_config import MorphMode, ObservationBo
obs_bounds = ObservationBoundsConfig().to_bounds_dict()
"""
Test for obs_processor.
Centralized: 40 features per agent:
disk_z_tilt scalar reshaped to (1,) 1 feat
joint_actuator_force 8 joints padded to 8 8 feat
joint_position 8 joints padded to 8 8 feat
joint_velocity 8 joints padded to 8 8 feat
robot_direction_to_target (x, y) 2 feat
segment_contact 4 segs, pre-padded by 9 13 feat (9 leading + 4)
"""
# Features per decentralized agent (one arm's data):
# disk_z_tilt → scalar → 1 feat
# joint_actuator_force → 4 segs × 2 joints → 8 feat
# joint_position → 4 segs × 2 joints → 8 feat
# joint_velocity → 4 segs × 2 joints → 8 feat
# robot_direction_to_target→ (x, y) → 2 feat
# segment_contact → 4 segs → 4 feat
# Total per agent: 1+8+8+8+2+4 = 31
NUM_ARMS = 5
SEGS_PER_ARM = 4 # healthy segments per arm
@ -28,7 +25,17 @@ SEGS_DAMAGED = [4, 4, 4, 4, 0] # arm 4 fully disabled
SEGS_DAMAGED_2 = [4, 0, 4, 2, 4] # arm 3 fully disabled
AGENT_INDICES = [0, 1, 2, 3, 4]
FEAT_PER_AGENT = 1 + 8 + 8 + 8 + 2 + 13 # = 40
FEAT_PER_AGENT = 1 + 8 + 8 + 8 + 2 + 4 # = 31
# Centralized flattening (needed_copies=1, one copy of global features):
# disk_z_tilt → repeated once → 1 feat
# joint_actuator_force → 5 arms × 8 joints → 40 feat
# joint_position → 5 arms × 8 joints → 40 feat
# joint_velocity → 5 arms × 8 joints → 40 feat
# robot_direction_to_target→ repeated once → 2 feat
# segment_contact → 5 arms × 4 segs → 20 feat
# Total: 1+40+40+40+2+20 = 143
FEAT_CENTRALIZED = 1 + 40 + 40 + 40 + 2 + 20 # = 143
def make_obs(segs_per_arm: list[int]) -> dict:
@ -73,10 +80,8 @@ def test_centralized_no_damage():
obs = batch_obs(obs)
global_state = proc(obs)
# shape test
assert global_state.shape == (1, 1, 188)
# TODO: more?
# Centralized: 5 agents flattened into 1 → shape (1, 1, 155)
assert global_state.shape == (1, 1, FEAT_CENTRALIZED)
def test_centralized_damaged_1_arm():
@ -86,7 +91,7 @@ def test_centralized_damaged_1_arm():
global_state = proc(obs)
# shape test
assert global_state.shape == (1, 1, 188)
assert global_state.shape == (1, 1, FEAT_CENTRALIZED)
def test_centralized_damaged_2_arms():
@ -96,7 +101,7 @@ def test_centralized_damaged_2_arms():
global_state = proc(obs)
# shape test
assert global_state.shape == (1, 1, 188)
assert global_state.shape == (1, 1, FEAT_CENTRALIZED)
def test_decentralized_fully_connected_no_damage():