Fix: Allow observation space dict in test_replanning

This commit is contained in:
Dominik Moritz Roth 2023-06-18 11:52:35 +02:00
parent fbba129034
commit f44f01b478

View File

@ -80,6 +80,9 @@ def test_learn_sub_trajectories(mp_type: str, env_wrap: Tuple[str, Type[RawInter
assert env.learn_sub_trajectories assert env.learn_sub_trajectories
assert env.traj_gen.learn_tau assert env.traj_gen.learn_tau
# This also verifies we are not adding the TimeAwareObservationWrapper twice # This also verifies we are not adding the TimeAwareObservationWrapper twice
if env.observation_space.__class__ in [spaces.Dict]:
assert spaces.flatten_space(env.observation_space) == env_step.observation_space
else:
assert env.observation_space == env_step.observation_space assert env.observation_space == env_step.observation_space
done = True done = True
@ -130,6 +133,9 @@ def test_replanning_time(mp_type: str, env_wrap: Tuple[str, Type[RawInterfaceWra
assert env.do_replanning assert env.do_replanning
assert callable(env.replanning_schedule) assert callable(env.replanning_schedule)
# This also verifies we are not adding the TimeAwareObservationWrapper twice # This also verifies we are not adding the TimeAwareObservationWrapper twice
if env.observation_space.__class__ in [spaces.Dict]:
assert spaces.flatten_space(env.observation_space) == env_step.observation_space
else:
assert env.observation_space == env_step.observation_space assert env.observation_space == env_step.observation_space
env.reset(seed=SEED) env.reset(seed=SEED)