[BugFix] RewardSum transform for multiple reward keys

[matteobettini]

This PR extends the RewardSum to work with multiple reward keys.

[vmoens]

How is this supposed to work with many done signals?
I think we should wait for #1539 to land first

[matteobettini]

This transform is supposed to work according to the MARL grouping api.

It will look for the _reset entry in the root and consider it the default reset
For each reward key then, if it finds a _reset in its tensordict it will use that instead.

so each reward key will be by default associated with the _reset in its td and, if that is not present, it will be associated with the _reset in root (if present)

the behavior aligns with the MARL api

[matteobettini]

it should not be dependent from #1539 if we follow this logic

[vmoens]

This transform is supposed to work according to the MARL grouping api.

It will look for the _reset entry in the root and consider it the default reset

I thought our default was that there wasn't a "_reset" at the root (only if there is a "done" at the root)?

[vmoens]

To be clear, the direction I understood we were moving towards:

TensorDict({"_reset": reset, "nested": {"done": done}), []) # no allowed! 
TensorDict({"nested": {"_reset": reset, "done": done}), []) # allowed
TensorDict({"_reset": reset, "done": done}, []) # allowed

The first is not allowed because:

• the "_reset" is not part of the env.reset_keys which MUST be an iterable that follows the grouped done
• it is unclear what this reset refers to and what to do with it in specific cases: does it prevail over sub-resets? do we allow any arbitrary shape? There are way too many decisions to be taken IMO

[matteobettini]

i have updated with what discussed, now i ll just have to test it

[vmoens]

If @matteobettini you're happy with 8da298f I'm good with merging this

[matteobettini]

LGTM some comments

[matteobettini]

Here i preferred having done keys rather than reset keys, this is because users are familiar with what a done key is and could not know about reset keys. Plus there is a 1:1 matching between the 2

[vmoens]

not the way it was done: if I pass env.done_keys there are some duplicates (eg, truncation / termination / done).
Having a "done_keys" list to me is more dangerous because of this, and eventually the only thing we're pointing is the tree structure where the reset_keys should be found. I personally prefer to pass reset_keys: it's what is needed, ie there is a lower risk that refactoring the done_keys mechanism in the future will break this transform. Per se asking users to pass a list X when we interpolate a list Y that is present within the env already as env.Y seems a convoluted way of doing things.

[vmoens]

Another thought about this: per se most users won't need to pass reset keys. We just support it if someone really wants to do nasty things like summing part of the rewards but not all etc. Advanced usage requires advanced understanding so it's fine to ask for reset_keys even this isn't something that is always user-facing.

[matteobettini]

LGTM