RL学习资料汇总 - 开源强化学习库TorchRL

TorchRL学习资源汇总 - 开源强化学习库

TorchRL是一个基于PyTorch的开源强化学习(RL)库,为RL研究和应用提供了丰富的工具和功能。本文将为您介绍TorchRL的主要特性,并汇总相关学习资源,帮助您快速掌握这个强大的RL工具。

TorchRL简介

TorchRL是一个高度模块化、高效、易于扩展的RL库,具有以下主要特点:

🐍 Python优先:以Python作为主要开发语言,使用简单灵活
⏱️ 高效:针对性能进行了优化,可支持大规模RL研究
🧮 模块化:高度模块化的架构,便于替换、转换或创建新组件
📚 文档完善:详细的文档确保用户可以快速理解和使用
✅ 测试严格:经过严格测试以确保可靠性和稳定性
⚙️ 可重用函数:提供一组高度可重用的损失函数、回报计算和数据处理函数

TorchRL架构图

核心功能

TorchRL提供了丰富的RL相关功能,包括:

通用环境接口:支持OpenAI Gym、DeepMind Control Suite等常见库,可进行并行执行。
数据收集器:支持多进程和分布式数据收集,可同步或异步工作。
高效的回放缓冲区:支持优先级回放等高级功能。
环境变换:提供跨库的环境变换,可在设备上以向量化方式执行。
分布式学习工具:如内存映射张量等。
丰富的模型架构:如Actor-Critic网络等。
探索包装器:便于在探索和利用之间切换。
损失模块和向量化回报计算。
通用训练器类:执行训练循环,支持日志记录和数据转换。

快速上手

以下是一个使用TorchRL实现PPO算法的简单示例:

import torch
from tensordict.nn import TensorDictModule
from tensordict.nn.distributions import NormalParamExtractor
from torch import nn

from torchrl.collectors import SyncDataCollector
from torchrl.data.replay_buffers import TensorDictReplayBuffer, LazyTensorStorage, SamplerWithoutReplacement
from torchrl.envs.libs.gym import GymEnv
from torchrl.modules import ProbabilisticActor, ValueOperator, TanhNormal
from torchrl.objectives import ClipPPOLoss
from torchrl.objectives.value import GAE

# 创建环境
env = GymEnv("Pendulum-v1")

# 定义策略网络
model = TensorDictModule(
  nn.Sequential(
      nn.Linear(3, 128), nn.Tanh(),
      nn.Linear(128, 128), nn.Tanh(),
      nn.Linear(128, 2),
      NormalParamExtractor()
  ),
  in_keys=["observation"],
  out_keys=["loc", "scale"]
)

# 定义价值网络
critic = ValueOperator(
  nn.Sequential(
      nn.Linear(3, 128), nn.Tanh(),
      nn.Linear(128, 128), nn.Tanh(),
      nn.Linear(128, 1),
  ),
  in_keys=["observation"],
)

# 定义Actor
actor = ProbabilisticActor(
  model,
  distribution_class=TanhNormal,
  distribution_kwargs={"low": -1.0, "high": 1.0},
  return_log_prob=True
)

# 创建数据收集器
collector = SyncDataCollector(
  env,
  actor,
  frames_per_batch=1000,
  total_frames=1_000_000,
)

# 定义PPO损失
loss_fn = ClipPPOLoss(actor, critic)
adv_fn = GAE(value_network=critic, gamma=0.99, lmbda=0.95)
optim = torch.optim.Adam(loss_fn.parameters(), lr=2e-4)

# 训练循环
for data in collector:
  for epoch in range(10):
      adv_fn(data)
      for sample in buffer:
          loss_vals = loss_fn(sample)
          loss_val = sum(value for key, value in loss_vals.items() if key.startswith("loss"))
          loss_val.backward()
          optim.step()
          optim.zero_grad()
  print(f"avg reward: {data['next', 'reward'].mean().item(): 4.4f}")