TD0Estimator

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

时序差分 (TD(0)) 对优势函数进行估计。

也称为自举时序差分或 1 步回报。

关键字参数:

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

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

• shifted (bool, optional) – 如果为 True，则值和下一个值将通过一次值网络调用进行估计。这速度更快，但仅在 (1) “下一个”值仅由一个时间步移位（例如，多步值估计时的情况并非如此）并且 (2) 在时间 t 和 t+1 使用的参数相同时（当使用目标参数时并非如此）才有效。默认为 False。

• average_rewards (bool, 可选) – 如果为 True，则在计算 TD 之前会对奖励进行标准化。

• differentiable (bool, optional) –

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

  注意

  使函数调用不可导的正确方法是使用 torch.no_grad() 上下文管理器/装饰器对其进行装饰，或为函数模块传递分离的参数。

• skip_existing (bool, optional) – 如果为 True，值网络将跳过输出已存在于 tensordict 中的模块。默认为 None，即 tensordict.nn.skip_existing() 的值不受影响。

• advantage_key (str 或 str 的元组, optional) – [已弃用] 优势条目的键。默认为 "advantage"。

• value_target_key (str 或 str 的元组, optional) – [已弃用] 优势条目的键。默认为 "value_target"。

• value_key (str 或 str 的元组, optional) – [已弃用] 要从输入 tensordict 读取的值键。默认为 "state_value"。

• device (torch.device, optional) – 将实例化缓冲区的位置。默认为 torch.get_default_device()。

• deactivate_vmap (bool, 可选) – 是否禁用 vmap 调用并用普通 for 循环替换它们。默认为 False。

forward(tensordict=None, *, params: TensorDictBase | None = None, target_params: TensorDictBase | None = None)

给定 tensordict 中的数据，计算 TD(0) 优势。

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

参数:

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

关键字参数:

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

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

返回:

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

示例

>>> from tensordict import TensorDict
>>> value_net = TensorDictModule(
...     nn.Linear(3, 1), in_keys=["obs"], out_keys=["state_value"]
... )
>>> module = TDEstimate(
...     gamma=0.98,
...     value_network=value_net,
... )
>>> 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, "terminated": terminated, "reward": reward}}, [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 = TDEstimate(
...     gamma=0.98,
...     value_network=value_net,
... )
>>> 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, target_params: TensorDictBase | None = None, next_value: torch.Tensor | None = None, **kwargs)

获取值估计，通常用作值网络的 target value。

如果状态值键存在于 tensordict.get(("next", self.tensor_keys.value)) 下，则将使用此值，而无需调用值网络。

参数:

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

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

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

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

Returns: 对应于状态值的张量。