StARformer
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:
python run_star_atari.py --seed 123 --data_dir_prefix [data_directory] --epochs 10 --num_steps 500000 --num_buffers 50 --batch_size 64 --seq_len 30 --model_type 'star' --game 'Breakout'
[data_directory]
is where you place the Atari dataset.
model_type
):
Variants ('star'
(imitation)'star_rwd'
(offline RL)'star_fusion'
(see Figure 4a in our paper)'star_stack'
(see Figure 4b in our paper)
Acknowledgement
This code is based on Decision-Transformer.