chore: import upstream snapshot with attribution

rllib/utils/tests/test_checkpointable.py

@@ -0,0 +1,117 @@
+import random
+import shutil
+import unittest
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import ray
+from ray.rllib.algorithms.ppo import PPO, PPOConfig
+from ray.rllib.utils.metrics import LEARNER_RESULTS
+
+
+class TestCheckpointable(unittest.TestCase):
+    """Tests the Checkpointable API."""
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        ray.shutdown()
+
+    def test_checkpoint_backward_compatibility(self):
+        """Tests backward compat. of checkpoints created with older versions of ray."""
+
+        # Get the directory of the current script
+        old_checkpoints_dir = Path(__file__).parent.resolve() / "old_checkpoints"
+
+        from ray.rllib.utils.tests.old_checkpoints.current_config import (
+            config as current_config,
+        )
+
+        for ray_version_dir in old_checkpoints_dir.iterdir():
+            import re
+
+            if not ray_version_dir.is_dir() or not re.search(
+                r"\Wray_[0-9_]+$", str(ray_version_dir)
+            ):
+                continue
+            # Unzip checkpoint for that ray version into a temp directory.
+            with TemporaryDirectory() as temp_dir:
+                temp_path = Path(temp_dir)
+                # Extract the zip file to the temporary directory
+                shutil.unpack_archive(ray_version_dir / "checkpoint.zip", temp_path)
+                # Restore the algorithm from the (old) msgpack-checkpoint, using the
+                # current Ray version's `config` object.
+                algo = PPO.from_checkpoint(path=temp_dir, config=current_config)
+                learner_res = algo.train()[LEARNER_RESULTS]
+                # Assert that the correct per-policy learning rates were used.
+                assert (
+                    learner_res["p0"]["default_optimizer_learning_rate"] == 0.00005
+                    and learner_res["p1"]["default_optimizer_learning_rate"] == 0.0001
+                )
+                algo.stop()
+
+                # Second experiment: Add all the policies to the config again that were
+                # present when the checkpoint was taken and try `from_checkpoint` again.
+                expanded_config = current_config.copy(copy_frozen=False)
+                all_pols = {"p0", "p1", "p2", "p3"}
+                expanded_config.multi_agent(
+                    policies=all_pols,
+                    # Create some completely new mapping function (that has nothing to
+                    # do with the checkpointed one).
+                    policy_mapping_fn=(
+                        lambda aid, eps, _p=tuple(all_pols), **kw: random.choice(_p)
+                    ),
+                    policies_to_train=all_pols,
+                )
+                expanded_config.rl_module(
+                    algorithm_config_overrides_per_module={
+                        "p1": PPOConfig.overrides(lr=0.001),
+                        "p2": PPOConfig.overrides(lr=0.002),
+                        "p3": PPOConfig.overrides(lr=0.003),
+                    }
+                )
+                # The checkpoint has metrics that have not been cleared (since were created with the old metrics APIs)
+                # We therefore test two things:
+                # 1. That the metrics are correct according to the new metrics APIs (meaning that we average over the logged default_optimizer_learning_rate values.)
+                # 2. That the metrics are correct if we simply reduce first (clearing the values that have been left in the checkpoint by the old metrics APIs)
+
+                # 1.
+                algo = PPO.from_checkpoint(path=temp_dir, config=expanded_config)
+                learner_res = algo.train()[LEARNER_RESULTS]
+                # Assert that the correct per-policy learning rates were used.
+                assert (
+                    learner_res["p0"]["default_optimizer_learning_rate"] == 0.00005
+                    and learner_res["p1"]["default_optimizer_learning_rate"]
+                    == (0.001 + 0.0001) / 2
+                    and learner_res["p2"]["default_optimizer_learning_rate"]
+                    == (0.002 + 0.0002) / 2
+                    and learner_res["p3"]["default_optimizer_learning_rate"]
+                    == (0.003 + 0.0003) / 2
+                )
+                algo.stop()
+
+                # 2.
+                algo = PPO.from_checkpoint(path=temp_dir, config=expanded_config)
+                algo.metrics.reduce()
+                learner_res = algo.train()[LEARNER_RESULTS]
+                # Assert that the correct per-policy learning rates were used.
+                assert (
+                    learner_res["p0"]["default_optimizer_learning_rate"] == 0.00005
+                    and learner_res["p1"]["default_optimizer_learning_rate"] == 0.001
+                    and learner_res["p2"]["default_optimizer_learning_rate"] == 0.002
+                    and learner_res["p3"]["default_optimizer_learning_rate"] == 0.003
+                )
+                algo.stop()
+
+        print(f"Algorithm restored and trained once. Learner results={learner_res}.")
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))