[RLlib] Algorithm/Policy checkpoint overhaul and Policy Model export (in native formats).

doc/source/serve/tutorials/rllib.md

@@ -47,8 +47,8 @@ def train_ppo_model():
     # Train for one iteration.
     algo.train()
     # Save state of the trained Algorithm in a checkpoint.
-    algo.save("/tmp/rllib_checkpoint")
-    return "/tmp/rllib_checkpoint/checkpoint_000001/checkpoint-1"
+    checkpoint_dir = algo.save("/tmp/rllib_checkpoint")
+    return checkpoint_dir
 
 
 checkpoint_path = train_ppo_model()

python/ray/train/rl/rl_checkpoint.py

@@ -1,9 +1,11 @@
 import os
+from packaging import version
 from typing import Optional
 
 from ray.air.checkpoint import Checkpoint
 import ray.cloudpickle as cpickle
 from ray.rllib.policy.policy import Policy
+from ray.rllib.utils.checkpoints import get_checkpoint_info
 from ray.rllib.utils.typing import EnvType
 from ray.util.annotations import PublicAPI
 
@@ -30,6 +32,16 @@ def get_policy(self, env: Optional[EnvType] = None) -> Policy:
         Returns:
             The policy stored in this checkpoint.
         """
+        # TODO: Deprecate this RLCheckpoint class (or move all our
+        #  Algorithm/Policy.from_checkpoint utils into here).
+        # If newer checkpoint version -> Use `Policy.from_checkpoint()` util.
+        checkpoint_info = get_checkpoint_info(checkpoint=self)
+        if checkpoint_info["checkpoint_version"] > version.Version("0.1"):
+            # Since we have an Algorithm checkpoint, will extract all policies in that
+            # Algorithm -> need to index into "default_policy" in the returned dict.
+            return Policy.from_checkpoint(checkpoint=self)["default_policy"]
+
+        # Older checkpoint version.
         with self.as_directory() as checkpoint_path:
             trainer_class_path = os.path.join(checkpoint_path, RL_TRAINER_CLASS_FILE)
             config_path = os.path.join(checkpoint_path, RL_CONFIG_FILE)

python/ray/tune/tests/test_tune_restore.py

@@ -46,7 +46,7 @@ def setUp(self):
 
         logdir = os.path.expanduser(os.path.join(tmpdir, test_name))
         self.logdir = logdir
-        self.checkpoint_path = recursive_fnmatch(logdir, "checkpoint-1")[0]
+        self.checkpoint_path = recursive_fnmatch(logdir, "algorithm_state.pkl")[0]
 
     def tearDown(self):
         shutil.rmtree(self.logdir)

rllib/BUILD

@@ -766,6 +766,13 @@ py_test(
     data = ["tests/data/cartpole/small.json"],
 )
 
+py_test(
+    name = "test_algorithm_export_checkpoint",
+    tags = ["team:rllib", "algorithms_dir", "algorithms_dir_generic"],
+    size = "medium",
+    srcs = ["algorithms/tests/test_algorithm_export_checkpoint.py"],
+)
+
 py_test(
     name = "test_callbacks",
     tags = ["team:rllib", "algorithms_dir", "algorithms_dir_generic"],
@@ -1876,6 +1883,13 @@ py_test(
     srcs = ["policy/tests/test_compute_log_likelihoods.py"]
 )
 
+py_test(
+    name = "policy/tests/test_export_checkpoint_and_model",
+    tags = ["team:rllib", "policy"],
+    size = "large",
+    srcs = ["policy/tests/test_export_checkpoint_and_model.py"]
+)
+
 py_test(
     name = "policy/tests/test_multi_agent_batch",
     tags = ["team:rllib", "policy"],
@@ -1919,6 +1933,14 @@ py_test(
 # Tag: utils
 # --------------------------------------------------------------------
 
+# Checkpoint Utils
+py_test(
+    name = "test_checkpoint_utils",
+    tags = ["team:rllib", "utils"],
+    size = "small",
+    srcs = ["utils/tests/test_checkpoint_utils.py"]
+)
+
 py_test(
     name = "test_errors",
     tags = ["team:rllib", "utils"],
@@ -2078,7 +2100,8 @@ py_test(
     name = "tests/backward_compat/test_backward_compat",
     tags = ["team:rllib", "tests_dir", "tests_dir_B"],
     size = "medium",
-    srcs = ["tests/backward_compat/test_backward_compat.py"]
+    srcs = ["tests/backward_compat/test_backward_compat.py"],
+    data = glob(["tests/backward_compat/checkpoints/**"]),
 )
 
 py_test(
@@ -2168,13 +2191,6 @@ py_test(
     srcs = ["tests/test_execution.py"]
 )
 
-py_test(
-    name = "tests/test_export",
-    tags = ["team:rllib", "tests_dir", "tests_dir_E"],
-    size = "medium",
-    srcs = ["tests/test_export.py"]
-)
-
 py_test(
     name = "tests/test_filters",
     tags = ["team:rllib", "tests_dir", "tests_dir_F"],
@@ -2940,6 +2956,16 @@ py_test(
     tags = ["team:rllib", "exclusive", "examples", "examples_E", "no_main"],
     size = "medium",
     srcs = ["examples/export/onnx_tf.py"],
+    args = ["--framework=tf"],
+)
+
+py_test(
+    name = "examples/export/onnx_tf2",
+    main = "examples/export/onnx_tf.py",
+    tags = ["team:rllib", "exclusive", "examples", "examples_E", "no_main"],
+    size = "medium",
+    srcs = ["examples/export/onnx_tf.py"],
+    args = ["--framework=tf2"],
 )
 
 py_test(