快捷方式

torchrl.objectives 包

TorchRL 提供了一系列可在训练脚本中使用的损失。目的是让损失易于重用/互换,并且具有简单的签名。

TorchRL 损失的主要特点是

  • 它们是状态对象:它们包含可训练参数的副本,因此 loss_module.parameters() 提供训练算法所需的一切。

  • 它们遵循 tensordict 约定:torch.nn.Module.forward() 方法将接收一个 tensordict 作为输入,该 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 提供了一个接口,允许从外部设置该 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 and 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 进行多头动作策略

注意

构建多头策略时要考虑的主要工具是: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 中(尽管它根据 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 的维度仍然匹配 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)

\(\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, *)

计算给定多个轨迹和回合结束的折扣累积奖励。

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 开发者文档

查看文档

教程

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

查看教程

资源

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

查看资源