[Feature] Add TD-MPC2 planner

[bsprenger]

Description

Note

This PR stacks on #4454 . I cannot find a good way to compare the diff against that base, since it is in my own fork.
Once it is merged, I will re-target this one to make the diff easier to read.

This PR introduces the central part of the TD-MPC2 algorithm: the planner module. The planner is responsible for rolling out candidate trajectories at inference time and planning which sequence of actions leads to the best outcome.

Motivation and Context

6th PR in the chain to address #4448 . The planning algorithm is described in depth in the original paper and associated implementation.

• I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

• New feature (non-breaking change which adds core functionality)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

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

• 📄 Preview Python docs built from this PR

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

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

[vmoens]

Let's use test classes whenever we can + short test names

[bsprenger]

Done, I refactored the whole chain of PRs to use test classes and short names. These changes are folded into the commits on the PRs that introduced each feature.

[vmoens]

not sure modules/functionals really belong here.

Also another finding from Codex:

Evaluation-time exploration cannot be disabled.

The planner always adds Gaussian noise to the selected action. The reference TD-MPC2 implementation skips this perturbation in evaluation mode. Add an explicit evaluation mode or respect TorchRL’s interaction type, with a regression test.

[bsprenger]

I moved the modules and helpers -- does the new layout look better to you?

And nice find from Codex. I missed that because in the official implementation here, the evaluate.py script actually does not set this eval mode. So I thought it was deprecated/unused. Looking closer , it is used in a few other places, so I added that here. Thanks!

[theap06]

Overall LGTM! Just one nitpick

[theap06]

This import is made at every step. Let's move this to the hot path.