[rllib] Improve test_single_agent_env_runner to prevent flaky tests (#58397)

## Description
In improving the `SingleEnvRunner.make_env`, I found that some of the
tests could be flaky.
This PR improves the testing, in particular, to `sample` to ensure that
the tests don't fail occasionally and the documentation to reflect this.

The primary flaky problem I found is that `sample(num_timesteps=X)` will
not always return a total of `X` timesteps, rather at least X timesteps
up to the number of environments more.
I'm updated the documentation to clarify this for users. 

In addition, I've added tests for when neither the number of timesteps
or episodes are given and for the `force_reset` argument

---------

Signed-off-by: Mark Towers <[email protected]>
Co-authored-by: Mark Towers <[email protected]>
Co-authored-by: simonsays1980 <[email protected]>

Author: 3 people

rllib/env/single_agent_env_runner.py

@@ -155,8 +155,15 @@ def sample(
     ) -> List[SingleAgentEpisode]:
         """Runs and returns a sample (n timesteps or m episodes) on the env(s).
 
+        If neither `num_timesteps` nor `num_episodes` are provided and the config
+        `batch_mode` is "truncate_episodes" then
+        `config.get_rollout_fragment_length(self.worker_index) * self.num_envs`
+        timesteps will be sampled.
+
         Args:
             num_timesteps: The number of timesteps to sample during this call.
+                The episodes returned will contain the total timesteps greater than or
+                equal to num_timesteps and less than num_timesteps + num_envs_per_env_runner.
                 Note that only one of `num_timesteps` or `num_episodes` may be provided.
             num_episodes: The number of episodes to sample during this call.
                 Note that only one of `num_timesteps` or `num_episodes` may be provided.
@@ -169,7 +176,7 @@ def sample(
             random_actions: If True, actions will be sampled randomly (from the action
                 space of the environment). If False (default), actions or action
                 distribution parameters are computed by the RLModule.
-            force_reset: Whether to force-reset all (vector) environments before
+            force_reset: Whether to force-reset all vectorized environments before
                 sampling. Useful if you would like to collect a clean slate of new
                 episodes via this call. Note that when sampling n episodes
                 (`num_episodes != None`), this is fixed to True.
@@ -203,6 +210,7 @@ def sample(
             # desired timesteps/episodes to sample and exploration behavior.
             if explore is None:
                 explore = self.config.explore
+
             if (
                 num_timesteps is None
                 and num_episodes is None
@@ -215,6 +223,7 @@ def sample(
 
             # Sample n timesteps.
             if num_timesteps is not None:
+                assert num_timesteps >= 0
                 samples = self._sample(
                     num_timesteps=num_timesteps,
                     explore=explore,
@@ -223,6 +232,7 @@ def sample(
                 )
             # Sample m episodes.
             elif num_episodes is not None:
+                assert num_episodes >= 0
                 samples = self._sample(
                     num_episodes=num_episodes,
                     explore=explore,

rllib/env/tests/test_single_agent_env_runner.py

@@ -1,17 +1,16 @@
 import unittest
-from functools import partial
 from unittest.mock import patch
 
 import gymnasium as gym
+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.env.utils import _gym_env_creator
 from ray.rllib.examples.envs.classes.simple_corridor import SimpleCorridor
-from ray.rllib.utils.test_utils import check
+from ray.tune.registry import ENV_CREATOR, _global_registry
 
 
 class TestSingleAgentEnvRunner(unittest.TestCase):
@@ -21,22 +20,29 @@ def setUpClass(cls) -> None:
 
         tune.register_env(
             "tune-registered",
-            lambda cfg: SimpleCorridor({"corridor_length": 10}),
+            lambda cfg: SimpleCorridor({"corridor_length": 10} | cfg),
         )
 
         gym.register(
             "TestEnv-v0",
-            partial(
-                _gym_env_creator,
-                env_context={"corridor_length": 10},
-                env_descriptor=SimpleCorridor,
-            ),
+            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")
+        gym.registry.pop("TestEnv-v0")
+        gym.registry.pop("TestEnv-v1")
+
     def test_distributed_env_runner(self):
         """Tests, whether SingleAgentEnvRunner can be distributed."""
 
@@ -68,7 +74,8 @@ def test_distributed_env_runner(self):
                 results = ray.get(results)
                 # Loop over individual EnvRunner Actor's results and inspect each.
                 for episodes in results:
-                    # Assert length of all fragments is  `rollout_fragment_length`.
+                    # 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),
                         [
@@ -79,13 +86,19 @@ def test_distributed_env_runner(self):
                         ],
                     )
 
-    def test_sample(self):
+    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=2,
-                rollout_fragment_length=64,
+                num_envs_per_env_runner=num_envs_per_env_runner,
+                rollout_fragment_length=rollout_fragment_length,
             )
         )
         env_runner = SingleAgentEnvRunner(config=config)
@@ -97,32 +110,86 @@ def test_sample(self):
                 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 (5 per env, because num_envs_per_env_runner=2)
+        # 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=10, random_actions=True)
-            check(len(episodes), 10)
+            episodes = env_runner.sample(
+                num_episodes=expected_episodes, random_actions=True
+            )
+            self.assertTrue(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 10 timesteps (5 per env)
+        # 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=10, random_actions=True)
+            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.
-            sum_ = sum(map(len, episodes))
-            self.assertTrue(sum_ in [10, 11])
+            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 128
-            # 2 (num_env_per_worker) * 64 (rollout_fragment_length).
-            sum_ = sum(map(len, episodes))
-            self.assertTrue(sum_ in [128, 129])
+            # 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))
 
     @patch(target="ray.rllib.env.env_runner.logger")
     def test_step_failed_reset_required(self, mock_logger):
@@ -172,29 +239,87 @@ def step(self, action):
 
         assert mock_logger.exception.call_count == 1
 
-    def test_vector_env(self):
+    def test_vector_env(self, num_envs_per_env_runner=5, rollout_fragment_length=10):
         """Tests, whether SingleAgentEnvRunner can run various vectorized envs."""
 
         for env in ["CartPole-v1", SimpleCorridor, "tune-registered"]:
             config = (
                 AlgorithmConfig()
                 .environment(env)
                 .env_runners(
-                    num_envs_per_env_runner=5,
-                    rollout_fragment_length=10,
+                    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.
-            episodes = env_runner.sample(random_actions=True)
-            self.assertEqual(
-                sum(len(e) for e in episodes),
-                config.num_envs_per_env_runner * config.rollout_fragment_length,
-            )
+            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,)
+
 
 if __name__ == "__main__":
     import sys