TD-MPC2

官方实现

TD-MPC2：用于连续控制的可扩展、鲁棒世界模型，由

Nicklas Hansen、Hao Su*、Xiaolong Wang*（加州大学圣地亚哥分校）共同开发

[官网] [论文] [模型] [数据集]

---

公告（2025年4月）：支持分幕任务！

我们在最新版本中新增了对分幕强化学习（具有终止状态的任务）的支持。可通过设置 episodic=true 来启用此功能，但默认仍为禁用状态，以确保不同版本之间结果的可复现性。

---

概述

TD-MPC2 是一种可扩展、鲁棒的基于模型的强化学习算法。它在跨越多个领域的 104 个连续控制任务中，仅使用 一套 超参数，便能与现有的无模型和基于模型的方法相媲美（右图）。此外，我们还展示了 TD-MPC2 的可扩展性：通过训练一个拥有 3.17 亿参数的智能体，在多个领域、不同机器人本体和动作空间上完成 80 个任务（左图）。

本仓库包含用于训练和评估单任务在线强化学习以及多任务离线强化学习 TD-MPC2 智能体的代码。我们还开源了 300+ 个 模型检查点（包括 12 个多任务模型），涵盖 4 个任务领域：DMControl、Meta-World、ManiSkill2 和 MyoSuite，以及用于训练多任务模型的我们的 30 任务和 80 任务数据集。我们的代码库同时支持状态观测和像素观测。我们希望本仓库能成为未来基于模型强化学习研究的有用社区资源。

---

快速入门

运行 TD-MPC2 的单任务在线强化学习需要一台配备 GPU 且至少有 12 GB 内存的机器；而使用我们提供的 80 任务数据集进行多任务离线强化学习则需 128 GB 内存。对于单任务在线强化学习以及评估提供的多任务模型（最高 3.17 亿参数），建议使用显存至少为 8 GB 的 GPU。训练 3.17 亿参数的模型则需要显存至少 24 GB 的 GPU。

我们提供了一个 Dockerfile 以便于安装。您可以通过以下命令构建 Docker 镜像：

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

该 Docker 镜像包含了运行 DMControl 所需的所有依赖项。我们还在 这里 提供了一个预构建的 Docker 镜像。

如果您更倾向于使用 conda 而不是 Docker，请先运行以下命令：

conda env create -f docker/environment.yaml

docker/environment.yaml 文件会安装训练 DMControl 任务所需的依赖。其他领域的依赖项可按照 docker/environment.yaml 中的说明进行安装。

若要运行 ManiSkill2，您还需要下载并链接必要的资源，方法是运行：

python -m mani_skill2.utils.download_asset all

这会将资源下载到 ./data 目录下。您可以将这些资源移动到任何位置。然后，在您的 ~/.bashrc 文件中添加以下行：

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

并重启终端。请注意，Meta-World 需要 MuJoCo 2.1.0 和 gym==0.21.0，而这些软件包的安装正变得越来越困难。我们在此托管不受限制的 MuJoCo 2.1.0 许可证（由 Google DeepMind 提供），地址为 https://www.tdmpc2.com/files/mjkey.txt。您可以通过运行以下命令下载许可证：

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

根据您现有的系统软件包情况，可能还需要安装其他依赖项。请参阅 docker/Dockerfile 以获取推荐的系统软件包列表。

---

支持的任务

本代码库支持我们论文中使用的来自 DMControl、Meta-World、ManiSkill2 和 MyoSuite 的全部 104 个连续控制任务。具体来说，它支持 DMControl 中的 39 个任务（包括 11 个自定义任务）、Meta-World 中的 50 个任务、ManiSkill2 中的 5 个任务以及 MyoSuite 中的 10 个任务，并覆盖了论文中涉及的所有任务。以下是各任务域的预期命名格式：

域	任务
dmcontrol	dog-run
dmcontrol	cheetah-run-backwards
metaworld	mw-assembly
metaworld	mw-pick-place-wall
maniskill	pick-cube
maniskill	pick-ycb
myosuite	myo-key-turn
myosuite	myo-key-turn-hard

这些任务可以通过在 evaluation.py 中指定 task 参数来运行。多任务训练和评估则分别通过设置 task=mt80 或 task=mt30 来指定 80 任务和 30 任务的数据集。虽然在训练或评估我们的多任务模型时，通常无需直接访问底层的任务 ID 或嵌入，但我们工作中使用的任务名称到任务嵌入的映射可以在 这里找到。截至 2025 年 4 月，我们的代码库也提供了对其他 MuJoCo/Box2d Gymnasium 任务的基本支持；有关任务列表，请参阅 envs 目录。按照 envs 中的示例，添加对自定义任务的支持应该相对简单。

注意： 我们还支持 DMControl 任务中的图像观测。如果您希望训练视觉策略，可以使用参数 obs=rgb。

使用示例

我们将在下面提供如何评估我们提供的 TD-MPC2 检查点，以及如何训练您自己的 TD-MPC2 智能体的示例。

评估

以下是一些关于如何评估已下载的单任务和多任务检查点的示例。

$ python evaluate.py task=mt80 model_size=48 checkpoint=/path/to/mt80-48M.pt
$ python evaluate.py task=mt30 model_size=317 checkpoint=/path/to/mt30-317M.pt
$ python evaluate.py task=dog-run checkpoint=/path/to/dog-1.pt save_video=true

所有单任务检查点都期望 model_size=5。多任务检查点则有多种模型尺寸可供选择。可用的参数为 model_size={1, 5, 19, 48, 317}。请注意，目前尚不支持对多任务检查点进行单任务评估。完整的参数列表请参阅 config.yaml。

训练

请参阅下方示例，了解如何在单任务（在线强化学习）和多任务数据集（离线强化学习）上训练 TD-MPC2。我们建议在 config.yaml 中配置 Weights and Biases (wandb) 以跟踪训练进度。

$ python train.py task=mt80 model_size=48 batch_size=1024
$ python train.py task=mt30 model_size=317 batch_size=1024
$ python train.py task=dog-run steps=7000000
$ python train.py task=walker-walk obs=rgb

对于单任务在线强化学习，我们推荐使用默认超参数，包括默认的 500 万参数模型规模（model_size=5）。多任务离线强化学习则受益于更大的模型规模，但更大的模型在训练和评估时的成本也会相应增加。可用的模型规模参数为 model_size={1, 5, 19, 48, 317}。完整参数列表请参阅 config.yaml。

---

引用

如果您认为我们的工作有所帮助，请考虑按以下方式引用我们的论文：

@inproceedings{hansen2024tdmpc2,
  title={TD-MPC2: 可扩展、鲁棒的连续控制世界模型}, 
  author={Nicklas Hansen 和 Hao Su 和 Xiaolong Wang},
  booktitle={国际表征学习大会 (ICLR)},
  year={2024}
}

以及原始的 TD-MPC 论文：

@inproceedings{hansen2022tdmpc,
  title={用于模型预测控制的时间差分学习},
  author={Nicklas Hansen 和 Xiaolong Wang 和 Hao Su},
  booktitle={国际机器学习大会 (ICML)},
  year={2022}
}

---

贡献

我们非常欢迎您为本项目做出贡献。如果您有任何建议或问题报告，欢迎随时提交议题或拉取请求，但在操作之前请先阅读我们的贡献指南。我们的目标是构建一个能够轻松扩展到新环境和新任务的代码库，也非常期待听到您的使用体验！

---

许可证

本项目采用 MIT 许可证授权——详情请参阅 LICENSE 文件。请注意，该仓库依赖于第三方代码，这些代码受其各自许可证的约束。

TD-MPC2 快速上手指南

TD-MPC2 是一种可扩展且鲁棒的基于模型的强化学习算法，适用于连续控制任务。本指南将帮助你快速配置环境并运行模型。

环境准备

系统要求

• GPU:
  • 单任务在线训练/评估：至少 8GB 显存。
  • 多任务离线训练（317M 参数模型）：至少 24GB 显存。
• 内存 (RAM):
  • 单任务在线训练：至少 12 GB。
  • 多任务离线训练（80 任务数据集）：至少 128 GB。
• 操作系统: Linux (推荐，Docker 支持最佳)。

前置依赖

• Docker (推荐) 或 Conda
• NVIDIA Driver 与 CUDA Toolkit (需与 PyTorch 版本兼容)
• Git

安装步骤

你可以选择使用 Docker（推荐，环境隔离好）或 Conda 进行安装。

方案一：使用 Docker (推荐)

1. 克隆仓库并进入 docker 目录：

   git clone https://github.com/nicklashansen/tdmpc2.git
   cd tdmpc2/docker
   

2. 构建镜像：

   docker build . -t <your_username>/tdmpc2:1.0.1
   

   注：你也可以直接拉取预构建镜像：docker pull nicklashansen/tdmpc2:1.0.1

3. 运行容器（需映射当前目录以便访问代码和数据）：

   docker run --gpus all -it --rm -v $(pwd)/..:/workspace <your_username>/tdmpc2:1.0.1 bash
   

方案二：使用 Conda

1. 创建虚拟环境：

   conda env create -f docker/environment.yaml
   conda activate tdmpc2
   

2. 额外配置 (针对 ManiSkill2 任务): 如果需要使用 ManiSkill2，需下载资源并配置环境变量：

   python -m mani_skill2.utils.download_asset all
   

   下载完成后，将资源路径添加到 ~/.bashrc：

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

   然后重启终端或执行 source ~/.bashrc。

3. 额外配置 (针对 Meta-World 任务): Meta-World 需要 MuJoCo 2.1.0。下载 License 文件：

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

   注意：确保已安装 gym==0.21.0 和 mujoco-py==2.1.0。

基本使用

以下示例展示如何评估预训练模型和训练新代理。请确保在激活的环境或 Docker 容器中运行。

1. 评估预训练模型 (Evaluation)

评估多任务模型 (80 任务, 48M 参数)：

python evaluate.py task=mt80 model_size=48 checkpoint=/path/to/mt80-48M.pt

评估单任务模型 (例如 dog-run) 并保存视频：

python evaluate.py task=dog-run checkpoint=/path/to/dog-1.pt save_video=true

注：单任务检查点默认 model_size=5；多任务检查点可选 {1, 5, 19, 48, 317}。

2. 训练新模型 (Training)

训练单任务在线 RL (默认 5M 参数模型)：

python train.py task=dog-run steps=7000000

训练视觉策略 (使用 RGB 图像观测)：

python train.py task=walker-walk obs=rgb

训练多任务离线 RL (80 任务数据集)：

python train.py task=mt80 model_size=48 batch_size=1024

提示：建议在 config.yaml 中配置 Weights and Biases (wandb) 以跟踪训练进度。