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

test(obs processor): simple tests for the output shapes of obs_processor

Browse source
This commit is contained in:
Robin Meersman 2026-05-10 12:32:57 +02:00
parent a0225eb48c
commit c597e272c7
1 changed files with 26 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

32
tests/test_obs_processor.py
View file

@ -55,11 +55,11 @@ def batch_obs(obs: dict):
return jax.tree_util.tree_map(lambda x: x[None, :], obs) return jax.tree_util.tree_map(lambda x: x[None, :], obs)
def make_processor(morph_mode: MorphMode, segments_per_arm: list[int]): def make_processor(morph_mode: MorphMode, needed_copies: int, segments_per_arm: list[int]):
return create_obs_processor( return create_obs_processor(
bounds_dict=obs_bounds, bounds_dict=obs_bounds,
num_arms=NUM_ARMS, num_arms=NUM_ARMS,
needed_copies=NUM_ARMS, needed_copies=needed_copies,
morph_mode=morph_mode, morph_mode=morph_mode,
segments_per_arm=segments_per_arm, segments_per_arm=segments_per_arm,
agent_indices=AGENT_INDICES, agent_indices=AGENT_INDICES,
@ -67,7 +67,7 @@ def make_processor(morph_mode: MorphMode, segments_per_arm: list[int]):
def test_centralized_no_damage(): def test_centralized_no_damage():
proc = make_processor(MorphMode.CENTRALIZED, SEGS_HEALTHY) proc = make_processor(MorphMode.CENTRALIZED, 1, SEGS_HEALTHY)
obs = make_obs(SEGS_HEALTHY) obs = make_obs(SEGS_HEALTHY)
obs = batch_obs(obs) obs = batch_obs(obs)
global_state = proc(obs) global_state = proc(obs)
@ -75,14 +75,34 @@ def test_centralized_no_damage():
# shape test # shape test
assert global_state.shape == (1, 1, FEAT_PER_AGENT * NUM_ARMS) assert global_state.shape == (1, 1, FEAT_PER_AGENT * NUM_ARMS)
# TODO: more?
def test_centralized_damaged_1_arm(): def test_centralized_damaged_1_arm():
pass proc = make_processor(MorphMode.CENTRALIZED, 1, SEGS_DAMAGED)
obs = make_obs(SEGS_DAMAGED)
obs = batch_obs(obs)
global_state = proc(obs)
# shape test
assert global_state.shape == (1, 1, FEAT_PER_AGENT * NUM_ARMS)
def test_decentralized_fully_connected_no_damage(): def test_decentralized_fully_connected_no_damage():
pass proc = make_processor(MorphMode.FULLY_CONNECTED, NUM_ARMS, SEGS_HEALTHY)
obs = make_obs(SEGS_HEALTHY)
obs = batch_obs(obs)
global_state = proc(obs)
# shape test
assert global_state.shape == (1, NUM_ARMS, FEAT_PER_AGENT)
def test_decentralized_fully_connected_damaged_1_arm(): def test_decentralized_fully_connected_damaged_1_arm():
pass proc = make_processor(MorphMode.FULLY_CONNECTED, NUM_ARMS, SEGS_DAMAGED)
obs = make_obs(SEGS_DAMAGED)
obs = batch_obs(obs)
global_state = proc(obs)
# shape test
assert global_state.shape == (1, NUM_ARMS, FEAT_PER_AGENT)