[Feature] Add TD-MPC2 loss module

[bsprenger]

Description

Note

This PR stacks on #4453 . 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 core loss objective for the TD-MPC2 algorithm, porting over the logic into the trainer framework.

Motivation and Context

5th PR in the chain to address #4448 . The design of the loss is 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/4454

• 📄 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! (2 Unrelated Failures)

As of commit 5842fcf with merge base 8a53f8c ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• Unit-tests on Linux / tests-gpu (3.12, 13.0, 2) / linux-job (gh) (similar failure)
  test/test_inference_server.py::TestProcessSlotTransport::test_server_batched_pass_on_cuda[True]

BROKEN TRUNK - The following job failed but was present on the merge base:

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

• Unit-tests on Linux / tests-stable-gpu (3.12, 13.0, 2) / linux-job (gh) (trunk failure)
  test/test_inference_server.py::TestProcessSlotTransport::test_server_batched_pass_on_cuda[True]

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

[theap06]

Overall looks pretty good. Just one nitpick: could you add a test case in tdmpc2 that checks via public api as the current tests cover shapes and the private methods?

[bsprenger]

Updated to test the public API for model_loss and actor_loss_from_latents. Thanks for the feedback!