Provably Efficient Reinforcement Learning with Multinomial Logit Function Approximation

Authors: long-fei  li
yu-jie  zhang
peng  zhao
zhi-hua  zhou
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 08:50:39

Abstract: We study a new class of MDPs that employs multinomial logit (MNL) function approximation to ensure valid probability distributions over the state space. Despite its benefits, introducing the non-linear function raises significant challenges in both *computational* and *statistical* efficiency. The best-known result of Hwang and Oh [2023] has achieved an $\widetilde{\mathcal{O}}(\kappa^{-1}dH^2\sqrt{K})$ regret, where $\kappa$ is a problem-dependent quantity, $d$ is the feature dimension, $H$ is the episode length, and $K$ is the number of episodes. While this result attains the same rate in $K$ as linear cases, the method requires storing all historical data and suffers from an $\mathcal{O}(K)$ computation cost per episode. Moreover, the quantity $\kappa$ can be exponentially small in the worst case, leading to a significant gap for the regret compared to linear function approximation. In this work, we first address the computational and storage issue by proposing an algorithm that achieves the same regret with only $\mathcal{O}(1)$ cost. Then, we design an enhanced algorithm that leverages local information to enhance statistical efficiency. It not only maintains an $\mathcal{O}(1)$ computation and storage cost per episode but also achieves an improved regret of $\widetilde{\mathcal{O}}(dH^2\sqrt{K} + d^2H^2\kappa^{-1})$, nearly closing the gap with linear function approximation. Finally, we establish the first lower bound for MNL function approximation, justifying the optimality of our results in $d$ and $K$.

Main paper: 
 introduction: 
 Reinforcement Learning (RL) with function approximation has achieved remarkable success in various applications involving large state and action spaces, such as games [Silver et al., 2016] , algorithm discovery [Fawzi et al., 2022] and large language models [Ouyang et al., 2022] . Therefore, establishing the theoretical foundation for RL with function approximation is of great importance. Recently, there have been many efforts devoted to understanding the linear function approximation, yielding numerous valuable results [Yang and Wang, 2019 , Jin et al., 2020 , Ayoub et al., 2020] .While these studies make important steps toward understanding RL with function approximation, there are still some challenges to be solved. In linear function approximation, transitions are assumed to be linear in specified feature mappings, such as P(s ′ |s, a) = ϕ(s ′ |s, a) ⊤ θ * for linear mixture MDPs and P(s ′ |s, a) = ϕ(s, a) ⊤ µ * (s ′ ) for linear MDPs. Here P(s ′ |s, a) is the probability from state s to state s ′ taking action a, ϕ(s ′ |s, a) and ϕ(s, a) are feature mappings, θ * and µ * (s ′ ) are unknown parameters. However, a transition function is a probability distribution over states, i.e., the magnitude falls within the range of [0, 1], and the aggregation equals 1. For certain feature mappings, linear transitions may not yield a valid probability distribution with arbitrary θ and µ as show by [Hwang and Oh, 2023] . While there are also some works exploring generalized linear [Wang et al., 2021] and general function approximation [Russo and Roy, 2013 , Foster et al., 2021 , Chen et al., 2023] , they made assumptions over value functions rather than transitions, hence do not tackle this challenge.38th Conference on Neural Information Processing Systems (NeurIPS 2024).Table 1 : Comparison between previous works and ours in terms of the regret, computation cost and storage cost. Here κ and κ * are exponentially small problem-dependent quantities defined in Assumption 1, d is the feature dimension, H is the episode length and K is the number of episodes. The computational cost and storage cost per episode indicate the dependence on episode count k. Hwang and Oh [2023] recently. Despite its benefits, the introduction of non-linear functions raises significant challenges in both computational and statistical efficiency. Specifically, the best-known approach of Hwang and Oh [2023] has achieved an O(κ -1 dH 2 √ K) regret. While matching the optimal regret of O(d √ H 3 K) as linear function approximation in terms of K [Zhou et al., 2021 , He et al., 2023] , their method requires storing all historical data and the computational cost per episode grows linearly with the episode count, which is expensive. Moreover, the quantity κ can be exponentially small, leading to a significant gap for the regret compared to the linear cases. Thus, a natural question arises:Is MNL function approximation more difficult than linear function approximation for RL?In this paper, we answer this question by making significant advancements in both computational and statistical efficiency for MDPs with MNL function approximation, nearly closing the gap with linear function approximation. Table 1 presents a comparison between previous studies and our work. Specifically, our contributions are summarized as follows:• We first propose the UCRL-MNL-OL algorithm based on a variant of Online Newton Step, attaining an O(κ -1 dH 2 √ K) regret with O(1) computational and storage costs per episode. This result matches the best-known regret by Hwang and Oh [2023] , yet achieves the same computational and storage efficiency as linear function approximation [Jin et al., 2020 , Zhou et al., 2021] .• As the quantity κ can be exponentially small, we propose an enhanced algorithm UCRL-MNL-LL, which leverages the local information to improve the statistical efficiency. It not only maintain O(1) computational and storage costs but achieve an improved regret of O(dH 2 √ K + d 2 H 2 κ -1 ). The higher-order term matches the optimal regret O(d √ H 3 K) for linear mixture MDPs [Zhou et al., 2021] and linear MDPs [He et al., 2023] in d and K, differing only by an O(H 1/2 ) factor.• We establish the first lower bound for MDPs with MNL function approximation by introducing a reduction to the logistic bandit problem. We prove a lower bound of Ω(dH √ Kκ * ) for MDPs with infinite action space. Though this does not constitute a strict lower bound for the finite actions case studied in this work, it suggests that our result may be optimal with respect to d and K. 1 From a technical perspective, inspired by recent advances in logistic bandits [Zhang et al., 2016 , Faury et al., 2020 , Zhang and Sugiyama, 2023] , we observe that the negative log-likelihood function is exponentially concave, which motivates us to apply online Newton step to estimate the unknown transition parameter in an online manner. Moreover, we employ the Bernstein-like inequalities and the self-concordant-like property [Bach, 2010] of the log-loss to achieve the better dependence on κ.Organization. We introduce the related work in Section 2 and present the setup in Section 3. Then, we design an algorithm with constant computational and storage costs in Section 4. Next, we present an algorithm with improved regret guarantee in Section 5. Finally, we establish the lower bound in Section 6. Section 7 concludes the paper. Due to space limits, we defer all proofs to the appendixes.Notations. Denote by [n] the set {1, . . . , n} and use [x] [a,b] to denote min(max(x, a), b). For X, Y ∈ R d×d , X ⪰ Y means X -Y is positive semi-definite. For a vector x ∈ R d and positive semi-definite matrix A ∈ R d×d , denote ∥x∥ A = √x ⊤ Ax. The O(•) hides polylogarithmic factors. 
related work: 
 In this section, we review related works from both setup and technical perspectives.RL with Generalized Linear Function Approximation. There are recent efforts devoted to investigating function approximation beyond the linear models. Wang et al. [2021] investigated RL with generalized linear function approximation. Notably, unlike our approach which models transitions using a generalized linear model, they apply this approximation directly to the value function. Another line of works [Chowdhury et al., 2021 , Li et al., 2022 , Ouhamma et al., 2023] has studied RL with exponential function approximation and also aimed to ensure that transitions constitute valid probability distributions. The MDP model can be viewed as an extension of bilinear MDPs in their work while our setting extends linear mixture MDPs. These studies are complementary to ours and not directly comparable. Moreover, these works also enter the computational and statistical challenges arising from non-linear function approximation that remain to be addressed. The most relevant work to ours is the recent work by Hwang and Oh [2023] , which firstly explored a similar setting to ours, where the transition is characterized using a multinomial logit model. We significantly improve upon their results by providing computationally and statistically more efficient algorithms.There have also been some works studies RL with general function approximation. Russo and Roy [2013] and Osband and Roy [2014] initiated the study on the minimal structural assumptions that render sample-efficient learning by proposing a structural condition called Eluder dimension. Recently, several works have investigated different conditions for sample-efficient interactive learning, such as Bellman Eluder (BE) dimension [Jin et al., 2021] , Bilinear classes [Du et al., 2021] , Decision-Estimation Coefficient (DEC) [Foster et al., 2021] , and Admissible Bellman Characterization (ABC) [Chen et al., 2023] . A notable difference is that they impose assumptions on the value functions while we study function approximation on the transitions to ensure valid probability distributions. Moreover, the goal of these works is to study the conditions for sample-efficient reinforcement learning, but not focus on the computational efficiency.Multinomial Logit Bandits. There are two types of multinomial logit bandits studied in the literature: the single-parameter model, where the parameter is a vector [Cheung and Simchi-Levi, 2017] and multiple-parameter model, where the parameter is a matrix [Amani and Thrampoulidis, 2021] . We focus on the single-parameter model, which are more relevant to our setting. The pioneering work by Cheung and Simchi-Levi [2017] achieved a Bayesian regret of O(κ -1 d √ K), where K denotes the number of rounds in bandits. This result was further enhanced by subsequent studies [Oh and Iyengar, 2019 , 2021 , Agrawal et al., 2023] . In particular, Périvier and Goyal [2022] significantly improved the dependence on κ, obtaining a regret of O(d √ κK + κ -1 ) in the uniform revenue setting. Most prior methods required storing all historical data and faced computational challenge. To address this issue, the most recent work by Lee and Oh [2024] proposed an algorithm with constant computational and storage costs building on recent advances in multiple-parameter model [Zhang and Sugiyama, 2023] . Their algorithm achieves the optimal regret of O(d√ κK + κ -1 ) and O(d √ K + κ -1) under uniform and non-uniform rewards respectively. However, although the underlying models of MNL bandits and MDPs share similarities, the challenges they present differ substantially, and techniques developed for MNL bandits cannot be directly applied to MNL MDPs. For example, in MNL bandits, the objective is to select a series of assortments with varying sizes that maximize the expected revenue, whereas in MNL MDPs, the goal is to choose one action at each stage to maximize the cumulative reward. Thus, it is necessary to design new algorithms tailored for MDPs to address these unique challenges. 
problem setup: 
 In this section, we present the problem setup of RL with multinomial logit function approximation.Inhomogeneous, Episodic MDPs. An inhomogeneous, episodic MDP instance can be denoted by a tuple M = (S, A, H, {P h } H h=1 , {r h } H h=1 ), where S is the state space, A is the action space, H is the length of each episode, P h : S × A × S → [0, 1] is the transition kernel with P h (s ′ | s, a) is being the probability of transferring to state s ′ from state s and taking action a at stage h, r h : S × A → [0, 1] is the deterministic reward function. A policy π = {π h } H h=1 is a collection of mapping π h , where each π h : S → ∆(A) is a function maps a state s to distributions over A at stage h. For any policy π and (s, a) ∈ S × A, we define the action-value function Q π h and value function V π h as follows:Q π h (s, a) = E H h ′ =h r h ′ (s h ′ , a h ′ ) s h = s, a h = a , V π h (s) = E a∼π h (•|s) [Q π h (s, a)] ,where the expectation of Q π h is taken over the randomness of the transition P and policy π. The optimal value function V * h and action-value functionQ * h given by V * h (s) = sup π V π h (s) and Q * h (s, a) = sup π Q π h (s, a). For any function V : S → R, we define [P h V ](s, a) = E s ′ ∼P h (•|s,a) V (s ′ ).The Bellman equation for policy π and Bellman optimality equation are given byEQUATIONs, a). Learning Protocol. In the online MDP setting, the learner interacts with the environment without the knowledge of the transition {P h } H h=1 . As learning rewards is no harder than transitions, we assume the reward {r h } H h=1 is known. The interaction proceeds in K episodes. At the beginning of episode k, the learner chooses a policyπ k = {π k,h } H h=1 . At each stage h ∈ [H],starting from the initial state s k,1 , the learner observes the state s k,h , chooses an action a k,h sampled fromπ k,h (• | s k,h ), obtains reward r h (s k,h , a k,h ) and transits to the next state s k,h+1 ∼ P h (• | s k,h , a k,h ) for h ∈ [H].The episode ends when s H+1 is reached. The goal of the learner is to minimize regret, defined asReg(K) = K k=1 V * 1 (s k,1 ) - K k=1 V π k 1 (s k,1 ),which is the difference between the cumulative reward of the optimal policy and the learner's policy.Multinomial Logit (MNL) Mixture MDPs. Although significant advances have been achieved for MDPs with linear function approximation, Hwang and Oh [2023] show that there exists a set of features such that no linear transition model (including bilinear and low-rank MDPs) can induce a valid probability distribution over the state space, which limits the expressiveness of function approximation. (s ′ | s, a) : S × A × S → R d with ∥ϕ(s ′ | s, a)∥ 2 ≤ 1 and unknown vectors {θ * h } H h=1 ∈ Θ with Θ = {θ ∈ R d , ∥θ∥ 2 ≤ B}, such that for all (s, a, h) ∈ S × A × [H] and s ′ ∈ S h,s,a , it holds that P h (s ′ | s, a) = exp(ϕ(s ′ | s, a) ⊤ θ * h ) s∈S h,s,a exp(ϕ( s | s, a) ⊤ θ * h ).Remark 2. While the work of Hwang and Oh [2023] focuses on the homogeneous setting, where the transitions remain the same across all stages (i.e., P 1 = ... = P H ), we address the more general inhomogeneous setting, allowing transitions to vary across different stages.For any θ ∈ R d , we define the induced transition asp s ′ s,a (θ) = exp(ϕ(s ′ | s,a) ⊤ θ)s∈Ss,a exp(ϕ( s | s,a) ⊤ θ) . We introduce the key problem-dependent quantities for this problem firstly introduced by Hwang and Oh [2023] . Assumption 1. There exists 0 < κ ≤ κ * < 1 such that for all (s, a, h) ∈ S × A × [H] and s ′ , s ′′ ∈ S h,s,a , it holds that inf θ∈Θ p s ′ s,a (θ)p s ′′ s,a (θ) ≥ κ and p s ′ s,a (θ * )p s ′′ s,a (θ * ) > κ * . Assumption 1 is similar to the standard assumption in generalized linear model literature [Filippi et al., 2010 , Wang et al., 2021] and logistic bandits literature [Faury et al., 2020 , Abeille et al., 2021 , Zhang and Sugiyama, 2023] to guarantee the Hessian matrix is non-singular in Property 1. 
computationally efficient algorithm by online learning: 
 In this part, we design an algorithm that achieves O(1) computational and storage costs per episode.h is unknown, we need to estimate it using the historical data. At episode k, we collect a trajectory {(s k,h , a k,h )} H h=1 , then define the variable: θ) . Then y k,h is a sample from the following multinomial distribution:EQUATIONy k,h ∼ multinomial(1, [p s1 k,h (θ * ), . . . , p s N k,h k,h (θ * )]), where the parameter 1 indicates that y k,h is a single-trial sample. Furthermore, we define the noise Hwang and Oh [2023] made the first step by proposing an approach using the maximum likelihood estimation (MLE). Specifically, their estimator θ k,h is defined asϵ s ′ k,h = y s ′ k,h -p s ′ k,h (θ * h ). It is clear that ϵ k,h ∈ [-1, 1] N k,h , E[ϵ k,h ] = 0 and s ′ ∈S k,h ϵ s ′ i,h = 0. Efficiency Concern.θ k,h = arg min θ∈R d k-1 i=1 s ′ ∈S i,h -y s ′ i,h log p s ′ i,h (θ) + λ 2 ∥θ∥ 2 2 .(1)Despite favorable theoretical guarantee, the computational and storage cost of this method is expensive. First, in episode k, the estimator θ k,h in (1) is computed using all samples collected in previous episodes, resulting an O(k) storage cost. Second, to solve this optimization problem, the projected gradient descent [Boyd and Vandenberghe, 2004] are usually applied. However, as discussed in Faury et al. [2022] , the optimization of the MLE problem typically require O(k log(1/ϵ)) iterations to achieve an ϵ-accurate solution, which is computationally expensive.To address this issue, we introduce a novel algorithm named UCRL-MNL with Online Learning (UCRL-MNL-OL), which achieves both the same regret with O(1) computational and storage cost. At a high level, our algorithm can be divided into two phases: (i) efficient online estimation, and (ii) efficient optimistic value function construction. We introduce the details in the following.Efficient Online Estimation. Instead of using all historical data, our algorithm updates the estimator in an online manner. Specifically, inspired by the works [Hazan et al., 2014 , Zhang et al., 2016] on logistic bandit, we find the negative log-likelihood function is exponentially concave, which motivates us to apply a variant of online Newton step [Zhang et al., 2016, Oh and Iyengar, 2021] .First, we define per-episodic loss function f k,h (θ), gradient g k,h (θ) and Hessian matrix H k,h (θ) asEQUATIONWe introduce the following key property of the loss function f k,h (θ) by a second-order Taylor expansion and mean value theorem, whose proof can be found in Lemma 5 in Appendix A. Property 1. (i). For any θ) .EQUATION(ii). For any θ ∈ Θ, it holds thatEQUATIONProperty 1 implies that f k,h (θ) is exponentially concave, which enables us to apply the online Newton step to construct an efficient estimator. Specifically, we construct θ k+1,h by solving the problem:EQUATIONwhereΣ k+1,h = Σ k,h + κ Algorithm 1 UCRL-MNL-OL Input: Regularization parameter λ, confidence width β k , confidence parameter δ. 1: Initialization: Σ 1,h = λI, θ 1,h = 0 for all h ∈ [H]. 2: for k = 1, . . . , K do 3: Compute Q k,h (•,•) in a backward way as in (7).for h = 1, . . . , H do 5:Observe current state s k,h , select action a k,h = arg max a∈A Q k,h (s k,h , a).6: Set Σ k+1,h = Σ k,h + κ 2 s ′ ∈S k,h ϕ s ′ k,h (ϕ s ′ k,h ) ⊤ . 7: Compute θ k+1,h = arg min θ∈Θ ⟨g k,h ( θ k,h ), θ -θ k,h ⟩ + 1 2 ∥θ -θ k,h ∥ 2 Σ k+1,h . 8:end for 9: end for We show the online estimator θ k,h in (4) enjoys computational and storage efficiency simultaneously. As the optimization problem exhibits a standard online mirror descent formulation, it can be solved with a projected gradient step with the following equivalent formulation byθ ′ k+1,h = θ k,h -Σ -1 k+1,h g k,h ( θ k,h ), and θ k+1,h = arg min θ∈Θ ∥θ -θ ′ k+1,h ∥ Σ k+1,h .This update enjoys a computational cost of only and Sugiyama, 2023, Lee and Oh, 2024] . As for storage costs, it avoids the need to store all historical data by updating the feature covariance matrix Σ k,h incrementally, requiring only O(d 2 ) storage cost. We note that a dependence on U is introduced in the computational cost. However, as we have discussed in Remark 1, U can be much smaller than the state space size S, which is acceptable. This dependency is typical for model-based methods that directly learn transitions, as it need to control the estimation error of transitions, which typically involves a total of U elements. Similar dependencies have been observed in the literature [Hwang and Oh, 2023] . Additionally, even for the model-free method [Yang and Wang, 2020, Zhou et al., 2021] which learn value functions, a common assumption is the value s ′ ∈S k,h ϕ s ′ k,h V (s ′ ) can be obtained by an Oracle, which depends on U implicitly.O(d 3 U ), independent of episode count k [ZhangThen, we show the estimator θ k,h is close to the true parameter θ * h by the following lemma. Lemma 1. For any δ ∈ (0, 1), define the confidence set asC k,h = θ ∈ Θ | ∥θ -θ k,h ∥ Σ k,h ≤ √ λB + 8 κ d log 1 + kU H dλδ ≜ β k .(5)Then, we have Pr[θ * h ∈ C k,h ] ≥ 1 -δ, ∀k ∈ [K], h ∈ [H]. Remark 3.Notably, a similar confidence set is achieved by Hwang and Oh [2023] by using the MLE defined in (1). In contrast, we obtain the same results by using an online estimator which only suffers O(1) computation and storage cost per round. Besides, we identify a technical issue in their analysis. Specifically, they bound the confidence set using the self-normalized concentration for vector-valued martingales in Lemma 15. However, it is crucial to recognize that the noise is not independent and satisfies s ′ ∈S i,h ϵ s ′ i,h = 0 since the learner visits each stage h exactly once per episode. Thus, the noise ϵ s ′ i,h becomes deterministic and non-zero given the remaining noise ϵ s ′′ i,h for s ′′ ̸ = s ′ . This contravenes the zero-mean sub-Gaussian condition in Lemma 15. We observe similar oversights also appear in the works on multinomial logit contextual bandits [Oh and Iyengar, 2019 , 2021 , Agrawal et al., 2023] . To our knowledge, this issue has not been explicitly identified in previous studies. We note that this issue can be resolved by a new self-normalized concentration with dependent noises in Lemma 1 of Li et al. [2024b] with only slight modifications in constant factors.Efficient Optimistic Value Function Construction. Given the confidence set C k,h , the most direct way is to adopt the principle of "optimism in the face of uncertainty" and construct the optimistic value function as the maximum expected reward over the confidence set. Specifically, we define the optimistic value function Q k,h (s, a) and V k,h (s) asQ k,h (s, a) = r h (s, a) + max θ∈ C k,h s ′ ∈S h,s,a p s ′ s,a (θ) V k,h+1 (s ′ ) [0,H] , V k,h (s) = max a∈A Q k,h (s, a). (6)However, this construction is not efficient as it requires solving a maximization problem over the confidence set. Actually, the algorithm only requires the estimated action-value function to be an upper bound for the true value functions. Thus, we can use a closed-form confidence bound instead of computing the maximal value over the confidence set. To this end, we present the following lemma that enables us to construct the optimistic value function more efficiently. Lemma 2. Suppose Lemma 1 holds. For anyV : S → [0, H] and (h, s, a) ∈ [H] × S × A, it holds s ′ ∈S h,s,a p s ′ s,a ( θ k,h )V (s ′ ) - s ′ ∈S h,s,a p s ′ s,a (θ * h )V (s ′ ) ≤ H β k max s ′ ∈S h,s,a ∥ϕ s ′ s,a ∥ Σ -1 k,h.Based on this lemma, we can replace the maximization problem in ( 6) with closed-form confidence bound and construct the optimistic value function Q k,h (s, a) asQ k,h (s, a) = r h (s, a) + s ′ ∈S h,s,a p s ′ s,a ( θ k,h ) V k,h+1 (s ′ ) + H β k max s ′ ∈S h,s,a ∥ϕ s ′ s,a ∥ Σ -1 k,h [0,H] . (7) At state s k,h , our algorithm chooses a k,h = arg max a∈A Q k,h (s k,h , a). Algorithm 1 presents the detailed procedure. We show that our algorithm achieves the following regret guarantee. Theorem 1. For any δ ∈ (0, 1), set β k as in Lemma 1 and λ = 1, with probability at least 1 -δ, Algorithm 1 ensures the following regret guaranteeReg(K) ≤ O κ -1 dH 2 √ K . Remark 4.Our algorithm attains the same regret of Hwang and Oh [2023] , yet for the more general inhomogeneous MDPs. Importantly, our algorithm only requires constant computational and storage costs, matching the computational efficiency of the linear cases [Jin et al., 2020 , Zhou et al., 2021] . 
statistically improved algorithm by local learning: 
 In this section, we present an enhanced algorithm, UCRL-MNL-LL, that leverages local information to improve the statistical efficiency for MDPs with MNL function approximation.While the UCRL-MNL-OL algorithm in Section 4 offers favorable computational and storage efficiency, it suffers from a dependence on κ -1 in the regret. By Assumption 1, the quantity κ may be exponentially small in the worst case, exhibiting an exponential dependence on the radius of the parameter set and linear in the number of reachable states, i.e., κ -1 = O((U exp(B)) 2 ). This creates a significant gap between MNL and linear function approximation.Recently, the improved dependence on κ has been achieved in the logistic bandit literature [Faury et al., 2020 , Abeille et al., 2021 , Périvier and Goyal, 2022] by the use of generalization of the Bernstein-like tail inequality [Faury et al., 2020] and the self-concordant-like property of the log loss [Bach, 2010] . Thus, a natural question then arises: Can we improve the statistical efficiency of MNL function approximation for MDPs? We answer this question affirmatively by proposing an enhanced algorithm that reduces the dependence on κ significantly through the use of local information. Though the achievement has been made in the logistic bandit, the extension to MDPs with MNL function approximation is non-trivial and new techniques specific to MDPs are required.We first analyze where the dependence on κ comes from. By the analysis of Theorem 1, we can see that the regret of the algorithm can be upper-bounded as follows,EQUATIONHere, the first term corresponds to the overestimation of the value function, and the second term represents the martingale difference sequence arising from the stochastic transition dynamics. The second term can be bounded using the Azuma-Hoeffding inequality, which is independent of κ. For the first term, by the construction of the confidence set in Lemma 1, β k represents the width of the confidence set, contributing to a κ -1/2 dependence. The feature covariance matrix Σ k,h defined in (4) reflects the degree of exploration across different states, which also results in a κ -1/2 dependence.To mitigate these dependencies, we design a local learning algorithm that: (i) constructs a confidence set independent of κ; and (ii) builds a κ-independent feature covariance matrix that effectively captures the exploration degree across different states. We introduce the details in the following.Algorithm 2 UCRL-MNL-LLStep size η, regularization parameter λ, confidence width β k , confidence parameter δ.1: Initialization: H 1,h = λI, θ 1,h = 0 for all h ∈ [H]. 2: for k = 1, . . . , K do 3: Compute Q k,h (•,•) in a backward way as in (11).4:for h = 1, . . . , H do 5: Observe state s k,h , select action a k,h = arg max a∈A Q k,h (s k,h , a). 6: Update H k,h = H k,h + ηH k,h ( θ k,h ). 7: Compute θ k+1,h = arg min θ∈Θ ⟨g k,h ( θ k,h ), θ -θ k,h ⟩ + 1 2η ∥θ -θ k,h ∥ H k,h . 8: Update H k,h = H k,h + H k,h ( θ k+1,h ). 9:end for 10: end forIn Property 1, we show the Hessian matrix H k,h is lower bounded by a positive definite matrixκ s ′ ∈S k,h ϕ s ′ k,h (ϕ s ′ k,h) ⊤ uniformly over the parameter domain, as in (3). This quantity measures the exploration degree across different states and is used to update the parameter estimation and construct the confidence set. However, the bound is not tight, as the Hessian matrix can be significantly larger in certain regions, away from the global minimum. This observation motivates us to design a local learning algorithm that can adaptively leverage local information for improved guarantee.Inspired by recent advances in multinomial logistic bandit (MLogB) [Zhang and Sugiyama, 2023] and multinomial logit contextual bandit (MNL) [Lee and Oh, 2024] , we run an online mirror descent algorithm to estimate the parameter θ * h . Differently from (4), we use the local Hessian matrix H k,h to update the estimation instead of the global lower bound. Specifically, we estimate θ k,h as follows:EQUATION)where η is the step size,H k,h = k-1 i=1 H i,h ( θ i+1,h ) + λI d and H k,h = H k,h + ηH k,h ( θ k,h). Note that both H k,h and H k,h can be updated incrementally. Similar to the update in (4), the optimization problem in (9) can be efficiently solved using the following two-step update:θ ′ k+1,h = θ k,h -η H -1 k,h g k,h (θ k,h ), and θ k+1,h = arg min θ∈Θ ∥θ -θ ′ k+1,h ∥ H k,h .This two-step update procedure incurs a computational cost of O(d 3 U ) and a storage cost of O(d 2 ) per episode, both independent of the episode count k.Based on this estimator, we can construct the κ-independent confidence set as follows.Lemma 3. For any δ ∈ (0, 1), define the confidence set asC k,h = θ ∈ Θ | ∥θ -θ k,h ∥ H k,h ≤ β k ,where Zhang and Sugiyama [2023] studied the multiple-parameter MLogB model, where the unknown parameter is a matrix. Consequently, the confidence set in Theorem 3 of their work exhibits a polynomial dependence on the number of possible outcomes, which corresponds to the number of reachable states U in our setting. This dependence is acceptable in the bandit setting while is not suitable for the MDP setting. In contrast, Lee and Oh [2024] focused on the single-parameter MNL model and revisited the self-concordant-like property, demonstrating that the log-loss of the single-parameter MNL model is 3 √ 2-self-concordant-like (Proposition B.1 in Lee and Oh [2024] ). This property is crucial for the improved confidence set in Lemma 3 that is independent of κ and U .β k = O( √ d log U log(KH/δ)). Then, we have Pr[θ * h ∈ C k,h ] ≥ 1 -δ, ∀k ∈ [K], h ∈ [H]. Remark 5.Based on the confidence set in Lemma 3, a natural choice for the optimistic value function construction is analogous to (7) and can be expressed as:EQUATIONHowever, though β K now is independent of κ and H -1 k,h preserves the local information in (10), the norm maxs ′ ∈S k,h ∥ϕ s ′ k,h ∥ H -1 k,his still in a global manner due to the maximum operation, leading to a κ -1/2 dependence. To address this issue, we propose a new construction of the optimistic value function. Specifically, we employ a second-order Taylor expansion to more accurately bound the value difference arising from transition estimation errors.Although the idea of using second-order Taylor expansion has been explored in bandits [Périvier and Goyal, 2022 , Zhang and Sugiyama, 2023 , Lee and Oh, 2024] , fundamental differences arise in the MDP setting. Specifically, Périvier and Goyal [2022] studied the uniform revenue setting, where the reward is identical for all actions. Zhang and Sugiyama [2023] and Lee and Oh [2024] focused on the non-uniform setting, but the rewards for different actions are known a priori. However, in the MDP setting, the value function is state-dependent and unknown to the learner, leading to a more challenging problem. Moreover, in MNL bandits, the objective is to select a series of assortments with varying sizes that maximize the expected revenue, whereas in MNL MDPs, the goal is to choose one action at each stage to maximize the cumulative reward. Due to these differences, the analyses used in the bandit setting cannot be directly applied to the MDP setting.For MDPs, we show the value difference arising from the transition estimation error as follows.Lemma 4. Suppose Lemma 3 holds. For anyV : S → [0, H] and (h, s, a) ∈ [H] × S × A, it holds s ′ ∈S h,s,a p s ′ s,a ( θ k,h )V (s ′ ) - s ′ ∈S h,s,a p s ′ s,a (θ * h )V (s ′ ) ≤ ϵ fst s,a + ϵ snd s,a .whereϵ fst s,a = H β k s ′ ∈S h,s,a p s ′ s,a ( θ k,h ) ϕ s ′ s,a - s ′′ ∈S h,s,a p s ′′ s,a ( θ k,h )ϕ s ′′ s,a H -1 k,h , ϵ snd s,a = 5 2 H β 2 k max s ′ ∈S h,s,a ∥ϕ s ′ s,a ∥ 2 H -1 k,h.Based on Lemma 4, we construct the optimistic value function as follows:EQUATIONIn contrast to the value function specified in (10), where the term maxs ′ ∈S h,s,a ∥ϕ s ′ s,a ∥ H -1 k,his utilized, the refined value function introduced in (11) substitutes this term with ϵ fst s,a + ϵ snd s,a . This adjustment better preserves local information, offering a more precise and κ-independent estimation error bound.At state s k,h , our algorithm chooses action a k,h = max Q k,h (s k,h , •). The detailed algorithm is presented in Algorithm 2. We show the regret guarantee of UCRL-MNL-LL in the following theorem.Theorem 2. For any δ ∈ (0, 1), setβ k = O( √ d log U log(KH/δ)), η = 1 2 log(1 + U ) + (B + 1) and λ = 84 √ 2η(B + d), with probability at least 1 -δ, UCRL-MNL-LL algorithm ensures the following regret guaranteeReg(K) ≤ O dH 2 √ K + κ -1 d 2 H 2 .Remark 6. The high-order term in Theorem 2 is now independent κ, significantly improving the statistical efficiency compared with Theorem 1. In comparison with the optimal regret of O(d √ H 3 K) for linear cases [Zhou et al., 2021 , He et al., 2023] , the higher-order term in Theorem 2 only differs by a factor of H 1/2 , almost matching the same computational and statistical efficiency simultaneously. 
lower bound: 
 In this section, we establish a lower bound for MNL mixture MDPs by presenting a novel reduction, which connects MNL mixture MDPs and the logistic bandit problem.Consider the following logistic bandit problem [Faury et al., 2020 , Abeille et al., 2021] : at each round k ∈ [K], the learner selects an action x k ∈ X and receives a reward r k sampled from Bernoulli distribution with mean µ(x ⊤ θ * ) = (1 + exp(-x ⊤ θ * )) -1 , where θ * ∈ {θ ∈ R d , ∥θ∥ 2 ≤ B} is the unknown parameter. The learner aims to to minimize the regret:Reg LogB (K) = max x∈X K k=1 µ(x ⊤ θ * ) - K k=1 µ(x ⊤ k θ * ).Theorem 3 (Lower Bound). For any logistic bandit problem B, there exists an MNL mixture MDP M such that learning M is as hard as learning H/2 independent instances of B simultaneously.Corollary 1. For any problem instance {θ * h } H h=1 and for K ≥ d 2 κ * , there exists an MNL mixture MDP with infinite action space such that Reg(K) ≥ Ω(dH√ Kκ * ).Remark 7. Corollary 1 can be proved by combining Theorem 3 and the Ω(d √ Kκ * ) lower bound for logistic bandits with infinite arms by Abeille et al. [2021] . To the best of our knowledge, a lower bound for logistic bandits with finite arms has not been established, which is beyond the scope of this work. This absence leaves the lower bound for MNL mixture MDPs with a finite action space open through this reduction. However, after the submission of our work to arXiv [Li et al., 2024a] , a follow up work by Park et al. [2024] proposed a new reduction that bridges MNL mixture MDPs with linear mixture MDPs by approximating multinomial logit functions to linear functions, employing the mean value theorem. Leveraging this new reduction, they established an Ω(dH 3/2 √ K) lower bound for MNL mixture MDPs with the finite action space. This achievement confirms that our result is indeed optimal with respect to the dependence on d and K, only differing by a O(H 1/2 ) factor.Dependence on H. By the discussion in Remark 7, we note that our result is optimal with respect to d and K, but loosing by a O(H 1/2 ) factor. We discuss the challenges in improving the dependence on H. Notably, MNL mixture MDPs can be viewed as a generalization of linear mixture MDPs [Ayoub et al., 2020 , Zhou et al., 2021] . The pioneering work by Ayoub et al. [2020] achieved a regret bound of O(dH 2 √ K) for linear mixture MDPs, which matches our results in Theorem 2, differing only on the lower order term. Later, Zhou et al. [2021] enhanced the dependence on H and attained an optimal regret bound of O(d √ H 3 K). This was made possible by recognizing that the value function in linear mixture MDPs is linear, allowing for direct learning of the value function while incorporating variance information. In contrast, the value function for MNL mixture MDPs does not conform to a specific structure, posing a significant challenge in using the variance information of value functions. Thus, it remains open whether similar improvements on H are attainable for MNL mixture MDPs. 
conclusion and future work: 
 In this work, we study MNL mixture MDPs that employ multinomial logit function approximation to ensure valid probability distributions over the state space. We address both the computational and statistical challenges for this problem. Specifically, we first propose an algorithm based on the online Newton step that attains the O(κ -1 dH 2 √ K) regret with O(1) computational and storage costs per episode. Next, we propose an enhanced algorithm that leverages local information to improve the statistical efficiency. It not only maintains O(1) computational and storage costs but also achieves an improved regret of O(dH 2 √ K + d 2 H 2 κ -1 ), nearly matching the result of linear function approximation from both computational and statistical perspectives. Finally, we establish the first lower bound for MNL mixture MDPs, justifying the optimality of our results in d and K.There are several interesting directions for future work. First, there still exists a gap between the upper and lower bounds and how to close this gap is an open problem. Besides, while this work focuses on stochastic rewards, extending this model to the non-stationary settings and studying the dynamic regret [Wei and Luo, 2021 , Zhao et al., 2022 , Li et al., 2023 ] is an another important direction. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: This is a theoretical paper and does not release any data or models that have a high risk for misuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [NA] Justification: This is a theoretical paper and does not use existing assets.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [NA] Justification: This is a theoretical paper and does not release new assets. Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: This section collects several key properties of the multinomial logit function used in the paper.Without loss of generality, we assume ∀S h,s,a , ∃s ′ ∈ S h,s,a such that ϕ(s ′ | s, a) = 0. Otherwise, we can always define a new feature mapping ϕ ′ (s ′′ | s, a) = ϕ(s ′ | s, a) -ϕ(s ′′ | s, a) for any s ′′ ∈ S h,s,a such that ϕ ′ (s ′ | s, a) = 0 and the transition induced by ϕ ′ is the same as that induced by ϕ. We denote this state as ṡh,s,a and Ṡh,s,a = S h,s,a \{ ṡh,s,a }.Recall the following definitions in the paper:f k,h (θ) = - s ′ ∈S k,h y s ′ k,h log p s ′ k,h (θ), g k,h (θ) = s ′ ∈S k,h (p s ′ k,h (θ) -y s ′ k,h )ϕ s ′ k,h H k,h (θ) = s ′ ∈S k,h p s ′ k,h (θ)ϕ s ′ k,h (ϕ s ′ k,h ) ⊤ - s ′ ∈S k,h s ′′ ∈S k,h p s ′ k,h (θ)p s ′′ k,h (θ)ϕ s ′ k,h (ϕ s ′′ k,h ) ⊤ . Σ k,h = κ 2 k-1 i=1 s ′ ∈S i,h ϕ s ′ i,h (ϕ s ′ i,h ) ⊤ + λI d , H k,h = k-1 i=1 H i,h ( θ i+1,h ) + λI d ,Lemma 5. The following statements hold for any k ∈[K], h ∈ [H]: H k,h (θ) ⪰ s ′ ∈ Ṡk,h p s ′ k,h (θ)p ṡk,h k,h (θ)ϕ s ′ k,h (ϕ s ′ k,h ) ⊤ ⪰ κ s ′ ∈ Ṡk,h ϕ s ′ k,h (ϕ s ′ k,h ) ⊤ . Proof. First, note that ∀x, y ∈ R d , (x -y)(x -y) ⊤ = xx ⊤ + yy ⊤ -xy ⊤ -yx ⊤ ⪰ 0 =⇒ xx ⊤ + yy ⊤ ⪰ xy ⊤ + yx ⊤ .Then, we haveH k,h (θ) = s ′ ∈S k,h p s ′ k,h (θ)ϕ s ′ k,h (ϕ s ′ k,h ) ⊤ - s ′ ∈S k,h s ′′ ∈S k,h p s ′ k,h (θ)p s ′′ k,h (θ)ϕ s ′ k,h (ϕ s ′′ k,h ) ⊤ = s ′ ∈ Ṡk,h p s ′ k,h (θ)ϕ s ′ k,h (ϕ s ′ k,h ) ⊤ - s ′ ∈ Ṡk,h s ′′ ∈ Ṡk,h p s ′ k,h (θ)p s ′′ k,h (θ)ϕ s ′ k,h (ϕ s ′′ k,h ) ⊤ = s ′ ∈ Ṡk,h p s ′ k,h (θ)ϕ s ′ k,h (ϕ s ′ k,h ) ⊤ - 1 2 s ′ ∈ Ṡk,h s ′′ ∈ Ṡk,h p s ′ k,h (θ)p s ′′ k,h (θ) ϕ s ′ k,h (ϕ s ′′ k,h ) ⊤ + ϕ s ′′ k,h (ϕ s ′ k,h ) ⊤ ⪰ s ′ ∈ Ṡk,h p s ′ k,h (θ)ϕ s ′ k,h (ϕ s ′ k,h ) ⊤ - 1 2 s ′ ∈ Ṡk,h s ′′ ∈ Ṡk,h p s ′ k,h (θ)p s ′′ k,h (θ) ϕ s ′ k,h (ϕ s ′ k,h ) ⊤ + ϕ s ′′ k,h (ϕ s ′′ k,h ) ⊤ = s ′ ∈ Ṡk,h p s ′ k,h (θ)ϕ s ′ k,h (ϕ s ′ k,h ) ⊤ - s ′ ∈ Ṡk,h s ′′ ∈ Ṡk,h p s ′ k,h (θ)p s ′′ k,h (θ)ϕ s ′ k,h (ϕ s ′ k,h ) ⊤ = s ′ ∈ Ṡk,h p s ′ k,h (θ) 1 - s ′′ ∈ Ṡk,h p s ′′ k,h (θ) ϕ s ′ k,h (ϕ s ′ k,h ) ⊤ = s ′ ∈ Ṡk,h p s ′ k,h (θ)p ṡk,h k,h (θ)ϕ s ′ k,h (ϕ s ′ k,h ) ⊤ ⪰ κ s ′ ∈ Ṡk,h ϕ s ′ k,h (ϕ s ′ k,h ) ⊤ ,where the last inequality holds by the definition of κ in Assumption 1. This finishes the proof.■ Lemma 6. Suppose λ ≥ 1, define ϕ s ′ k,h = ϕ s ′ k,h -s ′′ ∈S k,h p s ′′ k,h ( θ k+1,h )ϕ s ′′ k,h , for any k ∈ [K], h ∈ [H],the following statements hold:(I) k i=1 max s ′ ∈S i,h ∥ϕ s ′ i,h ∥ 2 Σ -1 i,h ≤ 4 κ d log 1 + kU λd (II) k i=1 s ′ ∈S i,h p s ′ i,h ( θ i+1,h )p ṡk,h i,h ( θ i+1,h )∥ϕ s ′ i,h ∥ 2 H -1 i,h ≤ 2d log 1 + k λd (III) k i=1 max s ′ ∈S i,h ∥ϕ s ′ i,h ∥ 2 H -1 i,h ≤ 2 κ d log 1 + k λd (IV) k i=1 s ′ ∈S i,h p s ′ i,h ( θ i+1,h )∥ ϕ s ′ i,h ∥ 2 H -1 i,h ≤ 2d log 1 + k λd (V) k i=1 max s ′ ∈S i,h ∥ ϕ s ′ i,h ∥ 2 H -1 i,h ≤ 2 κ d log 1 + k λdProof. We prove the five statements individually.k+1,h = Σ k,h + κ 2 s ′ ∈S k,h ϕ s ′ k,h (ϕ s ′ k,h ) ⊤ , we have det( Σ k+1,h ) = det( Σ k,h ) 1 + κ 2 s ′ ∈S k,h ∥ϕ s ′ k,h ∥ 2 Σ -1 k,h.Then, we getEQUATIONwhere the last inequality holds by determinant trace inequality in Lemma 16. Since λ ≥ 1, we havemax s ′ ∈S i,h ∥ϕ s ′ i,h ∥ Σ -1 i,h≤ 1, thus, we havek i=1 max s ′ ∈S i,h ∥ϕ s ′ i,h ∥ Σ -1 i,h ≤ 2 κ k i=1 min 1, κ 2 s ′ ∈S i,h ∥ϕ s ′ i,h ∥ Σ -1 i,h ≤ 4 κ k i=1 log 1 + κ 2 s ′ ∈S i,h ∥ϕ s ′ i,h ∥ 2 Σ -1 i,h ≤ 4 κ d log 1 + kU λd ,where the first inequality holds by the fact that z ≤ 2 log(1 + z) for any z ∈ [0, 1] and the last inequality holds by ( 12).Proof of statement (II). By Lemma 5, we haveH k,h (θ) ⪰ s ′ ∈ Ṡk,h p s ′ k,h (θ)p ṡk,h k,h (θ)ϕ s ′ k,h (ϕ s ′ k,h ) ⊤ . Thus, we have H k+1,h ⪰ H k,h + s ′ ∈ Ṡk,h p s ′ k,h ( θ k+1,h )p ṡk,h k,h ( θ k+1,h )ϕ s ′ k,h (ϕ s ′ k,h ) ⊤Then, we getdet(H i+1,h ) ≥ det(H i,h ) 1 + s ′ ∈ Ṡi,h p s ′ i,h ( θ i+1,h )p ṡi,h i,h ( θ i+1,h )∥ϕ s ′ i,h ∥ 2 H -1 i,h . Since λ ≥ 1, we have s ′ ∈ Ṡi,h p s ′ i,h ( θ i+1,h )p ṡi,h i,h ( θ i+1,h )∥ϕ s ′ i,h ∥ 2 H -1 i,h ≤ 1. Using the fact that z ≤ 2 log(1 + z) for any z ∈ [0, 1], we get k i=1 s ′ ∈ Ṡi,h p s ′ i,h ( θ i+1,h )p ṡi,h i,h ( θ i+1,h )∥ϕ s ′ i,h ∥ 2 H -1 i,h ≤ 2 k i=1 log 1 + s ′ ∈ Ṡi,h p s ′ i,h ( θ i+1,h )p ṡi,h i,h ( θ i+1,h )∥ϕ s ′ i,h ∥ 2 H -1 i,h ≤ 2 log det(H k+1,h ) det(H 1,h ) ≤ 2d log 1 + k λd .Proof of statement (III). By Lemma 5, we haveH k+1,h ⪰ H k+1,h + κ s ′ ∈ Ṡk,h ϕ s ′ k,h (ϕ s ′ k,h ) ⊤ Since λ ≥ 1, we have κ max s ′ ∈S i,h ∥ϕ s ′ i,h ∥ H -1 i,h ≤ κ. Using the fact that z ≤ 2 log(1 + z) for any z ∈ [0, 1]. By a similar analysis as the statement (2), we havek i=1 max s ′ ∈S i,h ∥ϕ s ′ i,h ∥ 2 H -1 i,h ≤ 2 κ k i=1 log 1 + max s ′ ∈S i,h ∥ϕ s ′ i,h ∥ H -1 i,h ≤ 2 κ log det(H k+1,h ) det(H 1,h ) ≤ 2 κ d log 1 + k λd .This finishes the proof.. By the definition of H k,h (θ), we haveEQUATIONThen, we getH k+1,h ⪰ H k,h + s ′ ∈S k,h p k,h ( θ k+1,h )( ϕ s ′ k,h )( ϕ s ′ k,h ) ⊤The remaining proof is similar to the proof of statement (2).Proof of statement (V). By (13), we haveH i,h ( θ i+1,h ) ⪰ s ′ ∈S i,h p i,h ( θ i+1,h )( ϕ s ′ i,h )( ϕ s ′ i,h ) ⊤ ⪰ κ s ′ ∈S i,h ( ϕ s ′ i,h )( ϕ s ′ i,h ) ⊤ .Then, the remaining proof is similar to the proof of statement (3). ■Proof. The main proof follows the proof of Theorem 1 in Zhang et al. [2016] . We definefk,h (θ) = E y k,h [f k,h (θ) | F k,h ], ḡk,h (θ) = E y k,h [g k,h (θ) | F k,h ].By Property 1, we have f k,h (θ) is exponential concave such thatf k,h ( θ k,h ) ≤ f k,h (θ * h ) + ⟨g k,h (θ * h ), θ k,h -θ * h ⟩ - κ 2 ∥( θ k,h -θ * h )∥ W k,h .Taking expectations on both sides, we haveEQUATIONBased the property of KL diverge, we ensure the true parameter θ * h is the minimizer of the expected loss function by the following lemma.Lemma 7. For any θ ∈ Θ, we have fk,h (θ * h ) ≤ fk,h (θ).Combining Lemma 7 and Equation ( 14), we haveEQUATIONBy standard analysis of OMD in Lemma 17, it holdsEQUATIONCombining Equation ( 15) and Equation ( 16), we haveEQUATIONFirst, we consider the first term . We show the gradient can be bounded by the following lemma.Lemma 8. For any positive semi-definite matrix Z, it holds thatEQUATION)Then, we consider the second term. Note thatΣ k+1,h = Σ k,h + κ 2 W k,h , we have ∥ θ k,h -θ * ∥ 2 Σ k+1,h = ∥ θ k,h -θ * ∥ 2 Σ k,h + κ 2 ∥ θ k,h -θ * ∥ 2 W k,h . (19) Next, we bound ⟨ḡ k,h ( θ k,h ) -g k,h ( θ k,h ), θ k,h -θ * h ⟩, which is a martingale difference sequence. Lemma 9. For any δ ∈ (0, 1) and θ 1 , . . . , θ k , θ * ∈ [0, B] d , with probability at least 1 -δ, for anyEQUATION)whereγ k = log 2k 2 H log(kU ) δ .Combining ( 17), ( 18), and ( 19), we haveEQUATION)where the first inequality holds by rearranging the terms in ( 17), the second holds by ( 18), ( 19) and ( 20). Summing ( 21) from i = 1 to k, we have∥ θ k+1,h -θ * h ∥ 2 Σ k+1,h ≤ λB 2 + 4 k i=1 max s ′ ∈S i,h ∥ϕ s ′ i,h ∥ 2 Σ -1 i+1,h + ( 8 κ + 16B)γ k ≤ λB 2 + 8 κ d log 1 + kU λd + ( 8 κ + 16B)γ k ,where the second inequality holds by ∥ϕs ′ i,h ∥ 2 ≤ 1 and Σ i,h ≥ I, ∀i ∈ [K],and the last inequality is by Lemma 6. This finishes the proof. ■In this section, we provide the proofs of the lemmas used in Appendix B.1.Proof. By the definition of fk,h (θ), we havefk,h (θ) -fk,h (θ * h ) = s ′ ∈S k,h p s ′ k,h (θ * h ) log p s ′ k,h (θ * h ) p s ′ k,h (θ) ≥ 0,where the last inequality is due tos ′ ∈S k,h p s ′ k,h (θ * h ) log p s ′ k,h (θ * h ) p s ′ k,h (θ)is the Kullback-Leibler divergence between p k,h (θ * h ) and p k,h (θ), which always is non-negative. ■Proof. For any positive semi-definite matrix Z,EQUATION) which implies x ⊤ i Zx i + x ⊤ j Zx j ≥ x ⊤ i Zx j + x ⊤ j Zx i , ∀x i , x j ∈ R d . Let x i = (p s ′ k,h (θ) -y s ′ k,h )ϕ s ′ k,h , we have ∥g k,h (θ)∥ 2 Z = s ′ ∈S k,h s ′′ ∈S k,h (p s ′ k,h (θ) -y s ′ k,h )(p s ′′ k,h (θ) -y s ′′ k,h )ϕ s ′ k,h Zϕ s ′′ k,h = s ′ ∈S k,h (p s ′ k,h (θ) -y s ′ k,h ) 2 ϕ s ′ k,h Zϕ s ′ k,h + 1 2 s ′ ∈S k,h s ′′ ∈S k,h (p s ′ k,h (θ) -y s ′ k,h )(p s ′′ k,h (θ) -y s ′′ k,h )(ϕ s ′ k,h Zϕ s ′′ k,h + ϕ s ′′ k,h Zϕ s ′ k,h ) ≤ s ′ ∈S k,h (p s ′ k,h (θ) -y s ′ k,h ) 2 ϕ s ′ k,h Zϕ s ′ k,h + s ′ ∈S k,h (p s ′ k,h (θ) -y s ′ k,h ) 2 ϕ s ′ k,h Zϕ s ′ k,h = 2 s ′ ∈S k,h (p s ′ k,h (θ) -y s ′ k,h ) 2 ϕ s ′ k,h Zϕ s ′ k,h ≤ 4 max s ′ ∈S k,h ∥ϕ s ′ k,h ∥ 2 Z .This finishes the proof. ■Proof. First, notice that(ḡ i,h (θ i ) -g i,h (θ i )) ⊤ (θ i -θ *) is a martingale difference sequence. Also, we have(ḡ i,h (θ i ) -g i,h (θ i )) ⊤ (θ i -θ * ) ≤ (ḡ i,h (θ i )) ⊤ (θ i -θ * ) + (g i,h (θ i )) ⊤ (θ i -θ * ) ≤ ∥ḡ i,h (θ i )∥ 2 ∥θ i -θ * ∥ 2 + ∥g i,h (θ i )∥ 2 ∥θ i -θ * ∥ 2 ≤ 4 √ 2B,where the last inequality holds by ∥ḡ i,h(θ i )∥ 2 = ∥ s ′ ∈S i,h (p s ′ i,h (θ) -y s ′ i,h )ϕ s ′ i,h ∥ 2 ≤ √ 2.We define the martingaleM k,h = k i=1 (ḡ i,h (θ i ) -g i,h (θ i )) ⊤ (θ i -θ * ), and define the conditional variance σ 2 i asσ 2 k,h = E[M 2 k,h | F i-1 ] = k i=1 E y k,h (ḡ i,h (θ i ) -g i,h (θ i )) ⊤ (θ i -θ * ) 2 ≤ k i=1 E y k,h (g i,h (θ i )) ⊤ (θ i -θ * ) 2 ≤ k i=1 s ′ ∈S i,h (ϕ s ′ i,h ) ⊤ (θ i -θ * ) 2 = k i=1 ∥θ i -θ * ∥ 2 W i,h ≜ A k,h .where the first inequality is due to the fact thatE[(ξ -E[ξ]) 2 ] ≤ E[ξ 2] for any random variable ξ.Note that A k,h is a random variable, so we cannot directly apply the Bernstein inequality to M k,h . Instead, we consider the following two cases:(i) A k,h ≤ 4B 2 kU and (ii) A k,h > 4B 2 kU . Case (i): A k,h = k i=1 ∥θ i -θ * ∥ 2 W i,h ≤ 4B 2 kU . Then, we have M k,h = k i=1 (ḡ i,h (θ i ) -g i,h (θ i )) ⊤ (θ i -θ * ) = k i=1 s ′ ∈S i,h (p s ′ i,h (θ i ) -y s ′ i,h )(ϕ s ′ i,h ) ⊤ (θ i -θ * ) ≤ k i=1 s ′ ∈S i,h |(ϕ s ′ i,h ) ⊤ (θ i -θ * )| ≤ kU k i=1 s ′ ∈S i,h ((ϕ s ′ i,h ) ⊤ (θ i -θ * )) 2 ≤ 2B.where the second equality is due to the definition of ḡi,h (θ i ) and g i,h (θ i ), the first inequality holds byp s ′ i,h (θ i ) -y s ′ i,h ∈ [-1, 1],the second inequality holds by the Cauchy-Schwarz inequality, and the last inequality holds by the condition of case (i).k,h = k i=1 ∥θ i -θ * ∥ 2 W j,h > 4B 2 kU .We have both a lower and upper bound for A k,h , i.e., 4B 2 kU < A k,h ≤ 4B 2 kU . Then, we can use the peeling process to bound M k,h as follows:Pr M k,h ≥ 2 τ k,h A k,h + 8Bτ k,h 3 = Pr M k,h ≥ 2 τ k,h A k,h + 8Bτ k,h 3 , 4B 2 kU < A k,h ≤ 4kU B 2 = Pr M k,h ≥ 2 τ k,h A k,h + 8Bτ k,h 3 , 4B 2 kU < A k,h ≤ 4kU B 2 , σ k,h ≤ A k,h ≤ m i=1 Pr M k,h ≥ 2 τ k,h A k,h + 8Bτ k,h 3 , 4B 2 2 i-1 kU < A k,h ≤ 4B 2 2 i kU , σ k,h ≤ A k,h ≤ m i=1 Pr M k,h ≥ 8B 2 2 i 3kU τ k,h + 8Bτ k,h 3 , σ k,h ≤ 4B 2 2 i kU ≤ m exp(-τ k,h ).where m = 2 log 2 (kU ), and the last inequality follows the Bernstein inequality for martingales.Combining above two cases, letting τ = log mk 2 δ/H and taking the union bound over k and h ∈ [H], we have with probability at least 1 -δ, for any k ∈EQUATIONwhereγ k = log 2k 2 H log(kU ) δ .Then, applying uv ≤ cu 2 + v 2 /(4c) for any c, u, v > 0 with c = 2/κ, we haveEQUATIONCombining ( 23) and ( 24), we havek i=1 (ḡ i,h (θ i ) -g i,h (θ i )) ⊤ (θ i -θ * ) ≤ κ 4 k i=1 ∥θ i -θ * ∥ 2 W i,h + ( 4 κ + 8B)γ k .This finishes the proof. ■Proof. The gradient of p s,a (θ) is given by∇p s ′ s,a (θ) = p s ′ s,a (θ)ϕ s ′ s,a -p s ′ s,a (θ) s ′′ ∈S h,s,a p s ′′ s,a (θ)ϕ s ′′ s,a .By the mean value theorem, there exists θ = νθEQUATION≤ H max s ′ ∈S h,s,a ∥ϕ s ′ s,a ∥ Σ -1 k,h ∥ θ k,h -θ * h ∥ Σ k,h ≤ Hβ k max s ′ ∈S h,s,a ∥ϕ s ′ s,a ∥ Σ -1 k,h, where the first inequality is by V : S → [0, H], the second inequality is by the definition of ϕ s ′ k,h , the third inequality is by the Holder's inequality, the last inequality is by Lemma 1. ■To prove the theorem, we first introduce the following lemma. Lemma 10. Suppose for all (h, s, a) ∈ [H] × S × A, it holds thatEQUATIONDefineEQUATIONThen, for any δ ∈ (0, 1], with probability at least 1 -δ, it holds thatReg(K) ≤ 2 K k=1 H h=1 Γ h,s k,h ,a k,h + H 2KH log(2/δ). Proof of Theorem 1. Substituting Γ h,s,a = H β k max s ′ ∈S h,s,a ∥ϕ s ′ s,a ∥ Σ -1 k,hinto Lemma 10, we haveReg(K) ≤ 2H β k K k=1 H h=1 max s ′ ∈S k,h ∥ϕ s ′ k,h ∥ Σ -1 k,h + H 2KH log(2/δ) ≤ 2H β k H h=1 K K k=1 max s ′ ∈S k,h ∥ϕ s ′ k,h ∥ 2 Σ -1 k,h + H 2KH log(2/δ) ≤ 2H 2 β k 4dK κ log 1 + KU dλ + H 2KH log(2/δ) ≤ O(κ -1 dH 2 √ K),where the second inequality holds by the Cauchy-Schwarz inequality and the third inequality holds by Lemma 6. This finishes the proof. ■In this section, we provide the proofs of the lemmas used in Appendix D.1.First, we introduce the following lemma. Lemma 11. Suppose (26) and (27) in Lemma 10 holds for all k ∈[K], h ∈ [H]. Then, for any (s, a, h) ∈ S × A × [H], it holds that Q * h (s, a) ≤ Q k,h (s, a) ≤ r h (s, a) + P h V k,h+1 (s, a) + 2Γ h,s,a .Proof. First, we prove the left-hand side of the lemma. We prove this by backward induction on h.For the stage h = H, by definition, we haveQ k,H (s, a) = r H (s, a) = Q * H (s, a), V k,H+1 (s) = 0 = V * H+1 (s). Suppose the statement holds for h + 1, we show it holds for h. By definition, if Q k,h (s, a) = H, this holds trivially. Otherwise, we haveQ k,h (s, a) = r h (s, a) + p s,a ( θ ) V k,h+1 + Γ h,s,a ≥ r h (s, a) + p s,a ( θ k,h )V * k + Γ h,s,a ≥ r h (s, a) + p s,a (θ * h )V * k = Q * h (s, a).where the first inequality is by the induction hypothesis, and the second inequality is by ( 26).Then, we prove the right-hand side of the lemma. By the definition of Q k,h (s, a), we haveQ k,h (s, a) = r h (s, a) + p s,a ( θ k,h ) V k,h+1 + Γ h,s,a ≤ r h (s, a) + p s,a (θ * h ) V k,h+1 + 2Γ h,s,a, where the inequality is by (26). ■Proof. By the definition ofReg(K) = K k=1 V * 1 (s k,1 ) - K k=1 V π k 1 (s k,1 ), we have K k=1 V * 1 (s k,1 ) - K k=1 V π k 1 (s k,1 ) = K k=1 Q * 1 (s k,1 , π * (s k,1 )) - K k=1 V π k 1 (s k,1 ) ≤ K k=1 Q 1 (s k,1 , π * (s k,1 )) - K k=1 V π k 1 (s k,1 ) ≤ K k=1 Q 1 (s k,1 , a k,1 ) - K k=1 V π k 1 (s k,1 ),where the first inequality is by Lemma 11 and the second is by a k,h = arg max a∈A Q k,h (s k,h , a).By the right-hand side of Lemma 11, we haveQ 1 (s k,1 , a k,1 ) -V π k 1 (s k,1 ) = r(s k,1 , a k,1 ) + P 1 V k,2 (s k,1 , a k,1 ) + +2Γ h,s k,1 ,a k,1 -r(s k,1 , a k,1 ) -P 1 V π k 2 (s k,1 , a k,1 ) ≤ P 1 ( V k,2 -V π k 2 )(s k,1 , a k,1 ) -( V k,2 -V π k 2 )(s k,2 ) + ( V k,2 -V π k 2 )(s k,2 ) + 2Γ h,s k,1 ,a k,1 ≤ P 1 ( V k,2 -V π k 2 )(s k,1 , a k,1 ) -( V k,2 -V π k 2 )(s k,2 ) + Q 2 (s k,2 , a k,2 ) -V π k 2 (s k,2 ) + 2Γ h,s k,1 ,a k,1 . Define M k,h = P h ( V k,h+1 -V π k h+1 )(s k,h , a k,h ) -( V k,h+1 -V π k h+1 )(s k,h+1). Applying this recursively, we haveQ 1 (s k,1 , a k,1 ) -V π k 1 (s k,1 ) ≤ 2 H h=1 Γ h,s k,h ,a k,h + H h=1 M k,hSumming over k, we have for any δ ∈ (0, 1], with probability at least 1 -δ, it holds thatReg(K) ≤ 2 K k=1 H h=1 Γ h,s k,h ,a k,h + K k=1 H h=1 M k,h ≤ 2 K k=1 H h=1 Γ h,s k,h ,a k,h + H 2KH log(2/δ)where the inequality holds by the Azuma-Hoeffding inequality as M k,h is a martingale difference sequence with M k,h ≤ 2H. This finishes the proof. ■Proof. The proof is similar to the proof of Theorem 3 in Zhang and Sugiyama [2023] and Lemma 1 of Lee and Oh [2024] . Definef k,h (θ) = f k,h (θ k,h ) + ⟨g k,h ( θ k,h ), θ -θ k,h ⟩ + 1 2 ∥θ -θ k,h ∥ H ( θ k,h ) ,which is a second order approximation of the original function f k,h (θ) at θ k,h . Then, the update rule in ( 9) can be rewritten asEQUATIONThen, we present the following lemma that bounds the estimation error of the update rule in (28).Lemma 12 (Lemma E.1 of Lee and Oh [2024] ). For the update rule defined in (28), it holds that∥ θ k+1,h -θ * h ∥ 2 H k+1,h ≤2η k i=1 f i,h (θ * h ) - k i=1 f i,h ( θ i+1,h ) + 4λB + 12 √ 2Bη k i=1 ∥ θ i+1,h -θ i,h ∥ 2 2 - k i=1 ∥ θ i+1,h -θ i,h ∥ 2 H i,h .We bound the right-hand side of the above lemma separately in the following.First, we bound the first term. Define the softmax function as follows.[σ k,h (z)] s = exp([z]) s 1 + s∈ Ṡk,h exp([z]) s , ∀s ∈ S k,h .Then, the loss function in (2) can be rewritten asf k,h (z k,h , y k,h ) = s ′ ∈S k,h 1[y s ′ k,h = 1] log 1 [σ k,h (z k,h )] s ′ .Define a pseudo-inverse function of σ k,h (•) as[σ -1 k,h (p)] s ′ = log [p] s ′ 1 -∥p∥ 1 , ∀p ∈ {p ∈ [0, 1] N k,h | ∥p∥ 1 < 1}.Then, we decompose the first term as follows.EQUATIONwherez k,h = σ -1 k,h (E θ∼P k,h [σ k,h ((ϕ s ′ k,h ) ⊤ θ) s ′ ∈ Ṡk,h ]), P k,h ≜ N ( θ k,h , (1 + cH -1 k,h )) is the Gaussian distribution with mean θ k,h and covariance (1 + cH -1 k,h) where c is a constant to be specified later. First, we bound the term (a) as follows.Lemma 13 (Lemma E.2 of Lee and Oh [2024] ). Let δ ∈ (0, 1] and λ ≥ 1, for all k ∈ [K], h ∈ [H], with probability at least 1 -δ, we havek i=1 f i,h (θ * h ) - k i=1 f i,h (z i,h , y i,h ) ≤ (3 log(1 + (U + 1)k) + 3) 17 16 λ + 2 √ λ log 2H √ 1 + 2k δ + 16 log 2H √ 1 + 2k δ 2 + 2.Then, we bound the term (b) as follows.Lemma 14 (Lemma E.3 of Lee and Oh [2024] ). For any c ≥ 0, let λ ≥ max{2, 72cd}, then for allEQUATIONH i,h + √ 6cd log 1 + k + 1 2λ .Now, we are ready to prove Lemma 3. Combining Lemma 12, Lemma 13, and Lemma 14, we have∥ θ k+1,h -θ * h ∥ 2 H k+1,h ≤ 2η (3 log(1 + (U + 1)k) + 3) 17 16 λ + 2 √ λ log 2H √ 1 + 2k δ + 16 log 2H √ 1 + 2k δ 2 +2 + √ 6cd log 1 + k + 1 2λ + 4λB + 12 √ 2Bη k i=1 ∥ θ i+1,h -θ i,h ∥ 2 2 + ( η c -1) k i=1 ∥ θ i+1,h + θ i,h ∥ 2 H i,h ≤ 2η (3 log(1 + (U + 1)k) + 3) 17 16 λ + 2 √ λ log 2H √ 1 + 2k δ + 16 log 2H √ 1 + 2k δ 2 + 2 + √ 6cd log 1 + k + 1 2λ + 4λB,where the second inequality holds by setting c = 7η/6 and λ ≥ max{84 √ 2ηB, 84dη}, we have12 √ 2Bη k i=1 ∥ θ i+1,h -θ i,h ∥ 2 2 + η c -1 k i=1 ∥ θ i+1,h + θ i,h ∥ 2 H i,h = 12 √ 2Bη k i=1 ∥ θ i+1,h -θ i,h ∥ 2 2 - 1 7 k i=1 ∥ θ i+1,h + θ i,h ∥ 2 H i,h ≤ 12 √ 2Bη - λ 7 k i=1 ∥ θ i+1,h -θ i,h ∥ 2 2 ≤ 0.Thus, by setting η = 1 2 log(U + 1) + (B + 1), λ = 84√ 2η(B + d), we have ∥ θ k+1,h -θ * h ∥ H k+1,h ≤ O √ d log U log(KH/δ) ≜ β k .This finishes the proof. ■Proof. By the second-order Taylor expansion, there existsθ = νθ * h + (1 -ν) θ k,h for some ν ∈ [0, 1], such that s ′ ∈S h,s,a p s ′ s,a ( θ k,h )V (s ′ ) - s ′ ∈S h,s,a p s ′ s,a (θ * h )V (s ′ ) = s ′ ∈S h,s,a ∇p s ′ s,a ( θ k,h ) ⊤ (θ * h -θ k,h )V (s ′ ) + 1 2 s ′ ∈S h,s,a ( θ k,h -θ * h ) ⊤ ∇ 2 p s ′ s,a ( θ)( θ k,h -θ * h )V (s ′ )The gradient of p s,a (θ) is given byEQUATIONs ′′ ∈S h,s,a p s ′′ s,a (θ)ϕ s ′′ s,a .For the first-order term, we haves ′ ∈S h,s,a ∇p s ′ s,a ( θ k,h ) ⊤ (θ * h -θ k,h )V (s ′ ) = s ′ ∈S h,s,a p s ′ s,a ( θ k,h )ϕ s ′ s,a (θ * h -θ k,h )V (s ′ ) - s ′ ∈S h,s,a p s ′ s,a ( θ k,h ) s ′′ ∈S h,s,a p s ′′ s,a ( θ k,h )ϕ s ′′ s,a (θ * h θ k,h )V (s ′ ) ≤ H s ′ ∈S + h,s,a p s ′ s,a ( θ k,h )   ϕ s ′ s,a (θ * h -θ k,h ) - s ′′ ∈S h,s,a p s ′′ s,a ( θ k,h )ϕ s ′′ s,a (θ * h -θ k,h )   ≤ H s ′ ∈S + h,s,a p s ′ s,a ( θ k,h )   ϕ s ′ s,a - s ′′ ∈S h,s,a p s ′′ s,a ( θ k,h )ϕ s ′′ s,a H -1 k,h θ * h -θ k,h H k,h   ≤ H β k s ′ ∈S + h,s,a p s ′ s,a ( θ k,h )   ϕ s ′ s,a - s ′′ ∈S h,s,a p s ′′ s,a ( θ k,h )ϕ s ′′ s,a H -1 k,h   ≤ H β k s ′ ∈S h,s,a p s ′ s,a ( θ k,h )   ϕ s ′ s,a - s ′′ ∈S h,s,a p s ′′ s,a ( θ k,h )ϕ s ′′ s,a H -1 k,h   (29)where in the first inequality, we denote S + h,s,a as the subset of S h,s,a such thatϕ s ′ s,a (θ * h -θ k,h ) - s ′′ ∈S h,s,a p s ′′ s,a ( θ k,h )ϕ s ′′ s,a (θ * h -θ k,h) is non-negative, the second inequality holds by the Holder's inequality, and the third inequality is by the confidence set in Lemma 3.For the second-order term, we define u s ′ s,a (θ) = ϕ s ′ s,a θ, and p s ′ s,a (u) = exp(u s ′ s,a ) 1+ s ′′ exp(u s ′′ s,a ) , further defineF (u) = s ′ ∈S h,s,a exp(u s ′ s,a ) 1 + s ′′ ∈S h,s,a exp(u s ′′ s,a ) , F (u) = s ′ ∈S h,s,a exp(u s ′ s,a )V (s ′ ) 1 + s ′′ ∈S h,s,a exp(u s ′′ s,a ).Then, we have 1 2s ′ ∈S h,s,a ( θ k,h -θ * h ) ⊤ ∇ 2 p s ′ s,a ( θ)( θ k,h -θ * h )V (s ′ ) = 1 2 u( θ k,h ) -u(θ * h ) ⊤ ∇ 2 F (u( θ)) u( θ k,h ) -u(θ * h ) = 1 2 s ′ ∈S h,s,a s ′′ ∈S h,s,a u s ′ s,a ( θ k,h ) -u s ′ s,a (θ * h ) ⊤ ∂ 2 F (u( θ)) ∂s ′ ∂s ′′ u s ′′ s,a ( θ k,h ) -u s ′′ s,a (θ * h ) ≤ H 2 s ′ ∈S h,s,a s ′′ ∈S h,s,a u s ′ s,a ( θ k,h ) -u s ′ s,a (θ * h ) ∂ 2 F (u( θ)) ∂s ′ ∂s ′′ u s ′′ s,a ( θ k,h ) -u s ′′ s,a (θ * h )where the inequality holds by V (s) ∈ [0, H], ∀s.According to Lemma 18, we have (omit the subscript S h,s,a for simplicity):EQUATIONTo bound the first term, by applying the AM-GM inequality, we obtainEQUATIONPlugging ( 31) into (30), we haveEQUATIONwhere the last inequality holds by Lemma 3. Combining ( 29) and ( 32) finishes the proof. ■Proof. Combining Lemma 4 and Lemma 10, we haveK k=1 V * 1 (s k,1 ) - K k=1 V π k 1 (s k,1 ) ≤ 2 K k=1 H h=1 (ϵ fst k,h + ϵ snd k,h ) + H 2KH log(2/δ)whereϵ fst k,h = H β k s ′ ∈S k,h p s ′ k,h ( θ k,h ) ϕ s ′ k,h - s ′′ ∈S k,h p s ′′ k,h ( θ k,h )ϕ s ′′ k,h H -1 k,h , ϵ snd k,h = 5 2 H β 2 k max s ′ ∈S k,h ∥ϕ s ′ k,h ∥ 2 H -1 k,h.Next, we bound ϵ fst k,h and ϵ snd k,h respectively. Bounding ϵ fst k,h . For simplicity, we denoteE θ [ϕ s ′ k,h ] = E s ′ ∼p s k,h (θ) [ϕ s ′ k,h ], φs ′ s,a = ϕ s ′ s,a -E θ k,h [ϕ s ′ k,h ], ϕ s ′ s,a = ϕ s ′ s,a -E θ k+1,h [ϕ s ′ k,h ]Then, we haveEQUATIONEQUATIONEQUATION.We bound these terms separately in the following.For the first term (c), we haveφs ′ k,h -ϕ s ′ k,h H -1 k,h = s ′′ ∈S k,h p s ′′ k,h ( θ k+1,h ) -p s ′′ k,h ( θ k,h ) ϕ s ′′ k,h H -1 k,h = s ′′ ∈S k,h ∇p s ′′ k,h (ξ k,h ) ⊤ ( θ k+1,h -θ k,h ) ϕ s ′′ k,h H -1 k,h s ′′ ∈S k,h ∇p s ′′ k,h (ξ k,h ) ⊤ ( θ k+1,h -θ k,h ) ϕ s ′′ k,h H -1 k,h = s ′′ ∈S k,h p s ′′ k,h (ξ k,h )ϕ s ′′ k,h -p s ′′ k,h (ξ k,h ) s ′′′ ∈S k,h p s ′′′ k,h (ξ k,h )ϕ s ′′′ k,h ⊤ ( θ k+1,h -θ k,h ) ϕ s ′′ k,h H -1 k,h ≤ s ′′ ∈S k,h p s ′′ k,h (ξ k,h ) ϕ s ′′ k,h ⊤ ( θ k+1,h -θ k,h ) ϕ s ′′ k,h H -1 k,h + s ′′ ∈S k,h p s ′′ k,h (ξ k,h ) ϕ s ′′ k,h H -1 k,h s ′′′ ∈S k,h p s ′′′ k,h (ξ k,h ) (ϕ s ′′′ k,h ) ⊤ ( θ k+1,h -θ k,h ) ≤ s ′′ ∈S k,h p s ′′ k,h (ξ k,h ) θ k+1,h -θ k,h H k,h ϕ s ′′ k,h 2 H -1 k,h + s ′′ ∈S k,h p s ′′ k,h (ξ k,h ) ϕ s ′′ k,h H -1 k,h s ′′′ ∈S k,h p s ′′′ k,h (ξ k,h ) ϕ s ′′′ k,h H -1 k,h θ k+1,h -θ k,h H k,h ≤ 2η √ λ s ′′ ∈S k,h p s ′′ k,h (ξ k,h ) ϕ s ′′ k,h 2 H -1 k,h + 2η √ λ s ′′ ∈S k,h p s ′′ k,h (ξ k,h ) ϕ s ′′ k,h H -1 k,h 2 ≤ 4η √ λ s ′′ ∈S k,h p s ′′ k,h (ξ k,h ) ϕ s ′′ k,h 2 H -1 k,h ≤ 4η √ λ max s ′′ ∈S k,h ϕ s ′′ k,h 2 H -1 k,hwhere and the fifth inequality is by Cauchy-Schwarz inequality and the fourth inequality is because by Lemma 17, since H k,h ⪰ H k,h ⪰ λI d , we haveθ k+1,h -θ k,h H k,h ≤ θ k+1,h -θ k,h H k,h ≤ 2η∥g k,h ( θ k,h )∥ H -1 k,h ≤ 2η √ λ ∥g k,h ( θ k,h )∥ 2 , and since g k,h (θ) = s ′ ∈S k,h (p s ′ k,h (θ) -y s ′ k,h )ϕ s ′ k,h , we have ∥g k,h ( θ k,h )∥ 2 ≤ s ′ ∈S k,h p s ′ k,h ( θ k,h )ϕ s ′ k,h 2 + s ′ ∈S k,h y s ′ k,h ϕ s ′ k,h 2 ≤ 2 max s ′ ∈S k,h ϕ s ′ k,h 2 ≤ 2.Therefore, we haveEQUATION)where the last inequality holds by Lemma 6.For the term (d), by similar analysis, we have(p s ′ k,h ( θ k,h ) -p s ′ k,h ( θ k+1,h )) ϕ s ′ k,h H -1 k,h = ∇p s ′ k,h (ξ k,h ) ⊤ ( θ k,h -θ k+1,h ) ϕ s ′ k,h H -1 k,h = p s ′ k,h (ξ k,h )ϕ s ′ k,h -p s ′ k,h (ξ k,h ) s ′′ ∈S k,h p s ′′ k,h (ξ )ϕ s ′′ k,h ⊤ ( θ k+1,h -θ k,h ) ϕ s ′ k,h H -1 k,h ≤ 2η √ λ   p s ′ k,h (ξ k,h )∥ϕ s ′ k,h ∥ H -1 k,h ϕ s ′ k,h H -1 k,h + p s ′ k,h (ξ k,h ) ϕ s ′ k,h H -1 k,h s ′′ ∈S k,h p s ′′ k,h (ξ k,h ) ϕ s ′′ k,h H -1 k,h   ≤ 2η √ λ max s ′′ ∈S k,h ∥ϕ s ′′ k,h ∥ H -1 k,h ϕ s ′′ k,h H -1 k,h + max s ′′ ∈S k,h ϕ s ′′ k,h H -1 k,h max s ′′′ ∈S k,h ϕ s ′′′ k,h H -1 k,h ≤ 2η √ λ max s ′′ ∈S k,h ϕ s ′′ k,h 2 H -1 k,h + ϕ s ′′ k,h 2 H -1 k,h 2 + 2η √ λ max s ′′ ∈S k,h ϕ s ′′ k,h H -1 k,h 2 + max s ′′′ ∈S k,h ϕ s ′′′ k,h H -1 k,h 2 2 ≤ 4η √ λ max max s ′′ ∈S k,h ϕ s ′′ k,h 2 H -1 k,h , max s ′′ ∈S k,h ϕ s ′′ k,h 2 H -1 k,h, where the third inequality holds by the AM-GM inequality. Thus, we haveEQUATIONwhere the last inequality holds by Lemma 6.Finally, we bound the term (e) as follows.EQUATIONwhere the first inequality holds by the Cauchy-Schwarz inequality and the last holds by Lemma 6.Thus, combining (33), (34), and (35), we haveEQUATIONFor the second-order term, by Lemma 6, we haveEQUATION)where the last inequality holds by Lemma 6.Combining ( 36) and (37), we haveReg(K) ≤ 2 K k=1 H h=1 (ϵ fst k,h + snd k,h ) + H 2KH log(2/δ) ≤ H 2 β K 2dK log 1 + K dλ + 16η κ √ λ d log(1 + K dλ ) + 5 κ H 2 β 2 K d log 1 + K dλ + H 2KH log 2 δ ≤ O dH 2 √ K + d 2 H 2 κ -1This finishes the proof. ■Proof. Our proof is similar to the proof of the lower bound for adversarial linear mixture MDPs with the unknown transition in Zhao et al. [2023] . We prove this lemma by reducing the MNL mixture MDP problem to a sequence of binary logistic bandit problems.In this section, we provide several supporting lemmas used in the proofs of the main results.Lemma 15 (Abbasi-Yadkori et al. [2011, Theorem 1] ). Let {F t } ∞ t=0 be a filtration. Let {η t } ∞ t=1 be a real-valued stochastic process such that η t is F t -measurable and η t is conditionally zero-mean R-sub-Gaussian for some R ≥ 0i.e. ∀λ ∈ R, E e ληt | F t-1 ≤ exp λ 2 R 2 /2 . Let {X t } ∞ t=1 be an R d -valued stochastic process such that X t is t-1 -measurable. Assume that V is a d × d positive definite matrix. For any t ≥ 0, define V t = V + t-1 s=1 X s X ⊤ s S t = t-1 s=1 η s X s .Then, for any δ > 0, with probability at least 1 -δ, for all t ≥ 0,∥S t ∥ 2 V -1 t ≤ 2R 2 log det (V t ) 1/2 det(V ) -1/2 δ .Lemma 16 (Abbasi-Yadkori et al. [2011, Lemma 10] ). Suppose x 1 , . . . , x t ∈ R d and for any1 ≤ s ≤ t, ∥x s ∥ 2 ≤ L. Let V t = λI d + t s=1 x s x ⊤ s for λ ≥ 0. Then, for any 1 ≤ s ≤ t, we have det(V t ) ≤ λ + tL 2 d d .Lemma 17 (Orabona [2019, Lemma 6.9] ). Let Z be a positive define matrix and X be a convex set,x = arg min x∈X ⟨g, x⟩ + 1 2η ∥x -x∥ 2 Z .Then, for all x * ∈ X , we have∥ x -x∥ Z ≤ 2η∥g∥ Z -1 . and ⟨g, x -x * ⟩ ≤ η 2 ∥g∥ 2 Z -1 + 1 2η ∥ x -x * ∥ 2 Z - 1 2η ∥ x -x * ∥ 2 Z .Lemma 18 (Lee and Oh [2024, Lemma D.3] ). Define Q : R K → R, such that for any u = (u 1 , . . . , u K ) ∈ R K , Q(u) = K i=1 exp(ui) v+ K k=1 exp(u k ) . Let p i (u) = exp(ui) v+ K k=1 exp(u k ) . Then, for all i ∈ [K], we have∂ 2 Q ∂i∂j ⩽ 3p i (u) if i = j, 2p i (u)p j (u) if i ̸ = j.generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).After the submission of our work to arXiv[Li et al.,2024a], a follow up work byPark et al. [2024] improved the lower bound to Ω(dH 3/2 √ K). This confirms that our result is indeed optimal with respect to d and K.s ′ ∈S k,h ϕ s ′ k,h (ϕ s ′ k,h) ⊤ is the feature covariance matrix. 


Decision: This work makes multiple contributions on bounding the regret in the recently introduced problem of reinforcement learning where the transition kernel is modeled as a multinomial logit. On the positive side, a first result obtains the same regret as an earlier result of Hwang and Oh but with much better efficiency in terms of space and computation. A next result greatly reduces the order of dependence on a commonly appearing, often very small problem-dependent constant ($\kappa$). A third result entirely avoids dependence on $\kappa$ in the leading term although paying linearly in $\kappa^{-1}$ in a non-leading term. Also, a lower bound from a somewhat related setting (but not identical, as the lower bound considers a setting with infinite action space) suggests the authors’ results are nearly optimal.

From a technical standpoint, this work seems to be reasonably strong (maybe not very strong, but in light of the contributions, it seems nice).

Another positive point which I only got from closely looking at the paper is that the authors identify technical issues in previous analyses and explain how to fix them. This is nice.

With the above said, there are some strong reservations about this work which should be very well addressed in the next revision:
- First, the regret bounds depend on $U$, the maximum number of reachable states. This dependence was not made clear in the presentations in the theorems in the main text, and the paper lacked discussion about this new dependence. In the discussion period, the authors argued that in many situations, $\kappa$ may already have some dependence on $U$. However, being up front about dependencies on problem-dependent constants, and discussing them, is very important.
- The authors’ work requires knowing the support of the next-state distribution. From the author discussion period, the examples provided by the authors to justify having knowledge of the support were not that convincing.

Overall, this is a worthwhile paper that would fit well in the conference proceedings. Despite the reservations stated above, I believe the paper should be accepted.

Reviews: review: paper_summary: paper_summary: The paper studies the recently proposed MDPs that use multinomial logit function approximation for state distribution validness. The results and algorithms improve the prior work of Hwang and Oh [2023] in multiple aspects, including computation efficiency, storage, and statistical dependence on the problem-dependent quantity $\kappa$ that can be exponentially small. In addition, the authors establish a matching lower bound on $d$, the feature space dimension, and $K$, the number of episodes.
strength_weakness: strength_weakness: - The primary high-level techniques and tools (seem to) come from existing works and relevant fields, such as MNL contextual bandits. The authors should put more effort into highlighting the technical challenges and novelties besides the previous comparisons. 
- It would be beneficial to include experiments on synthetic and real-world datasets and compare the results to existing baselines and relevant works. In particular, the new algorithms seem more involved than prior ones, which may affect their stability and adaptiveness.
- There is still a significant gap between the lower and upper bounds. Besides, I wonder how often $\kappa$ could be exponentially small in practical settings, though it's definitely of theoretical interest to approach the lower limits on parameter dependency.
questions: questions: Overall, I think the paper makes reasonable contributions to the problem, and I have no additional questions/comments besides the above.
limitations: limitations: The authors have made various comparisons and discussed the limitations of the results, which I'm satisfied with. I do not see any potential negative societal impact of their work.
paper_summary: paper_summary: The problem considered in this paper is online learning in MDPs where transition probabilities are modelled with a log-linear model (with "multinomial logit function approximation"). The finite horizon, time-inhomogenous setting is considered. The problem is motivated by allowing a nonlinear transformation in modeling the MDP and yet maintaining both computational and information theoretic tractability. Inspired by results in the analogous bandit problems and algorithms developed for them, a number of gradually more complex, but (statistically) better performing algorithms are considered. In particular, while naive approaches give a poor dependence on a problem parameter $\kappa$ that characterizes the "strength" of nonlinearity, by adopting previous ideas to the MDP setting, new algorithms are designed that eliminate this poor dependence. A lower bound is also established, which nearly matches the upper bound (but considers infinite action spaces, while the main paper considers finite action spaces).
strength_weakness: strength_weakness: 1. The novelty is limited by that we have seen the same story, same ideas playing out nicely in the closely related bandit setting. 
2. A new parameter, U, the number of next states that are reachable with positive probability in the worst case, appears in the analysis and will appear in the bounds.
3. It is an unpleasant surprise for the reader to discover this dependence only through carefully reading the paper, rather than being told upfront. It is not good that the opportunity to discuss whether this quantity needs to enter the regret bound, and that this quantity needs to be small for the algorithm to be tractable, is missed.
4. Line 83 and onward: The work of Uohamma is discussed but is mischaracterized. My reading of this work is that they do establish that their algorithm runs in polynomial time. It remains unclear why the exponential family model is incomparable with the one considered here; an explanation (with examples) is missing.
5. The paper could use some extra proofreading (e.g., the upper indices in the bottom of page 5, in the displayed equation are not correct); in line 149, in the definition of $U$, $|\cdot|$ is missing.
questions: questions: 1. Can you confirm that the regret and compute cost depend on U, the worst-case number of next states that one can transition to with positive probability? Do you think such dependencies are necessary? Are there any interesting examples where it is reasonable to expect that U is small, independently of the size of the state space?
2. What was the most challenging aspect of extending the bandit ideas to the MDP framework?
limitations: limitations: n.a.
paper_summary: paper_summary: This work studies the MNL function approximation inhomogeneous RL, achieves the $O(1)$ computation cost, and improves the regret guarantee with regard to $\kappa$. To improve the computation cost, this work employs the online Newton step instead of MLE estimation to estimate $\theta$. Then, they design a novel confidence set by making full use of local information to improve the dependence of $\kappa$.
strength_weakness: strength_weakness: 1.	[1] also use the online Newton step to improve the computation cost in the logit contextual bandits setting. It would be better to discuss the novelty of UCRL-MNL-OL.

[1] Oh, M. H., & Iyengar, G. (2021, May). Multinomial logit contextual bandits: Provable optimality and practicality. In Proceedings of the AAAI conference on artificial intelligence (Vol. 35, No. 10, pp. 9205-9213).
questions: questions: Question 1: This work achieves great results in the stochastic reward setting. Can you discuss the challenge when extending to the adversarial reward setting?
limitations: limitations: The authors adequately addressed the limitations.
paper_summary: paper_summary: In this paper, the author analyzes a Markov Decision Process (MDP) model with non-linear function approximation. Specifically, in the finite-time horizon inhomogeneous episodic MDPs setting, the transition dynamics are unknown but the reward function is known. The author proposes using a multinomial logit (MNL) function approximation to estimate transition dynamics, which is superior to the linear function approximation if the model is misspecified in \cite{hwang2023model}. Additionally, the author proposes *UCRL-MBL-OL*, which adapts the previous work that is model-based and has large computational and storage complexity, to an online style that only consumes constant computation and storage resources. Moreover, the author has proven that the regret bound of *UCRL-MBL-OL* matches the state-of-the-art in Theorem 1. Its regret bound achieves $\tilde{O}(\kappa^{-1} dH^2\sqrt{K})$, where $H$ is the time horizon length, $K$ is the number of total episodes and $\kappa$ is considered as a parameter to control the sparsity of the transition dynamics and $d$ is the hidden dimensionality. Ignoring the logarithmic factor and $\kappa$, such a regret bound has only a $\sqrt{H}$ gap compared to the lower bound. After that, with additional assumption, the author utilizes the local information to propose another two algorithms, *UCRL-MNL-LL* and *UCRL-MNL-LL+* to remove the dependence on $\kappa$ and get a tighter regret bound as well as maintain good properties of *UCRL-MNL-OL*.
strength_weakness: strength_weakness: 1. Although this paper focuses on reducing the computation complexity, I am curious about the sample complexity of *UCRL-MNL-OL*.
    
2. Since the algorithm builds up the estimation of the transition dynamics by using MNL function approximation, is it considered a model-based algorithm? More specifically, does it require storing the transition dynamics for each state-action pair in every step?
questions: questions: Please see the above "Weaknesses"
limitations: limitations: NA
paper_summary: paper_summary: This paper considers MDPs employing the MNL function for transition probability, following Hwang and Oh [2023]. The authors suggest efficient algorithms based on online Newton steps, inspired by [Hazan et al., 2014; Zhang et al., 2016; Oh and Iyengar, 2021]. Furthermore, to improve $\kappa$ dependency, they provide algorithms employing local learning with mirror descent inspired by [Zhang and Sugiyama, 2023; Lee and Oh, 2024]. The algorithms achieve $1/\sqrt{\kappa}$ or even detach the dependency of $\kappa$ from the leading term.
strength_weakness: strength_weakness: - Their suggested algorithms do not seem novel because they are based on previously proposed methods for logistic or MNL bandits. Specifically, the online Newton update is widely studied for MNL or logistic bandits [Oh and Iyengar, 2021; Zhang and Sugiyama, 2023].

- Furthermore, the improvement on $\kappa$ is based on the mirror descent algorithm proposed in [Zhang and Sugiyama, 2023; Lee and Oh, 2024], and the proofs seem to follow the steps in [Zhang and Sugiyama, 2023; Lee and Oh, 2024] in the appendix.

- Lastly, the MNL model for transition probability may have an inherent weakness: the number of achievable states for each (k,n) must be finite, and it is required to know the state space of $S_{k,n}$.

[1] Faury, Louis, et al. "Jointly efficient and optimal algorithms for logistic bandits." International Conference on Artificial Intelligence and Statistics. PMLR, 2022.
questions: questions: Are there any non-trivial technical novelties in utilizing the online mirror descent method for MDPs?
limitations: limitations: The authors discuss some interesting future work regarding regret bound. Additionally, I believe, as mentioned in Weaknesses, the MNL model for transition probability has an inherent weakness: the number of achievable states for each (k,n) must be finite, and it is required to know the state space of $S_{k,n}$.

Comments: title: Thanks for your reply.
Thanks for your valuable comment.
Thank you.
Thank you.
Thank you for your comment
Thanks for your reply.
Thank you for your comment.
Thank you for your response.
Thank you for your response.
comment: Thank you for your helpful suggestions. We will incorporate these discussions into the next version. Thanks again!
Thanks for your constructive feedback. We agree with the reviewer's observation that the support of the next state distribution plays a key role in the MNL MDP model. Investigating how to remove this prior knowledge is an important future direction.

Even though, there are some real-world applications beyond games where the support of the next state distribution is limited and known. For instance, in the robot navigating problem, the robot can only move to nearby locations, even though the overall state space may be extensive. Similarly, in language models, the current state can be conceptualized as the sequence of previously generated words, with the immediately accessible next states being the potential subsequent words. Although the vocabulary could be extensive, the feasible choices for the next word, dictated by grammar and context, are inherently limited and known.

We will add these discussions in the next version. We appreciate your helpful suggestions!
Thanks for the rebuttal; it was useful. 
Overall, my feelings towards the paper have not changed: This is a fine peace of work that should get published. I still have some reservations concerning the model; in fact, one of the reservations I have I forgot to include into my original review is that the algorithm needs to know the support of the next state distribution. This combined with that the size of this support appears in the bound, and that in the related linear kernel MDP setting, the size of the support does not appear in the bounds makes me feel uneasy about the paper. Why is knowing the support problematic? Well, if the size of the support was not part of the bound, one could just "max out" the support (all states are possible next states). If not, in the lack of results showing that the algorithm is robust in the face of misspecified next state support, one is afraid that knowing the support will actually be important. And I expect knowing the support is kinda tricky. Consider for example the SuperMario game, mentioned in the rebuttal. Here, knowing the next states (all 4 of them) encodes a tremendous amount of information about the game. Just consider how many states there are here and compare this to the number 4. How realistic is that one would actually know the support in scenarios like this. So I have doubts. I think an example, where the assumptions are more natural would tremendously help the paper. Also, why not think of some "real" application that people may care about (I have to say I have a hard time imagining anyone to seriously care about how to play SuperMario; and I also have a hard time imagining a more general problem that has the characteristics of SuperMario (ie I don't have a problem with simplified examples, but SuperMario and other games just feel like too arbitrary and unrelated to any "real world" applications).) 
In summary, notwithstanding these reservations, I am in support of accepting the paper, given that it looks at a somewhat reasonable setting and makes a nontrivial contribution.
Thank you for your helpful suggestion. We will include a discussion on sample complexity in the next version.
Thanks for your answers. It would be great to add sample complexity to the paper. I will not change the rating and will continue to support this work.
Thanks for your re-evaluation! We will carefully revise the paper to better highlight the technical contributions and the significance of our results.
Thank you for clarifying your contribution. I still believe that the technical contribution may not be significant, which is an important factor for this conference. This is why I consider the contribution of this work to be on the borderline (4 or 5). However, I may have overlooked the first contribution of this paper regarding the extension to RL theory using recent bandit algorithms and techniques. From that perspective, I have raised my score to 5.
We appreciate your feedback but have to disagree with this comment. We would like to offer some clarifications that underscore the contributions of our work to the field of RL.

**Technical relationship between bandits and MDPs.** Since bandits can be viewed as one-step MDPs, it's *common and reasonable* for methodologies and analyses in MDPs to draw inspiration from bandits (e.g., linear bandits vs. linear MDPs). To some extent, bandit techniques are fundamental to modern RL theory and many RL algorithms are built upon bandit algorithms. However, the most intriguing and challenging aspect of RL theory lies in leveraging its unique structure, particularly when addressing the intrinsic dimension in function approximation settings.

**Our unique technical challenge and innovations.** Our work builds on recent advancements in MNL bandits, which have been well acknowledged in our paper. However, it is crucial to note that *several non-trivial technical innovations are necessary* to achieve favorable regret bounds. While employing high-level ideas like the Taylor expansion is not entirely new, the corresponding terms of MNL MDPs are significantly different from the bandit setting, which requires *new and more sophisticated analyses*. 

**Significance and Impact of Our Results.** Beyond technical novelty, our results hold significant value and are important for the community. Understanding function approximation is one of the central challenges in RL theory and there are many efforts devoted to, yielding fruitful results. However, one limitation is the linear assumption may not guarantee valid transition probabilities. MNL MDPs address this limitation by incorporating non-linear function approximation, which brings significant challenges as well. Our work not only is *the first to achieve nearly the same statistical and computational efficiency* as linear function approximation but also establishes *the first lower bound* for this problem. Our work broadens the scope of function approximation greatly and makes a significant step forward in RL theory.

We hope this clarification highlights the uniqueness and importance of our contributions and addresses your concerns regarding the novelty and impact of our work. We believe our work is crucial for the community. We appreciate it if you could reconsider your evaluation in light of these clarifications.
Thanks for your careful response. The rebuttal addresses my concerns and it would be better to add the above content in the next version. I maintain my score to support this valuable work.
Overall, I agree that there are several adjustments needed to apply the mirror descent approach to RL. This work extends previous research on bandits to the RL setting. However, I have concerns about the significance of the technical novelty of this work. I believe that the second-order Taylor expansion mentioned by the author may not be unique to RL, as it is also used in bandit problems. Therefore, I'm maintaining my score.
Thank you for your response. We would like to take this opportunity to provide additional clarification on the technical contributions of our work.

While our algorithm and analysis share certain similarities with existing work on bandits (e.g., Zhang and Sugiyama, 2023; Lee and Oh, 2024), our approach introduces several novel elements specifically tailored to the MDP setting.

As the reviewer rightly pointed out (Point 2 of Weaknesses), the online estimation step and the analysis in Section 5.1 do draw from recent advances in bandits. However, it’s important to note that this step only mitigates the dependence to $O(\kappa^{-1/2})$, and an $O(\kappa^{-1/2})$ dependence still persists in the regret bound as shown in Theorem 2.

To further reduce the $O(\kappa^{-1/2})$ dependence, it is crucial to examine how local information is preserved in the regret analysis. To this end, Zhang and Sugiyama (2023) use a second-order Taylor expansion to construct the optimistic revenue (Proposition 1), but their approach is designed for a multi-parameter model, which differs fundamentally from our single-parameter model. Lee and Oh (2024) employ *a first-order upper bound* for each item and *specific properties of MNL bandit* to construct the optimistic expected revenue (Eq. (5) and (6) of their paper), which does not hold in our setting. 

To preserve local information in our work, we need to estimate the value function more accurately. Specifically, whereas Lemma 2 addresses value differences between the value functions, Lemma 4 refines this by incorporating a second-order Taylor expansion, **allowing us to maintain the local information $p(\theta)$ rather than applying the maximum operator as in Lemma 2**. Although this second-order Taylor expansion is inspired by the work of Zhang and Sugiyama (2023), we need new analysis to address the challenges unique to the single-parameter model, such as handling the negative term in our first-order term, which is not present in their work. Moreover, a naive application of their analysis would lead to a regret bound that scales linearly with the number of the reachable states $U$, which is undesired in the MDP setting.

We hope this clarifies the contributions of our approach. We are happy to provide further details if needed.
For the technical novelty, you mentioned that "In contrast, we employ the second-order Taylor expansion to achieve a more accurate approximation of the value function. This approach is non-trivial and requires a different analysis." Could you provide more details regarding this?