Reinforcement Learning with Multi-Step Transition Predictions: A Bellman Theory

Chenbei Lu; Zaiwei Chen; Tongxin Li; Chenye Wu; Adam Wierman

cs.LGOct 23, 2026

Reinforcement Learning with Multi-Step Transition Predictions: A Bellman Theory

Authors: Chenbei Lu, Zaiwei Chen, Tongxin Li, Chenye Wu, Adam Wierman

Organizations: Duffield College of Engineering, Cornell University · Edwardson School of Industrial Engineering, Purdue University · School of Data Science, The Chinese University of Hong Kong (Shenzhen) · School of Science and Engineering, The Chinese University of Hong Kong (Shenzhen) ·

Abstract

Multi-step predictions of future transition realizations can improve reinforcement learning, but their value remains difficult to quantify. Directly including these predictions in the state can lead to exponential state-space growth, making direct dynamic programming computationally expensive. We study discounted finite MDPs under a block-feedback protocol: at the start of each K-period block, the decision maker observes $K$ -step probabilistic transition predictions for a subset of actions and adapts decisions to observed states within the block. We show that the optimal value satisfies a Bayesian Bellman equation on the original state space. This representation reveals that predictive information creates value by recursively changing the local order of maximization and expectation in the Bellman expansion. We formalize this effect through a Bellman--Jensen analysis and quantify the effects of prediction horizon, error, action coverage, and within-block feedback. We propose block-feedback Bayesian offline learning with online adaptation (BF--BOLA), a model-based algorithm for learning near-optimal policies, and establish upper and lower sample-complexity bounds that match up to logarithmic factors in the stated regimes. The results show that predictions can reduce transition learning to uncovered rows and that learning a near-optimal policy can require fewer prediction samples than estimating the Bayesian value function.

Figures & tables

Protocol	Information timing	Value representation	Exact computation or storage bottleneck
FH	One prediction per block; open-loop block actions	Block-start value on $\mathcal{S}$	NP-hard in worst case
BF	One prediction per block; state-adaptive block actions	Block-start value on $\mathcal{S}$	$O(K\|\mathcal{S}\|^{2}\|\mathcal{A}\|)$ within-block DP per realized prediction
RH	Prediction refreshed every period	Augmented value on $\mathcal{S}\times\mathcal{Q}_{K}$	NP-hard for $K\geq 2$ ( Pla et al., 2026a )

Table 1: Comparison of prediction protocols. The last column reports the cost of exact planning or tabular value storage.

Target	Environment samples $D_{1}$	Prediction samples $D_{2}$
Near-optimal policy	$\displaystyle\widetilde{\Theta}\!\left(\frac{\|\mathcal{S}\|(\|\mathcal{A}\|-\|\mathcal{A}^{-}\|)}{(1-\gamma)^{3}\varepsilon_{\rm pol}^{2}}\right)$	$\displaystyle\widetilde{\Theta}\!\left(1+\frac{\gamma^{2K}}{(1-\gamma)^{2}(1-\gamma^{K})\varepsilon_{\rm pol}^{2}}\right)$
Bayesian value estimation	$\displaystyle\widetilde{\Theta}\!\left(\frac{\|\mathcal{S}\|(\|\mathcal{A}\|-\|\mathcal{A}^{-}\|)}{(1-\gamma)^{3}\varepsilon_{\rm est}^{2}}\right)$	$\displaystyle\widetilde{\Theta}\!\left(\frac{1}{(1-\gamma)^{2}(1-\gamma^{K})\varepsilon_{\rm est}^{2}}\right)$

Table 2: Statistical rates under block feedback, up to logarithmic factors, in the matching regimes of Theorems 5.3 and 5.4 . The

D_{1}

rates omit the additive

|\mathcal{S}||\mathcal{A}|

reward observations.

Panel of Figure 1 — Figure 1: Value of transition predictions in the production MDP. The two noise levels in (a) are $0$ and $0.15$ . In (b), solid curves have full coverage and dashed curves predict type 0 only. We use $K=3$ in (a), (c), and (d), and channel noise $0.15$ in (b) and (d). Each gain is measured against its baseline in the same MDP. Error bars and bands are $1.96$ standard errors across 24 MDP instances.

Panel of Figure 2 — Figure 2: Two-source learning in the production MDP. The top row varies $D_{2}$ under full coverage. The bottom row varies $D_{1}$ under type-0 coverage or no prediction, with the prediction law provided. The budget $D_{1}$ includes 12 reward observations. Policy success means loss at most $10^{-3}$ . Each budget uses 12 training runs on each of 10 MDP instances. Error bars and bands are $1.96$ standard errors across instance means.

Panel of Figure 3 — Figure 3: Optimal control in the data-calibrated wind-storage MDP, with capacity 4, per-period charge/discharge limit 1.5, and $\gamma=0.95$ . Panel (a) shows exact expected discounted partial sums under the infinite-horizon optimal policies. Panels (b,c) show paired percentage savings over the no-prediction policy for feedback protocols, including a one-step BF reference. Panels (a,b) use $K=3$ , except for that reference. Panels (a,c,d) use $\eta=0.15$ . Panel (d) times the model-specific Bellman solvers at the same value tolerance. Curves show means across 12 MDPs fitted to day-level bootstrap resamples. Bands and error bars show their 2.5th–97.5th percentiles.

Reinforcement Learning

Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.

See all 1,727 papers

Near-Optimal Reinforcement Learning with Multi-Step Transition Lookahead
Sep 10, 2026Corentin Pla, Hugo Richard, Marc Abeille +1Reinforcement Learning Approximation Algorithms
CREST, ENSAE · Criteo AI Lab FairPlay Joint Team
Optimal Multi-Reward Reinforcement Learning
Sep 29, 2026Zijun Chen, Zihan ZhangMarkov Decision Processes Sample-Efficient RL
Department of Computer Science and Engineering Hong Kong University of Science and Technology
End-to-End Efficient RL for Linear Bellman Complete MDPs with Deterministic Transitions
Mar 24, 2026Zakaria Mhammedi, Alexander Rakhlin, Nneka OkoloReinforcement Learning Linear Markov Decision Processes
Google Research, NYC · MIT

Sep 10, 2026J/K move · Enter open · S save

Near-Optimal Reinforcement Learning with Multi-Step Transition Lookahead

Corentin Pla, Hugo Richard, Marc Abeille, Vianney Perchet

CREST, ENSAE · Criteo AI Lab FairPlay Joint Team

Reinforcement Learning Approximation Algorithms Model-Based RL

Open paper AskPDF arXiv

We study reinforcement learning (RL) with transition look-ahead, where the agent may observe which states would be visited upon playing any sequence of actions before deciding its course of action. Although look-ahead can substantially improve achievable performance, [1] showed that optimal planning with multi-step transition look-ahead is NP-hard. However, this hardness was established using a discount factor close to one. It was therefore unknown whether the problem remains hard for every discount factor, and whether near-optimal planning can nevertheless be performed efficiently. We resolve both questions. First, we show that for every fixed discount factor, exact planning remains NP-hard. Second, we introduce a randomized polynomial-time approximation scheme for every fixed look-ahead depth. Third, we extend our approach to account for unknown transitions. We empirically validate the soundness of our results on the wind-farm storage-control benchmark of [2], showing that our approach, optimally accounting for -step look-ahead information, offers substantially better performance than existing algorithms. [1] Corentin Pla, Hugo Richard, Marc Abeille, Nadav Merlis, Vianney Perchet : On the Hardness of Reinforcement Learning with Transition Look-Ahead [2] Chenbei Lu, Zaiwei Chen, Tongxin Li, Chenye Wu, Adam Wierman : Reinforcement Learning with Imperfect Transition Predictions: A Bellman-Jensen Approach

Abstract

Reinforcement Learning

Explore similar work

Near-Optimal Reinforcement Learning with Multi-Step Transition Lookahead

Optimal Multi-Reward Reinforcement Learning

End-to-End Efficient RL for Linear Bellman Complete MDPs with Deterministic Transitions