chore: import upstream snapshot with attribution

rllib/algorithms/bc/tests/test_bc.py

@@ -0,0 +1,146 @@
+import unittest
+from pathlib import Path
+
+import ray
+from ray.rllib.algorithms.bc import BCConfig
+from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    EVALUATION_RESULTS,
+    LEARNER_RESULTS,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+
+
+class TestBC(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        ray.shutdown()
+
+    def test_bc_compilation_and_learning_from_offline_file(self):
+        # Define the data paths.
+        data_path = "offline/tests/data/cartpole/cartpole-v1_large"
+        base_path = Path(__file__).parents[3]
+        print(f"base_path={base_path}")
+        data_path = "local://" / base_path / data_path
+
+        print(f"data_path={data_path}")
+
+        # Define the BC config.
+        config = (
+            BCConfig()
+            .environment(env="CartPole-v1")
+            .learners(
+                num_learners=0,
+            )
+            .evaluation(
+                evaluation_interval=3,
+                evaluation_num_env_runners=1,
+                evaluation_duration=5,
+                evaluation_parallel_to_training=True,
+            )
+            # Note, the `input_` argument is the major argument for the
+            # new offline API.
+            .offline_data(
+                input_=[data_path.as_posix()],
+                dataset_num_iters_per_learner=1,
+            )
+            .training(
+                lr=0.0008,
+                train_batch_size_per_learner=2000,
+            )
+        )
+
+        num_iterations = 350
+        min_return_to_reach = 120.0
+
+        # TODO (simon): Add support for recurrent modules.
+        algo = config.build()
+        learnt = False
+        for i in range(num_iterations):
+            results = algo.train()
+            print(results)
+
+            eval_results = results.get(EVALUATION_RESULTS, {})
+            if eval_results:
+                episode_return_mean = eval_results[ENV_RUNNER_RESULTS][
+                    EPISODE_RETURN_MEAN
+                ]
+                print(f"iter={i}, R={episode_return_mean}")
+                if episode_return_mean > min_return_to_reach:
+                    print("BC has learnt the task!")
+                    learnt = True
+                    break
+
+        if not learnt:
+            raise ValueError(
+                f"`BC` did not reach {min_return_to_reach} reward from "
+                "expert offline data!"
+            )
+
+        algo.stop()
+
+    def test_bc_lr_schedule(self):
+        # Define the data paths.
+        data_path = "offline/tests/data/cartpole/cartpole-v1_large"
+        base_path = Path(__file__).parents[3]
+        data_path = "local://" / base_path / data_path
+
+        config = (
+            BCConfig()
+            .environment(env="CartPole-v1")
+            .learners(
+                num_learners=0,
+            )
+            .evaluation(
+                evaluation_interval=3,
+                evaluation_num_env_runners=1,
+                evaluation_duration=5,
+                evaluation_parallel_to_training=True,
+            )
+            # Note, the `input_` argument is the major argument for the
+            # new offline API.
+            .offline_data(
+                input_=[data_path.as_posix()],
+                dataset_num_iters_per_learner=1,
+            )
+            .training(
+                lr=[
+                    [0, 0.001],
+                    [3000, 0.01],
+                ],
+                train_batch_size_per_learner=2000,
+            )
+        )
+        algo = config.build()
+
+        done = False
+        while not done:
+            results = algo.train()
+            ts = results[NUM_ENV_STEPS_SAMPLED_LIFETIME]
+            assert ts > 0
+            lr = results[LEARNER_RESULTS][DEFAULT_POLICY_ID][
+                "default_optimizer_learning_rate"
+            ]
+            if ts < 3000:
+                # The learning rate should be linearly interpolated.
+                expected_lr = 0.001 + (ts / 3000) * (0.01 - 0.001)
+                self.assertAlmostEqual(lr, expected_lr, places=6)
+            else:
+                self.assertEqual(lr, 0.01)
+                done = True
+
+        algo.stop()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))