Beyond Markovian: Reflective Exploration via Bayes-Adaptive RL for LLM Reasoning

TL;DR

This paper introduces BARL, a Bayesian RL framework that enhances LLM reasoning by promoting reflective exploration, leading to improved performance and efficiency over traditional RL methods.

Contribution

The paper proposes a Bayesian RL approach for LLMs that encourages self-reflection and information gathering, addressing limitations of conventional RL in fostering reflective behaviors.

Findings

BARL outperforms traditional RL on reasoning tasks.

It achieves higher test accuracy and token efficiency.

Empirical results validate the effectiveness of Bayesian reflective exploration.

Abstract

Large Language Models (LLMs) trained via Reinforcement Learning (RL) have exhibited strong reasoning capabilities and emergent reflective behaviors, such as rethinking and error correction, as a form of in-context exploration. However, the Markovian policy obtained from conventional RL training does not give rise to reflective exploration behaviors since the policy depends on the history only through the state and therefore has no incentive to enrich identical states with additional context. Instead, RL exploration is only useful during training to learn the optimal policy in a trial-and-error manner. Therefore, it remains unclear whether reflective reasoning will emerge during RL, or why it is beneficial. To remedy this, we recast reflective exploration within a Bayesian RL framework, which optimizes the expected return under a posterior distribution over Markov decision processes induced by the training data. This Bayesian formulation admits uncertainty-adaptive policies that, through belief updates, naturally incentivize information-gathering actions and induce self-reflection behaviors. Our resulting algorithm, BARL, instructs the LLM to stitch and switch strategies based on the observed outcomes, offering principled guidance on when and how the model should reflectively explore. Empirical results on both synthetic and mathematical reasoning tasks demonstrate that BARL outperforms conventional RL approaches, achieving superior test-time performance and token efficiency. Our code is available at https://github.com/shenao-zhang/BARL.