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

在本文中，我们通过使用 PyTorch 和 TorchRL 库实现近端策略优化（PPO）算法，对倒立摆任务进行深入学习。从环境创建、模型设计、数据收集、损失函数实现、训练循环至模型评估，每一步均详尽介绍。首先，定义环境配置，创建倒立摆环境，并在 GPU 上运行以提升训练效率。接着，构建策略网络和价值网络，设计数据收集器，定义损失函数与优化器。最后，执行训练循环，评估模型性能，并讨论应用示例与结果。通过本文，您将深入了解 PPO 算法及其在强化学习领域的实际应用。

首先，我们需要准备好运行环境。我们将使用 MuJoCo 库中的倒立摆环境。确保已经安装了 gymnasium 库，并根据可用资源选择在 GPU 上运行训练。

!pip install gymnasium

定义超参数和配置环境：

import gymnasium as gymnasium

# 环境配置
device = 'cuda' if torch.cuda.is_available() else 'cpu'
total_frames = 50_000
frames_per_batch = 1000
sub_batch_size = 64
num_epochs = 10
clip_epsilon = 0.2
gamma = 0.99
lmbda = 0.95
entropy_eps = 1e-4

# 创建倒立摆环境
env = gymnasium.make('CartPole-v1', render_mode='human')
env = gymnasium.wrappers.TransformObservation(env, gymnasium.wrappers.NormalizeObservation())

接下来，我们将构建策略网络和价值网络。使用 TensorDictModule 和 ProbabilisticActor。

from torchrl.modules import (SequentialModule, TensorDictModule,
                             ProbabilisticActor, TanhNormal)

class PolicyModule(nn.Module):
    def __init__(self, input_size, output_size):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_size, 64),
            nn.ReLU(),
            nn.Linear(64, 64),
            nn.ReLU(),
            nn.Linear(64, output_size * 2)
        )
        self.distractor = NormalParamExtractor()

    def forward(self, x):
        out, distr_params = self.net(x).chunk(2, dim=-1)
        distr_params = out.view(out.size(0), 2, -1)
        return distr_params, out

policy_module = PolicyModule(env.observation_space.shape[0], env.action_space.n)

value_module = ValueOperator(
    module=nn.Sequential(
        nn.Linear(env.observation_space.shape[0], 64),
        nn.ReLU(),
        nn.Linear(64, 64),
        nn.ReLU(),
        nn.Linear(64, 1)
    ),
    in_keys=['observation']
)

actor_module = ProbabilisticActor(
    module=policy_module,
    spec=env.action_space,
    distribution_class=TanhNormal,
    distribution_kwargs={'min': -1.0, 'max': 1.0},
    return_log_prob=True
)

定义数据收集器，用于收集训练数据：

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

replay_buffer = ReplayBuffer(
    storage=LazyTensorStorage(max_size=frames_per_batch),
    sampler=SamplerWithoutReplacement()
)

collector = SyncDataCollector(
    env,
    policy_module,
    frames_per_batch=frames_per_batch,
    total_frames=total_frames,
    split_trajs=False,
    device=device
)

构建损失函数和优化器：

from torchrl.objectives import ClipPPOLoss
from torchrl.objectives.value import GAE

advantage_module = GAE(gamma=gamma, lmbda=lmbda, value_network=value_module, average_gae=True)

loss_module = ClipPPOLoss(
    actor_network=actor_module,
    critic_network=value_module,
    clip_epsilon=clip_epsilon,
    entropy_bonus=bool(entropy_eps),
    entropy_coef=entropy_eps,
    loss_critic_type="smooth_l1"
)

optimizer = torch.optim.Adam(loss_module.parameters(), lr=1e-3)

执行训练循环，并在每 10 批数据之后评估模型：

import tqdm

rewards = []
step_counts = []
lrs = []

for i in tqdm.trange(total_frames // frames_per_batch):
    tensordict_data = collector.collect()
    for _ in range(num_epochs):
        advantage_module(tensordict_data)
        data_view = tensordict_data.reshape(-1)
        replay_buffer.extend(data_view.to(device))
        for _ in range(frames_per_batch // sub_batch_size):
            subdata = replay_buffer.sample(sub_batch_size)
            loss = loss_module(subdata)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(loss_module.parameters(), 1)
            optimizer.step()
            optimizer.zero_grad()

    rewards.append(tensordict_data["next", "reward"].mean().item())
    step_counts.append(tensordict_data["step_count"].max().item())
    lrs.append(optimizer.param_groups[0]['lr'])

    # 评估模型
    if i % 10 == 0:
        with torch.no_grad():
            eval_reward, eval_step_count = evaluate_policy(actor_module, env, num_steps=1000)
        rewards.append(eval_reward)
        step_counts.append(eval_step_count)

在执行完上述步骤后，我们可以观察到模型在训练过程中的性能提升和策略的改进。结果讨论部分可以包括训练期间奖励和步数的可视化，以及在环境中的表现示例。此外，可以讨论模型的局限性、可能的优化方向以及与其他强化学习算法（如 A3C、DQN 等）的对比分析。

通过以上步骤，我们构建了一个完整的 PPO 系统，用于解决倒立摆任务，并在每一步中详细解释了每个组件的实现和功能。这不仅提供了一个实际的代码示例，还涵盖了强化学习中关键概念的深入理解，如策略优化、模型设计和训练过程优化。