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

Thank you for the thoughtful and encouraging review. We appreciate your positive comments on the motivation, clarity, novelty, and theoretical contributions of the work. Below, we address your questions in detail.

Q1. I’m a bit confused about the sample complexity D1. D1 is a function of |S|(|A| - |A-|) but also from what the text states |S||A|. The classical O(|S| |A|) state dependence is still in D1. Even if the first term vanishes, the second term still exists. Thus I’m unsure if the claim that BOLA provides better sample complexity guarantees than regular model-based RL algorithms is correct. I will ask for clarification. Maybe this is meant to say there is no epsilon scaling on that part of the dependence?

Response: Thank you for pointing this out, and we appreciate the opportunity to clarify. You are absolutely right: our claim refers specifically to the $\epsilon$ scaling in the sample complexity. The classical minimax sample complexity for model-based RL is $\mathcal{O}(|S||A|\epsilon^{-2})$ [1], with model-free methods typically requiring even more samples. In contrast, our term $D_1$ has a lower complexity $\mathcal{O}(|S|(|A| - |A^{-}|)\epsilon^{-2} + |S||A|)$. The key distinction is that only the first term in $D_1$ depends on $\epsilon^{-2}$, while the second term, $\mathcal{O}(|S||A|)$, does not scale with $\epsilon$ and is often negligible. We will make this clarification explicit in the revision to avoid any confusion.

[1] Azar, M. G., Munos, R., & Kappen, H. (2012). On the Sample Complexity of Reinforcement Learning with a Generative Model. In International Conference on Machine Learning.

Q2. A fundamental question is whether or not these algorithms provide any benefits over traditional RL algorithms. I think this answer could be a bit stronger. In the end, the upper bound in the worst case must still depend on the state-action space size which in practice is quite limiting.

Response: Yes, our algorithm can be more sample-efficient, especially when predictions are available (as discussed in our response to Q1).

• Theoretically, our sample complexity is $\mathcal{O}(|\mathcal{S}|(|\mathcal{A}| - |\mathcal{A}^-|)\epsilon^{-2})$, which is lower than the minimax sample complexity $\mathcal{O}(|\mathcal{S}||\mathcal{A}|\epsilon^{-2})$ in tabular MDPs [1]. Intuitively, for actions $a \in \mathcal{A}^-$, transition dynamics are directly obtained from predictions $\sigma$, so increasing $|\mathcal{A}^-|$ leads to improved efficiency. The $|\mathcal{S}||\mathcal{A}|$ scaling can not be avoided in statistical theory (because we at least need to sample all the reward function $r(s,a)$), but our framework can avoid the $|\mathcal{S}||\mathcal{A}|\epsilon^{-2}$ scaling of standard MDPs.

[1] Azar, M. G., Munos, R., & Kappen, H. (2012). On the Sample Complexity of Reinforcement Learning with a Generative Model. In International Conference on Machine Learning.

• Empirically, our framework scales naturally to large or continuous state spaces using function approximation, similar to how standard MDPs overcome the sample complexity dependence on state-action space size. The key idea is to approximate the Bayesian value function $V^{\text{Bayes}}(s)$ using function approximation (e.g., neural networks), trained via fitted value iteration. This is feasible because $V^{\text{Bayes}}(s)$ shares the same dimensionality and Bellman structure as the standard value function, with the only difference being the use of prediction $\sigma$. Once $V^{\text{Bayes}}(s)$ is learned, the optimal policy can be extracted via standard planning techniques. We provide an example of realizing the algorithm in our setting as follows:

Consider the case where the forecast horizon $K=1$. Given a dataset of transitions $\mathcal{D} = {(s_i, a_i, r_i, s_i')}$ and a dataset of one-step predictions $\mathcal{D}\sigma$, we aim to learn a Bayesian value function $\hat{V}^{\text{Bayes}}_\theta(s)$ using fitted value iteration:

• Step1: Sample a batch of states $s$ from $\mathcal{D}$
• Step 2: For each $s$, compute the Bellman target: $y(s) = \mathbb{E}{\sigma \sim \mathcal{D}\sigma} \left[ \max_{a \in \mathcal{A}} \left( r(s, a) + \gamma \cdot \frac{1}{M} \sum_{m=1}^M V_\theta(s^{(m,\sigma)}) \right) \right]$,

where $s^{(m,\sigma)} \sim \hat{P}\phi(\cdot \mid s, a, \sigma)$ are next states sampled from a learned transition model $\hat{P}\phi$ using forecast $\sigma$.

• Step 3: Update parameter $\theta$ of the value function by minimizing the squared error $\min_\theta \sum_{s \in \text{batch}} \left(V_\theta(s) - y(s)\right)^2$
• Step 4: Repeat from Steps 2-3 until convergence.

Q3. The analysis and limitations ignore the cost of obtaining the prediction oracle. In the worst case, it seems one may have to learn a full model of the transition dynamics to obtain such a model before the algorithm can be run. In this case, it seems this type of algorithm would be more expensive than a traditional RL algorithm.

Response: Thank you for raising this important point. In our analysis, we do not assume that prediction oracles are free; rather, we treat querying the oracle as incurring the same per-sample cost as interacting with the environment (e.g., one sample for each $\sigma$).

• In many practical settings, such as electricity markets, weather systems, or queueing networks, forecast information is externally available through domain-specific models and does not need to be learned by the RL agent. Moreover, generating such forecasts is often much less costly than estimating full transition dynamics via trial-based exploration. For instance, in a discrete-time M/M/1 queue, forecasting only the arrival probability $\lambda$ and service probability $\mu$ is sufficient to compute all transition probabilities analytically, such as $P(s_{t+1} = j+1 \mid s_t = j) = \lambda(1 - \mu)$ (i.e., an arrival occurs but no service).

• Our framework is also designed to flexibly support partial prediction (see Section 2.1), where only a subset of actions $\mathcal{A}^-$ are associated with transition forecasts. In this case, the agent only needs to learn the dynamics for the remaining actions $\mathcal{A} \setminus \mathcal{A}^-$, and we provide corresponding sample complexity guarantees.

• Even under the most conservative assumption where the prediction $\sigma$ must be queried for each transition entry (e.g., $|S||\mathcal{A}^-|$ samples for each $\sigma$), the total cost including both oracle queries ($D_2 = \mathcal{O}(|\mathcal{S}||\mathcal{A}^-|\epsilon^{-2})$) and environment samples ($D_1 = \mathcal{O}(|\mathcal{S}|(|\mathcal{A}| - |\mathcal{A}^-|)\epsilon^{-2})$) still matches the standard minimax sample complexity of model-based RL: $\mathcal{O}(|\mathcal{S}||\mathcal{A}|\epsilon^{-2})$. Thus, our framework is never worse than classical MDPs in theory, and it can provide substantial benefits when more forecasts are available in practice.

Minor: It might make sense for clarity purposes to rename the sub-gaussian variable since it shares notation with the predictions.

Response: Thanks for the suggestion, we will change the notation for clarity.

Again, thank you for your thoughtful review. We would be happy to further clarify any points if there are remaining questions during the discussion.

Dear authors,

thank you for clarifying my point about the state-action dependence. This was really my main concern that must be addressed in a final version of this paper as the claim as stated in the manuscript right now is not correct. I trust the authors will faithfully do so and am willing to raise my score to accept. I think this paper is an interesting contribution and I enjoyed learning something from it. Given that I believe I understood this point correctly in the first place, I raised my confidence to 4.

Now to the other points that really in my eyes don't really require discussion and should be seen as feedback. It simply might make sense to think about these and decide whether or not it is warranted to highlight them as limitations in some form of discussion.

The statement "our algorithm can be more sample efficient" misses the point I am making in my comment. One of the biggest challenges in empirical RL is dealing with large or infinitely sized state-action spaces. If our cost depends on the state-action space it is likely infeasible in practice irrespective of any epsilon dependence. This is a limitation of this approach. I understand that the proposed Bellman operator can potentially be iterated on using function approximation. Whether or not that is in fact more sample efficient is not demonstrated in this work and it is out of scope.

Secondly, I did not say that the paper argues that oracles are free. I said it ignores the fact that one needs to have an oracle. If it is easy to obtain an oracle, and querying an oracle is only computationally slightly more expensive, practitioners would likely simply run the computer for longer. In such cases, one might even argue that the cost of executing the oracle is lower than obtaining an actual data sample. That is the promise of model-based RL approaches after all. Yet in many interesting scenarios one needs to first learn an oracle from data. In practice, one could attempt to learn an oracle from data using neural networks for instance. Whether or not learning a forecaster and then applying an approach such as this is more or less useful that straight-up learning the standard value function is unclear to me. Again, I don't expect this to be answered in this manuscript.

We thank the reviewer for the detailed and thoughtful review. We're encouraged by your positive evaluation of our theoretical contributions and the clarity of our framework. We also appreciate your constructive suggestions. We address each of these points below.

Q1. Can other assumptions on forecasts, like parametric transition distributions, be accommodated within the current framework, perhaps as part of a broad spectrum of assumptions with parametric model on one extreme and full transition matrix forecast on the other?

Response: Thank you for raising this insightful point. Our framework is flexible enough to accommodate a wide range of forecasting assumptions from full transition matrices to parametric models. As long as the forecast can be expressed in terms of (or mapped to) some transition probabilities, either explicitly or via a compact parameterization, it can be seamlessly integrated into our approach.

• As the reviewer suggests, many real-world systems such as queueing and inventory control exhibit structured dynamics that can be characterized using only a few parameters. For example, in the classic discrete-time M/M/1 queue, the entire transition matrix can be fully determined by forecasting just two parameters: the arrival probability $\lambda$ and the service probability $\mu$. This allows us to compute all probabilities like $P(s_{t+1} = j+1 \mid s_t = j) = \lambda(1 - \mu)$ (i.e., an arrival occurs but no service), without needing to explicitly give the full matrix.

• Moreover, the required forecast model can be incomplete. our framework also supports partial forecasts with predictable action set $\mathcal{A}^-$, where only a subset of transitions can be predicted.

Q2. Does the framework support continuous or large state spaces (perhaps connected to question 1 above)? What modifications to BOLA will be needed to avoid the blowup in compute with value iteration?

Response: Yes, our framework can indeed be extended to large or even continuous state settings. The key is to learn the high-dimensional Bayesian value function $V^{{Bayes}}(s)$. To achieve this, we can use function approximation (e.g., neural networks) to learn $V^{Bayes}(s)$ via fitted value iteration. It is achievable because the Bayesian value function $V^{\text{Bayes}}(s)$ has the same dimensionality and quite similar definition (only different in prediction $\sigma$) of value function in standard MDPs. With the learned Bayesian value function, we can calculate the optimal policy follow a similar routine.

At a high level, our method remains applicable in more general settings because it preserves similar fundamental notions of Bellman optimality condition and optimal policies shared with classical MDPs.

Q3. The analysis on the Bellman-Jensen gap seems similar to the work of Mercier et. al. [1] on anticipatory sampling, which has been widely known in stochastic planning. It might be good to cite this (and related) work, and ideally discuss the main limitations of their analysis on the so-called "anticipatory gap" and why it is difficult to apply directly towards forecasts.

Response: Thank you for the helpful suggestion. The anticipatory gap studied by Mercier et al. [1] is indeed closely related to our Bellman-Jensen gap when the prediction window is limited to $K=1$. However, our setting is significantly more general and challenging, in several key ways:

• We allow arbitrary prediction windows $K \geq 1$;
• We consider maximizing infinite-horizon cumulative reward objectives, not just terminal rewards;
• We allow for incomplete predictions regarding subset $\mathcal{A}^-$ and prediction errors.

Moreover, the analysis of the gap in [1] is primarily case-based and lacks general theoretical characterization. As acknowledged by the authors themselves, “We are currently applying this method on more complex problems but proofs quickly become very cumbersome. As an alternative, we discuss another, empirical way to argue that the GAG of a problem is small”, they resort to empirical justification rather than formal analysis.

In contrast, our work provides a general, explicit bound on the Bellman-Jensen gap for arbitrary MDPs and any prediction window $K$. A key technical challenge lies in characterizing the recursive sum of coupled Jensen gaps across multi-step transitions. To address this, we introduce a novel dyadic horizon decomposition and a variance-based analysis (both in Appendix E) that enables us to decouple the multi-step dependencies in the Bellman operator.

Analytically quantifying the value of prediction is a long-standing challenge, especially for unstructured problems. Most existing results rely on strong assumptions—e.g., linear dynamics as in LQR settings [2]. To our knowledge, our work provides the first general theoretical analysis of the value of prediction in MDPs with arbitrary forecast lengths and incomplete information.

[1] Mercier, Luc, and Pascal Van Hentenryck. "Performance Analysis of Online Anticipatory Algorithms for Large Multistage Stochastic Integer Programs." IJCAI. 2007.
[2] Yu, C., Shi, G., Chung, S.-J., Yue, Y., and Wierman, A. (2021). The power of predictions in online control. Advances in Neural Information Processing Systems, 33:1994–2004

Dear authors,

Thank you for clarifying the points in the rebuttal.

In particular, I now understand the limitations of the Mercier et al., work and indeed I agree that the current framework and theoretical analysis applies in more general settings.

I think the current paper is technically solid, and I have learned a lot from the manuscript and subsequent discussions. so I will maintain my current score.

I think clarification of the most important points discussed here would significantly strengthen the overall write up, in addition to the asymptotic analysis clarification suggested by reviewer jBsU:

1. The learnability aspects and asymptotic bias arguments raised by reviewer DyeW and in the rebuttal are quite interesting. I wonder whether this is also related to the idea of an "uninformative action" in MDP parameter learning? Specifically, in your queueing example, it is impossible to acquire any new information about the arrival rate if no customers are admitted to the queue. Often dealing with such settings requires pretty strong assumptions on the structure of the MDP w.r.t which the current work is somewhat agnostic to. I think it would be interesting to discuss what the expected result of the algorithm and potential limitations would be.
2. The parameter learning setting as discussed still seems more practically useful, particularly in large state/action spaces. I think including this example up front, and more generally, common uses cases/problems the current framework was designed to handle, would strengthen the writeup.

Once again, thank you for addressing the points in your rebuttal.

Thank you very much for your positive and thoughtful review of our paper. We greatly appreciate your recognition regarding our novelty, theoretical results, experiments, and writing. Below, we address your minor points of confusion:

Q1. Equation 1 should be $P(s'|s,a,\sigma_k^*)$?

Response: Yes, that is correct. Thank you for catching this typo!

Q2. What does it mean to condition on $\sigma$ when $\sigma = {\sigma}^* +\epsilon$ is imperfect (for example, in the expression of Equation 5.) Do we treat $\sigma$ here as $\sigma^*$, ignoring the imperfection? In general it's not crystal clear what it means to condition on $\sigma$ (e.g. even in the value function definition of Equation 3). Can the authors be more formal about the probability kernels involved?

Response: Thank you for raising this point. In our framework, we condition on the imperfect transition prediction $\sigma$, not the perfect one ${\sigma}^*$. Specifically, throughout the paper (including in Equations 3 and 5), we use the imperfect transition probabilities defined as $P(s'|s,a,\sigma_k) = \sigma_k((s,a),s')$. This ensures that the agent makes decisions based solely on the imperfect predictions available to it (the agent never observes the perfect transitions $\sigma^*$). We agree that the confusion may have arisen due to the typo in Equation 1.

Q3. The current algorithm 1 samples transitions and predictions and does a plug-in estimator to learn hat{V}^{Bayes}_epsilon that approximates V^{Bayes}_epsilon. Intuition suggests that sampling more variables could eventually give us a good approximation to accurate transition prediction, because by concentration, $\sigma = \sigma^{*} + \epsilon$, and we should thus be learning unbiased value function. What is happening here?

Response: Thank you for the question. In our approach, we sample imperfect transition predictions $\sigma$ to approximate their distribution $f_{\sigma}$. This distribution reflects forecast errors of size $\epsilon$, and is different from the true distribution $f_*$ over accurate transitions $\sigma^*$.

Importantly, any accurate prediction $\sigma^*$ is a deterministic (binary) transition matrix, where each row is a one-hot vector. In contrast, the noisy forecasts $\sigma$ may contain non-binary probabilities (like 0.1, 0.9) due to errors. Thus, even with many samples from $f_{\sigma}$, we do not recover the distribution of accurate predictions. Hence, the learned Bayesian value function will converge, but to a biased solution. This is different from using transition samples to recover the accurate transition kernel using empirical average.

Two additional questions:

Q4. It feels like the BOLA algorithm can be extended to settings beyond the tabular one. At a high level the two components involve: estimating $P(\sigma)$, and estimating the rest of the model (r and $P(\sigma)$). Can we indeed say that this methodology works in any setting where the above quantities are learnable?

Response: Yes, the two-stage BOLA algorithm can indeed be extended beyond the tabular setting, as long as $P(\sigma)$ and other relevant system models (such as $r$ and $P(\sigma)$) are learnable. To extend to continuous state settings, we can use function approximation (e.g., neural networks) to learn the continuous Bayesian value function $V^{\text{Bayes}}(s)$. This can be done via fitted value iteration, leveraging the structural similarity between our Bayesian Bellman equation and that of standard MDPs.

At a high level, BOLA remains applicable in more general settings because it preserves the fundamental notions of value functions and optimal policies shared with classical MDPs.

Q5. There is a trade-off here between the statistical and computational costs of learning transitions over K steps (which gets more expensive due to an exponential scaling in K) and the quality of the learned policy (since more steps gives a better policy). Is there any interesting way to adaptively choose K to balance these considerations?

Response: This is an excellent question. As shown by both Theorem 4.1 and our experiments, the marginal benefit of extending the prediction horizon $K$ decreases exponentially, while the computational cost increases exponentially. To handle this trade-off, a simple yet effective strategy is to incrementally increase $K$ (e.g., starting from $K=1$) until the computational budget is met. Since the total trial cost $\sum_{k=1}^{K} D^k \approx D^K$ (where $D^k$ is the computational cost with prediction window $k$), this approach naturally identifies the optimal $K$ that balances reward performance and computational cost, and is nearly optimal in terms of total trial cost. In practice, this greedy search often achieves a good performance.