决策转换器
概述
决策转换器模型在 决策转换器:通过序列建模进行强化学习
由陈莉莉、陆凯文、拉杰斯瓦兰·阿拉文德、李基民、格罗弗·阿迪提亚、拉金·迈克尔、阿贝勒·皮特、斯里尼瓦斯·阿拉文德、莫达奇·伊戈尔提出。
论文中的摘要如下:
我们介绍了一种框架,将强化学习 (RL) 抽象为序列建模问题。这使我们能够利用 Transformer 架构的简单性和可扩展性,以及与之相关的语言建模方面的进展,例如 GPT-x 和 BERT。特别是,我们提出了决策转换器,这是一种将 RL 问题转化为条件序列建模的架构。与以前将价值函数拟合或计算策略梯度的方法不同,决策转换器只需通过利用因果掩蔽的 Transformer 来输出最佳动作。通过根据所需回报(奖励)、过去状态和动作对自回归模型进行条件化,我们的决策转换器模型可以生成实现所需回报的未来动作。尽管它很简单,但决策转换器在 Atari、OpenAI Gym 和 Key-to-Door 任务上与最先进的无模型离线 RL 基线相匹配或超过其性能。
此版本的模型适用于状态为向量的任务。
此模型由 edbeeching 贡献。原始代码可以在 这里找到。
DecisionTransformerConfig
class transformers.DecisionTransformerConfig
< 源代码 >( 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) — 环境中一个片段的最大长度 - action_tanh (
bool
, 可选, 默认为 True) — 是否在动作预测中使用 tanh 激活函数 - vocab_size (
int
, 可选, 默认为 50257) — GPT-2 模型的词汇量大小。定义调用 DecisionTransformerModel 时传递的inputs_ids
可以表示的不同标记数量。 - 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 比例。 - 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 的配置。它用于根据指定的参数实例化一个决策转换器模型,定义模型架构。使用默认值实例化一个配置将产生与标准决策转换器架构类似的配置。许多配置选项用于实例化作为架构的一部分使用的 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
forward
< 源代码 > ( input_ids: Optional = None past_key_values: Optional = None attention_mask: Optional = None token_type_ids: Optional = None position_ids: Optional = None head_mask: Optional = None inputs_embeds: Optional = None encoder_hidden_states: Optional = None encoder_attention_mask: Optional = None use_cache: Optional = None output_attentions: Optional = None output_hidden_states: Optional = None return_dict: Optional = None )
DecisionTransformerModel
class transformers.DecisionTransformerModel
< source >( config )
参数
- config (~DecisionTransformerConfig) — 模型配置类,包含模型的所有参数。使用配置文件初始化不会加载与模型关联的权重,仅加载配置。查看 from_pretrained() 方法加载模型权重。
Decision Transformer 模型这是一个 PyTorch torch.nn.Module 子类。把它当作一个普通的 PyTorch 模块,并在所有与一般使用和行为相关的方面参考 PyTorch 文档。
该模型基于 GPT2 架构,在离线强化学习设置中执行动作的自回归预测。有关更多详细信息,请参阅论文:https://arxiv.org/abs/2106.01345
forward
< source > ( states: Optional = None actions: Optional = None rewards: Optional = None returns_to_go: Optional = None timesteps: Optional = None attention_mask: Optional = None output_hidden_states: Optional = None output_attentions: Optional = None return_dict: Optional = None ) → transformers.models.decision_transformer.modeling_decision_transformer.DecisionTransformerOutput
or tuple(torch.FloatTensor)
参数
- states (
torch.FloatTensor
of shape(batch_size, episode_length, state_dim)
) — 轨迹中每一步的状态 - actions (
torch.FloatTensor
of shape(batch_size, episode_length, act_dim)
) — 当前状态下“专家”策略采取的动作,这些动作被屏蔽以进行自回归预测 - rewards (
torch.FloatTensor
of shape(batch_size, episode_length, 1)
) — 每个状态,动作的奖励 - returns_to_go (
torch.FloatTensor
of shape(batch_size, episode_length, 1)
) — 轨迹中每个状态的回报 - timesteps (
torch.LongTensor
of shape(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 正向方法,覆盖了 __call__
特殊方法。
虽然正向传递的配方需要在此函数中定义,但之后应该调用 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,
... )