Fix numpy test issue.

PiperOrigin-RevId: 683926353 Change-Id: I88a15bf0cb64ee60a95bd6567128e61272e4e39c
tensorflow · Oct 9, 2024 · a809aa8 · a809aa8
1 parent 9919a23
commit a809aa8
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tf_agents/bandits/agents/exp3_mixture_agent_test.py b/tf_agents/bandits/agents/exp3_mixture_agent_test.py
@@ -195,7 +195,9 @@ def testMixtureUpdate(
     reward_aggregates = self.evaluate(
         mixed_agent._variable_collection.reward_aggregates
     )
-    self.assertAllInSet(reward_aggregates[: num_agents - 1], [0.999])
+    self.assertAllClose(
+        reward_aggregates[: num_agents - 1], [0.999] * (num_agents - 1)
+    )
     agent_prob = 1 / num_agents
     est_rewards = 0.5 / agent_prob
     per_step_update = est_rewards