torchrl.objectives 包¶

TorchRL 提供了一系列可在训练脚本中使用的损失函数。其目标是提供易于复用/替换且签名简单的损失函数。

TorchRL 损失函数的主要特点如下：

它们是有状态对象：它们包含可训练参数的副本，因此 loss_module.parameters() 提供了训练算法所需的一切。
它们遵循 tensordict 约定：torch.nn.Module.forward() 方法将接收一个 tensordict 作为输入，其中包含返回损失值所需的所有信息。
它们输出一个 tensordict.TensorDict 实例，损失值以 "loss_<smth>" 为键写入，其中 smth 是描述损失的字符串。tensordict 中的其他键可能是训练期间有用的指标。

注意

我们返回独立损失的原因是为了让用户可以对不同的参数集使用不同的优化器。简单地将损失相加可以通过以下方式完成：

>>> loss_val = sum(loss for key, loss in loss_vals.items() if key.startswith("loss_"))

注意

损失函数中的参数初始化可以通过调用 get_stateful_net() 完成，该方法将返回网络的有状态版本，可以像其他任何模块一样初始化。如果修改是就地进行的 (in-place)，它将下游传递到使用相同参数集的任何其他模块（无论在损失函数内部还是外部）：例如，从损失函数修改 actor_network 参数也会修改收集器中的 actor。如果参数是非就地修改的 (out-of-place)，则可以使用 from_stateful_net() 将损失函数中的参数重置为新值。

torch.vmap 与随机性¶

TorchRL 损失模块大量调用 vmap()，以分摊在循环中调用多个相似模型的成本，并将其向量化。当需要在调用中生成随机数时，需要明确告诉 vmap 如何处理。为此，需要设置一个随机性模式，该模式必须是 “error”（默认，处理伪随机函数时出错）、“same”（跨批次复制结果）或 “different”（批次中的每个元素单独处理）之一。依赖默认设置通常会导致如下错误：

>>> RuntimeError: vmap: called random operation while in randomness error mode.

由于对 vmap 的调用隐藏在损失模块内部，TorchRL 提供了一个外部接口来设置 vmap 模式，通过 loss.vmap_randomness = str_value 完成，更多信息请参见 vmap_randomness()。

如果未检测到随机模块，LossModule.vmap_randomness 默认为 “error”，在其他情况下默认为 “different”。默认情况下，只有有限数量的模块被列为随机模块，但可以使用 add_random_module() 函数扩展该列表。

训练价值函数¶

TorchRL 提供了一系列价值估计器，例如 TD(0)、TD(1)、TD(\(\lambda\)) 和 GAE。简单来说，价值估计器是数据（主要是奖励和完成状态）和状态值（即，用于估计状态值的函数返回的值）的函数。要了解更多关于价值估计器的信息，请查阅 Sutton 和 Barto 的 RL 入门，特别是关于价值迭代和 TD 学习的章节。它根据数据和代理映射，对遵循某个状态或状态-动作对的折扣回报给出了一个有偏的估计。这些估计器用于两种场景：

为了训练价值网络来学习“真实”的状态值（或状态-动作值）映射，需要一个目标值来拟合。估计器越好（偏差越小，方差越小），价值网络就越好，这反过来可以显著加快策略训练。通常，价值网络的损失函数如下所示：
```
>>> value = value_network(states)
>>> target_value = value_estimator(rewards, done, value_network(next_state))
>>> value_net_loss = (value - target_value).pow(2).mean()
```
计算策略优化的“优势”信号。优势是价值估计值（来自估计器，即来自“真实”数据）与价值网络输出（即该值的代理）之间的差值。正优势可以看作是策略实际表现优于预期的信号，因此如果以该轨迹为例，则表示有改进的空间。相反，负优势表示策略表现不如预期。

事情并非总是像上面的例子那样简单，计算价值估计器或优势的公式可能比这稍微复杂一些。为了帮助用户灵活地使用不同的价值估计器，我们提供了一个简单的 API 来动态更改它。这里以 DQN 为例，但所有模块都遵循类似的结构：

>>> from torchrl.objectives import DQNLoss, ValueEstimators
>>> loss_module = DQNLoss(actor)
>>> kwargs = {"gamma": 0.9, "lmbda": 0.9}
>>> loss_module.make_value_estimator(ValueEstimators.TDLambda, **kwargs)

ValueEstimators 类枚举了可供选择的价值估计器。这使得用户可以轻松地依靠自动补全来做出选择。

LossModule(*args, **kwargs)

RL 损失函数的父类。

DQN¶

`DQNLoss`(args, *kwargs)	DQN 损失类。
`DistributionalDQNLoss`(args, *kwargs)	分布式 DQN 损失类。

DDPG¶

DDPGLoss(*args, **kwargs)

DDPG 损失类。

SAC¶

`SACLoss`(args, *kwargs)	TorchRL 实现的 SAC 损失函数。
`DiscreteSACLoss`(args, *kwargs)	离散 SAC 损失模块。

REDQ¶

REDQLoss(*args, **kwargs)

REDQ 损失模块。

CrossQ¶

CrossQLoss(*args, **kwargs)

TorchRL 实现的 CrossQ 损失函数。

IQL¶

`IQLLoss`(args, *kwargs)	TorchRL 实现的 IQL 损失函数。
`DiscreteIQLLoss`(args, *kwargs)	TorchRL 实现的离散 IQL 损失函数。

CQL¶

`CQLLoss`(args, *kwargs)	TorchRL 实现的连续 CQL 损失函数。
`DiscreteCQLLoss`(args, *kwargs)	TorchRL 实现的离散 CQL 损失函数。

GAIL¶

GAILLoss(*args, **kwargs)

TorchRL 实现的生成对抗模仿学习 (GAIL) 损失函数。

DT¶

`DTLoss`(args, *kwargs)	TorchRL 实现的在线决策 Transformer 损失函数。
`OnlineDTLoss`(args, *kwargs)	TorchRL 实现的在线决策 Transformer 损失函数。

TD3¶

TD3Loss(*args, **kwargs)

TD3 损失模块。

TD3+BC¶

TD3BCLoss(*args, **kwargs)

TD3+BC 损失模块。

PPO¶

`PPOLoss`(args, *kwargs)	PPO 损失父类。
`ClipPPOLoss`(args, *kwargs)	裁剪 PPO 损失。
`KLPENPPOLoss`(args, *kwargs)	KL 惩罚 PPO 损失。

将 PPO 与多头动作策略一起使用¶

注意

构建多头策略时要考虑的主要工具有：CompositeDistribution、ProbabilisticTensorDictModule 和 ProbabilisticTensorDictSequential。处理这些模块时，建议在脚本开头调用 tensordict.nn.set_composite_lp_aggregate(False).set()，以指示 CompositeDistribution 不应聚合对数概率，而应将其作为叶节点写入 tensordict 中。

在某些情况下，我们有一个单一的优势值，但采取了多个动作。每个动作都有自己的对数概率和形状。例如，动作空间可以结构化如下：

>>> action_td = TensorDict(
...     agents=TensorDict(
...         action0=Tensor(batch, n_agents, f0),
...         action1=Tensor(batch, n_agents, f1, f2),
...         batch_size=torch.Size((batch, n_agents))
...     ),
...     batch_size=torch.Size((batch,))
... )

其中 f0、f1 和 f2 是任意整数。

请注意，在 TorchRL 中，根 tensordict 的形状与环境的形状相同（如果环境是批量锁定的，否则其形状与正在运行的批量环境数相同）。如果 tensordict 是从缓冲区采样的，其形状将与回放缓冲区 batch_size 的形状相同。尽管 n_agent 维度对每个动作都常见，但它通常不会出现在根 tensordict 的批量大小中（尽管根据 MARL API，它会出现在包含智能体特定数据的子 tensordict 中）。

这种情况是有合理原因的：智能体数量可能会影响环境的某些规范，但不是全部。例如，有些环境在所有智能体之间共享一个完成状态。在这种情况下，一个更完整的 tensordict 可能看起来像这样：

>>> action_td = TensorDict(
...     agents=TensorDict(
...         action0=Tensor(batch, n_agents, f0),
...         action1=Tensor(batch, n_agents, f1, f2),
...         observation=Tensor(batch, n_agents, f3),
...         batch_size=torch.Size((batch, n_agents))
...     ),
...     done=Tensor(batch, 1),
...     [...] # etc
...     batch_size=torch.Size((batch,))
... )

请注意，done 状态和 reward 通常在最右侧伴随一个单例维度。请参阅文档的这部分，了解更多关于此限制的信息。

在给定其各自分布的情况下，我们的动作的对数概率可能看起来像这样：

>>> action_td = TensorDict(
...     agents=TensorDict(
...         action0_log_prob=Tensor(batch, n_agents),
...         action1_log_prob=Tensor(batch, n_agents, f1),
...         batch_size=torch.Size((batch, n_agents))
...     ),
...     batch_size=torch.Size((batch,))
... )

或

>>> action_td = TensorDict(
...     agents=TensorDict(
...         action0_log_prob=Tensor(batch, n_agents),
...         action1_log_prob=Tensor(batch, n_agents),
...         batch_size=torch.Size((batch, n_agents))
...     ),
...     batch_size=torch.Size((batch,))
... )

即，分布对数概率的维数通常从样本的维数到任何小于该维数的值不等，例如，如果分布是多元的（例如 Dirichlet）或是一个 Independent 实例。相反，tensordict 的维数仍然与环境/回放缓冲区的批量大小匹配。

在调用 PPO 损失时，损失模块将 схематически 执行以下一系列操作：

>>> def ppo(tensordict):
...     prev_log_prob = tensordict.select(*log_prob_keys)
...     action = tensordict.select(*action_keys)
...     new_log_prob = dist.log_prob(action)
...     log_weight = new_log_prob - prev_log_prob
...     advantage = tensordict.get("advantage") # computed by GAE earlier
...     # attempt to map shape
...     log_weight.batch_size = advantage.batch_size[:-1]
...     log_weight = sum(log_weight.sum(dim="feature").values(True, True)) # get a single tensor of log_weights
...     return minimum(log_weight.exp() * advantage, log_weight.exp().clamp(1-eps, 1+eps) * advantage)

要了解多头策略下的 PPO 流水线是什么样的，可以在库的示例目录中找到一个示例。

A2C¶

A2CLoss(*args, **kwargs)

TorchRL 实现的 A2C 损失函数。

Reinforce¶

ReinforceLoss(*args, **kwargs)

Reinforce 损失模块。

Dreamer¶

`DreamerActorLoss`(args, *kwargs)	Dreamer Actor 损失。
`DreamerModelLoss`(args, *kwargs)	Dreamer Model 损失。
`DreamerValueLoss`(args, *kwargs)	Dreamer Value 损失。

多智能体目标¶

这些目标是多智能体算法特有的。

QMixer¶

QMixerLoss(*args, **kwargs)

QMixer 损失类。

返回值¶

`ValueEstimatorBase`(args, *kwargs)	价值函数模块的抽象父类。
`TD0Estimator`(args, *kwargs)	时序差分 (TD(0)) 优势函数估计器。
`TD1Estimator`(args, *kwargs)	\(\infty\)-时序差分 (TD(1)) 优势函数估计器。
`TDLambdaEstimator`(args, *kwargs)	TD(\(\lambda\)) 优势函数估计器。
`GAE`(args, *kwargs)	广义优势估计函数的类封装器。
`functional.td0_return_estimate`(gamma, ...[, ...])	轨迹的 TD(0) 折扣回报估计。
`functional.td0_advantage_estimate`(gamma, ...)	轨迹的 TD(0) 优势估计。
`functional.td1_return_estimate`(gamma, ...[, ...])	TD(1) 回报估计。
`functional.vec_td1_return_estimate`(gamma, ...)	向量化 TD(1) 回报估计。
`functional.td1_advantage_estimate`(gamma, ...)	TD(1) 优势估计。
`functional.vec_td1_advantage_estimate`(gamma, ...)	向量化 TD(1) 优势估计。
`functional.td_lambda_return_estimate`(gamma, ...)	TD(\(\lambda\)) 回报估计。
`functional.vec_td_lambda_return_estimate`(...)	向量化 TD(\(\lambda\)) 回报估计。
`functional.td_lambda_advantage_estimate`(...)	TD(\(\lambda\)) 优势估计。
`functional.vec_td_lambda_advantage_estimate`(...)	向量化 TD(\(\lambda\)) 优势估计。
`functional.generalized_advantage_estimate`(...)	轨迹的广义优势估计。
`functional.vec_generalized_advantage_estimate`(...)	轨迹的向量化广义优势估计。
`functional.reward2go`(reward, done, gamma, *)	计算给定多条轨迹和情节结束的折扣累积奖励和。

工具函数¶

`HardUpdate`(loss_module, *[, ...])	用于 Double DQN/DDPG 中目标网络硬更新的类（与软更新相对）。
`SoftUpdate`(loss_module, *[, eps, tau])	用于 Double DQN/DDPG 中目标网络软更新的类。
`ValueEstimators`(value)	用于自定义估计器的价值函数枚举器。
`default_value_kwargs`(value_type)	默认价值函数关键字参数生成器。
`distance_loss`(v1, v2, loss_function[, ...])	计算两个张量之间的距离损失。
`group_optimizers`(*optimizers)	将多个优化器组合成一个。
`hold_out_net`(network)	将网络排除在计算图之外的上下文管理器。
`hold_out_params`(params)	将参数列表排除在计算图之外的上下文管理器。
`next_state_value`(tensordict[, operator, ...])	计算下一状态值（无梯度）以计算目标值。