fix: tests still assumed badly broadcasted segment_contacts
This commit is contained in:
parent
da06fa2f2e
commit
e21e5b1c3a
2 changed files with 29 additions and 22 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
Reference in a new issue