[Feature] Add TD-MPC2 optimization stepper

[bsprenger]

Description

Note

This PR stacks on #4455 . 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 optimization orchestrator (the optimization stepper) for the TD-MPC2 algorithm.

The optimization stepper is responsible for the key loss evaluation and parameter updates in the trainer framework. The optimization stepper for TD-MPC2 in particular has several key components for the world model and the actor prior, with separate optimizers.

Motivation and Context

7th PR in the chain towards closing #4448. The optimization procedure 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

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

• 📄 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 e189412 with merge base 65162f2 ():

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

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

• Examples Tests on Linux / tests (3.12, 13.0, 1) / linux-job (gh) (trunk failure)
  AttributeError: module 'ray.internal' has no attribute 'free'
• Examples Tests on Linux / tests (3.12, 13.0, 2) / linux-job (gh) (trunk failure)
  AttributeError: module 'ray.internal' has no attribute 'free'
• Unit-tests on Linux / tests-cpu (3.10, collectors) / linux-job (gh) (trunk failure)
  test/test_collectors.py::TestCollectorStats::test_ray_collector_stats
• Unit-tests on Linux / tests-cpu (3.13, collectors) / linux-job (gh) (trunk failure)
  test/test_collectors.py::TestCollectorStats::test_ray_collector_stats
• 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-gpu-distributed (3.12, 13.0) / linux-job (gh) (trunk failure)
  test/test_distributed.py::TestRayCollector::test_ray_collector_policy_constructor
• 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]
• Unit-tests on Linux / tests-stable-gpu-distributed (3.12, 13.0) / linux-job (gh) (trunk failure)
  test/test_distributed.py::TestRayCollector::test_ray_collector_policy_constructor

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