[Tune] Add export_formats option to export policy graphs (#3868)

In earlier PRs, PR#3585 and PR#3637, export_policy_model and export_policy_checkpoint were introduced for users to export TensorFlow model and checkpoint.

For Ray Tune users, these APIs are not accessible through YAML configurations.

In this pull request, export_formats option is provided to enable users to choose the desired export format.

python/ray/rllib/agents/agent.py

@@ -22,7 +22,7 @@ from ray.rllib.utils.annotations import override, PublicAPI, DeveloperAPI
 from ray.rllib.utils import FilterManager, deep_update, merge_dicts
 from ray.tune.registry import ENV_CREATOR, register_env, _global_registry
 from ray.tune.trainable import Trainable
-from ray.tune.trial import Resources
+from ray.tune.trial import Resources, ExportFormat
 from ray.tune.logger import UnifiedLogger
 from ray.tune.result import DEFAULT_RESULTS_DIR
 
@@ -602,6 +602,20 @@ class Agent(Trainable):
             input_evaluation_method=config["input_evaluation"],
             output_creator=output_creator)
 
+    @override(Trainable)
+    def _export_model(self, export_formats, export_dir):
+        ExportFormat.validate(export_formats)
+        exported = {}
+        if ExportFormat.CHECKPOINT in export_formats:
+            path = os.path.join(export_dir, ExportFormat.CHECKPOINT)
+            self.export_policy_checkpoint(path)
+            exported[ExportFormat.CHECKPOINT] = path
+        if ExportFormat.MODEL in export_formats:
+            path = os.path.join(export_dir, ExportFormat.MODEL)
+            self.export_policy_model(path)
+            exported[ExportFormat.MODEL] = path
+        return exported
+
     def __getstate__(self):
         state = {}
         if hasattr(self, "local_evaluator"):

python/ray/rllib/evaluation/tf_policy_graph.py

@@ -3,6 +3,7 @@ from __future__ import division
 from __future__ import print_function
 
 import os
+import errno
 import logging
 import tensorflow as tf
 import numpy as np
@@ -205,6 +206,12 @@ class TFPolicyGraph(PolicyGraph):
     @override(PolicyGraph)
     def export_checkpoint(self, export_dir, filename_prefix="model"):
         """Export tensorflow checkpoint to export_dir."""
+        try:
+            os.makedirs(export_dir)
+        except OSError as e:
+            # ignore error if export dir already exists
+            if e.errno != errno.EEXIST:
+                raise
         save_path = os.path.join(export_dir, filename_prefix)
         with self._sess.graph.as_default():
             saver = tf.train.Saver()

python/ray/rllib/test/test_checkpoint_restore.py

@@ -10,6 +10,7 @@ import numpy as np
 import ray
 
 from ray.rllib.agents.registry import get_agent_class
+from ray.tune.trial import ExportFormat
 
 
 def get_mean_action(alg, obs):
@@ -89,6 +90,15 @@ def test_ckpt_restore(use_object_store, alg_name, failures):
 
 
 def test_export(algo_name, failures):
+    def valid_tf_model(model_dir):
+        return os.path.exists(os.path.join(model_dir, "saved_model.pb")) \
+            and os.listdir(os.path.join(model_dir, "variables"))
+
+    def valid_tf_checkpoint(checkpoint_dir):
+        return os.path.exists(os.path.join(checkpoint_dir, "model.meta")) \
+            and os.path.exists(os.path.join(checkpoint_dir, "model.index")) \
+            and os.path.exists(os.path.join(checkpoint_dir, "checkpoint"))
+
     cls = get_agent_class(algo_name)
     if "DDPG" in algo_name:
         algo = cls(config=CONFIGS[name], env="Pendulum-v0")
@@ -102,16 +112,22 @@ def test_export(algo_name, failures):
     export_dir = "/tmp/export_dir_%s" % algo_name
     print("Exporting model ", algo_name, export_dir)
     algo.export_policy_model(export_dir)
-    if not os.path.exists(os.path.join(export_dir, "saved_model.pb")) \
-            or not os.listdir(os.path.join(export_dir, "variables")):
+    if not valid_tf_model(export_dir):
         failures.append(algo_name)
     shutil.rmtree(export_dir)
 
     print("Exporting checkpoint", algo_name, export_dir)
     algo.export_policy_checkpoint(export_dir)
-    if not os.path.exists(os.path.join(export_dir, "model.meta")) \
-            or not os.path.exists(os.path.join(export_dir, "model.index")) \
-            or not os.path.exists(os.path.join(export_dir, "checkpoint")):
+    if not valid_tf_checkpoint(export_dir):
+        failures.append(algo_name)
+    shutil.rmtree(export_dir)
+
+    print("Exporting default policy", algo_name, export_dir)
+    algo.export_model([ExportFormat.CHECKPOINT, ExportFormat.MODEL],
+                      export_dir)
+    if not valid_tf_model(os.path.join(export_dir, ExportFormat.MODEL)) \
+            or not valid_tf_checkpoint(os.path.join(export_dir,
+                                                    ExportFormat.CHECKPOINT)):
         failures.append(algo_name)
     shutil.rmtree(export_dir)