mirror of
https://github.com/Farama-Foundation/Gymnasium.git
synced 2025-08-26 00:07:41 +00:00
104 lines
3.2 KiB
Python
104 lines
3.2 KiB
Python
"""Test environment determinism by performing a rollout."""
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from tests.envs.spec_list import spec_list
|
|
|
|
|
|
@pytest.mark.parametrize("spec", spec_list, ids=[spec.id for spec in spec_list])
|
|
def test_env(spec):
|
|
"""Run a rollout with two environments and assert equality.
|
|
|
|
This test run a rollout of NUM_STEPS steps with two environments
|
|
initialized with the same seed and assert that:
|
|
|
|
- observation after first reset are the same
|
|
- same actions are sampled by the two envs
|
|
- observations are contained in the observation space
|
|
- obs, rew, done and info are equals between the two envs
|
|
|
|
Args:
|
|
spec (EnvSpec): Environment specification
|
|
|
|
"""
|
|
# Note that this precludes running this test in multiple
|
|
# threads. However, we probably already can't do multithreading
|
|
# due to some environments.
|
|
SEED = 0
|
|
NUM_STEPS = 50
|
|
|
|
env1, env2 = spec.make(), spec.make()
|
|
|
|
initial_observation1 = env1.reset(seed=SEED)
|
|
initial_observation2 = env2.reset(seed=SEED)
|
|
|
|
env1.action_space.seed(SEED)
|
|
env2.action_space.seed(SEED)
|
|
|
|
assert_equals(initial_observation1, initial_observation2)
|
|
|
|
for i in range(NUM_STEPS):
|
|
action1 = env1.action_space.sample()
|
|
action2 = env2.action_space.sample()
|
|
|
|
try:
|
|
assert_equals(action1, action2)
|
|
except AssertionError:
|
|
print("env1.action_space=", env1.action_space)
|
|
print("env2.action_space=", env2.action_space)
|
|
print("action_samples1=", action1)
|
|
print("action_samples2=", action2)
|
|
print(f"[{i}] action_sample1: {action1}, action_sample2: {action2}")
|
|
raise
|
|
|
|
# Don't check rollout equality if it's a a nondeterministic
|
|
# environment.
|
|
if spec.nondeterministic:
|
|
return
|
|
|
|
obs1, rew1, done1, info1 = env1.step(action1)
|
|
obs2, rew2, done2, info2 = env2.step(action2)
|
|
|
|
assert_equals(obs1, obs2, f"[{i}] ")
|
|
|
|
assert env1.observation_space.contains(obs1)
|
|
assert env2.observation_space.contains(obs2)
|
|
|
|
assert rew1 == rew2, f"[{i}] r1: {rew1}, r2: {rew2}"
|
|
assert done1 == done2, f"[{i}] d1: {done1}, d2: {done2}"
|
|
assert_equals(info1, info2, f"[{i}] ")
|
|
|
|
if done1: # done2 verified in previous assertion
|
|
env1.reset(seed=SEED)
|
|
env2.reset(seed=SEED)
|
|
|
|
env1.close()
|
|
env2.close()
|
|
|
|
|
|
def assert_equals(a, b, prefix=None):
|
|
"""Assert equality of data structures `a` and `b`.
|
|
|
|
Args:
|
|
a: first data structure
|
|
b: second data structure
|
|
prefix: prefix for failed assertion message for types and dicts
|
|
|
|
"""
|
|
assert type(a) == type(b), f"{prefix}Differing types: {a} and {b}"
|
|
if isinstance(a, dict):
|
|
assert list(a.keys()) == list(b.keys()), f"{prefix}Key sets differ: {a} and {b}"
|
|
|
|
for k in a.keys():
|
|
v_a = a[k]
|
|
v_b = b[k]
|
|
assert_equals(v_a, v_b)
|
|
elif isinstance(a, np.ndarray):
|
|
np.testing.assert_array_equal(a, b)
|
|
elif isinstance(a, tuple):
|
|
for elem_from_a, elem_from_b in zip(a, b):
|
|
assert_equals(elem_from_a, elem_from_b)
|
|
else:
|
|
assert a == b
|