Heuristic Transformer: Belief Augmented In-Context Reinforcement Learning

Oliver Dippel, Bei Peng, Alexei Lisitsa

Keywords: In-context reinforcement learning (ICRL), Transformer models, Bayes-adaptive Markov decision processes (BAMDP), Belief modeling, Variational autoencoder (VAE)

TL;DR: We propose a new in-context reinforcement learning method that combines belief modeling over reward functions with transformer-based policies.

Abstract: Transformers have demonstrated exceptional in-context learning (ICL) capabilities, enabling applications across natural language processing, computer vision, and sequential decision-making. In reinforcement learning, ICL reframes learning as a supervised problem, facilitating task adaptation without parameter updates. Building on prior work leveraging transformers for sequential decision-making, we propose Heuristic Transformer (HT), an in-context reinforcement learning (ICRL) approach that augments the in-context dataset with a belief distribution over rewards to achieve better decision-making. Using a variational auto-encoder (VAE), a low-dimensional stochastic variable is learned to represent the posterior distribution over rewards, which is incorporated alongside an in-context dataset and query states as prompt to the transformer policy. We assess the performance of HT across the Darkroom, Miniworld, and MuJoCo environments, showing that it consistently surpasses comparable baselines in terms of both effectiveness and generalization. Our method presents a promising direction to bridge the gap between belief-based augmentations and transformer-based decision-making.

Area: Learning and Adaptation (LEARN)

Generative A I: I acknowledge that I have read and will follow this policy.

Submission Number: 29