Value-Biased Maximum Likelihood Estimation for Model-based Reinforcement Learning in Discounted Linear MDPs

Q4: Empirical comparison of VBMLE and tabular RL methods as well as UCLK+, and test VBMLE on larger MDPs

• For a more extensive comparison, we further compare VBMLE (augmented with BO) with Posterior Sampling for Reinforcement Learning (PSRL) [5], a popular benchmark method for tabular RL, on an MDP with |S|=100 and |A|=4. Note that this size is already much larger than those used in the existing linear mixture MDP literature, e.g. Riverswim with |S|=6 in [6] . The results for larger MDPs are provided here: https://imgur.com/a/8XXHpYJ. The results demonstrate that VBMLE converges much faster than PSRL and thereby significantly outperforms PSRL [3] in regret performance, across both small and large MDPs.

• Additionally, we have also incorporated the regret performance of UCLK+ [4] into our experimental results. Note that UCLK+ still suffers from the computational bottleneck of Extended Value Iteration, and hence we could only test UCLK+ on the smaller MDP.

• We also thank the reviewer for providing the helpful reference [7] on the experiments of linear MDPs. We took a careful look at [7] and realized that [7] focuses mainly on linear MDPs and does not handle “linear mixture MDPs” (the difference between two settings are also mentioned in Q3). As “linear mixture MDPs” and linear MDPs are two different settings (and one could not recover the other), it would require a substantially different approach to adapt the technique in [7] to the linear mixture MDP setting. Accordingly, as an alternative, we test VBMLE on the larger MDP with |S|=100 as mentioned above.

[5] Ian Osband, Daniel Russo, and Benjamin Van Roy, “(More) Efficient Reinforcement Learning via Posterior Sampling,” NeurIPS 2013.

[6] Alex Ayoub, Zeyu Jia, Csaba Szepesvari, Mengdi Wang, and Lin F. Yang, “Model-Based Reinforcement Learning with Value-Targeted Regression,” ICML 2020.

[7] Tianjun Zhang et al., "Making linear MDPs practical via contrastive representation learning," ICML 2022.

Q5: Discuss any potential connections between VBMLE and the model-based RL via the optimistic approach (e.g., UCRL-based approaches)

Indeed, the VBMLE and UCRL approaches are related. UCRL chooses the most optimistic model within a confidence interval, while VBMLE maximizes an objective function which is a linear combination of the value function and the KL-penalty. One can observe that VBMLE optimization problem is a dual formulation of the UCRL optimization (with lagrangian $\frac{1}{\alpha(t)}$). UCRL is a special case of VBMLE for a specific choice of $\alpha(t)$. Interestingly, the superior empirical performance of RBMLE suggests that the confidence interval is too big. More detailed discussion on VBMLE vs UCRL can be found in [8] and [9].

[8] Akshay Mete, Rahul Singh, and P. R. Kumar, “Augmented RBMLE-UCB Approach for Adaptive Control of Linear Quadratic Systems,” NeurIPS 2022.

[9] Akshay Mete, Rahul Singh, and P. R. Kumar, "The RBMLE method for Reinforcement Learning," IEEE Annual Conference on Information Sciences and Systems (CISS), 2022.

Q6: Check the typo and notations without definitions (e.g., $\pi_\infty^{\text{MLE}}(s)$ in Section 4)

Thanks for catching this! We have added the definitions of $\pi_\infty^{\text{MLE}}(s)$ below (7) in our updated manuscript.

Q2: Compare the computational efficiency of VBMLE and UCLK

Thank the reviewer for the helpful suggestion. We provide a detailed comparison as follows:

1. UCLK suffers from expensive Extended Value Iteration: In UCLK, the main computational bottleneck lies in the large number of optimization runs needed in Extended Value Iteration (EVI), as shown in the pseudo code of Algorithm 2 in [3] (https://imgur.com/a/J00LX3L). Specifically, there is a quadratic constrained optimization in each Bellman update; Let the number of iterations is denoted by $U$, then each call of EVI will result in solving the quadratic constrained optimization problem for $|S| |A | U$ times. This is rather intractable in practice (e.g., in the experiments, we find that it takes more than 35 hours of wall clock time to finish one single step of UCLK). Similar issues also remain in the improved UCLK+ [4].

2. VBMLE can efficiently handle the non-concave objective by Bayesian optimization: Regarding VBMLE, the primary complexity of VBMLE arises from the non-concave constrained optimization in (10), which could be efficiently addressed by using Bayesian optimization (BO) techniques (e.g., GP-UCB), as also described in Q1 and Q2 above. Let $K$ denote the number of samples taken by GPUCB in each maximization run of finding VBMLE (we choose $H=25$ in the experiments). Then, the complexity of VBMLE with GP-UCB for finding $\theta^R_t$ is to solve the standard Value Iteration for only $H+1$ times (H is for BO, each sample requires one value iteration in our objective function, and another 1 for value iteration for $\theta^R_t$). This is a clear computational advantage over the EVI in UCLK.

Moreover, we also conduct experiments to compare the computation times of VBMLE and UCLK, as shown below: https://imgur.com/a/5ksC7Xd The detailed discussion is also provided in Appendix D in the updated manuscript.

[3] Dongruo Zhou, Jiafan He, and Quanquan Gu, “Provably Efficient Reinforcement Learning for Discounted MDPs with Feature Mapping,” ICML 2021.

[4] Dongruo Zhou, Quanquan Gu, and Csaba Szepesvari, “Nearly Minimax Optimal Reinforcement Learning for Linear Mixture Markov Decision Processes,” COLT 2021.

Q3: Clarify the related work on linear MDPs and linear kernel MDPs

In the original manuscript, we follow the terminology used in (Cai et al., 2020), one of the earliest works on this setting, and use the term “linear MDP” for the transition model parameterized as $P(s’|s,a)=<\phi(s’|s,a), \theta^*>$. However, we do agree with the reviewer that the term “linear MDP” could be overloaded as “linear MDP” also widely refers to the low-rank factorization where the transition model takes the form of the inner product between a state-action feature vector and the vector of unknown measures over states. Therefore, we agree that using the term “linear mixture MDP” could avoid the possible ambiguity and would be more consistent with other existing literature.

We have made the changes accordingly in the updated manuscript. Moreover, for clarity, we also reorganize the Related Work to provide a more thorough discussion on these lines of research works.

Q1: Explain how VBMLE handles the optimization problem in Equation (10)

The VBMLE maximization problem in (10) can be approximately solved by multiple approaches in practice:

(i) Solve VBMLE by Bayesian optimization: Due to the non-concavity of VBMLE, we propose to solve VBMLE by Bayesian optimization (BO), which is a powerful and generic method for provably maximizing (possibly non-concave) black-box objective functions. As a result, BO can provably find an $\epsilon$-optimal solution to VBMLE within finite iterations. Specifically:

• We have applied the GP-UCB algorithm, which is one classic BO algorithm and has been shown to provably find an $\epsilon$-optimal solution within $\tilde{\mathcal{O}}(1/\epsilon^2)$ iterations under smooth (possibly non-concave) objective functions [2]. Each sample taken by GP-UCB requires only one run of standard Value Iteration.
• To further demonstrate the compatibility of VBMLE and BO, we have extended the regret analysis of VBMLE to the case where only an $\epsilon$-optimal VBMLE solution is obtained (as also suggested by Reviewer p6f5). Specifically, let $K$ denote the number of samples taken by GPUCB in each maximization run of finding VBMLE. We show that VBMLE augmented with BO can achieve a regret bound of $\mathcal{O}(\max(\frac{d\sqrt{T}\log{T}}{p_{\text{min}^2}(1-\gamma)^2},\sqrt{T}\frac{(\log{K})^{d+1}}{\sqrt{K}} ))$. By using a moderate $K$, one could easily recover the same regret bound as that of VBMLE with an exact maximizer. In our experiments, we find that choosing $H=25$ is sufficient and also computationally efficient. The empirical regret of VBMLE augmented with BO and other benchmark methods can be found at https://imgur.com/a/8XXHpYJ and in Figure 1 of the updated manuscript. The computation times of VBMLE and other methods are provided below https://imgur.com/a/5ksC7Xd and also shown in Table 1 of the updated manuscript. We can see that VBMLE with BO enjoys both low empirical regret as well as low computation time.

We have also added the discussion on VBMLE with BO in Appendix C in the updated manuscript.

(ii) Approximately solve VBMLE by off-the-shelf constrained optimization solvers: In practice, another possibility is to use an off-the-shelf optimization solver. For example, we have also applied the trust region method (implementation available in SciPy), which is compatible with the VBMLE problem and only requires taking the VBMLE objective function as an input argument (and this is used for constructing the trust-region subproblem, e.g., please refer to [1]). However, as mentioned by the reviewer, one major issue is that due to the non-concavity of VBMLE objective, trust-region methods only ensure convergence to local optima. Despite this, from the experiments, we find that VBMLE augmented with a trust-region method achieves regret comparable to VBMLE with BO.

[1] Andrew R. Conn, Nicholas IM Gould, and Philippe L. Toint, “Trust region methods,” Society for Industrial and Applied Mathematics, 2000.

[2] Niranjan Srinivas, Andreas Krause, Sham M Kakade, and Matthias W Seeger, “Information-theoretic regret bounds for Gaussian process optimization in the bandit setting,” IEEE Transactions on Information Theory, 2012.

Q7: Explain why vector $\theta$ in UCLK is dependent on (s,a)?

As also described in Q3 above, UCLK involves Extended Value Iteration (EVI), and there is a quadratic constrained optimization in each Bellman update (please also refer to the pseudo code of EVI in UCLK: https://imgur.com/a/J00LX3L). As a result, if the number of iterations needed in each call of EVI is denoted by $U$, then each EVI shall result in solving the quadratic optimization problem for $|S| |A| U$ times.

Q8: Explain how VBMLE handles exploration-exploitation trade-off

The bias term in VBMLE controls the exploration and the part of MLE controls the exploitation. If we have $\alpha(t) = 0$, the VBMLE serves as MLE to be a greedy policy based on current observation. In VBMLE, we show that the choice $\alpha(t)=\sqrt{t}$ leads to an optimal trade-off between exploration and exploitation (cf. the statement of Theorem 2 in Appendix B).

Q9: Check the typo and notations without definitions

Thanks for catching this. We have fixed the typos and added the definitions of $\pi_\infty^{\text{MLE}}(s)$ below (7) in our updated manuscript.

Q4: Issues with $p_\text{min}$ in VBMLE

We would like to clarify the assumption about $p_\text{min}$ and the dependency of the regret bound on $p_{min}$ as follows:

1. The dependency of the regret bound on the true $p_{min}$: We agree with the reviewer that the regret bound is $p_\text{min}$-dependent and could be shown in the Big-O notation. As shown in the updated Theorem 2, we clearly point out that the regret bound is $\tilde{\mathcal{O}}(\frac{d\sqrt{T}}{p_{\text{min}}^4(1-\gamma)^2})$. Moreover, we would like to highlight that the above dependency on $1/p_{min}$ could be quite conservative in practice. Specifically, we evaluate VBMLE on two MDPs, each with a moderate $p_\text{min}$ or a very small $p_{min}$. The empirical regrets are shown in the figures below. https://imgur.com/a/Nw4Wp4I We could see that VBMLE could still achieve low empirical regret under a $p_\text{min}$ less than 0.001.

We have also added the discussion and the additional experimental results in Appendix F.

2. The assumption about the knowledge of the true $p_\text{min}$ could be lifted by using a strictly decreasing positive function $p(t)$ with $lim_{t \rightarrow \infty} p(t)=0$ as a surrogate for true $p_\text{min}$:
In the original manuscript, we use the knowledge about the true $p_\text{min}$ for two purposes: (i) To ensure that the parameter set $\mathbb{P}$ is not empty; (ii) To provide a lower bound on the term $\phi_t(s_{k+1})^\top \theta_t^{MLE}, \forall k \leq t$. Nevertheless, even without the knowledge of the true $p_\text{min}$, the above two requirements could still be met by using a strictly decreasing function $p(t)$ in the parameter set $\mathbb{P}$. Specifically, as $p(t)$ is strictly decreasing and goes to zero in the limit, there must exist some constant $T_0>0$ such that $p(t)< p_{min}$ for all $t>T_0$. As a result, we could remove the assumption about $p_\text{min}$ but with a tradeoff on an additional regret term $T_0$, which is independent of $T$, $d$, and $\gamma$. Under the choice of $p(t)=1/\log(t)$, VBMLE can be shown to achieve $\mathcal{O}\left(\frac{d\sqrt{T}(\log{T})^5}{(1-\gamma)^2} + \frac{T_0}{1-\gamma}\right)$ regret, without the need for the knowledge of the true $p_\text{min}$. For the detailed proof and description, please refer to Theorem 3, Remark 1, and Appendix C.

Q5: Clarify the related work on linear MDPs and linear kernel MDPs

We thank the reviewer for the suggestion on clarifying the terminology and the related works on these two lines of research. In the original manuscript, we follow the terminology used in (Cai et al., 2020), one of the earliest works on this setting, and use the term “linear MDP” for the transition model parameterized as $P(s’|s,a)=<\phi(s’|s,a), \theta^*>$. However, we do agree with the reviewer that the term “linear MDP” could be overloaded as “linear MDP” also widely refers to the low-rank factorization where the transition model takes the form of the inner product between a state-action feature vector and the vector of unknown measures over states. Therefore, we agree that using the term “linear mixture MDP” could avoid possible ambiguity and would be more consistent with other existing literature.

We have made the changes accordingly in the updated manuscript. Moreover, we also reorganize the Related Work to better highlight the difference between the two types of linear models.

Q6: Describe the difference between VBMLE and other RL works on using MLE to estimate transition kernels (e.g., FLAMBE [5])

We thank the reviewer for the suggestion and for the reference. The differences between VBMLE and FLAMBE mainly lie in the goal and the problem setting. Specifically,

• FLAMBE [5] and REP-UCB [6] focus on using MLE for the representation learning aspect of low-rank MDPs (i.e., learning the unknown features). As a result, these algorithms and their theoretical guarantees could help enable the downstream RL tasks (e.g., reward maximization or regret minimization), for any given reward function.

• By contrast, VBMLE, as a model-based RL algorithm, uses biased MLE as a subroutine of sequential decision making and focuses on regret minimization under linear mixture MDPs, with feature vectors already given. Therefore, VBMLE needs to address the exploration-exploitation issue in solving the linear mixture MDPs (while FLAMBE does not).

Based on the above, we can see that VBMLE and FLAMBE / REP-UCB focus on different aspects of RL under function approximation and therefore could potentially complement each other.

[5] A. Agarwal, S. Kakade, A. Krishnamurthy, and W. Sun, “FLAMBE: Structural complexity and representation learning of low rank MDPs,” NeurIPS 2020.

[6] Masatoshi Uehara, Xuezhou Zhang, and Wen Sun, “Representation learning for online and offline RL in low-rank MDPs,” ICLR 2021.

Q2: What if an empirical optimization algorithm is adopted to estimate the parameter? For example, the estimator is no longer the optimal one but may have some errors, say $|\hat{\theta}-\theta^{MLE}| \leq \epsilon$ for some $\epsilon$.

We thank the reviewer for the insightful suggestion. As mentioned in Q1 above, regarding the empirical optimization algorithm, we propose to apply Bayesian optimization (BO), which can provably find an $\epsilon$-optimal solution of a smooth (possibly non-concave) function within $\tilde{\mathcal{O}}(1/\epsilon^2)$ iterations, to find an approximate solution of VBMLE. To further demonstrate the compatibility of VBMLE and BO, we have extended the regret analysis of exact VBMLE to the case where only an $\epsilon$-optimal VBMLE solution is obtained, as suggested by the reviewer. Specifically, let $K$ denote the number of samples taken by GPUCB in each maximization run of finding VBMLE. As described in Theorem 4, we show that VBMLE augmented with BO can achieve a regret bound of $\mathcal{O}(\max(\frac{d\sqrt{T}\log{T}}{p_{\text{min}^2}(1-\gamma)^2},\sqrt{T}\frac{(\log{K})^{d+1}}{\sqrt{K}} ))$. The above regret bound holds for the following reasons: (i) The optimization error of $\theta^{R}$ can be addressed by completing the square in (103), ensuring that even if $\theta^{R}$ is not the exact maximizer, the RHS remains unchanged. (ii) VBMLE does not require the computation of $\theta^{MLE}$ (which only serves as an intermediate machinery for analysis), thus we do not need to take the optimization error for $\theta^{MLE}$ into account. (iii) The optimization error for the "value of the objective function" can be considered and nicely handled by Lemma 7. Moreover, this practical variant of VBMLE preserves high computational efficiency and also enjoys low empirical regret, as indicated in the figures below. https://imgur.com/a/8XXHpYJ https://imgur.com/a/5ksC7Xd

Q3: Compare the computational efficiency of VBMLE versus UCLK

Thank the reviewer for the helpful suggestion. We provide a detailed comparison as follows:

1. UCLK suffers from expensive Extended Value Iteration: In UCLK, the main computational bottleneck lies in the large number of optimization runs needed in Extended Value Iteration (EVI), as shown in the pseudo code of Algorithm 2 in [3] (https://imgur.com/a/J00LX3L). Specifically, there is a quadratic constrained optimization in each Bellman update; Let the number of iterations is denoted by $U$, then each call of EVI will result in solving the quadratic constrained optimization problem for $S \times A \times U$ times. This is rather intractable in practice (e.g., in the experiments, we find that it takes more than 35 hours of wall clock time to finish one single step of UCLK). Similar issues also remain in the improved UCLK+ [4].

2. VBMLE can efficiently handle the non-concave objective by Bayesian optimization: Regarding VBMLE, the primary complexity of VBMLE arises from the non-concave constrained optimization in (10), which could be efficiently addressed by using Bayesian optimization (BO) techniques (e.g., GP-UCB), as also described in Q1 and Q2 above. Let $K$ denote the number of samples taken by GPUCB in each maximization run of finding VBMLE (we choose $H=25$ in the experiments). Then, the complexity of VBMLE with GP-UCB for finding $\theta^R_t$ is to solve the standard Value Iteration for only $H+1$ times (H is for BO, each sample requires one value iteration in our objective function, and another 1 for value iteration for $\theta^R_t$). This is a clear computational advantage over the EVI in UCLK.

Moreover, we also conduct experiments to compare the computation times of VBMLE and UCLK, as shown below: https://imgur.com/a/5ksC7Xd The detailed discussion is also provided in Appendix D in the updated manuscript.

[3] Dongruo Zhou, Jiafan He, and Quanquan Gu, “Provably Efficient Reinforcement Learning for Discounted MDPs with Feature Mapping,” ICML 2021.

[4] Dongruo Zhou, Quanquan Gu, and Csaba Szepesvari, “Nearly Minimax Optimal Reinforcement Learning for Linear Mixture Markov Decision Processes,” COLT 2021.

Q1: Explain how to solve the optimization problem of VBMLE in Equation (10)

The VBMLE maximization problem in (10) can be approximately solved by multiple approaches in practice:

(i) Solve VBMLE by Bayesian optimization: Due to the non-concavity of VBMLE, we propose to solve VBMLE by Bayesian optimization (BO), which is a powerful and generic method for provably maximizing (possibly non-concave) black-box objective functions. As a result, BO can provably find an $\epsilon$-optimal solution to VBMLE within finite iterations. Specifically:

• We have applied the GP-UCB algorithm, which is one classic BO algorithm and has been shown to provably find an $\epsilon$-optimal solution within $\tilde{\mathcal{O}}(1/\epsilon^2)$ iterations under smooth (possibly non-concave) objective functions [2]. Each sample taken by GP-UCB requires only one run of standard Value Iteration.
• To further demonstrate the compatibility of VBMLE and BO, we have extended the regret analysis of VBMLE to the case where only an $\epsilon$-optimal VBMLE solution is obtained (as also suggested by Reviewer p6f5). Specifically, let $K$ denote the number of samples taken by GPUCB in each maximization run of finding VBMLE. We show that VBMLE augmented with BO can achieve a regret bound of $\mathcal{O}(\max(\frac{d\sqrt{T}\log{T}}{p_{\text{min}^2}(1-\gamma)^2},\sqrt{T}\frac{(\log{K})^{d+1}}{\sqrt{K}} ))$. By using a moderate $K$, one could easily recover the same regret bound as that of VBMLE with an exact maximizer. In our experiments, we find that choosing $H=25$ is sufficient and also computationally efficient. The empirical regret of VBMLE augmented with BO and other benchmark methods can be found at https://imgur.com/a/8XXHpYJ and in Figure 1 of the updated manuscript. The computation times of VBMLE and other methods are provided below https://imgur.com/a/5ksC7Xd and also shown in Table 1 of the updated manuscript. We can see that VBMLE with BO enjoys both low empirical regret as well as low computation time.

We have also added the discussion on VBMLE with BO in Appendix C in the updated manuscript.

(ii) Approximately solve VBMLE by off-the-shelf constrained optimization solvers: In practice, another possibility is to use an off-the-shelf optimization solver. For example, we have also applied the trust region method (implementation available in SciPy), which is compatible with the VBMLE problem and only requires taking the VBMLE objective function as an input argument (and this is used for constructing the trust-region subproblem, e.g., please refer to [1]). However, as mentioned by the reviewer, one major issue is that due to the non-concavity of VBMLE objective, trust-region methods only ensure convergence to local optima. Despite this, from the experiments, we find that VBMLE augmented with a trust-region method achieves regret comparable to VBMLE with BO.

[1] Andrew R. Conn, Nicholas IM Gould, and Philippe L. Toint, “Trust region methods,” Society for Industrial and Applied Mathematics, 2000.

[2] Niranjan Srinivas, Andreas Krause, Sham M Kakade, and Matthias W Seeger, “Information-theoretic regret bounds for Gaussian process optimization in the bandit setting,” IEEE Transactions on Information Theory, 2012.

Q1: Solving the optimization problem of VBMLE in Equation (10)

The VBMLE maximization problem in (10) can be approximately solved by multiple approaches in practice:

(i) Solve VBMLE by Bayesian optimization: Due to the non-concavity of VBMLE, we propose to solve VBMLE by Bayesian optimization (BO), which is a powerful and generic method for provably maximizing (possibly non-concave) black-box objective functions. As a result, BO can provably find an $\epsilon$-optimal solution to VBMLE within finite iterations. Specifically:

• We have applied the GP-UCB algorithm, which is one classic BO algorithm and has been shown to provably find an $\epsilon$-optimal solution within $\tilde{\mathcal{O}}(1/\epsilon^2)$ iterations under smooth (possibly non-concave) objective functions [2]. Each sample taken by GP-UCB requires only one run of standard Value Iteration.
• To further demonstrate the compatibility of VBMLE and BO, we have extended the regret analysis of VBMLE to the case where only an $\epsilon$-optimal VBMLE solution is obtained (as also suggested by Reviewer p6f5). Specifically, let $K$ denote the number of samples taken by GPUCB in each maximization run of finding VBMLE. We show that VBMLE augmented with BO can achieve a regret bound of $\mathcal{O}(\max(\frac{d\sqrt{T}\log{T}}{p_{\text{min}^2}(1-\gamma)^2},\sqrt{T}\frac{(\log{K})^{d+1}}{\sqrt{K}} ))$. By using a moderate $K$, one could easily recover the same regret bound as that of VBMLE with an exact maximizer. In our experiments, we find that choosing $H=25$ is sufficient and also computationally efficient. The empirical regret of VBMLE augmented with BO and other benchmark methods can be found at https://imgur.com/a/8XXHpYJ and in Figure 1 of the updated manuscript. The computation times of VBMLE and other methods are provided below https://imgur.com/a/5ksC7Xd and also shown in Table 1 of the updated manuscript. We can see that VBMLE with BO enjoys both low empirical regret as well as low computation time.

We have also added the discussion on VBMLE with BO in Appendix C in the updated manuscript.

(ii) Approximately solve VBMLE by off-the-shelf constrained optimization solvers: In practice, another possibility is to use an off-the-shelf optimization solver. For example, we have also applied the trust region method (implementation available in SciPy), which is compatible with the VBMLE problem and only requires taking the VBMLE objective function as an input argument (and this is used for constructing the trust-region subproblem, e.g., please refer to [1]). However, as mentioned by the reviewer, one major issue is that due to the non-concavity of VBMLE objective, trust-region methods only ensure convergence to local optima. Despite this, from the experiments, we find that VBMLE augmented with a trust-region method achieves regret comparable to VBMLE with BO.

[1] Andrew R. Conn, Nicholas IM Gould, and Philippe L. Toint, “Trust region methods,” Society for Industrial and Applied Mathematics, 2000.

[2] Niranjan Srinivas, Andreas Krause, Sham M Kakade, and Matthias W Seeger, “Information-theoretic regret bounds for Gaussian process optimization in the bandit setting,” IEEE Transactions on Information Theory, 2012.

Q2: Issues with $p_{min}$ in VBMLE

We would like to clarify the assumption about $p_{min}$ and the dependency of the regret bound on $p_{min}$ as follows:

1. The dependency of the regret bound on the true $p_{min}$: We agree with the reviewer that the regret bound is $p_\text{min}$-dependent and could be shown in the Big-O notation. As shown in the updated Theorem 2, we clearly point out that the regret bound is $\tilde{\mathcal{O}}(\frac{d\sqrt{T}}{p_{\text{min}}^4(1-\gamma)^2})$. Moreover, we would like to highlight that the above dependency on $1/p_{min}$ could be quite conservative in practice. Specifically, we evaluate VBMLE on two MDPs, each with a moderate $p_\text{min}$ or a very small $p_{min}$. The empirical regrets are shown in the figures below. https://imgur.com/a/Nw4Wp4I We could see that VBMLE could still achieve low empirical regret under a $p_\text{min}$ less than 0.001.

We have also added the above discussion and the additional experimental results in Appendix F.

2. The assumption about the knowledge of the true $p_\text{min}$ could be lifted by using a strictly decreasing positive function $p(t)$ with $lim_{t \rightarrow \infty} p(t)=0$ as a surrogate for true $p_\text{min}$:
In the original manuscript, we use the knowledge about the true $p_\text{min}$ for two purposes: (i) To ensure that the parameter set $\mathbb{P}$ is not empty; (ii) To provide a lower bound on the term $\phi_t(s_{k+1})^\top \theta_t^{MLE}, \forall k \leq t$. Nevertheless, even without the knowledge of the true $p_\text{min}$, the above two requirements could still be met by using a strictly decreasing function $p(t)$ in the parameter set $\mathbb{P}$. Specifically, as $p(t)$ is strictly decreasing and goes to zero in the limit, there must exist some constant $T_0>0$ such that $p(t)< p_{min}$ for all $t>T_0$. As a result, we could remove the assumption about $p_\text{min}$ but with a tradeoff on an additional regret term $T_0$, which is independent of $T$, $d$, and $\gamma$. Under the choice of $p(t)=1/log(t)$, VBMLE can be shown to achieve $\mathcal{O}\left(\frac{d\sqrt{T}(\log{T})^5}{(1-\gamma)^2} + \frac{T_0}{1-\gamma}\right)$ regret, without the knowledge of the true $p_\text{min}$. For the detailed proof and description, please refer to Theorem 3, Remark 1, and Appendix C.

Q3: Clarify the related work on linear MDPs and linear kernel MDPs

We thank the reviewer for the helpful suggestions on clarifying the terminology and for bringing some additional references to our attention. In the original manuscript, we follow the terminology used in (Cai et al., 2020), one of the earliest works on this setting, and use the term “linear MDP” for the transition model parameterized as $P(s’|s,a)=<\phi(s’|s,a), \theta^*>$. However, we do agree with the reviewer that the term “linear MDP” could be overloaded as “linear MDP” also widely refers to the low-rank factorization where the transition model takes the form of the inner product between a state-action feature vector and the vector of unknown measures over states. Therefore, we agree that using the term “linear mixture MDP” could avoid the possible ambiguity and would be more consistent with other existing literature.

We have made the changes accordingly in the updated manuscript. Moreover, for clarity, we also reorganize the Related Work to provide a more thorough discussion on these lines of research works.

Q4: Comparison of VBMLE and tabular RL methods and experiments on larger MDPs

We thank the reviewer for the helpful suggestion. For a more extensive evaluation, we further compare VBMLE (augmented with BO) with Posterior Sampling for Reinforcement Learning (PSRL) [3], a popular benchmark method for tabular RL, on an MDP with |S|=100 and |A|=4. Note that this size is already much larger than those used in the existing linear mixture MDP literature, e.g. Riverswim with |S|=6 in [4] . The results for larger MDPs are provided here: https://imgur.com/a/8XXHpYJ. The results demonstrate that VBMLE converges much faster than PSRL and thereby significantly outperforms PSRL [3] in regret performance, across both small and large MDPs.

[3] Ian Osband, Daniel Russo, and Benjamin Van Roy, “(More) Efficient Reinforcement Learning via Posterior Sampling,” NeurIPS 2013.

[4] Alex Ayoub, Zeyu Jia, Csaba Szepesvari, Mengdi Wang, and Lin F. Yang, “Model-Based Reinforcement Learning with Value-Targeted Regression,” ICML 2020.

Q3: Issues with $p_{\text{min}}$ in VBMLE

We would like to clarify the assumption about $p_{\text{min}}$ and the dependency of the regret bound on $p_{\text{min}}$ as follows:

1. The dependency of the regret bound on the true $p_{\text{min}}$: We agree with the reviewer that the regret bound is $p_\text{min}$-dependent and could be shown in the Big-O notation. As shown in the updated Theorem 2, we clearly point out that the regret bound is $\tilde{\mathcal{O}}(\frac{d\sqrt{T}}{p_{\text{min}}^4(1-\gamma)^2})$. Moreover, we would like to highlight that the above dependency on $1/p_{min}$ could be quite conservative in practice. Specifically, we evaluate VBMLE on two MDPs, each with a moderate $p_\text{min}=0.0939$ and a very small $p_{\text{min}}=0.00073$, respectively. The empirical regrets are shown in the figures below. https://imgur.com/a/Nw4Wp4I

We could see that VBMLE could still achieve low empirical regret under a $p_\text{min}$ less than 0.001.

We have also added the above discussion and the additional experimental results in Appendix F.

2. The assumption about the knowledge of the true $p_\text{min}$ could be lifted by using a strictly decreasing positive function $p(t)$ with $\lim_{t \rightarrow \infty} p(t)=0$ as a surrogate for true $p_\text{min}$:
In the original manuscript, we use the knowledge about the true $p_\text{min}$ for two purposes: (i) To ensure that the parameter set $\mathbb{P}$ is not empty; (ii) To provide a lower bound on the term $\phi_t(s_{k+1})^\top \theta_t^{MLE}, \forall k \leq t$. Nevertheless, even without the knowledge of the true $p_\text{min}$, the above two requirements could still be met by using a strictly decreasing function $p(t)$ in the parameter set $\mathbb{P}$. Specifically, as $p(t)$ is strictly decreasing and goes to zero in the limit, there must exist some constant $T_0>0$ such that $p(t)< p_{\text{min}}$ for all $t>T_0$. As a result, we could remove the assumption about $p_\text{min}$ but with a tradeoff on an additional regret term $T_0$, which is independent of $T$, $d$, and $\gamma$. Under the choice of $p(t)=1/log(t)$, VBMLE can be shown to achieve $\mathcal{O}\left(\frac{d\sqrt{T}(\log{T})^5}{(1-\gamma)^2} + \frac{T_0}{1-\gamma}\right)$ regret, without the knowledge of the true $p_\text{min}$. For the detailed proof and description, please refer to Theorem 3, Remark 1, and Appendix C.

Q4: The terminology of linear MDPs and linear mixture MDPs

We thank the reviewer for the suggestion on clarifying the terminology. In the original manuscript, we follow the terminology used in (Cai et al., 2020), one of the earliest works on this setting, and use the term “linear MDP” for the transition model parameterized as $P(s’|s,a)=<\phi(s’|s,a), \theta^*>$. However, we do agree with the reviewer that the term “linear MDP” could be overloaded as “linear MDP” also widely refers to the low-rank factorization where the transition model takes the form of the inner product between a state-action feature vector and the vector of unknown measures over states. Therefore, we agree that using the term “linear mixture MDP” could avoid the possible ambiguity and would be more consistent with other existing literature. We have made the changes accordingly in the updated manuscript.

Q5: Clarify the experimental and implementation details

• Regarding the parameter set $\mathbb{P}$: In the experiments, to showcase that VBMLE is competitive without the need for $p_\text{min}$, we opted not to utilize the parameter set defined in (14) for solving VBMLE; instead, we applied the standard simplex constraints, i.e., directly replacing $p_\text{min}$ with 0 in (14). Notably, the empirical results in Figure 1 in the original manuscript demonstrate that VBMLE indeed remains competitive without the knowledge of the true $p_\text{min}$.
• Additionally, we conduct another experiment in an environment with a relatively small $p_{\text{min}}$, and the results https://imgur.com/a/Nw4Wp4I indicate satisfactory regret performance, even without applying the knowledge of known $p_{\text{min}}$ in Assumption 1 to VBMLE.

For clarity and better reproducibility, we also add the above description and the experimental results to Appendix E in the updated manuscript.

Q2: Clarify the novelty of VBMLE

As highlighted by the reviewer, VBMLE and GPS-IDM [3] differ in two key aspects, and we would like to argue that these two aspects are actually fundamental, require totally different regret analysis, and lead to difference in efficiency:

1. Infinite-horizon vs Episodic (and hence the novel techniques involved in regret analysis): The problem settings are different, given that GPS-IDM focuses on episodic MDPs, while VBMLE is tailored to solving infinite-horizon MDPs. Without the reset capability of episodic MDPs, infinite-horizon discounted MDPs pose a unique challenge of tackling planning and exploration in one single trajectory. Indeed, in many applications, it is practically impossible to reset the state of the system. This difference has also been recognized by the existing literature (e.g., [4] and [5]).

In such a non-episodic setting, the regret analysis is more involved and needs to better handle the mixing behavior of the MDP. Specifically, to address this salient challenge of non-episodic setting, we introduce a new convergence result of MLE in Theorem 1 by (i) constructing a novel submartingale and (ii) connecting the linear mixture MDPs with Follow-the-Leader algorithm in online learning. These two techniques indeed handle the mixing behavior of infinite-horizon MDPs.

2. Computational advantage of VBMLE (due to its frequentist style): We agree with the reviewer that GPS-IDM can be viewed as the Bayesian counterpart of VBMLE (as they both use the value function for exploration). That being said, by taking the frequentist viewpoint, VBMLE could enjoy a computational advantage over its Bayesian counterpart. Specifically:

• In GPS-IDM, the update of posterior distribution can be computationally expensive in practice. As shown in Line 3 of Algorithm 3 in [3], each posterior update requires doing one run of standard Value Iteration for every possible model in the hypothesis class. Note that the hypothesis class could be fairly large in practice, especially under the $\epsilon$-bracket number technique for discretizing the continuous domain of model parameters (cf Appendix G in [3]), and hence the computational complexity could be substantial.

• By contrast, VBMLE takes the frequentist perspective and only requires a small number of runs of Value Iteration when it is augmented with BO (as described in Q1). This manifests the computational efficiency of VBMLE.

[3] Han Zhong et al., "GEC: A Unified Framework for Interactive Decision Making in MDP, POMDP, and beyond," arXiv preprint arXiv:2211.01962, 2022.

[4] Yuanzhou Chen, Jiafan He, and Quanquan Gu, “On the Sample Complexity of Learning Infinite-horizon Discounted Linear Kernel MDPs,” ICML 2022.

[5] Dongruo Zhou, Jiafan He, and Quanquan Gu, “Provably Efficient Reinforcement Learning for Discounted MDPs with Feature Mapping,” ICML 2021.

Q1: Explain the implementation of VBMLE in practice

We would like to highlight that the VBMLE maximization problem in (10) can be approximately solved by multiple approaches in practice:

(i) Solve VBMLE by Bayesian optimization: Due to the non-concavity of VBMLE, we propose to solve VBMLE by Bayesian optimization (BO), which is a powerful and generic method for provably maximizing (possibly non-concave) black-box objective functions. As a result, BO can provably find an $\epsilon$-optimal solution to VBMLE within finite iterations. Specifically:

• We have applied the GP-UCB algorithm, which is one classic BO algorithm and has been shown to provably find an $\epsilon$-optimal solution within $\tilde{\mathcal{O}}(1/\epsilon^2)$ iterations under smooth (possibly non-concave) objective functions [2]. Each sample taken by GP-UCB requires only one run of standard Value Iteration.
• To further demonstrate the compatibility of VBMLE and BO, we have extended the regret analysis of VBMLE to the case where only an $\epsilon$-optimal VBMLE solution is obtained (as also suggested by Reviewer p6f5). Specifically, let $K$ denote the number of samples taken by GPUCB in each maximization run of finding VBMLE. We show that VBMLE augmented with BO can achieve a regret bound of \mathcal{O}\(\max(\frac{d\sqrt{T}\log{T}}{p_{\text{min}^2}(1-\gamma)^2},\sqrt{T}\frac{(\log{K})^{d+1}}{\sqrt{K}} )). By using a moderate $K$, one could easily recover the same regret bound as that of VBMLE with an exact maximizer. In our experiments, we find that choosing $H=25$ is sufficient and also computationally efficient. The empirical regret of VBMLE augmented with BO and other benchmark methods can be found at https://imgur.com/a/8XXHpYJ and in Figure 1 of the updated manuscript. The computation times of VBMLE and other methods are provided below https://imgur.com/a/5ksC7Xd and also shown in Table 1 of the updated manuscript. We can see that VBMLE with BO enjoys both low empirical regret as well as low computation time.

We have also added the discussion on VBMLE with BO in Appendix C in the updated manuscript.

(ii) Approximately solve VBMLE by off-the-shelf constrained optimization solvers: In practice, another possibility is to use an off-the-shelf optimization solver. For example, we have also applied the trust region method (implementation available in SciPy), which is compatible with the VBMLE problem and only requires taking the VBMLE objective function as an input argument (and this is used for constructing the trust-region subproblem, e.g., please refer to [1]). However, as mentioned by the reviewer, one major issue is that due to the non-concavity of VBMLE objective, trust-region methods only ensure convergence to local optima. Despite this, from the experiments, we find that VBMLE augmented with a trust-region method achieves regret comparable to VBMLE with BO.

[1] Andrew R. Conn, Nicholas IM Gould, and Philippe L. Toint, “Trust region methods,” Society for Industrial and Applied Mathematics, 2000.

[2] Niranjan Srinivas, Andreas Krause, Sham M Kakade, and Matthias W Seeger, “Information-theoretic regret bounds for Gaussian process optimization in the bandit setting,” IEEE Transactions on Information Theory, 2012.