[BugFix] action_spec_unbatched whenever necessary

[vmoens]

Stack from ghstack (oldest at bottom):

• [Refactor] Use <spec>_unbatched in VMAS #2593
• -> [BugFix] action_spec_unbatched whenever necessary #2592

cc @matteobettini we now have a more generic single_action_spec that is somewhat similar to unbached_action_spec (name is borrowed from gymnasium)
There is something a bit fishy about unbatched_action_spec though which is that it's the unbatched_full_action_spec, and it doesn't (always) match action spec in its sturcture (it will always be composite).

LMK if these changes make sense (implemented these changes to fix tutorials/sphinx-tutorials/multiagent_competitive_ddpg.py which was broken)

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• [DomainsOnly] Jobs fail with GLIBC version not found

❌ 3 New Failures, 1 Cancelled Job, 15 Unrelated Failures

As of commit 98a3b39 with merge base a47b32c ():

NEW FAILURES - The following jobs have failed:

• Continuous Benchmark (PR) / CPU Pytest benchmark (gh)
  FAILED ../../../../../../tmp/test_objectives_benchmarks.py::test_iql_speed[True-None] - torch._dynamo.exc.Unsupported: Graph break under GenericContextWrappingVariable
• Continuous Benchmark (PR) / GPU Pytest benchmark (gh)
  FAILED ../../../../tmp/test_objectives_benchmarks.py::test_iql_speed[True-None] - torch._dynamo.exc.Unsupported: Graph break under GenericContextWrappingVariable
• SOTA Tests on Linux / tests (3.9, 12.1) / linux-job (gh)
  RuntimeError: Command docker exec -t 45d00801dc209b9d4eeb8b68c6125c8092b0db32e3985f2ef76dbd5d324d273b /exec failed with exit code 1

CANCELLED JOB - The following job was cancelled. Please retry:

• Generate documentation / build-docs (3.10, 12.1) / linux-job (gh)
  ##[error]The operation was canceled.

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• Unit-tests on Linux / tests-cpu-oldget (3.12) / linux-job (gh) (similar failure)
  test/test_trainer.py::TestRecorder::test_recorder

BROKEN TRUNK - The following jobs failed but was present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Build Linux Wheels / pytorch/rl (pytorch/rl, test/smoke_test.py, torchrl) / upload / manywheel-py3_9-rocm6_1 (gh) (trunk failure)
  ##[error]Unable to find an artifact with the name: pytorch_rl__3.9_rocm6.1_x86_64
• Build Linux Wheels / pytorch/rl (pytorch/rl, test/smoke_test.py, torchrl) / upload / manywheel-py3_9-rocm6_2 (gh) (trunk failure)
  ##[error]Unable to find an artifact with the name: pytorch_rl__3.9_rocm6.2_x86_64
• Build Windows Wheels / pytorch/rl / upload / wheel-py3_9-cpu (gh) (trunk failure)
  ##[error]Unable to find any artifacts for the associated workflow
• Build Windows Wheels / pytorch/rl / upload / wheel-py3_9-cuda11_8 (gh) (trunk failure)
  ##[error]Unable to find any artifacts for the associated workflow
• Build Windows Wheels / pytorch/rl / upload / wheel-py3_9-cuda12_4 (gh) (trunk failure)
  ##[error]Unable to find any artifacts for the associated workflow
• Habitat Tests on Linux / tests (3.9, 12.1) / linux-job (gh) (trunk failure)
  AttributeError: _ARRAY_API not found
• Libs Tests on Linux / unittests-jumanji (3.9, 12.1) / linux-job (gh) (trunk failure)
  test/test_libs.py::TestJumanji::test_jumanji_rendering[batch_size1-RubiksCube-v0]
• Unit-tests on Linux / tests-cpu (3.10) / linux-job (gh) (trunk failure)
  test/test_trainer.py::TestRecorder::test_recorder
• Unit-tests on Linux / tests-cpu (3.11) / linux-job (gh) (trunk failure)
  test/test_trainer.py::TestRecorder::test_recorder
• Unit-tests on Linux / tests-cpu (3.12) / linux-job (gh) (trunk failure)
  test/test_trainer.py::TestRecorder::test_recorder
• Unit-tests on Linux / tests-cpu (3.9) / linux-job (gh) (trunk failure)
  test/test_trainer.py::TestRecorder::test_recorder
• Unit-tests on Linux / tests-gpu (3.11, 12.1) / linux-job (gh) (trunk failure)
  test/test_trainer.py::TestRecorder::test_recorder
• Unit-tests on Linux / tests-stable-gpu (3.10, 11.8) / linux-job (gh) (trunk failure)
  test/test_trainer.py::TestRecorder::test_recorder
• Unit-tests on Windows / unittests-cpu / windows-job (gh) (trunk failure)
  ##[error]Process completed with exit code 1.

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

[matteobettini]

Got it.

Love that there is a way to get the unbatched specs easily now. This is super helpful.

I guess the prefix single_ has already been decided but I will give my 2 cents on this anyway as I think there are some important points.

I think single_ is a really confusing name. In gymnasium I guess it can work as the library is not multi-agent/multi-task. but in torchrl it sounds super confusing to me:

• it could be meaning single agent
• or single spec (as opposed to full spec). this is why single_full_spec sounds oximoric to me
• it does not correlate with the concept of a batch_size

To me the prefix unbatched_ (or variations of it) drives the concept home better without any of the confusions above

[matteobettini]

Suggested change

	spec=env.single_action_spec,
	spec=env.single_full_action_spec,

Wherever unbatched_action_spec is used, it should be changed to single_full_action_spec

This is valid for all the places here. As the components that were receiving unbatched_ are expecting a composite

[vmoens]

We can name it unbatched but it will clash with VMAS unbatched (which are somewhat different).
It's a nightly feature IIRC so we can do it

[matteobettini]

We can name it unbatched but it will clash with VMAS unbatched (which are somewhat different). It's a nightly feature IIRC so we can do it

Yeah, the vmas ones were there for the absence of this feature. Now that this exists the vmas ones could be removed.

Your are right that the vmas name implied full so there would be a clash there.

We could think about how to best deal with this but i think torchrl should not be prevented to use the better name cause of this

[matteobettini]

so the difference between full and normal is just present when there is only one agent group.

What we could do is that for a period of time we could warn vmas users that request unbatched_spec in envs with just one group that this is now returning the non-composite version and if they want the old composite one there is unbatched_full_spec

[vmoens]

That's bc breaking, we need to warn them the behaviour will change and they can make the warning disappear by using the full version

[vmoens]

Intermediate solution: single_action_spec becomes action_spec_unbatched
that way there is no conflict, and we can just raise a warning in VMAS to let people know about the new API. Wdyt?

[matteobettini]

Intermediate solution: single_action_spec becomes action_spec_unbatched that way there is no conflict, and we can just raise a warning in VMAS to let people know about the new API. Wdyt?

Genius! Love it

[vmoens]

Ok I also simplified the mock envs that had unbatched specs, LMK if that makes sense!

[matteobettini]

I think that makes sense. My comments from the first review still hold tho. I would just make this a refactoring and no logic change

so unbatched_spec in vmas should be translated to full_spec_unbatched

[vmoens]

I would just make this a refactoring and no logic change

There is not logic change, see my comment above

[vmoens]

@matteobettini I'll wait until the tests pass but it should be good.
I did not make any deprecation in VMAS, not sure how you want to procede there.

[matteobettini]

Maybe we can put a warning to use the new feature?