Trajectory Transformer

## [](#overview)Overview

The Trajectory Transformer model was proposed in [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://arxiv.org/abs/2106.02039) by Michael Janner, Qiyang Li, Sergey Levine.

The abstract from the paper is the following:

_Reinforcement learning (RL) is typically concerned with estimating stationary policies or single-step models, leveraging the Markov property to factorize problems in time. However, we can also view RL as a generic sequence modeling problem, with the goal being to produce a sequence of actions that leads to a sequence of high rewards. Viewed in this way, it is tempting to consider whether high-capacity sequence prediction models that work well in other domains, such as natural-language processing, can also provide effective solutions to the RL problem. To this end, we explore how RL can be tackled with the tools of sequence modeling, using a Transformer architecture to model distributions over trajectories and repurposing beam search as a planning algorithm. Framing RL as sequence modeling problem simplifies a range of design decisions, allowing us to dispense with many of the components common in offline RL algorithms. We demonstrate the flexibility of this approach across long-horizon dynamics prediction, imitation learning, goal-conditioned RL, and offline RL. Further, we show that this approach can be combined with existing model-free algorithms to yield a state-of-the-art planner in sparse-reward, long-horizon tasks._

Tips:

This Transformer is used for deep reinforcement learning. To use it, you need to create sequences from actions, states and rewards from all previous timesteps. This model will treat all these elements together as one big sequence (a trajectory).

This model was contributed by [CarlCochet](https://huggingface.co/CarlCochet). The original code can be found [here](https://github.com/jannerm/trajectory-transformer).

## [](#transformers.TrajectoryTransformerConfig)TrajectoryTransformerConfig

### class transformers.TrajectoryTransformerConfig

[](#transformers.TrajectoryTransformerConfig)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/trajectory_transformer/configuration_trajectory_transformer.py#L31)

( vocab\_size = 100action\_weight = 5reward\_weight = 1value\_weight = 1block\_size = 249action\_dim = 6observation\_dim = 17transition\_dim = 25n\_layer = 4n\_head = 4n\_embd = 128embd\_pdrop = 0.1attn\_pdrop = 0.1resid\_pdrop = 0.1learning\_rate = 0.0006max\_position\_embeddings = 512initializer\_range = 0.02layer\_norm\_eps = 1e-12kaiming\_initializer\_range = 1use\_cache = Truepad\_token\_id = 1bos\_token\_id = 50256eos\_token\_id = 50256\*\*kwargs )

This is the configuration class to store the configuration of a [TrajectoryTransformerModel](/docs/transformers/v4.30.0/en/model_doc/trajectory_transformer#transformers.TrajectoryTransformerModel). It is used to instantiate an TrajectoryTransformer model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of the TrajectoryTransformer [CarlCochet/trajectory-transformer-halfcheetah-medium-v2](https://huggingface.co/CarlCochet/trajectory-transformer-halfcheetah-medium-v2) architecture.

Configuration objects inherit from [PretrainedConfig](/docs/transformers/v4.30.0/en/main_classes/configuration#transformers.PretrainedConfig) and can be used to control the model outputs. Read the documentation from [PretrainedConfig](/docs/transformers/v4.30.0/en/main_classes/configuration#transformers.PretrainedConfig) for more information.

```
>>> from transformers import TrajectoryTransformerConfig, TrajectoryTransformerModel

>>> 
>>> configuration = TrajectoryTransformerConfig()

>>> 
>>> model = TrajectoryTransformerModel(configuration)

>>> 
>>> configuration = model.config```

## [](#transformers.TrajectoryTransformerModel)TrajectoryTransformerModel

### class transformers.TrajectoryTransformerModel

[](#transformers.TrajectoryTransformerModel)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/trajectory_transformer/modeling_trajectory_transformer.py#L406)

( config )

Parameters

-   [](#transformers.TrajectoryTransformerModel.config)**config** ([TrajectoryTransformerConfig](/docs/transformers/v4.30.0/en/model_doc/trajectory_transformer#transformers.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 [from\_pretrained()](/docs/transformers/v4.30.0/en/main_classes/model#transformers.PreTrainedModel.from_pretrained) method to load the model weights.

The bare TrajectoryTransformer Model transformer outputting raw hidden-states without any specific head on top. 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.

the full GPT language model, with a context size of block\_size

#### forward

[](#transformers.TrajectoryTransformerModel.forward)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/trajectory_transformer/modeling_trajectory_transformer.py#L467)

( trajectories: typing.Optional\[torch.LongTensor\] = Nonepast\_key\_values: typing.Optional\[typing.Tuple\[typing.Tuple\[torch.Tensor\]\]\] = Nonetargets: typing.Optional\[torch.FloatTensor\] = Noneattention\_mask: typing.Optional\[torch.FloatTensor\] = Noneuse\_cache: typing.Optional\[bool\] = Noneoutput\_attentions: typing.Optional\[bool\] = Noneoutput\_hidden\_states: typing.Optional\[bool\] = Nonereturn\_dict: typing.Optional\[bool\] = None ) → `transformers.models.trajectory_transformer.modeling_trajectory_transformer.TrajectoryTransformerOutput` or `tuple(torch.FloatTensor)`

The [TrajectoryTransformerModel](/docs/transformers/v4.30.0/en/model_doc/trajectory_transformer#transformers.TrajectoryTransformerModel) forward method, overrides the `__call__` special method.

Although the recipe for forward pass needs to be defined within this function, one should call the `Module` instance afterwards instead of this since the former takes care of running the pre and post processing steps while the latter silently ignores them.

Examples:

```
>>> 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,
... )```