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
Describe the bug
MaskedCategoricalexposes bothmodeanddeterministic_sample(the latter returnsself.mode).LLMMaskedCategoricalimplementsmode(argmax of the masked logits) but has nodeterministic_sample.Callers that query
deterministic_sample(TensorDictCompositeDistribution,ExplorationType.DETERMINISTIC) therefore miss the attribute and fall back onmodewith a warning, or raiseAttributeErroron a direct access.This is the same gap
MaskedCategoricalhad, but only on the LLM distribution.MaskedCategoricalitself is unchanged.To Reproduce
Expected behavior
dist.deterministic_sampleexists and equalsdist.mode(argmax of the masked logits), matchingMaskedCategorical.Reason and Possible fixes
Add next to
modeintorchrl/modules/distributions/discrete.py:Checklist