Decision Transformer

概述

Decision Transformer 模型在 Decision Transformer: Reinforcement Learning via Sequence Modeling 中被提出
作者为 Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, Igor Mordatch。

论文的摘要如下

我们介绍了一个框架，将强化学习 (RL) 抽象为一个序列建模问题。这使我们能够利用 Transformer 架构的简洁性和可扩展性，以及语言建模方面的相关进展，如 GPT-x 和 BERT。特别是，我们提出了 Decision Transformer，一种将 RL 问题转化为条件序列建模的架构。与先前拟合价值函数或计算策略梯度的 RL 方法不同，Decision Transformer 仅通过利用因果掩码 Transformer 输出最优动作。通过将自回归模型调节到期望的回报（奖励）、过去的状态和动作，我们的 Decision Transformer 模型可以生成实现期望回报的未来动作。尽管 Decision Transformer 非常简单，但它在 Atari、OpenAI Gym 和 Key-to-Door 任务上匹配或超过了最先进的无模型离线 RL 基线的性能。

此模型版本适用于状态为向量的任务。

此模型由 edbeeching 贡献。原始代码可以在这里找到。

DecisionTransformerConfig

class transformers.DecisionTransformerConfig

< source >

( state_dim = 17 act_dim = 4 hidden_size = 128 max_ep_len = 4096 action_tanh = True vocab_size = 1 n_positions = 1024 n_layer = 3 n_head = 1 n_inner = None activation_function = 'relu' resid_pdrop = 0.1 embd_pdrop = 0.1 attn_pdrop = 0.1 layer_norm_epsilon = 1e-05 initializer_range = 0.02 scale_attn_weights = True use_cache = True bos_token_id = 50256 eos_token_id = 50256 scale_attn_by_inverse_layer_idx = False reorder_and_upcast_attn = False **kwargs )

参数

state_dim (int, 可选, 默认为 17) — RL 环境的状态大小
act_dim (int, 可选, 默认为 4) — 输出动作空间的大小
hidden_size (int, 可选, 默认为 128) — 隐藏层的大小
max_ep_len (int, 可选, 默认为 4096) — 环境中一个 episode 的最大长度
action_tanh (bool, 可选, 默认为 True) — 是否在动作预测上使用 tanh 激活函数
vocab_size (int, 可选, 默认为 50257) — GPT-2 模型的词汇表大小。定义了在调用 DecisionTransformerModel 时传递的 inputs_ids 可以表示的不同 token 的数量。
n_positions (int, 可选, 默认为 1024) — 此模型可能使用的最大序列长度。通常设置为较大的值以防万一（例如，512 或 1024 或 2048）。
n_layer (int, 可选, 默认为 3) — Transformer 编码器中隐藏层的数量。
n_head (int, 可选, 默认为 1) — Transformer 编码器中每个注意力层的注意力头数。
n_inner (int, 可选) — 内部前馈层的维度。如果未设置，则默认为 n_embd 的 4 倍。
activation_function (str, 可选, 默认为 "gelu") — 激活函数，可以从列表 ["relu", "silu", "gelu", "tanh", "gelu_new"] 中选择。
resid_pdrop (float, 可选, 默认为 0.1) — 嵌入、编码器和池化器中所有全连接层的 dropout 概率。
embd_pdrop (int, 可选, 默认为 0.1) — 嵌入的 dropout 比率。
attn_pdrop (float, 可选, 默认为 0.1) — 注意力的 dropout 比率。
layer_norm_epsilon (float, 可选, 默认为 1e-5) — 在层归一化层中使用的 epsilon 值。
initializer_range (float, 可选, 默认为 0.02) — 用于初始化所有权重矩阵的截断正态分布初始化器的标准差。
scale_attn_weights (bool, 可选, 默认为 True) — 通过除以 sqrt(hidden_size) 来缩放注意力权重。
use_cache (bool, 可选, 默认为 True) — 模型是否应返回最后的键/值注意力 (并非所有模型都使用)。
scale_attn_by_inverse_layer_idx (bool, 可选, 默认为 False) — 是否额外地通过 1 / layer_idx + 1 缩放注意力权重。
reorder_and_upcast_attn (bool, 可选, 默认为 False) — 在使用混合精度训练时，是否在计算注意力 (点积) 之前缩放键 (K)，并将注意力点积/softmax 向上转换为 float()。

这是用于存储 DecisionTransformerModel 配置的配置类。它用于根据指定的参数实例化 Decision Transformer 模型，定义模型架构。使用默认值实例化配置将产生与标准 DecisionTransformer 架构类似的配置。许多配置选项用于实例化用作架构一部分的 GPT2 模型。

配置对象继承自 PretrainedConfig，可用于控制模型输出。阅读 PretrainedConfig 的文档以获取更多信息。

示例

>>> from transformers import DecisionTransformerConfig, DecisionTransformerModel

>>> # Initializing a DecisionTransformer configuration
>>> configuration = DecisionTransformerConfig()

>>> # Initializing a model (with random weights) from the configuration
>>> model = DecisionTransformerModel(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config

DecisionTransformerGPT2Model

class transformers.DecisionTransformerGPT2Model

< source >

( config )

forward

< source >

( input_ids: typing.Optional[torch.LongTensor] = None past_key_values: typing.Optional[typing.Tuple[typing.Tuple[torch.Tensor]]] = None attention_mask: typing.Optional[torch.FloatTensor] = None token_type_ids: typing.Optional[torch.LongTensor] = None position_ids: typing.Optional[torch.LongTensor] = None head_mask: typing.Optional[torch.FloatTensor] = None inputs_embeds: typing.Optional[torch.FloatTensor] = None encoder_hidden_states: typing.Optional[torch.Tensor] = None encoder_attention_mask: typing.Optional[torch.FloatTensor] = None use_cache: typing.Optional[bool] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None )

DecisionTransformerModel

class transformers.DecisionTransformerModel

< source >

( config )

参数

config (~DecisionTransformerConfig) — 包含模型所有参数的模型配置类。使用配置文件初始化不会加载与模型关联的权重，仅加载配置。查看 from_pretrained() 方法来加载模型权重。

Decision Transformer 模型。此模型是 PyTorch torch.nn.Module 子类。像使用常规 PyTorch 模块一样使用它，并参阅 PyTorch 文档以了解与通用用法和行为相关的所有事项。

该模型基于 GPT2 架构构建，用于在离线 RL 设置中执行动作的自回归预测。有关更多详细信息，请参阅论文：https://arxiv.org/abs/2106.01345

forward

< source >

( states: typing.Optional[torch.FloatTensor] = None actions: typing.Optional[torch.FloatTensor] = None rewards: typing.Optional[torch.FloatTensor] = None returns_to_go: typing.Optional[torch.FloatTensor] = None timesteps: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.FloatTensor] = None output_hidden_states: typing.Optional[bool] = None output_attentions: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None ) → transformers.models.decision_transformer.modeling_decision_transformer.DecisionTransformerOutput or tuple(torch.FloatTensor)

参数

states (torch.FloatTensor，形状为 (batch_size, episode_length, state_dim)) — 轨迹中每个步骤的状态
actions (torch.FloatTensor，形状为 (batch_size, episode_length, act_dim)) — “专家”策略为当前状态采取的动作，这些动作在自回归预测中被掩盖
rewards (torch.FloatTensor，形状为 (batch_size, episode_length, 1)) — 每个状态、动作的奖励
returns_to_go (torch.FloatTensor，形状为 (batch_size, episode_length, 1)) — 轨迹中每个状态的回报
timesteps (torch.LongTensor，形状为 (batch_size, episode_length)) — 轨迹中每个步骤的时间步
attention_mask (torch.FloatTensor，形状为 (batch_size, episode_length)) — 掩码，用于在执行自回归预测时掩盖动作

返回值

transformers.models.decision_transformer.modeling_decision_transformer.DecisionTransformerOutput 或 tuple(torch.FloatTensor)

一个 transformers.models.decision_transformer.modeling_decision_transformer.DecisionTransformerOutput 或一个 torch.FloatTensor 元组 (如果传递了 return_dict=False 或当 config.return_dict=False 时)，包含各种元素，具体取决于配置 (DecisionTransformerConfig) 和输入。

last_hidden_state (torch.FloatTensor，形状为 (batch_size, sequence_length, hidden_size)) — 模型最后一层输出的隐藏状态序列。
state_preds (torch.FloatTensor，形状为 (batch_size, sequence_length, state_dim)) — 环境状态预测
action_preds (torch.FloatTensor，形状为 (batch_size, sequence_length, action_dim)) — 模型动作预测
return_preds (torch.FloatTensor，形状为 (batch_size, sequence_length, 1)) — 每个状态的预测回报
hidden_states (tuple(torch.FloatTensor)，可选，当传递 output_hidden_states=True 或当 config.output_hidden_states=True 时返回) — torch.FloatTensor 的元组 (每个嵌入输出一个 + 每层输出一个)，形状为 (batch_size, sequence_length, hidden_size)。

模型在每层输出端的隐藏状态，加上初始嵌入输出。
attentions (tuple(torch.FloatTensor)，可选，当传递 output_attentions=True 或当 config.output_attentions=True 时返回) — torch.FloatTensor 的元组 (每层一个)，形状为 (batch_size, num_heads, sequence_length, sequence_length)。

注意力 softmax 之后的注意力权重，用于计算自注意力头中的加权平均。

DecisionTransformerModel 的 forward 方法，覆盖了 __call__ 特殊方法。

尽管 forward 传递的配方需要在此函数中定义，但应该在之后调用 Module 实例，而不是直接调用此函数，因为前者负责运行预处理和后处理步骤，而后者会静默地忽略它们。

示例

>>> from transformers import DecisionTransformerModel
>>> import torch

>>> model = DecisionTransformerModel.from_pretrained("edbeeching/decision-transformer-gym-hopper-medium")
>>> # evaluation
>>> model = model.to(device)
>>> model.eval()

>>> env = gym.make("Hopper-v3")
>>> state_dim = env.observation_space.shape[0]
>>> act_dim = env.action_space.shape[0]

>>> state = env.reset()
>>> states = torch.from_numpy(state).reshape(1, 1, state_dim).to(device=device, dtype=torch.float32)
>>> actions = torch.zeros((1, 1, act_dim), device=device, dtype=torch.float32)
>>> rewards = torch.zeros(1, 1, device=device, dtype=torch.float32)
>>> target_return = torch.tensor(TARGET_RETURN, dtype=torch.float32).reshape(1, 1)
>>> timesteps = torch.tensor(0, device=device, dtype=torch.long).reshape(1, 1)
>>> attention_mask = torch.zeros(1, 1, device=device, dtype=torch.float32)

>>> # forward pass
>>> with torch.no_grad():
...     state_preds, action_preds, return_preds = model(
...         states=states,
...         actions=actions,
...         rewards=rewards,
...         returns_to_go=target_return,
...         timesteps=timesteps,
...         attention_mask=attention_mask,
...         return_dict=False,
...     )

< > 更新 on GitHub

Transformers

Decision Transformer

概述

DecisionTransformerConfig

class transformers.DecisionTransformerConfig

DecisionTransformerGPT2Model

class transformers.DecisionTransformerGPT2Model

forward

DecisionTransformerModel

class transformers.DecisionTransformerModel

forward