DI-engine is a comprehensive decision intelligence engine for PyTorch and JAX, designed to standardize and accelerate research and development in deep reinforcement learning (DRL). It caters to researchers and practitioners by providing a modular framework that supports a vast array of DRL algorithms, from basic to multi-agent, imitation learning, offline RL, model-based RL, and LLM+RL.

How It Works

DI-engine employs a python-first, asynchronous-native approach with task and middleware abstractions. It modularly integrates core decision-making components: Env, Policy, and Model. A key innovation is its use of TreeTensor as a general data container, simplifying data handling across diverse modules and enabling seamless extension of PyTorch tensor operations to nested data structures. This design promotes code reusability and efficient large-scale RL training.

Quick Start & Requirements

• Install: pip install DI-engine
• Prerequisites: PyTorch or JAX. Docker images are available for various environments (Atari, MuJoCo, etc.).
• Resources: Official documentation and tutorials are available, including a 3-minute kickoff Colab notebook.

Highlighted Details

• Supports over 50 DRL algorithms, including DQN, PPO, SAC, QMIX, GAIL, CQL, DreamerV3, and LLM+RL methods like PromptPG.
• Extensive environment support, covering Atari, MuJoCo, PyBullet, PettingZoo (MARL), SMAC, GFootball, and real-world applications like autonomous driving (DI-drive) and traffic light control (DI-smartcross).
• Features TreeTensor for efficient handling of complex, nested data structures, simplifying pipeline implementation.
• Includes specialized modules for distributed RL (DI-orchestrator, DI-hpc) and efficient data storage (DI-store).

Maintenance & Community

The project is actively maintained by OpenDILab. Community engagement is encouraged via GitHub issues, Discord, Slack, and WeChat. A roadmap is available for future development.

Licensing & Compatibility

DI-engine is released under the Apache 2.0 license, permitting commercial use and integration with closed-source projects.

Limitations & Caveats

While comprehensive, the sheer breadth of algorithms and environments means some may be less mature or have fewer examples than others. Users should consult specific algorithm documentation for detailed performance characteristics and potential limitations.