开始模型优化¶

注意

如需在笔记本中运行本教程，请在开头添加一个安装单元，内容为：

!pip install tensordict
!pip install torchrl

在 TorchRL 中，我们尝试以 PyTorch 中惯用的方式处理优化，即使用专门设计的损失模块，其唯一目的就是优化模型。该方法能高效地将策略执行与其训练解耦，使我们能够设计出与传统监督学习示例中常见的训练循环相似的训练流程。

因此，典型的训练循环看起来像这样：

>>> for i in range(n_collections):
...     data = get_next_batch(env, policy)
...     for j in range(n_optim):
...         loss = loss_fn(data)
...         loss.backward()
...         optim.step()

在本简明教程中，您将快速了解损失模块。由于基础用法的 API 通常较为直观，因此本教程内容简洁。

RL目标函数¶

在强化学习（RL）中，创新通常体现为探索用于优化策略（即新算法）的全新方法，而非像其他领域那样聚焦于新架构。在 TorchRL 中，这些算法被封装在损失模块（loss modules）中。损失模块负责协调算法的各个组件，并生成一组可用于反向传播以训练相应组件的损失值。

在这个教程中，我们将以一个流行的离策略算法为例， DDPG。

要构建一个损失模块，唯一需要的是一组定义为 :class:`~tensordict.nn.TensorDictModule`s 的网络。大多数情况下，其中一个模块将是策略。可能还需要其他辅助网络，例如 Q-Value 网络或某种形式的评估器。让我们看看这在实践中是什么样子：DDPG 需要一个从观察空间到动作空间的确定性映射，以及一个预测状态-动作对价值的价值网络。DDPG 损失将尝试找到输出最大化给定状态下价值的动作的策略参数。

要构建损失函数，我们需要同时使用策略网络（actor）和价值网络（value）。如果它们是按照DDPG的要求构建的，那么这就足以获得一个可训练的损失模块：

from torchrl.envs import GymEnv

env = GymEnv("Pendulum-v1")

from torchrl.modules import Actor, MLP, ValueOperator
from torchrl.objectives import DDPGLoss

n_obs = env.observation_spec["observation"].shape[-1]
n_act = env.action_spec.shape[-1]
actor = Actor(MLP(in_features=n_obs, out_features=n_act, num_cells=[32, 32]))
value_net = ValueOperator(
    MLP(in_features=n_obs + n_act, out_features=1, num_cells=[32, 32]),
    in_keys=["observation", "action"],
)

ddpg_loss = DDPGLoss(actor_network=actor, value_network=value_net)

就这样！我们的损失模块现在可以使用来自环境的数据运行了（我们省略了探索、存储及其他功能，以专注于损失功能）：

rollout = env.rollout(max_steps=100, policy=actor)
loss_vals = ddpg_loss(rollout)
print(loss_vals)

TensorDict(
    fields={
        loss_actor: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        loss_value: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        pred_value: Tensor(shape=torch.Size([100]), device=cpu, dtype=torch.float32, is_shared=False),
        pred_value_max: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        target_value: Tensor(shape=torch.Size([100]), device=cpu, dtype=torch.float32, is_shared=False),
        target_value_max: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        td_error: Tensor(shape=torch.Size([100]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

损失模块的输出¶

如您所见，我们从损失函数中获得的值并非单个标量，而是一个包含多个损失项的字典。

原因很简单：因为可能同时训练多个网络，而且部分用户希望在不同的步骤中分别优化各个模块，因此 TorchRL 的目标函数将返回包含各种损失项的字典。

这种格式还允许我们将元数据与损失值一起传递。通常，我们确保只有损失值是可微的，以便您可以简单地对字典中的值求和以获得总损失。如果您希望完全控制正在发生的事情，可以仅对键以 "loss_" 前缀开头的条目求和：

total_loss = 0
for key, val in loss_vals.items():
    if key.startswith("loss_"):
        total_loss += val

训练一个LossModule¶

鉴于所有这些，训练模块与在任何其他训练循环中所做的并没有太大不同。因为它封装了模块，获取可训练参数列表的最简单方法是查询 parameters() 方法。

我们需要一个优化器（或者，如果您选择的话，每个模块各配一个优化器）。

from torch.optim import Adam

optim = Adam(ddpg_loss.parameters())
total_loss.backward()

以下各项通常会出现在您的训练循环中：

optim.step()
optim.zero_grad()

进一步考虑：目标参数¶

另一个需要考虑的重要方面是目标参数在离策略算法（如DDPG）中的存在。目标参数通常代表随时间延迟或平滑的参数版本，在策略训练期间的价值估计中起着至关重要的作用。利用目标参数进行策略训练通常比使用当前价值网络参数配置更为高效。一般来说，目标参数的管理由损失模块处理，用户无需直接关注。然而，根据具体需求更新这些值仍然是用户的职责。TorchRL提供了几个更新器，即 HardUpdate 和 SoftUpdate, 可以轻松实例化而无需深入了解损失模块的底层机制。

from torchrl.objectives import SoftUpdate

updater = SoftUpdate(ddpg_loss, eps=0.99)

在训练循环中，您需要在每次优化步骤或每次采集步骤时更新目标参数：

updater.step()

这就是开始使用损失模块所需了解的一切！

要进一步探讨这个主题，请参阅：

损失函数模块参考页；
编写DDPG损失函数教程；
损失在 PPO 或 DQN 中的应用。

脚本总运行时间： (0 分钟 26.921 秒)

估计内存使用量： 320 MB

通过 Sphinx-Gallery 生成的画廊