Adds Gin configs for NeuralFalconAgent

PiperOrigin-RevId: 474813720 Change-Id: Ib6e9a0694cc4b9b9c28c79e9bc82e987ffbe48eb
tensorflow · Sep 16, 2022 · 3409e73 · 3409e73
1 parent b4ac366
commit 3409e73
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tf_agents/bandits/policies/falcon_reward_prediction_policy.py b/tf_agents/bandits/policies/falcon_reward_prediction_policy.py
@@ -144,6 +144,7 @@ def __init__(self,
   def _get_exploitation_coefficient(self) -> types.FloatOrReturningFloat:
     coef = self._exploitation_coefficient() if callable(
         self._exploitation_coefficient) else self._exploitation_coefficient
+    coef = tf.cast(coef, dtype=tf.float32)
     return tf.maximum(coef, 0.0)
 
   @property