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
Describe the bug
ActionScalingandFlattenActionwrapin_keys_invwith:A NestedKey
("agent", "action")is a tuple, so it is treated as two keys"agent"and"action".ActionScalingthen raisesValueError(single action key per instance).FlattenActionsilently registers"agent"and"action"as two keys.Existing nested-key tests pass a list (
in_keys_inv=[("agent", "action")]) and miss this.LastActionin the same file usesisinstance(value, (str, tuple))and treats a NestedKey tuple as one key.Expected behavior
A NestedKey tuple is one key.
ActionScaling(in_keys_inv=("agent", "action"))andFlattenAction(..., in_keys_inv=("agent", "action"))should store[("agent", "action")].Checklist