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

2025-05-26

AI Summary

Traditional Markovian reinforcement learning (RL) suffers from restrictive assumptions of state Markovness and deterministic policy optimization, hindering modeling of dynamic reflective behavior and resulting in poor backtracking and error-correction capabilities at test time. This work introduces Bayesian adaptive RL to large language model (LLM) inference training—its first application in this context—formalizing reflective exploration as online Bayesian updating of environment beliefs and dynamic policy switching or stitching. By optimizing posterior Markov decision processes (MDPs) and implementing belief-driven policy control, our approach explicitly models and incentivizes reflection during training, overcoming limitations of history compression. Evaluated on synthetic and mathematical reasoning tasks, our method significantly outperforms standard Markovian RL: it improves token efficiency and exploration effectiveness at test time, while demonstrating superior backtracking, error correction, and contextual adaptation capabilities.

Technology Category

• Machine Learning: Reinforcement Learning
• Natural Language Processing: (Large) Language Models
• Humans and AI: Human-Aware Planning and Behavior Prediction

Application Category

• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

Large Language Models (LLMs) trained via Reinforcement Learning (RL) have exhibited strong reasoning capabilities and emergent reflective behaviors, such as backtracking and error correction. However, conventional Markovian RL confines exploration to the training phase to learn an optimal deterministic policy and depends on the history contexts only through the current state. Therefore, it remains unclear whether reflective reasoning will emerge during Markovian RL training, or why they are beneficial at test time. To remedy this, we recast reflective exploration within the Bayes-Adaptive RL framework, which explicitly optimizes the expected return under a posterior distribution over Markov decision processes. This Bayesian formulation inherently incentivizes both reward-maximizing exploitation and information-gathering exploration via belief updates. 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 standard Markovian RL approaches at test time, achieving superior token efficiency with improved exploration effectiveness. Our code is available at https://github.com/shenao-zhang/BARL.

Problem

• Overcoming Markovian RL's limitation in reflective reasoning emergence
• Optimizing exploration and exploitation via Bayes-Adaptive RL framework
• Enhancing LLM reasoning efficiency with strategic reflective exploration

Innovation

• Bayes-Adaptive RL for reflective exploration
• Optimizes return under posterior MDP distribution
• Stitches strategies based on observed outcomes