TorchRL: 强化学习的新利器

TorchRL简介

TorchRL是一个基于PyTorch的开源强化学习库,旨在为研究人员和开发者提供高效、模块化和可扩展的工具,以简化强化学习实验和应用开发过程。作为PyTorch生态系统的一部分,TorchRL继承了PyTorch的易用性和灵活性,同时针对强化学习的特殊需求进行了优化设计。

主要特性

TorchRL具有以下几个突出的特点:

1. 以Python为中心:以Python作为主要语言,确保了易用性和灵活性。

2. 高效性能:针对强化学习研究应用进行了性能优化。

3. 模块化设计:高度模块化的架构允许轻松替换、转换或创建新组件。

4. 文档完善:详尽的文档确保用户能快速理解和使用该库。

5. 测试严格:经过严格测试,保证了可靠性和稳定性。

6. 可重用函数:提供了一系列高度可重用的函数,用于成本函数、回报计算和数据处理。

设计理念

TorchRL的设计遵循以下原则:

• 与PyTorch生态系统对齐:遵循流行PyTorch库的结构和约定。
• 最小依赖:仅依赖Python标准库、NumPy和PyTorch,可选依赖常见环境库和数据集。

TensorDict:简化强化学习代码编写

TorchRL通过TensorDict这一便捷的数据结构,极大地简化了强化学习代码的编写。TensorDict允许用户以流线型方式编写RL代码,使得跨不同设置(如在线到离线、基于状态到基于像素的学习)重用代码变得容易。

例如,使用TensorDict,我们可以用不到100行代码完成一个完整的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, 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, 128), nn.Tanh(),
      nn.Linear(128, 1),
  ),
  in_keys=["observation"],
)
actor = ProbabilisticActor(
  model,
  in_keys=["loc", "scale"],
  distribution_class=TanhNormal,
  distribution_kwargs={"low": -1.0, "high": 1.0},
  return_log_prob=True
  )
buffer = TensorDictReplayBuffer(
  storage=LazyTensorStorage(1000),
  sampler=SamplerWithoutReplacement(),
  batch_size=50,
  )
collector = SyncDataCollector(
  env,
  actor,
  frames_per_batch=1000,
  total_frames=1_000_000,
)
loss_fn = ClipPPOLoss(actor, critic)
adv_fn = GAE(value_network=critic, average_gae=True, gamma=0.99, lmbda=0.95)
optim = torch.optim.Adam(loss_fn.parameters(), lr=2e-4)

for data in collector:  # collect data
  for epoch in range(10):
      adv_fn(data)  # compute advantage
      buffer.extend(data)
      for sample in buffer:  # consume data
          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}")

这种设计使得代码重用变得简单,大大提高了开发效率。

TorchRL的主要功能

1. 环境接口

TorchRL提供了一个通用的环境接口,支持常见的库(如OpenAI Gym、DeepMind Control Suite等)和无状态执行(如基于模型的环境)。批处理环境容器允许并行执行,提高了效率。

env_make = lambda: GymEnv("Pendulum-v1", from_pixels=True)
env_parallel = ParallelEnv(4, env_make)  # 创建4个并行环境
tensordict = env_parallel.rollout(max_steps=20, policy=None)  # 随机rollout(未给定策略)
assert tensordict.shape == [4, 20]  # 4个环境,20步rollout
env_parallel.action_spec.is_in(tensordict["action"])  # 规格检查返回True

2. 数据收集器

TorchRL提供了多进程和分布式数据收集器,可以同步或异步工作。通过使用TensorDict,TorchRL的训练循环变得非常类似于监督学习中的常规训练循环:

env_make = lambda: GymEnv("Pendulum-v1", from_pixels=True)
collector = MultiaSyncDataCollector(
    [env_make, env_make],
    policy=policy,
    devices=["cuda:0", "cuda:0"],
    total_frames=10000,
    frames_per_batch=50,
    ...
)
for i, tensordict_data in enumerate(collector):
    loss = loss_module(tensordict_data)
    loss.backward()
    optim.step()
    optim.zero_grad()
    collector.update_policy_weights_()

3. 高效的回放缓冲区

TorchRL提供了高效且通用的回放缓冲区,具有模块化存储:

storage = LazyMemmapStorage(  # 内存映射(物理)存储
    cfg.buffer_size,
    scratch_dir="/tmp/"
)
buffer = TensorDictPrioritizedReplayBuffer(
    alpha=0.7,
    beta=0.5,
    collate_fn=lambda x: x,
    pin_memory=device != torch.device("cpu"),
    prefetch=10,  # 多线程采样
    storage=storage
)

对于离线强化学习,回放缓冲区还可以作为常见数据集的包装器:

from torchrl.data.replay_buffers import SamplerWithoutReplacement
from torchrl.data.datasets.d4rl import D4RLExperienceReplay
data = D4RLExperienceReplay(
    "maze2d-open-v0",
    split_trajs=True,
    batch_size=128,
    sampler=SamplerWithoutReplacement(drop_last=True),
)
for sample in data:  # 或者 sample = data.sample()
    fun(sample)

4. 环境变换

TorchRL提供了跨库的环境变换,可以在设备上以向量化方式执行,处理和准备来自环境的数据以供代理使用:

env_make = lambda: GymEnv("Pendulum-v1", from_pixels=True)
env_base = ParallelEnv(4, env_make, device="cuda:0")  # 创建4个并行环境
env = TransformedEnv(
    env_base,
    Compose(
        ToTensorImage(),
        ObservationNorm(loc=0.5, scale=1.0)),  # 一次性在设备上执行变换
)
tensordict = env.reset()
assert tensordict.device == torch.device("cuda:0")

其他变换还包括:奖励缩放(RewardScaling)、形状操作(张量连接、维度扩展等)、连续操作的连接(CatFrames)、调整大小(Resize)等。

与其他库不同,这些变换被堆叠为一个列表(而不是相互包装),这使得添加和删除变换变得容易:

env.insert_transform(0, NoopResetEnv())  # 在索引0处插入NoopResetEnv变换

5. 模型和架构

TorchRL提供了各种架构和模型(例如actor-critic):

# 创建一个nn.Module
common_module = ConvNet(
    bias_last_layer=True,
    depth=None,
    num_cells=[32, 64, 64],
    kernel_sizes=[8, 4, 3],
    strides=[4, 2, 1],
)
# 将其包装在SafeModule中,指示要读取的键和输出写入的位置
common_module = SafeModule(
    common_module,
    in_keys=["pixels"],
    out_keys=["hidden"],
)
# 将策略模块包装在NormalParamsWrapper中,使输出张量分为loc和scale,并将scale映射到正空间
policy_module = SafeModule(
    NormalParamsWrapper(
        MLP(num_cells=[64, 64], out_features=32, activation=nn.ELU)
    ),
    in_keys=["hidden"],
    out_keys=["loc", "scale"],
)
# 使用SafeProbabilisticTensorDictSequential组合SafeModule和SafeProbabilisticModule,
# 指示如何构建torch.distribution.Distribution对象以及如何使用它
policy_module = SafeProbabilisticTensorDictSequential(  # 随机策略
    policy_module,
    SafeProbabilisticModule(
        in_keys=["loc", "scale"],
        out_keys="action",
        distribution_class=TanhNormal,
    ),
)
value_module = MLP(
    num_cells=[64, 64],
    out_features=1,
    activation=nn.ELU,
)
# 将策略和价值函数包装在一个公共模块中
actor_value = ActorValueOperator(common_module, policy_module, value_module)
# 从中获取独立策略
standalone_policy = actor_value.get_policy_operator()

6. 探索包装器和模块

TorchRL提供了探索包装器和模块,可以轻松在探索和利用之间切换:

policy_explore = EGreedyWrapper(policy)
with set_exploration_type(ExplorationType.RANDOM):
    tensordict = policy_explore(tensordict)  # 将使用epsilon-greedy
with set_exploration_type(ExplorationType.DETERMINISTIC):
    tensordict = policy_explore(tensordict)  # 不会使用epsilon-greedy

7. 损失模块和回报计算

TorchRL提供了一系列高效的损失模块和高度向量化的函数用于回报和优势计算:

# 损失模块
from torchrl.objectives import DQNLoss
loss_module = DQNLoss(value_network=value_network, gamma=0.99)
tensordict = replay_buffer.sample(batch_size)
loss = loss_module(tensordict)

# 优势计算
from torchrl.objectives.value.functional import vec_td_lambda_return_estimate
advantage = vec_td_lambda_return_estimate(gamma, lmbda, next_state_value, reward, done, terminated)

8. 通用训练器类

TorchRL提供了一个通用的训练器类,可以执行上述训练循环。通过钩子机制,它还支持在任何给定时间进行任何日志记录或数据转换操作。

示例和教程

TorchRL提供了一系列示例和教程,展示了库的功能:

• DQN
• DDPG
• IQL
• CQL
• TD3
• TD3+BC
• A2C
• PPO
• SAC
• REDQ
• Dreamer
• Decision Transformers
• RLHF

这些示例涵盖了强化学习中的多种算法和应用场景,为用户提供了丰富的参考。

安装

TorchRL的安装非常简单,可以通过pip直接安装最新的稳定版本:

pip install torchrl

对于想要使用最新功能的用户,可以安装每晚构建版本:

pip3 install torchrl-nightly

TorchRL还提供了额外的依赖项,可以根据需要安装:

pip3 install "torchrl[atari,dm_control,gym_continuous,rendering,tests,utils,marl,checkpointing]"

结语

TorchRL作为一个强大而灵活的强化学习库,为研究人员和开发者提供了丰富的工具和资源。通过其模块化设计、高效性能和与PyTorch生态系统的紧密集成,TorchRL大大简化了强化学习实验和应用开发的过程。无论是初学者还是经验丰富的研究人员,都能在TorchRL中找到适合自己需求的工具和功能。