GAE

class torchrl.objectives.value.GAE(*args, **kwargs)

围绕广义优势估计函数的类包装器。

有关更多上下文，请参阅“使用广义优势估计的高维连续控制” https://arxiv.org/pdf/1506.02438.pdf。

参数:

• gamma (标量) – 指数平均折扣。

• lmbda (标量) – 轨迹折扣。

• value_network (TensorDictModule) – 用于检索价值估计的价值算子。

• average_gae (布尔值) – 如果 True，则结果 GAE 值将被标准化。默认值为 False。

• differentiable (布尔值, 可选) –

  如果 True，则梯度将通过价值函数的计算进行传播。默认值为 False。

  注意

  使函数调用不可微分的方式是将其在 torch.no_grad() 上下文管理器/装饰器中装饰，或为函数式模块传递分离的参数。

• vectorized (布尔值, 可选) – 是否使用 lambda 返回的矢量化版本。默认值为 True。

• skip_existing (布尔值, 可选) – 如果 True，则价值网络将跳过输出已存在于 tensordict 中的模块。默认为 None，即，tensordict.nn.skip_existing() 的值不受影响。默认为“state_value”。

• advantage_key (字符串 或 字符串元组, 可选) – [已弃用] 优势项的键。默认为 "advantage"。

• value_target_key (字符串 或 字符串元组, 可选) – [已弃用] 优势项的键。默认为 "value_target"。

• value_key (字符串 或 字符串元组, 可选) – [已弃用] 要从输入 tensordict 读取的价值键。默认为 "state_value"。

• shifted (布尔值, 可选) – 如果 True，则通过对价值网络的一次调用来估计价值和下一个价值。这更快，但仅在 (1) "next" 值仅按一个时间步移位（例如，在多步价值估计中并非如此）并且 (2) 在时间 t 和 t+1 使用的参数相同（当要使用目标参数时并非如此）的情况下才有效。默认为 False。

• device (torch.device, 可选) – 模块的设备。

• time_dim (整数, 可选) – 输入 tensordict 中对应时间的维度。如果未提供，则默认为带有 "time" 名称的维度（如果有），否则默认为最后一个维度。可以在调用 value_estimate() 期间覆盖。相对于输入 tensordict 的负维度将被考虑。

GAE 将返回一个包含优势值的 "advantage" 项。它还将返回一个 "value_target" 项，其中包含用于训练价值网络的回报值。最后，如果 gradient_mode 为 True，则将返回另一个可微分的 "value_error" 项，它仅表示回报与价值网络输出之间的差异（即，应将额外的距离损失应用于该有符号值）。

注意

与其他优势函数一样，如果 value_key 已经存在于输入 tensordict 中，则 GAE 模块将忽略对价值网络的调用（如果有）并使用提供的价值。

forward(tensordict: TensorDictBase, *, params: List[Tensor] | None = None, target_params: List[Tensor] | None = None, time_dim: int | None = None) → TensorDictBase

计算给定 tensordict 中数据的 GAE。

如果提供了函数式模块，则可以将包含参数（以及相关目标参数）的嵌套 TensorDict 传递给该模块。

参数:

tensordict (TensorDictBase) – 包含数据的 TensorDict（一个观察键、"action"、("next", "reward")、("next", "done")、("next", "terminated") 和 "next" tensordict 状态，由环境返回），用于计算价值估计和 GAE。传递给该模块的数据应结构化为 [*B, T, *F]，其中 B 是批次大小、T 是时间维度，而 F 是特征维度。tensordict 必须具有 [*B, T] 的形状。

关键字参数:

• params (TensorDictBase, optional) – 一个嵌套的 TensorDict，包含要传递给函数值网络模块的参数。

• target_params (TensorDictBase, optional) – 一个嵌套的 TensorDict，包含要传递给函数值网络模块的目标参数。

• time_dim (int, optional) – 输入 tensordict 中对应于时间的维度。如果未提供，则默认为带有 "time" 名称的维度（如果有），否则默认为最后一个维度。负维度相对于输入 tensordict 考虑。

返回值:

一个更新的 TensorDict，其中包含在构造函数中定义的优势和 value_error 键。

示例

>>> from tensordict import TensorDict
>>> value_net = TensorDictModule(
...     nn.Linear(3, 1), in_keys=["obs"], out_keys=["state_value"]
... )
>>> module = GAE(
...     gamma=0.98,
...     lmbda=0.94,
...     value_network=value_net,
...     differentiable=False,
... )
>>> obs, next_obs = torch.randn(2, 1, 10, 3)
>>> reward = torch.randn(1, 10, 1)
>>> done = torch.zeros(1, 10, 1, dtype=torch.bool)
>>> terminated = torch.zeros(1, 10, 1, dtype=torch.bool)
>>> tensordict = TensorDict({"obs": obs, "next": {"obs": next_obs}, "done": done, "reward": reward, "terminated": terminated}, [1, 10])
>>> _ = module(tensordict)
>>> assert "advantage" in tensordict.keys()

该模块还支持非 tensordict（即解包的 tensordict）输入

示例

>>> value_net = TensorDictModule(
...     nn.Linear(3, 1), in_keys=["obs"], out_keys=["state_value"]
... )
>>> module = GAE(
...     gamma=0.98,
...     lmbda=0.94,
...     value_network=value_net,
...     differentiable=False,
... )
>>> obs, next_obs = torch.randn(2, 1, 10, 3)
>>> reward = torch.randn(1, 10, 1)
>>> done = torch.zeros(1, 10, 1, dtype=torch.bool)
>>> terminated = torch.zeros(1, 10, 1, dtype=torch.bool)
>>> advantage, value_target = module(obs=obs, next_reward=reward, next_done=done, next_obs=next_obs, next_terminated=terminated)

value_estimate(tensordict, params: TensorDictBase | None = None, target_params: TensorDictBase | None = None, time_dim: int | None = None, **kwargs)

获取价值估计，通常用作价值网络的目标值。

如果状态值键存在于 tensordict.get(("next", self.tensor_keys.value)) 下，则将使用此值，而不会求助于价值网络。

参数:

• tensordict (TensorDictBase) – 包含要读取数据的 tensordict。

• target_params (TensorDictBase, optional) – 一个嵌套的 TensorDict，包含要传递给函数值网络模块的目标参数。

• next_value (torch.Tensor, optional) – 下一个状态或状态-动作对的值。与 target_params 互斥。

• **kwargs – 要传递给价值网络的关键字参数。

返回值：一个对应于状态值的张量。