ray/rllib/env/wrappers/tests/test_recsim_wrapper.py
Balaji Veeramani 7f1bacc7dc
[CI] Format Python code with Black (#21975)
See #21316 and #21311 for the motivation behind these changes.
2022-01-29 18:41:57 -08:00

41 lines
1.5 KiB
Python

import gym
import unittest
from ray.rllib.examples.env.recsim_recommender_system_envs import (
InterestEvolutionRecSimEnv,
)
from ray.rllib.env.wrappers.recsim import MultiDiscreteToDiscreteActionWrapper
from ray.rllib.utils.error import UnsupportedSpaceException
class TestRecSimWrapper(unittest.TestCase):
def test_observation_space(self):
env = InterestEvolutionRecSimEnv()
obs = env.reset()
self.assertTrue(
env.observation_space.contains(obs),
f"{env.observation_space} doesn't contain {obs}",
)
new_obs, _, _, _ = env.step(env.action_space.sample())
self.assertTrue(env.observation_space.contains(new_obs))
def test_action_space_conversion(self):
env = InterestEvolutionRecSimEnv({"convert_to_discrete_action_space": True})
self.assertIsInstance(env.action_space, gym.spaces.Discrete)
env.reset()
action = env.action_space.sample()
self.assertTrue(env.action_space.contains(action))
new_obs, _, _, _ = env.step(action)
self.assertTrue(env.observation_space.contains(new_obs))
def test_double_action_space_conversion_raises_exception(self):
env = InterestEvolutionRecSimEnv({"convert_to_discrete_action_space": True})
with self.assertRaises(UnsupportedSpaceException):
env = MultiDiscreteToDiscreteActionWrapper(env)
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", __file__]))