Paper: Fast Value Tracking for Deep Reinforcement Learning (ICLR-2024)

Reviewer comment (Theory):
In equation 2, the introduction of $\pi$ seems to lack motivation. It is unclear why it can tackle the issues suffered by KTD.

Author response:
KTD suffers from inaccurate linearization and resulted operations on large-scale matrices. The proposed algorithm employs a state augmentation approach, see equation (5) of the paper, to avoid linearization. Furthermore, it accelerates computation  using the forecast-analysis procedure developed in ensemble Kalman filter, which avoids direct decomposition of the preconditioned matrix given in equation (11) of the paper. Refer to Zhang et al. [1] for more discussions on this issue.

[1] Peiyi Zhang, Qifan Song, and Faming Liang. A langevinized ensemble kalman filter for large-
scale static and dynamic learning. Statistica Sinica, 2023. doi: 10.5705/ss.202022.0172.

Reviewer comment (Writing):
The meaning of the term "stage" needs to be clarified.

Author response:
The term "stage" refers to  a time step of the state space model, where a batch of training samples are collected based on the current value of $\theta$. In the revision, we will change it to ``step'' as used in some other RL papers.

Reviewer comment (Writing):
It is problematic to introduce a notation without explanation the first time it appears, such as $h(\cdot)$.

Author response:
In the revision, we will add explanations for all newly introduced notations. Specifically, $h(x_t,\theta_t)$ can be interpreted as a stochastic estimator of $E_{\rho} [r(s,a)|s,a]$ based on the new sample $s_{t+1}$ or  $(s_{t+1},a_{t+1})$, depending on our interest in $V$- or $Q$-functions.

Reviewer comment (Writing):
The term "Optimal policy exploration" should be further explained.

Author response:
We use term to express the property of the proposed algorithm: It can potentially sample all optimal policies with correct proportions. We illustrates this property using Figure 4, which indicates the proposed algorithm can sample optimal policies in approximately correct proportions, while Adam fails to explore different optimal policies.

Reviewer comment (Writing):
It is hard to get the main contribution of the paper from the introduction.

Author response:
In the revision, we will particularly emphasize that (i) we provide a new formulation for RL, and (ii) provide an associated sampling algorithm and study its theoretical guarantee, enabling the uncertainty associated with the RL system to be properly quantified. However, the existing methods fail to do so.
We note that our use of SGLD in this paper is new and should be classified as one of our proposed algorithms for RL. In the rebuttal, we have implemented an existing version of SGLD (Vanilla SGLD) for RL. on the mean square Bellman error, which cannot provide faithful confidence intervals as ours. 
The following figure shows the mean coverage rate of 95% prediction interval over 100 experiments.
https://pasteboard.co/xHNnZ0kGEyoX.png

Reviewer comment (Writing):
The phrase "such as e.g.?" is awkward.

Author response:
The text will be changed accordingly.

Reviewer comment (Experiments):
The uncertainty quantification for Deep RL is not restricted to Kalman TD frameworks. In the literature of Bayesian RL, we do have works that adopt other approximate sampling techniques, e.g. ensembling in Bootstrapped DQN (Osband et.al. 2016), "Deep Exploration via Randomized Value Functions" (Osband et.al. 2019), and RLSVI with general function approximation (Ishfaq et.al., 2021). Although the above methods do not adopt MCMC methods as approximate sampling scheme, could authors compare the proposed method with the existing Bayesian RL methods, not just restricted to the Kalman TD context?

Author response:
In our most recent experiment, we compared our LKTD and the reformulated SGLD algorithms with DQN, BootDQN, and Vanilla SGMCMC algorithms, which do not incorporate our state space reformulation. The box plot of the coverage rate in Figure 1 (https://pasteboard.co/xHNnZ0kGEyoX.png) clearly shows that only the LKTD and reformulated SGLD algorithms achieve a reliable 95% prediction interval. In contrast, the other algorithms failed to converge to the true Q-value function. This failure is attributed to the semi-gradient derived from the DQN structure, which represents a biased gradient for the mean square Bellman equation.

Reviewer comment (Novelty):
Though the authors provide a clear comparison to the existing Kalman TD frameworks, it is desirable to see a more thorough related work study in Bayesian RL / uncertainty quantification in deep RL, which would be beneficial.

Author response:
Bayesian DQN deploys Thompson sampling on the approximated posterior of the Q-function. However, the posterior distribution is conditional on the last hidden layer of the deep Q network. Since the network parameters are constantly changing throughout the training phase, it is unclear where the Q-function converges. On the other hand, LKTD guarantees the Q-network converges to a distribution around the optimal Q-function. Therefore, we can draw samples during the training phase and perform uncertainty quantification in a single experiment. Thanks for your advice, in the revision, we will include a more detailed explanation of the difference between our algorithm and related works.

Reviewer comment (Theory):
The proposed method requires a proper prior distribution to achieve the performance guarantee in LEnKF. It is unclear how the quality of the prior affects the performance through the current theorems. Is it possible to capture such quantity in the upper bounds provided in Theorem 1 or Theorem 2?

Author response:
Thanks for your thoughtful question. The prior can generally affect   
    the constants $A$, $L_U$, $m_U$, $M_U$, and $B$, which are involved in the upper bound provided 
    in Theorem 1. However, as implied by equation (9) in Algorithm 1, the prior effect  on these constants is of the order $O(\frac{1}{\mathcal{N}})$, which can be diminished by increasing the pseudo population size $\mathcal{N}$.

Reviewer comment (Experiments):
The experiment environments are relatively simple, especially for the control tasks. It is suggested authors consider more complicated locomotion tasks to demonstrate the effectiveness of the proposed method.

Author response:
As suggested by another reviewer, we will include behavior suite environment in our revision.

Reviewer comment (Theory):
The proofs of Theorem 1 and 2 rely on a significant number of assumptions. It is suggested authors comment on whether those assumptions are imposed in order to achieve the convergence guarantee in sampling schemes, and how such assumptions can be satisfied in practice.

Author response:
Thank you for your thoughtful question. In Lemma 1, we show that the Kalman temporal difference (KTD) framework used to solve the RL problem can be formulated as a preconditioned SGLD algorithm. This formulation is general without specific assumptions involved. 
    Then we establish the convergence of the preconditioned SGLD algorithm under the standard assumptions for stochastic gradient MCMC algorithms. We refer these assumptions to Raginsky et al. (2017) [1] (cited in the paper).  As explained below, these assumptions can be easily satisfied in practice.
Assumption A1-(C1) concerns the stationarity of the Markov decision process (MDP), a regular condition for almost all RL studies. 
     
Assumption A1-(C3) basically requires the log-prior density function $\log \pi(\theta)$ 
    and the neural network function $h(x,\theta)$ (with $\theta$ representing the weights 
    of the neural network) to be twice continuously differentiable. In this paper, we set $\pi(\theta)$ 
    as  a mixture Gaussian distribution (see equation (A34)). Therefore, A1-(C3) is satisfied.

Intuitively, the assumption A1-(C4) is to bound the dynamics inside a bounded domain; if $\theta_k$ is far away from the origin, then the dynamics are forced to get back around the origin. Although the neural network likelihood function is multimodal, the mixture Gaussian prior used in the paper forces the resulting posterior to be well behaved around the origin and thus to satisfy 
    A1-(C4) and A1-(C2). 

Assumption A1-(C5) regulates the variation of the stochastic gradient involved 
     at each iteration of the algorithm. This assumption can be satisfied by a pre-fixed value 
     of $\mathcal{N}$.

Assumption A1-(C6) concerns the distribution of the initial weights of the neural network. It can be satisfied by initializing weights in a bounded space.

Assumption A2 concerns the setting of the learning rate, which can be satisfied by setting $\epsilon_k=O(1/k^{\alpha})$ for some $1/2 < \alpha<1$. Assumption A3 is taken from Chen et al. (2015) [2] (cited in the paper), a regular assumption for test functions. Assumption A4 concerns the eigenvalues of the preconditioning matrix. It is automatically satisfied by the special structure 
     of $\Sigma_t= \frac{n}{\mathcal{N}}[I-\epsilon_t H_t^T(\epsilon_t H_t H_t^T+R_t)^{-1} H_t]$, following from an extension of Ostrowski's theorem [3].

The above discussions will be incorporated into the revised manuscript.

[1] Maxim Raginsky, Alexander Rakhlin, and Matus Telgarsky. Non-convex learning via
stochastic gradient Langevin dynamics: a nonasymptotic analysis. In Proceedings of the
2017 Conference on Learning Theory, pp. 1674–1703, 2017.}

[2] Changyou Chen, Nan Ding, and Lawrence Carin. On the convergence of stochastic gradient
mcmc algorithms with high-order integrators. In Advances in Neural Information Processing
Systems, pp. 2278–2286, 2015.

[3] Nicholas J. Higham and Sheung Hun Cheng, Modifying the Inertia of Matrices Arising in Optimization, Linear Algebra Appl. 275–276, 261-279, 1998.

Reviewer comment (Experiments):
Comparison with existing posterior sampling value-based algorithms for exploration is missing. Say ensemble sampling, Bootstrapped DQN or HyperDQN.

Author response:
We have implemented Bootstrapped DQN to construct 95\% prediction interval for the escape environment.  Due to the use of semi-gradient in DQN structure, the Q-functions construct by BootDQN is biased and thus the resulting confidence interval is less faithful. 
In figure 1 (https://pasteboard.co/xHNnZ0kGEyoX.png), the box plot of coverage rate shows that BootDQN failed to cover the true Q-value at the nominal 95\% level. In figure 2 (https://pasteboard.co/kBcqXjvKh6pP.png) , the MSE of Q-value estimates indicates that BootDQN exhibits a much larger bias compare to LKTD and SGLD (with pseudo population). In figure 3 (https://pasteboard.co/1mRkGnk2YJ0N.png), there is a clear decrease in interval length as the pseudo population size increases, which agrees with the statement in Remark 2. In our implementation for Bootstrap DQN, we set the learning rate to $10^{-4}$, and Bernoulli masking probability to 0.5. We demonstrated the results for two settings: head=5 and head=20.

Reviewer comment (Experiments):
Comparison for uncertainty quantification in deep neural network is missing.

Author response:
In this rebuttal, we have included the results of SGLD and SGHMC, which were used to train deep neural networks but in the conventional  formulation. That is, we apply both algorithms directly to the mean square temporal difference loss function without state space reformulation. 
    \begin{equation}
        \sum_{i=1}^n (Q(s,a|\theta) - r -\gamma \max_{a'} Q(s',a'|\theta))^2 + ||\theta||^2
    \end{equation}
    As shown by the plots given in the common reply, they could not produce correct coverage for Q-values. This experiment indicates the importance of our new formulation. For Vanilla SGLD and Vanilla SGHMC, both learning rate and temperature are set to $10^{-4}$ and 0.1 respectively. And the momentum coefficient for SGHMC is set to be 0.9.