使用 TorchRL 进行强化学习 (PPO) 教程 — PyTorch 教程 2.4.0+cu124 文档

在本教程中，我们将深入探讨如何使用 PyTorch 和 torchrl 库进行强化学习，特别是通过近端策略优化 (PPO) 算法解决经典倒立摆问题。借助于 PPO 这种高效、在线的策略梯度方法，它尤其适用于解决具有连续动作空间的强化学习问题。借助 torchrl 库的丰富功能，我们将构建一个完整的 PPO 实现，包括环境定义、策略与价值网络设计、数据收集与重放缓冲区管理以及训练循环的详细步骤。我们还将介绍如何评估模型性能，并最终展示如何在 Google Colab 环境中运行整个教程。通过这个教程，您将具备构建和应用强化学习模型，尤其是使用 PPO 解决复杂问题的实操能力。

强化学习是一种机器学习方法，其中智能体（agent）通过与环境交互来学习最优策略。近端策略优化（PPO）是一种高效、在线的策略梯度方法，广泛用于解决连续动作空间的强化学习问题。本教程将使用 torchrl 库，一个用于构建强化学习算法的 PyTorch 包，来实现 PPO 算法。

首先，我们需要定义环境和对观察结果进行预处理的一些转换器。我们将使用 Gym 创建我们的环境。

import torch
import gymnasium as gym
from torchrl.envs import TransformedEnv, GymEnv, ObservationNorm
from torchrl.envs.utils import DoubleToFloat

# 创建 Gym 环境
env = TransformedEnv(
    GymEnv("CartPole-v1"), 
    transforms=[ObservationNorm(), DoubleToFloat()]
)

# 环境的元数据，包括观察、动作、奖励和输入的规格
print(env.observation_spec)
print(env.action_spec)
print(env.reward_spec)
print(env.input_spec)

# 检查环境规范
check_env_specs(env)

PPO 实现的关键组件是策略 (Policy) 和价值函数 (Value Function)。我们将设计一个基于神经网络的策略，它将以连续动作空间为输出，并使用价值函数来估计长期回报。

import torch.nn as nn
from torchrl.modules import TanhNormal, ProbabilisticActor, ValueOperator, TensorDictModule

class ActorCritic(nn.Module):
    def __init__(self, state_size, action_size):
        super(ActorCritic, self).__init__()
        self.actor = nn.Sequential(
            nn.Linear(state_size, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, 2 * action_size),
        )
        self.critic = nn.Sequential(
            nn.Linear(state_size, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, 1),
        )
        self.actor_output = TanhNormal(
            loc=self.actor.output, scale=1.0, min=-1.0, max=1.0
        )
        self.critic_output = ValueOperator(self.critic.output)

    def forward(self, state):
        action_params = self.actor(state)
        action_dist = self.actor_output(action_params)
        value = self.critic_output(state)
        return action_dist, value

# 初始化策略与价值网络
state_size = env.observation_spec.shape[0]
action_size = env.action_spec.shape[0]
model = ActorCritic(state_size, action_size)

数据收集器和重放缓冲区是 PPO 过程中的关键组成部分。我们将使用 SyncDataCollector 和 ReplayBuffer。

from torchrl.collectors import SyncDataCollector
from torchrl.data.replay_buffers import ReplayBuffer
from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement

# 创建数据收集器和重放缓冲区
collector = SyncDataCollector(env, model, total_frames=50000, frames_per_batch=1000)
buffer = ReplayBuffer(storage=ReplayBuffer.Storage.LazyTensorStorage(max_size=10000),
                     sampler=SamplerWithoutReplacement())

训练循环是 PPO 实现的核心部分。我们将实现一个循环，用于收集数据、计算优势、更新策略和价值函数，并记录性能指标。

import time
from tqdm import tqdm
import torch.optim as optim
from torchrl.objectives import ClipPPOLoss
from torchrl.objectives.value import GAE

# 定义超参数
clip_eps = 0.2
gamma = 0.99
lmbda = 0.95
batch_size = 128

# 初始化损失函数和优化器
gae_module = GAE(gamma=gamma, lmbda=lmbda)
ppo_loss = ClipPPOLoss(model, gae_module, clip_eps=clip_eps, entropy_bonus=True)
optimizer = optim.Adam(model.parameters(), lr=0.001)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'max', patience=5)

# 训练循环
start_time = time.time()
episode_rewards = []
for episode in tqdm(range(100)):
    observations, _ = env.reset()
    episode_reward = 0
    while True:
        action_dist, _ = model(observations)
        action = action_dist.sample()
        observations, reward, terminated, truncated, _ = env.step(action)
        episode_reward += reward
        if terminated or truncated:
            episode_rewards.append(episode_reward)
            break

    # 更新收集的数据到重放缓冲区
    buffer.extend(observations)
    # 从重放缓冲区中采样数据用于训练
    for _ in range(100):
        batch = buffer.sample(batch_size)
        loss = ppo_loss(batch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        scheduler.step(loss)

    # 记录指标
    current_time = time.time()
    elapsed_time = current_time - start_time
    print(f"Episode: {episode}, Reward: {episode_reward}, Time: {elapsed_time:.2f}s")

通过遵循上述步骤，您已经了解了如何使用 torchrl 库构建一个 PPO 解决方案来解决倒立摆问题。本教程强调了使用 torchrl 的关键组件来实现强化学习算法的灵活性和效率。如果您希望深入探索强化学习的更多应用或理论，可以参考 PyTorch 的官方文档或相关学术文献。同时，推荐网站如慕课网等平台提供了丰富的强化学习课程资源，适合不同学习阶段的开发者。

通过实践和理论学习的结合，您将能够构建更复杂的强化学习模型，并将其应用于各种实际场景，如游戏控制、机器人控制、医疗健康等领域。强化学习领域的持续发展为解决复杂问题提供了强大的技术支撑，而 torchrl 库作为 PyTorch 的一部分，为这一过程提供了高效且易于使用的工具。