[mottopanikeiku]

Description

Categorical.expand discarded its mask, allowing forbidden random actions. MultiCategorical.expand added an extra dimension to its mask and raised on valid input. I now pass the existing mask to the constructors, which already broadcast it to the new shape. I also documented that expansion preserves outcome masks.

A small reproducer:

import torch
from torchrl.data import Categorical, MultiCategorical

torch.manual_seed(0)
spec = Categorical(3, mask=torch.tensor([False, True, False]))
print(spec.expand(8).rand().tolist())
spec = MultiCategorical([3, 2], mask=torch.tensor([False, True, False, True, False]))
print(spec.expand(2, 2).rand().tolist())

	main	This PR
Categorical sample	[2, 0, 2, 0, 1, 0, 1, 1]	[1, 1, 1, 1, 1, 1, 1, 1]
MultiCategorical sample	RuntimeError: Cannot expand mask to the desired shape.	[[1, 0], [1, 0]]

Tests:

• Four new cases in TestExpand fail on main. They cover singleton expansion, a new leading batch dimension, per-entry masks, and sampled action values.
• With the fix, test/test_specs.py: 1824 passed, 9 skipped.
• pre-commit passes on the changed files.
• I tested the Python sources without building the optional C++ extension; these spec operations do not use it.

Motivation and Context

Batching a masked action spec should not enable forbidden actions or fail while broadcasting its mask.

close #4588

• I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

• Bug fix (non-breaking change which fixes an issue)

Checklist

• I have read the CONTRIBUTION guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.

I used AI assistance (Claude and GPT-6 Astra) to find, fix and test this bug.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4590

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⚠️ 16 Awaiting Approval

As of commit 9b01f22 with merge base e7a05a7 ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build Aarch64 Linux Wheels (gh)
• Build Linux Wheels (gh)
• Build M1 Wheels (gh)
• Build Windows Wheels (gh)
• Continuous Benchmark (PR) (gh)
• Examples Tests on Linux (gh)
• Generate documentation (gh)
• Habitat Tests on Linux (gh)
• Libs Tests on Linux (gh)
• Lint (gh)
• LLM Tests on Linux (gh)
• Push Binary Nightly (gh)
• SOTA Tests on Linux (gh)
• Tutorials Tests on Linux (gh)
• Unit-tests on Linux (gh)
• Validate Test Partitioning (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.