TD-MPC2

TD-MPC2: 可扩展、强大的持续控制世界模型的官方实现

作者: Nicklas Hansen, Hao Su*, Xiaolong Wang* (加州大学圣地亚哥分校)

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

概述

TD-MPC2是一种可扩展、强大的基于模型的强化学习算法。它在覆盖多个领域的104个连续控制任务中表现优于现有的基于模型和无模型的方法,使用了单一组超参数(右图)。我们还通过训练一个317M参数的代理人来执行跨多个领域、多个构架和多个动作空间的80个任务,展示了TD-MPC2的可扩展性(左图)。

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

入门

您需要一台配备GPU和至少12GB RAM的机器来进行TD-MPC2的单任务在线强化学习,以及配备128GB RAM的机器来进行我们提供的80任务数据集的多任务离线强化学习。对于单任务在线强化学习和评估提供的多任务模型(最多317M参数),推荐使用至少8GB显存的GPU。训练317M参数模型需要至少24GB显存的GPU。

我们提供了一个Dockerfile来进行轻松安装。您可以通过运行以下命令来构建Docker镜像:

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

这个Docker镜像包含了运行DMControl、Meta-World和ManiSkill2实验所需的所有依赖项。

如果您更喜欢手动安装依赖项,可以通过运行以下命令来安装conda依赖项:

conda env create -f docker/environment.yaml
pip install gym==0.21.0

environment.yaml文件安装了在DMControl任务上训练所需的依赖项。其他领域可以通过按照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。我们提供了谷歌DeepMind的无限制MuJoCo 2.1.0许可证，地址为https://www.tdmpc2.com/files/mjkey.txt。您可以通过运行以下命令下载该许可证:

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

如果您还没有安装MuJoCo 2.1.0,请参考docker/Dockerfile中的安装说明。MyoSuite需要gym==0.13.0,这与Meta-World和ManiSkill2不兼容。如果需要,请使用pip install myosuite单独安装。根据您现有的系统包,您可能需要安装其他依赖项。请参考docker/Dockerfile中推荐的系统包列表。

支持的任务

该代码库目前支持来自DMControl、Meta-World、ManiSkill2和MyoSuite的104个连续控制任务。具体来说,它支持39个DMControl任务(包括11个自定义任务),50个Meta-World任务,5个ManiSkill2任务和10个MyoSuite任务,涵盖了论文中使用的所有任务。以下是每个任务域的预期名称格式:

可以通过为evaluation.py指定task参数来运行这些任务。多任务训练和评估可以通过设置task=mt80或task=mt30来指定80任务集和30任务集。

截至2023年12月27日,TD-MPC2代码库还支持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

我们建议对单任务在线强化学习使用默认超参数,包括默认的5M参数模型大小(model_size=5)。多任务离线强化学习受益于更大的模型尺寸,但更大的模型也越来越昂贵。可用参数为model_size={1, 5, 19, 48, 317}.有关完整的参数列表,请参见config.yaml。

截至2024年1月7日,TD-MPC2代码库还支持多GPU训练,用于多任务离线强化学习实验;使用分支distributed和参数world_size=N在N个GPU上进行训练。我们无法保证分布式训练会产生相同的结果,但根据我们有限的测试,它们似乎类似。

引用

如果您觉得我们的工作有用,请考虑引用我们的论文:

@inproceedings{hansen2024tdmpc2,
  title={TD-MPC2: Scalable, Robust World Models for Continuous Control}, 
  author={Nicklas Hansen and Hao Su and Xiaolong Wang},
  booktitle={International Conference on Learning Representations (ICLR)},
  year={2024}
}

以及原始的TD-MPC论文:

@inproceedings{hansen2022tdmpc,
  title={Temporal Difference Learning for Model Predictive Control},
  author={Nicklas Hansen and Xiaolong Wang and Hao Su},
  booktitle={International Conference on Machine Learning (ICML)},
  year={2022}
}

贡献

您非常欢迎为这个项目做出贡献。如果您有任何建议或错误报告,请随时提出问题或拉取请求,但请先查看我们的指南。我们的目标是建立一个可轻松扩展到新环境和任务的代码库,我们很乐意听到您的体验!

许可证

该项目采用MIT许可证 - 有关详细信息,请参见LICENSE文件。请注意,该存储库依赖于第三方代码,这些代码受其各自的许可证约束。