Yes, Q-learning Helps Offline In-Context RL

Official implementation accompanying Yes, Q-learning Helps Offline In-Context RL.

The code trains transformer-based in-context agents with either Algorithm Distillation (AD) or reinforcement-learning objectives. The included core experiments cover discrete GridWorld tasks and continuous MuJoCo tasks with IC-DQN, IC-CQL, IC-IQL, IC-TD3, and IC-TD3+BC.

This repository is a cleaned release of the submitted supplementary-code snapshot. It contains:

• Dark Room, Dark Key-to-Door, Janus, and supporting GridWorld environments;
• HalfCheetah-Velocity, Ant-Direction, Hopper-Parameters, and Walker-Parameters environments;
• tabular Q-learning and SAC dataset collection;
• AD and RL-objective transformer training;
• the aggregate result files and plotting code used for the included figures and tables.

Two limitations are important for exact reproduction of the latest paper revision:

• XLand-MiniGrid and the visual MW-DR9 implementation are not present in the submitted snapshot.
• The code registers K2D at sizes 9 and 13, whereas the latest paper appendix refers to K2D9 and K2D19. Verify this version mismatch before reporting an exact reproduction of the latest paper.

The paper experiments used NVIDIA H100 GPUs. Training requires Linux, an NVIDIA GPU, CUDA, and FlashAttention. CPU execution is not supported by the attention implementation.

docker build -t yes-q-learning-icrl .
docker run --gpus all --rm -it \
  -v "$PWD":/workspace/yes-q-learning-icrl \
  yes-q-learning-icrl

Python 3.10 is recommended.

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install torch==2.3.0 torchvision==0.18.0 torchaudio==2.3.0 \
  --index-url https://download.pytorch.org/whl/cu121
python -m pip install -r requirements.txt
python -m pip install flash-attn==2.5.9.post1 --no-build-isolation

Weights & Biases is used for experiment logging. Authenticate with wandb login, or run locally with WANDB_MODE=offline.

The discrete datasets contain complete learning histories collected with tabular Q-learning. This example constructs the complete DR9 dataset with 70 training targets and five histories per target using the paper's collector hyperparameters:

python generate_data.py \
  --env_name QDark-Room-9x9-v0 \
  --num_train_goals 70 \
  --num_histories 350 \
  --num_episodes 200 \
  --lr 0.9933 \
  --discount 0.9 \
  --savedir data/DR9_70_5

Existing non-empty output directories are protected. Pass --overwrite true only when the old dataset should be replaced.

Split a complete dataset into early, mid, and late expertise thirds:

python build_datasets.py --full_dataset data/DR9_70_5

Create a lower-coverage dataset by sampling a fixed number of histories:

python subsample_data.py \
  --full_dataset data/DR9_70_5 \
  --target_size 20 \
  --seed 0

Continuous-control histories are collected with the CleanRL-derived SAC collector:

python sac_train.py \
  --env-id HalfCheetahVel-v0 \
  --num-goals 100 \
  --histories-per-goal 1 \
  --num-workers 8 \
  --total-timesteps 100000 \
  --gamma 0.99 \
  --learning-starts 2000 \
  --savedir data/HCV_100_1

Dataset names follow {environment}_{number of targets}_{histories per target}, with _early, _mid, or _late appended for expertise subsets.

Ready-to-edit configurations are provided in configs/:

WANDB_MODE=offline python train_tuples_mapstyle.py \
  --config_path configs/ad_dr9.yaml

WANDB_MODE=offline python train_tuples_rl_mapstyle.py \
  --config_path configs/ic_cql_dr9.yaml

The RL entry point selects algorithms with the following fields:

For complete datasets, the paper uses subsample 4 for DR and ANT, 8 for K2D, and 2 for the other continuous tasks. For incomplete datasets it uses 1 for DR and continuous tasks and 2 for K2D. See Appendix D and Appendix I of the paper for the full tuning protocol and per-environment values.

Pre-aggregated scores are stored in plotting/bin/. Generate the paper overview plots with:

python plotting/plot.py --suite overview

Available suites are overview, janus, discrete, continuous, order, and all. Outputs are written to plotting/out/ and plotting/tables/. The larger suites use bootstrap confidence intervals and can take several minutes.

configs/                       paper-aligned example configurations
plotting/bin/                  aggregate score inputs
plotting/plot.py               figure and table reproduction
src/envs/                      task environments
src/models.py                  AD and RL-objective transformer models
src/utils/data.py              learning-history dataset loader
generate_data.py               discrete Q-learning collection
sac_train.py                   continuous SAC collection
build_datasets.py              expertise-level dataset splits
train_tuples_mapstyle.py       Algorithm Distillation training
train_tuples_rl_mapstyle.py    IC-DQN/CQL/IQL/TD3/TD3+BC training

@article{tarasov2025yes,
  title={Yes, Q-learning helps offline in-context RL},
  author={Tarasov, Denis and Nikulin, Alexander and Zisman, Ilya and Klepach, Albina and Polubarov, Andrei and Lyubaykin, Nikita and Derevyagin, Alexander and Kiselev, Igor and Kurenkov, Vladislav},
  journal={arXiv preprint arXiv:2502.17666},
  year={2025}
}

This project is licensed under the Apache License 2.0.