RLLTE: 强化学习的长期演进项目

RLLTE简介

RLLTE (Reinforcement Learning Long-Term Evolution)是一个旨在推动强化学习(RL)研究和应用长期发展的开源项目。该项目的灵感来源于电信领域的长期演进(LTE)标准项目,旨在为RL研究和应用提供开发组件和标准。RLLTE不仅提供了顶级算法的实现,还作为开发算法的工具包,为研究人员和开发者提供了一个强大而灵活的平台。

RLLTE的主要特点

RLLTE具有以下几个显著特点:

长期演进:提供最新的算法和技巧,持续推动RL技术的发展。
完整生态系统:涵盖任务设计、模型训练、评估和部署(如TensorRT、CANN等)的全过程。
模块化设计:采用面向模块的设计,实现RL算法的完全解耦。
优化工作流:实现全硬件加速的优化工作流程。
支持自定义:允许用户自定义环境和模块。
多设备支持:支持GPU和NPU等多种计算设备。
丰富的基准测试:提供大量可重复使用的基准测试(RLLTE Hub)。
AI助手:集成了基于大型语言模型的智能助手(RLLTE Copilot)。

RLLTE项目结构图

上图展示了RLLTE的整体项目结构,包括核心组件、环境、算法、工具等各个模块。这种模块化的设计使得RLLTE具有极高的灵活性和可扩展性。

快速开始

安装

RLLTE提供了多种安装方式,推荐使用pip进行安装:

conda create -n rllte python=3.8 # 创建虚拟环境
pip install rllte-core # 基础安装
pip install rllte-core[envs] # 安装预定义环境

对于希望从源码安装的用户,可以通过git克隆仓库:

git clone https://github.com/RLE-Foundation/rllte.git
pip install -e . # 基础安装
pip install -e .[envs] # 安装预定义环境

使用内置算法进行快速训练

RLLTE提供了多种经典RL算法的实现,并提供了简单的接口用于构建应用。以下是一个使用DrQ-v2算法解决DeepMind Control Suite任务的示例:

from rllte.env import make_dmc_env 
from rllte.agent import DrQv2

if __name__ == "__main__":
    device = "cuda:0"
    env = make_dmc_env(env_id="cartpole_balance", device=device)
    eval_env = make_dmc_env(env_id="cartpole_balance", device=device)
    agent = DrQv2(env=env, eval_env=eval_env, device=device, tag="drqv2_dmc_pixel")
    agent.train(num_train_steps=500000, log_interval=1000)

运行这段代码,您将看到类似下图的训练过程:

RL训练过程

RLLTE还支持在华为NPU上进行训练,只需将device参数从"cuda:0"改为"npu:0"即可。

创建自定义RL智能体

RLLTE的设计理念之一是让开发者能够快速实现自己的RL算法。通常只需要三个步骤就可以创建一个自定义的RL智能体:

选择原型
选择必要的模块来构建智能体
合并这些模块并编写update函数

以下是一个实现Advantage Actor-Critic (A2C)算法来解决Atari游戏的示例:

from rllte.common.prototype import OnPolicyAgent
from rllte.xploit.encoder import MnihCnnEncoder
from rllte.xploit.policy import OnPolicySharedActorCritic
from rllte.xploit.storage import VanillaRolloutStorage
from rllte.xplore.distribution import Categorical
from torch import nn
import torch as th

class A2C(OnPolicyAgent):
    def __init__(self, env, tag, seed, device, num_steps) -> None:
        super().__init__(env=env, tag=tag, seed=seed, device=device, num_steps=num_steps)
        # 创建模块
        encoder = MnihCnnEncoder(observation_space=env.observation_space, feature_dim=512)
        policy = OnPolicySharedActorCritic(observation_space=env.observation_space,
                                          action_space=env.action_space,
                                          feature_dim=512,
                                          opt_class=th.optim.Adam,
                                          opt_kwargs=dict(lr=2.5e-4, eps=1e-5),
                                          init_fn="xavier_uniform"
                                          )
        storage = VanillaRolloutStorage(observation_space=env.observation_space,
                                        action_space=env.action_space,
                                        device=device,
                                        storage_size=self.num_steps,
                                        num_envs=self.num_envs,
                                        batch_size=256
                                        )
        dist = Categorical()
        # 设置所有模块
        self.set(encoder=encoder, policy=policy, storage=storage, distribution=dist)
    
    def update(self):
        for _ in range(4):
            for batch in self.storage.sample():
                # 评估采样的动作
                new_values, new_log_probs, entropy = self.policy.evaluate_actions(obs=batch.observations, actions=batch.actions)
                # 策略损失部分
                policy_loss = - (batch.adv_targ * new_log_probs).mean()
                # 价值损失部分
                value_loss = 0.5 * (new_values.flatten() - batch.returns).pow(2).mean()
                # 更新
                self.policy.optimizers['opt'].zero_grad(set_to_none=True)
                (value_loss * 0.5 + policy_loss - entropy * 0.01).backward()
                nn.utils.clip_grad_norm_(self.policy.parameters(), 0.5)
                self.policy.optimizers['opt'].step()

这个例子展示了RLLTE的模块化设计如何使得实现复杂的RL算法变得简单直观。开发者可以轻松地混合和匹配不同的组件来创建新的算法或改进现有算法。

算法解耦与模块替换

RLLTE的一大特点是允许开发者替换已实现算法中的特定模块,以进行性能比较和算法改进。例如,如果想比较不同编码器的效果,只需调用.set函数即可:

from rllte.xploit.encoder import EspeholtResidualEncoder
encoder = EspeholtResidualEncoder(...)
agent.set(encoder=encoder)

这种高度的开放性和灵活性使得RLLTE成为RL研究和开发的理想平台。

RLLTE生态系统

RLLTE不仅仅是一个算法库,它还提供了一个完整的生态系统来支持RL项目的各个方面:

Hub: 快速训练API和可重用的基准测试。
Evaluation: 用于算法评估的合理可靠指标。
Env: 打包好的环境,可快速调用。
Deployment: 便捷的模型部署API。
Pre-training: RL中的预训练方法。
Copilot: 基于大型语言模型的智能助手。

这些组件共同构成了一个强大的RL开发和研究平台,能够满足从算法设计到模型部署的各种需求。

贡献指南

RLLTE是一个开源项目,欢迎社区成员贡献代码和想法。在开始编写代码之前,请阅读CONTRIBUTING.md文件,了解贡献指南。

引用项目

如果您在研究中使用了RLLTE,请使用以下格式引用:

@article{yuan2023rllte,
  title={RLLTE: Long-Term Evolution Project of Reinforcement Learning}, 
  author={Mingqi Yuan and Zequn Zhang and Yang Xu and Shihao Luo and Bo Li and Xin Jin and Wenjun Zeng},
  year={2023},
  journal={arXiv preprint arXiv:2309.16382}
}

致谢

RLLTE项目得到了香港理工大学、东方高等研究院和FLW基金会的支持。EIAS HPC提供了GPU计算平台,而华为昇腾社区提供了NPU计算平台用于测试。

项目的部分代码借鉴或受到了其他优秀项目的启发,RLLTE团队对此表示衷心的感谢。详细致谢名单可以在ACKNOWLEDGMENT.md文件中查看。

结语

RLLTE作为一个长期演进项目,旨在推动强化学习领域的持续发展。它不仅提供了先进的算法实现,还构建了一个完整的生态系统,为研究人员和开发者提供了强大的工具和平台。通过模块化设计和灵活的接口,RLLTE使得实现和改进RL算法变得更加简单高效。无论您是RL领域的研究人员、学生还是工程师,RLLTE都能为您的工作提供有力支持。我们期待看到更多基于RLLTE的创新成果,共同推动强化学习技术的进步和应用。