Files
Gymnasium/tests/vector/test_vector_env_wrapper.py
pseudo-rnd-thoughts 640c509109 Rename to gymnasium
2022-09-08 10:10:07 +01:00

32 lines
978 B
Python

import numpy as np
from gymnasium.vector import VectorEnvWrapper, make
class DummyWrapper(VectorEnvWrapper):
def __init__(self, env):
self.env = env
self.counter = 0
def reset_async(self, **kwargs):
super().reset_async()
self.counter += 1
def test_vector_env_wrapper_inheritance():
env = make("FrozenLake-v1", asynchronous=False)
wrapped = DummyWrapper(env)
wrapped.reset()
assert wrapped.counter == 1
def test_vector_env_wrapper_attributes():
"""Test if `set_attr`, `call` methods for VecEnvWrapper get correctly forwarded to the vector env it is wrapping."""
env = make("CartPole-v1", num_envs=3)
wrapped = DummyWrapper(make("CartPole-v1", num_envs=3))
assert np.allclose(wrapped.call("gravity"), env.call("gravity"))
env.set_attr("gravity", [20.0, 20.0, 20.0])
wrapped.set_attr("gravity", [20.0, 20.0, 20.0])
assert np.allclose(wrapped.get_attr("gravity"), env.get_attr("gravity"))