Files
Gymnasium/tests/wrappers/test_flatten_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

23 lines
573 B
Python

import pytest
import numpy as np
import gym
from gym.wrappers import FlattenObservation
from gym import spaces
@pytest.mark.parametrize("env_id", ["Blackjack-v1"])
def test_flatten_observation(env_id):
env = gym.make(env_id)
wrapped_env = FlattenObservation(env)
obs = env.reset()
wrapped_obs = wrapped_env.reset()
space = spaces.Tuple((spaces.Discrete(32), spaces.Discrete(11), spaces.Discrete(2)))
wrapped_space = spaces.Box(0, 1, [32 + 11 + 2], dtype=np.int64)
assert space.contains(obs)
assert wrapped_space.contains(wrapped_obs)