Skip to content

[BUG] ActionScaling and FlattenAction treat NestedKey tuples as sequences of keys #4464

Description

@YeonwooSung

Describe the bug

ActionScaling and FlattenAction wrap in_keys_inv with:

if not isinstance(in_keys_inv, (list, tuple)):
    in_keys_inv = [in_keys_inv]

A NestedKey ("agent", "action") is a tuple, so it is treated as two keys "agent" and "action".

  • ActionScaling then raises ValueError (single action key per instance).
  • FlattenAction silently registers "agent" and "action" as two keys.

Existing nested-key tests pass a list (in_keys_inv=[("agent", "action")]) and miss this. LastAction in the same file uses isinstance(value, (str, tuple)) and treats a NestedKey tuple as one key.

from torchrl.envs.transforms import ActionScaling, FlattenAction

ActionScaling(in_keys_inv=("agent", "action"))
# ValueError: ActionScaling only supports a single action key per instance.

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

Expected behavior

A NestedKey tuple is one key. ActionScaling(in_keys_inv=("agent", "action")) and FlattenAction(..., in_keys_inv=("agent", "action")) should store [("agent", "action")].

Checklist

  • I have checked that there is no similar issue in the repo (required)
  • I have read the documentation (required)
  • I have provided a minimal working example to reproduce the bug (required)

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions