TD-MPC2: a model-based reinforcement learning algorithm you can actually run

What the paper claims, and what the code contains

The project is the official implementation of a paper, and the readme states the claims plainly.

TD-MPC2 is a scalable, robust model-based reinforcement learning algorithm, and it compares favourably to existing model-free and model-based methods across 104 continuous control tasks spanning multiple domains, with a single set of hyperparameters. The scalability claim is demonstrated by training a single 317 million parameter agent to perform 80 tasks across multiple domains, embodiments and action spaces.

The word single is doing the work in both claims. One set of hyperparameters for 104 tasks is a much stronger statement than a per-task tuned result, because the usual failure of a reinforcement learning paper is that each task needs its own settings. And one agent doing 80 tasks is a different claim from 80 agents, because the parameters have to be shared across domains with different embodiments and different action spaces.

The repository contains code for training and evaluating both single-task online reinforcement learning and multi-task offline reinforcement learning agents, so the two halves of the claim are both runnable.

The paper itself is on arXiv, the project has its own website with separate pages for the models and the datasets, and the authors are from UC San Diego, with the two later authors marked as equal contributors.

Four domains, a hundred and four tasks

The supported task list is specific enough to be useful, and it is worth knowing the breakdown.

The codebase supports all 104 continuous control tasks from DMControl, Meta-World, ManiSkill2 and MyoSuite used in the paper. Specifically, 39 tasks from DMControl including 11 custom ones, 50 from Meta-World, 5 from ManiSkill2 and 10 from MyoSuite.

The 11 custom DMControl tasks are a detail that shows the work behind the claim, since the standard suite is not the whole set. The domains are also a deliberate spread: DMControl is physics simulation with images, Meta-World is robot manipulation from a single policy, ManiSkill2 is a newer simulation stack, and MyoSuite is biomechanical muscle control, which is the most physically awkward of the four.

The readme gives the task naming convention with a table, because the names are what you pass on the command line: the domain prefix and the task, such as a run task and a backwards-running cheetah in DMControl, an assembly and a pick-place-wall task in Meta-World, a cube and a YCB object task in ManiSkill, and a key-turn task in MyoSuite with a hard variant.

Multi-task training and evaluation is selected by setting the task argument to mt80 or mt30 for the two dataset sizes. There is also basic support for other MuJoCo and Box2d Gymnasium tasks as of April 2025, listed in the envs directory, and the readme says that adding a custom task should be straightforward by following the examples in there. Image observations are supported in the DMControl tasks, for visual policies, with an argument to enable them.

Checkpoints and datasets, the parts you cannot recreate

The most valuable thing in the repository is not the code.

More than 300 model checkpoints are open-sourced, including 12 multi-task models, across the four task domains. There are also the 30-task and 80-task datasets used to train the multi-task models, published on the project website.

Those checkpoints come in several sizes. The available model sizes are 1, 5, 19, 48 and 317 million parameters, and all single-task checkpoints expect the size of 5. The evaluation examples in the readme show the shape of the command: the task, the model size and the checkpoint path, with a multi-task evaluation, a single-task one, and a third that saves a video.

Two limitations are stated and both matter in practice. Single-task evaluation of a multi-task checkpoint is currently not supported, so the 317 million parameter model cannot be evaluated on one task at a time, and the full list of arguments lives in the configuration file rather than on the command line.

The datasets matter for a specific reason. The multi-task offline route is the part of the paper that is hardest to reproduce, because it needs a training set that covers 80 tasks with a shared representation, and the fact that the project ships it turns a research reproduction into a download.

The codebase supports both state and pixel observations, with image observations supported in the DMControl tasks.

The hardware requirements are the first thing to read

The getting started section states the requirements before anything else, and they are the first thing to check.

You need a machine with a GPU and at least 12 gigabytes of RAM for single-task online reinforcement learning. For multi-task offline reinforcement learning on the provided 80-task dataset you need 128 gigabytes of RAM, which is a lot and is the number that will decide whether the multi-task half of the paper is reachable for you.

A GPU with at least 8 gigabytes of memory is recommended for single-task online reinforcement learning and for evaluation of the provided multi-task models, which go up to 317 million parameters. Training that 317 million parameter model needs a GPU with at least 24 gigabytes of memory.

So there are three distinct hardware profiles, and it is worth picking the one you actually need. Evaluation of a released checkpoint is the cheapest by a wide margin, at 8 gigabytes of GPU memory and no mention of a large memory requirement. Single-task online training is the middle case at 12 gigabytes of RAM. The multi-task offline training is the expensive one.

The readme also says the project provides a Dockerfile for easy installation, and a pre-built image is published, so the environment is the least of your problems if the hardware is right.

Setting up with Docker or conda, and the MuJoCo licence

There are two installation routes and the Docker one is the short one.

cd docker && docker build . -t <user>/tdmpc2:1.0.1

That image contains all the dependencies needed for running DMControl, and a pre-built image is published for people who would rather not build it.

The conda route creates the environment from a file in the same directory:

conda env create -f docker/environment.yaml

and the readme explains that this installs the dependencies for training on DMControl, with the other domains installed by following the instructions in the same file.

ManiSkill2 needs an extra step, because it downloads its assets separately:

python -m mani_skill2.utils.download_asset all

and then the asset directory has to be pointed at, with an environment variable added to your shell profile:

export MS2_ASSET_DIR=<path>/<to>/<data>

The dependency that the readme warns about is Meta-World, which requires MuJoCo 2.1.0 and a specific version of the gym library, and which the author says is becoming increasingly difficult to install. The project hosts an unrestricted MuJoCo 2.1.0 licence, courtesy of Google DeepMind, and you can fetch it with a single command:

wget https://www.tdmpc2.com/files/mjkey.txt -O ~/.mujoco/mjkey.txt

That is a good catch for an evaluator. A repository whose environment cannot be recreated is a repository whose results cannot be checked, and hosting the licence is the difference.

Episodic tasks, and why the new feature is off by default

There is an announcement at the top of the readme about episodic reinforcement learning, and the detail in it is the interesting part.

Support for episodic tasks, meaning tasks with terminations, was added in the latest release, and it can be enabled with the episodic flag. It remains disabled by default, and the stated reason is to ensure reproducibility of results across releases.

That is a good decision and a rare one. A change that alters how episodes terminate changes the learning problem, so a default that silently changed the semantics of every existing configuration would invalidate previously reported numbers. Keeping it off by default means the older behaviour is still what you get, and the new behaviour is something you opt into with a flag and a fair expectation that the results will differ.

For somebody reproducing the paper's numbers, that matters directly, because the task set includes 50 Meta-World tasks and 10 MyoSuite tasks, and whether terminations are handled changes what those tasks are.

The rest of the configuration surface is a file, and the readme's advice is to consult it for the full list of arguments rather than trying to memorise the flags from the examples. The evaluation and training commands are consistent about their shape, with the task, the model size and the checkpoint path, and the multi-task and single-task forms are distinguished by which of the task and model size arguments you set.

The repository layout and where the project stands

The top level of the tree shows the shape of a research release.

There is the package directory with the implementation, an envs directory holding the environment and task definitions, including the list of the additional Gymnasium tasks, a docker directory with the Dockerfile and the conda environment file, a datasets directory, a results directory, an assets directory for the images in the readme, and a directory named nautilus. There is a contributing guide, the licence file, the readme and the assets.

Inside the package, the readme points at a common module for the task to task embedding mapping used in the multi-task work, and the configuration file holds the full argument list. There is a separate evaluation entry point, and the training examples follow the same command shape as the evaluation ones.

The licence is MIT, the default branch is main, and the project has a website that hosts the models, the datasets and the MuJoCo licence file. There are no GitHub releases to point at, and the last push was on 2026-07-13 with the repository not archived, so the code is being worked on even though the versioned artefacts live on the website.

The readme closes by saying the hope is that the repository will serve as a useful community resource for future research on model-based reinforcement learning, and for a project that ships 300 checkpoints and two datasets, that is the more likely outcome: the code will be read, but the artefacts will be the thing other people actually build on.

Editorial conclusion

TD-MPC2 suits a reinforcement learning researcher who wants a modern model-based method with the checkpoints and datasets already produced, because a single set of hyperparameters across 104 tasks and a multi-task agent at 80 tasks are exactly the artefacts that are expensive to reproduce. It is a poor fit for a small GPU or a small memory budget, since the multi-task offline route asks for 128 gigabytes of memory, and a poor fit if you want to run a different algorithm family, since the repository is an implementation of one paper rather than a framework. Before you start, read the hardware requirements, and if you only want to evaluate a released model, remember that the single-task evaluation path expects a specific model size and that evaluating a multi-task checkpoint on a single task is not supported.

Frequently asked questions

Official sources

1. Issues
2. License: MIT
3. nicklashansen/tdmpc2 on GitHub
4. Project website
5. README