OpenRL: 一个统一的强化学习框架

OpenRL: 打造统一简洁的强化学习框架

强化学习(Reinforcement Learning, RL)作为人工智能的一个重要分支,近年来在游戏、机器人、自然语言处理等领域取得了突破性进展。然而,相比于机器学习的其他分支,RL仍然相对小众,学习资源较为缺乏。为了让更多人能够方便地入门和使用RL,OpenRL-Lab团队开发了OpenRL这一开源强化学习框架。

OpenRL是基于PyTorch开发的通用强化学习研究框架,旨在为RL研究社区提供一个简单易用、灵活高效且可持续的平台。它支持单智能体、多智能体、离线RL、自我对弈和自然语言等多种任务的训练,采用模块化设计和高层抽象,让用户可以通过统一且友好的接口完成各种任务的训练。

OpenRL的主要特点

OpenRL具有以下突出特点:

统一简洁的接口:支持所有任务/环境的训练
支持单智能体和多智能体任务
支持使用专家数据集进行离线RL训练
支持自我对弈训练
支持自然语言任务(如对话)的强化学习训练
支持DeepSpeed加速训练
提供Arena功能,方便评估各种智能体
支持从Hugging Face导入模型和数据集
提供用户自定义环境集成教程
支持LSTM、GRU、Transformer等多种模型
多种训练加速方法
支持用户自定义训练模型、奖励模型、训练数据和环境
支持gymnasium环境
提供Callbacks功能实现日志、保存、早停等功能
支持字典观察空间
支持wandb、tensorboardX等可视化工具
串行或并行环境训练,保证结果一致性
提供中英文文档
提供单元测试和代码覆盖率测试
遵循Black代码风格指南和类型检查

OpenRL与其他RL库的对比

相比于其他流行的RL库,OpenRL在功能覆盖面上更加全面。下表对比了OpenRL与一些主流RL库的功能支持情况:

库	NLP/RLHF	多智能体	自我对弈训练	离线RL	DeepSpeed
OpenRL	✔️	✔️	✔️	✔️	✔️
Stable Baselines3	❌	❌	❌	❌	❌
Ray/RLlib	❌	✔️	✔️	✔️	❌
DI-engine	❌	✔️	部分支持	✔️	❌
Tianshou	❌	部分支持	部分支持	✔️	❌
MARLlib	❌	✔️	部分支持	❌	❌
MAPPO Benchmark	❌	✔️	❌	❌	❌
RL4LMs	✔️	❌	❌	❌	❌
trlx	✔️	❌	❌	❌	✔️
trl	✔️	❌	❌	❌	✔️

可以看到,OpenRL在功能覆盖上更加全面,尤其是在NLP/RLHF、多智能体、自我对弈训练等方面具有明显优势。

OpenRL的安装与使用

OpenRL的安装非常简单,用户可以直接通过pip安装:

pip install openrl

使用Anaconda或Miniconda的用户也可以通过conda安装:

conda install -c openrl openrl

想要修改源代码的用户可以从源码安装:

git clone https://github.com/OpenRL-Lab/openrl.git && cd openrl
pip install -e .

安装完成后,可以通过命令行检查OpenRL的版本:

openrl --version

OpenRL还提供了Docker镜像,方便用户快速部署:

# 无GPU加速
sudo docker pull openrllab/openrl-cpu
# GPU加速 
sudo docker pull openrllab/openrl

OpenRL的快速上手

OpenRL为强化学习初学者提供了简单易用的接口。以下是使用PPO算法训练CartPole环境的示例代码:

from openrl.envs.common import make
from openrl.modules.common import PPONet as Net
from openrl.runners.common import PPOAgent as Agent

env = make("CartPole-v1", env_num=9)  # 创建环境,设置环境并行数为9
net = Net(env)  # 创建神经网络
agent = Agent(net)  # 初始化智能体
agent.train(total_time_steps=20000)  # 开始训练,设置环境运行总步数为20000步

使用OpenRL训练智能体只需四个简单步骤:创建环境 => 初始化模型 => 初始化智能体 => 开始训练。

对于训练好的智能体,用户也可以轻松进行测试:

agent = Agent(Net(make("CartPole-v1", env_num=9)))  # 初始化训练器
agent.train(total_time_steps=20000)
# 创建测试环境,设置环境并行数为9,设置渲染模式为group_human
env = make("CartPole-v1", env_num=9, render_mode="group_human")
agent.set_env(env)  # 智能体需要交互环境
obs, info = env.reset()  # 初始化环境获得初始观察和环境信息
while True:
    action, _ = agent.act(obs)  # 智能体根据环境观察预测下一步动作
    # 环境根据动作前进一步,获得下一步观察、奖励、是否结束和环境信息
    obs, r, done, info = env.step(action)
    if any(done): break
env.close()  # 关闭测试环境

在普通笔记本电脑上执行上述代码,只需几秒钟就能完成训练。

CartPole训练效果

对于多智能体和自然语言处理等训练任务,OpenRL同样提供了简单易用的接口。

OpenRL的应用与发展

目前,OpenRL已经支持了大量强化学习算法和环境,包括:

算法:

PPO, A2C, DDPG, SAC, DQN等主流RL算法
MAPPO等多智能体算法
GAIL, BC等模仿学习算法
多臂老虎机算法
基于模型的RL算法
离线RL算法

环境:

Gymnasium, MuJoCo, PettingZoo, MPE等标准环境
Atari, StarCraft II, SMACv2等游戏环境
Omniverse Isaac Gym等物理仿真环境
自然语言对话环境

OpenRL正在被越来越多的研究项目使用,已经成为一个成熟的强化学习框架。OpenRL-Lab团队将持续维护和更新OpenRL,欢迎更多人加入开源社区,为强化学习的发展贡献力量。

总结

OpenRL作为一个统一的强化学习框架,具有接口简洁、功能全面、易于上手等优点,为强化学习的研究和应用提供了便利。无论是RL初学者还是资深研究者,都可以从OpenRL中受益。随着社区的不断发展,相信OpenRL会在未来发挥更大的作用,推动强化学习技术的进步。

如果OpenRL对您的工作有所帮助,欢迎引用:

@article{huang2023openrl,
  title={OpenRL: A Unified Reinforcement Learning Framework},
  author={Huang, Shiyu and Chen, Wentse and Sun, Yiwen and Bie, Fuqing and Tu, Wei-Wei},
  journal={arXiv preprint arXiv:2312.16189},
  year={2023}
}

OpenRL是一个开源项目,欢迎更多人参与贡献。无论是提出建议、报告问题,还是直接贡献代码,都可以帮助OpenRL变得更好。让我们共同努力,为强化学习的发展贡献自己的力量!