模型优化入门

作者: Vincent Moens

注意

要在 notebook 中运行本教程，请在开头添加一个安装单元格，其中包含

!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 通常很简单，本教程将保持简洁。

强化学习目标函数

在强化学习中，创新通常涉及探索优化策略的新方法（即，新算法），而不是像其他领域中那样关注新架构。在 TorchRL 中，这些算法被封装在损失模块中。损失模块协调算法的各个组件，并产生一组损失值，这些损失值可以反向传播以训练相应的组件。

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

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

要构建损失函数，我们需要 actor 网络和价值网络。如果它们是按照 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)

LossModule 的输出

如您所见，我们从损失函数收到的值不是单个标量，而是一个包含多个损失的字典。

原因很简单：因为可能一次训练多个网络，并且由于某些用户可能希望在不同的步骤中分离每个模块的优化，因此 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 中的实际损失函数。

估计内存使用量： 320 MB