[BugFix] RewardSum key check

[matteobettini]

In the RewardSum transform, reset keys where checked to be valid only when not directly provided by the user.

rl/torchrl/envs/transforms/transforms.py

Line 4773 in 07fcfb1

if len(reset_keys) != len(self.in_keys) or not _check_match(

This led to silent errors when the users provides reset keys manually but there are not as many as the in_keys gathered automatically from the reward_spec.

This PR runs the key check in both the case of user provided or autoimatically gathered reset keys, making the function more resilient.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (8 Unrelated Failures)

As of commit 290ba1e with merge base 07fcfb1 ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• Continuous Benchmark (PR) / CPU Pytest benchmark (gh)
• Continuous Benchmark (PR) / GPU Pytest benchmark (gh)
• Unit-tests on Linux GPU / tests (3.8, 12.1) / linux-job (gh)
• Unit-tests on MacOS CPU / tests (3.11) / macos-job (gh)
• Unit-tests on MacOS CPU / tests (3.8) / macos-job (gh)

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

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

• Habitat Tests on Linux / tests (3.9, 11.6) / linux-job (gh)
• Unit-tests on Linux GPU, latest stable release / tests (3.8, 11.8) / linux-job (gh)
• Unit-tests on Windows GPU / unittests / windows-job (gh)

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

[vmoens]

Is this needed?

If we look at the message in the error, we specifically tell the user that we cannot retrieve them automatically and hence the user should pass them.
If we pass them manually, the check should be different, e.g., we could just test that there are as many in_keys as reset_keys.

[matteobettini]

I think having both checks is useful and at worst it cannot hurt.
The second part of the check will help preventing the user to pass "("smth","_reset") when the in_key is "reward"

The checks are run just once

[vmoens]

The current check makes sure that the root of in_keys has a similar root for resets. If resets are passed manually, they could be located anywhere IMO, not at the sample place as in_keys. I guess we want manual resets to be as flexible as possible. If we restrict them as we do for automatically gathered reset entries, what do we gain with the manual version?

[matteobettini]

I thought the manual version was just for when the auto does not work, but i see the point.
I ll update to just check the length of the key list in both cases as what to i want to avoid is the zip silently failing due to diff lenghts (which just happened to me)

[vmoens]

At this point we need a better set of tests for these error messages! It isn't easy to make sure we get them right just by looking at the code

[matteobettini]

I have added a test that would fail on main that checks all the errors related to the length of keys in rewardsum

[vmoens]

What do you mean "that would fail on main"? In the sense that they test the new feature or that the feature is bc breaking?

[vmoens]

LGTM thanks

[matteobettini]

What do you mean "that would fail on main"? In the sense that they test the new feature or that the feature is bc breaking?

They test that where before there was a silent error, now there is an explicit exeption.