[kjgggggg]

Description

Indexing a masked MultiOneHot spec drops its mask, and indexing a masked OneHot spec with an ellipsis indexes the wrong dim of the mask:

import torch
from torchrl.data import Composite, MultiOneHot, OneHot

batch = torch.arange(12).view(3, 4)
# every element of the (3, 4) batch allows a single outcome per group
mask = torch.cat(
    [torch.nn.functional.one_hot(batch % 2, 2), torch.nn.functional.one_hot(batch % 3, 3)], -1
).bool()
spec = MultiOneHot([2, 3], shape=(3, 4, 5), mask=mask)
print(spec[0].mask is None)

mask = torch.nn.functional.one_hot(batch % 5, 5).bool()
spec = OneHot(5, shape=(3, 4, 5), mask=mask)
spec[..., 1]                                  # RuntimeError on main
Composite(action=spec, shape=(3, 4))[..., 1]  # RuntimeError on main

On main, spec[0].mask is None prints True (so spec[0].rand() draws from all outcomes) and both ellipsis indexings raise RuntimeError: Cannot expand mask to the desired shape.. When the number of outcomes equals the last batch dim, OneHot(4, shape=(3, 4, 4), mask=m)[..., 1] does not raise but gets m[..., 1], a slice along the outcome dim, as its mask. With this PR, spec[0].mask is mask[0], and spec[..., 1].mask is the mask of the batch elements [:, 1].

Causes, in torchrl/data/tensor_specs.py:

• MultiOneHot.__getitem__ rebuilds the spec without passing mask.
• OneHot.__getitem__ computes the new shape with _shape_indexing(self.shape[:-1], idx), which never indexes the last (outcome) dim, but indexes the mask with self.mask[idx]. The mask has the outcome dim as its last dim, so an ellipsis in idx expands up to the outcome dim instead of the last batch dim. Composite.__getitem__ forwards the index to its leaves, so a composite spec with a masked OneHot leaf fails the same way.

Fix: a helper _index_outcome_mask moves the outcome dim of the mask in front, applies (slice(None), *idx) so that the index only spans the batch dims, and moves the outcome dim back. OneHot.__getitem__ uses it, and MultiOneHot.__getitem__ now passes its indexed mask on. Categorical (whose mask has the same layout) is left alone here, since #4624 changes its __getitem__; MultiCategorical.__getitem__ also drops its mask, but it lets the index reach its last (sub-action) dim, so its mask cannot be indexed the same way and is left for a separate change.

Tests:

• New test_indexing_preserves_mask (8 indices, among which ellipses, None, slices and a tensor, for OneHot and MultiOneHot): the indexed spec has the mask obtained by indexing each outcome of the original mask as a tensor of the batch shape, and its samples are the only allowed outcomes. On main, all 8 MultiOneHot cases and the 2 ellipsis OneHot cases fail; the 6 other OneHot cases pass.
• python -m pytest test/test_specs.py -n 4: 1836 passed, 9 skipped.
• python -m pytest test/test_specs.py test/envs test/transforms test/modules test/test_exploration.py -n 6 (CPU): 9060 passed, 1022 skipped, 1 failed; the failure, test/envs/test_parallel.py::TestLibThreading::test_auto_num_threads, comes from OMP_NUM_THREADS=2, which this run set (it fails the same way on the branch with that variable and passes without it: TestLibThreading 2 passed). Same command on main: 9045 passed, 1022 skipped.
• ufmt check and flake8 --config=setup.cfg (pre-commit versions) pass on the changed files.

Motivation and Context

An indexed action spec must not enable actions the mask forbids, and indexing with an ellipsis must work on masked specs as it does on unmasked ones.

close #4632

• 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.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 1a53248 with merge base 380dd55 ():

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.

[meta-cla]

Hi @kjgggggg!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!