Add trajectory transformer

[CarlCochet]

What does this PR do?

Add Trajectory Transformer to transformers.

This is a transformer model used for offline deep reinforcement learning, taking observations, actions and rewards as one big sequence.
The model is adapted from the code here.
The caching mechanism as well as attentions outputs and hidden states outputs are additions not present in the original code.

The adapted forward pass produces the same output tensors as the forward pass in the original model.

Example usage for random environment data:

from transformers import TrajectoryTransformerModel
import torch
import numpy as np

config = TrajectoryTransformerConfig.from_pretrained("CarlCochet/trajectory-transformer-halfcheetah-medium-v2")
model = TrajectoryTransformerModel.from_pretrained("CarlCochet/trajectory-transformer-halfcheetah-medium-v2", config=config)

batch_size = 1
# sequence made of [ actions, observations, reward ]
seq_length = config.action_dim + config.observation_dim + 1 

trajectories = torch.LongTensor([np.random.permutation(seq_length) for _ in range(batch_size)])
outputs = model(trajectories)

prediction_logits = outputs.logits

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

TODO

• Pass all relevant tests ; need to disable some irrelevant tests as well
• Make a demo notebook

Who can review?

Anyone in the community is free to review the PR once the tests have passed.
@LysandreJik @thomwolf @edbeeching

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[CarlCochet]

Needed to add this for Windows because fix-copies and fixup were incorrectly modifying index.mdx and serialization.mdx (in docs/source/en)

[LysandreJik]

This looks quite good already! I have only left nits. I think an in-depth example would really help comprehension.

[LysandreJik]

Before merging, it would be nice to upload these to the author's organization. Have you been in contact with the authors? It would be nice to ask them if they want to have an organization on the hub and to host these weights there.

Let me know if you'd like some pointers.

[CarlCochet]

I have indeed been in contact with the authors and could probably ask if they want to do that, however their pretrained checkpoints are not compatible with this model. I had to modify them on my end to make them compatible with transformers, so I'm not sure how that would work.

[LysandreJik]

Given that it's a brand new model with an API that is very different from the original implementation, I think it would significantly help usage if we could add some in-depth examples here, either in the markdown or as a link to a colab notebook.

[CarlCochet]

If I am to put an in-depth example, it'll have to be in a collab as it'll probably have a large amount of code. Should I put a link to a collab in this place?

[LysandreJik]

Pinging @sgugger for a second review!

[sgugger]

Thanks a lot for adding this new model!

[LysandreJik]

LGTM, thanks @CarlCochet!