[RLlib] TFPolicy.export_model: Add timestep placeholder to model's signature, if needed. (#13988)

rllib/policy/tf_policy.py

@@ -709,9 +709,14 @@ class TFPolicy(Policy):
             input_signature["prev_reward"] = \
                 tf1.saved_model.utils.build_tensor_info(
                     self._prev_reward_input)
+
         input_signature["is_training"] = \
             tf1.saved_model.utils.build_tensor_info(self._is_training)
 
+        if self._timestep is not None:
+            input_signature["timestep"] = \
+                tf1.saved_model.utils.build_tensor_info(self._timestep)
+
         for state_input in self._state_inputs:
             input_signature[state_input.name] = \
                 tf1.saved_model.utils.build_tensor_info(state_input)

rllib/tests/test_export.py

@@ -6,8 +6,11 @@ import unittest
 
 import ray
 from ray.rllib.agents.registry import get_trainer_class
+from ray.rllib.utils.framework import try_import_tf
 from ray.tune.trial import ExportFormat
 
+tf1, tf, tfv = try_import_tf()
+
 CONFIGS = {
     "A3C": {
         "explore": False,
@@ -105,6 +108,11 @@ def export_test(alg_name, failures):
             or not valid_tf_checkpoint(os.path.join(export_dir,
                                                     ExportFormat.CHECKPOINT)):
         failures.append(alg_name)
+
+    # Test loading the exported model.
+    model = tf.saved_model.load(os.path.join(export_dir, ExportFormat.MODEL))
+    assert model
+
     shutil.rmtree(export_dir)