codingutf8 2022 the trajectory transformers paper s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license pytorch trajectorytransformer model import math import os from dataclasses import dataclass from typing import optional tuple union import numpy as np import torch import torch utils checkpoint from torch import nn from torch nn import functional as f from modelingutils import pretrainedmodel from utils import modeloutput addstartdocstrings addstartdocstringstomodelforward logging replacereturndocstrings from configurationtrajectorytransformer import trajectorytransformerconfig logger logging getloggername checkpointfordoc carlcochettrajectorytransformerhalfcheetahmediumv2 configfordoc trajectorytransformerconfig trajectorytransformerpretrainedmodelarchivelist carlcochettrajectorytransformerhalfcheetahmediumv2 see all trajectorytransformer models at https huggingface comodels filtertrajectorytransformer def loadtfweightsintrajectorytransformermodel config tfcheckpointpath load weights from tf model adamv and adamm are variables used in adamweightdecayoptimizer to calculated m and v which are not required for using pretrained model base class for model s outputs that also contains a pooling of the last hidden states args loss torch floattensor of shape 1 optional returned when labels is provided language modeling loss logits torch floattensor of shape batchsize sequencelength config vocabsize prediction scores of the language modeling head scores for each vocabulary token before softmax pastkeyvalues tupletupletorch tensor optional returned when usecachetrue is passed or when config usecachetrue tuple of length config nlayers containing tuples of tensors of shape batchsize numheads sequencelength embedsizeperhead contains precomputed hiddenstates key and values in the attention blocks that can be used see pastkeyvalues input to speed up sequential decoding hiddenstates tupletorch floattensor optional returned when outputhiddenstatestrue is passed or when config outputhiddenstatestrue tuple of torch floattensor one for the output of the embeddings one for the output of each layer of shape batchsize sequencelength hiddensize hiddenstates of the model at the output of each layer plus the initial embedding outputs attentions tupletorch floattensor optional returned when outputattentionstrue is passed or when config outputattentionstrue tuple of torch floattensor one for each layer of shape batchsize numheads sequencelength sequencelength gpt2attentions weights after the attention softmax used to compute the weighted average in the selfattention heads an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models trajectorytransformerinputsdocstring r args trajectories torch longtensor of shape batchsize sequencelength batch of trajectories where a trajectory is a sequence of states actions and rewards pastkeyvalues tupletupletorch tensor of length config nlayers optional contains precomputed hiddenstates key and values in the attention blocks as computed by the model see pastkeyvalues output below can be used to speed up sequential decoding the inputids which have their past given to this model should not be passed as inputids as they have already been computed targets torch longtensor of shape batchsize sequencelength optional desired targets used to compute the loss attentionmask torch floattensor of shape batchsize sequencelength optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossaryattentionmask usecache bool optional if set to true pastkeyvalues key value states are returned and can be used to speed up decoding see pastkeyvalues outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail returndict bool optional whether or not to return a utils modeloutput instead of a plain tuple args input torch floattensor of shape b nmodels inputdim the input to the layer batchsize x nmodels x outputdim key query value projections for all heads regularization output projection causal mask to ensure that attention is only applied to the left in the input sequence mask previous value estimates calculate query key values for all heads in batch and move head forward to be the batch dim batchsize x nheads x sequencelength x headdim causal selfattention batchsize x nheads x sequencelength x sequencelength batchsize x sequencelength x embeddingdim reassemble all head outputs side by side output projection mlp the full gpt language model with a context size of blocksize def initself config super initconfig input embedding stem 1 for stop token self tokemb nn embeddingconfig vocabsize config transitiondim 1 config nembd self posemb nn parametertorch zeros1 config blocksize config nembd self drop nn dropoutconfig embdpdrop transformer self blocks nn modulelistblockconfig for in rangeconfig nlayer decoder head self lnf nn layernormconfig nembd self head einlinearconfig transitiondim config nembd config vocabsize 1 biasfalse self vocabsize config vocabsize self stoptoken config vocabsize config transitiondim self blocksize config blocksize self observationdim config observationdim self actiondim config actiondim self transitiondim config transitiondim self embeddingdim config nembd self actionweight config actionweight self rewardweight config rewardweight self valueweight config valueweight self gradientcheckpointing false self postinit def getblocksizeself return self blocksize def offsettokensself trajectories sequencelength trajectories shape nstates intnp ceilsequencelength self transitiondim offsets torch arangeself transitiondim self vocabsize offsets offsets repeatnstates totrajectories device offsettrajectories trajectories offsets sequencelength offsettrajectoriestrajectories self vocabsize self stoptoken return offsettrajectories def padtofullobservationself hiddenstates batchsize sequencelength hiddenstates shape npad self transitiondim sequencelength self transitiondim self transitiondim padding torch zerosbatchsize npad self embeddingdim devicehiddenstates device batchsize x paddedsequencelength x embeddingdim hiddenstatespad torch cathiddenstates padding dim1 hiddenstatespad hiddenstatespad view1 self transitiondim self embeddingdim return hiddenstatespad npad addstartdocstringstomodelforward trajectorytransformerinputsdocstring formatbatchsize sequencelength replacereturndocstringsoutputtypetrajectorytransformeroutput configclassconfigfordoc def forward self trajectories optionaltorch longtensor none pastkeyvalues optionaltupletupletorch tensor none targets optionaltorch floattensor none attentionmask optionaltorch floattensor none usecache optionalbool none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniontupletorch tensor trajectorytransformeroutput r returns examples python from transformers import trajectorytransformermodel import torch model trajectorytransformermodel frompretrained carlcochettrajectorytransformerhalfcheetahmediumv2 model todevice model eval observationsdim actiondim batchsize 17 6 256 seqlength observationsdim actiondim 1 trajectories torch longtensornp random permutationself seqlength for in rangebatchsize to device targets torch longtensornp random permutationself seqlength for in rangebatchsize todevice outputs model trajectories targetstargets usecachetrue outputattentionstrue outputhiddenstatestrue returndicttrue batchsize x sequencelength x embeddingdim forward the gpt model batchsize x sequencelength x embeddingdim if we are given some desired targets also calculate the loss make weights coding utf 8 2022 the trajectory transformers paper s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license pytorch trajectorytransformer model see all trajectorytransformer models at https huggingface co models filter trajectory_transformer load tf checkpoints in a pytorch model load weights from tf model adam_v and adam_m are variables used in adamweightdecayoptimizer to calculated m and v which are not required for using pretrained model base class for model s outputs that also contains a pooling of the last hidden states args loss torch floattensor of shape 1 optional returned when labels is provided language modeling loss logits torch floattensor of shape batch_size sequence_length config vocab_size prediction scores of the language modeling head scores for each vocabulary token before softmax past_key_values tuple tuple torch tensor optional returned when use_cache true is passed or when config use_cache true tuple of length config n_layers containing tuples of tensors of shape batch_size num_heads sequence_length embed_size_per_head contains pre computed hidden states key and values in the attention blocks that can be used see past_key_values input to speed up sequential decoding hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each layer of shape batch_size sequence_length hidden_size hidden states of the model at the output of each layer plus the initial embedding outputs attentions tuple torch floattensor optional returned when output_attentions true is passed or when config output_attentions true tuple of torch floattensor one for each layer of shape batch_size num_heads sequence_length sequence_length gpt2attentions weights after the attention softmax used to compute the weighted average in the self attention heads an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models this model is a pytorch torch nn module https pytorch org docs stable nn html torch nn module sub class use it as a regular pytorch module and refer to the pytorch documentation for all matter related to general usage and behavior parameters config trajectorytransformerconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the pretrainedmodel from_pretrained method to load the model weights args trajectories torch longtensor of shape batch_size sequence_length batch of trajectories where a trajectory is a sequence of states actions and rewards past_key_values tuple tuple torch tensor of length config n_layers optional contains precomputed hidden states key and values in the attention blocks as computed by the model see past_key_values output below can be used to speed up sequential decoding the input_ids which have their past given to this model should not be passed as input_ids as they have already been computed targets torch longtensor of shape batch_size sequence_length optional desired targets used to compute the loss attention_mask torch floattensor of shape batch_size sequence_length optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask use_cache bool optional if set to true past_key_values key value states are returned and can be used to speed up decoding see past_key_values output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple args input torch floattensor of shape b n_models input_dim the input to the layer batch_size x n_models x output_dim key query value projections for all heads regularization output projection causal mask to ensure that attention is only applied to the left in the input sequence mask previous value estimates calculate query key values for all heads in batch and move head forward to be the batch dim batch_size x n_heads x sequence_length x head_dim causal self attention batch_size x n_heads x sequence_length x sequence_length batch_size x sequence_length x embedding_dim re assemble all head outputs side by side output projection mlp the full gpt language model with a context size of block_size input embedding stem 1 for stop token transformer decoder head batch_size x padded_sequence_length x embedding_dim returns examples python from transformers import trajectorytransformermodel import torch model trajectorytransformermodel from_pretrained carlcochet trajectory transformer halfcheetah medium v2 model to device model eval observations_dim action_dim batch_size 17 6 256 seq_length observations_dim action_dim 1 trajectories torch longtensor np random permutation self seq_length for _ in range batch_size to device targets torch longtensor np random permutation self seq_length for _ in range batch_size to device outputs model trajectories targets targets use_cache true output_attentions true output_hidden_states true return_dict true batch_size x sequence_length x embedding_dim forward the gpt model each index maps to a learnable vector each position maps to a learnable vector batch_size x sequence_length x embedding_dim if we are given some desired targets also calculate the loss make weights
import math
import os
from dataclasses import dataclass
from typing import Optional, Tuple, Union

import numpy as np
import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import functional as F

from ....modeling_utils import PreTrainedModel
from ....utils import (
    ModelOutput,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from .configuration_trajectory_transformer import TrajectoryTransformerConfig


logger = logging.get_logger(__name__)

_CHECKPOINT_FOR_DOC = "CarlCochet/trajectory-transformer-halfcheetah-medium-v2"
_CONFIG_FOR_DOC = "TrajectoryTransformerConfig"

TRAJECTORY_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "CarlCochet/trajectory-transformer-halfcheetah-medium-v2",
    
]


def load_tf_weights_in_trajectory_transformer(model, config, tf_checkpoint_path):
    
    try:
        import re

        import numpy as np
        import tensorflow as tf
    except ImportError:
        logger.error(
            "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
            "https://www.tensorflow.org/install/ for installation instructions."
        )
        raise
    tf_path = os.path.abspath(tf_checkpoint_path)
    logger.info(f"Converting TensorFlow checkpoint from {tf_path}")
    
    init_vars = tf.train.list_variables(tf_path)
    names = []
    arrays = []
    for name, shape in init_vars:
        logger.info(f"Loading TF weight {name} with shape {shape}")
        array = tf.train.load_variable(tf_path, name)
        names.append(name)
        arrays.append(array)

    for name, array in zip(names, arrays):
        name = name.split("/")
        
        
        if any(
            n in ["adam_v", "adam_m", "AdamWeightDecayOptimizer", "AdamWeightDecayOptimizer_1", "global_step"]
            for n in name
        ):
            logger.info(f"Skipping {'/'.join(name)}")
            continue
        pointer = model
        for m_name in name:
            if re.fullmatch(r"[A-Za-z]+_\d+", m_name):
                scope_names = re.split(r"_(\d+)", m_name)
            else:
                scope_names = [m_name]
            if scope_names[0] == "kernel" or scope_names[0] == "gamma":
                pointer = getattr(pointer, "weight")
            elif scope_names[0] == "output_bias" or scope_names[0] == "beta":
                pointer = getattr(pointer, "bias")
            elif scope_names[0] == "output_weights":
                pointer = getattr(pointer, "weight")
            elif scope_names[0] == "squad":
                pointer = getattr(pointer, "classifier")
            else:
                try:
                    pointer = getattr(pointer, scope_names[0])
                except AttributeError:
                    logger.info(f"Skipping {'/'.join(name)}")
                    continue
            if len(scope_names) >= 2:
                num = int(scope_names[1])
                pointer = pointer[num]
        if m_name[-11:] == "_embeddings":
            pointer = getattr(pointer, "weight")
        elif m_name == "kernel":
            array = np.transpose(array)
        try:
            if pointer.shape != array.shape:
                raise ValueError(f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched")
        except AssertionError as e:
            e.args += (pointer.shape, array.shape)
            raise
        logger.info(f"Initialize PyTorch weight {name}")
        pointer.data = torch.from_numpy(array)
    return model


@dataclass
class TrajectoryTransformerOutput(ModelOutput):
    

    loss: Optional[torch.FloatTensor] = None
    logits: torch.FloatTensor = None
    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None


class TrajectoryTransformerPreTrainedModel(PreTrainedModel):
    

    config_class = TrajectoryTransformerConfig
    load_tf_weights = load_tf_weights_in_trajectory_transformer
    base_model_prefix = "trajectory_transformer"
    main_input_name = "trajectories"
    supports_gradient_checkpointing = True

    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if isinstance(module, nn.Linear) and module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)
        elif isinstance(module, EinLinear):
            for i in range(module.n_models):
                nn.init.kaiming_uniform_(module.weight[i], a=math.sqrt(5) / self.config.kaiming_initializer_range)
                if module.bias is not None:
                    fan_in, _ = nn.init._calculate_fan_in_and_fan_out(module.weight[i])
                    bound = (1 / math.sqrt(fan_in)) * self.config.initializer_range
                    nn.init.uniform_(module.bias[i], -bound, bound)


TRAJECTORY_TRANSFORMER_START_DOCSTRING = r

TRAJECTORY_TRANSFORMER_INPUTS_DOCSTRING = r


class EinLinear(nn.Module):
    def __init__(self, n_models, in_features, out_features, bias):
        super().__init__()
        self.n_models = n_models
        self.out_features = out_features
        self.in_features = in_features
        self.weight = nn.Parameter(torch.Tensor(n_models, out_features, in_features))
        if bias:
            self.bias = nn.Parameter(torch.Tensor(n_models, out_features))
        else:
            self.register_parameter("bias", None)

    def reset_parameters(self):
        for i in range(self.n_models):
            nn.init.kaiming_uniform_(self.weight[i], a=math.sqrt(5))
            if self.bias is not None:
                fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight[i])
                bound = 1 / math.sqrt(fan_in)
                nn.init.uniform_(self.bias[i], -bound, bound)

    def forward(self, input):
        
        
        output = torch.einsum("eoi,bei->beo", self.weight, input)
        if self.bias is not None:
            raise RuntimeError()
        return output


class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()

        if config.n_embd % config.n_head != 0:
            raise ValueError(f"n_head ({config.n_head}) should be a divisor of n_embd ({config.n_embd})")

        
        self.key = nn.Linear(config.n_embd, config.n_embd)
        self.query = nn.Linear(config.n_embd, config.n_embd)
        self.value = nn.Linear(config.n_embd, config.n_embd)

        
        self.attn_drop = nn.Dropout(config.attn_pdrop)
        self.resid_drop = nn.Dropout(config.resid_pdrop)

        
        self.proj = nn.Linear(config.n_embd, config.n_embd)

        
        self.register_buffer(
            "mask",
            torch.tril(torch.ones(config.block_size, config.block_size)).view(
                1, 1, config.block_size, config.block_size
            ),
            persistent=False,
        )

        
        joined_dim = config.observation_dim + config.action_dim + 2
        self.mask.squeeze()[:, joined_dim - 1 :: joined_dim] = 0

        self.n_head = config.n_head

    def forward(
        self,
        hidden_states: Optional[Tuple[torch.FloatTensor]],
        layer_past: Optional[Tuple[torch.Tensor]] = None,
        use_cache: Optional[bool] = False,
        output_attentions: Optional[bool] = False,
    ):
        batch_size, sequence_length, embedding_dim = hidden_states.size()

        
        
        key = (
            self.key(hidden_states)
            .view(batch_size, sequence_length, self.n_head, embedding_dim // self.n_head)
            .transpose(1, 2)
        )
        query = (
            self.query(hidden_states)
            .view(batch_size, sequence_length, self.n_head, embedding_dim // self.n_head)
            .transpose(1, 2)
        )
        value = (
            self.value(hidden_states)
            .view(batch_size, sequence_length, self.n_head, embedding_dim // self.n_head)
            .transpose(1, 2)
        )

        if layer_past is not None:
            past_key, past_value = layer_past
            key = torch.cat((past_key, key), dim=-2)
            value = torch.cat((past_value, value), dim=-2)

        if use_cache is True:
            present = (key, value)
        else:
            present = None

        
        
        attn_weights = (torch.matmul(query, key.transpose(-2, -1))) * (1.0 / math.sqrt(key.size(-1)))
        attn_weights = attn_weights.masked_fill(
            self.mask[:, :, :sequence_length, :sequence_length] == 0, torch.finfo(attn_weights.dtype).min
        )
        attn_weights = F.softmax(attn_weights, dim=-1)
        self._attn_map = attn_weights.clone()
        attn_weights = self.attn_drop(attn_weights)

        output = torch.matmul(attn_weights, value)
        
        
        output = output.transpose(1, 2).contiguous().view(batch_size, sequence_length, embedding_dim)

        
        output = self.resid_drop(self.proj(output))

        outputs = (output, present)
        if output_attentions:
            outputs += (attn_weights,)

        return outputs


class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln1 = nn.LayerNorm(config.n_embd)
        self.ln2 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)

        
        self.l1 = nn.Linear(config.n_embd, 4 * config.n_embd)
        self.act = nn.GELU()
        self.l2 = nn.Linear(4 * config.n_embd, config.n_embd)
        self.drop = nn.Dropout(config.resid_pdrop)

    def forward(
        self,
        hidden_states: Optional[Tuple[torch.FloatTensor]],
        layer_past: Optional[Tuple[torch.Tensor]] = None,
        use_cache: Optional[bool] = False,
        output_attentions: Optional[bool] = False,
    ):
        residual = hidden_states
        hidden_states = self.ln1(hidden_states)

        attn_outputs = self.attn(
            hidden_states, layer_past=layer_past, use_cache=use_cache, output_attentions=output_attentions
        )
        attn_output = attn_outputs[0]
        outputs = attn_outputs[1:]
        hidden_states = attn_output + residual

        residual = hidden_states
        hidden_states = self.ln2(hidden_states)
        hidden_states = self.l1(hidden_states)
        hidden_states = self.act(hidden_states)
        hidden_states = self.l2(hidden_states)
        hidden_states = residual + self.drop(hidden_states)

        if use_cache:
            outputs = (hidden_states,) + outputs
        else:
            outputs = (hidden_states,) + outputs[1:]

        return outputs


@add_start_docstrings(
    "The bare TrajectoryTransformer Model transformer outputting raw hidden-states without any specific head on top.",
    TRAJECTORY_TRANSFORMER_START_DOCSTRING,
)
class TrajectoryTransformerModel(TrajectoryTransformerPreTrainedModel):
    

    def __init__(self, config):
        super().__init__(config)

        
        self.tok_emb = nn.Embedding(config.vocab_size * config.transition_dim + 1, config.n_embd)

        self.pos_emb = nn.Parameter(torch.zeros(1, config.block_size, config.n_embd))
        self.drop = nn.Dropout(config.embd_pdrop)
        
        self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layer)])
        
        self.ln_f = nn.LayerNorm(config.n_embd)
        self.head = EinLinear(config.transition_dim, config.n_embd, config.vocab_size + 1, bias=False)

        self.vocab_size = config.vocab_size
        self.stop_token = config.vocab_size * config.transition_dim
        self.block_size = config.block_size

        self.observation_dim = config.observation_dim
        self.action_dim = config.action_dim
        self.transition_dim = config.transition_dim
        self.embedding_dim = config.n_embd

        self.action_weight = config.action_weight
        self.reward_weight = config.reward_weight
        self.value_weight = config.value_weight

        self.gradient_checkpointing = False

        self.post_init()

    def get_block_size(self):
        return self.block_size

    def offset_tokens(self, trajectories):
        _, sequence_length = trajectories.shape

        n_states = int(np.ceil(sequence_length / self.transition_dim))

        offsets = torch.arange(self.transition_dim) * self.vocab_size
        offsets = offsets.repeat(n_states).to(trajectories.device)

        offset_trajectories = trajectories + offsets[:sequence_length]
        offset_trajectories[trajectories == self.vocab_size] = self.stop_token
        return offset_trajectories

    def pad_to_full_observation(self, hidden_states):
        batch_size, sequence_length, _ = hidden_states.shape

        n_pad = (self.transition_dim - sequence_length % self.transition_dim) % self.transition_dim
        padding = torch.zeros(batch_size, n_pad, self.embedding_dim, device=hidden_states.device)

        
        hidden_states_pad = torch.cat([hidden_states, padding], dim=1)
        hidden_states_pad = hidden_states_pad.view(-1, self.transition_dim, self.embedding_dim)

        return hidden_states_pad, n_pad

    @add_start_docstrings_to_model_forward(
        TRAJECTORY_TRANSFORMER_INPUTS_DOCSTRING.format("batch_size, sequence_length")
    )
    @replace_return_docstrings(output_type=TrajectoryTransformerOutput, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        trajectories: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
        targets: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], TrajectoryTransformerOutput]:
        r
        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
        )

        if past_key_values is None:
            past_key_values = tuple([None] * len(self.blocks))

        batch_size, sequence_length = trajectories.size()

        if sequence_length > self.block_size:
            raise ValueError("Cannot forward, model block size is exhausted.")

        offset_trajectories = self.offset_tokens(trajectories)
        
        
        token_embeddings = self.tok_emb(offset_trajectories)  
        position_embeddings = self.pos_emb[:, :sequence_length, :]  

        hidden_states = self.drop(token_embeddings + position_embeddings)

        if self.gradient_checkpointing and self.training:
            if use_cache:
                logger.warning_once(
                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
                )
                use_cache = False

        presents = () if use_cache else None
        all_self_attentions = () if output_attentions else None
        all_hidden_states = () if output_hidden_states else None

        for i, (block, layer_past) in enumerate(zip(self.blocks, past_key_values)):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            if self.gradient_checkpointing and self.training:
                outputs = self._gradient_checkpointing_func(
                    block.__call__,
                    hidden_states,
                    layer_past,
                    use_cache,
                    output_attentions,
                )
            else:
                outputs = block(hidden_states, layer_past, use_cache, output_attentions)

            hidden_states = outputs[0]
            if use_cache is True:
                presents = presents + (outputs[1],)

            if output_attentions:
                all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)

        
        hidden_state = self.ln_f(hidden_states)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        hidden_states_pad, n_pad = self.pad_to_full_observation(hidden_state)

        logits = self.head(hidden_states_pad)
        logits = logits.reshape(batch_size, sequence_length + n_pad, self.vocab_size + 1)
        logits = logits[:, :sequence_length]

        
        if targets is not None:
            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.view(-1), reduction="none")
            if self.action_weight != 1 or self.reward_weight != 1 or self.value_weight != 1:
                
                n_states = int(np.ceil(sequence_length / self.transition_dim))
                weights = torch.cat(
                    [
                        torch.ones(self.observation_dim, device=trajectories.device),
                        torch.ones(self.action_dim, device=trajectories.device) * self.action_weight,
                        torch.ones(1, device=trajectories.device) * self.reward_weight,
                        torch.ones(1, device=trajectories.device) * self.value_weight,
                    ]
                )
                weights = weights.repeat(n_states)
                weights = weights[1:].repeat(batch_size, 1)
                loss = loss * weights.view(-1)
            loss = (loss * attention_mask.view(-1)).mean()
        else:
            loss = None

        if not return_dict:
            return tuple(v for v in [loss, logits, presents, all_hidden_states, all_self_attentions] if v is not None)

        return TrajectoryTransformerOutput(
            loss=loss,
            logits=logits,
            past_key_values=presents,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
        )