[BugFix] Make MultiCategorical.to_one_hot accept a batched spec and the scalar sample of a shape [1] spec

[Nicholas022400701]

Description

MultiCategorical.to_one_hot (torchrl/data/tensor_specs.py) iterates for i, n in enumerate(self.nvec). For a spec with batch dims (MultiCategorical([3, 2], shape=(4, 2)), what env.action_spec looks like for a batched env or after spec.expand(...)) nvec has shape (4, 2), n is a row and one_hot(val[..., i], n) raises TypeError: one_hot(): argument 'num_classes' (position 2) must be int, not Tensor. rand, is_in, to_one_hot_spec and MultiOneHot.to_categorical all handle the batched spec, so the conversion worked in one direction only. Second path: MultiCategorical([5]) has shape [1] and remove_singleton=True, so rand() is a scalar and to_one_hot(scalar) raises IndexError: too many indices for tensor of dimension 0.

The widths now come from the last dim of nvec (nvec.flatten(0, -2)[0], the reduction _split_self uses) and a scalar val is unsqueezed first, as is_in and _project do.

Rows with different widths (MultiCategorical([[2, 4], [3, 2]]), #2486) follow to_one_hot_spec, which stacks such rows into a StackedMultiOneHot whose samples are nested tensors: to_one_hot encodes each row with its own row spec and returns torch.nested.nested_tensor of the rows, so spec.to_one_hot_spec().is_in(spec.to_one_hot(val)) holds there too. Rows that are homogeneous below the first dim encode to dense entries ([[[3, 2], [3, 2]], [[4, 5], [4, 5]]] gives entries of shape (2, 5) and (2, 9)). A value with extra leading dims, and rows that differ beyond the first dim, raise the RuntimeError of the stacked spec's rand ("Cannot create a nested tensor with a stack dimension other than 0"), since no nested tensor along dim 0 exists for them. Review round 1 asked for this instead of the ValueError the first push raised.

Tests: test_multi_discrete_conversion_batched_spec (nvec [5] and [5, 2, 3], spec shapes [3] and [4, 5], both conversion directions) and test_multi_discrete_conversion_singleton fail on main (4 TypeError, 1 IndexError); test_multi_discrete_conversion_heterogeneous (rows [2, 4] and [3, 2], with and without a batch dim of 3) checks the nested result against the stacked spec's is_in, each entry against its row spec's encoding, and the two error cases. pytest test/test_specs.py: 1805 passed, 9 skipped. black 22.3.0, usort 1.0.3 and flake8 clean. The heterogeneous test filters torch's one-time "nested tensors are in prototype stage" warning inside the test, the module turns warnings into errors and a function-level filterwarnings mark does not override that.

#4517 also addresses the heterogeneous rows of #2486.

Motivation and Context

Fixes #4493. Every other method of the spec already supports the batched spec; to_one_hot was the one that did not.

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

AI disclosure: I used an AI coding agent to help write this patch, the tests and this description. I have read the change and the tests myself and I will answer review comments personally.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 8 New Failures

As of commit ae18f27 with merge base b8b6f0c ():

NEW FAILURES - The following jobs have failed:

• Examples Tests on Linux / tests (3.12, 13.0, 1) / linux-job (gh)
  AttributeError: module 'ray.internal' has no attribute 'free'
• Examples Tests on Linux / tests (3.12, 13.0, 2) / linux-job (gh)
  AttributeError: module 'ray.internal' has no attribute 'free'
• Unit-tests on Linux / tests-cpu (3.10, collectors) / linux-job (gh)
  test/test_collectors.py::TestCollectorStats::test_ray_collector_stats
• Unit-tests on Linux / tests-cpu (3.13, collectors) / linux-job (gh)
  test/test_collectors.py::TestCollectorStats::test_ray_collector_stats
• Unit-tests on Linux / tests-gpu (3.12, 13.0, 2) / linux-job (gh)
  test/test_inference_server.py::TestProcessSlotTransport::test_server_batched_pass_on_cuda[True]
• Unit-tests on Linux / tests-gpu-distributed (3.12, 13.0) / linux-job (gh)
  test/test_distributed.py::TestRayCollector::test_ray_collector_policy_constructor
• Unit-tests on Linux / tests-stable-gpu (3.12, 13.0, 2) / linux-job (gh)
  test/test_inference_server.py::TestProcessSlotTransport::test_server_batched_pass_on_cuda[True]
• Unit-tests on Linux / tests-stable-gpu-distributed (3.12, 13.0) / linux-job (gh)
  test/test_distributed.py::TestRayCollector::test_ray_collector_policy_constructor

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

[bsprenger]

Thanks for tackling this! Let's make sure we add support for heterogeneous specs though, I am not currently seeing a reason to disallow them

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[bsprenger]

LGTM, I added one additional test case.

In the future, let's also avoid putting so many inline comments. They bloat the diff, and generally we should strive for the code to be self-explanatory. A comment is more useful if it explains why we do something, rather than what the code is doing.