We open-source a total of 324 TD-MPC2 model checkpoints, including 12 multi-task models (ranging from 1M to 317M parameters) trained on 80, 70, and 30 tasks, respectively. We are excited to see what the community will do with these models, and hope that our release will encourage other research labs to open-source their checkpoints as well. This section aims to provide further details about the released models.
As one of the first major releases of model checkpoints for reinforcement learning, use of our TD-MPC2 checkpoints is fairly open-ended. We envision that our checkpoints will be useful for researchers interested in training, finetuning, evaluating, and analyzing models on any of the 104 continuous control tasks that we release models for. However, we also expect the community to discover new use cases for these checkpoints.
Direct Use
Model checkpoints can be loaded using the official implementation, and then be used to reproduce our results and/or generate trajectories for any of the supported tasks.
Out-of-Scope Use
We do not expect our model checkpoints to generalize to new (unseen) tasks as is. Such model usage will most likely require some amount of fine-tuning on target task data.
We describe the training procedure for single-task and multi-task model checkpoints in the following.
Training Procedure (Single-task)
Single-task checkpoints are trained using the official implementation with default hyperparameters. All models have 5M parameters. Most, but not all, models are trained until convergence. Refer to the individual task curves in our paper for a detailed breakdown of model performance on each task.
Training Procedure (Multi-task)
Multi-task checkpoints are trained using the official implementation with batch_size=1024 and otherwise default hyperparameters. We release checkpoints trained on the 80-task and 30-task datasets provided here, as well as a 70-task dataset that is obtained by filtering the 80-task dataset based on task IDs. We release model checkpoints ranging from 1M to 317M parameters.
If you find our work useful, please consider citing the paper as follows:
BibTeX:
@misc{hansen2023tdmpc2,
title={TD-MPC2: Scalable, Robust World Models for Continuous Control},
author={Nicklas Hansen and Hao Su and Xiaolong Wang},
year={2023},
eprint={2310.16828},
archivePrefix={arXiv},
primaryClass={cs.LG}
}
