Title: Operator World Models for Reinforcement Learning
Venue: NeurIPS 2024
Accepted: yes
Keywords: Reinforcement Learning

Transfer Operators

World Models

Policy Gradient

Conditional Mean Embeddings

Mirror Descent
Abstract: Policy Mirror Descent (PMD) is a powerful and theoretically sound methodology for sequential decision-making. However, it is not directly applicable to Reinforcement Learning (RL) due to the inaccessibility of explicit action-value functions. We address this challenge by introducing a novel approach based on learning a world model of the environment using conditional mean embeddings. Leveraging tools from operator theory we derive a closed-form expression of the action-value function in terms of the world model via simple matrix operations. Combining these estimators with PMD leads to POWR, a new RL algorithm for which we prove convergence rates to the global optimum. Preliminary experiments in finite and infinite state settings support the effectiveness of our method.

Operator World Models for Reinforcement Learning

Abstract

Policy Mirror Descent (PMD) is a powerful and theoretically sound methodology for sequential decision-making. However, it is not directly applicable to Reinforcement Learning (RL) due to the inaccessibility of explicit action-value functions. We address this challenge by introducing a novel approach based on learning a world model of the environment using conditional mean embeddings. Leveraging tools from operator theory we derive a closed-form expression of the action-value function in terms of the world model via simple matrix operations. Combining these estimators with PMD leads to POWR, a new RL algorithm for which we prove convergence rates to the global optimum. Preliminary experiments in finite and infinite state settings support the effectiveness of our method 1 .

1 Introduction

In recent years, Reinforcement Learning (RL) [1] has seen significant progress, with methods capable of tackling challenging applications such as robotic manipulation [2], playing Go [3] or Atari games [4] and resource management [5] to name but a few. The central challenge in RL settings is to balance the trade-off between exploration and exploitation, namely to improve upon previous policies while gathering sufficient information about the environment dynamics. Several strategies have been proposed to tackle this issue, such as Q-learning-based methods [4], policy optimization [6, 7] or actor-critics [8] to name a few. In contrast, when full information about the environment is available, sequential decision-making methods need only to focus on exploitation. Here, strategies such as policy improvement or policy iteration [9] have been thoroughly studied from both the algorithmic and theoretical standpoints. Within this context, the understanding of Policy Mirror Descent (PMD) methods has recently enjoyed a significant step forward, with results guaranteeing convergence to a global optimum with associated rates [10, 11, 12].

In their original formulation, PMD methods require explicit knowledge of the action-value functions for all policies generated during the optimization process. This is clearly inaccessible in RL applications. Recently, [12] showed how PMD convergence rates can be extended to settings in which inexact estimators of the action-value function are used (see [13] for a similar result from a regret-based perspective). The resulting convergence rates, however, depend on uniform norm bounds on the approximation error, usually guaranteed only under unrealistic and inefficient assumptions such as the availability of a (perfect) simulator to be queried on arbitrary state-action pairs. Moreover, these strategies require repeating this sampling/learning process for any policy generated by the PMD algorithm, which is computationally expensive and demands numerous interactions with the

environment. A natural question, therefore, is whether PMD approaches can be efficiently deployed in RL settings while enjoying the same strong theoretical guarantees.

In this work, we address these issues by proposing a novel approach to estimating the action-value function. Unlike previous methods that directly approximate the action-value function from samples, we first learn the transition operator and reward function associated with the Markov decision process (MDP). To model the transition operator, we adopt the Conditional Mean Embedding (CME) framework [14, 15]. We then leverage an operatorial characterization of the action-value function to express it in terms of these estimated quantities. This strategy draws a peculiar connection with world model methods and can be interpreted as world model learning via CMEs. The notion of world models for RL has been popularized by Ha and Schmidhuber in [16] distilling ideas from the early nineties [17, 18]. Traditional world model methods such as [16, 19] emphasize learning an implicit model of the environment in the form of a simulator. The simulator can be sampled directly in the latent representation space, which is usually of moderate dimension, resulting in a compressed and high-throughput model of the environment. This approach, however, requires extensive sampling for application to PMD and incurs into two sources of error in estimating the action-value function: model and sampling error. In contrast, CMEs can be used to estimate expectations without sampling and incur only in model error, for which learning bounds are available [20, 21]. One of our key results shows that by modeling the transition operator as a CME between suitable Sobolev spaces, we can compute estimates of the action-value function of any sufficiently smooth policy in closed form via efficient matrix operations.

Combining our estimates of the action-value function with the PMD framework we obtain a novel RL algorithm that we dub Policy mirror descent with Operator World-models for Reinforcement learning (POWR) . A byproduct of adopting CMEs to model the transition operator is that we can naturally extend PMD to infinite state space settings. We leverage recent advancements in characterizing the sample complexity of CME estimators to prove convergence rates for the proposed algorithm to the global maximum of the RL Problem. Our approach is similar in spirit to [22], which proposed a value iteration strategy based on CMEs. We extend these ideas to PMD strategies and refine previous results on convergence rates. Learning the transition operator with a least-squares based estimator was also recently considered in [23] and [24]. The latter proposed an optimistic strategy to prove near-optimal regret bounds in linear mixture MDP settings [25]. In contrast, in this work, we cast our problem within a linear MDP setting with possibly infinite latent dimension. We validate our approach on simple environments from the Gym library [26] both in finite and infinite state settings, reporting promising evidence in support of our theoretical analysis.

Contributions . The main contributions of this paper are: i ) a CME-based world model framework, which enables us to generate estimators for the action-value function of a policy in closed form via matrix operations. ii ) An (inexact) PMD algorithm combining the learned CMEs world models with mirror descent update steps to generate improved policies. iii ) Showing that the algorithm is well-defined when learning the world model as an operator between a suitable family of Sobolev spaces. iv ) Showing convergence rates of the proposed approach to the global maximum of the RL problem, under regularity assumptions on the MDP. v ) Empirically testing the proposed approach in practice, comparing it with well-established baselines.

2 Problem Formulation and Policy Mirror Descent

We consider a Markov Decision Process (MDP) over a state space X and action space A , with transition kernel τ . We assume X and A to be Polish, τ : Ω → P ( X ) to be a Borel measurable function from the joint space Ω = X × A to the space P ( X ) of Borel probability measures on X . We define a policy to be a Borel measurable function π : X → P ( A ) . When A (respectively X ) is a finite set, the space P ( A ) = ∆( A ) ⊆ R | A | (respectively P ( X ) = ∆( X ) ⊆ R | X | ) corresponds to the probability simplex. Given a discount factor γ > 0 , an initial state distribution ν ∈ P ( X ) and a Borel measurable bounded and non-negative reward 2 function r : Ω → R we denote by

the (discounted) expected return of the policy π applied to the MDP, yielding the Markov process ( X t , A t ) t ∈ N , where X 0 is distributed according to ν and for each t ∈ N the action A t is distributed according to π ( ·| X t ) and X t +1 according to τ ( ·| X t , A t ) .

In sequential decision settings, the goal is to find the optimal policy π ∗ maximizing (1) over the space of all measurable policies. In reinforcement learning, one typically assumes that knowledge of the transition τ , the reward r , and (possibly) the starting distribution ν is not available. It is only possible to gather information about these quantities by interacting with the MDP to sample state-action pairs ( x t , a t ) and corresponding rewards r ( x t , a t ) and transitions x t +1 .

Policy Mirror Descent (PMD) . In so-called tabular settings - in which both X and A are finite sets the policy optimization problem amounts to maximizing (1) over the space Π = ∆( A ) ⊗ R |X| of column substochastic matrices, namely matrices M ∈ R |A|×|X| with non-negative entries and whose columns sum up to one, namely M ∗ 1 A = 1 X , with 1 denoting the vector with all entries equal to one on the appropriate space. Borrowing from the convex optimization literature - where mirror descent algorithms offer a powerful approach to minimize a convex functional over a convex constraint set [27, 28] - recent work proposed to adopt mirror descent also for policy optimization, a strategy known as policy mirror descent (PMD) [10]. Even though the objective in (1) is not convex (or concave, since we are maximizing it), it turns out that mirror ascent can nevertheless enjoy global convergence to the maximum, with sublinear [11] or even linear rates [12], at the cost of dimension-dependent constants.

Starting from an initial policy π 0 , PMD generates a sequence ( π t ) t ∈ N according to the update step

for any x ∈ X , with η > 0 a step size, D a suitable Bregman divergence [28] and q π : Ω → R the so-called action-value function of a policy π , see also [12, Sec. 4]. The action-value function

is the discounted return obtained by taking action a ∈ A in state x ∈ X and then following the policy π . The solution to (2) crucially depends on the choice of D . For example, in [10] the authors observed that if D is the Kullback-Leibler divergence, PMD corresponds to the Natural Policy Gradient originally proposed in [29] while [12] showed that if D is the squared euclidean distance, PMD recovers the Projected Policy Gradient method from [11].

PMD in Reinforcement Learning . A clear limitation to adopting PMD in RL settings is that (3) needs exact knowledge of the action-value functions q π t associated to each iterate π t of the algorithm. This requires evaluating the expectation in (3), which is not possible in RL where we do not know the reward r and MDP transition distribution τ in advance. While sampling strategies can be adopted to estimate q π t , a key question is how the approximation error affects PMD.

The work in [12] provides an answer to this question, extending the analysis of PMD to the case where estimates ˆ q π t are used in place of the true action-value function in (2). We recall here an informal version of the result for the case of sublinear convergence rates for PMD. We postpone a more rigorous statement of the theorem and its assumptions to Sec. 5, where we extend it to infinite state spaces X .

Theorem 1 (Inexact PMD (Sec. 5 in [12]) - Informal) . In the tabular setting, let ( π t ) t ∈ N be a sequence of policies obtained by applying the PMD update in (2) with functions ˆ q π t : Ω → R in place of q π t and D a suitable Bregman divergence. For any T ∈ N and ε > 0 , if ∥ ˆ q π t -q π t ∥ ∞ ≤ ε for all t = 1 , . . . , T , then

Thm. 1 implies that inexact PMD retains the convergence rates of its exact counterpart, provided that the approximation error for each action-value function is of order 1 /T in uniform norm ∥·∥ ∞ . While this result supports estimating the action-value function in RL, implementing this strategy in practice poses two main challenges, even in tabular settings. First, approximating the expectation in (3) in ∥·∥ ∞ norm via sampling requires 'starting' the MDP from each state x ∈ X , multiple times. This

is often not possible in RL, where we do not have control over the starting distribution ν . Second, repeating this sampling process to learn a ˆ q π t for each policy π t can become extremely expensive in terms of both the number of computations and interactions with the environment.

In this work, we propose a new strategy to tackle the problems above. Instead of re-sampling the MDP to estimate ˆ q π t at each iteration t , we learn estimators ˆ r and ˆ τ for the reward and transition distribution, respectively. For any policy π , we then leverage the relation between these quantities in (3) to generate an estimator ˆ q for q π . This approach tackles the above challenges since 1) it enables us to control the approximation error on any action-value function in terms of the approximation error of ˆ r and ˆ τ ; 2) it does not require sampling the MDP to learn a new ˆ q π t for each π t generated by PMD.

3 Operator World Models

In this section, we present an operator-based formulation of the problem introduced in Sec. 2 (see also [11]). This will be instrumental in extending the PMD theory to arbitrary state spaces X , to quantify the approximation error of the action-value function in terms of the approximation error of the reward and transition distribution, and to motivate conditional mean embeddings as the tool to learn these latter quantities.

Conditional Expectation Operators . We start by defining the transfer operator T associated with the MDP transition distribution τ . Let B b ( X ) denote the space of bounded Borel measurable functions on a space X . Formally, T : B b ( X ) → B b (Ω) is the linear operator such that, for any f ∈ B b ( X )

where X ′ is sampled according to τ ( ·| x, a ) . Note that T is the Markov operator [30, Ch. 19] encoding the dynamics of the MDP and its conjugate T ∗ : M (Ω) →M ( X ) is the operator mapping signed Borel measures µ ∈ M (Ω) to their transition via τ as ( T ∗ µ )( B ) = ∫ B× Ω τ ( dx ′ | x, a ) µ ( dx, da ) for any measurable B ⊆ X . For any policy π we define the operator P π : B b (Ω) → B b ( X ) such that for all g ∈ B b (Ω)

where the expectation is taken over the action A sampled according to π ( ·| x ) . Also P π is a Markov operator and its conjugate P ∗ π : M ( X ) →M (Ω) is the operator mapping any ν ∈ M ( X ) to its joint measure with π , namely ( P ∗ π ν )( C ) = ∫ C π ( da | x ) ν ( dx ) for any measurable C ⊆ Ω .

Operator Formulation of RL . With these two operators in place, we can characterize the expected reward after a single interaction between a policy π and the MDP as ( TP π r )( x, a ) = E [ r ( X ′ , A ′ ) | X 0 = x, A 0 = a ] . This observation can be applied recursively, yielding the operatorial characterization of the action-value function from (3)

where the last equality follows from T and P π being Markov operators [30, Ch. 19] ( ∥ T ∥ = ∥ P π ∥ = 1 ), making the Neumann series convergent. Analogously, we can reformulate the RL objective introduced in (1) as the pairing

for ν ∈ P ( X ) a starting distribution. In both (7) and (8) the operatorial formulation encodes the cumulative reward collected through the (possibly infinitely many) interactions of the policy with the MDP in closed form, as the inversion ( Id -γ TP π ) -1 r . This characterization motivates us to learn T and r from data and then express any action-value function as the interaction of these two terms with the policy π as in (7), rather than learning each q π independently for any π .

Learning the World Model via Conditional Mean Embeddings . Conditional Mean Embeddings (CME) offer an effective tool to model and learn conditional expectation operators from data [15]. They cast the problem of learning T by studying the restriction of its action on a suitable family

of functions. Let φ : X → F and ψ : Ω → G two feature maps with values into the Hilbert spaces F and G . With some abuse of notation (which is justified by them being Hilbert spaces), we interpret F and G as subspaces of functions in B b ( X ) and B b (Ω) of the form f ( x ) = ⟨ f, φ ( x ) ⟩ and g ( x, a ) = ⟨ g, ψ ( x, a ) ⟩ for any f ∈ F and g ∈ G and any ( x, a ) ∈ Ω . We say that the linear MDP assumption holds with respect to ( φ, ψ ) if

Assumption 1 (Linear MDP - Well-specified CME) . The restriction of T to F is a Hilbert-Schmidt operator T | F ∈ HS ( F , G ) .

In CME settings, the assumption above is known as requiring the CME of τ to be well-specified . The following result, proved in Appendix A.2, clarifies this aspect and establishes the relation of Asm. 1 with the standard definition of linear MDP.

Proposition 2 (Well-specified CME) . Under Asm. 1, ( T | F ) ∗ = ( T ∗ ) | G and, for any ( x, a ) ∈ Ω

Proposition 2 shows that (9) is equivalent to the standard linear MDP assumption [31, Ch. 8] when X is a finite set (taking φ the one-hot encoding) while being weaker in infinite settings. From the CME perspective, the proposition characterizes the action of ( T | F ) ∗ as sending evaluation vectors in G to the conditional expectation of evaluation vectors in F with respect to τ , the definition of conditional mean embedding of τ [32, 15]. This characterization also suggests a learning strategy: (9) characterizes the action of T as evaluating the conditional expectation of a vector φ ( X ′ ) given ( x, a ) . Given a set of points ( x i , a i ) n i =1 and corresponding x ′ sampled from τ ( ·| x i , a i ) , this can be learned by minimizing the squared loss, yielding the estimator (see [15, Sec 4.2])

When F and G are finite dimensional, S n and Z n are matrices with n rows, each corresponding respectively to the vectors ψ ( x i , a i ) and φ ( x ′ i ) for i = 1 , . . . , n . In the infinite setting, they generalize to operators S n : G → R n and Z n : F → R n . The matrix K λ = S n S ∗ n + nλ Id n ∈ R n × n is the regularized Gram (or kernel) matrix with ( i, j ) -th entry corresponding to

We conclude our discussion on learning world models via CMEs by noting that in most RL settings, the reward function is unknown, too. Analogously to what we have described for T n and following the standard practice in supervised settings, we can learn an estimator for r solving a problem akin to (10). This yields a function of the form r n = S ∗ n b = ∑ n i =1 b i ψ ( x i , a i ) as the linear combination of the embedded training points with the entries of the vector b = K -1 λ y where y ∈ R n is the vector with entries y i = r ( x i , a i ) .

Estimating the Action-value Function q π . We now propose our strategy to generate an estimator for the action-value function q π of a given policy π in terms of an estimator for the reward r and a world model for T learned in terms of the restriction to G and F . To this end, we need to introduce the notion of compatibility between a policy π and the pair ( G , F ) .

Definition 1 ( ( G , F ) -compatibility) . A policy π : X → P ( A ) is compatible with two subspaces F ⊆ B b ( X ) and G ⊆ B b (Ω) if the restriction P π to G is a bounded linear operator with range ⊆ F , that is ( P π ) | G : G → F .

Definition 1 is analogous to the linear MDP Asm. 1 in that it requires the restriction of P π to G to take values in the associated space F . However, it is slightly weaker since it requires this restriction to be bounded (and linear) rather than being an HS operator. We will discuss in Sec. 4 how this difference will allow us to show that a wide range of policies (in particular those generated by our POWR method) is ( G , F ) -compatible for our choice of function spaces. Definition 1 is the key condition that enables us to generate an estimator for q π , as characterized by the following result.

Algorithm 1 POWR: POLICY MIRROR DESCENT WITH OPERATOR WORLD-MODELS FOR RL

Input: Dataset ( x i , a i , x ′ i , r i ) n i =1 , discount factor γ ∈ (0 , 1) , step size η > 0 , kernel function k ( x, x ′ ) = ⟨ ϕ ( x ) , ϕ ( x ′ ) ⟩ with ϕ : X → H as in Proposition 4, initial weights C 0 = 0 ∈ R n ×|A| .

return

π

T

:

X →

∆(

A

)

such that

π

T

(

x

) =

SOFTMAX

(

η H

x

C

T

)

with

H

x

= (

k

(

x, x

i

))

n

i

=1

∈

R

.

Proposition 3. Let T n = S ∗ n BZ n ∈ HS ( F , G ) and r n = S ∗ n b ∈ G for respectively a B ∈ R n × n and b ∈ R n . Let π be ( G , F ) -compatible. Then,

where M π = Z n P π S ∗ n ∈ R n × n is the matrix with entries

Proposition 3 leverages a kernel trick argument to express the estimator for the action-value function of π as the linear combination ˆ q π = ∑ n i =1 c i ψ ( x i , a i ) of the (embedded) training points ψ ( x i , a i ) and the entries c i of the vector c = ( Id -γBM π ) -1 b ∈ R n . We prove the result in Appendix A.4. We note that in (12) both B and M π are n × n matrices and therefore the characterization of ˆ q π amounts to solving a n × n linear system. For settings where n is large, one can adopt random projection methods such as Nyström approximation to learn T n and r n [33]. These strategies have been recently shown to significantly reduce the computational load of learning while retaining the same empirical and theoretical performance as their non-approximated counterparts [34, 35].

We conclude this section noting how (13) implies that we only need to be able to evaluate π , but we do not need explicit knowledge of P π as operator. As we shall see, this property will be instrumental to prove generalization bounds for our proposed PMD algorithm in Sec. 4.

4 Proposed Algorithm: POWR

We are ready to describe our algorithm for world model-based PMD. In the following, we restrict to the case where |A| < ∞ is a finite set. As introduced in Sec. 2, policy mirror descent methods are mainly characterized by the choice of Bregman divergence D used for the mirror descent update and, in the case of inexact methods, the estimator ˆ q π t of the action-value function q π t for the intermediate policies generated by the algorithm.

In POWR, we combine the CME world model presented in Sec. 3 with mirror descent steps using the Kullback-Leibler divergence D KL ( p ; p ′ ) = ∑ a ∈A p a log( p a /p ′ a ) in the update of (2). It was shown in [10] that in this case PMD corresponds to Natural Policy Gradient [29]. As showed in [36, Example 9.10], the solution to (2) can be written in closed form for any x ∈ X as

where we used the estimated action-value function ˆ q π from Proposition 3. Additionally, the formula above can be applied recursively, expressing π t +1 as the softmax operator applied to the discounted

n

sum of the action-value functions up to the current iteration

Choice of the Feature Maps . A key question to address before adopting the action-value estimators introduced in Sec. 3 is choosing the two spaces F and G to perform world model learning. Specifically, to apply Proposition 3 and obtain proper estimators ˆ q π t , we need to guarantee that all policies generated by the PMD update (14) are ( G , F ) -compatible (Definition 1). The following result describes a suitable family of such spaces.

Proposition 4 (Separable Spaces) . Let ϕ : X → H be a feature map into a Hilbert space H . Let F = H ⊗ H and G = R |A| ⊗ H with feature maps respectively φ ( x ) = ϕ ( x ) ⊗ ϕ ( x ) and ψ ( x, a ) = ϕ ( x ) ⊗ e a , with e a ∈ R |A| the one-hot encoding of action a ∈ A . Let π : X → ∆( A ) be a policy such that π ( a |· ) = ⟨ p a , ϕ ( · ) ⟩ with p a ∈ H for any a ∈ A . Then, π is ( G , F ) -compatible.

Proposition 4 (proof in Appendix A.5) states that for the specific choice of function spaces F = H⊗H and G = R |A| ⊗H , we can guarantee ( G , F ) -compatibility, provided that H is rich enough to 'contain' all π ( a |· ) for a ∈ A . We postpone the discussion on identifying a suitable spaces H for PMD to Sec. 5, since ( G , F ) -compatibility is not needed to mechanically apply Proposition 3 and obtain an estimator ˆ q π . This is because (12) exploits a kernel-trick to bypass the need to know P π explicitly and rather requires only to be able to evaluate π on the training data. The latter is possible for π t +1 , thanks to the explicit form of the PMD update in (14). We can therefore present our algorithm.

POWR . Alg. 1 describes Policy mirror descent with Operator World-models for Reinforcement learning (POWR) . Following the intuition of Proposition 4, the algorithm assumes to work with separable spaces. During an initial phase, we learn the world model T n = S ∗ n K -1 λ Z n and the reward r n = S ∗ n b fitting the conditional mean embedding described in (10) on a dataset ( x i , a i , x ′ i , r i ) n i =1 (e.g. obtained via experience replay [37]). Once the world model has been learned, we optimize the policy and perform PMD iterations via (15). In this second phase, we first evaluate the past (cumulative) action-value estimators ∑ t s =0 ˆ q π s to obtain the policy π t +1 ( ·| x i ) by (inexact) PMD via the softmax operator in (15). We use the newly obtained policy to compute the matrix M π t +1 defined in (13), which is a key component to obtain the estimator ˆ q π t +1 for q π t +1 . We note that in the case of the separable spaces of Proposition 4, this matrix reduces to the n × n matrix with entries k ( x ′ i , x j ) π t +1 ( a j | x ′ i ) , where k ( x ′ i , x j ) = ⟨ ϕ ( x ′ i ) , ϕ ( x j ) ⟩ is the kernel matrix between initial and evolved states. Finally, we obtain c = ( Id -γK -1 λ M ) -1 b and model ˆ q π t +1 = S ∗ n c according to Proposition 3.

Clearly, the world model learning and PMD phases can be alternated in POWR, essentially finding a trade-off between exploration and exploitation. This could possibly lead to a refinement of the world model as more observations are integrated into the estimator. While in this work we do not investigate the effects of such alternating strategy, Thm. 7 offers relevant insights in this sense. The result characterizes the behavior of the PMD algorithm when combined with a varying (possibly increasing) accuracy in the estimation of the action-valued function (see Sec. 5 for more details).

5 Theoretical Analysis

We now show that POWR converges to the global maximizer of the RL problem in (1). To this end, we first identify a family of function spaces guaranteed to be compatible with the policies generated by Alg. 1. Then, we provide an extension of the result in [12] for inexact PMD to infinite state spaces X , showing the impact of the action-value approximation error on the convergence rates. For the estimator introduced in Proposition 3, we relate this error to the approximation errors of T n and r n leveraging the Simulation Lemma A.6. Finally, we use recent advancements in the characterization of CMEs' fast learning rates to bound the sample complexity of these latter quantities, yielding error bounds for POWR.

POWR is Well-defined . To properly apply Proposition 3 to estimate the action-value function of any PMD iterate π t , we need to guarantee that every iterate belongs to the space H according to Proposition 4. The following result provides such a family of spaces.

Theorem 5. Let X ⊂ R d be a compact set and let H = W 2 ,s ( X ) be the Sobolev space of smoothness s > 0 (see e.g. [38]). Let π t ( a |· ) and ˆ q π t ( · , a ) belong to H for any a ∈ A and π t ( a | x ) > 0 for any x ∈ X . Then the policy π t +1 solution to the PMD update in (14) belongs to H .

According to Thm. 5, Sobolev spaces offer a viable choice for compatibility with PMD-generated policies. This observation is further supported by the fact that Sobolev spaces of smoothness s > d/ 2 are so-called reproducing kernel Hilbert spaces (rkhs) (see e.g. [39, Ch. 10]). We recall that rkhs are always naturally equipped with a ϕ : X → H such that the inner product ⟨ ϕ ( x ) , ϕ ( x ′ ) ⟩ = k ( x, x ′ ) defines a so-called reproducing kernel, namely a positive definite function that is (usually) efficient to evaluate, even if ϕ ( x ) is high or infinite dimensional. For example, H = W 2 ,s ( X ) with s = ⌈ d 2 ⌉ has associated kernel k ( x, x ′ ) = e -∥ x -x ′ ∥ /σ with bandwidth σ > 0 [39]. By applying Thm. 5 to the iterates generated by Alg. 1 we have the following result.

Corollary 6. With the hypothesis of Proposition 4 let H = W 2 ,s ( X ) with s > d/ 2 . Let T n ∈ HS ( F , G ) and r n ∈ G characterized as in Proposition 3. Let π 0 ( a |· ) ∝ e ηq 0 ( · ,a ) for q 0 such that q 0 ( · , a ) ∈ H any a ∈ A . Then, for any t ∈ N the PMD iterates π t generated by Alg. 1 are such that π t ( a |· ) ∈ H and hence are ( G , F ) -compatible.

The above corollary guarantees us that if we are able to learn our estimates for the action-value function in a suitably regular Sobolev space H , then POWR is well-defined. This is a necessary condition to then being able to study its theoretical behavior in our main result. We report the proofs of Thm. 5 and Cor. 6 in Appendix C.1.

Inexact PMD Converges . We now present a more rigorous version of the characterization of the convergence rates of the inexact PMD algorithm discussed informally in Thm. 1.

Theorem 7 (Convergenge of Inexact PMD) . Let ( π t ) t ∈ N be a sequence of policies generated by Alg. 1 that are all ( G , F ) -compatible. If the action-value functions ˆ q π t are estimated with an error ∥ q π t -ˆ q π t ∥ ∞ ≤ ε t , the iterates of Alg. 1 converge to the optimal policy as

where π ∗ : X → ∆( A ) is a measurable maximizer of (8) .

Thm. 7 shows that inexact PMD can behave comparably to its exact version provided that 1) the action value functions ˆ q π t are estimated with increasing accuracy, and 2) that the sequence of policies is ( G , F ) -compatible, for example in the Sobolev-based setting of Cor. 6. Specifically, if ∥ q π t -ˆ q π t ∥ ∞ ≤ O (1 /t ) for any t ∈ N , the convergence rate of inexact PMD is of order O (log T/T ) , only a logarithmic factor slower than exact PMD. This means that we do not necessarily need a good approximation of the world model from the beginning but rather a strategy to improve upon such approximation as we perform more PMD iterations. This suggests adopting an alternating strategy between exploration (world model learning) and exploitation (PMD steps), as suggested in Sec. 4. We do not investigate this question in this work.

The demonstration technique used to prove Thm. 7 follows closely [12, Thm. 8 and 13]. We provide a proof in Appendix B.1 since the original result did not allow for a decreasing approximation error but rather assumed a constant one. Moreover, extending it to the case of infinite X requires taking care of additional details related to potential measurability issues.

Action-value approximation error in terms of World Model estimates . Thm. 7 highlights the importance of studying the error of the estimator for the action-value functions produced by Alg. 1. These objects are obtained via the formula described in Proposition 3 in terms of T n , r n and P π . The exact q π has an analogous closed-form characterization in terms of T , r and P π as expressed in (7), and motivating our operator-based approach. The following result compares these quantities in terms of the approximation errors of the world model and the reward function.

Lemma 8 (Implications of the Simulation Lemma) . Let T n and r n the empirical estimators of the transfer operator T and reward function r as defined in Proposition 3, respectively. If T satisfies Asm. 1, r ∈ G , and γ ∥ ∥ T n ∥ ∥ < γ ′ < 1 , then, for every ( G , F ) -compatible policy π

In the result above, when applied to a function in G , such as r n , the uniform norm is to be interpreted as the uniform norm of the evaluation of such function, namely ∥ r n ∥ ∞ = sup ( x,a ) ∈ Ω | ⟨ r n , ψ ( x, a ) ⟩ | , and analogously for T n . The proof, reported in Appendix C.2, follows by first decomposing the difference q π -ˆ q π with the simulation lemma [31, Lemma 2.2] and then applying the triangular inequality for the uniform norm.

POWRconverges . We are ready to state the convergence result for Alg. 1. We consider the setting where the dataset used to learn T n (and r n ) is made of i.i.d. triplets ( x i , a i , x ′ i ) with ( x i , a i ) sampled from a distribution ρ ∈ P (Ω) supported on all Ω (such as the state occupancy measure (see e.g. [11] or Appendix A.3) of the uniform policy π ( ·| x ) = 1 / |A| ) and x ′ i sampled from τ ( ·| x i , a i ) . To guarantee bounds in uniform norm, the result makes a further regularity assumption, of the transfer operator (and the reward function)

Assumption 2 (Strong Source Condition) . Let ρ ∈ P (Ω) and C ρ the covariance operator ∑ a ∈A ∫ X ρ ( dx, a ) ψ ( x, a ) ⊗ ψ ( x, a ) . The transition operator T and the reward function r are such that T | F ∈ HS ( F , G ) and r ∈ G . Further, ∥ ∥ ( T | F ) ∗ C -β ρ ∥ ∥ HS < ∞ and ∥ ∥ C -β ρ r ∥ ∥ G < ∞ for some β > 0 .

Asm. 2 imposes a strong requirement to the so-called source condition , a quantity that describes how well the target objective of the learning process (here T and r ) 'interact' with the sampling distribution. The assumption is always satisfied when the hypothesis space is finite dimensional (e.g. in the tabular RL setting) and imposes additional smoothness on T and r when belonging to a Sobolev space. Equipped with this assumption, we can now state the convergence theorem for Alg. 1.

Theorem 9. Let ( π t ) t ∈ N be a sequence of policies generated by Alg. 1 in the same setting of Cor. 6. If the action-value functions ˆ q π t are estimated from a dataset ( x i , a i ; x ′ i ) n i =1 with ( x i , a i ) ∼ ρ ∈ P (Ω) such that Asm. 2 holds with parameter β , the iterates of Alg. 1 converge to the optimal policy as

with probability not less than 1 -4 e -δ . Here, α ∈ ( β 2+2 β , β 1+2 β ) and π ∗ : X → ∆( A ) is a measurable maximizer of (8) .

The proof of Thm. 9 is reported in Appendix C and combines the results discussed in this section with fast convergence rates for the least-squares [40] and CME [21] estimators. In particular we first use Thm. 5 to guarantee that the policies produced by Alg. 1 are all ( G , F ) -compatible and therefore that applying Proposition 3 to obtain an estimator for the action-value function is well-defined. Then, we use Lemma 8 to study the approximation error of these estimators in terms of our estimates for the world model and the reward function. Bounds on these quantities are then used in the result for inexact PMD convergence in Thm. 7. We note here that since the latter results require convergence in uniform norm, we cannot leverage standard results for least-squares and CME convergence, which characterize convergence in L 2 (Ω , µ ) and would only require Asm. 1 (Linear MDP) to hold. Rather, we need to impose Asm. 2 to guarantee faster rates in uniform norm.

6 Experimental results

We empirically evaluated POWR on classical Gym environments [26], ranging from discrete ( FrozenLake-v1 , Taxi-v3 ) to continuous state spaces ( MountainCar-v0 ). To ensure balancing between exploration and exploitation of our method, we alternated between running the environment with the current policy to collect samples for world model learning and running Alg. 1 for a number of steps to generate a new policy. Appendix D provides implementation details regarding this process as well as additional results.

Fig. 1 compares our approach with the performance of well-established baselines including A2C [41], DQN [4], TRPO [7], and PPO [6]. The figure reports the average cumulative reward obtained by the models on test environments with respect to the number of interactions with the MDP ( timesteps in log scale in the figure) across 7 different training runs. In all plots, the horizontal dashed line represents the 'success' threshold for the corresponding environment, according to official guidelines. We observe that our method outperforms all competitors by a significant margin in terms of sample complexity, that is, the reward achieved after a given number of timesteps. In the case of the Taxi-v3

10 3

10 4

10 5

Timestep (logscale)

0.0

0.2

0.4

0.6

0.8

1.0

Reward

A2C

DQN

TRPO

PPO

(a)

FrozenLake-v1

10 3

10 4

10 5

Timestep (logscale)

800

700

600

500

400

300

200

100

0

Reward

(b)

Taxi-v3

10 3

10 4

10 5

10 6

Timestep (logscale)

200

180

160

140

120

100

Reward

(c)

MountainCar-v0

environment, it avoids converging to a local optimum, in contrast every other method with the exception of DQN. On the downside, we note that our method exhibits less stability than other approaches, particularly during the initial stages of the training process. This is arguably due to a sub-optimal interplay between exploration and exploitation, which will be the subject of future work.

7 Conclusions and Future Work

Motivated by recent advancements in policy mirror descent (PMD), this work introduced a novel reinforcement learning (RL) algorithm leveraging these results. Our approach operates in two, possibly alternating, phases: learning a world model and planning via PMD. During exploration, we utilize conditional mean embeddings (CMEs) to learn a world model operator, showing that this procedure is well-posed when performed over suitable Sobolev spaces. The planning phase involves PMD steps for which we guarantee convergence to a global optimum at a polynomial rate under specific MDP regularities.

Our analysis opens avenues for further exploration. Firstly, extending PMD to infinite action spaces remains a challenge. While we introduced the operatorial perspective on RL for infinite state space settings, the PMD update with KL divergence requires approximation methods (e.g., Monte Carlo) whose impact on convergence requires investigation. Secondly, scalability to large environments requires adopting approximated yet efficient CME estimators like Nystrom [35] or reduced-rank regressors [42, 43]. Thirdly, a question we touched upon only empirically, is whether alternating world model learning with inexact PMD updates benefits the exploration-exploitation trade-off. Studying this strategy's impact on convergence is a promising future direction. Finally, a crucial question is generalizing our policy compatibility results beyond Sobolev spaces. Ideally, a representation learning process would identify suitable feature maps that guarantee compatibility with the PMD-generated policies while allowing for added flexibility in learning the world model.

Acknowledgments and Disclosure of Funding

We acknowledge financial support from NextGenerationEU and MUR PNRR project PE0000013 CUP J53C22003010006 'Future Artificial Intelligence Research (FAIR)', EU grant ELISE (GA no 951847) and EU Project ELIAS (GA no 101120237).

References

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018.

OpenAI: Marcin Andrychowicz, Bowen Baker, Maciek Chociej, Rafal Jozefowicz, Bob McGrew, Jakub Pachocki, Arthur Petron, Matthias Plappert, Glenn Powell, Alex Ray, et al. Learning dexterous in-hand manipulation. The International Journal of Robotics Research , 39(1):3-20, 2020.

David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. Nature , 529(7587):484-489, 2016.

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv1312.5602 , 2013.

Hongzi Mao, Mohammad Alizadeh, Ishai Menache, and Srikanth Kandula. Resource management with deep reinforcement learning. In Proceedings of the 15th ACM workshop on hot topics in networks , pages 50-56, 2016.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv1707.06347 , 2017.

John Schulman, Sergey Levine, Philipp Moritz, Michael I. Jordan, and Pieter Abbeel. Trust region policy optimization, 2017.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor, 2018.

Dimitri Bertsekas. Dynamic Programming and Optimal Control: Volume I , volume 4. Athena scientific, 2012.

Lior Shani, Yonathan Efroni, and Shie Mannor. Adaptive trust region policy optimization: Global convergence and faster rates for regularized mdps. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 34, pages 5668-5675, 2020.

Alekh Agarwal, Sham M. Kakade, Jason D. Lee, and Gaurav Mahajan. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. Journal of Machine Learning Research , 22(98):1-76, 2021.

Lin Xiao. On the convergence rates of policy gradient methods. Journal of Machine Learning Research , 23(282):1-36, 2022.

Matthieu Geist, Bruno Scherrer, and Olivier Pietquin. A theory of regularized markov decision processes. In International Conference on Machine Learning , pages 2160-2169. PMLR, 2019.

Kenji Fukumizu, Francis R. Bach, and Michael I. Jordan. Dimensionality reduction for supervised learning with reproducing kernel Hilbert spaces. Journal of Machine Learning Research , 5:73-99, 2004.

Krikamol Muandet, Kenji Fukumizu, Bharath Sriperumbudur, Bernhard Schölkopf, et al. Kernel mean embedding of distributions: A review and beyond. Foundations and Trends® in Machine Learning , 10(1-2):1-141, 2017.

David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. Advances in neural information processing systems , 31, 2018.

Richard S Sutton. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin , 2(4):160-163, 1991.

Jürgen Schmidhuber. Making the world differentiable: on using self supervised fully recurrent neural networks for dynamic reinforcement learning and planning in non-stationary environments , volume 126. Inst. für Informatik, 1990.

Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104 , 2023.

Carlo Ciliberto, Lorenzo Rosasco, and Alessandro Rudi. A general framework for consistent structured prediction with implicit loss embeddings. Journal of Machine Learning Research , 21(98):1-67, 2020.

Zhu Li, Dimitri Meunier, Mattes Mollenhauer, and Arthur Gretton. Optimal rates for regularized conditional mean embedding learning. Advances in Neural Information Processing Systems , 35:4433-4445, 2022.

Steffen Grünewälder, Guy Lever, Luca Baldassarre, Massimiliano Pontil, and Arthur Gretton. Modelling transition dynamics in mdps with rkhs embeddings. In Proceedings of the 29th International Conference on International Conference on Machine Learning , pages 1603-1610, 2012.

Preston Rozwood, Edward Mehrez, Ludger Paehler, Wen Sun, and Steven L Brunton. Koopmanassisted reinforcement learning. arXiv preprint arXiv:2403.02290 , 2024.

Antoine Moulin and Gergely Neu. Optimistic planning by regularized dynamic programming. In International Conference on Machine Learning , pages 25337-25357. PMLR, 2023.

Alex Ayoub, Zeyu Jia, Csaba Szepesvari, Mengdi Wang, and Lin Yang. Model-based reinforcement learning with value-targeted regression. In International Conference on Machine Learning , pages 463-474. PMLR, 2020.

Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arxiv. arXiv preprint arXiv:1606.01540 , 10, 2016.

Amir Beck and Marc Teboulle. Mirror descent and nonlinear projected subgradient methods for convex optimization. Operations Research Letters , 31(3):167-175, 2003.

Sébastien Bubeck. Convex optimization: Algorithms and complexity. Foundations and Trends® in Machine Learning , 8(3-4):231-357, 2015.

Sham M. Kakade. A natural policy gradient. Advances in Neural Information Processing Systems , 14, 2001.

C.D. Aliprantis and K.C. Border. Infinite Dimensional Analysis: A Hitchhiker's Guide . Studies in Economic Theory. Springer, 1999.

Alekh Agarwal, Nan Jiang, Sham M. Kakade, and Wen Sun. Reinforcement learning: Theory and algorithms. 2021. URL https://rltheorybook.github.io , 2022.

Steffen Grünewälder, Guy Lever, Luca Baldassarre, Sam Patterson, Arthur Gretton, and Massimilano Pontil. Conditional mean embeddings as regressors. In Proceedings of the 29th International Conference on International Conference on Machine Learning , pages 1803-1810, 2012.

Christopher Williams and Matthias Seeger. Using the Nyström method to speed up kernel machines. Advances in Neural Information Processing Systems , 13, 2000.

Alessandro Rudi, Raffaello Camoriano, and Lorenzo Rosasco. Less is more: Nyström computational regularization. Advances in Neural Information Processing Systems , 28, 2015.

Giacomo Meanti, Antoine Chatalic, Vladimir R. Kostic, Pietro Novelli, Massimiliano Pontil, and Lorenzo Rosasco. Estimating Koopman operators with sketching to provably learn large scale dynamical systems. Advances in Neural Information Processing Systems , 36, 2023.

Amir Beck. First-Order Methods in Optimization . Society for Industrial and Applied Mathematics, 2017.

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Bellemare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature , 518(7540):529-533, 2015.

Robert A. Adams and John J. F. Fournier. Sobolev Spaces . Elsevier, 2003.

Holger Wendland. Scattered data approximation , volume 17. Cambridge University Press, 2004.

Simon Fischer and Ingo Steinwart. Sobolev norm learning rates for regularized least-squares algorithms. Journal of Machine Learning Research , 21(205):1-38, 2020.

Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. arXiv1602.01783 , 2016.

Vladimir R. Kostic, Pietro Novelli, Andreas Maurer, Carlo Ciliberto, Lorenzo Rosasco, and Massimiliano Pontil. Learning dynamical systems via Koopman operator regression in reproducing kernel Hilbert spaces. Advances in Neural Information Processing Systems , 35:4017-4031, 2022.

Giacomo Turri, Vladimir Kostic, Pietro Novelli, and Massimiliano Pontil. A randomized algorithm to solve reduced rank operator regression. arXiv preprint arXiv:2312.17348 , 2023.

Gene H. Golub and Charles F. Van Loan. Matrix Computations . Johns Hopkins University Press, 2013.

O. Kallenberg. Foundations of Modern Probability . Probability and Its Applications. Springer New York, 2002.

Giulia Luise, Saverio Salzo, Massimiliano Pontil, and Carlo Ciliberto. Sinkhorn barycenters with free support via frank-wolfe algorithm. Advances in Neural Information Processing Systems , 32, 2019.

Vladimir R. Kostic, Karim Lounici, Pietro Novelli, and Massimiliano Pontil. Sharp spectral rates for koopman operator learning. Advances in Neural Information Processing Systems , 36, 2023.

Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines3: Reliable reinforcement learning implementations. Journal of Machine Learning Research , 22(268):1-8, 2021.

Appendix

The appendices are organized as follows:

Appendix A discuss the operatorial formulation of RL and show how to derive the operatorbased results in this work.

Appendix B focuses on policy mirror descent (PMD) and its convergence rate in the inexact setting.

Appendix C proves the main result of this work, namely the theoretical analysis of POWR.

Appendix D provide details on the experiments reported in this work.

A Operatorial Results

A.1 Auxiliary Lemma

We recall here a corollary of the Sherman-Woodbury identity [44].

Lemma A.1. Let A and B two conformable linear operators such that ( I + AB ) -1 is invertible. Then ( I + AB ) -1 A = A ( I + BA ) -1

Proof. The result is obvious if A is invertible. More generally, we consider the following two applications of the Sherman-Woodbury [44] formula

Multiplying the two equations by A respectively to the right and to the left, we obtain the desired result.

A.2 Markov operators and their properties

We recall here the notion of Markov operators, which is central for a number of results in the following. We refer to [30, Chapter 19] for more details on the topic.

Definition A.1 (Markov operators) . Let X and Y be Polish spaces. A bounded linear operator L ( B b ( X ) , B b ( Y )) is a Markov operator if is positive and maps the unit function to itself, that is:

where 1 X : X → R (respectively 1 Y ) denotes the function taking constant value equal to 1 on X (respectively Y ).

We recall that Markov operators are a convex subset of L ( B b ( X ) , B b ( Y )) . Here we denote this space as L M ( B b ( X ) , B b ( Y )) . Direct inspection of (5) and (6) shows that the transfer operator T associated to an MDP and the policy operator P π associated to a policy π are both Markov operators.

Markov Operators and Policy Operators . In (6) we defined the policy operator P π associated to a policy π . It turns out that the converse is also true, namely that any such Markov operator is a policy operator.

Proposition A.2. Let P ∈ L M ( B b ( X ) , B b (Ω)) be a Markov operator. Then there exists π P , such that the associated policy operator corresponds to P , namely P π P = P .

Proof. Define the map π P : X → M ( A ) taking value in the space of bounded Borel measures over A such that, for any x ∈ X and any B ⊆ A Borel measurable subset

and

We need to guarantee that for every x ∈ X the function π P ( ·| x ) is a signed measure. To show this, first note that the operation defined by π P is well-defined, since for any measurable set B the function 1 X×B is also measurable, making π P ( B| x ) well defined as well. Moreover, since 1 ∅ ( a ) = 0 for any a ∈ A , it implies that 1 X×∅ = 0 and therefore π P ( ∅| x ) = 0 for any x ∈ X . Finally, σ -additivity follows from the definition of indicator functions, namely 1 ⋃ ∞ i =1 B i = ∑ ∞ i =1 1 B i for any family of pair-wise disjoint sets ( B i ) n i =1 , which implies π P ( ⋃ ∞ i =1 B i | x ) = ∑ ∞ i =1 π P ( B i | x ) for any x ∈ X .

We now apply the two properties of Markov operators to show that π P takes values in P ( A ) , namely it is a non-negative measure that sums to 1 . Since Markov operators map non-negative functions in non-negative functions and since 1 X×B ≥ 0 for any B ⊆ X , we have π ( ·| x ) ≥ 0 as well for any x ∈ X . Moreover, since Ω = X × A and P 1 Ω = 1 X , we have

for any x ∈ X . Therefore π P ( ·| x ) is a probability measure for any x ∈ X . Direct application of (6) shows that the associated policy operator corresponds to P , namely P π P = P as desired.

Given the correspondence between policies and their Markov operator according to (6) and Proposition A.2, in the following we will denote the policy operator associated to a policy π only P where clear from context.

With the definition of the Markov operator in place, we can now prove the following result introduced in the main paper.

Proposition 2 (Well-specified CME) . Under Assumption 1, ( T | F ) ∗ = ( T ∗ ) | G and, for any ( x, a ) ∈ Ω

Proof. Recall that since they are Hilbert spaces F ∼ = F ∗ and G ∼ = G ∗ are isometric to their dual and therefore we can interpret any f ∈ F as the function f ( · ) = ⟨ f, φ ( · ) ⟩ with some abuse of notation, where clear from context. By Assumption 1 we have that T | F takes values in G . This means that ( T | F f ) ∈ G or, in other words

Since the above equality holds for any f ∈ F (9) holds, as desired.

We note that the result can be extended to the setting where T | F ( F ) ⊆ G , namely the image of T | F is contained in G , namely a sort of ( F , G ) -compatibility for the transition operator (see Definition 1).

A.3 The operatorial formulation of RL

According to the operatorial characterization in (7), the action value function of a policy π is directly related to the action of the associated policy operator P . To highlight this relation, we will adopt the following notation:

· Action-value (or Q-)function.

from which we obtain

Value function.

Cumulative reward. The RL objective functional

State visitation (or State occupancy) measure. By the characterization of the adjoints of P and T (see discussion in Section 3 we can represent the evolution of a state distribution ν t at time t to the next state distribution as ν t +1 = T ∗ P ∗ ν t . Applying this relation recursively, we recover the state visitation probability associated to the starting state distribution ν 0 = ν ∈ P ( X ) , the MDP with transition T and the policy P as

where the (1 -γ ) γ t is a normalizing factor to guarantee that the series corresponds to a convex combination of the probability distributions ν t , hence guaranteeing d ν ( P ) to be well-defined (namely it belongs to P ( X ) ).

Previous well-known RL results in operator form . Under the operatorial formulation of RL, we can recover several well-known results from the reinforcement literature with concise proofs. We recall here a few of these results that will be useful in the following.

Remark A.1. Algebraic manipulation of the cumulative expected reward J ( P ) implies

where we used Lemma A.1 and d ν ( P ) is the state visitation distribution starting from ν and following the policy P .

The following result, known as Performance Difference Lemma [see e.g. 11, Lemma 1.16], will be instrumental to prove the convergence rates for PMD in Theorem 7.

Lemma A.3 (Performance difference) . Let P 1 , P 2 two policy operators. The following equality holds

Proof. Using the definition of J ( P 1 ) and Lemma A.1 one gets

Adirect consequence of the operator formulation of the performance difference lemma is the following operator-based characterization of the differential behavior of the RL objective. The result can be found in [11] for the case of finite state and action spaces, however here the operatorial formulation allows for a much more concise proof.

Corollary A.4 (Directional derivatives) . For any two Markov P 1 , P 2 : B b ( X ) → B b (Ω) , we have that the directional derivative in P 1 towards P 2 is

Proof. The result follows by recalling that the space of Markov operators is convex, namely for any h ∈ [0 , 1] the term P h = P 1 + h ( P 2 -P 1 ) is still a Markov operator. Therefore, we can apply Lemma A.3 to obtain

We can therefore divide the above quantity by h and send h → 0 . The result follows by observing that d ν ( P h ) = ( Id -γ TP h ) -∗ ν → ( Id -γ TP 1 ) -∗ ν = d ν ( P 1 ) for h → 0 , since ∥ P h ∥ = 1 for any h ∈ [0 , 1] and the function M ↦→ ( I -γ TM ) -1 is continuous on the open ball of radius 1 /γ > 1 in L ( B b (Ω) , B b ( X )) with respect to the operator norm.

Properties of ( Id -γ PT ) -1 . The quantity ( Id -γ PT ) -1 (note, not ( Id -γ TP ) -1 ) plays a central role in the study of POWR. We prove here a few properties that will be useful in the following. Lemma A.5 (Properties of Id -γ PT ) . The following facts are true:

For any f ≥ 0 ∈ B b ( X ) it holds ( Id -γ PT ) -1 f ≥ f .

The operator (1 -γ )( Id -γ PT ) -1 is a Markov operator.

For any positive measure ν ∈ M ( B b ( X )) it holds (1 -γ ) ∥ ( Id -γ PT ) -∗ ν ∥ TV = ∥ ν ∥ TV .

For any positive measure ν ∈ M ( B b ( X )) it holds ∥ P ∗ ν ∥ TV = ∥ ν ∥ TV .

For any bounded linear operator X , policy operator P and discount factor γ < ∥ X ∥ , it holds ∥ ∥ ( Id -γ XP ) -1 ∥ ∥ ∞ ≤ 1 / (1 -γ ∥ X ∥ ) .

Proof. Since both T and P are Markov operators by construction, it immediately follows that their composition is a Markov operator as well. Using the Neumann series representation of ( Id -γ PT ) -1 it follows that for all f ≥ 0 ∈ B b ( X )

proving (1). Further,

showing that (1 -γ )( Id -γ PT ) -1 1 X = 1 X and proving (2). Finally, since (1 -γ )( Id -γ PT ) -1 and P are Markov operators, (3) and (4) follow from the direct application of [30, Theorem 19.2]. For the last point (5), let f ∈ B b (Ω) . As P is a conditional expectation operator, it holds that

Where the inequality is just the conditional version of Jensen's inequality [45, Chapter 5] applied on the (convex) ∥·∥ ∞ function, while the equality comes from the fact that TP is a Markov operator. Then, we have

Simulation Lemma . We report here the Simulation lemma, since it will be key to bridging the gap between Policy Mirror Descent and Conditional Mean Embeddings in Theorem 9 through Lemma 8.

Lemma A.6 (Simulation Lemma [31]-Lemma 2.2) . Let γ > 0 and let T 1 , T 2 two linear operators with operator norm strictly less than γ . Let P be a policy operator. Denote by q ( P , T ) = ( Id -γ TP ) -1 r the (generalized) action-value function associated to these terms and v ( P , T ) = P q ( P , T ) the corresponding value function. Then the following equality holds

Proof. Using the same technique of the proof of Lemma A.3 one has

where we have used fact that for any two invertible operators M and P it holds M -1 -P -1 = M -1 ( P -M ) P -1 for the second equation and applied the operatorial characterization of the value function to conclude the proof.

We then have the following result, which hinges on a generalization of the standard Simulation lemma in [31, Lemma 2.2] where we account also for the reward function to vary.

Corollary A.7. Let γ > 0 and let T 1 , T 2 two linear operators with operator norm strictly less than γ . Let r 1 and r 2 be two reward functions and P a policy operator. Denote by q ( P , T , r ) = ( Id -γ TP ) -1 r the (generalized) action-value function associated to these terms and v ( P , T , r ) = P q ( P , T , r ) the corresponding value function. Then the following equality holds

Proof. The difference between action-value functions can be written as

where we added and removed a term ( Id -γ T 1 P ) -1 r 2 . The result follows by plugging in the Simulation Lemma A.6 for the second term of the right hand side.

The corollary above will be useful in Appendix C to control the approximation error of the estimates ˆ q π t appearing in the convergence rates for inexact PMD in Theorem 7.

A.4 Action-value Estimator for ( G , F ) -compatible Policies

We can leverage the notation introduced in this section to prove the following form for the world model-based estimator of the action-value function.

Proposition 3. Let T n = S ∗ n BZ n ∈ HS ( F , G ) and r n = S ∗ n b ∈ G for respectively a B ∈ R n × n and b ∈ R n . Let π be ( G , F ) -compatible. Then,

where M π = Z n P π S ∗ n ∈ R n × n is the matrix with entries

Proof. By hypothesis

Eq. (12) follows by applying Lemma A.1. Eq. (13) can be verified by direct calculation. Denote by ( e i ) i m =1 the vectors of the canonical basis in R n . Then, for any i, j = 1 , . . . , n

Now, we recall that the two operators S n : G → R n and Z n : F → R n are the evaluation operators for respectively the points ( x i , a i ) n i =1 and ( x ′ i ) n i =1 . Namely, for any vector v ∈ R n

This implies that

Since P π is ( G , F ) -compatible by hypothesis, we can leverage the same reasoning used in Proposition 2 to show that

for any x ′ ∈ X . By plugging this equation in the previous characterization for ( M π ) ij we have

as required.

A.5 Separable Spaces

We show here the sufficient condition for ( G , F ) -compatibility of a policy in the case of the separable spaces introduced in Section 4.

Proposition 4 (Separable Spaces) . Let ϕ : X → H be a feature map into a Hilbert space H . Let F = H ⊗ H and G = R |A| ⊗ H with feature maps respectively φ ( x ) = ϕ ( x ) ⊗ ϕ ( x ) and ψ ( x, a ) = ϕ ( x ) ⊗ e a , with e a ∈ R |A| the one-hot encoding of action a ∈ A . Let π : X → ∆( A ) be a policy such that π ( a |· ) = ⟨ p a , ϕ ( · ) ⟩ with p a ∈ H for any a ∈ A . Then, π is ( G , F ) -compatible.

Proof. The proposition follows from observing that for any v ∈ R |A| and h ∈ H , applying P π according to (6) to the function g ( x, a ) = ⟨ h, ϕ ( x ) ⟩ ⟨ v, De a ⟩ yields

Hence ( P π g )( x ) = ⟨ f, φ ( x ) ⟩ with f = h ⊗ h ′ ∈ H ⊗ H = F and h ′ = ∑ a ∈A ⟨ v, De a ⟩ p a ∈ H . Therefore, the restriction of P π to G takes value in F as desired.

B Policy Mirror Descent

In this section we briefly review the tools needed to formulate the PMD method and discuss the convergence rates for inexact PMD. Most of the discussion follows the presentation in [12] formulated within the notation used in this work.

Let D : ∆( A ) × rint ∆( A ) → R a Bregman divergence [36, Definition 9.2] over the probability simplex, where rint ∆( A ) denotes the relative interior of ∆( A ) . In the following, for any t ∈ N we will denote by π t the policy produced at iteration t by a PMD algorithm according to the update (2) (with either the exact action-value function or an estimator, as discussed in Section 3) with divergence D and step-size η > 0 . We denote P t = P π t the associated operator. We recall here the PMD update from (2), highlighting the dependency on the policy operator P t via the action-value function q ( P t ) .

While this point-wise characterization is sufficient to define the updated policy π t +1 : X → ∆( A ) from the previous π t and its action-value function q ( P t ) , we need to guarantee that π t +1 is measurable. If that were not the case, we would not be able to guarantee the existence of a P t +1 associated with it, possibly affecting the well-definiteness of iteratively applying the mirror descent update (B.1). The following result addresses this issue.

Lemma B.1 (Measurability of the Mirror Descent updates) . Let D : ∆( A ) × rint ∆( A ) → R be a Bregman divergence continuous in its first argument. There exists a measurable policy π t +1 : X → ∆( A ) that satisfies (B.1) for all x ∈ X .

Proof. The proof follows from the Measurable Maximum Theorem [30, Theorem 18.19]. Let us denote f t : X × ∆( A ) → R the function

̸

Let also κ : X ↠ ∆( A ) be the constant correspondance x ↦→ ∆( A ) for all x ∈ X . κ clearly has nonempty compact values, and it is also weakly measurable since for any open set G ⊂ ∆( A ) , its lower inverse κ ℓ ( G ) := { x ∈ X : κ ( x ) ∩ G = ∅ } = X belongs to the Borel sigma-algebra of X . Finally, since q ( P t ) ∈ B b (Ω) , and by assumption D is continuous in its first argument, then we have that f t is a Carathéodory function. Then, by [30, Theorem 18.19] we have that the correspondence of minimizers µ : X ↠ ∆( A ) defined as

admits a measurable selector, which we denote π t +1 : X → ∆( A ) , proving the statement of the Lemma.

The previous Lemma is the key technical step enabling us to extend the convergence rates of Mirror Descent proved in [12] to non-tabular settings. We now state and prove fa ew Lemmas instrumental to prove Theorem 7.

Lemma B.2 (Three-points lemma) . Let π t +1 : X → ∆( A ) a measurable minimizer of (B.2) and P t +1 its associated operator. For every measurable policy π : X → ∆( A ) (alongside its associated operator P ) it holds

Proof. The function f t ( x, p ) in (B.2) is convex and differentiable in p as it is a sum of a linear function and a (strictly convex) Bregman divergence. By the first-order optimality condition [36, Corollay 3.68], a minimizer p ∗ of f t ( x, · ) satisfies, for all p ∈ ∆( A )

Since π t +1 ( x ) is a minimizer of f t ( x, · ) by assumption, letting p ∗ = π t +1 ( x ) , the first order optimality condition (B.4) becomes

Where in the first line we used the definition of Bregman divergence [36, Definition 9.2] D ( p ; q ) := ψ ( p ) -ψ ( q ) -⟨∇ ψ ( q ) , p -q ⟩ for a suitable Legendre function ψ : ∆( A ) → R , the first implication follows from the three-points property of Bregman divergences [36, Lemma 9.11], and the last implication from the positivity of D ( π t +1 ( x ); π t ( x )) .

Corollary B.3 (MD Iterations are monotonically increasing) . This Corollary is essentially a restatement of [12, Lemma 7]. Let ( P t ) t ∈ N be the sequence of policy operators associated to the measurable minimizers of (B.1) for all t ∈ N . For all x ∈ X it holds

i.e. the objective function is always increased by a mirror descent iteration. Further, if ˜ q ( P t ) ∈ B b (Ω) is such that ∥ q ( P t ) -˜ q ( P t ) ∥ ∞ ≤ ε t , then (B.5) holds inexactly on ˜ q ( P t ) as

Proof. By setting π ( x ) = π t ( x ) in (B.3), and recalling that D ( p ; q ) ≥ 0 with equality if and only if p = q , it follows that

giving (B.5). Integrating (B.5) over ( Id -γ P t +1 T ) -∗ ν and using the Performance Difference Lemma A.3 one gets (B.6). Finally, we get (B.7) from

Where the first inequality follows from (B.5), and the latter from the fact that policy operators are Markov operators and have norm 1, and ∥ ∥ P t +1 -P t ∥ ∥ ≤ ∥ ∥ P t +1 ∥ ∥ + ∥ ∥ P t ∥ ∥ = 2 .

B.1 Convergence rates of PMD

We are finally ready to prove the convergence rates for the Policy Mirror Descent algorithm (B.1). The proof technique is loosely based on [12, Theorem 8, Lemma 12], and extends them to the case of general state spaces through the key Lemma B.1 and using a fully operatorial formalism.

Theorem 7 (Convergenge of Inexact PMD) . Let ( π t ) t ∈ N be a sequence of policies generated by Algorithm 1 that are all ( G , F ) -compatible. If the action-value functions ˆ q π t are estimated with an error ∥ q π t -ˆ q π t ∥ ∞ ≤ ε t , the iterates of Algorithm 1 converge to the optimal policy as

where π ∗ : X → ∆( A ) is a measurable maximizer of (8) .

Proof. As usual, in this proof we denote the estimated and exact action-value functions as q n ( P t ) := ˆ q π t = ( Id -γ T n P t ) -1 r n and q ( P t ) := q π t = ( Id -γ TP t ) -1 r , respectively. From hypothesis, Algorithm 1 is well-defined since all policies it generates are ( G , F ) -compatible. The resulting sequence of policies ( π t ) t ∈ N are generated via the update rule (14) on the inexact action-value functions q n ( P t ) , as defined in (12). As the update rule (14) is a (measurable) minimizer of (B.1) when D equals the Kullback-Leibler divergence, the three-points Lemma B.2 with π ( x ) = π ∗ ( x ) yields

and

Adding and subtracting the term [( P ∗ -P t +1 ) q ( P t )] ( x ) , and bounding the remaining difference as [( P ∗ -P t +1 )( q ( P t ) -q n ( P t ))] ( x ) ≤ 2 ε t - see the derivation of (B.8) - one gets

Adding and subtracting P t q ( P t ) on the left side gives

and integrating with respect to the positive measure ( Id -γ P ∗ T ) -∗ ν and using the performance difference Lemma A.3 on the left hand side one has

where we used (A.8) on the right-hand-side terms. Since ( P t +1 -P t ) q ( P t ) + 2 ε t ≥ 0 because of (B.7), we can use fact (1) from Lemma A.5 with ( Id -γ P t +1 T ) -1 and the performance difference Lemma A.3 to get

Substituting this bound in (B.9) and summing from t = 0 . . . T -1 one gets to

Using facts (3) and (4) from Lemma A.5 we have that the terms ⟨ P q ( P ) , d ν ( P ∗ ) ⟩ on the right hand side can be bounded as

while -⟨ D ( π ∗ ; π T ) , d ν ( P ∗ ) ⟩ can be dropped due to the positivity of Bregman divergences yielding

Now notice that for all t < T it holds

J

(

P

t

) =

⟨

P

t

q

(

P

t

)

, ν

⟩

so that

Combining this with (B.10) we obtain

leading to the desired bound.

C POWRConvergence Rates

In this section, we prove the convergence of POWR. To do so, we need to first show that under the choice of spaces F and G proposed in this work, the resulting PMD iterations are well defined. Then, we need to bound the approximation error of the estimates for the action-value functions of the iterates produced by the inexact PDM algorithm, which appear in the rates of Theorem 7.

C.1 POWR is Well-defined

In order to guarantee that the iterations of POWR generate policies π t for which we can compute an estimator according to the formula in Proposition 3, we need to guarantee that all such policies are ( G , F ) -compatible. In particular, we restrict to the case of the separable spaces introduced in Proposition 4, for which it turns out that it is sufficient to show that all policies belong to the space H characterizing F = H⊗H and G = R |A| ⊗H . The following results provide a candidate for choosing such a space.

Theorem 5. Let X ⊂ R d be a compact set and let H = W 2 ,s ( X ) be the Sobolev space of smoothness s > 0 (see e.g. [38]). Let π t ( a |· ) and ˆ q π t ( · , a ) belong to H for any a ∈ A and π t ( a | x ) > 0 for any x ∈ X . Then the policy π t +1 solution to the PMD update in (14) belongs to H .

Proof. We recall that Sobolev spaces [38] over a compact subset X of R D are closed with respect to the operations of sum, multiplication, exponentiation or inversion (if the function is supported on the entire domain X ), namely for any two f, f ′ ∈ H , f + f ′ , f f ′ , e f ∈ H and, if f ( x ) > 0 for all x ∈ X , 1 /f ∈ H . This follows by applying the chain rule and the boundedness of derivatives over the compact X (see for instance [46, Lemma E.2.2]). The proof follows by observing that the one-step update π t +1 in (14) is expressed precisely in terms of these operations and the hypothesis that π t ( a |· ) and ˆ q π t ( · , a ) belong to H for any a ∈ A .

Combining the choice of space H according to the above result and combining with the PMD iterations of Algorithm 1 we have the following corollary.

Corollary 6. With the hypothesis of Proposition 4 let H = W 2 ,s ( X ) with s > d/ 2 . Let T n ∈ HS ( F , G ) and r n ∈ G characterized as in Proposition 3. Let π 0 ( a |· ) ∝ e ηq 0 ( · ,a ) for q 0 such that q 0 ( · , a ) ∈ H any a ∈ A . Then, for any t ∈ N the PMD iterates π t generated by Algorithm 1 are such that π t ( a |· ) ∈ H and hence are ( G , F ) -compatible.

Proof. We proceed by induction. Since ¯ q ( · , a ) ∈ H we can apply the same reasoning in Theorem 5 to guarantee that π 0 ( a |· ) ∈ H for any a ∈ A . Moreover, π 0 ( a | cot) > 0 for any a ∈ A since it is the (normalized) exponential of a function. Hence π 0 is ( G , F ) -compatible. Therefore, ˆ q 0 obtained according to Proposition 3 is well defined and belongs to G , implying ˆ q 0 ( · , a ) ∈ H for any a ∈ A . Now, assume by the inductive hypothesis that the policy π t ( a |· ) generated by POWR at time t and the corresponding estimator ˆ q π t ( · , a ) of the action value function belong to H and that π t ( a | x ) > 0 for any ( x, a ) ∈ Ω . Then, by Theorem 5 we have that also π t +1 the solution to the PMD update in (14) belongs to H (and is therefore ( G , F ) -compatible). Additionally, since π t +1 can be expressed as the softmax of a (finite) sum of functions in H , we have also π t +1 ( a | x ) > 0 for al ( x, a ) ∈ Ω , proving the inductive hypothesis and concluding the proof.

The above corollary guarantees us that if we are able to learn our estimates for the action-value function in H a suitably regular Sobolev space, then POWR is well-defined. This is a necessary condition to then being able to study its theoretical behavior in our main result.

C.2 Controlling the Action-value Estimation Error

Wenowshowhowto control the estimation error for the action-value function. We start by considering the following application of the (generalized) Simulation lemma in Corollary A.7.

Lemma 8 (Implications of the Simulation Lemma) . Let T n and r n the empirical estimators of the transfer operator T and reward function r as defined in Proposition 3, respectively. If T

satisfies Assumption 1, r ∈ G , and γ ∥ ∥ T n ∥ ∥ < γ ′ < 1 , then, for every ( G , F ) -compatible policy π

Proof. Recall that in the notation of these appendices, the action value of a policy and its estimator via the world model CME framework are denoted q π = q ( P ) and ˆ q π = q n ( P ) respectively. We can apply Corollary A.7 to obtain

Then, by Lemma A.5, point 5, we have

where v ( P ) := P ( Id -γ TP ) -1 r is the value function of the MDP, and we used that γ ∥ T n ∥ < γ ′ . Because of Assumption 1, r ∈ G , and P being ( G , F ) -compatible, it holds that v ( P ) ∈ F , while Proposition 3 implies r n ∈ G , and ( T -T n ) v ( P ) ∈ G as well. Therefore, using the reproducing property

where we assumed a bounded kernel ⟨ ψ ( x, a ) , ψ ( x, a ) ⟩ ≤ C ψ for all ( x, a ) ∈ Ω . Similarly, for the term depending on T n -T we have

Combining the previous two bounds, we get to

as desired.

According to the result above, we can control the approximation error for the action value function in terms of the approximation errors ∥ r n -r ∥ G and ∥ ∥ T | F -T n ∥ ∥ HS . This can be done by leveraging state-of-the-art statistical learning rates for the ridge regression and CME estimators from [40, 21, 47]. The following lemma connects Assumption 2 with the notation used in [40] which enables us to use the required result.

Lemma C.1 (Relation between (A.8) and [40]'s definition) . The following two facts are equivalent

g ∈ G satisfies the strong source condition Assumption 2 with parameter β on the probability distribution ρ .

2. g ∈ [ G ] 1+2 β ρ as in the notation of [40].

Proof. Using the same notations as in [40], we have

For 1 = ⇒ 2 we have that that for g ∈ [ G ] 1+2 β ρ

whose G -norm is ∥ ∥ C -β ρ g ∥ ∥ G = ∥ ∥ ( a i ) i ∈ N ∥ ∥ ℓ 2 < ∞ .

For 2 = ⇒ 1 , let g = ∑ i ∈ N b i µ 1 2 i e i with ∥ ∥ ( b i ) i ∈ N ∥ ∥ ℓ 2 < ∞ (since). Now, ∥ ∥ C -β ρ g ∥ ∥ G < ∞ is equivalent to ∥ ∥ ( b i µ -β i ) i ∈ N ∥ ∥ ℓ 2 < ∞ . By letting b i = µ β i a i we have that ∥ ∥ ( a i ) i ∈ N ∥ ∥ ℓ 2 < ∞ and that

that is g ∈ [ G ] 1+2 β ρ .

With the connection between [40] and Assumption 2 in place we can characterize the bound on the approximation error for the world model-based estimation of the action-value function.

Proposition C.2. Let T n and r n the empirical estimators of the transfer operator T and reward function r as defined in Proposition 3, respectively. When P is a ( G , F ) -compatible policy as in Definition 1 and the strong source condition Assumption 2 is attained with parameter β , it holds

with rates α ∈ ( β 2+2 β , β 1+2 β ) and probability not less than 1 -4 e -δ .

Proof. We use Lemma C.1 to apply Theorem 3.1 (ii) from [40] to show that under Assumption 2 with parameter β it holds, with probability not less than 1 -4 e -δ ,

The rate α r ∈ ( β 2+2 β , β 1+2 β ) is determined by the properties of the inclusion G ↪ → B b (Ω) , and the constant c r > 0 is independent of n and δ . Similarly, point (2.) of [21, Theorem 2] shows that under Assumption 2

again with probability not less than 1 -4 e -δ , rates α T ∈ ( β 2+2 β , β 1+2 β ) and with c T > 0 independent of n and δ . Combining every bound and denoting α := min( α r , α T ) , we conclude

C.3 Convergence Rates for POWR

With a bound on the estimation error of the action-value function by Algorithm 1, we are finally ready to state the complexity bounds for POWR.

Theorem 9. Let ( π t ) t ∈ N be a sequence of policies generated by Algorithm 1 in the same setting of Corollary 6. If the action-value functions ˆ q π t are estimated from a dataset ( x i , a i ; x ′ i ) n i =1 with ( x i , a i ) ∼ ρ ∈ P (Ω) such that Assumption 2 holds with parameter β , the iterates of Algorithm 1 converge to the optimal policy as

with probability not less than 1 -4 e -δ . Here, α ∈ ( β 2+2 β , β 1+2 β ) and π ∗ : X → ∆( A ) is a measurable maximizer of (8) .

Proof. Since the setting of Corollary 6 implies that P t are ( G , F ) -compatible for all t , and Assumption 2 is holding, then q ( P t ) and q n ( P t ) belong to G for all ( P t ) t ∈ N . This assures that we can use the statistical learning bounds Proposition C.2 into Theorem 7, yielding the final bound.

D Experimental details

D.1 Additional Results

In Fig. 2 we show the average timestep at which a reward threshold is met during the training phase. The testing environments are the same as introduced previously, with reward thresholds being the standard ones given in [26], except for the Taxi-v3 environment, where it is marginally lower. Interestingly, in this environment, only DQN and our algorithm are capable of achieving the original threshold within 1 . 5 × 10 6 timesteps during the training. On the other hand, the new lower threshold is also reached by the PPO algorithm.

Our approach attain the desired reward quicker than the competing algorithms. Furthermore, the timestep at which POWRreaches the threshold exhibits a lower variance compared to other techniques. This implies that our approach requires a stable amount of timesteps to learn how to solve a specific environment.

(a)

FrozenLake-v1

(b)

Taxi-v3

(c)

MountainCar-v0

2.0

eritep (x 10)

x) deg5e

X) de5

1.0

0.5

A2C

DQN

Algarithm

PPO

TRPO

POWR (0urs)

C

DQN

Algorithm

PPO

TRPO

C

DQN

Algorithm

PPO

TRPO

D.2 Other methods

We compare the performance of our algorithm with several baselines. In particular, we considered A2C [41], DQN [4], TRPO [7] and PPO [6], which we implemented using the stable baselines library [48]. We used the standard hyperparameters in [48].

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: we provide theoretical and experimental proofs of our claims.

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: see the conclusion section.

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: We give all the proofs in the Appendix.

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: In the paper, we provided the algorithm description and analysis, including the hyperparameters of the baselines and our proposed algorithm.

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: We provided the link to the code and all the hyperparameters we tested. 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: We provided a description of the open-source environments, in which we tested our POWR and the baselines.

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 showed seven independent runs for each experiment we reported to prove the effectiveness of our method. Moreover, in the plot, we can observe the bands based on the minimum and maximum values reached by the different runs.

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

Justification: We provided the code.

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 respected the ethics guidelines for the development of this work.

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: We believe that this work doesn't have any societal impact.

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: This work doesn't have 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 implemented our approach from scratch. We mentioned the baselines and the environments for testing involved in this work, that are open-source

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

Justification: We provided the code of our algorithm.

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: This work does not involve crowdsourcing nor research with human subjects.

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: This paper does not involve crowdsourcing nor research with human subjects.

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:\nThis paper extends policy mirror descent to the RL setting. They solve the problem of requiring a exact value function by formulating an approximate value function based on operators over the transition and reward functions, which yields a closed form solution. Then they use this approximate value function for policy improvement. Lastly, they provide theoretical analysis on convergence guarantees and error bounds while explicitly laying out assumptions. They also perform experiments on toy RL problems and show that their algorithm outperforms typical RL baselines on these low-dimensional problems.\n\nSTRENGTHS:\n1) This paper's presentation is great. How this work relates to prior work is clearly discussed. Key concepts from this paper are well explained and well motivated. \n2) Theoretical analysis is sound and shows promising results.\n3) Comparison to RL baselines, even on toy problems, is much appreciated. Results are encouraging. \n4) Discussions of limitations and assumptions are detailed. \n5) Key results of this paper, IE value functions with closed form solutions via operators, are novel and significant.\n\nWEAKNESSES:\n1) I would like to see experiments on more complicated MDPs, with large state-action spaces, to test the scaling properties. I recognize that this paper already makes large contributions, so this is not necessary for now, but it is a great future direction. \n2) I would like to see extensions to continuous action spaces. Again not necessary for now, but it seems like it should be possible given the formulation with additional approximations. \n\nMinor: \nTypo on line 335\n\nQUESTIONS:\nHow well can this algorithm scale to higher-dimensional settings?\n\nLIMITATIONS:\nAuthors have adequately addressed limitation and assumptions. \n\nI see no potential negative societal impacts as a result of this work.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes an approach (the first PMD approach adapted to RL setting) that can learn a world model using conditional mean embeddings (CME). The operatorial formulation of RL is used to express the action-value function in closed form via matrix operations. The proposed algorithm is proved to converge to global optimum, and validated by simple/preliminary experiments.\n\nSTRENGTHS:\nTo the best of my knowledge, this work is the first one to apply PMD to RL setting with provable global convergence guarantee.\n\nThe references and previous work are cited and well-discussed.\n\nThe theoretical contribution of this work is sound and solid.\n\nThe proposed method is validated to be more sample efficient than the current state-of-the-art RL algorithms like PPO, DQN, TRPO\n\nWEAKNESSES:\nLine 88: \"Borrowing from the convex optimization literature\", please give the exact reference.\n\nLine 97: at the end, lack of a \"space\"\n\nFrom Line 91 to the end of the manuscript..., it seems like the hyper-link of references does not work. They lead me to incorrect place, and make it hard to follow the idea.\n\nThe experimental validation is limited to simple/toy environments in OpenAI Gym.\n\nIn the caption of Figure 1, I would change \"dark lines\" to \"solid line\". The authors should also mention the reward threshold is the \"success\" reward in the caption, not just in the main texts. It is probably better to add the lable of dashed line in the plot.\n\nSince the Figure 1 is in logscale, the authors should show the learning curves starting from 0 or 10^1, so that we can observe the full learning process.\n\nQUESTIONS:\nIn Figure 1, it is a bit surprising that the variance of the proposed method in (a) and (b) is 0 (or almost 0). I know that general RL algorithms will still have some fluctuation. Maybe the authors can give some intuition why the proposed method achieves such small variance?\n\nIn Figure 1(c), why the proposed method stops at 10^4 step? I would expect to see the full learning curve, otherwise it is not convincing that the algorithm is convergent.\n\nDoes POWR work in continuous action space? Even simple tasks in Gym including Pendulum, MountainCarContinuous, LunarLanderContinuous, etc.\n\nLIMITATIONS:\nThe major limitation of this work is the empirical validation, where the authors only consider three toy/simple environments in OpenAI Gym.\n\nAlso, the experiments are limited to the discrete action space. It will be better to see if the proposed algorithm works in continuous action space like robotics, etc.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-13"}

comments: SUMMARY:
(This review has been updated after revert of desk reject)

STRENGTHS:
(updated after revert of desk reject)
It's a solid theoretical paper with nice writing and experimental results on continuous control benchmarks.

WEAKNESSES:
(updated after revert of desk reject)
The experimental results and more interpretation/justification of the results are limited. It would be stronger if the authors could further connect the experiments with the theoretical results.

QUESTIONS:
N/A

LIMITATIONS:
yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13

comments: SUMMARY:
The paper presents a practical implementation of Policy Mirror Descent (PMD) for Reinforcement Learning (RL). PMD requires knowledge of the action value function of the current policy at each iteration. Existing methods that approximate the action-value function depend on the ability to restart the Markov Decision Process (MDP) from each state multiple times. Instead, this paper proposes learning estimators of the reward function and the transition model and combining them to approximate the action-value function. To facilitate this, an RL formulation using operators is introduced, enabling the use of conditional mean embeddings for approximation. The authors study the convergence of the proposed algorithm to the global optimum and compare it with other RL methods on classic Gym environments.

STRENGTHS:
- The paper is well-written, and the authors thoroughly introduce their method in relation to existing techniques.

WEAKNESSES:
- Considering the connection between the proposed method and world models, the authors should discuss world models more thoroughly [1-3].

[1] Schmidhuber, J. Making the world differentiable: On using supervised learning recurrent neural
networks for dynamic reinforcement learning and planning in non-stationary environments. (1990)

[2] Schmidhuber, J. An on-line algorithm for dynamic reinforcement learning and planning in reactive environments. (1990)

[3] Ha, D., and Schmidhuber, J. World Models (2018)

QUESTIONS:
See above. 

The paper is rather technical, and several sections are beyond my expertise. I am looking forward to discussing it further with other reviewers to potentially increase the score for this submission.

LIMITATIONS:
Limitations are discussed in the paper.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-19