fix: clear cached probs in MaskableCategorical.apply_masking to avoid Simplex validate_args error (torch 2.9+) (#326)

* fix: clear cached probs in MaskableCategorical.apply_masking to avoid Simplex validate_args error (torch 2.9+, issue #322)

* fix: clear cached probs in apply_masking, add regression test, update changelog (#322)

* style: apply black formatting to test_distributions.py

* Update mask PPO test and add comments

* Simplify implementation

* Apply suggestions from code review

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update comments

---------

Co-authored-by: Antonin RAFFIN <antonin.raffin@ensta.org>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: 3 people

docs/misc/changelog.md

@@ -16,9 +16,10 @@
 
 ### Bug Fixes:
 
-- Fix `MaskablePPO` and `RecurrentPPO` inaccurate `n_updates` counting when `target_kl` early exits the training loop
-- Fix `RecurrentPPO` and `MaskablePPO` forward and predict do not reshape action before clip it (@immortal-boy)
+- Fixed `MaskablePPO` and `RecurrentPPO` inaccurate `n_updates` counting when `target_kl` early exits the training loop
+- Fixed `RecurrentPPO` and `MaskablePPO` `forward` and `predict` not reshaping the action before clipping it (@immortal-boy)
 - Do not call `forward()` method directly in `RecurrentPPO` (@immortal-boy)
+- Fixed `MaskableCategorical.apply_masking()` crashing with `ValueError: Simplex` when cached `probs` deviate from sum=1 in float32 with large action spaces (torch 2.9+) (@kirann-05)
 
 ### Deprecations:
 

sb3_contrib/common/maskable/distributions.py

@@ -7,7 +7,6 @@
 from stable_baselines3.common.distributions import Distribution
 from torch import nn
 from torch.distributions import Categorical
-from torch.distributions.utils import logits_to_probs
 
 SelfMaskableCategoricalDistribution = TypeVar("SelfMaskableCategoricalDistribution", bound="MaskableCategoricalDistribution")
 SelfMaskableMultiCategoricalDistribution = TypeVar(
@@ -65,10 +64,10 @@ def apply_masking(self, masks: MaybeMasks) -> None:
             logits = self._original_logits
 
         # Reinitialize with updated logits
-        super().__init__(logits=logits)
-
-        # self.probs may already be cached, so we must force an update
-        self.probs = logits_to_probs(self.logits)
+        # Clear cached probs before reinit to avoid validate_args Simplex error
+        # when stale float32 probs deviate from sum=1 by >1e-6 (torch 2.9+, many categories)
+        self.__dict__.pop("probs", None)
+        super().__init__(logits=logits, validate_args=self._validate_args)
 
     def entropy(self) -> th.Tensor:
         if self.masks is None:

tests/test_distributions.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 import torch as th
+from torch.distributions.utils import logits_to_probs
 
 from sb3_contrib.common.maskable.distributions import (
     MaskableBernoulliDistribution,
@@ -16,8 +17,8 @@ def test_applying_mask(self):
         Show that probs change as a result of masking
         """
 
-        starting_probs = th.Tensor([[0.2, 0.2, 0.6], [1, 0, 0]])
-        expected_probs = th.Tensor([[0, 0.25, 0.75], [0, 0.5, 0.5]])
+        starting_probs = th.tensor([[0.2, 0.2, 0.6], [1, 0, 0]])
+        expected_probs = th.tensor([[0, 0.25, 0.75], [0, 0.5, 0.5]])
         mask = np.array([[False, True, True], [False, True, True]])
 
         distribution = MaskableCategorical(probs=starting_probs)
@@ -29,8 +30,8 @@ def test_modifying_mask(self):
         Show that masks apply independently of each other
         """
 
-        starting_probs = th.Tensor([[0.2, 0.2, 0.6], [1, 0, 0]])
-        expected_probs = th.Tensor([[0.5, 0.5, 0], [0, 1, 0]])
+        starting_probs = th.tensor([[0.2, 0.2, 0.6], [1, 0, 0]])
+        expected_probs = th.tensor([[0.5, 0.5, 0], [0, 1, 0]])
         first_mask = np.array([[False, True, True], [False, True, True]])
         second_mask = np.array([[True, True, False], [False, True, False]])
 
@@ -52,7 +53,7 @@ def test_removing_mask(self):
         Show that masking may be unapplied to recover original probs
         """
 
-        starting_probs = th.Tensor([[0.2, 0.2, 0.6], [1, 0, 0]])
+        starting_probs = th.tensor([[0.2, 0.2, 0.6], [1, 0, 0]])
         mask = np.array([[False, True, True], [False, True, True]])
 
         distribution = MaskableCategorical(probs=starting_probs)
@@ -64,14 +65,47 @@ def test_removing_mask(self):
     def test_masking_affects_entropy(self):
         # All outcomes equally likely
         NUM_DIMS = 3
-        logits = th.Tensor([[0] * NUM_DIMS])
+        logits = th.zeros((1, NUM_DIMS), dtype=th.float32)
         dist = MaskableCategorical(logits=logits)
 
-        # For each possible number of valid actions v, show that e^entropy == v
-        for v in range(1, NUM_DIMS + 1):
-            masks = [j < v for j in range(NUM_DIMS)]
+        # For each possible number of valid actions valid_action, show that e^entropy == valid_action
+        for valid_action in range(1, NUM_DIMS + 1):
+            masks = [j < valid_action for j in range(NUM_DIMS)]
             dist.apply_masking(masks)
-            assert int(dist.entropy().exp()) == v
+            assert int(dist.entropy().exp()) == valid_action
+
+    def test_apply_masking_no_simplex_error_with_cached_probs(self):
+        """
+        Regression test for issue #322.
+        apply_masking() should not raise a Simplex validate_args error
+        when probs are cached before masking is re-applied (torch 2.9+, large action spaces).
+        """
+        n = 992
+        delta = 17
+        # Set all logits to -delta, then set the first one to 0 to make it the most likely
+        logits = th.full((1, n), -delta, dtype=th.float32)
+        logits[0, 0] = 0.0
+
+        # Expected probs: first dim is close to one, rest are almost zero after the softmax
+        expected_proba = th.cat((th.ones(1, 1), th.zeros(1, n - 1)), dim=1)
+
+        distribution = MaskableCategorical(logits=logits, validate_args=True)
+        _ = distribution.probs  # cache probs on the instance
+
+        th.testing.assert_close(distribution.probs, expected_proba, rtol=5e-5, atol=5e-5)
+        # Apply softmax, should be within the default tolerance
+        th.testing.assert_close(distribution.probs, logits_to_probs(logits))
+
+        mask = th.zeros((1, n), dtype=th.bool)
+        # Only the first action is valid, mask the rest of the actions,
+        # the rtol and atol when comparing to expected_proba can be kept to the default values
+        mask[0, 0] = True
+
+        # Should not raise ValueError: Simplex constraint
+        distribution.apply_masking(mask.numpy())
+        th.testing.assert_close(distribution.probs, expected_proba)
+        # After masking, it should deviate from the original probs
+        assert not th.allclose(distribution.probs, logits_to_probs(logits))
 
 
 class TestMaskableCategoricalDistribution:
@@ -133,7 +167,7 @@ def test_dim_masking(self):
         NUM_DIMS = 2
         dist = MaskableCategoricalDistribution(NUM_DIMS)
 
-        logits = th.Tensor([[0] * NUM_DIMS])
+        logits = th.zeros((1, NUM_DIMS), dtype=th.float32)
         dist.proba_distribution(logits)
 
         assert (dist.distribution.probs == 0.5).all()
@@ -216,7 +250,7 @@ def test_dim_masking(self):
         NUM_CATS = 3
         dist = MaskableMultiCategoricalDistribution([DIMS_PER_CAT] * NUM_CATS)
 
-        logits = th.Tensor([[0] * DIMS_PER_CAT * NUM_CATS])
+        logits = th.tensor([[0] * DIMS_PER_CAT * NUM_CATS], dtype=th.float32)
         dist.proba_distribution(logits)
 
         assert len(dist.distributions) == NUM_CATS
@@ -304,7 +338,7 @@ def test_dim_masking(self):
         BINARY_STATES = 2
         dist = MaskableBernoulliDistribution(NUM_DIMS)
 
-        logits = th.Tensor([[0] * BINARY_STATES * NUM_DIMS])
+        logits = th.tensor([[0] * BINARY_STATES * NUM_DIMS], dtype=th.float32)
         dist.proba_distribution(logits)
 
         assert len(dist.distributions) == NUM_DIMS