ValueOperator

class torchrl.modules.tensordict_module.ValueOperator(*args, **kwargs)

RL 中價值函數的通用類別。

ValueOperator 類別帶有 in_keys 和 out_keys 參數的預設值（分別為 [“observation”] 和 [“state_value”] 或 [“state_action_value”]，取決於 “action” 鍵是否為 in_keys 列表的一部分）。

參數:

• module (nn.Module) – 用於將輸入映射到輸出參數空間的 torch.nn.Module。

• in_keys (字串的可迭代物件, 選用) – 要從輸入 tensordict 讀取並傳遞到模組的鍵。如果它包含多個元素，則值將按照 in_keys 可迭代物件給定的順序傳遞。預設為 ["observation"]。

• out_keys (字串的可迭代物件) – 要寫入輸入 tensordict 的鍵。out_keys 的長度必須與嵌入式模組傳回的張量數量相符。使用 “_” 作為鍵可避免將張量寫入輸出。預設為 ["state_value"] 或 ["state_action_value"]，如果 "action" 是 in_keys 的一部分。

範例

>>> import torch
>>> from tensordict import TensorDict
>>> from torch import nn
>>> from torchrl.data import Unbounded
>>> from torchrl.modules import ValueOperator
>>> td = TensorDict({"observation": torch.randn(3, 4), "action": torch.randn(3, 2)}, [3,])
>>> class CustomModule(nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.linear = torch.nn.Linear(6, 1)
...     def forward(self, obs, action):
...         return self.linear(torch.cat([obs, action], -1))
>>> module = CustomModule()
>>> td_module = ValueOperator(
...    in_keys=["observation", "action"], module=module
... )
>>> td = td_module(td)
>>> print(td)
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([3, 2]), device=cpu, dtype=torch.float32, is_shared=False),
        observation: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        state_action_value: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)