DecisionTransformer

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

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