Fix SquareCBExploration.act for batched values (broadcast crash + per-row probability sum)

pearl/policy_learners/exploration_modules/contextual_bandits/squarecb_exploration.py

@@ -79,19 +79,23 @@ def act(
         values = values.view(-1, action_space.n)  # (batch_size, action_space.n)
         values = self.clamp(values)
         max_val, max_indices = torch.max(values, dim=1)
-        max_val.repeat(1, action_space.n)
-        empirical_gaps = max_val - values
+        empirical_gaps = max_val.unsqueeze(1) - values
 
         # Construct probability distribution over actions and sample from it
         selected_actions = torch.zeros((values.size(dim=0),), dtype=torch.int)
-        prob_policy = self.get_unnormalize_prob(empirical_gaps, max_val, action_space.n)
         for batch_ind in range(values.size(dim=0)):
+            # Build the unnormalized policy for this row. Passing the scalar row
+            # maximum keeps this compatible with subclasses such as
+            # FastCBExploration whose get_unnormalize_prob branches on max_val.
+            prob_policy = self.get_unnormalize_prob(
+                empirical_gaps[batch_ind, :], max_val[batch_ind], action_space.n
+            )
             # Get sum of all the probabilities besides the maximum
-            prob_policy[batch_ind, max_indices[batch_ind]] = 0.0
+            prob_policy[max_indices[batch_ind]] = 0.0
             complementary_sum = torch.sum(prob_policy)
-            prob_policy[batch_ind, max_indices[batch_ind]] = 1.0 - complementary_sum
+            prob_policy[max_indices[batch_ind]] = 1.0 - complementary_sum
             # Sample from SquareCB update rule
-            dist_policy = Categorical(prob_policy[batch_ind, :])
+            dist_policy = Categorical(prob_policy)
             selected_actions[batch_ind] = dist_policy.sample()
 
         return selected_actions.squeeze(-1)

test/unit/with_pytorch/test_squarecb_exploration.py

@@ -0,0 +1,100 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+#
+
+# pyre-strict
+
+import unittest
+
+import torch
+from pearl.policy_learners.exploration_modules.contextual_bandits.squarecb_exploration import (
+    FastCBExploration,
+    SquareCBExploration,
+)
+from pearl.utils.instantiations.spaces.discrete_action import DiscreteActionSpace
+
+
+class TestSquareCBExploration(unittest.TestCase):
+    """Tests for SquareCBExploration.act over batched value tensors."""
+
+    def test_act_batched_states_does_not_crash(self) -> None:
+        # batch_size (2) != action_count (3) exercises the gap broadcasting:
+        # ``empirical_gaps = max_val.unsqueeze(1) - values`` must align the
+        # per-row max with the (batch_size, action_count) values.
+        action_space = DiscreteActionSpace(
+            actions=[torch.tensor([i]) for i in range(3)]
+        )
+        exploration = SquareCBExploration(gamma=10.0)
+        values = torch.tensor([[0.10, 0.20, 0.90], [0.80, 0.30, 0.10]])
+        torch.manual_seed(0)
+        actions = exploration.act(
+            subjective_state=torch.zeros(2, 4),
+            action_space=action_space,
+            values=values,
+        )
+        self.assertEqual(actions.shape[0], 2)
+        self.assertTrue(int(actions.min()) >= 0)
+        self.assertTrue(int(actions.max()) < action_space.n)
+
+    def test_act_probabilities_are_per_row_valid_distributions(self) -> None:
+        # The greedy action's residual probability must be computed from the
+        # current row only (sum over that row), so every row of the sampling
+        # distribution sums to 1 and the greedy action carries the most mass.
+        action_space = DiscreteActionSpace(
+            actions=[torch.tensor([i]) for i in range(3)]
+        )
+        exploration = SquareCBExploration(gamma=10.0)
+        values = torch.tensor([[0.10, 0.20, 0.90], [0.80, 0.30, 0.10]])
+
+        # Reconstruct, per row, the distribution act() builds via the module's
+        # own get_unnormalize_prob (no randomness involved).
+        max_val, max_indices = torch.max(values, dim=1)
+        empirical_gaps = max_val.unsqueeze(1) - values
+        rows = []
+        for b in range(values.size(0)):
+            prob = exploration.get_unnormalize_prob(
+                empirical_gaps[b, :], max_val[b], action_space.n
+            )
+            prob[max_indices[b]] = 0.0
+            prob[max_indices[b]] = 1.0 - torch.sum(prob)
+            rows.append(prob)
+        prob = torch.stack(rows)
+
+        self.assertTrue(torch.allclose(prob.sum(dim=1), torch.ones(2), atol=1e-6))
+        # Greedy action (argmax of values) should be the most probable per row.
+        self.assertTrue(torch.equal(prob.argmax(dim=1), values.argmax(dim=1)))
+
+    def test_act_single_state(self) -> None:
+        action_space = DiscreteActionSpace(
+            actions=[torch.tensor([i]) for i in range(3)]
+        )
+        exploration = SquareCBExploration(gamma=10.0)
+        torch.manual_seed(0)
+        action = exploration.act(
+            subjective_state=torch.zeros(1, 4),
+            action_space=action_space,
+            values=torch.tensor([[0.10, 0.20, 0.90]]),
+        )
+        self.assertTrue(0 <= int(action) < action_space.n)
+
+    def test_fastcb_act_batched_states(self) -> None:
+        # FastCBExploration inherits act() and overrides get_unnormalize_prob
+        # with a branch on max_val; act() must therefore feed it a scalar row
+        # maximum so batched input does not raise on an ambiguous truth value.
+        action_space = DiscreteActionSpace(
+            actions=[torch.tensor([i]) for i in range(3)]
+        )
+        exploration = FastCBExploration(gamma=10.0)
+        values = torch.tensor([[0.10, 0.20, 0.90], [0.80, 0.30, 0.10]])
+        torch.manual_seed(0)
+        actions = exploration.act(
+            subjective_state=torch.zeros(2, 4),
+            action_space=action_space,
+            values=values,
+        )
+        self.assertEqual(actions.shape[0], 2)
+        self.assertTrue(int(actions.min()) >= 0)
+        self.assertTrue(int(actions.max()) < action_space.n)