diff --git a/rllib/policy/tf_policy.py b/rllib/policy/tf_policy.py
index 516b7288a..9a6164845 100644
--- a/rllib/policy/tf_policy.py
+++ b/rllib/policy/tf_policy.py
@@ -493,8 +493,10 @@ class TFPolicy(Policy):
 
         # build output signatures
         output_signature = self._extra_output_signature_def()
-        output_signature["actions"] = \
-            tf.saved_model.utils.build_tensor_info(self._sampled_action)
+        for i, a in enumerate(tf.nest.flatten(self._sampled_action)):
+            output_signature["actions_{}".format(i)] = \
+                tf.saved_model.utils.build_tensor_info(a)
+
         for state_output in self._state_outputs:
             output_signature[state_output.name] = \
                 tf.saved_model.utils.build_tensor_info(state_output)
diff --git a/rllib/tests/test_checkpoint_restore.py b/rllib/tests/test_checkpoint_restore.py
index 2f35fdfe7..903d135da 100644
--- a/rllib/tests/test_checkpoint_restore.py
+++ b/rllib/tests/test_checkpoint_restore.py
@@ -80,7 +80,7 @@ def test_ckpt_restore(use_object_store, alg_name, failures):
     else:
         alg2.restore(alg1.save())
 
-    for _ in range(10):
+    for _ in range(5):
         if "DDPG" in alg_name or "SAC" in alg_name:
             obs = np.clip(
                 np.random.uniform(size=3),