{"id":19682739,"url":"https://github.com/ml-jku/ra-dt","last_synced_at":"2025-04-29T05:30:33.652Z","repository":{"id":259742906,"uuid":"870340687","full_name":"ml-jku/RA-DT","owner":"ml-jku","description":"Retrieval-Augmented Decision Transformer: External Memory for In-context RL","archived":false,"fork":false,"pushed_at":"2024-10-27T16:10:03.000Z","size":349,"stargazers_count":16,"open_issues_count":0,"forks_count":1,"subscribers_count":3,"default_branch":"main","last_synced_at":"2025-04-05T13:23:34.197Z","etag":null,"topics":["decision-transformers","multi-task-learning","reinforcement-learning","retrieval-augmented-generation"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/ml-jku.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-10-09T21:30:40.000Z","updated_at":"2025-03-21T18:43:29.000Z","dependencies_parsed_at":"2024-10-27T18:58:06.018Z","dependency_job_id":"b2af0db6-1e04-4eee-859f-3d684f032a49","html_url":"https://github.com/ml-jku/RA-DT","commit_stats":null,"previous_names":["ml-jku/ra-dt"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ml-jku%2FRA-DT","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ml-jku%2FRA-DT/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ml-jku%2FRA-DT/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ml-jku%2FRA-DT/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/ml-jku","download_url":"https://codeload.github.com/ml-jku/RA-DT/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":251444076,"owners_count":21590410,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["decision-transformers","multi-task-learning","reinforcement-learning","retrieval-augmented-generation"],"created_at":"2024-11-11T18:11:58.019Z","updated_at":"2025-04-29T05:30:32.689Z","avatar_url":"https://github.com/ml-jku.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Retrieval-Augmented Decision Transformer: External Memory for In-context RL\n[![arXiv](https://img.shields.io/badge/arXiv-2410.07071-b31b1b.svg)](https://arxiv.org/abs/2410.07071)\n[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)\n\nThomas Schmied\u003csup\u003e**1**\u003c/sup\u003e, Fabian Paischer\u003csup\u003e**1**\u003c/sup\u003e, Vihang Patil\u003csup\u003e**1**\u003c/sup\u003e, Markus Hofmarcher\u003csup\u003e**2**\u003c/sup\u003e, Razvan Pascanu\u003csup\u003e**3,4**\u003c/sup\u003e, Sepp Hochreiter\u003csup\u003e**1,5**\u003c/sup\u003e \n\n\u003csup\u003e**1**\u003c/sup\u003eELLIS Unit Linz and LIT AI Lab, Institute for Machine Learning, JKU Linz, Austria\\\n\u003csup\u003e**2**\u003c/sup\u003eExtensity AI, \u003csup\u003e**3**\u003c/sup\u003eGoogle DeepMind\\, \u003csup\u003e**4**\u003c/sup\u003eUCL\\, \u003csup\u003e**5**\u003c/sup\u003eNXAI\n\nThis repository contains the source code for **\"Retrieval-Augmented Decision Transformer: External Memory for In-context RL\"**.\nThe paper is available on [ArXiv](https://arxiv.org/abs/2410.07071). \n\n![Retrieval-Augmented Decision Transformer](./figures/radt.png) \n\n## Overview\nThis code-based is built on [L2M](https://github.com/ml-jku/L2M) and relies on open-source frameworks, including: \n- [PyTorch](https://github.com/pytorch/pytorch)\n- [Huggingface transformers](https://github.com/huggingface/transformers)\n- [stable-baselines3](https://github.com/DLR-RM/stable-baselines3)\n- [wandb](https://github.com/wandb/wandb)\n- [Hydra](https://github.com/facebookresearch/hydra)\n\n\nWhat is in this repository?\n```\n.\n├── configs                    # Contains all .yaml config files for Hydra to configure agents, envs, etc.\n│   ├── agent_params            \n│   ├── wandb_callback_params\n│   ├── env_params\n│   ├── eval_params\n│   ├── run_params\n│   └── config.yaml            # Main config file for Hydra - specifies log/data/model directories.\n├── continual_world            # Submodule for Continual-World.\n├── dmc2gym_custom             # Custom wrapper for DMControl.\n├── figures             \n├── src                        # Main source directory.\n│   ├── algos                  # Contains agent/model/prompt classes.\n│   ├── augmentations          # Image augmentations.\n│   ├── buffers                # Contains replay trajectory buffers.\n│   ├── callbacks              # Contains callbacks for training (e.g., WandB, evaluation, etc.).\n│   ├── data                   # Contains data utilities.\n│   ├── envs                   # Contains functionality for creating environments.\n│   ├── exploration            # Contains exploration strategies.\n│   ├── optimizers             # Contains optimizers.\n│   ├── schedulers             # Contains learning rate schedulers.\n│   ├── tokenizers_custom      # Contains custom tokenizers for discretizing states/actions.\n│   ├── utils                  \n│   └── __init__.py\n├── LICENSE\n├── README.md\n├── environment.yaml\n├── requirements.txt\n├── precompute_img_embeds.py   # Script for pre-computing image embeddings.\n├── evaluate.py                # Entry point for evaluating agents.\n└── main.py                    # Entry point for training agents.\n```\n## Installation\nEnvironment configuration and dependencies are available in `environment.yaml` and `requirements.txt`.\n\nFirst, create the conda environment.\n```\nconda env create -f environment.yaml\nconda activate radt\n```\n\nThen install the remaining requirements (with MuJoCo already downloaded, if not see [here](#MuJoCo-installation)): \n```\npip install -r requirements.txt\n```\n\nInit the `continualworld` submodule and install: \n```\ngit submodule init\ngit submodule update\ncd continual_world\npip install .\n```\nInstall `meta-world`:\n```\npip install git+https://github.com/rlworkgroup/metaworld.git@18118a28c06893da0f363786696cc792457b062b\n```\n\nInstall custom version of [dmc2gym](https://github.com/denisyarats/dmc2gym). Our version makes `flatten_obs` optional, \nand, thus, allows us to construct the full observation space of all DMControl envs. \n```\ncd dmc2gym_custom\npip install -e .\n```\n\nInstalling `minihack` without sudo rights. In this case, follow instructions provided [here](https://github.com/facebookresearch/nle/issues/246).\n\nTo install the GPU version of `faiss` (which we used in our experiments), use: \n```\npip uninstall faiss-cpu\nconda install -c conda-forge faiss-gpu=1.7.4\n```\nIn case this installation causes issues, we recommend installing via `mamba`. \n\n### MuJoCo installation\nFor the installation of MuJoCo and potential tips on troubleshooting, we refer to the L2M repository: https://github.com/ml-jku/L2M\n\n## Setup\n\n### Experiment configuration\nThis codebase relies on [Hydra](https://github.com/facebookresearch/hydra), which configures experiments via `.yaml` files. \nHydra automatically creates the log folder structure for a given run, as specified in the respective `config.yaml` file.\n\nThe `config.yaml` is the main configuration entry point and contains the default parameters. The file references the respective default parameter files under the block\n`defaults`. In addition, `config.yaml` contains 4 important constants that configure the directory paths: \n```\nLOG_DIR: ../logs\nDATA_DIR: ../data\nSSD_DATA_DIR: ../data\nMODELS_DIR: ../models\n```\n\n### Datasets\nThe datasets we generated for grid-worlds are available on the Huggingface Hub 🤗, and can be downloaded using the `huggingface-cli`: \n```\n# Dark-Room\nhuggingface-cli download ml-jku/dark_room --local-dir=./dark_room --repo-type dataset\n# Dark Key-Door\nhuggingface-cli download ml-jku/dark_keydoor --local-dir=./dark_keydoor --repo-type dataset\n# MazeRunner\nhuggingface-cli download ml-jku/mazerunner --local-dir=./mazerunner --repo-type dataset\n```\n\nIn addition, the datasets for grid-worlds and our Procgen datasets are available on our webserver. \nDownload using:  \n```\n# Dark-Room + Dark KeyDoor\nwget --recursive --no-parent --no-host-directories --cut-dirs=2 -R \"index.html*\" https://ml.jku.at/research/RA-DT/downloads/minihack/dark_room\n# MazeRunner\nwget --recursive --no-parent --no-host-directories --cut-dirs=3 -R \"index.html*\" https://ml.jku.at/research/RA-DT/downloads/mazerunner/\n# Procgen 2M\nwget --recursive --no-parent --no-host-directories --cut-dirs=3 -R \"index.html*\" https://ml.jku.at/research/RA-DT/downloads/procgen_2M/\n# Procgen 20M\nwget --recursive --no-parent --no-host-directories --cut-dirs=3 -R \"index.html*\" https://ml.jku.at/research/RA-DT/downloads/procgen_20M/\n```\n\n## Running experiments\nIn the following, we provide illustrative examples of how to run the experiments in the paper. \n\n### Dark-Room\nTo train RA-DT (domain-specific), RA-DT (domain-agnosstic), and AD on Dark-Room 10x10 for 3 seeds each, run: \n```\n# RA-DT\npython main.py -m seed=42,43,44 experiment_name=darkroom_10x10_radt_v1 env_params=dark_room agent_params=radt_disc_icl run_params=finetune eval_params=pretrain_icl agent_params.load_path.file_name=dt_medium_64_new +agent_params.reinit_policy=True\n\n# RA-DT - domain-agnostic\npython main.py -m seed=42,43,44 experiment_name=darkroom_10x10_radt_lm_v1 env_params=dark_room agent_params=radt_disc_icl run_params=finetune eval_params=pretrain_icl agent_params.load_path=null +agent_params/retriever_kwargs=discrete_s_r_rtg +agent_params.retriever_kwargs.beta=10\n\n# AD\npython main.py -m seed=42,43,44 experiment_name=darkroom_10x10_ad_v1 env_params=dark_room agent_params=ad_gridworlds run_params=finetune eval_params=pretrain_icl +agent_params.replay_buffer_kwargs.n_ep_later=100\n```\n\nNote that domain-specific RA-DT requires to load a pre-trained checkpoint. Pre-train a DT using: \n```\npython main.py -m seed=42 experiment_name=darkroom_10x10_dt_pretrain_v1 env_params=dark_room agent_params=dt_gridworlds run_params=finetune eval_params=pretrain_icl_grids agent_params.huggingface.use_fast_attn=False +wandb_callback_params=pretrain\n```\nAfterwards, the model path can be passed to RA-DT via `agent_params.load_path`.\n\n### Dark Key-Door\nFor the same methods on Dark Key-Door 10x10, run:\n```\n# RADT + deduplication\npython main.py -m seed=42,43,44 experiment_name=dark_keydoor_10x10_radt_v1 env_params=dark_keydoor agent_params=radt_disc_icl agent_params/data_paths=dark_keydoor_10x10_train run_params=finetune eval_params=pretrain_icl agent_params.load_path='${MODELS_DIR}/minihack/dark_keydoor_10x10/dt_medium_64.zip' +agent_params.reinit_policy=True\n\n# RADT + domain-agnostic\npython main.py -m seed=42,43,44 experiment_name=dark_keydoor_10x10_radt_lm_v1 env_params=dark_keydoor agent_params=radt_disc_icl agent_params/data_paths=dark_keydoor_10x10_train run_params=finetune eval_params=pretrain_icl agent_params.load_path=null +agent_params/retriever_kwargs=discrete_s_r_rtg +agent_params.retriever_kwargs.beta=10\n\n# AD\npython main.py -m seed=42,43,44 experiment_name=dark_keydoor_10x10_ad_v1 env_params=dark_keydoor agent_params=ad_gridworlds agent_params/data_paths=dark_keydoor_10x10_train run_params=finetune eval_params=pretrain_icl +agent_params.replay_buffer_kwargs.n_ep_later=100\n```\n\n### MazeRunner\nFor the same methods on MazeRunner 15x15, run:\n\n```\n# RA-DT\npython main.py -m seed=42,43,44 experiment_name=mazerunner_radt_v1 env_params=mazerunner agent_params=radt_mazerunner run_params=finetune eval_params=pretrain_icl +agent_params.reinit_policy=True eval_params.eval_freq=100000 +eval_params.first_step=False eval_params.n_eval_episodes=30\n\n# RA-DT domain-agnostic\npython main.py -m seed=42,43,44 experiment_name=mazerunner_radt_lm_v1 env_params=mazerunner agent_params=radt_mazerunner run_params=finetune eval_params=pretrain_icl agent_params.load_path=null +agent_params/retriever_kwargs=discrete_r_rtg_mazerunner +agent_params.retriever_kwargs.beta=10 eval_params.eval_freq=100000 +eval_params.first_step=False eval_params.n_eval_episodes=30\n\n# AD\npython main.py -m seed=42,43,44 experiment_name=mazerunner_radt_ad_v1 env_params=mazerunner agent_params=ad_gridworlds agent_params/data_paths=mazerunner15x15 run_params=finetune eval_params=pretrain_icl agent_params.eval_context_len=800 agent_params.huggingface.max_length=400 agent_params.huggingface.n_positions=2400 +agent_params.replay_buffer_kwargs.n_ep_later=1000 eval_params.eval_freq=100000 +eval_params.first_step=False eval_params.n_eval_episodes=30\n```\n\n## Citation\nIf you find this useful, please consider citing our work: \n```\n@article{schmied2024retrieval,\n  title={Retrieval-Augmented Decision Transformer: External Memory for In-context RL},\n  author={Schmied, Thomas and Paischer, Fabian and Patil, Vihang and Hofmarcher, Markus and Pascanu, Razvan and Hochreiter, Sepp},\n  journal={ArXiv},\n  year={2024},\n  url={https://arxiv.org/abs/2410.07071}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fml-jku%2Fra-dt","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fml-jku%2Fra-dt","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fml-jku%2Fra-dt/lists"}