[RLlib] Upgrade tf eager code to no longer use `experimental_relax_sh…

…apes` (but `reduce_retracing` instead). (ray-project#29214) Signed-off-by: Artur Niederfahrenhorst <artur@anyscale.com>
clarng · Jan 31, 2023 · 91111bd · 91111bd
1 parent d370bbf
commit 91111bd
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/rllib/policy/eager_tf_policy.py b/rllib/policy/eager_tf_policy.py
@@ -180,7 +180,7 @@ def compute_actions_from_input_dict(
                     tf.function(
                         super(TracedEagerPolicy, self)._compute_actions_helper,
                         autograph=False,
-                        experimental_relax_shapes=True,
+                        reduce_retracing=True,
                     )
                 )
                 self._traced_compute_actions_helper = True
@@ -206,7 +206,7 @@ def learn_on_batch(self, samples):
                     tf.function(
                         super(TracedEagerPolicy, self)._learn_on_batch_helper,
                         autograph=False,
-                        experimental_relax_shapes=True,
+                        reduce_retracing=True,
                     )
                 )
                 self._traced_learn_on_batch_helper = True
@@ -226,7 +226,7 @@ def compute_gradients(self, samples: SampleBatch) -> ModelGradients:
                     tf.function(
                         super(TracedEagerPolicy, self)._compute_gradients_helper,
                         autograph=False,
-                        experimental_relax_shapes=True,
+                        reduce_retracing=True,
                     )
                 )
                 self._traced_compute_gradients_helper = True
@@ -246,7 +246,7 @@ def apply_gradients(self, grads: ModelGradients) -> None:
                     tf.function(
                         super(TracedEagerPolicy, self)._apply_gradients_helper,
                         autograph=False,
-                        experimental_relax_shapes=True,
+                        reduce_retracing=True,
                     )
                 )
                 self._traced_apply_gradients_helper = True