12345678910111213141516171819202122232425262728293031323334353637 |
- 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__]))
|