Files
Gymnasium/tests/wrappers/test_time_aware_observation.py
Ariel Kwiatkowski 947b857bd4 Test refactoring (#2427)
* Move tests to root with automatic PyCharm import refactoring. This will likely fail some tests

* Changed entry point for a registration test env.

* Move a stray lunar_lander test to tests/envs/...

* black

* Change the version from which importlib_metadata is replaced with importlib.metadata. Also requiring installing importlib_metadata for python 3.8 now.

???????????

* Undo last commit
2021-09-28 19:53:30 -04:00

34 lines
1.0 KiB
Python

import pytest
import gym
from gym.wrappers import TimeAwareObservation
@pytest.mark.parametrize("env_id", ["CartPole-v1", "Pendulum-v1"])
def test_time_aware_observation(env_id):
env = gym.make(env_id)
wrapped_env = TimeAwareObservation(env)
assert wrapped_env.observation_space.shape[0] == env.observation_space.shape[0] + 1
obs = env.reset()
wrapped_obs = wrapped_env.reset()
assert wrapped_env.t == 0.0
assert wrapped_obs[-1] == 0.0
assert wrapped_obs.shape[0] == obs.shape[0] + 1
wrapped_obs, _, _, _ = wrapped_env.step(env.action_space.sample())
assert wrapped_env.t == 1.0
assert wrapped_obs[-1] == 1.0
assert wrapped_obs.shape[0] == obs.shape[0] + 1
wrapped_obs, _, _, _ = wrapped_env.step(env.action_space.sample())
assert wrapped_env.t == 2.0
assert wrapped_obs[-1] == 2.0
assert wrapped_obs.shape[0] == obs.shape[0] + 1
wrapped_obs = wrapped_env.reset()
assert wrapped_env.t == 0.0
assert wrapped_obs[-1] == 0.0
assert wrapped_obs.shape[0] == obs.shape[0] + 1