快捷方式

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() 来初始化损失中的参数,它将返回一个可像任何其他模块一样初始化的网络的状态化版本。如果修改是就地进行的,它将向下传递给使用相同参数集的任何其他模块(在损失内部和外部):例如,从损失中修改 actor_network 参数也将修改收集器中的 actor。如果参数是就地修改的,则可以使用 from_stateful_net() 将损失中的参数重置为新值。

torch.vmap 和随机性

TorchRL 损失模块有很多对 vmap() 的调用,以摊销在循环中调用多个类似模型的成本,而是将这些操作向量化。 vmap 需要被明确告知如何在生成随机数时进行处理。为此,需要设置随机模式,并且该模式必须是以下之一:“error”(默认,处理伪随机函数时出错)、“same”(在批次之间复制结果)或“different”(批次的每个元素被单独处理)。依赖默认设置通常会导致此类错误:

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

由于 vmap 的调用埋藏在损失模块的深处,TorchRL 提供了一个接口,可以通过 loss.vmap_randomness = str_value 从外部设置该 vmap 模式,有关更多信息,请参阅 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 损失的父类。

add_random_module(module)

将一个随机模块添加到将被 vmap_randomness() 检测为随机的模块列表中。

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 和多头动作策略

注意

构建多头策略时要考虑的主要工具有:CompositeDistributionProbabilisticTensorDictModuleProbabilisticTensorDictSequential。在处理这些时,建议在脚本开头调用 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,))
... )

其中 f0f1f2 是任意整数。

请注意,在 TorchRL 中,根 tensordict 的形状与环境相同(如果环境是批次锁定的,则与运行的批次环境数量相同)。如果 tensordict 是从缓冲区采样的,它还将具有重放缓冲区 batch_size 的形状。尽管 n_agent 维度对于每个动作都通用,但它通常不会出现在根 tensordict 的 batch-size 中(尽管它出现在包含代理特定数据的子 tensordict 中,根据 MARL API)。

这有一个合乎逻辑的原因:代理的数量可能条件了环境的某些但不是全部规范。例如,某些环境对所有代理共享一个完成状态。在这种情况下,一个更完整的 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 的维度仍然与 env 的/重放缓冲区 的 batch-size 匹配。

在 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 损失类。

Returns

ValueEstimatorBase(*args, **kwargs)

价值函数模块的抽象父类。

TD0Estimator(*args, **kwargs)

优势函数的时序差分 (TD(0)) 估计。

TD1Estimator(*args, **kwargs)

优势函数的无穷时序差分 (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, *)

计算给定多个轨迹和情节结束的折扣累计奖励总和。

Utils

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, ...])

计算下一个状态值(无梯度),以计算目标值。

文档

访问全面的 PyTorch 开发者文档

查看文档

教程

为初学者和高级开发者提供深入的教程

查看教程

资源

查找开发资源并让您的问题得到解答

查看资源