Reinforcement Learning with Imperfect Transition Predictions: A Bellman-Jensen Approach

TL;DR

This work addresses reinforcement learning with imperfect, multi-step transition predictions by introducing a tractable Bayesian value function that marginalizes over prediction randomness, avoiding exponential state augmentation. It develops a Bellman-Jensen Gap framework to quantify the value of look-ahead predictions and derive explicit performance bounds under finite horizon, partial predictability, and prediction errors. Building on these insights, the authors propose BOLA, a two-stage model-based algorithm that learns offline Bayes values and online adapts to real-time predictions, with finite-sample guarantees and improved sample efficiency when predictions are informative. Experiments on synthetic MDPs and a wind-energy storage control task validate the theory and show practical gains from longer horizons and robust performance across prediction error regimes.

Abstract

Traditional reinforcement learning (RL) assumes the agents make decisions based on Markov decision processes (MDPs) with one-step transition models. In many real-world applications, such as energy management and stock investment, agents can access multi-step predictions of future states, which provide additional advantages for decision making. However, multi-step predictions are inherently high-dimensional: naively embedding these predictions into an MDP leads to an exponential blow-up in state space and the curse of dimensionality. Moreover, existing RL theory provides few tools to analyze prediction-augmented MDPs, as it typically works on one-step transition kernels and cannot accommodate multi-step predictions with errors or partial action-coverage. We address these challenges with three key innovations: First, we propose the \emph{Bayesian value function} to characterize the optimal prediction-aware policy tractably. Second, we develop a novel \emph{Bellman-Jensen Gap} analysis on the Bayesian value function, which enables characterizing the value of imperfect predictions. Third, we introduce BOLA (Bayesian Offline Learning with Online Adaptation), a two-stage model-based RL algorithm that separates offline Bayesian value learning from lightweight online adaptation to real-time predictions. We prove that BOLA remains sample-efficient even under imperfect predictions. We validate our theory and algorithm on synthetic MDPs and a real-world wind energy storage control problem.

Figures (5)

• Figure 1: Wind Farm Storage Control
• Figure 2: Storage Control Performance. (a) Cumulative cost reduction for different prediction horizons: longer prediction horizon yields greater savings over the no-prediction baseline. (b) Robustness to prediction noise: cost savings decline roughly linearly with error, yet all predictive policies outperform the baseline even at 30% noise.
• Figure 3: Wind Farm Storage Control
• Figure 4: Prediction Helps More in Low-Value States. (a) The expected return at a low-value state increases significantly with prediction horizon $K$. (b) The return at a high-value state also improves with $K$, but the marginal gain is much smaller.
• Figure 5: Sample Efficiency of Learning with Predictions. (a) Convergence of (Bayesian) value function estimation error. (b) Policy performance measured by expected return. BOLA with predictions achieves similar convergence and consistently higher performance across varying sample sizes compared to classical model-based RL.

Theorems & Definitions (15)

• Theorem 3.1: Bellman Optimality Equation for Bayesian Value Function
• Corollary 3.2: Optimal Policy with Bayesian Value Function and Transition Predictions
• Definition 4.1: Maximal Bellman-Jensen Gap
• Theorem 4.1: Bellman-Jensen Performance Bound
• Corollary 4.2
• Theorem 5.1
• Lemma E.1: Proof in Appendix \ref{['proof-thm2-lemma1']}
• Lemma E.2: Proof in Appendix \ref{['proof-thm2-lemma2']}
• Lemma E.3: Proof in Appendix \ref{['proof to lemma4']}
• Lemma E.4: Proof in Appendix \ref{['Proof of Lemma Key_lemma_inf_off']}