Official codebase for SICQL: Scalable In-Context Q-Learning.

Implementation Details of SICQL.

conda env create -f environment.yaml

The datasets used in our experiments are located in the datasets/ directory. The collected dataset can be downloaded here

Running experiments based our code could be quite easy, you can execute the provided shell script to start:

sh run.sh

Or if you wish to skip training, you can directly load our pretrained model using the provided script:

python3 eval.py --env HalfCheetahVel-v0 --horizon 200 --context_horizon 4 --lr 3e-4 --n_layer 4 --m_layer 2 --head 1  --seed 1  --rollin_type 'mix' --epoch 400000  --freq 10000 --beta 1000.0 --n_embd 128  --context_hidden_dim 128 --context_dim 16  --c_layer 1  --context_epoch 200 --device 'cuda:0'