Sample-efficient Learning of Infinite-horizon Average-reward MDPs with General Function Approximation

Thank you for your comments and assessments!

What are the novelties in the technical or algorithmic level in this paper? What are the new challenges for exploration in infinite horizon averaged reward setting?

• Please see our general response for details.

It seems to me the main contribution of this paper is to transfer and generalize some existing techniques and results (especially in finite horizon setting) to infinite horizon averaged return setting. The definition of AGEC Complexity Measure, FLOP algorithm and the techniques for analyzing lazy update rules seem share much similarity with previous literature like [1], [2]. I didn't find much novelty in algorithm design or technique analysis.

• While we admit that the structural assumption over the Bellman dominance shares a similar form as the standard GEC (Zhong et al., 2022), the transferability assumption is new. Moreover, our algorithm design is different from the posterior-sampling-based approach presented in Zhong et al. (2022). We opt for an optimism-based algorithm for establishing a low policy switch updating, which is crucial in the achievement of sublinear regret for AMDPs.
• Compared with Xiong et al. (2023), we adopt a weaker assumption over the Bellman dominance and provide a more rigorous definition of the eluder condition with nuanced adpation to the average-reward setting. Moreover, the average-reward RL adopts a completely different notion of regret, and the Q- or V- function does not naturally exist due to the infinite horizon without discount. we consider a broad subclass where a modified version of the Bellman optimality equation holds and present a novel regret decomposition for theoretical analysis of the explorability of AMDPs, resulting in significantly different regret analysis from the aforementioned work.

Is the "lazy policy update" really necessary? Although the authors explain the motivation for low policy switching is because of the additional cost in regret analysis. I'm curious whether it can be avoidable or it reveals some fundamental difficulty.

• That's a really interesting question. There are two primary approaches to address online learning of AMDPs: the reduction to episodic MDPs, and the utilization of a low policy switching algorithm. However, achieving $\mathcal{\tilde{O}}(\sqrt{T})$ has proven to be challenging for reduction-based algorithms in existing works under the online learning framework (Wei et al., 2019, Wei et al., 2021), unless fortified with strong assumptions, such as ergodic AMDPs, where every policy is guaranteed to reach every state after a sufficient number of steps. In our pursuit of a near-optimal algorithm without relying on such stringent assumptions, we introduce an additional yet mild structural assumption to ensure the validity of a low-switching algorithm for general function classes.

• At out best knowledge, all existing algorithms for online learning of AMDPs depend on such lazy update strategy to achieve a near optimal regret beyond strong assumptions, from the classical UCRL2 algorithm for tabular AMDPs (Jaksch et al., 2010), to FOPO for linear AMDPs (Wei et al., 2021) and (UCRL2-VTR for linear mixture MDPs (Wu et al., 2022). Furthermore, we speculate that the underlying mechanism behind this phenomenon lies in the challenge of accurately evaluating $\pi_t$ when imposing a step-by-step policy update, i.e., $\pi_{t}\neq\pi_{t+1}$ for all $t\in[T]$. This difficulty arises due to the intricate dependency structure in the historical data $\{(s_i,a_i)\}_{i\leq[t]}$.

• Specifically, in the episodic setting, the learner executes the policy $\pi_t$ throughout the epside, and the value function can be computed by bellman equation starting from a zero function. Thus, the error can be controlled and is bounded by the cummulative error across episode following decoposition ${\rm Reg}(T)\leq\sum_{t=1}^T\sum_{h=1}^H\mathbb{E}[\mathcal{E}(f_t)(s_h^t,a_h^t)\vert\pi_t]$. This means that the estimation error of the current policy can be controlled by the sum of uncertainty. But for average reward case, if you update policy in each iteration, we cannot bound the suboptimality of the current policy by a sum of estimation errors as ${\rm Reg}(T)\leq\sum_{t=1}^T\mathcal{E}(f_t)(s_t,a_t)+\sum_{t=1}^T\left(\mathbb{E}[V_t(s_{t+1})]-V_t(s_t)\right)$.

• The low-switching policy update in our algorithm ensures that the learner scrutinizes the policy based on data collected from new epochs. This corresponds to the update condition in our algorithm, where we only update when a substantial increase in cumulative discrepancy surpassing $3\beta$ has occurred since the last update.

Thank you for your comments and assessments!

The definition of AGEC (Definition 3), which is the central object and contribution in this paper, lacks clarity.

• Thank you for pointing out the lack of rigor in our definitions and your comments have been well-taken. In the original version, we used $\mathcal{O}$ notation to hide the contants in the definition, which may not adhere to mathematical rigor. In the revised version, we substitute the $\mathcal{O}$ notation with explicit definitions of absolute constants (highlighted in blue), which does not hurt the correctness of the result.
• Regarding the dicrepancy function, we would like to highlight that the definition of the discrepancy function is inherently problem-specific, making it challenging to offer a unified definition applicable across all cases. To ensure clarity in our presentation, we provide a detailed definition of the problem-specific discrepancy function in each corresponding section such as (3.1), (3.3), (C.2).

The function approximation error can be critical in RL, as its multiplying factor usually depends on the exploration performance in various forms. If we remove the realizability assumption, how does the additional term depend on the complexity measure defined in this paper?

• We would like to clarify that our complexity measure characterizes the difficulty in exploring the best-in-class model within the designated hypothesis class. On the other hand, the approximation error stems from the disparity between the best-in-class approximator and the ground-truth optimal function, which is not theoretically provable. Consequently, the removal of the realizability assumption will sure introduce additional error terms; however, it does not depend on the proposed complexity measure.
• A generalization of realizability for value-based problems is that there exists $f\in\mathcal{F}$ and $\epsilon_{\rm real}>0$ such that $\Vert Q^*-Q_f\Vert_\infty\leq\epsilon_{\rm real}$. The condition requires that the optimal bias state-action value function lies in the function class approximately with $\epsilon_{\rm real}$ error, enabling a more nuanced analysis on approximation. In this case, we conjecture that our algorithm, designed with a larger $\beta$ to accommodate misspecification, is capable of finding an $(\epsilon + \mathrm{poly}(d, \mathrm{sp}(V^*)) \cdot \epsilon_{\rm real})$-optimal policy using polynomial (in $\epsilon$) samples. It is indeed a very interesting future direction to design a unified sample-efficient algorithm considering the approximation error under general function approximation. We leave more meticulous analysis as future work.

In Equation 2.1, what is $V^*(s,a)$. Should it be $Q^*$ ?

• Thank you for pointing that out. We have rectified the typo (Page 3), which should be $J^*+Q^*(s,a)=r(s,a)+\mathbb{E}_{s'\sim\mathbb{P}(\cdot|s,a)}[V^*(s')].$

In the abstract and in multiple places in the paper, ${\rm sp}(v^*)$ appears with $v^*$. Should it be ${\rm sp}(V^*)$ ? In the paper, it is assumed that ${\rm sp}(V^*)$ is known. Is this knowledge necessary?

• Thank you for pointing out the confusion in the notations. We have now replaced ${\rm sp}(v^*)$ by ${\rm sp}(V^*)$ for a clearer presentation in the revised version. Regarding the span ${\rm sp}(V^*)$, we assert that the knowledge of the span of the optimal state bias function is crucial for establishing a more rigorous regret guarantee due to theoretical necessity. It is noteworthy that, unlike the episodic setting, there is no inherent upper bound for the optimal function $V^*$ in the average-reward setting, leading to a hypothesis class with uncontrollable complexity. The assumption of a known span implies a bounded span for the optimal bias function, with ${\rm sp}(V^)$ being available to the learner. This knowledge allows for the construction of a more stringent optimistic parameter. Alternatively, in practical scenarios, the learner can choose a conservative estimation of the span, such as a sufficiently large constant $B_{sp}>0$, ensuring $B_{sp}\geq{\rm sp}(V^*)$. The substitution of ${\rm sp}(V^*)$ with $B_{sp}>0$ in the construction of $\beta$ ensures the persistence of the optimistic planning in Algorithm 1, and results in an inflation of regret in squared root.

We would like to thank the reviewer for the positive assessment of this work! Here's our response to the question.

Can authors provide a brief comparison between this work and existing works in terms of the regret results?

• Please see our general response for details.

Thank you for your comments and assessments!

The sample complexity is based on GEC and is not entirely novel. Specifically, Definition 3 (Bellman Dominance) in this paper is of the same form with Definition 3.4 of Zhong et al. (2022b), where is defined such that the sum of the Bellman error being less than the in-sample training error plus the burn-in cost.

• While we admit that the structural assumption over the Bellman dominance shares a similar form as the standard GEC (Zhong et al., 2022b) with nuanced adaption in average-reward terms, the transferability assumption is a novel addition. This additional structural constraint presents the unique challenges of AMDPs, and plays a pivotal role in facilitating the design of a low-switching algorithm. This is particularly essential in addressing the inherent difficulty in exploring AMDPs.
• Moreover, our algorithm design diverges from the posterior-sampling-based approach presented in Zhong et al. (2022b). We embrace an optimism-based algorithm to establish a low policy switch updating mechanism, a crucial element in achieving sublinear regret for AMDPs. As a consequence, the technical details of our regret analysis differ significantly from the aforementioned work. Please refer to our general response more desciptions of technical novelties.

However, there seems a lack of discussion regarding alternative choices for the discrepancy function, such as the Hellinger distance-based discrepancy in Zhong et al. (2022b).

• We appreciate the reviewer's suggestion regarding the importance of discussing alternative choices for the discrepancy function, including the Hellinger distance-based discrepancy. Discussions about MLE-AGEC concerning an alternative discrepancy functions are presented in Appendix C (highlighted in blue) with a different transferability assumption over the structure. Here, we choose the total variation (TV)-based discrepancy function instead, which results from the slight difference between the choice of discrepancy: $l_{f'}(f,g,\zeta_t)=\frac{1}{2}|P_f(s_{t+1}|s_t,a_t)/P_{f^*}(s_{t+1}|s_t,a_t)-1|$ (ours) and $l_{f'}(f,\zeta_t)=\frac{1}{2}(P_f(s_{t+1}|s_t,a_t)/P_{f^*}(s_{t+1}|s_t,a_t)-1)^2$ (Zhong et al., 2022b). We remark that both choices of discrepancy is sufficient to cover MLE-based algorithms.

The algorithm is based on upper confidence bound, which the confidence region chosen based on the discrepancy function. This approach is closely related to Algorithm 1 in Chen et al. (2022b) and Algorithm 1 in Jin et al. (2021). The proposed algorithm is usually impractical to implement, since it involves solving a global constrained optimization.

• While acknowledging the resemblance of our algorithm design to previous lines of work on optimism-based algorithms for general function approximation, we notably diverge by employing a low-switch updating strategy coupled with a discrepancy-based updating condition.
• Moreover, we would like to comment that our algorithm is in some sense information theoretical -- our work identifies a statistically tractable subclass of AMDPs, and the algorithm serves as a certificate for the regret upper bound. The main focus of our paper is to address the problem of minimal structural assumptions that empower sample-efficient learning in AMDPs, which facilitates algorithm design beyond specific problems such as linear AMDPs and linear mixture AMDPs.

Could the authors elaborate on the primary technical challenges they encountered while deriving their theoretical results, when adapt the GEC in Zhong et al. (2022b) to the infinite-horizon average-reward setting, with a constrained algorithm, which seems to be a generalization of the framework established by Jin et al. (2021)?

• Please see our general response for details.

Can the authors elaborate the optimality of the regret in Theorem 3, when restricted to each specific instances? i.e. linear mixture AMDP, linear AMDP, etc.

• Please see our general response for details.

[A1] Han Zhong, Wei Xiong, Sirui Zheng, Liwei Wang, Zhaoran Wang, Zhuoran Yang, and Tong Zhang. Gec: A unified framework for interactive decision making in mdp, pomdp, and beyond. arXiv preprint arXiv:2211.01962, 2022a.

Besides, we provide a discussion about the performance of our algorithm on specific instances in the following, and please refer to Appendix F.5 for detailed calculations.

• Performance on concrete examples: Our algorithm achieves a $\tilde{\mathcal{O}}\Big({\rm sp}(V^*)^\frac{3}{2}d^\frac{3}{2}\sqrt{T}\Big)$ regret for both linear AMDPs and linear mixture AMDPs, and attains $\tilde{\mathcal{O}}({\rm sp}(V^*)d_{\rm E}\sqrt{T\beta})$ for model-based problems. In comparasion, FOPO (Wei et al., 2021) achieves $\tilde{\mathcal{O}}\Big({\rm sp}(V^*)d^\frac{3}{2}\sqrt{T}\Big)$ for linear AMDPs, UCRL2-VTR (Wu et al., 2022) achieves $\mathcal{\tilde{O}}(d\sqrt{DT})$ for linear mixture AMDPs, and SIM-TO-REAL (Chen et al., 2022a) attains $\tilde{\mathcal{O}}(Dd_{\rm E}\sqrt{T\beta})$ for model-based problems. In Wu et al., (2022), a lower bound of the order $\mathcal{{O}}(d\sqrt{DT})$ is proved for linear mixture AMDPs, where the diameter $D$ is closely related but incomparable with the span ${\rm sp}(V^*)$ considered in our paper. Based on the argument, we posit that our algorithm is nearly optimal concerning the episode length $T$.
• Suboptimality: While conceding that our algorithm introduces additional multiplying factors when compared to the minimax rate and existing algorithms, we assert that this is an unavoidable trade-off inherent in developing an algorithm for general function classes. The augmented regret is a consequence of a more rigorous design of the optimistic parameter tailored for specific problems. For instance, in linear AMDPs, the learner can structure the optimistic bonus term for each step as $\| b_t \|_{\Lambda_t}\leq\beta$ based on the empirical covariance matrix, which is inattainable for general function classes. We remark that such suboptimality is fairly common in general function approximation literature (Jin et al., 2021, Zhong et al., 2022).
• The main focus of our paper is to address the problem of minimal structural assumptions that empower sample-efficient learning in AMDPs, rather than striving to enhance regret results. Our algorithm is in some sense information theoretical -- our work identifies a statistically tractable subclass of AMDPs, and the algorithm serves as a certificate for the regret upper bound.

Finally, we would like to emphasize the significance of our work. Our proposed complexity measure, AGEC, not only unifies all known tractable AMDPs but all includes some new AMDPs identified by us. Furthermore, our generic algoirthm, applicable to all AMDPs with low AGEC, positions our work as the first comprehensive framework for AMDPs. Additionally, our study establishes a connection between sample-efficient learning for AMDPs and lazy-updating algorithms, supported by a general theoretical analysis pipeline. We firmly believe that these contributions are paramount in advancing the understanding of AMDPs, a category of MDP that has, to some extent, been overlooked when compared to finite-horizon MDPs.

[A1] Chi Jin, Qinghua Liu, and Sobhan Miryoosefi. Bellman eluder dimension: New rich classes of rl problems, and sample-efficient algorithms. Advances in neural information processing systems,34:13406–13418, 2021.

[A2] Han Zhong, Wei Xiong, Sirui Zheng, Liwei Wang, Zhaoran Wang, Zhuoran Yang, and Tong Zhang. Gec: A unified framework for interactive decision making in mdp, pomdp, and beyond. arXiv preprint arXiv:2211.01962, 2022a.

[A3] Chen-Yu Wei, Mehdi Jafarnia Jahromi, Haipeng Luo, and Rahul Jain. Learning infinite-horizon average-reward mdps with linear function approximation. In International Conference on Artificial Intelligence and Statistics, pp. 3007–3015. PMLR, 2021.

[A4] Yue Wu, Dongruo Zhou, and Quanquan Gu. Nearly minimax optimal regret for learning infinite horizonaverage-reward mdps with linear function approximation. In International Conference on Artificial Intelligence and Statistics, pp. 3883–3913. PMLR, 2022.

[A5] Xiaoyu Chen, Jiachen Hu, Chi Jin, Lihong Li, and Liwei Wang. Understanding domain randomization for sim-to-real transfer. International Conference on Learning Representations, 2022a.