[Feature] Add TD-MPC2 Q-network ensemble

[bsprenger]

Description

Note

This PR stacks on #4452 . 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 TD-MPC2 Q ensemble, which is a key part of the algorithm. The Q-ensemble consists of $N$ parallel Q-functions to compute TD targets during training of the model and to be used during the planning stage of the algorithm.

Motivation and Context

This is the 4th PR in the chain intended to address #4448 . The Q-ensemble 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

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

• 📄 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 d6e7f5e with merge base 71a4143 ():

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

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

• Unit-tests on Linux / tests-gpu (3.12, 13.0, 2) / linux-job (gh) (trunk failure)
  test/test_inference_server.py::TestProcessSlotTransport::test_server_batched_pass_on_cuda[True]
• 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.

[vmoens]

Advertised num_bins modes are rejected.

The Q-ensemble documentation says 0 provides identity regression and 1 provides symlog regression, but the constructor rejects both. The config factories and loss also only support categorical mode.

Either implement the documented/reference behavior or explicitly narrow the API and documentation to num_bins > 1.

[bsprenger]

I updated the documented behaviour. Thanks!

[bsprenger]

@theap06 I've rebased this one after merging its base branch, looking for a re-review / CI, thanks!

[theap06]

LGTM!