From aec03656d5b3c90b17152afd037e7257e6ced46e Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Thu, 27 Feb 2020 00:22:54 +0100 Subject: [PATCH] [RLlib] TupleActions cannot be exported by Policy: Fixes issues 7231 and 5593. #7333 --- rllib/policy/tf_policy.py | 6 ++++-- rllib/tests/test_checkpoint_restore.py | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) 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),