---

更新于 2024年12月23日 DI-engine-v0.5.3

DI-engine 简介

文档 | 中文文档 | 教程 | 特性 | 任务与中间件 | TreeTensor | 路线图

DI-engine 是一个基于 PyTorch 和 JAX 的通用决策智能引擎。

它提供了 以 Python 为中心 且 原生异步 的任务和中间件抽象，并模块化地集成了几个最重要的决策概念：环境、策略和模型。基于上述机制，DI-engine 支持 多种深度强化学习算法，具有卓越的性能、高效的运行效率、结构清晰的 文档 以及完善的 单元测试：

• 最基础的 DRL 算法：如 DQN、Rainbow、PPO、TD3、SAC、R2D2、IMPALA
• 多智能体 RL 算法：如 QMIX、WQMIX、MAPPO、HAPPO、ACE
• 仿生学习算法（BC/IRL/GAIL）：如 GAIL、SQIL、Guided Cost Learning、Implicit BC
• 离线 RL 算法：BCQ、CQL、TD3BC、Decision Transformer、EDAC、Diffuser、Decision Diffuser、SO2
• 基于模型的 RL 算法：SVG、STEVE、MBPO、DDPPO、DreamerV3
• 探索性算法：HER、RND、ICM、NGU
• LLM + RL 算法：PPO-max、DPO、PromptPG、PromptAWR
• 其他算法：如 PER、PLR、PCGrad
• MCTS + RL 算法：AlphaZero、MuZero，请参考 LightZero
• 生成模型 + RL 算法：Diffusion-QL、QGPO、SRPO，请参考 GenerativeRL

DI-engine 的目标是 标准化不同的决策智能环境和应用，既支持学术研究，也适用于原型应用开发。此外，还支持各种训练流水线和定制化的决策 AI 应用程序：

(点击收起)

• 传统学术环境
  • DI-zoo：结合 DI-engine 的多种决策智能演示与基准环境。
• 教程课程
  • PPOxFamily：PPO x Family 强化学习教程课程。
• 现实世界中的决策 AI 应用
  • DI-star：星际争霸 II 中的决策 AI。
  • PsyDI：面向心理评估的多模态交互式聊天机器人。
  • DI-drive：自动驾驶平台。
  • DI-sheep：三子棋游戏中的决策 AI。
  • DI-smartcross：交通信号灯控制中的决策 AI。
  • DI-bioseq：生物序列预测与搜索中的决策 AI。
  • DI-1024：深度强化学习 + 1024 游戏。
• 研究论文
  • InterFuser：[CoRL 2022] 基于可解释传感器融合 Transformer 的安全增强自动驾驶。
  • ACE：[AAAI 2023] ACE：具有双向动作依赖性的合作式多智能体 Q 学习。
  • GoBigger：[ICLR 2023] 多智能体决策智能环境。
  • DOS：[CVPR 2023] ReasonNet：基于时序与全局推理的端到端驾驶。
  • LightZero：[NeurIPS 2023 Spotlight] 轻量高效 MCTS/AlphaZero/MuZero 算法工具包。
  • SO2：[AAAI 2024] 关于离线到在线强化学习中 Q 值估计的一个视角。
  • LMDrive：[CVPR 2024] LMDrive：基于大型语言模型的闭环端到端驾驶。
  • SmartRefine：[CVPR 2024] SmartRefine：用于高效运动预测的场景自适应精炼框架。
  • ReZero：通过反向视图和全缓冲区重新分析提升基于 MCTS 的算法性能。
  • UniZero：利用可扩展潜在世界模型实现通用且高效的规划。
• 文档与教程
  • DI-engine-docs：教程、最佳实践及 API 参考。
  • awesome-model-based-RL：精选的基于模型的强化学习资源列表。
  • awesome-exploration-RL：精选的探索型强化学习资源列表。
  • awesome-decision-transformer：精选的决策 Transformer 资源列表。
  • awesome-RLHF：精选的强化学习与人类反馈资源列表。
  • awesome-multi-modal-reinforcement-learning：精选的多模态强化学习资源列表。
  • awesome-diffusion-model-in-rl：精选的强化学习中扩散模型资源列表。
  • awesome-ui-agents：精选的 UI 代理资源列表，涵盖 Web、App、操作系统等。
  • awesome-AI-based-protein-design：人工智能辅助蛋白质设计的研究论文合集。
  • awesome-end-to-end-autonomous-driving：精选的端到端自动驾驶资源列表。
  • awesome-driving-behavior-prediction：关于驾驶行为预测的研究论文合集。

在底层，DI-engine 提供了一系列高度可复用的模块，包括 RL 优化函数、PyTorch 工具以及 辅助工具。

此外，DI-engine 还针对高效、稳健的大规模强化学习训练进行了一些特殊的 系统优化与设计：

(点击查看详情)

• treevalue：树形嵌套数据结构。
• DI-treetensor：树形嵌套 PyTorch 张量库。
• DI-toolkit：一个用于决策智能的简单工具包。
• DI-orchestrator：强化学习 Kubernetes 自定义资源与 Operator 库。
• DI-hpc：强化学习 HPC OP 库。
• DI-store：强化学习对象存储。

尽情享受探索与利用的乐趣吧。

大纲

• DI-engine 简介
• 大纲
• 安装
• 快速入门
• 特性
  • 算法多样性
  • 环境多样性
  • 通用数据容器：TreeTensor
• 反馈与贡献
• 支持者
  • ↳ 星标用户
  • ↳ 分叉用户
• 引用
• 许可证

安装

您可以通过以下命令从 PyPI 简单地安装 DI-engine：

pip install DI-engine

有关安装的更多信息，请参阅 安装指南。

我们的 Docker Hub 仓库地址为 这里，我们准备了包含常见 RL 环境的 基础镜像 和 环境镜像。

(点击查看详情)

• 基础镜像：opendilab/ding:nightly
• RPC 镜像：opendilab/ding:nightly-rpc
• Atari 镜像：opendilab/ding:nightly-atari
• Mujoco 镜像：opendilab/ding:nightly-mujoco
• DMC 镜像：opendilab/ding:nightly-dmc2gym
• MetaWorld 镜像：opendilab/ding:nightly-metaworld
• SMAC 镜像：opendilab/ding:nightly-smac
• GRF 镜像：opendilab/ding:nightly-grf
• CityFlow 镜像：opendilab/ding:nightly-cityflow
• EvoGym 镜像：opendilab/ding:nightly-evogym
• D4RL 镜像：opendilab/ding:nightly-d4rl

详细文档托管在 doc | 中文文档 上。

快速入门

3分钟入门

3分钟入门 (colab)

DI-engine Huggingface 入门 (colab)

如何迁移一个新的强化学习环境 | 如何迁移一个新的强化学习环境

如何定制策略使用的神经网络模型 | 如何定制策略使用的神经网络模型

测试/部署 强化学习策略 的样例

新老 pipeline 的异同对比

功能特性

算法多样性

(点击收起)

 离散动作空间，常见于常规强化学习算法中（1-23）

 连续动作空间，常见于常规强化学习算法中（1-23）

 混合动作空间（离散+连续）（1-23）

 分布式强化学习｜分布式强化学习

 多智能体强化学习｜多智能体强化学习

 强化学习中的探索机制｜强化学习中的探索机制

 模仿学习｜模仿学习

 离线强化学习｜离线强化学习

 基于模型的强化学习｜基于模型的强化学习

 其他细分方向的算法，通常作为整个流程中的插件式组件

注：Runnable Demo 中的 .py 文件可在 dizoo 目录下找到。

序号	算法	标签	文档与实现	可运行Demo
1	DQN		DQN文档 DQN中文文档 policy/dqn	python3 -u cartpole_dqn_main.py / ding -m serial -c cartpole_dqn_config.py -s 0
2	C51		C51文档 policy/c51	ding -m serial -c cartpole_c51_config.py -s 0
3	QRDQN		QRDQN文档 policy/qrdqn	ding -m serial -c cartpole_qrdqn_config.py -s 0
4	IQN		IQN文档 policy/iqn	ding -m serial -c cartpole_iqn_config.py -s 0
5	FQF		FQF文档 policy/fqf	ding -m serial -c cartpole_fqf_config.py -s 0
6	Rainbow		Rainbow文档 policy/rainbow	ding -m serial -c cartpole_rainbow_config.py -s 0
7	SQL		SQL文档 policy/sql	ding -m serial -c cartpole_sql_config.py -s 0
8	R2D2		R2D2文档 policy/r2d2	ding -m serial -c cartpole_r2d2_config.py -s 0
9	PG		PG文档 policy/pg	ding -m serial -c cartpole_pg_config.py -s 0
10	PromptPG		policy/prompt_pg	ding -m serial_onpolicy -c tabmwp_pg_config.py -s 0
11	A2C		A2C文档 policy/a2c	ding -m serial -c cartpole_a2c_config.py -s 0
12	PPO/MAPPO		PPO文档 policy/ppo	python3 -u cartpole_ppo_main.py / ding -m serial_onpolicy -c cartpole_ppo_config.py -s 0
13	PPG		PPG文档 policy/ppg	python3 -u cartpole_ppg_main.py
14	ACER		ACER文档 policy/acer	ding -m serial -c cartpole_acer_config.py -s 0
15	IMPALA		IMPALA文档 policy/impala	ding -m serial -c cartpole_impala_config.py -s 0
16	DDPG/PADDPG		DDPG文档 policy/ddpg	ding -m serial -c pendulum_ddpg_config.py -s 0
17	TD3		TD3文档 policy/td3	python3 -u pendulum_td3_main.py / ding -m serial -c pendulum_td3_config.py -s 0
18	D4PG		D4PG文档 policy/d4pg	python3 -u pendulum_d4pg_config.py
19	SAC/[MASAC]		SAC文档 policy/sac	ding -m serial -c pendulum_sac_config.py -s 0
20	PDQN		policy/pdqn	ding -m serial -c gym_hybrid_pdqn_config.py -s 0
21	MPDQN		policy/pdqn	ding -m serial -c gym_hybrid_mpdqn_config.py -s 0
22	HPPO		policy/ppo	ding -m serial_onpolicy -c gym_hybrid_hppo_config.py -s 0
23	BDQ		policy/bdq	python3 -u hopper_bdq_config.py
24	MDQN		policy/mdqn	python3 -u asterix_mdqn_config.py
25	QMIX		QMIX文档 policy/qmix	ding -m serial -c smac_3s5z_qmix_config.py -s 0
26	COMA		COMA文档 policy/coma	ding -m serial -c smac_3s5z_coma_config.py -s 0
27	QTran		policy/qtran	ding -m serial -c smac_3s5z_qtran_config.py -s 0
28	WQMIX		WQMIX文档 policy/wqmix	ding -m serial -c smac_3s5z_wqmix_config.py -s 0
29	CollaQ		CollaQ文档 policy/collaq	ding -m serial -c smac_3s5z_collaq_config.py -s 0
30	MADDPG		MADDPG文档 policy/ddpg	ding -m serial -c ptz_simple_spread_maddpg_config.py -s 0
31	GAIL		GAIL文档 reward_model/gail	ding -m serial_gail -c cartpole_dqn_gail_config.py -s 0
32	SQIL		SQIL文档 entry/sqil	ding -m serial_sqil -c cartpole_sqil_config.py -s 0
33	DQFD		DQFD文档 policy/dqfd	ding -m serial_dqfd -c cartpole_dqfd_config.py -s 0
34	R2D3		R2D3文档 R2D3中文文档 policy/r2d3	python3 -u pong_r2d3_r2d2expert_config.py
35	Guided Cost Learning		Guided Cost Learning中文文档 reward_model/guided_cost	python3 lunarlander_gcl_config.py
36	TREX		TREX文档 reward_model/trex	python3 mujoco_trex_main.py
37	Implicit Behavorial Cloning (DFO+MCMC)		policy/ibc model/template/ebm	python3 d4rl_ibc_main.py -s 0 -c pen_human_ibc_mcmc_config.py
38	BCO		entry/bco	python3 -u cartpole_bco_config.py
39	HER		HER文档 reward_model/her	python3 -u bitflip_her_dqn.py
40	RND		RND文档 reward_model/rnd	python3 -u cartpole_rnd_onppo_config.py
41	ICM		ICM文档 ICM中文文档 reward_model/icm	python3 -u cartpole_ppo_icm_config.py
42	CQL		CQL文档 policy/cql	python3 -u d4rl_cql_main.py
43	TD3BC		TD3BC文档 policy/td3_bc	python3 -u d4rl_td3_bc_main.py
44	Decision Transformer		policy/dt	python3 -u d4rl_dt_mujoco.py
45	EDAC		EDAC文档 policy/edac	python3 -u d4rl_edac_main.py
46	QGPO		QGPO文档 policy/qgpo	python3 -u ding/example/qgpo.py
47	MBSAC(SAC+MVE+SVG)		policy/mbpolicy/mbsac	python3 -u pendulum_mbsac_mbpo_config.py \ python3 -u pendulum_mbsac_ddppo_config.py
48	STEVESAC(SAC+STEVE+SVG)		policy/mbpolicy/mbsac	python3 -u pendulum_stevesac_mbpo_config.py
49	MBPO		MBPO文档 world_model/mbpo	python3 -u pendulum_sac_mbpo_config.py
50	DDPPO		world_model/ddppo	python3 -u pendulum_mbsac_ddppo_config.py
51	DreamerV3		world_model/dreamerv3	python3 -u cartpole_balance_dreamer_config.py
52	PER		worker/replay_buffer	rainbow demo
53	GAE		rl_utils/gae	ppo demo
54	ST-DIM		torch_utils/loss/contrastive_loss	ding -m serial -c cartpole_dqn_stdim_config.py -s 0
55	PLR		PLR文档 data/level_replay/level_sampler	python3 -u bigfish_plr_config.py -s 0
56	PCGrad		torch_utils/optimizer_helper/PCGrad	python3 -u multi_mnist_pcgrad_main.py -s 0
57	AWR		policy/ibc	python3 -u tabmwp_awr_config.py

环境适应性

(点击收起)

序号	环境	标签	可视化	代码与文档链接
1	Atari			dizoo链接 环境教程 环境指南
2	box2d/bipedalwalker			dizoo链接 环境教程 环境指南
3	box2d/lunarlander			dizoo链接 环境教程 环境指南
4	classic_control/cartpole			dizoo链接 环境教程 环境指南
5	classic_control/pendulum			dizoo链接 环境教程 环境指南
6	competitive_rl			dizoo链接 环境指南
7	gfootball			dizoo链接 环境教程 环境指南
8	minigrid			dizoo链接 环境教程 环境指南
9	MuJoCo			dizoo链接 环境教程 环境指南
10	PettingZoo			dizoo链接 环境教程 环境指南
11	overcooked			dizoo链接 环境教程
12	procgen			dizoo链接 环境教程 环境指南
13	pybullet			dizoo链接 环境指南
14	smac			dizoo链接 环境教程 环境指南
15	d4rl			dizoo链接 环境指南
16	league_demo			dizoo链接
17	pomdp atari			dizoo链接
18	bsuite			dizoo链接 环境教程 环境指南
19	ImageNet			dizoo链接 环境指南
20	slime_volleyball			dizoo链接 环境教程 环境指南
21	gym_hybrid			dizoo链接 环境教程 环境指南
22	GoBigger			dizoo链接 环境教程 环境指南
23	gym_soccer			dizoo链接 环境指南
24	multiagent_mujoco			dizoo链接 环境指南
25	bitflip			dizoo链接 环境指南
26	sokoban			dizoo链接 环境教程 环境指南
27	gym_anytrading			dizoo链接 环境教程
28	mario			dizoo链接 环境教程 环境指南
29	dmc2gym			dizoo链接 环境教程 环境指南
30	evogym			dizoo链接 环境教程 环境指南
31	gym-pybullet-drones			dizoo链接 环境指南
32	beergame			dizoo链接 环境指南
33	classic_control/acrobot			dizoo链接 环境指南
34	box2d/car_racing			dizoo链接 环境指南
35	metadrive			dizoo链接 环境指南
36	cliffwalking			dizoo链接 env tutorial 环境指南
37	tabmwp			dizoo链接 env tutorial 环境指南
38	frozen_lake			dizoo链接 env tutorial 环境指南
39	ising_model			dizoo链接 env tutorial 环境指南
40	taxi			dizoo链接 环境教程 环境指南

表示离散动作空间

表示连续动作空间

表示混合（离散+连续）动作空间

表示多智能体强化学习环境

表示与探索相关且奖励稀疏的环境

表示离线强化学习环境

表示模仿学习或监督学习数据集

表示允许智能体之间对战的环境

附注：Atari 中的一些环境，例如 MontezumaRevenge，也属于稀疏奖励类型。

通用数据容器：TreeTensor

DI-engine 在各个组件中使用 TreeTensor 作为基础数据容器，它易于使用，并且在环境定义、数据处理和 DRL 优化等不同代码模块中保持一致。以下是一些具体的代码示例：

• TreeTensor 可以轻松地将 torch.Tensor 的所有操作扩展到嵌套数据中：

  (点击查看详情)

  import treetensor.torch as ttorch
  
  
  # 创建随机张量
  data = ttorch.randn({'a': (3, 2), 'b': {'c': (3, )}})
  # 克隆并分离梯度
  data_clone = data.clone().detach()
  # 通过属性访问树形结构
  a = data.a
  c = data.b.c
  # 堆叠/拼接/分割
  stacked_data = ttorch.stack([data, data_clone], 0)
  cat_data = ttorch.cat([data, data_clone], 0)
  data, data_clone = ttorch.split(stacked_data, 1)
  # 重塑
  data = data.unsqueeze(-1)
  data = data.squeeze(-1)
  flatten_data = data.view(-1)
  # 索引
  data_0 = data[0]
  data_1to2 = data[1:2]
  # 执行数学计算
  data = data.sin()
  data.b.c.cos_().clamp_(-1, 1)
  data += data ** 2
  # 反向传播
  data.requires_grad_(True)
  loss = data.arctan().mean()
  loss.backward()
  # 打印形状
  print(data.shape)
  # 结果
  # <Size 0x7fbd3346ddc0>
  # ├── 'a' --> torch.Size([1, 3, 2])
  # └── 'b' --> <Size 0x7fbd3346dd00>
  #     └── 'c' --> torch.Size([1, 3])
  

• TreeTensor 可以让经典的深度强化学习流程实现起来既简单又高效。

  (点击查看详情)

  import torch
  import treetensor.torch as ttorch
  
  B = 4
  
  
  def get_item():
      return {
          'obs': {
              'scalar': torch.randn(12),
              'image': torch.randn(3, 32, 32),
          },
          'action': torch.randint(0, 10, size=(1,)),
          'reward': torch.rand(1),
          'done': False,
      }
  
  
  data = [get_item() for _ in range(B)]
  
  
  # 执行 `stack` 操作
  - def stack(data, dim):
  -     elem = data[0]
  -     if isinstance(elem, torch.Tensor):
  -         return torch.stack(data, dim)
  -     elif isinstance(elem, dict):
  -         return {k: stack([item[k] for item in data], dim) for k in elem.keys()}
  -     elif isinstance(elem, bool):
  -         return torch.BoolTensor(data)
  -     else:
  -         raise TypeError("不支持该元素类型: {}".format(type(elem)))
  - stacked_data = stack(data, dim=0)
  + data = [ttorch.tensor(d) for d in data]
  + stacked_data = ttorch.stack(data, dim=0)
  
  # 验证
  - assert stacked_data['obs']['image'].shape == (B, 3, 32, 32)
  - assert stacked_data['action'].shape == (B, 1)
  - assert stacked_data['reward'].shape == (B, 1)
  - assert stacked_data['done'].shape == (B,)
  - assert stacked_data['done'].dtype == torch.bool
  + assert stacked_data.obs.image.shape == (B, 3, 32, 32)
  + assert stacked_data.action.shape == (B, 1)
  + assert stacked_data.reward.shape == (B, 1)
  + assert stacked_data.done.shape == (B,)
  + assert stacked_data.done.dtype == torch.bool
  

反馈与贡献

• 在 Github 上 提交问题

• 打开或参与我们的 论坛

• 在 DI-engine 的 Discord 服务器 上讨论

• 在 DI-engine 的 Slack 沟通频道 上讨论

• 在 DI-engine 的微信交流群中讨论（添加我们微信：ding314assist）

  

• 联系我们的邮箱 (opendilab@pjlab.org.cn)

• 参与我们的未来计划 Roadmap

我们非常感谢所有关于算法和系统设计方面的反馈与贡献。CONTRIBUTING.md 文件提供了必要的信息。

支持者

↳ 星标用户

↳ 分支用户

引用

@misc{ding,
    title={DI-engine: 用于决策智能的通用 AI 系统/引擎},
    author={Niu, Yazhe 和 Xu, Jingxin 和 Pu, Yuan 和 Nie, Yunpeng 和 Zhang, Jinouwen 和 Hu, Shuai 和 Zhao, Liangxuan 和 Zhang, Ming 和 Liu, Yu},
    publisher={GitHub},
    howpublished={\url{https://github.com/opendilab/DI-engine}},
    year={2021},
}

许可证

DI-engine 采用 Apache 2.0 许可证发布。

DI-engine 快速上手指南

DI-engine 是一个基于 PyTorch 和 JAX 的通用决策智能引擎，支持多种深度强化学习（DRL）算法、多智能体强化学习、模仿学习及离线强化学习等。本指南将帮助你快速完成环境配置并运行第一个示例。

1. 环境准备

在开始之前，请确保你的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 18.04+) 或 macOS。Windows 用户建议使用 WSL2。
• Python 版本: 3.8, 3.9, 3.10 或 3.11。
• GPU 驱动: 如需使用 GPU 加速训练，请确保已安装正确的 NVIDIA 驱动和 CUDA 工具包。
• 包管理器: 推荐使用 conda 管理虚拟环境，也可使用 pip。

2. 安装步骤

方法一：使用 pip 安装（推荐）

你可以直接从 PyPI 安装最新稳定版。为了获得更快的下载速度，国内用户建议使用清华或阿里镜像源。

# 创建并激活虚拟环境 (可选但推荐)
conda create -n di-env python=3.9
conda activate di-env

# 使用国内镜像源安装 DI-engine
pip install DI-engine -i https://pypi.tuna.tsinghua.edu.cn/simple

方法二：使用 Conda 安装

如果你偏好使用 conda 通道进行安装：

conda create -n di-env python=3.9
conda activate di-env
conda install -c opendilab di-engine

验证安装

安装完成后，运行以下命令检查是否成功：

python -c "import ding; print(ding.__version__)"

如果输出了版本号且无报错，则说明安装成功。

3. 基本使用

DI-engine 的核心设计理念是模块化，主要包含 Env (环境)、Policy (策略) 和 Model (模型)。以下是一个最简化的代码示例，展示如何初始化一个 DQN 策略并在 CartPole 环境中运行一步。

from ding import DQNPolicy, CommandModeCollector, EnvManager
from ding.utils import set_pkg_seed
from gymnasium.wrappers import TimeLimit
import gymnasium as gym

# 1. 设置随机种子以保证复现性
set_pkg_seed(0, use_cuda=False)

# 2. 创建环境 (以 CartPole-v1 为例)
# DI-engine 通常需要对原生 gym 环境进行简单封装
env = gym.make('CartPole-v1')
env = TimeLimit(env, max_episode_steps=200)

# 3. 初始化策略配置
# 这里使用 DQN 的默认配置，实际使用中可加载自定义 config 文件
policy = DQNPolicy(cfg=dict(
    model=dict(
        obs_shape=4,
        action_shape=2,
        encoder_hidden_size_list=[128, 128],
    ),
    learn=dict(
        update_per_collect=1,
        batch_size=64,
        learning_rate=0.001,
    ),
    collect=dict(
        n_sample=100,
        unroll_len=1,
    ),
))

# 4. 初始化收集器 (Collector 负责与环境交互收集数据)
collector = CommandModeCollector(
    policy=policy,
    env=EnvManager([lambda: env]),
    trajectory_space_name='trajectory',
)

# 5. 执行一次数据采集
# collect() 返回 collected_data (收集到的数据) 和 other_info (其他信息)
collected_data, other_info = collector.collect(n_step=10, random_fraction=0.5)

print(f"成功收集了 {len(collected_data)} 条轨迹数据")
print("第一步快速上手完成！接下来可以查看官方文档进行完整训练。")

下一步建议

• 完整训练: 参考 di-engine 自带的配置文件（位于 ding/example 目录），使用 ding entry 命令行工具启动完整训练任务。
• 查阅文档: 访问 中文文档 获取详细的算法列表、API 参考及进阶教程。