DecisionTransformer

class torchrl.modules.DecisionTransformer(state_dim, action_dim, config: dict | DTConfig = None, device: torch.device | None = None)

線上 Decision Transformer。

描述於 https://arxiv.org/abs/2202.05607 。

如果使用者沒有提供特定的配置，transformer 會使用預設配置來建立 GPT2 模型。 default_config = { … “n_embd”: 256, … “n_layer”: 4, … “n_head”: 4, … “n_inner”: 1024, … “activation”: “relu”, … “n_positions”: 1024, … “resid_pdrop”: 0.1, … “attn_pdrop”: 0.1, }

引數:

• state_dim (int) – 狀態空間的維度

• action_dim (int) – 動作空間的維度

• config (DTConfig 或 dict，可選) – transformer 架構配置，用於從 transformers 建立 GPT2Config。預設為 default_config。

示例

>>> config = DecisionTransformer.default_config()
>>> config.n_embd = 128
>>> print(config)
DTConfig(n_embd: 128, n_layer: 4, n_head: 4, n_inner: 1024, activation: relu, n_positions: 1024, resid_pdrop: 0.1, attn_pdrop: 0.1)
>>> # alternatively
>>> config = DecisionTransformer.DTConfig(n_embd=128)
>>> model = DecisionTransformer(state_dim=4, action_dim=2, config=config)
>>> batch_size = [3, 32]
>>> length = 10
>>> observation = torch.randn(*batch_size, length, 4)
>>> action = torch.randn(*batch_size, length, 2)
>>> return_to_go = torch.randn(*batch_size, length, 1)
>>> output = model(observation, action, return_to_go)
>>> output.shape
torch.Size([3, 32, 10, 128])

class DTConfig(n_embd: Any = 256, n_layer: Any = 4, n_head: Any = 4, n_inner: Any = 1024, activation: Any = 'relu', n_positions: Any = 1024, resid_pdrop: Any = 0.1, attn_pdrop: Any = 0.1)

DecisionTransformer 的預設配置。

forward(observation: Tensor, action: Tensor, return_to_go: Tensor)

定義每次呼叫時執行的計算。

所有子類都應重寫此方法。

注意

儘管前向傳播的實現需要在此函式中定義，但您應該在之後呼叫 Module 例項而不是此函式，因為前者會處理註冊的鉤子，而後者則會靜默忽略它們。