mirror of
https://github.com/vale981/ray
synced 2025-03-06 18:41:40 -05:00
37 lines
1.3 KiB
Python
37 lines
1.3 KiB
Python
import gym
|
|
import unittest
|
|
|
|
from ray.rllib.env.wrappers.recsim_wrapper import (
|
|
MultiDiscreteToDiscreteActionWrapper, make_recsim_env)
|
|
from ray.rllib.utils.error import UnsupportedSpaceException
|
|
|
|
|
|
class TestRecSimWrapper(unittest.TestCase):
|
|
def test_observation_space(self):
|
|
env = make_recsim_env(config={})
|
|
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 = make_recsim_env(
|
|
config={"convert_to_discrete_action_space": True})
|
|
self.assertIsInstance(env.action_space, gym.spaces.Discrete)
|
|
env.reset()
|
|
action = env.action_space.sample()
|
|
env.step(action)
|
|
|
|
def test_double_action_space_conversion_raises_exception(self):
|
|
env = make_recsim_env(
|
|
config={"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__]))
|