faster-trajectory-transformer - GitLite

KO | EN

gitlite — search

#javascript #typescript #go #golang #machine-learning #deep-learning #linux #nodejs #macos #cli #kubernetes #security #mcp #open-source

faster-trajectory-transformer

(★ 118)

Implementation of Trajectory Transformer with attention caching and batched beam search

파일 탐색기

최종 버전 다운로드 (.zip)

- - halfcheetah_medium.yaml
  - hopper_medium.yaml
  - walker2d_medium.yaml
- eval_base.yaml
- halfcheetah.gif
- hopper.gif
- walker2d.gif
- - config.yaml
  - discretizer.pt
  - model_last.pt
- - config.yaml
  - discretizer.pt
  - model_last.pt
- - config.yaml
  - discretizer.pt
  - model_last.pt
- eval.py
- train.py
- - d4rl_dataset.py
  - get_d4rl.py
- - beam.py
  - sample.py
- - common.py
  - discretization.py
  - env.py
  - scheduler.py
.gitignore
environment.yml
LICENSE
main.py
README.md
requirements.txt

# 설치 가이드

1. 코드 내려받기

git clone https://github.com/Howuhh/faster-trajectory-transformer

깃허브에서 프로젝트 코드 전체를 내 컴퓨터로 내려받습니다.

cd faster-trajectory-transformer

방금 내려받은 프로젝트 폴더 안으로 이동합니다.

2. Python

쉬움 추천

사전 준비물

Git GitHub에서 프로젝트 코드를 내려받으려면 필요합니다.
Python 3 설치 시 'Add Python to PATH' 옵션을 꼭 체크하세요 (Windows).

pip install -r requirements.txt

requirements.txt 등에 명시된 파이썬 라이브러리를 설치합니다.

python main.py

파이썬 스크립트(또는 모듈)를 실행합니다.

✅ 에러 메시지 없이 실행되고 터미널에 안내 문구가 출력되면 정상입니다.

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# Faster Trajectory Transformer

![video](https://raw.githubusercontent.com/Howuhh/faster-trajectory-transformer/main/media/halfcheetah.gif)
![video](https://raw.githubusercontent.com/Howuhh/faster-trajectory-transformer/main/media/walker2d.gif)
![video](https://raw.githubusercontent.com/Howuhh/faster-trajectory-transformer/main/media/hopper.gif)

This is reimplementation of Trajectory Transformer, introduced in **Offline Reinforcement Learning
as One Big Sequence Modeling Problem** [paper](https://arxiv.org/abs/2106.02039).

The original implementation has few problems with inference speed, namely quadratic attention during
inference and sequential rollouts. The former slows down planning a lot, while the latter does not
allow to do rollouts in parallel and utilize GPU to the full.

Still, even after all changes, it is not that fast compared to traditional methods such as PPO or SAC/DDPG.
However, the gains are huge, what used to take hours now takes a dozen minutes (25 rollouts, 1k steps each, for example).
Training time remains the same, though.

## Changes

### Attention caching

During beam search we're only predicting one token at a time. So with the naive implementation model will make a lot of
unnecessary computations to recompute attention maps for full past context. However it is not necessary,
as it was already computed when the previous token was predicted. **All we need is to cache it**!

Actually, attention caching is a common thing in NLP field, but a lot of RL practitioners may
not be familiar with NLP, so the code also can be educational.

### Vectorized rollouts

Vectorized environments allow batching beam search planning and select actions in parallel, which is a lot faster
if you need to evaluate agent on number of episodes (or seeds) during training.

## Training

I trained it on D4RL medium datasets to validate that everything is OK. Scores seem to be very close to the original.
Pretrained models are [available](pretrained).

All training parameters can be seen in training [configs](configs/medium).
Also, all datasets for [D4RL](https://sites.google.com/view/d4rl/home) Gym tasks are supported.

```bash
python scripts/train.py --config="configs/medium/halfcheetah_medium.yaml" --device="cuda" --seed="42"
```

## Evaluation

Available evaluation parameters can be seen in validation [config](configs/eval_base.yaml).
Here parameters are set to match evaluation configs from original implementation by [@jannerm](https://github.com/jannerm/trajectory-transformer).

```bash
# you can override every config value from command line

# halfcheetah-medium-v2
python scripts/eval.py \
    --config="configs/eval_base.yaml" --device="cuda" --seed="42" \
    checkpoints_path="pretrained/halfcheetah" \
    beam_context=5 \
    beam_steps=5 \
    beam_width=32

# hopper-medium-v2
python scripts/eval.py \
    --config="configs/eval_base.yaml" --device="cuda" --seed="42" \
    checkpoints_path="pretrained/hopper" \
    beam_context=5 \
    beam_steps=15 \
    beam_width=128

# walker2d-medium-v2
python scripts/eval.py \
    --config="configs/eval_base.yaml" --device="cuda" --seed="42" \
    checkpoints_path="pretrained/walker2d" \
    beam_context=5 \
    beam_steps=15 \
    beam_width=128
```

## References
```
@inproceedings{janner2021sequence,
  title = {Offline Reinforcement Learning as One Big Sequence Modeling Problem},
  author = {Michael Janner and Qiyang Li and Sergey Levine},
  booktitle = {Advances in Neural Information Processing Systems},
  year = {2021},
}
```

Please, if you use this code in your research, cite it:
```
@misc{fastertt2022,
  title={Faster Trajectory Transformer},
  author={Alexander Nikulin},
  year={2022},
  url={https://github.com/Howuhh/faster-trajectory-transformer},
}
```