Latent Trajectory: A New Framework for Actor-Critic Reinforcement Learning with Uncertainty Quantification

Authors: Frank Shih, Faming Liang
Primary area: reinforcement learning
Keywords: Reinforcement learning, Stochastic gradient MCMC, Bayesian sampling, Uncertainty quantification

Abstract:
Uncertainty quantification for deep neural networks is crucial for building reliable modern AI models. This challenge is particularly pronounced in deep reinforcement learning, where agents continuously learn from their interactions with stochastic environments, and the uncertainty of the value function is a key concern for ensuring reliable and robust RL applications. The complexity increases in actor-critic methods, as the training process alternates between optimizing the actor and critic networks, whose optimization nature makes the uncertainty of the value function hard to be quantified. 
To address this issue, we introduce a novel approach to RL training that conceptualizes transition trajectories as latent variables. Building on this framework, we propose an adaptive Stochastic Gradient Markov Chain Monte Carlo (SGMCMC) algorithm for training deep actor-critic models. This new training method allows for the implicit integration of latent transition trajectories, resulting in a trajectory-independent training process. We provide theoretical guarantees for the convergence of our algorithm and offer empirical evidence showing improvements in both performance and robustness of the deep actor-critic model under our Latent Trajectory Framework (LTF). Furthermore, this framework enables accurate uncertainty quantification for the value function of the RL system, paving the way for more reliable and robust RL applications.

Decision: reject
Reviewer scores: 3, 5, 3, 3

Reviews and discussion:
<thread>
creation date:
1730189682514

thread_type:
official_review

<thread_object>
creator: reviewer_8ewq

<content>
summary:
This paper introduces a Latent Trajectory Framework (LTF) to improve uncertainty quantification in deep actor-critic reinforcement learning, addressing the challenge of value function uncertainty in stochastic environments. Using an adaptive Stochastic Gradient Markov Chain Monte Carlo (SGMCMC) algorithm, the method enables trajectory-independent training, backed by theoretical convergence guarantees and empirical performance improvements. This approach enhances both the robustness and reliability of RL applications by integrating latent transition trajectories.

strengths:
Theoretical Analysis: The paper provides a theoretical analysis of the Latent Trajectory Framework, attempting to establish convergence and benefits for uncertainty quantification in RL.

weaknesses:
**$.1**
- L47: Why is the actor considered "unknown"? We can access its weights during training, so we should be able to evaluate it on any state-action pair. Even if the actor were unknown, how does that affect uncertainty quantification?
General: Could you clarify the type of uncertainty you’re addressing?

**$.2**
- L100: What is π(x∣θ)? I thought π represented the policy, a distribution over the action space.
- L103: What is π(ψ∣θ)?
- L105: What does "pseudo-population size" mean? Is N not equal to the batch size n?
- L107: Similar to above, this line is unclear.
- Eq(3): Could you provide the intuition behind this learning objective?

**$.3**
- How does this approach differ from the SGMCMC method discussed in Shih & Liang (2024)?

**$.4**
- $4.1: The writing lacks organization. For instance, metrics are introduced at L323, but the computation details are only explained two paragraphs later. Why is coverage rate the chosen metric for uncertainty quantification?
- $4.2: Since actor-critic methods are typically used for continuous action spaces, why not use Mujoco benchmarks for Fig. 7? Additionally, can LT be extended to SAC or other recent methods?

**Overall**
- This paper heavily relies on prior work for explanations and notations, which makes it challenging for readers unfamiliar with the domain to follow.

questions:
none
</content>
</thread_object>

</thread>
<thread>
creation date:
1730372630583

thread_type:
official_review

<thread_object>
creator: reviewer_y7uj

<content>
summary:
The authors propose a novel method for characterising uncertainty in value functions using a stochastic gradient MCMC algorithm. They carry out a convergence analysis of their method before evaluating PyBullet and Gridworld environments.

strengths:
The theoretical analysis of the proposed algorithm seems sounds from a cursory readying. Convergence guarantees are always welcome in papers.

weaknesses:
My main concern relates to lack of positioning of the paper. Bayesian RL offers a precise way to characterise uncertainty in an MDP. At every timestep, the posterior over uncertain variables updates according to a Bayesian Bellman operator. Uncertainty can be characterised in the state-reward transitions as in model-based approaches or other sufficient variables like value functions or Bellman operators as in model-free approaches[4].  There is already a wealth of literature in uncertainty quantification in value functions. See my question below to authors.

The authors claim `Notably, uncertainty quantification for value-functions is generally beyond the reach of conventional iterative optimization algorithms used to train actor-critic models'. This is not true. Methods such optimistic actor-critic [1], BBAC[2] and EVE[3] (to name but a few) have been able to quantify uncertainty in value functions when used in continuous control for some time now as uncertainty quantification is essential to their exploration methods. There also exist analyses of the various approximate inference tools used to quantify uncertainty [5]. See [6] for a recent comparison of state of the art continuous control using uncertainty quantification evaluated in a variety of domains. 

Empirical weaknesses:

There is a significant lack of comparison to other methods that quantify uncertainty in value functions. A comparison of these methods seems essential to evaluate the contribution of the proposed method. Moreover, the authors don't indicate number of timesteps in their evaluations so it is difficult to gauge the worth of their approach in comparison to similar Bayesian methods. 

[1] Coisek et al., Better Exploration with Optimistic Actor-Critic, 2019, https://arxiv.org/pdf/1910.12807
[2] Fellows et al., Bayesian Bellman Operators, 2023, https://arxiv.org/pdf/2106.05012
[3] Schmitt et al., Exploration via Epistemic Value Estimation, 2023, https://arxiv.org/pdf/2303.04012
[4] Fellows et al., Bayesian Exploration Networks, 2024 https://arxiv.org/pdf/2308.13049
[5] Coisek et al., Conservative Uncertainty Estimation By Fitting Prior Networks, 2020 https://openreview.net/pdf?id=BJlahxHYDS
[6]  Tasdighi et al., Deep Exploration with PAC-Bayes, 2025, https://arxiv.org/pdf/2402.03055

questions:
How does the proposed method fit into the general model-free Bayesian RL framework that precisely characterises uncertainty in value functions? How does $\pi_\mathcal{N}(\psi\vert \theta)$ relate to the posterior over $g_\psi$ under these frameworks? If it differs from existing characterisations, ie [1]-[4], what theoretical, algorithmic or empirical advantages does it offer over a full Bayesian approach for characterising uncertainty in value functions? 

Can the authors approach be used to derived Bayes-optimal policies? If not, why does their uncertainty quantification prevent this? 

Can the authors extend their empirical evaluation to include other methods that quantify uncertainty in their value functions?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730676534745

thread_type:
official_review

<thread_object>
creator: reviewer_5ds5

<content>
summary:
Uncertainty quantification of the value function is crucial for a robust reinforcement learning algorithm. This work considers the challenging actor-critic setting. The introduced latent trajectory framework (LTF) is built upon the adaptive Stochastic Gradient Markov chain Monte Carlo (SGMCMC) by treating the transition trajectory and the value function parameter as latent variables for policy optimization. The proposed method is theoretically proved to be able to converge under mild conditions. The experiments on indoor escape environments  and PyBullet environment show that the proposed method has better performance compared with baseline A2C algorithms.

strengths:
- Theoretical results: The proposed LTF is grounded by theoretical convergence results. The paper contains most of the proof details, and the logic in the writing is easy to follow. 
- Experiments: The experiments show multiple metrics for evaluating the performance of the proposed LTF. For instance, the KL and MSE can help evaluate the performance from different perspectives. The results in HalfCheetah show that the proposed method can effectively improve the performance compared with A2C.

weaknesses:
- The proposed LTF only compares with the vanilla A2C method, while there have been many  AC-based methods proposed, such as [W1,W2]. I recognize that the proposed LTF introduces a new perspective by treating the value parameters as a latent variable, while the experiments are lacking. In particular, the experiments in Section 4.1 are conducted on rather new environments (by Liang et al.). The lack of comparison with other AC methods (e.g., in experiments and/or related works section) makes the results and performance gain less convincing. If the comparison is not possible or not necessary, please clarify the reasons. 
- The proposed LTF is largely based on SGMCMC algorithms that are proposed in Liang et al., 2022a; Deng et al.,2019, while the major contribution of the LTF is less clear. From my understanding, the main contribution is on the A2C settings, which pose unique challenges for uncertainty quantification. It will be very beneficial for the authors to explicitly state the key challenges of applying SGMCMC to A2C settings and how their approach addresses these challenges. 




[W1] Zhou et al. "Natural actor-critic for robust reinforcement learning with function approximation." Advances in neural information processing systems 36 (2024).

[W2] Wu, et al. "Uncertainty weighted actor-critic for offline reinforcement learning." arXiv preprint arXiv:2105.08140 (2021).

questions:
- What is the parameter $\delta_j$ in Eqn. 10?
- What are the relation between $\epsilon_{k,l}$ and $\epsilon_{k} $ in Algorithm 1
- How to select $\mathcal{N}$ in practice?
- Why the variance in Figure 7 for LT-A2C is not significantly reduced if the uncertainty is effectively evaluated during update (e.g., HopperBullet, Evaluation reward)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730685879522

thread_type:
official_review

<thread_object>
creator: reviewer_pdcr

<content>
summary:
This paper introduces the latent trajectory framework (LTF) that implicitly models the uncertainty of Q-functions by drawing multiple samples of critic parameters, essentially forming a distribution over Q-values.

strengths:
- The paper provides theroretical justification for the convergence of the proposed method.
- The insight of conditional independence between the critic parameters and past actor parameters given the current state trajectory is particularly interesting.

weaknesses:
## Lack of related works and comparison
It is hard to position this paper in the context of prior work as it lacks a discussion of how this work relates to existing RL algorithms that model the uncertainty of Q-values. The paper should discuss how this work is different from distributional RL methods, which also model the uncertainty of Q-values. The paper should also discuss why SGMCMC is the suitable method for the problem at hand, in contrast to prior work. This made the paper particularly hard to understand as a reader.

## The problem of uncertainty quantification is not well-motivated
The paper does not provide a clear motivation for why it is important to model the uncertainty of Q-values. It mentions "accurately quantifying the uncertainty of the value function has been a critical concern for ensuring reliable and robust RL applications", but does not provide any concrete examples of why this is the case or precisely in what scenarios this is important. Ideally, the paper should analyze the limitations of existing RL algorithms that do not model the uncertainty of Q-values and provide examples of scenarios where this leads to suboptimal performance, and how the proposed method addresses these limitations in experiments.

## No results on PPO
The paper mentions that PPO suffers from severe miscalibration issues in both the actor and critic, but does not provide any results on PPO to demonstrate this.

## Connecting uncertainty quantification to performance
While the escape environment discusses the relationship of how LTF leads to better MSE of value functions, how this translates to better performance in the escape environment is not clear. Conversely, on other environments, the paper does not provide a clear analysis of how the uncertainty quantification of Q-values leads to better performance.

It is also not so clearly apparent that LT-A2C has smaller seed variability than A2C, as the paper claims. The performance difference is largely imperceptible in the plots, considering the confidence intervals.

questions:
Can you build a connection to prior works in uncertainty quantification?

How would you make this paper more approachable to a wider audience? Currently, it requires knowing a lot of prior work to make sense of the motivation, algorithm, and convergence proofs. The paper should be largely self-sufficient when reading.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732389355495

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Meta reply for all reviewers

comment:
### **Meta Reply**

We sincerely thank you for your valuable comments and suggestions. We understand that your primary concern lies in how our framework compares to existing Bayesian methods for addressing uncertainties in reinforcement learning. In this paper, we addresses a fundamental limitation of Bayesian RL methods when applied to the dynamic RL processes.

**1. Misalignment Between Bayesian Philosophy and Online RL Training:**
> In online RL, an infinite stream of transition data is available. Consequently, the uncertainties of network parameters arise primarily from the stochastic nature of the environment, not from a fixed dataset. Bayesian methods, however, rely on posterior distributions conditioned on a static dataset, leading to a theoretical mismatch. This misalignment raises questions about the appropriateness of Bayesian RL methods for online RL training.

**2. Dynamic Nature of RL Processes Is Often Overlooked:**
> RL training involves two interdependent steps: policy control (actor network update) and policy evaluation (critic network update). The actor network’s parameters, which generate new data, change dynamically at every time step. Existing Bayesian RL methods often disregard this evolving nature, treating historical transition tuples in the replay memory as identically distributed. This assumption weakens their theoretical rigor and their applicability to dynamically changing systems.

**3. Challenges with Uncertainty Quantification in Existing Bayesian Methods:**
> While Bayesian methods aim to quantify uncertainties in Q-functions, the convergence of the Q-network’s distribution remains unclear. Moreover, their ability to provide reliable uncertainty quantification is often unproven through numerical experiments. For instance, in the indoor escape environment [1], popular methods like Bootstrapped DQN (BootDQN) and Distributional RL (QR-DQN) fail to correctly estimate the Q-network's  distribution, resulting in unfaithful coverage rates. Additionally, random prior networks (RPNs) [2] are unable to construct accurate confidence intervals for true Q-values, for which posterior consistency is not theoretically guaranteed.

> Using publicly available code, we conducted experiments to evaluate coverage rates and mean squared error (MSE). Results in Table 1 (table summarizing 100 experiments with mean (standard deviation)) demonstrate that existing Bayesian methods fail to provide reliable uncertainty quantification in this setting. Furthermore, current Bayesian actor-critic algorithms heavily rely on these Bayesian DQN-style methods, and introducing an additional actor network only makes the interval predictions worse.

**4. Challenges in Extending Bayesian Theory to Deep Neural Networks:**
> While Bayesian methods can theoretically provide accurate uncertainty quantification, their applicability is constrained by restrictive assumptions. For instance, Gaussian Process Temporal Difference (GPTD) and Kalman Temporal Difference (KTD) methods require linearity to uphold Gaussian assumptions, limiting their use in more complex scenarios. In the context of deep neural networks, common approaches like bootstrapping with ensemble networks increase computational complexity and memory usage significantly, making them less practical for real-world applications. Additionally, techniques such as random prior functions (RPFs) are often employed in methods like Bayesian Bellman operators to introduce stochasticity. However, these approaches lack theoretical guarantees, further questioning their reliability in Bayesian reinforcement learning.

Our framework provides the following key contributions, overcoming the limitations of Bayesian RL while maintaining theoretical rigor and practical efficiency:

**1. Unified Actor-Critic Framework with Latent Trajectories:**
> We integrate the two RL steps—policy control and policy evaluation—into a unified dynamic system. By conceptualizing transition trajectories as latent variables, we use a variant of Fisher’s identity to implicitly integrate out these trajectories during policy gradient evaluation (see Eq. (4) in our paper). This approach fully accounts for the dynamic nature of RL and provides theoretical guarantees for the convergence of both actor and critic networks. Additionally, we establish convergence for the entire RL training process, offering a rigorous foundation for dynamic uncertainty estimation.

**2. Efficient and Flexible Framework Without Architectural Changes:** 
>Unlike Bayesian RL methods, which often rely on ensemble networks or Gaussian process assumptions, our latent trajectory framework is lightweight and adaptable. It requires no changes to the network architecture, making it compatible with any critic network structure. Our approach seamlessly integrates with scalable SGMCMC sampling, providing both computational efficiency and ease of implementation.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732927419467

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Reminder for rebuttal review

comment:
Dear reviewers,

We hope this message finds you well. This is a gentle reminder to review our rebuttal and revised manuscript, which we have submitted in response to your valuable comments. In the revised version, we have included a new “Related Work” section to better address the limitations of existing Bayesian methods and to clearly highlight the contributions of our work.

If you have any further questions or require additional clarifications, please feel free to reach out. We greatly appreciate your feedback and look forward to your response.

Best Regards
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper reframes actor critic methods as finding the zero of an integral equation by regarding transition tuple and value function parameters as latent variables. Some theoretical and empirical results are provided. First, I am concerned about the motivation of this work. Actor critic is not simply to find the stationary point of the actor networks. The dominating actor critic methods like TRPO, PPO are designed based on a monotonic performance improvement theorem. Even for the classical policy gradient method, the goal is still to maximize the expected total reward. Ending up with a stationary point is merely a consequence (due the complexity of nonlinear optimization), not the goal. Second, it is not clear what kind of uncertainty the authors refer to -- the authors never formally define uncertainty. In RL, we can have two entirely different uncertainties, the epistemic uncertainty and the aleatoric uncertainty. More discussion about this is needed to better place the paper in the literature. Third, the empirical improvements of LTA2C is really marginal, even just compared with a pretty weak baseline. Moreover, the selection of tasks is not well justified. It is claimed that the performance improvement is from better exploration but there is no evidence supporting this, and the used tasks are not exploration desperate.