Skip to content

[BUG] Sparse MaskedCategorical entropy includes padding or returns NaN #4512

Description

@AHMETHAKANBEZIR1

Bug description

On current main 8f92a50cb1e8543d6d7545605e4e5b399224928b, sparse MaskedCategorical.entropy() does not exclude padded indices. With default neg_inf=-inf, a padded logit produces 0 * -inf = NaN. With finite neg_inf, padding incorrectly contributes to entropy instead of being excluded, unlike the dense-mask path and the method's documented valid-outcome behavior.

Reproduction

import torch
from torchrl.modules.distributions import MaskedCategorical

logits = torch.tensor([1., 2., 3., 4.])
indices = torch.tensor([0, 2, -1])
for neg_inf in (float('-inf'), -10.):
    dist = MaskedCategorical(logits=logits, indices=indices,
                             padding_value=-1, neg_inf=neg_inf)
    print(dist.entropy())  # NaN or an incorrect finite value
print(torch.distributions.Categorical(logits=logits[[0, 2]]).entropy())
# tensor(0.3653): entropy over the two valid actions

Environment: Windows, Python 3.12.14 / PyTorch 2.10 CPU, TorchRL Python source checkout, TensorDict 0.14.2 official wheel. The optional TorchRL C++ extension is not built; this reproduces in the Python distribution implementation.

I intend to mask padding/non-finite zero-probability contributions before the existing entropy normalization, and test entropy and gradients against independent compact torch.distributions.Categorical references. This is separate from sample/log_prob shape handling in #4511 and mode ties in #4509. Investigated with Codex assistance; reproduction run locally.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Labels

No labels
No labels

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions