rllte - 强化学习研究和应用的长期演进项目

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

项目概述

RLLTE 项目受到电信领域长期演进技术（LTE标准项目）的启发，旨在为强化学习（RL）的研究和应用提供开发组件和标准。除了提供高质量的算法实现外，RLLTE 还作为开发算法的工具包，帮助研究人员和工程师更高效地进行强化学习算法的开发和应用。

主要功能和特点包括：

🧬长远演进：提供最新的算法与小技巧。
🏞️完整生态系统：集合任务设计、模型训练、评估与部署的全流程支持。
🧱模块化设计：实现RL算法的完全解耦。
🚀优化工作流：提供全面的硬件加速支持。
⚙️支持自定义环境与模块：用户可以根据需要调整使用的组件。
🖥️支持多种计算设备：如GPU与NPU。
💾大量可复用的基准测试：如RLLTE Hub。
🤖大语言模型协作工具：RLLTE Copilot的帮助。

快速开始

安装

RLLTE的安装推荐使用pip，或者可以通过git从Github克隆代码库进行安装。

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

或者使用git克隆代码库：

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

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

RLLTE提供了已被广泛认可的RL算法实现，用户可以通过简单的接口进行使用。比如，要在NVIDIA GPU上使用DrQ-v2解决DeepMind控制套件的问题，简化为以下train.py脚本：

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)

如需在HUAWEI NPU上训练，只需将device改为npu:0即可。

三步创建你的RL代理

通过RLLTE，开发人员只需经过三个简单的步骤即可实现一个RL算法。以Advantage Actor-Critic (A2C) 代理为例：

选择一个原型：

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

合并这些模块并编写.update函数：

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(...)
        storage = VanillaRolloutStorage(...)
        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()

最后，通过如下代码训练代理：

from rllte.env import make_atari_env
if __name__ == "__main__":
    device = "cuda"
    env = make_atari_env("PongNoFrameskip-v4", num_envs=8, seed=0, device=device)
    agent = A2C(env=env, tag="a2c_atari", seed=0, device=device, num_steps=128)
    agent.train(num_train_steps=10000000)

通过这些步骤，用户可以轻松地用RLLTE创建强化学习代理。

算法解耦与模块替换

RLLTE允许开发人员替换已实现算法中的指定模块，实现性能比较和算法改进。用户可以使用内置模块或自定义模块，只需调用.set功能即可完成。例如，要比较不同编码器的效果，可如下设置：

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

更多的教程和细节可以参看相关文档与指南。RLLTE是一个非常开放的框架，为开发人员提供了丰富的探索与实验空间。