Title: Policy Mirror Descent with Lookahead
Venue: NeurIPS 2024
Accepted: yes
Keywords: Policy mirror descent

policy gradient methods

policy iteration

multi-step greedy policy improvement

policy optimization
Abstract: Policy Mirror Descent (PMD) stands as a versatile algorithmic framework encompassing several seminal policy gradient algorithms such as natural policy gradient, with connections with state-of-the-art reinforcement learning (RL) algorithms such as TRPO and PPO. PMD can be seen as a soft Policy Iteration algorithm implementing regularized 1-step greedy policy improvement. However, 1-step greedy policies might not be the best choice and recent remarkable empirical successes in RL such as AlphaGo and AlphaZero have demonstrated that greedy approaches with respect to multiple steps outperform their 1-step counterpart. In this work, we propose a new class of PMD algorithms called $h$-PMD which incorporates multi-step greedy policy improvement with lookahead depth $h$ to the PMD update rule. To solve discounted infinite horizon Markov Decision Processes with discount factor $\gamma$, we show that $h$-PMD which generalizes the standard PMD enjoys a faster dimension-free $\gamma^h$-linear convergence rate, contingent on the computation of multi-step greedy policies. We propose an inexact version of $h$-PMD where lookahead action values are estimated. Under a generative model, we establish a sample complexity for $h$-PMD which improves over prior work. Finally, we extend our result to linear function approximation to scale to large state spaces. Under suitable assumptions, our sample complexity only involves dependence on the dimension of the feature map space instead of the state space size.

Policy Mirror Descent with Lookahead

Abstract

Policy Mirror Descent (PMD) stands as a versatile algorithmic framework encompassing several seminal policy gradient algorithms such as natural policy gradient, with connections with state-of-the-art reinforcement learning (RL) algorithms such as TRPO and PPO. PMD can be seen as a soft Policy Iteration algorithm implementing regularized 1-step greedy policy improvement. However, 1-step greedy policies might not be the best choice and recent remarkable empirical successes in RL such as AlphaGo and AlphaZero have demonstrated that greedy approaches with respect to multiple steps outperform their 1-step counterpart. In this work, we propose a new class of PMD algorithms called h -PMD which incorporates multi-step greedy policy improvement with lookahead depth h to the PMD update rule. To solve discounted infinite horizon Markov Decision Processes with discount factor γ , we show that h -PMD which generalizes the standard PMD enjoys a faster dimension-free γ h -linear convergence rate, contingent on the computation of multi-step greedy policies. We propose an inexact version of h -PMD where lookahead action values are estimated. Under a generative model, we establish a sample complexity for h -PMD which improves over prior work. Finally, we extend our result to linear function approximation to scale to large state spaces. Under suitable assumptions, our sample complexity only involves dependence on the dimension of the feature map space instead of the state space size.

1 Introduction

Policy Mirror Descent (PMD) is a general class of algorithms for solving reinforcement learning (RL) problems. Motivated by the surge of interest in understanding popular Policy Gradient (PG) methods, PMD has been recently investigated in a line of works [25, 53, 29, 20]. Notably, PMD encompasses several PG methods as particular cases via its flexible mirror map. A prominent example is the celebrated Natural PG method. PMD has also close connections to state-of-the-art methods such as TRPO and PPO [50] which have achieved widespread empirical success [41, 42], including most recently in fine-tuning Large Language Models via RL from human feedback [35]. Interestingly, PMD has also been inspired by one of the most fundamental algorithms in RL: Policy Iteration (PI). While PI alternates between policy evaluation and policy improvement, PMD regularizes the latter step to address the instability issue of PI with inexact policy evaluation.

Policy Iteration and its variants have been extensively studied in the literature, see e.g. [5, 32, 39]. In particular, PI has been studied in conjunction with the lookahead mechanism [12], i.e. using multistep greedy policy improvement instead of single-step greedy policies. Intuitively, the idea is that the application of the Bellman operator multiple times before computing a greedy policy leads to a more accurate approximation of the optimal value function. Implemented via Monte Carlo Tree Search (MCTS), multi-step greedy policy improvement 2 contributed to the empirical success of some of the

most impressive applications of RL including AlphaGo, AlphaZero and MuZero [44, 46, 45, 40]. Besides this practical success, a body of work [12, 13, 49, 52, 51] has investigated the role of lookahead in improving the performance of RL algorithms, reporting a convergence rate speed-up when enhancing PI with h -step greedy policy improvement with a reasonable computational overhead.

In this work, we propose to cross-fertilize the PMD class of algorithms with multi-step greedy policy improvement to obtain a novel class of PMD algorithms enjoying the benefits of lookahead. Our contributions are as follows:

We propose a novel class of algorithms called h -PMD enhancing PMD with multi-step greedy policy updates where h is the depth of the lookahead. This class collapses to standard PMD when h = 1 . When the stepsize parameter goes to infinity, we recover the PI algorithm with multiple-step greedy policy improvement previously analyzed in [12]. When solving a Markov Decision Process problem with h -PMD, we show that the value suboptimality gap converges to zero geometrically with a contraction factor of γ h . This rate is faster than the standard convergence rate of PI and the similar rate for PMD which was recently established in [20] for h ≥ 2 . This rate improvement requires the computation of a h -step lookahead value function at each iteration which can be performed using planning methods such as tree-search. These results for exact h -PMD are exposed in Sec. 4, when value functions can be evaluated exactly.

We examine the inexact setting where the h -step action value functions are not available. In this setting, we propose a Monte Carlo sample-based procedure to estimate the unknown h -step lookahead action-value function at each state-action pair. We provide a sample complexity for this Monte Carlo procedure, improving over standard PMD thanks to the use of lookahead. Larger lookahead depth translates into a better sample complexity and a faster suboptimality gap convergence rate. However, this improvement comes at a more intensive computational effort to perform the lookahead steps. Sec. 5 discusses this inexact setting and the aforementioned tradeoff.

We extend our results to the function approximation setting to address the case of large state spaces where tabular methods are not tractable. In particular, we design a h -PMD algorithm where the h -step lookahead action-value function is approximated using a linear combination of state-action feature vectors in the policy evaluation step. Under linear function approximation and using a generative sampling model, we provide a performance bound for our h -PMD algorithm with linear function approximation. Our resulting sample complexity only depends on the dimension of the feature map space instead of the size of the state space and improves over prior work analyzing PMD with function approximation.

We perform simulations to verify our theoretical findings empirically on the standard DeepSea RL environment from Deep Mind's bsuite [34]. Our experiments illustrate the convergence rate improvement of h -PMD with increasing lookahead depth h in both the exact and inexact settings.

2 Preliminaries

Notation. For any integer n ∈ N \ { 0 } , the set of integers from 1 to n is denoted by [ n ] . For a finite set X , we denote by ∆( X ) the probability simplex over the set X . For any d ∈ N , we endow the Euclidean space R d with the norm ∥ · ∥ ∞ defined for every v ∈ R d by ∥ v ∥ ∞ := max i ∈ [ d ] | v i | where v i is the i -th coordinate. The relative interior of a set Z is denoted by ri( Z ) .

Value functions and optimal policy. For any policy π ∈ Π , we define the state value function V π : S → R for every s ∈ S by V π ( s ) = E π [ ∑ ∞ t =0 γ t r ( s t , a t ) | s 0 = s ] . The stateaction value function Q π : S × A → R can be similarly defined for every s ∈ S , a ∈ A by Q π ( s, a ) = E π [ ∑ ∞ t =0 γ t r ( s t , a t ) | s 0 = s, a 0 = a ] where E π is the expectation over the state-

Markov Decision Process (MDP). We consider an infinite horizon discounted Markov Decision Process M = ( S , A , r, γ, P, ρ ) where S and A are finite sets of states and actions respectively with cardinalities S = |S| and A = |A| respectively, P : S × A → ∆( S ) is the Markovian probability transition kernel, r : S × A → [0 , 1] is the reward function, γ ∈ (0 , 1) is the discount factor and ρ ∈ ∆( S ) is the initial state distribution. A randomized stationary policy is a mapping π : S → ∆( A ) specifying the probability π ( a | s ) of selecting action a in state s for any s ∈ S , a ∈ A . The set of all such policies is denoted by Π . At each time step, the learning agent takes an action a with probability π ( a | s ) , the environment transitions from the current state s ∈ S to the next state s ′ ∈ S with probability P ( s ′ | s, a ) and the agent receives a reward r ( s, a ) .

action Markov chain ( s t , a t ) induced by the MDP and the policy π generating actions. The goal is to find a policy π maximizing V π . A classic result shows that there exists an optimal deterministic policy π ⋆ ∈ Π maximizing V π simultaneously for all the states [36]. In this work, we focus on searching for an ϵ -optimal policy, i.e. a policy π satisfying ∥ V ⋆ -V π ∥ ∞ ≤ ϵ where V ⋆ := V π ⋆ .

The Bellman optimality operator T and the expected Bellman operator T π for any given policy π ∈ Π are γ -contraction mappings w.r.t. the ∥ · ∥ ∞ norm and hence admit unique fixed points V ⋆ and V π respectively. In particular, these operators satisfy the following inequalities: ∥T π V -V π ∥ ∞ ≤ γ ∥ V -V π ∥ ∞ and ∥T V -V ⋆ ∥ ∞ ≤ γ ∥ V -V ⋆ ∥ ∞ for any V ∈ R S . Moreover, we recall that for e = (1 , 1 , . . . , 1) ∈ R S and any c ∈ R ≥ 0 , V ∈ R S , we have T π ( V + c e ) = T π V + γc e and T ( V + c e ) = T V + γc e . We recall that the set G ( V ⋆ ) := { π ∈ Π : T π V ⋆ = T V ⋆ } of 1 -step greedy policies w.r.t. the optimal value V ⋆ coincides with the set of stationary optimal policies.

Bellman operators and Bellman equations. Wewill often represent the reward function r by a vector in R SA and the transition kernel by the operator P ∈ R SA × S acting on vectors v ∈ R S (which can also be seen as functions defined over S ) as follows: ( Pv )( s, a ) = ∑ s ′ ∈S P ( s ′ | s, a ) v ( s ′ ) for any s ∈ S , a ∈ A . We further define the mean operator M π : R SA → R S mapping any vector Q ∈ R SA to a vector in R S whose components are given by ( M π Q )( s ) = ∑ a ∈A π ( a | s ) Q ( s, a ) , and the maximum operator M ⋆ : R SA → R S defined by ( M ⋆ Q )( s ) = max a ∈A Q ( s, a ) for any s ∈ S , Q ∈ R SA . Using these notations, we introduce for any given policy π ∈ Π the associated expected Bellman operator T π : R S → R S defined for every V ∈ R S by T π V := M π ( r + γPV ) and the Bellman optimality operator T : R S → R S defined by T V := M ⋆ ( r + γPV ) for every V ∈ R S . Using the notations r π := M π r and P π := M π P , we can also simply write T π V = r π + γP π V and T V = max π ∈ Π T π V for any V ∈ R S .

3 A Refresher on PMD and PI with Lookahead

3.1 PMDand its Connection to PI

Policy Gradient methods consist in performing gradient ascent updates with respect to an objective E ρ [ V π ( s )] w.r.t. the policy π parametrized by its values π ( a | s ) for s, a ∈ S × A in the case of a tabular policy parametrization. Interestingly, policy gradient methods have been recently shown to be closely connected to a class of Policy Mirror Descent algorithms by using dynamical reweighting of the Bregman divergence with discounted visitation distribution coefficients in a classical mirror descent update rule using policy gradients. We refer the reader to section 4 in [53] or section 3.1 in [20] for a detailed exposition. The resulting PMD update rule is given by

where we use the shorthand notations π k s = π k ( ·| s ) and π s = π ( ·| s ) , where ( η k ) is a sequence of positive stepsizes, Q π k s ∈ R A is a vector containing state action values for π k at state s , and D ϕ is the Bregman divergence induced by a mirror map ϕ : dom ϕ → R such that ∆( A ) ⊂ dom ϕ, i.e. for any π, π ′ ∈ Π , s ∈ S ,

where we suppose throughout this paper that the function ϕ is of Legendre type, i.e. strictly convex and essentially smooth in the relative interior of dom ϕ (see [37], section 26). The mirror map choice gives rise to a large class of algorithms. Two special cases are noteworthy: (a) When ϕ is the squared 2 -norm, the Bregman divergence is the squared Euclidean distance and the resulting algorithm is a projected Q -ascent and (b) when ϕ is the negative entropy, the corresponding Bregman divergence is the Kullback-Leibler (KL) divergence and (PMD) becomes the popular Natural Policy Gradient algorithm (see e.g. [53], section 4 for details).

Now, using our operator notations, one can observe that the (PMD) update rule above is equivalent to the following update rule (see Lemma A.1 for a proof),

where D ϕ ( π, π k ) ∈ R S is a vector whose components are given by D ϕ ( π s , π k s ) for s ∈ S . Note that the maximization can be carried out independently for each state component.

Connection to Policy Iteration. As previously noticed in the literature [53, 20], PMD can be interpreted as a 'soft' PI. Indeed, taking infinite stepsizes η k in PMD (or a null Bregman divergence) immediately leads to the following update rule:

which corresponds to synchronous Policy Iteration (PI). The method alternates between a policy evaluation step to estimate the value function V π k of the current policy π k and a policy improvement step implemented as a one-step greedy policy with respect to the current value function estimate.

3.2 Policy Iteration with Lookahead

As discussed in [12], Policy Iteration can be generalized by performing h ≥ 1 greedy policy improvement steps at each iteration instead of a single step. The resulting h -PI update rule is:

where h ∈ N \ { 0 } . The h -greedy policy π k +1 w.r.t V π k selects the first optimal action of a non-stationary h -horizon optimal control problem. It can also be seen as the 1 -step greedy policy w.r.t. T h -1 V π k . In the rest of this paper, we will denote by G h ( V ) the set of h -step greedy policies w.r.t. any value function V ∈ R S , i.e. G h ( V ) := argmax π ∈ Π T π T h -1 V = { π ∈ Π : T π T h -1 V = T h V } . Notice that G 1 ( V ) = G ( V ) . Overall, the h -PI method alternates between finding a h -greedy policy and estimating the value of this policy. Interestingly, similarly to PI, a h -greedy policy guarantees monotonic improvement, i.e. V π ′ ≥ V π component-wise for any π ∈ Π , π ′ ∈ G h ( V π ) . Moreover, since T h is a γ h -contraction, it can be shown that the sequence ( ∥ V ⋆ -V π k ∥ ∞ ) k is contracting with coefficient γ h . See section 3 in [12] for further details about h -PI.

Remark 3.1 . While the iteration complexity always improves with larger depth h (see Theorem 3 in [12]), each iteration becomes more computationally demanding. As mentioned in [14], when the model is known, the h -greedy policy can be computed with Dynamic Programming (DP) in linear time in the depth h .

Another possibility is to implement a tree-search of depth h starting from the root state s to compute the h -greedy policy in deterministic MDPs. We refer to [14] for further details. Sampling-based tree search methods such as Monte Carlo Tree Search (MCTS) (see [9]) can be used to circumvent the exponential complexity in the depth h of tree search methods.

4 Policy Mirror Descent with Lookahead

In this section we propose our novel Policy Mirror Descent (PMD) algorithm which incorporates h -step lookahead for policy improvement: h -PMD.

4.1 h -PMD: Using h -Step Greedy Updates

Similarly to the generalization of PI to h -PI discussed in the previous section, we obtain our algorithm h -PMD by incorporating lookahead to PMD as follows:

Notice that we recover (1), i.e. (PMD), by setting h = 1 . We now provide an alternative way to write the h -PMD update rule which is more convenient for its implementation. We introduce a few additional notations for this purpose. For any policy π ∈ Π , let V π h := T h -1 V π for any integer h ≥ 1 . Consider the corresponding action value function defined for every s, a ∈ S × A by:

Using these notations, it can be easily shown that the h -PMD rule above is equivalent to:

Before moving to the analysis of this scheme, we provide two concrete examples for the implementation of h -PMD' corresponding to the choice of two standard mirror maps.

(1) Projected Q h -ascent. When the mirror map ϕ is the squared 2 -norm, the corresponding Bregman divergence is the squared Euclidean distance and h -PMD' becomes

for every s ∈ S and Proj ∆( A ) is the projection operator on the simplex ∆( A ) .

(2) Multiplicative Q h -ascent. When the mirror map ϕ is the negative entropy, the Bregman divergence D ϕ is the Kullback-Leibler divergence and for every ( s, a ) ∈ S × A ,

Connection to AlphaZero. Before switching gears to the analysis, we point out an interesting connection between h -PMD and the successful class of AlphaZero algorithms [46, 45, 40] which are based on heuristics.

It has been shown in [16] that AlphaZero can be seen as a regularized policy optimization algorithm, drawing a connection to the standard PMD algorithm (with h = 1 ). We argue that AlphaZero is even more naturally connected to our h -PMD algorithm. Indeed, the Q values in [16, Eq. (7) p. 3] correspond more closely to the lookahead action value function Q π h approximated via tree search instead of the standard Q π function with h = 1 . This connection clearly delineates the dependence on the lookahead depth (or tree search depth) h which was not clear in [16]. We have implemented a version of our h -PMD algorithm with Deep Mind's MCTS implementation (see section D.6 and the code provided for details). Conducting further experiments to compare our algorithm to AlphaZero on similar large scale settings would be interesting. We are not aware of any theoretical convergence guarantee for AlphaZero which relies on many heuristics. In contrast, our h -PMD algorithm enjoys theoretical guarantees that we establish in the next sections.

4.2 Convergence Analysis for Exact h -PMD

Using the contraction properties of the Bellman operators, it can be shown that the suboptimality gap ∥ V π k -V ∗ ∥ ∞ for PI iterates converges to zero at a linear rate with a contraction factor γ , regardless of the underlying MDP. This instance-independent convergence rate was generalized to PMD in [20]. In this section, we establish a linear convergence rate for the subobtimality value function gap of the exact h -PMD algorithm with a contraction factor of γ h where h is the lookahead depth. We assume for now that the value function V π k and a greedy policy π k +1 in ( h -PMD) can be computed exactly. These assumptions will be relaxed in the next section dealing with inexact h -PMD.

Theorem 4.1 (Exact h -PMD) . Let ( c k ) be a sequence of positive reals and let the stepsize η k in ( h -PMD) satisfy η k ≥ 1 c k ∥ min π ∈G h ( V π k ) D ϕ ( π, π k ) ∥ ∞ where we recall that G h ( V π k ) is the set of greedy policies with respect to T h -1 V π k and that the minimum is computed component-wise. Initialized at π 0 ∈ ri(Π) , the iterates ( π k ) of ( h -PMD) with h ∈ N \ { 0 } satisfy for every k ∈ N ,

A few remarks are in order regarding this result:

Compared to [20], our new algorithm achieves a faster γ h -rate where h is the depth of the lookahead. It should be noted that the approach in [20] is claimed to be optimal over all methods that are guaranteed to increase the probability of the greedy action at each timestep. Adding lookahead to the policy improvement step circumvents this restriction.

Unlike prior work [53, 25] featuring distribution mismatch coefficients which can scale with the size of the state space, our convergence rate does not depend on any other instance-dependent quantities. This is thanks to the analysis which circumvents the use of the performance difference lemma similarly to [20].

Choosing c t = γ 2 h ( t +1) c 0 for a positive constant c 0 for every integer t , the bound becomes (after bounding the geometric sum): ∥ V ⋆ -V π k ∥ ∞ ≤ γ hk ( ∥ V ⋆ -V π 0 ∥ ∞ + c 0 (1 -γ ) 2 ) . As c 0 → 0 we recover the linear convergence result of h -PI established in [12].

This faster rate comes at the cost of a more complex value function computation at each iteration. However, our experiments suggest that the benefits of the faster convergence rate greatly outweigh the extra cost of computing the lookahead, in terms of both overall running time until convergence and sample complexity (in the inexact case). See section 7 and Appendix D for evidence.

The cost of computing the adaptive stepsizes is typically simply that of computing a Bregman divergence between two policies. See Appendix A.3, D.4 for a more detailed discussion.

We defer the proof of Theorem 4.1 to Appendix A. Our proof highlights the relationship between h -PMD and h -PI through the explicit use of Bellman operators. Notice that even in the particular case of h = 1 , our proof is more compact compared to the one in [20] (see sec. 6 and Lemma A.1 to A.3 therein for comparison) using our convenient notations.

5 Inexact and Stochastic h -Policy Mirror Descent

In this section, we discuss the case where the lookahead action value Q π k h defined in (2) is unknown. We propose a procedure to estimate it using Monte Carlo sampling and we briefly discuss an alternative using a tree search method. The resulting inexact h -PMD update rule for every s ∈ S is

where ˆ Q π k h is the estimated version of the lookahead action value Q π k h induced by policy π k . We conclude this section by providing a convergence analysis for inexact h -PMD and discussing its sample complexity using the proposed Monte Carlo estimator.

5.1 Monte Carlo h -Greedy Policy Evaluation

Estimating the lookahead action value function Q π h for a given policy π involves solving a h -horizon planning problem using samples from the MDP. Our estimation procedure combines a standard planning method with Monte Carlo estimation under a generative model of the MDP. We give an algorithmic description of the procedure below, and defer the reader to Appendix B.1 for a precise definition of the procedure. Applying the recursive algorithm below with k := h returns an estimate for the action value function Q π h ( s, a ) at a given state-action pair ( s, a ) ∈ S × A .

Note that the base case of the recursion estimates the value function using Monte Carlo rollouts, see Appendix B.1.

Remark 5.1 . Notice that actions are exhaustively selected in our procedure like in a planning method. Using bandit ideas to guide Monte Carlo planning (see e.g. [23]) would be interesting to investigate for more efficiency. We leave the investigation of such selective action sampling and exploration for future work. In practice, lookahead policies are often computed using tree search techniques [12]. Remark 5.2 . When h = 1 , the procedure collapses to a simple Monte Carlo estimate of the Q-value function Q π k at each iteration k of the h -PMD algorithm.

5.2 Analysis of Inexact h -PMD

The next theorem is a generalization of Theorem 4.1 to the inexact setting in which the lookahead function Q π k h can be estimated with some errors at each iteration k of h -PMD.

Theorem 5.3 (Inexact h -PMD) . Suppose there exists b ∈ R + s.t. ∥ ˆ Q π k h -Q π k h ∥ ∞ ≤ b where the maximum norm is over both the state and action spaces. Let ˜ G h = argmax π ∈ Π M π ˆ Q π k h and let ( c k ) be a sequence of positive reals and let the stepsize η k in ( h -PMD) satisfy η k ≥ 1 c k ∥ min π ∈ ˜ G h D ϕ ( π, π k ) ∥ ∞

where we recall that the minimum is computed component-wise. Initialized at π 0 ∈ ri(Π) , the iterates ( π k ) of inexact ( h -PMD) with h ∈ N \ { 0 } satisfy for every k ∈ N ,

The proof of this result is similar to the proof of Theorem 4.1 and consists in propagating the error b in the analysis. We defer it to Appendix B.

As can be directly seen from this formulation, the γ h convergence rate from the exact setting generalizes to the inexact setting. Therefore, higher lookahead depths lead to faster convergence rates. Another improvement with relation to [20] is that the additive error term has a milder dependence on the effective horizon 1 -γ : the term 2 b (1 -γ )(1 -γ h ) is smaller for all values of h ≥ 2 than 4 b (1 -γ ) 2 . Larger lookahead depths yield a smaller asymptotic error in terms of the suboptimality gap.

We are now ready to establish the sample complexity of inexact h -PMD under a generative model using Theorem 5.3 together with concentration results for our Monte Carlo lookahead action value estimator.

Theorem 5.4 (Sample complexity of h -PMD) . Assume that inexact h -PMD is run for a number of iterations K > 1 h (1 -γ ) log( 4 ϵ (1 -γ )(1 -γ h ) ) , using the choice of stepsize defined by the sequence ( c k ) := ( γ 2 h ( k +1) ) . Additionally, suppose we are given a target suboptimality value ϵ > 0 , and a probability threshold δ > 0 . Finally, assume the lookahead value function ˆ Q π k h is approximated at each iteration with the Monte Carlo estimation procedure described in section 5.1, with the following parameter values: M 0 = ˜ O ( γ 2 h (1 -γ ) 4 (1 -γ h ) 2 ε 2 ) and for all j ∈ [ h ] , M j = M = ˜ O ( 1 (1 -γ ) 6 ε 2 ) . Then, with probability at least 1 -δ , the suboptimality at all iterations k satisfy the following bound:

Using this procedure, h -PMD uses at most KM 0 HS + hKMSA samples in total. The overall sample complexity of the inexact h -PMD is then given by ˜ O ( S hϵ 2 (1 -γ ) 6 (1 -γ h ) 2 + SA ϵ 2 (1 -γ ) 7 ) where the notation ˜ O hides at most polylogarithmic factors.

Compared to Theorem 5.1 in [20], the dependence on the effective horizon improves by a factor of the order of 1 / (1 -γ ) for h > A -1 (1 -γ ) -1 thanks to our Monte Carlo lookahead estimator. Their sample complexity is of order ˜ O ( 1 ϵ 2 (1 -γ ) 8 ) , whereas for h > 1 A (1 -γ ) ours becomes ˜ O ( 1 ϵ 2 (1 -γ ) 7 ) .

Our sample complexity does not depend on quantities such as the distribution mismatch coefficient used in prior work [53, 25], which may scale with the state space size.

6 Extension to Function Approximation

In MDPs with prohibitively large state action spaces, we represent each state-action pair ( s, a ) ∈ S × A by a feature vector ψ ( s, a ) ∈ R d where typically d ≪ SA and where ψ : S × A → R d is a feature map, also represented as a matrix Ψ ∈ R SA × d . Assuming that this feature map holds enough 'information' about each state action pair (more rigorous assumptions will be provided later), it is possible to approximate any value function using a matrix vector product Ψ θ , where θ ∈ R d . In this section, we propose an inexact h -PMD algorithm using action value function approximation.

6.1 Inexact h -PMD with Function Approximation

Using the notations above, approximating Q π k h by Ψ θ k , the inexact h -PMD update rule becomes:

We note that the policy update above can be implicit to avoid computing a policy for every state. It is possible to implement this algorithm using our Monte Carlo planning method in a similar fashion

to [27]: use the procedure in section 5.1 to estimate Q π k h ( s, a ) for all state action pairs ( s, a ) in some set C ⊂ S × A , and use these estimates as targets for a least squares regression to approximate Q π k h by Ψ θ k at time step k . See Appendix C for further details. We conclude this section with a convergence analysis of h -PMD with linear function approximation.

6.2 Convergence Analysis of h -PMD with Linear Function Approximation

Our analysis follows the approach of [27]. We make the following standard assumptions.

Assumption 6.1. The feature matrix Ψ ∈ R SA × d where d ≤ SA is full rank.

Assumption 6.2 (Approximate Universal value function realizability) . There exists ϵ FA > 0 s.t. for any π ∈ Π , inf θ ∈ R d ∥ Q π h -Ψ θ ∥ ∞ ≤ ϵ FA .

We defer a discussion about Assumption 6.2 to Appendix C.2.

Theorem 6.3 (Convergence of h -PMD with linear function approximation) . Let ( θ k ) be the sequence of iterates produced by h -PMD with linear function approximation, run for K iterations, starting with policy π 0 ∈ ri(Π) , using stepsizes defined by ( c k ) . Assume that targets ˆ Q π k h were computed using the procedure described in sec. 5.1 such that with probability at least 1 -δ, ∀ k ≤ K, ∀ z ∈ C ⊂ S × A the targets satisfy | ˆ Q π k h ( z ) -Q π k h ( z ) | ≤ ϵ . 3

Then, there exists a choice of C for which the policy iterates satisfy at each iteration:

The proof of Theorem 6.3 can be found in Appendix C.3. Notice that our performance bound does not depend on the state space size like in sec. 5 and depends on the dimension d of the feature space instead. We should notice though that the choice of the set C is crucial for this result and given by the Kiefer-Wolfowitz theorem [22] which guarantees that |C| = O ( d 2 ) as discussed in [27]. Moreover, our method for function approximation yields a sample complexity that is instance independent. Existing results for PMD with function approximation [3, 54] depend on distribution mismatch coefficients which can scale with the size of the state space.

In terms of computational complexity, at each iteration, h -PMD uses O (( MA ) h ) elementary operations where M is the size of the minibatch of trajectories used, A is the size of the action space and h is the lookahead depth. This computational complexity which scales exponentially in h is inherent to tree search methods. Despite this seemingly prohibitive computational scaling, tree search has been instrumental in practice for achieving state of the art results in some environments [40, 19] and modern implementations (notably using GPU-based search on vectorized environments) make tree search much more efficient in practice.

7 Simulations

We conduct simulations to investigate the effect of the lookahead depth on the convergence rate and illustrate our theoretical findings. We run the h -PMD algorithm for different values of h in both exact and inexact settings on the DeepSea environment from DeepMind's bsuite [34] using a grid size of 64 by 64, and a discount factor γ = 0 . 99 . Additional experiments are provided in Appendix D. Our codebase where all our experiments can be replicated is available here: https: //gitlab.com/kimon.protopapa/pmd-lookahead .

Exact h -PMD. We run the exact h -PMD algorithm for 100 iterations for increasing values of h using the KL divergence. Similar results were observed for the Euclidean divergence. We tested two different stepsize schedules: (a) in dotted lines in Fig. 1 (left), η k equal to its lower bound in sec. 4, with the choice c k := γ 2 h ( k +1) (note the dependence on h ); and (b) in solid lines, η k identical stepsize schedule across all values of h with c k := γ 2( k +1) to isolate the effect of the lookahead. We clearly observe under both stepsize schedules that h -PMD converges in fewer iterations when lookahead value functions are used instead of regular Q functions.

0

10

20

30

40

Iteration

0

.

0

0

.

1

0

.

2

0

.

3

0

.

4

0

.

5

‖

V /star

-

V π

k

‖∞

0

25

50

75

100

Iteration

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

0

10

20

30

Runtime (s)

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

h

=1

h

=5

h

=10

h

=15

h

=20

Inexact h -PMD. In this setting, we estimated the value function Q h using the vanilla Monte Carlo planning procedure detailed in section 5 in a stochastic variant of the DeepSea environment, and all the parameters except for h were kept identical across runs. As predicted by our results, h -PMD converges in fewer iterations when h is increased (see Fig. 1 (middle)). We also observed in our simulations that h -PMD uses less samples overall (see Appendix D for the total number of samples used at each iteration), and usually converges after less overall computation time (see Fig. 1 (right)). We also refer the reader to Appendix D where we have performed additional experiments with a larger lookahead depth h = 100 . In this case, the algorithm converges in a single iteration. This is theoretically expected as computing the lookahead values with very large h boils down to computing the optimal values like in value iteration with a large number of iterations. We also performed additional experiments in continuous control tasks to illustrate the general applicability of our algorithm (see Fig. 8 in Appendix D).

Remark 7.1 . ( Choice of the lookahead depth h . ) The lookahead depth is a hyperparameter of the algorithm and can be tuned similarly to other hyperparameters such as the step size of the algorithm. Of course, the value would depend on the environment and the structure of the reward at hand. Sparse and delayed reward settings will likely benefit from lookahead with larger depth values. We have performed several simulations with different values of h for each environment setting and the performance can potentially improve drastically with a better lookahead depth value (see also appendix D for further simulations). In addition, we observe that larger lookahead depth is not always better: see Fig. 7 in the appendix for an example where large lookahead depth becomes slower and does not perform better. Note that in this more challenging practical setting the best performance is not obtained for higher values of h : intermediate values of h perform better. This illustrates the tradeoff in choosing the depth h between an improved convergence rate and the computational cost induced by a larger h . We believe further investigation regarding the selection of the depth parameter might be useful to further improve the practical performance of the algorithm.

8 Related Work

Policy Mirror Descent and Policy Gradient Methods. Motivated by their empirical success [41, 42], the analysis of PG methods has recently attracted a lot of attention [1, 6, 21, 7, 53, 25, 20]. Among these works, a line of research focused on the analysis of PMD [53, 25, 28, 29, 20, 3] as a flexible algorithmic framework, and its particular instances such as the celebrated natural policy gradient [21, 54]. In particular, for solving unregularized MDPs (which is our main focus), Xiao [53] established linear convergence of PMD in the tabular setting with a rate depending on the instance dependent distribution mismatch coefficients, and extended their results to the inexact setting using a generative model. Similar results were established for the particular case of the natural policy gradient algorithm [6, 21]. More recently, Johnson, Pike-Burke, and Rebeschini [20] showed a dimension-free

rate for PMD with exact policy evaluation using adaptive stepsizes, closing the gap with PI. Notably, they further show that this rate is optimal and their analysis inspired from PI circumvents the use of the performance difference lemma which was prevalent in prior work. We improve over these results using lookahead and we extend our analysis beyond the tabular setting by using linear function approximation. To the best of our knowledge, the use of multi-step greedy policy improvement in PMD and its analysis are novel in the literature.

Policy Iteration with Multiple-Step Policy Improvement. Multi-step greedy policies have been successfully used in applications such as the game of Go and have been studied in a body of works [12, 13, 52, 15, 49]. In particular, multi-step greedy policy improvement has been studied in conjunction with PI in [12, 52]. Efroni et al. [12] introduced h -PI which incorporates h -lookahead to PI and generalized existing analyses on PI with 1-step greedy policy improvement to h -step greedy policies. However, their work requires access to an h -greedy policy oracle, and the h -PI algorithm is unstable in the stochastic setting even for h = 1 because of potentially large policy updates. We address all these issues in the present work with our h -PMD algorithm. Winnicki and Srikant [52] showed that a first-visit version of a PI scheme using a single sample path for policy evaluation converges to the optimal policy provided that the policy improvement step uses lookahead instead of a 1 -step greedy policy. Unlike our present work, the latter work assumes access to a lookahead greedy policy oracle. More recently, Alegre et al. [2] proposed to use lookahead for policy improvement in a transfer learning setting. Extending Generalized Policy Improvement (GPI) [4] which identifies a new policy that simultaneously improves over all previously learned policies each solving a specific task, they introduce h -GPI which is a multi-step extension of GPI. We rather consider a lookahead version of PMD including policy gradient methods instead of PI. We focus on solving a single task and provide policy optimization guarantees in terms of iteration and sample complexities.

Lookahead-Based Policy Improvement with Tree Search. Lookahead policy improvement is usually implemented via tree search methods [9, 43, 40, 24, 18, 14]. Recently, Dalal et al. [11] proposed a method combining policy gradient and tree search to reduce the variance of stochastic policy gradients. The method relies on a softmax policy parametrization incorporating a tree expansion. Our general h -PMD framework which incorporates a different lookahead policy improvement step also brings together PG and tree search methods. Morimura et al. [31] designed a method with a mixture policy of PG and MCTS for non-Markov Decision Processes. Grill et al. [17] showed that AlphaZero as well as several other MCTS algorithms compute approximate solutions to a family of regularized policy optimization problems which bear similarities with our h -PMD algorithm update rule for h = 1 (see sec. 3.2 therein).

9 Conclusion

In this work, we introduced a novel class of PMD algorithms with lookahead inspired by the success of lookahead policies in practice. We have shown an improved γ h -linear convergence rate depending on the lookahead depth in the exact setting. We proposed a stochastic version of our algorithm enjoying an improved sample complexity. We further extended our results to scale to large state spaces via linear function approximation with a performance bound independent of the state space size. Our paper offers several interesting directions for future research. A possible avenue for future work is to investigate the use of more general function approximators such as neural networks to approximate the lookahead value functions and scale to large state-action spaces. Recent work for PMD along these lines [3] might be a good starting point. Enhancing our algorithm with exploration mechanisms (see e.g. the recent work [33]) is an important future direction. Our PMD update rule might offer some advantage in this regard as was recently observed in the literature [28]. Constructing more efficient fully online estimators for the lookahead action values using MCTS and designing adaptive lookahead strategies to select the tree search horizon [38] are promising avenues for future work. Our algorithm brings together two popular and successful families of RL algorithms: PG methods and tree search methods. Looking forward, we hope our work further stimulates research efforts to investigate practical and efficient methods enjoying the benefits of both classes of methods.

Acknowledgments and Disclosure of Funding

We would like to thank the anonymous area chair and reviewers for their useful comments which helped us to improve our paper. Most of this work was completed when A.B. was affiliated with ETH Z¨ urich as a postdoctoral fellow. This work was partially supported by an ETH Foundations of Data Science (ETH-FDS) postdoctoral fellowship.

References

A. Agarwal, S. M. Kakade, J. D. Lee, and G. Mahajan. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. The Journal of Machine Learning Research , 22(1):4431-4506, 2021. 9

L. N. Alegre, A. Bazzan, A. Nowe, and B. C. da Silva. Multi-step generalized policy improvement by leveraging approximate models. In Advances in Neural Information Processing Systems , volume 36, pages 38181-38205, 2023. 10

C. Alfano, R. Yuan, and P. Rebeschini. A novel framework for policy mirror descent with general parameterization and linear convergence. In Thirty-seventh Conference on Neural Information Processing Systems , 2023. 8, 9, 10

A. Barreto, S. Hou, D. Borsa, D. Silver, and D. Precup. Fast reinforcement learning with generalized policy updates. Proceedings of the National Academy of Sciences , 117(48):3007930087, 2020. 10

D. Bertsekas and J. N. Tsitsiklis. Neuro-dynamic programming . Athena Scientific, 1996. 1

J. Bhandari and D. Russo. On the linear convergence of policy gradient methods for finite mdps. In International Conference on Artificial Intelligence and Statistics , pages 2386-2394. PMLR, 2021. 9

J. Bhandari and D. Russo. Global optimality guarantees for policy gradient methods. Operations Research , 2024. 9

G. Brockman, V. Cheung, L. Pettersson, J. Schneider, J. Schulman, J. Tang, and W. Zaremba. Openai gym, 2016. 31

C. B. Browne, E. Powley, D. Whitehouse, S. M. Lucas, P. I. Cowling, P. Rohlfshagen, S. Tavener, D. Perez, S. Samothrakis, and S. Colton. A survey of monte carlo tree search methods. IEEE Transactions on Computational Intelligence and AI in games , 4(1):1-43, 2012. 4, 10

G. Chen and M. Teboulle. Convergence analysis of a proximal-like minimization algorithm using bregman functions. SIAM Journal on Optimization , 3(3):538-543, 1993. 17

G. Dalal, A. Hallak, S. Mannor, and G. Chechik. Softtreemax: Policy gradient with tree search, 2022. 10

Y. Efroni, G. Dalal, B. Scherrer, and S. Mannor. Beyond the one-step greedy approach in reinforcement learning. In Proceedings of the 35th International Conference on Machine Learning , volume 80 of Proceedings of Machine Learning Research , pages 1387-1396. PMLR, 10-15 Jul 2018. 1, 2, 4, 5, 6, 10

Y. Efroni, G. Dalal, B. Scherrer, and S. Mannor. Multiple-step greedy policies in approximate and online reinforcement learning. In Advances in Neural Information Processing Systems , volume 31, 2018. 2, 10

Y. Efroni, G. Dalal, B. Scherrer, and S. Mannor. How to combine tree-search methods in reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 33, pages 3494-3501, 2019. 4, 10

Y. Efroni, M. Ghavamzadeh, and S. Mannor. Online planning with lookahead policies, 2020. 10

J.-B. Grill, F. Altch´ e, Y. Tang, T. Hubert, M. Valko, I. Antonoglou, and R. Munos. Montecarlo tree search as regularized policy optimization. In International Conference on Machine Learning , pages 3769-3778. PMLR, 2020. 5

J.-B. Grill, F. Altch´ e, Y. Tang, T. Hubert, M. Valko, I. Antonoglou, and R. Munos. Monte-Carlo tree search as regularized policy optimization. In H. D. III and A. Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 3769-3778. PMLR, 13-18 Jul 2020. 10

A. Hallak, G. Dalal, S. Dalton, I. Frosio, S. Mannor, and G. Chechik. Improve agents without retraining: Parallel tree search with off-policy correction, 2023. 10

M. Hessel, I. Danihelka, F. Viola, A. Guez, S. Schmitt, L. Sifre, T. Weber, D. Silver, and H. van Hasselt. Muesli: Combining improvements in policy optimization, 2022. 8

E. Johnson, C. Pike-Burke, and P. Rebeschini. Optimal convergence rate for exact policy mirror descent in discounted markov decision processes. In Thirty-seventh Conference on Neural Information Processing Systems , 2023. 1, 2, 3, 4, 5, 6, 7, 9, 16, 17, 18, 19, 29

S. Khodadadian, P. R. Jhunjhunwala, S. M. Varma, and S. T. Maguluri. On linear and superlinear convergence of natural policy gradient algorithm. Systems & Control Letters , 164:105214, 2022. 9

J. Kiefer and J. Wolfowitz. The equivalence of two extremum problems. Canadian Journal of Mathematics , 12:363-366, 1960. 8, 27

L. Kocsis and C. Szepesv´ ari. Bandit based monte-carlo planning. In Machine Learning: ECML 2006 , pages 282-293, Berlin, Heidelberg, 2006. Springer Berlin Heidelberg. 6

L. Kocsis and C. Szepesv´ ari. Bandit based monte-carlo planning. In J. F¨ urnkranz, T. Scheffer, and M. Spiliopoulou, editors, Machine Learning: ECML 2006 , pages 282-293, Berlin, Heidelberg, 2006. Springer Berlin Heidelberg. 10

G. Lan. Policy mirror descent for reinforcement learning: Linear convergence, new sampling complexity, and generalized problem classes. Mathematical programming , 198(1):1059-1106, 2023. 1, 5, 7, 9, 18

R. T. Lange. gymnax: A JAX-based reinforcement learning environment library, 2022. 30

T. Lattimore and C. Szepesvari. Learning with good feature representations in bandits and in rl with a generative model. In International Conference on Machine Learning , 2019. 8, 27, 28

Y. Li and G. Lan. Policy mirror descent inherently explores action space. arXiv preprint arXiv:2303.04386 , 2023. 9, 10

Y. Li, G. Lan, and T. Zhao. Homotopic policy mirror descent: policy convergence, algorithmic regularization, and improved sample complexity. Mathematical Programming , pages 1-57, 2023. 1, 9, 18

J. Mei, B. Dai, A. Agarwal, M. Ghavamzadeh, C. Szepesv´ ari, and D. Schuurmans. Orderingbased conditions for global convergence of policy gradient methods. Advances in Neural Information Processing Systems , 36, 2024. 27

T. Morimura, K. Ota, K. Abe, and P. Zhang. Policy gradient algorithms with monte-carlo tree search for non-markov decision processes, 2022. 10

R. Munos. Error bounds for approximate policy iteration. In Internation Conference on Machine Learning , volume 3, pages 560-567. Citeseer, 2003. 1

B. O'Donoghue. Efficient exploration via epistemic-risk-seeking policy optimization, 2023. 10

I. Osband, Y. Doron, M. Hessel, J. Aslanides, E. Sezener, A. Saraiva, K. McKinney, T. Lattimore, C. Szepesv´ ari, S. Singh, B. Van Roy, R. Sutton, D. Silver, and H. van Hasselt. Behaviour suite for reinforcement learning. In International Conference on Learning Representations , 2020. 2, 8, 28, 30, 31

L. Ouyang, J. Wu, X. Jiang, D. Almeida, C. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems , 35:27730-27744, 2022. 1

M. L. Puterman. Markov decision processes: discrete stochastic dynamic programming . John Wiley & Sons, 2014. 3

R. T. Rockafellar. Convex Analysis . Princeton University Press, Princeton, 1970. 3

A. Rosenberg, A. Hallak, S. Mannor, G. Chechik, and G. Dalal. Planning and learning with adaptive lookahead. Proceedings of the AAAI Conference on Artificial Intelligence , 37(8):96069613, Jun. 2023. 10

B. Scherrer. Improved and generalized upper bounds on the complexity of policy iteration. Advances in Neural Information Processing Systems , 26, 2013. 1

J. Schrittwieser, I. Antonoglou, T. Hubert, K. Simonyan, L. Sifre, S. Schmitt, A. Guez, E. Lockhart, D. Hassabis, T. Graepel, T. P. Lillicrap, and D. Silver. Mastering atari, go, chess and shogi by planning with a learned model. Nature , 588:604 - 609, 2019. 2, 5, 8, 10, 31

J. Schulman, S. Levine, P. Abbeel, M. Jordan, and P. Moritz. Trust region policy optimization. In International conference on machine learning , pages 1889-1897. PMLR, 2015. 1, 9

J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization algorithms, 2017. 1, 9

D. Shah, Q. Xie, and Z. Xu. Non-asymptotic analysis of monte carlo tree search, 2020. 10

D. Silver, A. Huang, C. J. Maddison, A. Guez, L. Sifre, G. Van Den Driessche, J. Schrittwieser, I. Antonoglou, V. Panneershelvam, M. Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature , 529(7587):484-489, 2016. 2

D. Silver, T. Hubert, J. Schrittwieser, I. Antonoglou, M. Lai, A. Guez, M. Lanctot, L. Sifre, D. Kumaran, T. Graepel, et al. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science , 362(6419):1140-1144, 2018. 2, 5

D. Silver, J. Schrittwieser, K. Simonyan, I. Antonoglou, A. Huang, A. Guez, T. Hubert, L. Baker, M. Lai, A. Bolton, et al. Mastering the game of go without human knowledge. nature , 550(7676):354-359, 2017. 2, 5

R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction . MIT press, 2018. 1

C. Szepesv´ ari. CMPUT 653: Theoretical Foundations of Reinforcement Learning. Lecture 6. Online planning - Part II., 2023. 21, 23, 26

M. Tomar, Y. Efroni, and M. Ghavamzadeh. Multi-step greedy reinforcement learning algorithms. In International Conference on Machine Learning , pages 9504-9513. PMLR, 2020. 2, 10

M. Tomar, L. Shani, Y. Efroni, and M. Ghavamzadeh. Mirror descent policy optimization. In International Conference on Learning Representations , 2022. 1

A. Winnicki, J. Lubars, M. Livesay, and R. Srikant. The role of lookahead and approximate policy evaluation in reinforcement learning with linear value function approximation. arXiv preprint arXiv:2109.13419 , 2021. 2

A. Winnicki and R. Srikant. On the convergence of policy iteration-based reinforcement learning with monte carlo policy evaluation. In International Conference on Artificial Intelligence and Statistics , pages 9852-9878. PMLR, 2023. 2, 10

L. Xiao. On the convergence rates of policy gradient methods. The Journal of Machine Learning Research , 23(1):12887-12922, 2022. 1, 3, 4, 5, 7, 9, 15, 17, 18

R. Yuan, S. S. Du, R. M. Gower, A. Lazaric, and L. Xiao. Linear convergence of natural policy gradient methods with log-linear policies. In The Eleventh International Conference on Learning Representations , 2023. 8, 9

Contents

A Additional Details and Proofs for Section 4: Exact h -PMD

Lemma A.1. For any policies π, π ′ ∈ Π and every state s ∈ S , we have ⟨ Q π ′ s , π s ⟩ = T π V π ′ ( s ) . As a consequence, the update rules (PMD) and (1) in the main part are equivalent.

Proof. Let π, π ′ ∈ Π . For every s ∈ S , we have

The second part of the statement holds by setting π ′ = π k .

A.1 Auxiliary Lemma for the Proof of Theorem 4.1: Exact h -PMD

Lemma A.2 (Triple Point Descent Lemma for h -PMD) . Let h ∈ N \ { 0 } . Initialized at π 0 ∈ ri(Π) , the iterates of h -step greedy PMD remain in ri(Π) and satisfy the following inequality for any k ∈ N and any policy π ∈ Π ,

Proof. Recall the definition of the h -step greedy action value function denoted by Q π h : S × A → R in (2). Then, recall that for each state s ∈ S and every k ∈ N , we have

Applying the three point descent lemma (see Lemma 6 in [53] with ϕ ( π ) = ⟨ Q π k h ( s, · ) , π ⟩ for any π ∈ Π and C = ∆( A ) with the notation therein), we obtain for any policy π ∈ Π ,

We obtain the desired result by rearranging the above last inequality and using (9).

Lemma A.3. Let h ∈ N \ { 0 } . Initialized at π 0 ∈ ri(Π) , the iterates of h -step greedy PMD satisfy for every k ∈ N , where ( c k ) is the sequence of positive integers introduced in Theorem 4.1 such that ˜ d k η k ≤ c k and ˜ d k = ∥ min π ∈G h ( V π k ) D ϕ ( π, π k ) ∥ ∞ .

Proof. Let k ∈ N . Recall that any policy π ∈ G h ( V π k ) satisfies T π T h -1 V π k = T h V π k . Using Lemma A.2 for π ∈ G h ( V π k ) together with the previous identity and the nonnegativity of the Bregman divergence implies that for any for any policy π ∈ G h ( V π k ) ,

where the last inequality follows from the assumption on the stepsizes η k . Wenowprove by recurrence that for every n ∈ N ,

with ∑ n i =1 γ i = 0 for n = 0 . The base case for n = 0 is immediate. We now suppose that the result holds for n and we show that it remains true for n +1 . We have

T π k +1 T h -1 V π k = T π k +1 T h -1 T π k V π k ( T π k V π k = V π k ) ≤ T π k +1 T h V π k ( T π k V ≤ T V for any V ∈ R S and monotonicity of T π k +1 ) ≤ T π k +1 ( T π k +1 T h -1 V π k + c k e ) (using (10) and monotonocity of T π k +1 ) = ( T π k +1 ) 2 T h -1 V π k + γc k e ( T π ( V + c e ) = T π V + γc e for any V ∈ R S , c ∈ R ) ≤ T π k +1 ( ( T π k +1 ) n +1 T h -1 V π k + n ∑ i =1 γ i c k e ) + γc k e (using recurrence hypothesis and monotonicity of T π k +1 ) n

which concludes the recurrence proof of (11) for every n ∈ N . To conclude, we take the limit in (11) and use the fact that for any V ∈ R S , the sequence (( T π k +1 ) n V ) n converges to the fixed point V π k +1 of the contractive operator T π k +1 . Hence, we obtain

This concludes the proof.

Remark A.4 . We note that although in our work the stepsise η k is chosen to be the same for all states for simplicity, the same result holds for a stepsize that is different for each state as long as η k ( s ) ≥ 1 c k D ϕ ( π s , π k s ) for each s ∈ S . The same is true for all results depending on this lemma (or the inexact variant).

A.2 Proof of Theorem 4.1

As a warmup, we start by reporting a convergence result for the standard PMD algorithm, recently established in [20]. Using the contraction properties of the Bellman operators, it can be shown that the suboptimality gap ∥ V π k -V ∗ ∥ ∞ for PI iterates converges to zero at a linear rate with a contraction factor γ , regardless of the underlying MDP. This instance-independent convergence rate was generalized to PMD in [20]. Before stating this result, we report a useful three-point descent lemma to quantify the improvement of the PMD policy resulting from the proximal step compared to an arbitrary policy.

Lemma A.5 (Three Point Descent Lemma) . Using the PMD update rule, we have for every k ∈ N and every π ∈ Π ,

This lemma results from an application of Lemma 6 in [53] which is a variant of Lemma 3.2 in [10]. This same result was also used in [20].

Theorem A.6 (Exact PMD, [20]) . Let ( c k ) be a sequence of positive reals and let the stepsize η k in (PMD) satisfy η k ≥ 1 c k ∥ min π ∈G ( V π k ) D ϕ ( π, π k ) ∥ ∞ where the minimum is computed componentwise. Initialized at π 0 ∈ ri(Π) , the iterates ( π k ) of (PMD) satisfy for every k ∈ N ,

We provide here a complete proof of this result as a warmup for our upcoming main result. Although equivalent, notice that our proof is more compact compared to the one provided in [20] (see section 6 and Lemma A.1 to A.3 therein) thanks to the use of our convenient notations, and further highlights the relationship between PMD and PI through the explicit use of Bellman operators.

Proof. Using Lemma A.5 with a policy π ∈ G ( V π k ) for which T π V π k = T V π k , we have

Setting d k := ∥ min π ∈G ( V π k ) D ϕ ( π, π k ) ∥ ∞ and recalling that D ϕ ( π, π k ) ≥ 0 for any policy π ∈ Π , we obtain

Notice that for Policy Iteration we have T V π k = T π k +1 V π k i.e. π k +1 ∈ G ( V π k ) . Instead, for PMD we have that π k +1 is approximately greedy (see (16)).

Applying Lemma A.5 again with π = π k , recalling that T π k V π k = V π k and that D ϕ ( π, π ′ ) ≥ 0 for any π, π ′ ∈ Π , we immediately obtain V π k ≤ T π k +1 V π k . Iterating this relation and using the monotonicity of the Bellman operator, we obtain V π k ≤ T π k +1 V π k ≤ ( T π k +1 ) 2 V π k ≤ · · · ≤ ( T π k +1 ) n V π k ≤ · · · ≤ V π k +1 where the last inequality follows from the fact that T π k +1 has a unique fixed point V π k +1 . We are now ready to control the suboptimality gap as follows:

where the first inequality follows from the inequality V π k +1 ≥ T π k +1 V π k proved above, the second one stems from (16), and the last one is a consequence of our assumption d k /η k ≤ c k . We conclude the proof by taking the max norm, applying the triangle inequality and using the contractiveness of the Bellman optimality operator T .

Proof of Theorem 4.1

Proof. Similarly to Lemma A.5, we show that for any π ∈ G h ( V π k ) ,

see Lemma A.2 for a proof. Setting ˜ d k := ∥ min π ∈G h ( V π k ) D ϕ ( π, π k ) ∥ ∞ (note the difference with d k in the proof of Theorem A.6), and recalling that π ∈ G h ( V π k ) implies that π ∈ G ( T h -1 V π k ) , we obtain from (17) an inequality similar to (16),

Then, we observe that

where the second inequality follows from using (17) and monotonicity of T π k +1 and the last inequality uses the assumption on the stepsizes. Then we show by recursion by iterating (19), taking the limit and using the fact that T π k +1 has a unique fixed point V π k +1 , that

see Lemma A.3 for a complete proof. We now control the suboptimality gap as follows,

where the first inequality stems from (20) and the last one from reusing (18). Similarly to Theorem A.6, we conclude the proof by taking the max norm, applying the triangle inequality and using the contractiveness of the Bellman optimality operator h times to obtain

A recursion gives the desired result.

A.3 Additional Comments About the Stepsizes in Theorem 4.1

About the cost of computing adaptive stepsizes. There is typically only one greedy policy at each iteration, and in the case that there is more than one it is possible to just pick any arbitrary greedy policy and the lower bound condition on the step size will still be satisfied. The cost of computing the step size is typically the cost of computing a Bregman divergence between two policies. We elaborate on this in the following using the discussion in [20] p. 7 in 'Computing the step size' paragraph which also applies to our setting. We use the step size condition η k ≥ 1 c k ∥ min π ∈ G h ( V π k ) D ϕ ( π, π k ) ∥ ∞ . We have a minimum in the condition which means we can just pick a single greedy policy ˜ π ∈ G h ( V π k ) (which can easily be computed given a lookahead value function in the exact setting or its estimate in an inexact setting). As for the maximum over the state space, this is because we use the same step size in all the states. This condition can readily be replaced by a state dependent step size η k ( s ) ≥ 1 c k D ϕ (˜ π ( ·| s ) , π k ( ·| s )) which is enough for our results to hold.

About the condition on the stepsizes. First, notice that PMD has been analyzed using an increasing stepsize in the exact setting (see e.g. [25, 53, 29]) and this is natural as this corresponds to emulating Policy Iteration. In the present work, notice that we do not require the stepsize to go to infinity and the stepsize does not typically go to infinity under our condition. Recall that we only require the stepsize to satisfy the condition: η k ≥ 1 c k ∥ min π ∈G h ( V π k ) D ϕ ( π, π k ) ∥ ∞ . As the policy gets closer to the optimal one it should also get closer to a greedy policy (the optimal policy is greedy with respect to its value function) and the term containing a Bregman divergence will vanish (since the optimal policy is already greedy with respect to its own value function). This effect balances the decreasing c k . Notice that we typically choose c k = γ 2 h ( k +1) .

We refer the reader to section D.4 for additional observations regarding stepsizes in experiments.

B Additional Details and Proofs for Section 5: Inexact h -PMD

B.1 Details and Precise Notation for Lookahead Computation

We describe here our approach to estimate Q π h in more detail and introduce the notation necessary for the analysis in the following sections. Assume that we are estimating Q π h ( s, a ) at a given fixed state-action pair ( s, a ) ∈ S × A . We define a sequence ( S k ) of sets of visited states as follows: Set S h = { s } and define for every k ∈ [0 , h -1] , for every ˜ s ∈ S k +1 and for ever ˜ a ∈ A the multiset S k (˜ s, ˜ a ) := { s ′ j ∼ P ( ·| ˜ s, ˜ a ) : j ∈ [ M k +1 ] } (including duplicated states) where M k +1 is the number of states sampled from each state ˜ s ∈ S k +1 , and S k = ⋃ ˜ s ∈S , ˜ a ∈A S k (˜ s, ˜ a ) ∪ S k +1 (without duplicated states). The estimation procedure is recursive. First we define the value function ˆ V 1 by ˆ V 1 (˜ s ) := ˆ V π (˜ s ) , ∀ ˜ s ∈ S 0 , where ˆ V π (˜ s ) is an estimate of V π (˜ s ) computed using M 0 rollouts of the policy π from state ˜ s . More precisely, we sample M 0 trajectories ( τ j (˜ s )) j ∈ [ M 0 ] of length H ≥ 1 rolling out policy π , i.e. τ j (˜ s ) := { ( s j k , a j k , r ( s j k , a j k )) } 0 ≤ k ≤ H -1 with s j 0 = ˜ s and a j k ∼ π ( ·| s j k ) . Then the Monte Carlo estimator is computed as ˆ V π (˜ s ) = 1 M 0 ∑ M 0 j =1 ∑ H -1 k =0 γ k r ( s j k , a j k ) . For each stage k ∈ [1 , h ] for h ≥ 1 , we define the following estimates for every ˜ s ∈ S k , ˜ a ∈ A ,

The desired estimate ˆ Q π h ( s, a ) is obtained with k = h . The recursive procedure described above can be represented as a (partial incomplete) tree where the root node is the state-action pair ( s, a ) and each sampled action leads to a new state node s ′ followed by a new state-action node ( s ′ , a ′ ) . To trigger the recursion, we assign estimated values ˆ V π (˜ s ) of the true value function V π (˜ s ) to each one of the leaf state nodes ˜ s of this tree, i.e. the state level h -1 of this tree which corresponds to the last states visited in the sampled trajectories. The desired estimate ˆ Q π h ( s, a ) is given by the estimate obtained at the root state-action node ( s, a ) . Recalling that Q π h ( s, a ) := r ( s, a ) + γ ( PV π h )( s, a ) for any s ∈ S , a ∈ A and V π h := T h -1 V π , the overall estimation procedure consists in using Monte Carlo rollouts to estimate the value function V π before approximating V π h := T h -1 V π using approximate Bellman operator steps and finally computing the estimate ˆ Q π h ( s, a ) using a sampled version of Q π h ( s, a ) := r ( s, a ) + γ ( PV π h )( s, a ) .

B.2 Proof of Theorem 5.3: Inexact h -PMD

Throughout this section, we suppose that the assumptions of Theorem 5.3 hold.

Lemma B.1. Let π ∈ Π , k ∈ N . Suppose that the estimator ˆ Q π k h is such that ∥ ˆ Q π k h -Q π h ∥ ∞ ≤ b for some positive scalar b . Then the approximate h -step lookahead value function ˆ V π k h ( π ) := M π ˆ Q π k h is a b -approximation of T π T h -1 V π k , i.e.,

Proof. The H¨ older inequality implies that for any policy π ∈ Π :

Lemma B.2 (Approximate h -Step Greedy Three-Point Descent Lemma) . For any policy π ∈ Π , the iterates π k of approximate h -PMD satisfy the following inequality for every k ≥ 0 :

Proof. We apply again the Three-Point Descent Lemma (see e.g. Lemma A.1 in [20]) to obtain:

which by rearranging and expressing in vector form gives:

Lemma B.1 implies that ˆ V π k h ( π ) ≥ T π T h -1 V π k -b e and ˆ V π k h ( π k +1 ) ≤ T π k +1 T h -1 V π k + b e yielding the lemma.

The following lemma is an analogue of (18) in the proof of Theorem 4.1 above, modified for the inexact case.

Lemma B.3. Assuming access at each iteration to an approximate Q h value function ˆ Q π k h with ∥ ˆ Q π k h -Q π h ∥ ∞ ≤ b , the iterates of inexact h -PMD satisfy for every k ≥ 0 :

Proof. Let ¯ π be any true h -step greedy policy ¯ π ∈ G h ( V π k ) , and let ˜ π be any approximate h -step greedy policy ˜ π ∈ argmax π ∈ Π ˆ V π k h ( π ) . By definition of ˜ π and equation (26) we have:

Finally we have the desired result using Lemma B.1, recalling the argument from the end of Lemma B.2 that ˆ V π k h (¯ π ) ≥ T ¯ π T h -1 V π k -b e = T h V π k -b e and ˆ V π k h ( π k +1 ) ≤ T π k +1 T h -1 V π k + b e .

Lemma B.4. Given a ˆ Q π k h with ∥ ˆ Q π k h -Q π h ∥ ∞ ≤ b , let ˜ G h = argmax π ∈ Π ˆ V π k h ( π ) denote the set of approximate h-step greedy policies and ˜ d k = ∥ min ˜ π ∈ ˜ G h D (˜ π, π k ) ∥ ∞ . Then the iterates of inexact h-PMD satisfy the following bound:

Proof. Analogously to Lemma A.3 we proceed by induction using Lemma B.3. We write for every k ≥ 0 ,

where (a) follows from the monotonicity of both Bellman operators T π k +1 , T and the fact that V π k ≤ T V π k , (b) uses Lemma B.3 and monotonicity of the Bellman operators. The last inequality stems from the fact that for every V ∈ R S , lim n → + ∞ ( T π k +1 ) n V = V π k +1 since V π k +1 is the unique fixed point of the Bellman operator T π k +1 which is a γ -contraction.

End of the proof of Theorem 5.3. We are now ready to conclude the proof of the theorem. Letting ˜ d k = ∥ min ˜ π ∈ ˜ G h D (˜ π, π k ) ∥ ∞ as in the theorem, we have:

where (i) follows from B.3 and (ii) from Lemma B.4. By the triangle inequality and the contraction property of T we obtain a recursion similar to the previous theorems:

Unfolding this recursive inequality concludes the proof.

B.3 Proof of Theorem 5.4

We first recall the notations used in our lookahead value function estimation procedure before providing the proof. The procedure we use is a standard online planning method and our proof follows similar steps to the proof provided for example in the lecture notes [48]. Notice though that we are not approximating the optimal Q function but rather the lookahead value function Q π h induced by a given policy π ∈ Π at any state action pair ( s, a ) ∈ S × A .

Recall that for any k ∈ [2 , h -1] , ˆ V k and ˆ Q k are defined as follows:

At the bottom of the recursion, recall that ˆ V 1 is defined for every ˜ s ∈ S 0 by

using M 0 trajectories ( τ j (˜ s )) j ∈ [ M 0 ] of length H ≥ 1 rolling out policy π , i.e. τ j (˜ s ) := { ( s j k , a j k , r ( s j k , a j k )) } 0 ≤ k ≤ H -1 with s j 0 = ˜ s and a j k ∼ π ( ·| s j k ) .

Note that the cardinality of each set S k is evidently upper bounded by S but also by A h -k ∏ i ∈ [ k +1 ,h ] M i as the number of states in each set grows by a factor of at most AM k at step k . In this section we prove some concentration bounds on the above estimators, which will be useful for the rest of the analysis.

Lemma B.5 (Concentration Bounds of Lookahead Value Function Estimators) . Let δ ∈ (0 , 1) . For any state ˜ s ∈ S 0 , with probability 1 -δ we have:

For any k ∈ [1 , h ] , for state ˜ s ∈ S k and for any action ˜ a ∈ A , with probability 1 -δ we have:

Finally, for any k ∈ [1 , h -1] , for any ˜ s ∈ S k , also with probability 1 -δ we have:

Proof. The proof of this result is standard and relies on Hoeffding's inequality. We provide a proof for completeness. Let Pr be the probability measure induced by the interaction of the planner and the MDP simulator on an adequate probability space. We denote by E the corresponding expectation.

First we show the bound for the value function estimate ˆ V 1 . This estimate is biased due to the truncation of the rollouts at a fixed horizon H , which is the source of the first error term in the bound. For any state ˜ s ∈ S 0 , the bias in ˆ V 1 (˜ s ) is given by

The second term is due to Hoeffding's inequality: each trajectory is i.i.d, and the estimator is bounded with probability 1 by 1 1 -γ . Therefore ˆ V 1 (˜ s ) concentrates around its mean, and with probability 1 -δ we have:

Now we prove the bound for the ˆ Q k function. Assume k ∈ [1 , h ] is arbitrary, and ˜ s ∈ S k . For any action ˜ a ∈ A , we can bound the error of ˆ Q k as follows:

Finally we bound the ˆ V k +1 function for k > 1 . Assume k ∈ [1 , h -1] is arbitrary, and ˜ s ∈ S k . With probability at least 1 -δ we have:

To show this last result, we use a standard union bound. For every t ≥ 0 , we have:

B.3.1 Estimation Error Bound for ˆ Q h

We start by defining some useful Bellman operators. The Bellman operator J : R S → R S × A and its sampled version ˆ J : R S → R S × A are defined for every V ∈ R S and any ( s, a ) ∈ S × A by

where C ( s, a ) is a set of m states sampled i.i.d from the distribution P ( ·| s, a ) .

Lemma B.6. The operators T and ˆ T satisfy for every U, V ∈ R S , s ∈ S , a ∈ A ,

Proof. The proof is immediate and is hence omitted.

We now show how to bound the error of ˆ Q k for any k ∈ [2 , h ] . We can bound the total error on ˆ Q k by decomposing the error into parts, accounting individually for the contribution of each stage in the estimation procedure. We will make use of the operators ˆ J as defined above and we will index them by the sets S k ( k ∈ [1 , h ] ) for clarity. We use the notation ∥ · ∥ X for the infinity norm over any finite set X . Recall that Q π k = r + γP T k -1 V π = J T k -1 V π for every k ∈ [1 , h ] . The following decomposition holds for any k ∈ [2 , h ] using the aforementioned contraction properties:

which, after defining the terms ∆ k = max a ∈A ∥ ˆ Q k ( · , a ) -Q π k ( · , a ) ∥ S k and and unrolling this recursion leads to the following bound on ∆ h :

We begin by bounding ∆ 1 by partially applying the recursion above:

which from Lemma B.5 we know to be bounded with probability at most 1 -δ |S 0 | by:

Finally, each term E j for j ≤ k can be bounded using the following lemma below. The proof relies on a union bound over the states reached during the construction of the search tree, see the lemma towards the end of [48] for a detailed explanation and a proof.

Lemma B.7. Let n ( j ) = |S j | for j ∈ [ k ] and any k ∈ [ h ] , let S 1 be the root state of the tree search and for each 1 < i ≤ n ( j ) let S i ∈ S j be the state reached at the i -1 th simulator call such that S j = { S i } n ( j ) i =1 . For 0 < δ ≤ 1 with probability 1 -An ( j ) δ for any 1 ≤ i ≤ n ( j ) we have:

Combining this result with (50) and the unrolled recursion gives the final concentration bound with probability at least 1 -δ V -δ J for δ V , δ J ∈ (0 , 1) ,

by a union bound over j ≤ k , assuming M j = M for all 1 ≤ j ≤ h -1 .

B.3.2 End of the Proof of Theorem 5.4

For any state-action pair ( s, a ) ∈ S × A , define the set |S h | = { s } and the sets |S k | accordingly for 0 ≤ k ≤ h -1 . Here we set M t := M > 0 , ∀ t ∈ [1 , h ] . Assuming the worst case growth of | S k | ≤ | S k +1 | AM k +1 , we can bound |S 0 | deterministically as |S 0 | ≤ A h M h . Using the bound in the previous section, it is possible to bound the error of ˆ Q h ( s, a ) as follows:

with probability at least 1 -δ V -δ J .

Assume now we want to estimate the value function ˆ Q ( s, a ) for all state action pairs ( s, a ) in some set C ⊆ S × A : we want to achieve max ( s,a ) ∈C | ˆ Q ( s, a ) -Q ( s, a ) | ≤ b for some positive accuracy b with probability at least 1 -δ V -δ J . This can be achieved by choosing the following parameters:

Picking C = S × A gives ∥ ˆ Q h -Q π h ∥ ∞ ≤ b . If estimates for { V k } h k =1 are reused across state action pairs, and the components of each value function are only computed at most once for every state s ∈ S , then the overall number of samples is of the order of KM 0 H |S 0 | + Kh |S 0 | · |C| M . In particular, with C = S × A , the number of samples becomes of the order of KM 0 HS + KhS · AM where we reuse estimators (which are in this case computed for each state-action pair ( s, a ) ∈ S × A . Let K > 1 h (1 -γ ) log( 4 ϵ (1 -γ )(1 -γ h ) ) as in Theorem 5.4. Choosing δ V = δ J = δ/ 2 for some δ ∈ (0 , 1) ϵ (1 -γ )(1 -γ h )

and b = 4 yields the following bound with probability at least 1 -δ (by Theorem 5.3):

To conclude, note that when k = K , we have:

C Additional Details and Proofs for Section 6: h -PMD with Function Approximation

To incorporate function approximation to h -PMD, the overall approach is as follows: compute an estimate of the lookahead value function by fitting a linear estimator using least squares, then perform an update step as in inexact PMD, taking advantage of the same error propagation analysis which is

agnostic to the source of errors in the value function. However a problem that arises is in storing the new policy, since a classic tabular storage method would require O ( S ) memory, which we would like to avoid to scale to large state spaces. A solution to this problem will be presented later on in this section.

Since we are using lookahead in our h -PMD algorithm, we modify the targets in our least squares regression procedure in order to approximate Q π h rather than Q π as in the standard PMD algorithm (with h = 1 ).

C.1 Inexact h -PMD with Linear Function Approximation

Here we provide a detailed description of the h -PMD algorithm using function approximation to estimate the Q h -function. Firstly we consider the lookahead depth h as a fixed input to the algorithm. We assume that the underlying MDP has finite state and action spaces S and A , with transition kernel P : S × A → ∆( A ) , reward vector R ∈ [0 , 1] S × A and discount factor γ . We also assume we are given a feature map ψ : S × A → R d , also written as Ψ ∈ R SA × d .

We adapt the inexact h -PMD algorithm from section 5 as follows:

Take as input any policy π 0 ∈ rint (Π) , for example one that is uniform over actions in all states.

Initialize our memory containing π 0 , and an empty list Θ which will hold the parameters of our value function estimates computed at each iteration.

At each iteration k = 0 , . . . , K , perform the following steps:

(a) Compute the targets ˆ R ( z ) := ˆ Q h ( z ) for all z ∈ C using the procedure described in section 5.1, using π k as a policy.

(b) Compute θ k := argmin θ ∈ R d ∑ z ∈C ρ ( z )( ψ ( z ) ⊤ θ -ˆ R ( z )) 2 and append it to the list Θ .

(c) Using Ψ θ k as an estimate for Q π k h , compute π k +1 using the usual update from inexact h -PMD. This step can be merged with step 3a, the details are described below.

Computing the full policy π k for every state at every iteration would require Ω( S ) operations, although the number of samples to the environment still depends only on d . If this is not a limitation the policy update step can be performed as in inexact h -PMD simply by using Ψ θ k as an estimate for Q π k h , as described in step 3c. Otherwise, we describe a method below to compute π k on demand when it is needed in the computation of the targets in step 3a, eliminating step 3c entirely.

The policy π 0 is already given for all states.

For any k > 0 given π k we can compute π k +1 at any state s using the following:

(a) Compute any greedy policy ˜ π ( ·| s ) using Ψ θ k , for example ˜ π ( a ′ | s ) ∝ ✶ { a ′ ∈ argmax a ∈A { ⟨ ψ ( s, a ) ⊤ θ k } } (b) Compute η k := 1 γ 2 h ( k +1) D ϕ (˜ π ( ·| s ) , π k ( ·| s ))

(c) Finally, compute π k +1 ( ·| s ) ∈ argmax π ∈ Π { η k ⟨ (Ψ θ k ) s , π ( ·| s ) ⟩ -D ϕ ( π ( ·| s ) , π k ( ·| s )) }

Remark C.1 . This procedure can be used to compute π k +1 ( ·| s ) 'on demand' for any k > 0 and s ∈ S by recursively computing π k ( ·| s ′ ) at any states s ′ required for the computation of Ψ θ k (since π 0 is given). The memory cost requirement for this method is of the order of O ( KA min {|S 0 | H,S } ) in total which is always less than computing the full tabular policy at each iteration.

Remark C.2 . Above, (Ψ θ k ) s denotes a vector in R A that can be computed in O ( Ad ) operations: each component in a is equal to ψ ( s, a ) ⊤ θ k which can be computed in O ( d ) operations. Also, the stepsize η k is chosen to be equal to its lower bound in the analysis of 5.3, with the value of c k chosen to be equal to γ 2 h ( k +1) , as it is needed for linear convergence. We also note that by memorizing policies as they are computed, we manage to perform just as many policy updates as are needed by the algorithm.

C.2 Discussion of Assumption 6.2

We make two preliminary comments before discussing the relevance of our assumption and alternatives to relax it below:

First, notice that we do not require Q π h to be represented as a linear function. Our assumption is an approximate action value function approximation.

When the lookahead depth is h = 1 , we recover the standard PMD algorithm. Notice that in this case our assumption is standard in the RL literature for linear function approximation.

Now, using our notations, notice that for any policy π ∈ Π , Q π h = r + γP T h -1 V π = r + γP T h -1 M π Q π = J T h -1 M π Q π . We have two alternatives:

Directly approximate Q π h as we do in our assumption. In that case the targets used in our regression procedure are estimates ˆ Q h of Q π h . Note also that this approach is meaningful from a practical point as we may directly use a neural network to approximate the lookahead Q-function which is used in our h -PMD algorithm rather than approximating the Q-function first and then trying to approximate from it the lookahead Q-function Q π h .

Approximate Q π with linear function approximation and use our propagation analysis (see section 5.1 and proof in appendix B.3.1, B.3.2) based on the link between Q π h and Q π highlighted in the formula above. This approach which propagates the function approximation error on Q π to Q π h constitutes an alternative to Assumption 6.2 on Q π h by an assumption on Q π . We discuss it further in the rest of this section.

In the following exposition we provide more details regarding estimating Q π using linear function approximation and relaxing assumption 6.2 using the second approach discussed above.

Let Assumption 6.1 hold (as in Theorem 6.3). Instead of assumption 6.2, we introduce the following assumption on the Q -function Q π rather than the h -step lookahead Q-function Q π h .

We modify the h -PMD with function approximation algorithm as follows: instead of directly approximating Q π h using linear function approximation, we first use linear function approximation to approximate Q π from which we construct an estimate of Q π h similarly to our method for inexact h -PMD. This requires us to replace the regression targets with simpler ones, which are Monte Carlo estimates of Q π instead of Q π h . Specifically, these new targets are defined as:

where ( s ( i ) t , a ( i ) t ) are state action pairs sampled according to the policy π , z = ( s, a ) ∈ C ⊆ S × A and ( s ( i ) 0 , a ( i ) 0 ) = z .

For any policy π we can use the same function approximation procedure as described in our paper (see Appendix C.1) to yield a ˆ θ ∈ R d with ∥ Q π -Ψ ˆ θ ∥ ∞ ≤ ϵ 0 , (for some value of ϵ 0 to be defined below) by simply plugging in h = 1 . Note that this particular case was previously analysed in [48] (Lecture 8). We reuse an intermediate result therein (equation (7)) which guarantees the bound:

with probability at least 1 -δ for some 0 < δ < 1 .

Now we reuse most of our estimation procedure in Section 5.1, simply by changing the definition of ˆ V 1 to ˆ V 1 ( s ) := ∑ a ∈A π ( a | s )(Ψ ˆ θ ) s,a for all s ∈ S 0 . This evidently retains the guarantee that | ˆ V 1 ( s ) -V π ( s ) | ≤ ϵ 0 ( δ ) with probability at least 1 -δ for each s ∈ S 0 . Finally, using the same error propagation analysis as in Appendix B.3.1, replacing C V 1 ( δ V 1 |S 0 | ) with ϵ 0 ( δ V 1 ) we obtain the following

bound for any individual state action pair ( s, a ) :

with probability at least 1 -δ V 1 -δ V t -δ Q h .

Finally, the update rule as described in Appendix C.1 is replaced by π k +1 ( ·| s ) ∈ argmax π s ∈ ∆( A ) ( η k ⟨ ˆ Q π k h ( s, · ) , π s ⟩ -D ϕ ( π s , π k s )) . Note that we only compute π k in the states where we need it to compute π k +1 as described in Appendix C.1.

Unfortunately, it is not clear how to directly obtain a result similar to Theorem 6.3 from this approach (controlling ∥ V ⋆ -V π k ∥ ∞ at each iteration, i.e. uniformly over all states) beyond the bound that we provide above. A naive union bound over the state action space would result in a bound with a logarithmic dependence on the size of the state space, unlike Theorem 6.3 which results in a bound that is completely independent of the size of the state space.

We would like to add an additional comment regarding possible ways to relax Assumption 6.2 even for the standard case where h = 1 . Recently, Mei et al. [30] proposed an interesting alternative approach relying on order-based conditions (instead of approximation error based ones) to analyze policy gradient methods with linear function approximation. We would like to highlight though that their work only applies to the bandit setting to the best of our knowledge. The extension to MDPs is far from being straightforward and would require some specific investigation. It would be interesting to see if one can relax Assumption 6.2 using similar ideas.

C.3 Convergence Analysis of Inexact h -PMD using Function Approximation

Our analysis follows the same lines as the proof in [27] under our assumptions from section 6.

We compute ˆ Q h ( z ) for a subset C of state action pairs using the vanilla Monte Carlo Planning algorithm presented in section 5. These will be our targets for estimating the value function using least squares. In order to control the extrapolation of these estimates to the full Q π h function, we will need the following theorem.

Theorem C.4 ([22]) . Assume Z is finite and ψ : Z → R d is such that the underlying matrix Ψ ∈ R |Z|× d is of rank d . Then there exist C ⊆ Z and a distribution ρ : C → [0 , 1] (i.e. ∑ z ∈C ρ ( z ) = 1 ) satisfying the following conditions:

|C| ≤ d ( d +1) / 2 ,

Let C ⊆ S × A , ρ : C → [0 , 1] be as described in the Kiefer-Wolfowitz theorem. Define ˆ θ as follows:

Assuming that G ρ is nonsingular, a closed form solution for the above least squares problem is given by ˆ θ = G -1 ρ ∑ z ∈C ρ ( z ) ˆ R m ( z ) ϕ ( z ) . This will be used to estimate our lookahead value function as Ψ ˆ θ . We can control the extrapolation error of this estimate as follows:

Theorem C.5. Let δ ∈ (0 , 1) and let θ ∈ R d such that ∥ Q π h -Φ θ ∥ ∞ ≤ ϵ π . Assume that we estimate the targets ˆ Q h using the Monte Carlo planning procedure described in section 5.1 with parameters, M and M 0 . There exists a subset C ⊆ Z and a distribution ρ : C → R ≥ 0 for which, with probability at least 1 -δ , the least squares estimator ˆ θ satisfies:

where

Proof. The proof of this theorem is identical to the one in [27], up to using our new targets and error propagation analysis instead of the typical function approximation targets for Q -function estimation (without lookahead). The subset C ⊆ S × A and distribution ρ : C → [0 , 1] are given by the Kiefer-Wolfowitz theorem stated above. We suppose we have access to such a set. Note that this assumption is key to controlling the extrapolation error of a least squares estimate Ψ ˆ θ .

We have now fully described how to estimate the lookahead value function Q π h using a number of trajectories that does not depend on the size of the state space. In the previous section, we described how to use this estimate in an h -PMD style update, while maintaining a memory size that does not depend on S . Therefore we have an h -PMD algorithm that does not use memory, computation or number of samples depending on S .

D Further Simulations

D.1 Simulations for h -PMD with Euclidean Regularization

Here we provide details of an additional experiment running h -PMD with Euclidean regularization. As in section 7, experiments were run on the bsuite DeepSea environment [34] with γ = 0 . 99 and a grid size of 64x64. The figures are identical in structure to the ones in 7, 32 runs were performed for each value of h in 1, 5, 10, 15, and 20. In Fig. 3, we report the number of samples to the environment at each iteration for the same runs. When compared to the figures showing the convergence, it is clearly seen that running the algorithm with higher values of lookahead results in convergence using fewer overall samples.

0

10

20

30

40

Iteration

0

.

0

0

.

1

0

.

2

0

.

3

0

.

4

0

.

5

‖

V /star

-

V π

k

‖

0

20

40

60

Iteration

0

.

0

0

.

2

0

.

4

0

.

6

h

=1

h

=5

h

=10

h

=15

h

=20

D.2 Sample Complexity Comparison for Different Lookahead Depths

We have recorded the sample complexity of our algorithm for different values of h in Figure 3. Notice that the number of samples used at each iteration increases with h . However, notice that this increase is small: h -PMD with h = 20 requires only 1.5 times more samples per iteration than h = 1 as shown in Figure 3. Furthermore, increasing h greatly improves the convergence rate of the algorithm as described by our theory, which results in a much smaller number of iterations required until convergence (for example h = 20 converges in at least 10 times fewer iterations than h = 1 (see figure 1 (right) in section 7). Overall, the algorithm is more sample efficient for higher values of h .

D.3 Total Running Time Comparison for Different Lookahead Depths

We provide a plot accounting for running time instead of the number of iterations (see Figure 4). Observe that the overall performance is very similar to the same plot with the number of iterations on the x-axis.

0

20

40

60

80

Iteration

0

2

4

6

8

×

10

7

h

=1

h

=5

h

=10

h

=15

h

=20

0

2

4

6

8

10

Time (s)

0

.

0

0

.

1

0

.

2

0

.

3

0

.

4

0

.

5

‖

V /star

-

V π

k

‖∞

0

10

20

30

40

50

60

70

80

Iterations

0

.

0

0

.

1

0

.

2

0

.

3

0

.

4

0

.

5

‖

V /star

-

V π

k

‖∞

h

=1

h

=5

h

=10

h

=15

h

=20

D.4 About the Stepsizes in Simulations

In our experiments we use the exact same stepsizes described in our theory. We also observe in our simulations that the stepsize does not go to infinity but rather decreases at some point and/or stabilizes. This is a consequence of the adaptivity of the stepsize discussed above. We also note that the simulations reported in [20] (appendix D p. 18) also validate the fact that the stepsizes do not go to infinity (see the right plots therein for the stepsizes). We also observe a similar behavior for the stepsizes in our experiments (see Figure 5).

D.5 Larger lookahead depths

In order to practically investigate the tradeoff of increasing the lookahead depth h , we performed additional experiments where h was increased to 100 in the DeepSea bsuite environment. The plots in Figure 6 show that increasing h past the values in our initial experiments still results in better performance.

0

25

50

75

100

125

150

175

200

Iterations

10

-

2

10

-

1

10

0

10

1

10

2

10

3

ηk

h

=1

h

=5

h

=10

h

=15

h

=20

0

20

40

60

80

100

Iterations

0

.

0

0

.

1

0

.

2

0

.

3

0

.

4

0

.

5

‖

V /star

-

V π

k

‖∞

0

2

4

6

8

10

Time (s)

0

.

0

0

.

1

0

.

2

0

.

3

0

.

4

0

.

5

‖

V /star

-

V π

k

‖∞

h

=1

h

=5

h

=10

h

=15

h

=20

h

=100

D.6 Experiments with MCTS-based Lookahead Estimation

We have performed additional experiments in an even more challenging and practical setting, on two environments from the bsuite [34] set: DeepSea and Catch. We have implemented our PMD algorithm with a Monte Carlo Tree Search (MCTS) algorithm (which we call PMD-MCTS) to compute the lookahead Q-function. We used DeepMind's MCTS implementation (MCTX) 4 in order to run this algorithm on any deterministic gym style environments implemented in JAX (we use gymnax [26] to implement these environments). Note that any tree search algorithm implementation can be used to estimate the lookahead function and plugged into our PMD method. See figures below. Note that this training procedure is more practical: we do not evaluate the value function at each state-action pair anymore at each time step, we only evaluate the value function (and so only update the policy) in states that we encounter using the current policy. This significantly relaxes the generative model assumption. Note that in this more challenging setting the best performance is not obtained for higher values of h : intermediate values of h perform better, illustrating the tradeoff in choosing the depth h .

0

100

200

300

400

500

Time (s)

0

.

0

0

.

5

1

.

0

1

.

5

‖

V /star

-

V π

k

‖∞

0

200

400

600

800

1000

Iterations

10

-

2

10

-

1

10

0

‖

V /star

-

V π

k

‖∞

h

=1

h

=2

h

=3

h

=4

h

=5

0

500

1000

1500

2000

2500

3000

3500

Time (s)

0

.

0

0

.

2

0

.

4

0

.

6

‖

V /star

-

V π

k

‖∞

0

50

100

150

200

250

300

350

400

Iterations

0

.

0

0

.

2

0

.

4

0

.

6

‖

V /star

-

V π

k

‖∞

h

=1

h

=2

h

=3

h

=4

h

=5

D.7 Continuous Control Experiments

Finally, we performed experiments in classic continuous control environments in order to investigate the performance of h -PMD when applied to MDPs with large/infinite state spaces. For this continuous state space setting, we implemented a fully online variant of h -PMD based off of DeepMind's MuZero [40] algorithm. Experiments show that increasing the lookahead depth can lead to better performance up to a certain point, and the tradeoff seems to be environment-specific.

0

25000

50000

75000

100000

125000

150000

175000

200000

Steps

0

100

200

300

400

500

600

Reward

CartPole-v1

0

100000

200000

300000

400000

500000

Steps

500

400

300

200

100

Reward

Acrobot-v1

h

=1

h

=2

h

=3

h

=4

h

=5

D.8 Additional experiments

We conclude with results from a larger set of experiments designed to verify the robustness of our experimental results to different environment parameters. We report the performance of the h -PMD algorithm in the inexact setting in two different environments, DeepSea from bsuite [34], and the

classic Gridworld environment. The effect of γ is as expected: higher values of γ result in slower convergence, lower values of γ yield faster convergence. The size of the environment ('chain length' for DeepSea and number of columns/rows for Gridworld) directly affects the number of iterations until convergence predictably: larger environments take more iterations to solve. Note that in all cases the effect of h is as expected, i.e. higher h leads to faster convergence.

0

20

40

60

80

100

Iterations

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

‖

V /star

-

V π

k

‖∞

DeepSea: chain length 32,

γ

= 0

.

99

h

=1

h

=5

h

=10

h

=15

h

=20

0

20

40

60

80

100

Iterations

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

‖

V /star

-

V π

k

‖∞

DeepSea: chain length 16,

γ

= 0

.

99

h

=1

h

=5

h

=10

h

=15

h

=20

0

20

40

60

80

100

Iterations

-

0

.

5

0

.

0

0

.

5

1

.

0

1

.

5

‖

V /star

-

V π

k

‖∞

DeepSea: chain length 32,

γ

= 0

.

9

h

=1

h

=5

h

=10

h

=15

h

=20

0

20

40

60

80

100

Iterations

0

.

0

0

.

1

0

.

2

0

.

3

‖

V /star

-

V π

k

‖∞

DeepSea: chain length 32,

γ

= 0

.

999

h

=1

h

=5

h

=10

h

=15

h

=20

0

5

10

15

20

25

30

Iterations

0

20

40

60

80

‖

V /star

-

V π

k

‖∞

Gridworld: 32 columns, 16 rows,

γ

= 0

.

99

h

=1

h

=5

h

=10

h

=15

h

=20

0

5

10

15

20

25

30

Iterations

0

20

40

60

‖

V /star

-

V π

k

‖∞

Gridworld: 16 columns, 8 rows,

γ

= 0

.

99

h

=1

h

=5

h

=10

h

=15

h

=20

0

5

10

15

20

25

30

Iterations

-

0

.

5

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

2

.

5

‖

V /star

-

V π

k

‖∞

Gridworld: 32 columns, 16 rows,

γ

= 0

.

9

h

=1

h

=5

h

=10

h

=15

h

=20

0

5

10

15

20

25

30

Iterations

0

200

400

600

‖

V /star

-

V π

k

‖∞

Gridworld: 32 columns, 16 rows,

γ

= 0

.

999

h

=1

h

=5

h

=10

h

=15

h

=20

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Please see the abstract, the introduction and the results.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Please see the future work discussion in the conclusion and the discussion about the computational tradeoff in Section 4. Assumptions for function approximation are discussed in details in the Appendix.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate 'Limitations' section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Please see the complete proofs in the different appendix sections.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Details provided in the main part in the simulation sections and in the appendix, we also provide a link to an anonymous repository for the code.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Please see link to an anonymous repository for the code.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: This is discussed in the description of simulations in the main part and in the appendix as well as in the code provided.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We provide error bars (shaded area) based on runnning our simulations using different seeds and specify it is computed using the standard deviation.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer 'Yes' if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We used CPUs and we also tracked running time (see relevant figures in the appendix). We performed several simulations for tuning which are not all reported in the paper. We reported several of them to show the robustness of our results to varying hyperparameters.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We checked the code and our paper conforms to it.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: Our contributions are mainly theoretical and focus on improving the performance of existing policy gradient methods.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to 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).

11. 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: Our work does not involve such risks.

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: [Yes]

Justification: We mention whenever we use specific existing code from some repositories and give credit appropriately.

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.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: does not apply to our paper.

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.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: does not apply to our work.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: does not apply to our work.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies policy mirror descent (PMD) where the policy improvement step is modified to include an action-value with $h$-step lookahead (contains $h-1$ applications of the Bellman optimality operator on the value of the policy from the previous iteration). In the exact tabular setting where the value functions can be computed exactly, this leads to an improvement of the linear convergence rate of standard PMD from $\\gamma$ to $\\gamma^h$. In the inexact tabular setting where the the lookahead Q-functions are estimated through Monte-Carlo rollouts, this gives a sample-complexity for finding a $\\varepsilon$-optimal policy of $\\tilde{O}(1/\\varepsilon^2(1-\\gamma)^7)$, improving by a factor of $1/(1-\\gamma)$ with respect to prior work. Finally, they extend their results beyond the tabular setting to a linear function approximation setting.\n\nSTRENGTHS:\n- The paper is quite well written and easy to follow.\n- The paper presents an interesting extension of PMD to include look-ahead leading to faster convergence. The algorithm and results cover different settings from exact tabular to inexact with linear function approximation. Numerical experiments support these claims both in terms of iterations and runtime.\n\nWEAKNESSES:\n- The Analysis of Theorem 4.1 closely follows that of [16] while the remaining results (Theorem 5.4 and Theorem 6.3) follow using analyses similar to prior works (except for the estimation of the lookahead Q-function). Nevertheless, I recognise the merit and novelty in combining these with the idea of lookahead from [9]. Note - it would be good to explicitly state where the proof of Theorem 6.3 can be found in the Appendix.\n- It seems the only disadvantage of $h$-PMD over $1$-PMD is in the computational cost of computing a more complex value function (lines 212-215) but it is unclear exactly what this cost is (other than what is conveyed in the experiments). It would be nice to explicitly quantify the additional computation (e.g. in the tabular setting, computing the true value function can be done by finding the fixed point of $\\mathcal{T}^\\pi$ which involves inverting a $S\\times S$ matrix, then how much extra computation is required to achieve the value function with lookahead ?). Note: in lines 194-195, the assumption is not only that the value function $V^{\\pi_k}$ can be computed exactly but also the value function with lookahead $V^{\\pi_k}_h$ right ?\n- The rationale behind how to choose $h$ is not clear to me (see also questions below). For example, in Theorem 5.4 (line 270), if the overall sample complexity is $\\tilde{O} (\\frac{S}{h\\varepsilon^2(1-\\gamma)^6(1-\\gamma^h)^2} + \\frac{SA}{\\varepsilon^2 (1-\\gamma)^7})$, this seems to suggest that we should take $h \\rightarrow \\infty$. Besides the computational infeasibility of this, if the lookahead Q-functions are estimated using Algorithm 1, then taking $h \\rightarrow \\infty$ should also result in a sample complexity going to $\\infty$, it seems odd that this does not appear in the sample complexity. It would be nice to have an explanation of this.\n\n\nTypos:\n- Line 128: you are missing a “the” in front of squared Euclidean distance.\n- Line 564-565, in equation (12), $\\mathcal{T}^{\\pi_{k+1}}$ should be to the power of $n+2$ not squared.\n- Line 344: succesfully -> successfully\n\nQUESTIONS:\nThis questions are related to choosing $h$ as discussed in the weaknesses.\n- It seems that one step of $h$-PI is equivalent to $h$ steps of value iteration with a policy produced at the end by acting greedily with respect to the last value. Is this the case ? If so, is there a way to choose a good value of $h$ that balances the benefits of VI and PI/PMD ?\n- In the simulations, it seems like taking $h$ bigger provides better convergence without the cost of longer runtime. Is this the case for all values of $h$ ? If not, it would be nice to see this reflected in the experiments (i.e. taking much larger values of $h$ which begin to impede on the runtime). If it is true for all values of $h$, is it essentially saying that value iteration is computationally more efficient than PI/PMD in the setting you consider ? And then it would be nice to consider settings where this perhaps is not the case ?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-23"}

comments: SUMMARY:
The authors proposed a version of the policy mirror descent algorithm that uses a multi-step greedy policy improvement operator. Afterward, the authors showed the theoretical benefits of the proposed method by a better contraction rate $\gamma^h$ instead of $\gamma$-contraction for a usual 1-step greedy policy improvement algorithm. Finally, the inexact version of the algorithm was proposed with an improved sample complexity in the finite MDP setting and, under a linear functional approximation setup, guarantees that are independent of the state space size.

STRENGTHS:
- The proposed combination of PMD and $h$-step greedy updates implies interesting theoretical results, such as reducing sample complexity with increasing the planning horizon in a very explicit way;
- Strong result with linear functional approximation that gives an implementable algorithm with a polynomial (in $h$) running time;

WEAKNESSES:
- The experimental part of the paper might be improved by running continuous control experiments.

QUESTIONS:
- I think the paper may benefit from a theoretical example of h-PI with greedy updates not convergent whereas h-PMD converges.

LIMITATIONS:
The paper is mostly of theoretical nature and thus does not imply any negative societal impact.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

comments: SUMMARY:
The author propose a multi-step greedy approach for Policy Mirror Descent. Combing PMD with multiple greedy policy updates results in a faster $\gamma^{h}$ rate improved the previously thought optimal $\gamma$ rate. Additionally, the authors extend their analysis to the stochastic setting and when using function approximation (with linear function approximation).

STRENGTHS:
The idea of combining lookahead to PMD is well-motivated and interesting.  The motivation and presentation of the results are explained well. All results also do not rely distribution mismatch coefficients, which is something that appears quite often in prior work

WEAKNESSES:
It seems that the analysis is done in the functional representation of the policy. It would be good to also discuss how the methods could extended for specific policy parameterization such as softmax policies.
Since the resulting bounds do not rely on concentrability, it would be nice to see experiments where the state space is much larger.

QUESTIONS:
How well does $h$-PMD perform against PI in the deterministic setting?
Could you provide some intuition on why distribution mismatch coefficients does not appear in the stochastic setting? It's surprising to me since there isn't  any explicit exploration.

LIMITATIONS:
No concerns.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces h-PMD, an extension of the Policy Mirror Descent (PMD) algorithm, which incorporates multi-step lookahead to improve policy updates in reinforcement learning. PMD is a general framework that includes several policy gradient methods and relates to advanced algorithms like TRPO and PPO. Recognizing that multi-step greedy policies often outperform their single-step counterparts, the authors propose h-PMD, which integrates multi-step lookahead depth into PMD. This new class of algorithms achieves a faster convergence rate for solving discounted infinite horizon MDPs, under both exact and inexact settings. The paper also extends these results to linear function approximation, demonstrating improved sample complexity that depends on the dimension of the feature map space rather than the state space size.\n\nSTRENGTHS:\n1. The paper introduces a novel extension to the PMD framework by incorporating multi-step lookahead, enhancing the algorithm's performance and convergence rate.\n\n2. The proposed h-PMD algorithm achieves a faster 𝛾^ℎ-linear convergence rate, which is an improvement over standard PMD and Policy Iteration methods.\n\n3. The paper addresses both exact and inexact settings, providing a sample complexity analysis that demonstrates improved performance over previous methods, especially with increasing lookahead depth.\n\n4. By extending h-PMD to linear function approximation, the authors make the algorithm applicable to large state spaces, which is crucial for practical applications in complex environments.\n\n5. The theoretical findings are supported by empirical results from simulations on the DeepSea RL environment, illustrating the benefits of the h-PMD approach\n\nWEAKNESSES:\n1. The h-PMD algorithm, while improving convergence rates, also introduces higher computational complexity due to the multi-step lookahead, which can be demanding for large-scale problems.\n\n2. The paper might lack detailed guidance on implementing h-PMD in various practical scenarios, making it challenging for practitioners to adopt and utilize the algorithm effectively.\n\n3. The results are contingent on certain assumptions, such as the availability of multi-step greedy policies and the use of a generative model. These assumptions might limit the generalizability of the findings to all RL problems.\n\n4. The benefits of the h-PMD algorithm are closely tied to the lookahead depth ℎ. Determining the optimal ℎ in practice could be non-trivial and might require extensive experimentation.\n\n5. While the paper demonstrates the theoretical and empirical advantages of h-PMD, it may not provide a comprehensive comparative analysis against a wide range of existing RL algorithms, which would strengthen the case for its superiority.\n\nQUESTIONS:\n1. Can you provide more detailed guidelines on how to implement the h-PMD algorithm in various practical settings, including both exact and inexact scenarios?\n\n2. How does the computational complexity of h-PMD compare to standard PMD and other state-of-the-art RL algorithms like TRPO and PPO in practice? Are there strategies to mitigate the increased computational burden?\n\n3. What methods or heuristics do you recommend for determining the optimal lookahead depth in different environments? How sensitive is the performance of h-PMD to the choice of ℎ?\n\n4. How critical is the assumption of a generative model for the theoretical guarantees provided in the paper? Can h-PMD be effectively applied in settings where a generative model is not available?\n\n5. How well does the h-PMD algorithm with linear function approximation perform in very high-dimensional or continuous state spaces? Have you explored other types of function approximation beyond linear?\n\n6. Can you provide more extensive empirical results across a variety of RL environments to demonstrate the robustness and general applicability of h-PMD?\n\n7. How does h-PMD compare to other advanced RL algorithms like AlphaZero, in terms of performance and computational efficiency? Have you conducted any comparative studies?\n\n8. What are the scalability limits of h-PMD in terms of state and action space sizes? Are there practical cases where h-PMD might not be feasible due to its complexity?\n\n9. Can you elaborate on the trade-offs involved in the inexact setting of h-PMD? How does the estimation of lookahead action values impact the overall performance and convergence rate?\n\n10. Have you applied h-PMD to any real-world RL problems or industrial applications? If so, what were the outcomes and challenges faced?\n\nLIMITATIONS:\n1. The multi-step lookahead in h-PMD increases the computational burden significantly, potentially making it impractical for real-time or resource-constrained applications.\n\n2. The performance of h-PMD is dependent on the chosen lookahead depth ℎ. Finding the optimal depth can be challenging and may require considerable computational resources for tuning.\n\n3. The sample complexity improvements are based on the assumption of a generative model, which may not always be available or practical in many real-world scenarios.\n\n4. While the extension to linear function approximation addresses scalability to large state spaces, the practical implementation and effectiveness of this approach in very high-dimensional or continuous spaces remain uncertain.\n\n5. The empirical validation is limited to the DeepSea RL environment. Broader validation across diverse and more complex environments would be necessary to confirm the general applicability and robustness of the proposed h-PMD algorithm.\n\n6. The paper might not provide sufficient practical implementation details, making it difficult for practitioners to apply the h-PMD algorithm to their specific problems without further guidance and experimentation.\n\n7. Some theoretical results assume exact policy evaluation, which may not be feasible in many practical settings where only approximate evaluations are possible.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}

comments: SUMMARY:
This paper applies the idea of lookahead in policy improvement procedures under the PMD setting. They prove the $\gamma^h$-linear convergence rate. They also propose the inexact version of h-PMD and extend to the function approximation case.

STRENGTHS:
The writing of the paper is pretty good. The main idea and the results of the paper are easy to follow. 

The paper is also complete. They propose the theoretical analysis in both tabular and function approximation settings.

WEAKNESSES:
Although the paper proved $\gamma^h$-contraction property, it's not out of expectation.

I'm still not completely convinced that $h$-horizon lookahead policy improvement could bring additional benefit in the function approximation case. Even in tabular cases, additional computational effort are required.





Minor problems:

Line 70: bsuite.

QUESTIONS:
- It's interesting to see that larger lookahead leads to faster training speed based on Figure 1. How about the sample complexities? Could you provide the used samples for each method?
- The meaning of the dotted lines should be clarified in Figure 1 (Left).
- Based on the theoretical result of the paper, does that mean larger $h$ leads to better performance? For Figure 1, what if we have a larger $h$?

LIMITATIONS:
The paper discusses some potential improvements in future work on the practical algorithm. However, they didn't discuss the limitations of the existing theoretical/practical results.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-16