Updated on 2024.06.27 DI-engine-v0.5.2

Introduction to DI-engine

Documentation | 中文文档 | Tutorials | Feature | Task & Middleware | TreeTensor | Roadmap

DI-engine is a generalized decision intelligence engine for PyTorch and JAX.

It provides python-first and asynchronous-native task and middleware abstractions, and modularly integrates several of the most important decision-making concepts: Env, Policy and Model. Based on the above mechanisms, DI-engine supports various deep reinforcement learning algorithms with superior performance, high efficiency, well-organized documentation and unittest:

• Most basic DRL algorithms: such as DQN, Rainbow, PPO, TD3, SAC, R2D2, IMPALA
• Multi-agent RL algorithms: such as QMIX, WQMIX, MAPPO, HAPPO, ACE
• Imitation learning algorithms (BC/IRL/GAIL): such as GAIL, SQIL, Guided Cost Learning, Implicit BC
• Offline RL algorithms: BCQ, CQL, TD3BC, Decision Transformer, EDAC, Diffuser, Decision Diffuser, SO2
• Model-based RL algorithms: SVG, STEVE, MBPO, DDPPO, DreamerV3
• Exploration algorithms: HER, RND, ICM, NGU
• LLM + RL Algorithms: PPO-max, DPO, PromptPG, PromptAWR
• Other algorithms: such as PER, PLR, PCGrad
• MCTS + RL algorithms: AlphaZero, MuZero, please refer to LightZero
• Generative Model + RL algorithms: Diffusion-QL, QGPO, SRPO, please refer to GenerativeRL

DI-engine aims to standardize different Decision Intelligence environments and applications, supporting both academic research and prototype applications. Various training pipelines and customized decision AI applications are also supported:

On the low-level end, DI-engine comes with a set of highly re-usable modules, including RL optimization functions, PyTorch utilities and auxiliary tools.

BTW, DI-engine also has some special system optimization and design for efficient and robust large-scale RL training:

Have fun with exploration and exploitation.

Outline

• Introduction to DI-engine
• Outline
• Installation
• Quick Start
• Feature
  • Algorithm Versatility
  • Environment Versatility
  • General Data Container: TreeTensor
• Feedback and Contribution
• Supporters
  • ↳ Stargazers
  • ↳ Forkers
• Citation
• License

Installation

You can simply install DI-engine from PyPI with the following command:

pip install DI-engine

If you use Anaconda or Miniconda, you can install DI-engine from conda-forge through the following command:

conda install -c opendilab di-engine

For more information about installation, you can refer to installation.

And our dockerhub repo can be found here，we prepare base image and env image with common RL environments.

The detailed documentation are hosted on doc | 中文文档.

Quick Start

3 Minutes Kickoff

3 Minutes Kickoff (colab)

DI-engine Huggingface Kickoff (colab)

How to migrate a new RL Env | 如何迁移一个新的强化学习环境

How to customize the neural network model | 如何定制策略使用的神经网络模型

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

新老 pipeline 的异同对比

Feature

Algorithm Versatility

Environment Versatility

General Data Container: TreeTensor

DI-engine utilizes TreeTensor as the basic data container in various components, which is ease of use and consistent across different code modules such as environment definition, data processing and DRL optimization. Here are some concrete code examples:

• TreeTensor can easily extend all the operations of torch.Tensor to nested data:

  (Click for Details)

  ```python import treetensor.torch as ttorch # create random tensor data = ttorch.randn({'a': (3, 2), 'b': {'c': (3, )}}) # clone+detach tensor data_clone = data.clone().detach() # access tree structure like attribute a = data.a c = data.b.c # stack/cat/split stacked_data = ttorch.stack([data, data_clone], 0) cat_data = ttorch.cat([data, data_clone], 0) data, data_clone = ttorch.split(stacked_data, 1) # reshape data = data.unsqueeze(-1) data = data.squeeze(-1) flatten_data = data.view(-1) # indexing data_0 = data[0] data_1to2 = data[1:2] # execute math calculations data = data.sin() data.b.c.cos_().clamp_(-1, 1) data += data ** 2 # backward data.requires_grad_(True) loss = data.arctan().mean() loss.backward() # print shape print(data.shape) # result # # ├── 'a' --> torch.Size([1, 3, 2]) # └── 'b' --> # └── 'c' --> torch.Size([1, 3]) ```

• TreeTensor can make it simple yet effective to implement classic deep reinforcement learning pipeline

  (Click for Details)

  ```diff 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)] # execute `stack` op - 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("not support elem type: {}".format(type(elem))) - stacked_data = stack(data, dim=0) + data = [ttorch.tensor(d) for d in data] + stacked_data = ttorch.stack(data, dim=0) # validate - 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 ```

Feedback and Contribution

• File an issue on Github
• Open or participate in our forum
• Discuss on DI-engine discord server
• Discuss on DI-engine slack communication channel

• Discuss on DI-engine's WeChat group (i.e. add us on WeChat: ding314assist)

  
• Contact our email (opendilab@pjlab.org.cn)
• Contributes to our future plan Roadmap

We appreciate all the feedbacks and contributions to improve DI-engine, both algorithms and system designs. And CONTRIBUTING.md offers some necessary information.

Supporters

↳ Stargazers

↳ Forkers

Citation

@misc{ding,
    title={DI-engine: A Universal AI System/Engine for Decision Intelligence},
    author={Niu, Yazhe and Xu, Jingxin and Pu, Yuan and Nie, Yunpeng and Zhang, Jinouwen and Hu, Shuai and Zhao, Liangxuan and Zhang,  Ming and Liu, Yu},
    publisher={GitHub},
    howpublished={\url{https://github.com/opendilab/DI-engine}},
    year={2021},
}

License

DI-engine released under the Apache 2.0 license.