chore: import upstream snapshot with attribution

rllib/callbacks/tests/test_callbacks_old_api_stack.py

@@ -0,0 +1,216 @@
+import unittest
+from collections import Counter
+
+import ray
+from ray.rllib.algorithms.callbacks import DefaultCallbacks, make_multi_callbacks
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.examples.envs.classes.random_env import RandomEnv
+
+
+class EpisodeAndSampleCallbacks(DefaultCallbacks):
+    def __init__(self):
+        super().__init__()
+        self.counts = Counter()
+
+    def on_episode_start(self, *args, **kwargs):
+        self.counts.update({"start": 1})
+
+    def on_episode_step(self, *args, **kwargs):
+        self.counts.update({"step": 1})
+
+    def on_episode_end(self, *args, **kwargs):
+        self.counts.update({"end": 1})
+
+    def on_sample_end(self, *args, **kwargs):
+        self.counts.update({"sample": 1})
+
+
+class OnSubEnvironmentCreatedCallback(DefaultCallbacks):
+    def on_sub_environment_created(
+        self, *, worker, sub_environment, env_context, **kwargs
+    ):
+        # Create a vector-index-sum property per remote worker.
+        if not hasattr(worker, "sum_sub_env_vector_indices"):
+            worker.sum_sub_env_vector_indices = 0
+        # Add the sub-env's vector index to the counter.
+        worker.sum_sub_env_vector_indices += env_context.vector_index
+        print(
+            f"sub-env {sub_environment} created; "
+            f"worker={worker.worker_index}; "
+            f"vector-idx={env_context.vector_index}"
+        )
+
+
+class OnEpisodeCreatedCallback(DefaultCallbacks):
+    def __init__(self):
+        super().__init__()
+        self._reset_counter = 0
+
+    def on_episode_created(
+        self, *, worker, base_env, policies, env_index, episode, **kwargs
+    ):
+        print(f"Sub-env {env_index} is going to be reset.")
+        self._reset_counter += 1
+
+        # Make sure the passed in episode is really brand new.
+        assert episode.env_id == env_index
+        assert episode.length == -1
+        assert episode.worker is worker
+
+
+class TestCallbacks(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def test_episode_and_sample_callbacks(self):
+        config = (
+            PPOConfig()
+            .api_stack(
+                enable_rl_module_and_learner=False,
+                enable_env_runner_and_connector_v2=False,
+            )
+            .environment("CartPole-v1")
+            .env_runners(num_env_runners=0)
+            .callbacks(EpisodeAndSampleCallbacks)
+            .training(train_batch_size=50, minibatch_size=50, num_epochs=1)
+        )
+        algo = config.build()
+        algo.train()
+        algo.train()
+        callback_obj = algo.env_runner.callbacks
+        self.assertGreater(callback_obj.counts["sample"], 0)
+        self.assertGreater(callback_obj.counts["start"], 0)
+        self.assertGreater(callback_obj.counts["end"], 0)
+        self.assertGreater(callback_obj.counts["step"], 0)
+        algo.stop()
+
+    def test_on_sub_environment_created(self):
+
+        config = (
+            PPOConfig()
+            .api_stack(
+                enable_rl_module_and_learner=False,
+                enable_env_runner_and_connector_v2=False,
+            )
+            .environment("CartPole-v1")
+            # Create 4 sub-environments per remote worker.
+            # Create 2 remote workers.
+            .env_runners(num_envs_per_env_runner=4, num_env_runners=2)
+        )
+
+        for callbacks in (
+            OnSubEnvironmentCreatedCallback,
+            make_multi_callbacks([OnSubEnvironmentCreatedCallback]),
+        ):
+            config.callbacks(callbacks)
+
+            algo = config.build()
+            # Fake the counter on the local worker (doesn't have an env) and
+            # set it to -1 so the below `foreach_env_runner()` won't fail.
+            algo.env_runner.sum_sub_env_vector_indices = -1
+
+            # Get sub-env vector index sums from the 2 remote workers:
+            sum_sub_env_vector_indices = algo.env_runner_group.foreach_env_runner(
+                lambda w: w.sum_sub_env_vector_indices
+            )
+            # Local worker has no environments -> Expect the -1 special
+            # value returned by the above lambda.
+            self.assertTrue(sum_sub_env_vector_indices[0] == -1)
+            # Both remote workers (index 1 and 2) have a vector index counter
+            # of 6 (sum of vector indices: 0 + 1 + 2 + 3).
+            self.assertTrue(sum_sub_env_vector_indices[1] == 6)
+            self.assertTrue(sum_sub_env_vector_indices[2] == 6)
+            algo.stop()
+
+    def test_on_sub_environment_created_with_remote_envs(self):
+        config = (
+            PPOConfig()
+            .api_stack(
+                enable_rl_module_and_learner=False,
+                enable_env_runner_and_connector_v2=False,
+            )
+            .environment("CartPole-v1")
+            .env_runners(
+                # Make each sub-environment a ray actor.
+                remote_worker_envs=True,
+                # Create 2 remote workers.
+                num_env_runners=2,
+                # Create 4 sub-environments (ray remote actors) per remote
+                # worker.
+                num_envs_per_env_runner=4,
+            )
+        )
+
+        for callbacks in (
+            OnSubEnvironmentCreatedCallback,
+            make_multi_callbacks([OnSubEnvironmentCreatedCallback]),
+        ):
+            config.callbacks(callbacks)
+
+            algo = config.build()
+            # Fake the counter on the local worker (doesn't have an env) and
+            # set it to -1 so the below `foreach_env_runner()` won't fail.
+            algo.env_runner.sum_sub_env_vector_indices = -1
+
+            # Get sub-env vector index sums from the 2 remote workers:
+            sum_sub_env_vector_indices = algo.env_runner_group.foreach_env_runner(
+                lambda w: w.sum_sub_env_vector_indices
+            )
+            # Local worker has no environments -> Expect the -1 special
+            # value returned by the above lambda.
+            self.assertTrue(sum_sub_env_vector_indices[0] == -1)
+            # Both remote workers (index 1 and 2) have a vector index counter
+            # of 6 (sum of vector indices: 0 + 1 + 2 + 3).
+            self.assertTrue(sum_sub_env_vector_indices[1] == 6)
+            self.assertTrue(sum_sub_env_vector_indices[2] == 6)
+            algo.stop()
+
+    def test_on_episode_created(self):
+        # 1000 steps sampled (2.5 episodes on each sub-environment) before training
+        # starts.
+        config = (
+            PPOConfig()
+            .api_stack(
+                enable_rl_module_and_learner=False,
+                enable_env_runner_and_connector_v2=False,
+            )
+            .environment(
+                RandomEnv,
+                env_config={
+                    "max_episode_len": 200,
+                    "p_terminated": 0.0,
+                },
+            )
+            .env_runners(num_envs_per_env_runner=2, num_env_runners=1)
+            .callbacks(OnEpisodeCreatedCallback)
+        )
+
+        algo = config.build()
+        algo.train()
+        # Two sub-environments share 4000 steps in the first training iteration
+        # (train_batch_size=4000).
+        # -> 4000 / 2 [sub-envs] = 2000 [per sub-env]
+        # -> 1 episode = 200 timesteps
+        # -> 10 episodes per sub-env
+        # -> 11 episodes created [per sub-env] = 22 episodes total
+        self.assertEqual(
+            22,
+            algo.env_runner_group.foreach_env_runner(
+                lambda w: w.callbacks._reset_counter,
+                local_env_runner=False,
+            )[0],
+        )
+        algo.stop()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))