Skip to content

[BUG] LLMMaskedCategorical is missing deterministic_sample #4465

Description

@YeonwooSung

Describe the bug

MaskedCategorical exposes both mode and deterministic_sample (the latter returns self.mode). LLMMaskedCategorical implements mode (argmax of the masked logits) but has no deterministic_sample.

Callers that query deterministic_sample (TensorDict CompositeDistribution, ExplorationType.DETERMINISTIC) therefore miss the attribute and fall back on mode with a warning, or raise AttributeError on a direct access.

This is the same gap MaskedCategorical had, but only on the LLM distribution. MaskedCategorical itself is unchanged.

To Reproduce

import torch
from torchrl.modules.distributions import LLMMaskedCategorical

logits = torch.tensor([[[1.0, 3.0, 2.0], [4.0, 0.0, 1.0]]])
mask = torch.ones(1, 2, dtype=torch.bool)
dist = LLMMaskedCategorical(logits=logits, mask=mask)
print(dist.mode)
print(dist.deterministic_sample)
AttributeError: 'LLMMaskedCategorical' object has no attribute 'deterministic_sample'

Expected behavior

dist.deterministic_sample exists and equals dist.mode (argmax of the masked logits), matching MaskedCategorical.

Reason and Possible fixes

Add next to mode in torchrl/modules/distributions/discrete.py:

@property
def deterministic_sample(self):
    return self.mode

Checklist

  • I have checked that there is no similar issue in the repo (required)
  • I have read the documentation (required)
  • I have provided a minimal working example to reproduce the bug (required)

Activity

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

Metadata

Metadata

Assignees

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