# Learning How to Infer Partial MDPs for In-Context Adaptation and Exploration

arXiv: https://arxiv.org/abs/2302.04250 · HTML: https://mlsanity.com/paper/2302.04250/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 37 references that are in the map.
- First posted on arXiv 8 Feb 2023.

## Builds on (most cited first, 29 of 29)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Meta-Learning in Neural Networks: A Survey](https://mlsanity.com/paper/2004.05439/) (2020)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [Deep Exploration via Bootstrapped DQN](https://mlsanity.com/paper/1602.04621/) (2016)
- [RL2: Fast Reinforcement Learning via Slow Reinforcement Learning](https://mlsanity.com/paper/1611.02779/) (2016)
- [A Tutorial on Thompson Sampling](https://mlsanity.com/paper/1707.02038/) (2017)
- [Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables](https://mlsanity.com/paper/1903.08254/) (2019)
- [Learning to reinforcement learn](https://mlsanity.com/paper/1611.05763/) (2016)
- [Transformers Can Do Bayesian Inference](https://mlsanity.com/paper/2112.10510/) (2021)
- [Learning to Adapt in Dynamic, Real-World Environments Through Meta-Reinforcement Learning](https://mlsanity.com/paper/1803.11347/) (2018)
- [Noisy Networks for Exploration](https://mlsanity.com/paper/1706.10295/) (2017)
- [(More) Efficient Reinforcement Learning via Posterior Sampling](https://mlsanity.com/paper/1306.0940/) (2013)
- [In-context Reinforcement Learning with Algorithm Distillation](https://mlsanity.com/paper/2210.14215/) (2022)
- [Value Iteration Networks](https://mlsanity.com/paper/1602.02867/) (2016)
- [Deep Exploration via Randomized Value Functions](https://mlsanity.com/paper/1703.07608/) (2017)
- [Value Prediction Network](https://mlsanity.com/paper/1707.03497/) (2017)
- [Self-Attention Between Datapoints: Going Beyond Individual Input-Output Pairs in Deep Learning](https://mlsanity.com/paper/2106.02584/) (2021)
- [The Predictron: End-To-End Learning and Planning](https://mlsanity.com/paper/1612.08810/) (2016)
- [Learning to Induce Causal Structure](https://mlsanity.com/paper/2204.04875/) (2022)
- [Efficient Exploration through Bayesian Deep Q-Networks](https://mlsanity.com/paper/1802.04412/) (2018)
- [Retrieval-Augmented Reinforcement Learning](https://mlsanity.com/paper/2202.08417/) (2022)
- [Meta Reinforcement Learning with Latent Variable Gaussian Processes](https://mlsanity.com/paper/1803.07551/) (2018)
- [Alchemy: A benchmark and analysis toolkit for meta-reinforcement learning agents](https://mlsanity.com/paper/2102.02926/) (2021)
- [BBQ-Networks: Efficient Exploration in Deep Reinforcement Learning for Task-Oriented Dialogue Systems](https://mlsanity.com/paper/1711.05715/) (2017)
- [Structure Learning in Motor Control:A Deep Reinforcement Learning Model](https://mlsanity.com/paper/1706.06827/) (2017)

## Built on by (most cited first, 2 of 2)

- [Passive learning of active causal strategies in agents and language models](https://mlsanity.com/paper/2305.16183/) (2023)
- [Dynamics as Prompts: In-Context Learning for Sim-to-Real System Identifications](https://mlsanity.com/paper/2410.20357/) (2024)

## Abstract (arXiv)

To generalize across tasks, an agent should acquire knowledge from past tasks that facilitate adaptation and exploration in future tasks. We focus on the problem of in-context adaptation and exploration, where an agent only relies on context, i.e., history of states, actions and/or rewards, rather than gradient-based updates. Posterior sampling (extension of Thompson sampling) is a promising approach, but it requires Bayesian inference and dynamic programming, which often involve unknowns (e.g., a prior) and costly computations. To address these difficulties, we use a transformer to learn an inference process from training tasks and consider a hypothesis space of partial models, represented as small Markov decision processes that are cheap for dynamic programming. In our version of the Symbolic Alchemy benchmark, our method's adaptation speed and exploration-exploitation balance approach those of an exact posterior sampling oracle. We also show that even though partial models exclude relevant information from the environment, they can nevertheless lead to good policies.