Heuristic transformer: belief augmented in-context reinforcement learning

Transformers have recently shown that reinforcement learning can be reframed as an in-context prediction problem, allowing agents to adapt to new tasks without updating their parameters. However, existing in-context RL approaches rely solely on past trajectories as prompts, leaving the model to implicitly infer uncertainty about the reward structure from raw experience alone. We introduce the Heuristic Transformer (HT), an in-context RL method that explicitly augments the prompt with a learned belief over rewards. A low-dimensional stochastic latent variable captures the posterior distribution over rewards and is provided to the transformer alongside trajectories and query states. This enables the model to reason directly over uncertainty rather than inferring it indirectly from data. Across Darkroom, Miniworld, and MuJoCo benchmarks, HT consistently outperforms existing in-context RL baselines in both performance and generalization, particularly in stochastic settings. Our results show that combining belief-based representations with transformer policies is a powerful mechanism for improving in-context decision-making, and suggest a new direction for integrating probabilistic reasoning into transformer-based RL.