[Feature Request] Port TDMPC2 as a trainer #4448

Motivation

TD-MPC2 is one of the current state of the art algorithms for continuous control. It would be a great world model algorithm to include in torchRL's trainer framework, because:

• Unlike DreamerV3, it is decoder free
• It demonstrates state of the art performance on continuous control, outperforming models like DreamerV3 for single tasks
• It demonstrates state of the art performance on multi-task learning as well.

The official repository is easy to read and already uses pytorch, including some pieces of torchRL such as replay buffers.

I'm particularly interested to try this out on the Microduck environment for multi-task learning!

Solution

Have a working trainer implementation of TD-MPC2, on a single task and state-based observations to start. Future work can add multi-task learning and pixel observations. The first implementation should just focus on getting the core pieces in place.

Alternatives

Additional context

Checklist

• I have checked that there is no similar issue in the repo (required)