import unittest from unittest.mock import patch import gymnasium as gym from gymnasium.envs.classic_control.cartpole import CartPoleVectorEnv from gymnasium.envs.mujoco.swimmer_v4 import SwimmerEnv import ray from ray import tune from ray.rllib.algorithms.algorithm_config import AlgorithmConfig from ray.rllib.env.env_runner import StepFailedRecreateEnvError from ray.rllib.env.single_agent_env_runner import SingleAgentEnvRunner from ray.rllib.examples.envs.classes.simple_corridor import SimpleCorridor from ray.rllib.examples.envs.classes.ten_step_error_env import TenStepErrorEnv from ray.tune.registry import ENV_CREATOR, _global_registry class TestSingleAgentEnvRunner(unittest.TestCase): @classmethod def setUpClass(cls) -> None: ray.init() tune.register_env( "tune-registered", lambda cfg: SimpleCorridor({"corridor_length": 10} | cfg), ) tune.register_env( "tune-registered-vector", lambda cfg: CartPoleVectorEnv(**cfg), ) gym.register( "TestEnv-v0", entry_point=SimpleCorridor, kwargs={"corridor_length": 10}, ) gym.register( "TestEnv-v1", entry_point=SwimmerEnv, kwargs={"forward_reward_weight": 2.0, "reset_noise_scale": 0.2}, ) @classmethod def tearDownClass(cls) -> None: ray.shutdown() _global_registry.unregister(ENV_CREATOR, "tune-registered") _global_registry.unregister(ENV_CREATOR, "tune-registered-vector") gym.registry.pop("TestEnv-v0") gym.registry.pop("TestEnv-v1") def test_distributed_env_runner(self): """Tests, whether SingleAgentEnvRunner can be distributed.""" remote_class = ray.remote(num_cpus=1, num_gpus=0)(SingleAgentEnvRunner) # Test with both parallelized sub-envs and w/o. async_vectorization_mode = [False, True] for async_ in async_vectorization_mode: for env_spec in ["tune-registered", "CartPole-v1", SimpleCorridor]: config = ( AlgorithmConfig().environment(env_spec) # Vectorize x5 and by default, rollout 10 timesteps per individual # env. .env_runners( num_env_runners=5, num_envs_per_env_runner=5, rollout_fragment_length=10, remote_worker_envs=async_, ) ) array = [ remote_class.remote(config=config) for _ in range(config.num_env_runners) ] # Sample in parallel. results = [a.sample.remote(random_actions=True) for a in array] results = ray.get(results) # Loop over individual EnvRunner Actor's results and inspect each. for episodes in results: # Assert length of all fragments >= `rollout_fragment_length * num_envs_per_env_runner` and # < rollout_fragment_length * (num_envs_per_env_runner + 1) self.assertIn( sum(len(e) for e in episodes), [ config.num_envs_per_env_runner * config.rollout_fragment_length + i for i in range(config.num_envs_per_env_runner) ], ) def test_sample( self, num_envs_per_env_runner=5, expected_episodes=10, expected_timesteps=20, rollout_fragment_length=64, ): config = ( AlgorithmConfig() .environment("CartPole-v1") .env_runners( num_envs_per_env_runner=num_envs_per_env_runner, rollout_fragment_length=rollout_fragment_length, ) ) env_runner = SingleAgentEnvRunner(config=config) # Expect error if both num_timesteps and num_episodes given. self.assertRaises( AssertionError, lambda: env_runner.sample( num_timesteps=10, num_episodes=10, random_actions=True ), ) # Verify that an error is raised if a negative number is used self.assertRaises( AssertionError, lambda: env_runner.sample(num_timesteps=-1, random_actions=True), ) self.assertRaises( AssertionError, lambda: env_runner.sample(num_episodes=-1, random_actions=True), ) # Sample 10 episodes (2 per env, because num_envs_per_env_runner=5) # Repeat 100 times for _ in range(100): episodes = env_runner.sample( num_episodes=expected_episodes, random_actions=True ) self.assertGreaterEqual(len(episodes), expected_episodes) # Since we sampled complete episodes, there should be no ongoing episodes # being returned. self.assertTrue(all(e.is_done for e in episodes)) self.assertTrue(all(e.t_started == 0 for e in episodes)) # Sample 20 timesteps (4 per env) # Repeat 100 times env_runner.sample(random_actions=True) # for the `e.t_started > 0` for _ in range(100): episodes = env_runner.sample( num_timesteps=expected_timesteps, random_actions=True ) # Check the sum of lengths of all episodes returned. total_timesteps = sum(len(e) for e in episodes) self.assertTrue( expected_timesteps <= total_timesteps <= expected_timesteps + num_envs_per_env_runner ) self.assertTrue(any(e.t_started > 0 for e in episodes)) # Sample a number of timesteps that's not a factor of the number of environments # Repeat 100 times expected_uneven_timesteps = expected_timesteps + num_envs_per_env_runner // 2 for _ in range(100): episodes = env_runner.sample( num_timesteps=expected_uneven_timesteps, random_actions=True ) # Check the sum of lengths of all episodes returned. total_timesteps = sum(len(e) for e in episodes) self.assertTrue( expected_uneven_timesteps <= total_timesteps <= expected_uneven_timesteps + num_envs_per_env_runner, ) self.assertTrue(any(e.t_started > 0 for e in episodes)) # Sample rollout_fragment_length=64, 100 times # Repeat 100 times for _ in range(100): episodes = env_runner.sample(random_actions=True) # Check, whether the sum of lengths of all episodes returned is 320 # 5 (num_env_per_worker) * 64 (rollout_fragment_length). total_timesteps = sum(len(e) for e in episodes) self.assertTrue( num_envs_per_env_runner * rollout_fragment_length <= total_timesteps <= ( num_envs_per_env_runner * rollout_fragment_length + num_envs_per_env_runner ) ) self.assertTrue(any(e.t_started > 0 for e in episodes)) # Test that force_reset will create episodes from scratch even with `num_timesteps` episodes = env_runner.sample( num_timesteps=expected_timesteps, random_actions=True, force_reset=True ) self.assertTrue(all(e.t_started == 0 for e in episodes)) episodes = env_runner.sample( num_timesteps=expected_timesteps, random_actions=True, force_reset=False ) self.assertTrue(any(e.t_started > 0 for e in episodes)) def test_sample_with_env_error(self): config = ( AlgorithmConfig() .environment(TenStepErrorEnv) # Vectorize x2 and by default, rollout 64 timesteps per individual env. .env_runners(num_envs_per_env_runner=2, rollout_fragment_length=64) .fault_tolerance(restart_failed_sub_environments=True) ) env_runner = SingleAgentEnvRunner(config=config) # Sample first episode. # Since both environments are reset at the same step, we should get 2 episodes. episodes = env_runner.sample(num_episodes=2, random_actions=True) self.assertEqual(len(episodes), 2) self.assertEqual(len(episodes[0]), 10) self.assertListEqual( [info["last_eps_errored"] for info in episodes[0].infos], [False] * 11 ) # Sample second episode. # This should reset the env under the hood and the sample from a new env. episodes = env_runner.sample(num_episodes=2, random_actions=True) self.assertEqual(len(episodes), 2) self.assertEqual(len(episodes[0]), 10) self.assertListEqual( [info["last_eps_errored"] for info in episodes[0].infos], [False] * 11 ) # Sample timesteps episodes = env_runner.sample(num_timesteps=10, random_actions=True) self.assertEqual(len(episodes), 2) self.assertEqual(len(episodes[0]), 5) self.assertEqual(len(episodes[1]), 5) # Because both envs have been reset, last_eps_errored should be true self.assertListEqual( [info["last_eps_errored"] for info in episodes[0].infos], [True] * 6 ) # Sample timesteps episodes = env_runner.sample(num_timesteps=10, random_actions=True) self.assertEqual(len(episodes), 2) self.assertEqual(len(episodes[0]), 5) self.assertEqual(len(episodes[1]), 5) self.assertListEqual( [info["last_eps_errored"] for info in episodes[0].infos], [False] * 6 ) @patch(target="ray.rllib.env.env_runner.logger") def test_step_failed_reset_required(self, mock_logger): """Tests, whether SingleAgentEnvRunner can handle StepFailedResetRequired.""" # Define an env that raises StepFailedResetRequired class ErrorRaisingEnv(gym.Env): def __init__(self, config=None): # As per gymnasium standard, provide observation and action spaces in your # constructor. self.observation_space = gym.spaces.Discrete(2) self.action_space = gym.spaces.Discrete(2) self.exception_type = config["exception_type"] def reset(self, *, seed=None, options=None): return self.observation_space.sample(), {} def step(self, action): raise self.exception_type() config = ( AlgorithmConfig() .environment( ErrorRaisingEnv, env_config={"exception_type": StepFailedRecreateEnvError}, ) .env_runners(num_envs_per_env_runner=1, rollout_fragment_length=10) .fault_tolerance(restart_failed_sub_environments=True) ) env_runner = SingleAgentEnvRunner(config=config) # Check that we don't log the error on the first step (because we don't raise StepFailedResetRequired) # We need two steps because the first one naturally raises ResetNeeded because we try to step before the env is reset. env_runner._try_env_reset() env_runner._try_env_step(actions=[None]) assert mock_logger.exception.call_count == 0 config.environment(ErrorRaisingEnv, env_config={"exception_type": ValueError}) env_runner = SingleAgentEnvRunner(config=config) # Check that we don't log the error on the first step (because we don't raise StepFailedResetRequired) # We need two steps because the first one naturally raises ResetNeeded because we try to step before the env is reset. env_runner._try_env_reset() env_runner._try_env_step(actions=[None]) assert mock_logger.exception.call_count == 1 def test_vector_env(self, num_envs_per_env_runner=5, rollout_fragment_length=10): """Tests, whether SingleAgentEnvRunner can run various vectorized envs.""" # "ALE/Pong-v5" works but ale-py is not installed on microcheck for env in ["CartPole-v1", SimpleCorridor, "tune-registered"]: config = ( AlgorithmConfig() .environment(env) .env_runners( num_envs_per_env_runner=num_envs_per_env_runner, rollout_fragment_length=rollout_fragment_length, ) ) env_runner = SingleAgentEnvRunner(config=config) # Sample with the async-vectorized env. for i in range(100): episodes = env_runner.sample(random_actions=True) total_timesteps = sum(len(e) for e in episodes) self.assertTrue( num_envs_per_env_runner * rollout_fragment_length <= total_timesteps <= ( num_envs_per_env_runner * rollout_fragment_length + num_envs_per_env_runner ) ) env_runner.stop() def test_env_context(self): """Tests, whether SingleAgentEnvRunner can pass kwargs to the environments correctly.""" # default without env configs config = AlgorithmConfig().environment("Swimmer-v4") env_runner = SingleAgentEnvRunner(config=config) assert env_runner.env.env.get_attr("_forward_reward_weight") == (1.0,) assert env_runner.env.env.get_attr("_reset_noise_scale") == (0.1,) # Test gym registered environment env with kwargs config = AlgorithmConfig().environment( "Swimmer-v4", env_config={"forward_reward_weight": 2.0, "reset_noise_scale": 0.2}, ) env_runner = SingleAgentEnvRunner(config=config) assert env_runner.env.env.get_attr("_forward_reward_weight") == (2.0,) assert env_runner.env.env.get_attr("_reset_noise_scale") == (0.2,) # Test gym registered environment env with pre-set kwargs config = AlgorithmConfig().environment("TestEnv-v1") env_runner = SingleAgentEnvRunner(config=config) assert env_runner.env.env.get_attr("_forward_reward_weight") == (2.0,) assert env_runner.env.env.get_attr("_reset_noise_scale") == (0.2,) # Test using a mixture of registered kwargs and env configs config = AlgorithmConfig().environment( "TestEnv-v1", env_config={"forward_reward_weight": 3.0} ) env_runner = SingleAgentEnvRunner(config=config) assert env_runner.env.env.get_attr("_forward_reward_weight") == (3.0,) assert env_runner.env.env.get_attr("_reset_noise_scale") == (0.2,) # Test env-config with Tune registered or callable # default config = AlgorithmConfig().environment("tune-registered") env_runner = SingleAgentEnvRunner(config=config) assert env_runner.env.env.get_attr("end_pos") == (10.0,) # tune-registered config = AlgorithmConfig().environment( "tune-registered", env_config={"corridor_length": 5.0} ) env_runner = SingleAgentEnvRunner(config=config) assert env_runner.env.env.get_attr("end_pos") == (5.0,) # callable config = AlgorithmConfig().environment( SimpleCorridor, env_config={"corridor_length": 5.0} ) env_runner = SingleAgentEnvRunner(config=config) assert env_runner.env.env.get_attr("end_pos") == (5.0,) def test_vectorize_mode(self): """Test different vectorize mode for creating the environment.""" # default config = ( AlgorithmConfig() .environment("CartPole-v1") .env_runners(num_envs_per_env_runner=3) ) env_runner = SingleAgentEnvRunner(config=config) assert isinstance(env_runner.env.env, gym.vector.SyncVectorEnv) # different vectorize mode options contained in gymnasium registry for env_name, mode, expected_env_type in [ ("CartPole-v1", "sync", gym.vector.SyncVectorEnv), ("CartPole-v1", gym.VectorizeMode.SYNC, gym.vector.SyncVectorEnv), ("CartPole-v1", "async", gym.vector.AsyncVectorEnv), ("CartPole-v1", gym.VectorizeMode.ASYNC, gym.vector.AsyncVectorEnv), ("CartPole-v1", "vector_entry_point", CartPoleVectorEnv), ("CartPole-v1", gym.VectorizeMode.VECTOR_ENTRY_POINT, CartPoleVectorEnv), # TODO (mark) re-add with ale-py 0.11 support # ("ALE/Pong-v5", "vector_entry_point", AtariVectorEnv), # ("ALE/Pong-v5", gym.VectorizeMode.VECTOR_ENTRY_POINT, AtariVectorEnv), ]: config = ( AlgorithmConfig() .environment(env_name) .env_runners(gym_env_vectorize_mode=mode, num_envs_per_env_runner=3) ) env_runner = SingleAgentEnvRunner(config=config) assert isinstance(env_runner.env.env, expected_env_type) # test with tune registered vector environment config = ( AlgorithmConfig() .environment( "tune-registered-vector", env_config={"sutton_barto_reward": True} ) .env_runners( gym_env_vectorize_mode="vector_entry_point", num_envs_per_env_runner=3 ) ) env_runner = SingleAgentEnvRunner(config=config) assert isinstance(env_runner.env.env, CartPoleVectorEnv) assert env_runner.env.env._sutton_barto_reward is True # test with callable vector environment config = ( AlgorithmConfig() .environment( lambda cfg: CartPoleVectorEnv(**cfg), env_config={"sutton_barto_reward": True}, ) .env_runners( gym_env_vectorize_mode="vector_entry_point", num_envs_per_env_runner=3 ) ) env_runner = SingleAgentEnvRunner(config=config) assert isinstance(env_runner.env.env, CartPoleVectorEnv) assert env_runner.env.env._sutton_barto_reward is True # check passing the env config with a gym_env_vectorize_mode config = ( AlgorithmConfig() .environment("CartPole-v1", env_config={"sutton_barto_reward": True}) .env_runners(gym_env_vectorize_mode="sync", num_envs_per_env_runner=3) ) env_runner = SingleAgentEnvRunner(config=config) assert env_runner.env.env.get_attr("_sutton_barto_reward") == (True, True, True) config = ( AlgorithmConfig() .environment("CartPole-v1", env_config={"sutton_barto_reward": True}) .env_runners( gym_env_vectorize_mode="vector_entry_point", num_envs_per_env_runner=3 ) ) env_runner = SingleAgentEnvRunner(config=config) assert env_runner.env.env._sutton_barto_reward is True if __name__ == "__main__": import sys import pytest sys.exit(pytest.main(["-v", __file__]))