StARformer: Transformer with State-Action-Reward Representations

This repository contains the PyTorch implementation for our paper titled StARformer: Transformer with State-Action-Reward Representations.
We learn local State-Action-Reward representations (StAR-representations) to improve (long) sequence modeling for reinforcement learning (and imitation learning).

Results

Installation

Dependencies can be installed by Conda:

conda env create -f my_env.yml

And install Atari ROMs.

Datasets

Please follow this instruction for datasets.

Example usage

See run.sh or below: