class_docstr: The model builds upon the GPT2 architecture to perform autoregressive prediction of actions in an offline RL
  setting. Refer to the paper for more details: https://arxiv.org/abs/2106.01345
class_docstr_tok: ["The","model","builds","upon","the","GPT2","architecture","to","perform","autoregressive","prediction","of","actions","in","an","offline","RL","setting",".","Refer","to","the","paper","for","more","details",":","https",":","\/\/arxiv.org\/abs\/2106.01345"]
docstr: Returns:
  
  Examples:
  
  ```python
  >>> 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,
  ...     )
  ```
docstr_tok: ["Returns",":","Examples",":","``","`","python",">",">",">","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",",","...",")","``","`"]
code: def forward(
      self,
      states: Optional[torch.FloatTensor] = None,
      actions: Optional[torch.FloatTensor] = None,
      rewards: Optional[torch.FloatTensor] = None,
      returns_to_go: Optional[torch.FloatTensor] = None,
      timesteps: Optional[torch.LongTensor] = None,
      attention_mask: Optional[torch.FloatTensor] = None,
      output_hidden_states: Optional[bool] = None,
      output_attentions: Optional[bool] = None,
      return_dict: Optional[bool] = None,
  ) -> Union[Tuple[torch.FloatTensor], DecisionTransformerOutput]:
      r"""
      Returns:
  
      Examples:
  
      ```python
      >>> 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,
      ...     )
      ```"""
  
      output_attentions = (
          output_attentions
          if output_attentions is not None
          else self.config.output_attentions
      )
      output_hidden_states = (
          output_hidden_states
          if output_hidden_states is not None
          else self.config.output_hidden_states
      )
      return_dict = (
          return_dict
          if return_dict is not None
          else self.config.use_return_dict
      )
  
      batch_size, seq_length = states.shape[0], states.shape[1]
  
      if attention_mask is None:
          # attention mask for GPT: 1 if can be attended to, 0 if not
          attention_mask = torch.ones(
              (batch_size, seq_length), dtype=torch.long
          )
  
      # embed each modality with a different head
      state_embeddings = self.embed_state(states)
      action_embeddings = self.embed_action(actions)
      returns_embeddings = self.embed_return(returns_to_go)
      time_embeddings = self.embed_timestep(timesteps)
  
      # time embeddings are treated similar to positional embeddings
      state_embeddings = state_embeddings + time_embeddings
      action_embeddings = action_embeddings + time_embeddings
      returns_embeddings = returns_embeddings + time_embeddings
  
      # this makes the sequence look like (R_1, s_1, a_1, R_2, s_2, a_2, ...)
      # which works nice in an autoregressive sense since states predict actions
      stacked_inputs = (
          torch.stack(
              (returns_embeddings, state_embeddings, action_embeddings),
              dim=1,
          )
          .permute(0, 2, 1, 3)
          .reshape(batch_size, 3 * seq_length, self.hidden_size)
      )
      stacked_inputs = self.embed_ln(stacked_inputs)
  
      # to make the attention mask fit the stacked inputs, have to stack it as well
      stacked_attention_mask = (
          torch.stack(
              (attention_mask, attention_mask, attention_mask), dim=1
          )
          .permute(0, 2, 1)
          .reshape(batch_size, 3 * seq_length)
      )
      device = stacked_inputs.device
      # we feed in the input embeddings (not word indices as in NLP) to the model
      encoder_outputs = self.encoder(
          inputs_embeds=stacked_inputs,
          attention_mask=stacked_attention_mask,
          position_ids=torch.zeros(
              stacked_attention_mask.shape,
              device=device,
              dtype=torch.long,
          ),
          output_attentions=output_attentions,
          output_hidden_states=output_hidden_states,
          return_dict=return_dict,
      )
      x = encoder_outputs[0]
  
      # reshape x so that the second dimension corresponds to the original
      # returns (0), states (1), or actions (2); i.e. x[:,1,t] is the token for s_t
      x = x.reshape(
          batch_size, seq_length, 3, self.hidden_size
      ).permute(0, 2, 1, 3)
  
      # get predictions
      return_preds = self.predict_return(
          x[:, 2]
      )  # predict next return given state and action
      state_preds = self.predict_state(
          x[:, 2]
      )  # predict next state given state and action
      action_preds = self.predict_action(
          x[:, 1]
      )  # predict next action given state
      if not return_dict:
          return (state_preds, action_preds, return_preds)
  
      return DecisionTransformerOutput(
          last_hidden_state=encoder_outputs.last_hidden_state,
          state_preds=state_preds,
          action_preds=action_preds,
          return_preds=return_preds,
          hidden_states=encoder_outputs.hidden_states,
          attentions=encoder_outputs.attentions,
      )
  
imports: import math
  import os
  from dataclasses import dataclass
  from typing import Optional, Tuple, Union
  import torch
  import torch.utils.checkpoint
  from torch import nn
  from torch.cuda.amp import autocast
  from ...activations import ACT2FN
  from ...modeling_outputs import BaseModelOutputWithPastAndCrossAttentions
  from ...modeling_utils import PreTrainedModel
  from ...pytorch_utils import Conv1D, find_pruneable_heads_and_indices, prune_conv1d_layer
  from ...utils import ModelOutput, add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
  from .configuration_decision_transformer import DecisionTransformerConfig
label_desc: Use image node_id 2 for calling the DecisionTransformerModel obj's underlying member method code with example usage: obj.forward(states, actions, rewards, returns_to_go, timesteps, attention_mask, output_hidden_states, output_attentions, return_dict) and returns: DecisionTransformerOutput, state_preds,  action_preds,  return_preds