轨迹 Transformer

针对 将离线强化学习视为一个大型序列建模问题 的代码已发布。

新增： 也请查看 Alexander Nikulin 的分支，其中包含了注意力缓存和向量化回放缓冲区！

安装

所有 Python 依赖项均列在 environment.yml 中。可通过以下命令安装：

conda env create -f environment.yml
conda activate trajectory
pip install -e .

为确保可复现性，我们还在 Dockerfile 中包含了系统要求（参见 安装说明），不过 Conda 安装通常适用于大多数标准的 Linux 机器。

使用

使用以下命令训练 Transformer：python scripts/train.py --dataset halfcheetah-medium-v2

要复现离线强化学习结果：python scripts/plan.py --dataset halfcheetah-medium-v2

默认情况下，这些命令将使用 config/offline.py 中的超参数。您也可以通过运行时参数覆盖它们：

python scripts/plan.py --dataset halfcheetah-medium-v2 \
	--horizon 5 --beam_width 32

由于离散化策略的变化，部分超参数与论文中列出的不同。这些超参数将在下一版 arXiv 论文中更新，以匹配当前代码库中的设置。

预训练模型

我们提供了 16 个数据集的预训练模型：{halfcheetah, hopper, walker2d, ant}-{expert-v2, medium-expert-v2, medium-v2, medium-replay-v2}。可通过 ./pretrained.sh 下载这些模型。

模型将保存在 logs/$DATASET/gpt/pretrained 目录下。要使用这些模型进行规划，可以使用 gpt_loadpath 标志指定：

python scripts/plan.py --dataset halfcheetah-medium-v2 \
	--gpt_loadpath gpt/pretrained

pretrained.sh 还会从每个模型下载 15 个 规划方案，保存到 logs/$DATASET/plans/pretrained。您可以使用 python plotting/read_results.py 来读取这些规划方案。

\begin{table*}[h]
\centering
\small
\begin{tabular}{llrrrrrr}
\toprule
\multicolumn{1}{c}{\bf 数据集} & \multicolumn{1}{c}{\bf 环境} & \multicolumn{1}{c}{\bf BC} & \multicolumn{1}{c}{\bf MBOP} & \multicolumn{1}{c}{\bf BRAC} & \multicolumn{1}{c}{\bf CQL} & \multicolumn{1}{c}{\bf DT} & \multicolumn{1}{c}{\bf TT (我们的方法)} \\
\midrule
Medium-Expert & HalfCheetah & $59.9$ & $105.9$ & $41.9$ & $91.6$ & $86.8$ & $95.0$ \scriptsize{\raisebox{1pt}{$\pm 0.2$}} \\
Medium-Expert & Hopper & $79.6$ & $55.1$ & $0.9$ & $105.4$ & $107.6$ & $110.0$ \scriptsize{\raisebox{1pt}{$\pm 2.7$}} \\
Medium-Expert & Walker2d & $36.6$ & $70.2$ & $81.6$ & $108.8$ & $108.1$ & $101.9$ \scriptsize{\raisebox{1pt}{$\pm 6.8$}} \\
Medium-Expert & Ant & $-$ & $-$ & $-$ & $-$ & $-$ & $116.1$ \scriptsize{\raisebox{1pt}{$\pm 9.0$}} \\
\midrule
Medium & HalfCheetah & $43.1$ & $44.6$ & $46.3$ & $44.0$ & $42.6$ & $46.9$ \scriptsize{\raisebox{1pt}{$\pm 0.4$}} \\
Medium & Hopper & $63.9$ & $48.8$ & $31.3$ & $58.5$ & $67.6$ & $61.1$ \scriptsize{\raisebox{1pt}{$\pm 3.6$}} \\
Medium & Walker2d & $77.3$ & $41.0$ & $81.1$ & $72.5$ & $74.0$ & $79.0$ \scriptsize{\raisebox{1pt}{$\pm 2.8$}} \\
Medium & Ant & $-$ & $-$ & $-$ & $-$ & $-$ & $83.1$ \scriptsize{\raisebox{1pt}{$\pm 7.3$}} \\
\midrule
Medium-Replay & HalfCheetah & $4.3$ & $42.3$ & $47.7$ & $45.5$ & $36.6$ & $41.9$ \scriptsize{\raisebox{1pt}{$\pm 2.5$}} \\
Medium-Replay & Hopper & $27.6$ & $12.4$ & $0.6$ & $95.0$ & $82.7$ & $91.5$ \scriptsize{\raisebox{1pt}{$\pm 3.6$}} \\
Medium-Replay & Walker2d & $36.9$ & $9.7$ & $0.9$ & $77.2$ & $66.6$ & $82.6$ \scriptsize{\raisebox{1pt}{$\pm 6.9$}} \\
Medium-Replay & Ant & $-$ & $-$ & $-$ & $-$ & $-$ & $77.0$ \scriptsize{\raisebox{1pt}{$\pm 6.8$}} \\
\midrule
\multicolumn{2}{c}{\bf 平均值（不包括 Ant）} & 47.7 & 47.8 & 36.9 & 77.6 & 74.7 & 78.9 \hspace{.6cm} \\
\multicolumn{2}{c}{\bf 平均值（所有设置）} & $-$ & $-$ & $-$ & $-$ & $-$ & 82.2 \hspace{.6cm} \\
\bottomrule
\end{tabular}
\label{table:d4rl}
\end{table*}

Docker

将您的 MuJoCo 密钥复制到 Docker 构建上下文中，并构建容器：

cp ~/.mujoco/mjkey.txt azure/files/
docker build -f azure/Dockerfile . -t trajectory

测试容器：

docker run -it --rm --gpus all \
	--mount type=bind,source=$PWD,target=/home/code \
	--mount type=bind,source=$HOME/.d4rl,target=/root/.d4rl \
	trajectory \
	bash -c \
	"export PYTHONPATH=$PYTHONPATH:/home/code && \
	python /home/code/scripts/train.py --dataset hopper-medium-expert-v2 --exp_name docker/"

在 Azure 上运行

设置

1. 在 Azure 上启动作业还需要一个额外的 Python 依赖项：

pip install git+https://github.com/JannerM/doodad.git@janner

2. 将上一节 [Docker] 中构建的镜像打标签并推送到 Docker Hub：

export DOCKER_USERNAME=$(docker info | sed '/Username:/!d;s/.* //')
docker tag trajectory ${DOCKER_USERNAME}/trajectory:latest
docker image push ${DOCKER_USERNAME}/trajectory

3. 更新 azure/config.py，可以直接修改文件，也可以设置相应的 环境变量。要设置 AZURE_STORAGE_CONNECTION 变量，需前往您的存储账户的“访问密钥”部分，单击“显示密钥”，然后复制“连接字符串”。

4. 下载 azcopy：./azure/download.sh

使用

使用 python azure/launch_train.py 启动训练作业，使用 python azure/launch_plan.py 启动规划作业。

这些脚本不接受运行时参数，而是通过 params_to_sweep 中的参数组合来运行对应的脚本（分别为 scripts/train.py 和 scripts/plan.py）。

查看结果

要将结果从 Azure 存储容器同步到本地，运行 ./azure/sync.sh。

挂载存储容器的步骤如下：

1. 使用 ./azure/make_fuse_config.sh 创建 blobfuse 配置文件。
2. 运行 ./azure/mount.sh 将存储容器挂载到 ~/azure_mount。

卸载容器时，请运行 sudo umount -f ~/azure_mount; rm -r ~/azure_mount。

参考文献

@inproceedings{janner2021sequence,
  title = {将离线强化学习视为一个大型序列建模问题},
  author = {迈克尔·詹纳和李启阳和谢尔盖·列文},
  booktitle = {神经信息处理系统进展},
  year = {2021},
}

致谢

GPT 的实现源自安德烈·卡帕西的 minGPT 仓库。

Trajectory Transformer 快速上手指南

Trajectory Transformer 是将离线强化学习（Offline RL）转化为序列建模问题的开源实现。本指南帮助开发者快速搭建环境并运行基础示例。

环境准备

• 操作系统: 推荐 Linux (Ubuntu/CentOS)，Windows 用户建议使用 WSL2 或 Docker。
• 核心依赖:
  • Conda (推荐 Miniconda)
  • GPU 驱动及 CUDA 工具包 (用于加速训练，可选但推荐)
  • MuJoCo 许可证 (运行 D4RL 数据集必需)
• 前置知识: 熟悉 Python 虚拟环境管理及基础的强化学习概念。

注意：国内用户若遇到 conda 或 pip 下载缓慢，可临时切换至清华或阿里镜像源：

conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/
pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

1. 克隆仓库

   git clone https://github.com/jannerm/trajectory-transformer.git
   cd trajectory-transformer
   

2. 创建并激活 Conda 环境 使用项目提供的配置文件安装所有 Python 依赖：

   conda env create -f environment.yml
   conda activate trajectory
   

3. 安装本地包 以可编辑模式安装当前项目：

   pip install -e .
   

4. 配置 MuJoCo (关键) 确保已将 MuJoCo 密钥文件 (mjkey.txt) 放置在 ~/.mujoco/ 目录下。若未安装 MuJoCo，请先完成其官方安装流程。

基本使用

1. 训练模型

使用默认超参数在 halfcheetah-medium-v2 数据集上训练 Transformer 模型：

python scripts/train.py --dataset halfcheetah-medium-v2

2. 执行规划 (推理)

加载训练好的模型进行离线策略规划，复现论文结果：

python scripts/plan.py --dataset halfcheetah-medium-v2

3. 自定义参数

可通过命令行标志覆盖默认配置（例如修改预测步长和束搜索宽度）：

python scripts/plan.py --dataset halfcheetah-medium-v2 \
	--horizon 5 --beam_width 32

4. 使用预训练模型 (可选)

若不想从头训练，可下载作者提供的 16 个数据集预训练权重：

# 下载预训练模型和示例计划
./pretrained.sh

# 使用预训练模型进行规划
python scripts/plan.py --dataset halfcheetah-medium-v2 \
	--gpt_loadpath gpt/pretrained

模型文件将保存至 logs/$DATASET/gpt/pretrained，生成的计划数据位于 logs/$DATASET/plans/pretrained。