[BugFix] Treat FlattenAction NestedKey tuples as one key

[YeonwooSung]

Description

ActionScaling and FlattenAction wrapped in_keys_inv with isinstance(..., (list, tuple)). A NestedKey ("agent", "action") is a tuple, so it was treated as two keys "agent" and "action".

ActionScaling then raised (single action key per instance). FlattenAction defaulted in_keys / out_keys / out_keys_inv from list(in_keys_inv) and registered the two strings.

Both constructors now wrap with isinstance(..., (str, tuple)), matching LastAction and Transform. Existing tests that pass a list of NestedKeys still work. in_keys, out_keys, and out_keys_inv were already handled by the Transform setters.

Code example

from torchrl.envs.transforms import ActionScaling, FlattenAction

t = ActionScaling(in_keys_inv=("agent", "action"))
assert t.in_keys_inv == [("agent", "action")]

t = FlattenAction(first_dim=-2, last_dim=-1, in_keys_inv=("agent", "action"))
assert t.in_keys_inv == [("agent", "action")]
assert t.in_keys == [("agent", "action")]

Motivation and Context

close #4464

• 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/4476

• 📄 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 58945fb with merge base de28012 ():

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.

[vmoens]

What about insinstance(key, NestedKey)?

[YeonwooSung]

isinstance(key, NestedKey) is a runtime check: a string or nested tuple is a NestedKey, a list of keys is not. ActionScaling and FlattenAction now wrap in_keys_inv with that.