GAE

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

廣義優勢估計函式的類包裝器。

Refer to “HIGH-DIMENSIONAL CONTINUOUS CONTROL USING GENERALIZED ADVANTAGE ESTIMATION” https://arxiv.org/pdf/1506.02438.pdf for more context.

引數:

• gamma (scalar) – exponential mean discount.

• lmbda (scalar) – trajectory discount.

• value_network (TensorDictModule, optional) – 用於檢索值估計的值運算子。如果為 None，此模組將期望 "state_value" 鍵已填充，並且不會呼叫值網路來生成它。

• average_gae (bool) – 如果為 True，則結果的 GAE 值將進行標準化。預設為 False。

• differentiable (bool, optional) –

  if True, gradients are propagated through the computation of the value function. Default is False.

  注意

  The proper way to make the function call non-differentiable is to decorate it in a torch.no_grad() context manager/decorator or pass detached parameters for functional modules.

• vectorized (bool, optional) – 是否使用 lambda 返回值的向量化版本。如果未編譯，則預設為 True。

• skip_existing (bool, optional) – 如果為 True，則值網路將跳過其輸出已存在於 tensordict 中的模組。預設為 None，即 tensordict.nn.skip_existing() 的值不受影響。預設為“state_value”。

• advantage_key (str or tuple of str, optional) – [Deprecated] the key of the advantage entry. Defaults to "advantage".

• value_target_key (str or tuple of str, optional) – [已棄用] 優勢項的鍵。預設為 "value_target"。

• value_key (str or tuple of str, optional) – [已棄用] 從輸入 tensordict 讀取的值鍵。預設為 "state_value"。

• shifted (bool, optional) – 如果設定為 True，值和下一個值將透過對值網路的單次呼叫來估計。這更快，但僅在以下情況下有效：(1) "next" 值僅偏移一步（例如，對於多步值估計則不適用），並且 (2) 在時間 t 和 t+1 使用的引數相同（在使用目標引數時則不適用）。預設為 False。

• device (torch.device, optional) – 緩衝區將被例項化的裝置。預設為 torch.get_default_device()。

• time_dim (int, optional) – 輸入 tensordict 中對應時間的維度。如果未提供，則預設為標記有 "time" 名稱的維度（如果存在），否則預設為最後一個維度。可以在呼叫 value_estimate() 時覆蓋。負維度相對於輸入 tensordict 進行考慮。

• auto_reset_env (bool, optional) – 如果為 True，則該回合的最後一個 "next" 狀態無效，因此 GAE 計算將使用 value 而不是 next_value 來引導截斷的回合。

• deactivate_vmap (bool, optional) – 如果為 True，則不會使用 vmap 呼叫，向量化對映將替換為簡單的 for 迴圈。預設為 False。

GAE 將返回一個包含優勢值的 "advantage" 條目。它還將返回一個 "value_target" 條目，其中包含用於訓練值網路的返回。最後，如果 gradient_mode 為 True，則將返回一個額外的、可微分的 "value_error" 條目，它簡單地表示返回與值網路輸出之間的差值（即，應將額外的距離損失應用於此帶符號值）。

注意

與其他優勢函式一樣，如果 value_key 已存在於輸入 tensordict 中，則 GAE 模組將忽略對值網路的呼叫（如果有），而是使用提供的值。

注意

GAE 可以與依賴於迴圈神經網路的值網路一起使用，前提是初始化標記（“is_init”）和終止/截斷標記已正確設定。如果 shifted=True，則軌跡批次將被展平，並且每個軌跡的最後一步將放置在扁平 tensordict 中，位於根的最後一步之後，以便每個軌跡有 T+1 個元素。如果 shifted=False，則將堆疊根和 “next” 軌跡，並且值網路將使用 vmap 呼叫堆疊的軌跡。由於 RNN 需要相當多的控制流，因此它們目前不與 torch.vmap 相容，因此在這些情況下必須啟用 deactivate_vmap 選項。同樣，如果 shifted=False，則根 tensordict 的 “is_init” 條目將被複制到 “next” 條目的 “is_init” 上，以便根和 “next” 資料都能得到良好的分隔。

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

根據 tensordict 中的資料計算 GAE。

If a functional module is provided, a nested TensorDict containing the parameters (and if relevant the target parameters) can be passed to the module.

引數:

tensordict (TensorDictBase) – 一個 TensorDict，包含計算值估計和 GAE 所需的資料（一個觀察鍵、"action"、("next", "reward")、("next", "done")、("next", "terminated") 和 "next" tensordict 狀態，如環境返回）。傳遞給此模組的資料應結構化為 [*B, T, *F]，其中 B 是批次大小，T 是時間維度，F 是特徵維度。tensordict 的形狀必須為 [*B, T]。

關鍵字引數:

• params (TensorDictBase, optional) – A nested TensorDict containing the params to be passed to the functional value network module.

• target_params (TensorDictBase, optional) – A nested TensorDict containing the target params to be passed to the functional value network module.

• time_dim (int, optional) – 輸入 tensordict 中對應時間的維度。如果未提供，則預設為標記有 "time" 名稱的維度（如果存在），否則預設為最後一個維度。負維度相對於輸入 tensordict 進行考慮。

返回:

An updated TensorDict with an advantage and a value_error keys as defined in the constructor.

示例

>>> 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()

The module supports non-tensordict (i.e. unpacked tensordict) inputs too

示例

>>> 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)

Gets a value estimate, usually used as a target value for the value network.

如果狀態值鍵存在於 tensordict.get(("next", self.tensor_keys.value)) 下，則將使用此值，而無需呼叫值網路。

引數:

• tensordict (TensorDictBase) – the tensordict containing the data to read.

• target_params (TensorDictBase, optional) – A nested TensorDict containing the target params to be passed to the functional value network module.

• next_value (torch.Tensor, optional) – 下一個狀態或狀態-動作對的值。與 target_params 互斥。

• **kwargs – the keyword arguments to be passed to the value network.

Returns: a tensor corresponding to the state value.