DecisionTransformer

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

在线决策 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 实例，而不是直接调用此函数，因为前者负责运行已注册的钩子，而后者则会默默忽略它们。