1
Fork 0
This commit is contained in:
Robin Meersman 2026-05-06 19:26:32 +02:00
parent af68a0064c
commit 5e7fd27b61
4 changed files with 63 additions and 34 deletions

View file

@ -22,14 +22,13 @@ def test_centralized_forward_pass_with_padding():
}
segments_per_arm = jnp.array((4, 0, 4, 2, 4))
num_segments = segments_per_arm.sum().item()
num_arms = jnp.where(segments_per_arm > 0, 1, 0).sum().item()
# 2. Process and Pad Observation
masks = compute_padding_masks(segments_per_arm=list(segments_per_arm))
obs_processor = create_obs_processor(
bounds_dict={},
num_segments=num_segments,
needed_copies=1,
num_arms=num_arms,
padding_masks=masks,
morph_mode=MorphMode.CENTRALIZED,

View file

@ -45,12 +45,11 @@ def test_processor_converts_to_egocentric_direction():
env = BrittleStarJaxEnvWrapper.default(num_envs=1, backend=Backend.MJX)
segments_per_arm = jnp.array((4, 4, 4, 4, 4))
num_segments = segments_per_arm.sum().item()
num_arms = jnp.where(segments_per_arm > 0, 1, 0).sum().item()
obs_processor = create_obs_processor(
bounds_dict=cfg.obs_bounds.to_bounds_dict(),
num_segments=num_segments,
needed_copies=1,
num_arms=num_arms,
padding_masks=env.padding_masks,
morph_mode=MorphMode.CENTRALIZED,