Operator World Models for Reinforcement Learning

Authors: pietro  novelli
marco  pratticò
massimiliano  pontil
carlo  ciliberto
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 19:27:33

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.

Main paper: 
 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] , which proposed an optimistic strategy to prove near-optimal regret bounds in linear mixture MDP settings [24] . 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 [25] 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.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. 
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 reward2 function r : Ω → R we denote byEQUATIONthe (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 setsthe 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 [26, 27] -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 stepEQUATIONfor any x ∈ X , with η > 0 a step size, D a suitable Bregman divergence [27] and q π : Ω → R the so-called action-value function of a policy π, see also [12, Sec. 4 ]. The action-value functionEQUATIONis 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 [28] 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 , thenEQUATIONThm. 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.ρ as in the notation of [39] .Proof. Using the same notations as in [39] , we haveg ∈ [G] 1+2β ρ ⇐⇒ g = i∈N a i µ 1 2 +β i e i and (a i ) i∈N ℓ 2 < ∞.1 =⇒ 2 we have that that for g ∈[G] 1+2β ρ C -β ρ g = i∈N a i µ 1 2 i e i (C.1) 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 g = i∈N a i µ 1 2 +β i e i , that is g ∈ [G] 1+2β ρ .With the connection between [39] 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 q n (P) -q(P)∞ ≤ O(δ 2 n -α ), (C.2)with rates α ∈ β 2+2β , β 1+2β and probability not less than 1 -4e -δ .Proof. We use Lemma C.1 to apply Theorem 3.1 (ii) from [39] to show that under Assumption 2 with parameter β it holds, with probability not less thanEQUATIONis 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 2T| F -T n HS(F ,G) ≤ δ 2 c T n -α T (C.4)again with probability not less than 1-4e -δ , rates α T ∈ β 2+2β , β 1+2β and with c T > 0 independent of n and δ. Combining every bound and denoting α := min(α r , α T ), we concludeq n (P) -q(P) ∞ ≤ δ 2 1 -γ ′ C ψ c r + γc T r ∞ 1 -γ n -α = O(δ 2 n -α ). (C.5)as required.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. 
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 :EQUATIONwhere X ′ is sampled according to τ (•|x, a). Note that T is the Markov operator [29, 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 operatorEQUATIONwhere 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)EQUATIONwhere the last equality follows from T and P π being Markov operators [29, Ch. 19 ] (∥T∥ = ∥P π ∥ = 1), making the Neumann series convergent. Analogously, we can reformulate the RL objective introduced in (1) as the pairingEQUATIONfor ν ∈ 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. EQUATIONProposition 2 shows that ( 9) is equivalent to the standard linear MDP assumption [30, 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 τ [31, 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] )EQUATIONWhen 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 toEQUATIONWe 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.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 weightsEQUATION)/* Policy Mirror Descent */ for t = 0, 1, . . . , T -1 do:π t+1 = SOFTMAX(ηH C t ) ∈ R n×|A| ▷ PMD Step (15) M πt+1 = H ⊙ (π t+1 E ⊤ ) ∈ R n×n▷ Proposition 3, Eq. ( 13)EQUATION)whereEQUATIONProposition 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 [32] . 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 [33, 34] .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.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes]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). 
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 [28] . As showed in [35, Example 9 .10], the solution to (2) can be written in closed form for any x ∈ X asEQUATIONwhere 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 sum of the action-value functions up to the current iterationEQUATION)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.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 modelT 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 [36] ). 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 entriesk(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). 
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. [37] ). 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. [38, 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 [38] . 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 ηq0(•,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 asEQUATIONwhere π * : 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 πqπ -q π ∞ ≤ 1 1 -γ ′ const ψ r n -r G + γ r ∞ 1 -γ T| F -T n HS .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 [30, Lemma 2.2] and then applying the triangular inequality for the uniform norm.POWR converges. 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 asJ(π * ) -J(π T ) ≤ O 1 T + δ 2 n -αwith probability not less than 1 -4e -δ . Here, α ∈EQUATIONand π * : 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 [39] 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. 
experimental results: 
 We empirically evaluated POWR on classical Gym environments [25] , 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 [40] , 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 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. 
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 [34] or reduced-rank regressors [41, 42] . 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.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, namely1 ∞ i=1 Bi = ∞ i=1 1 Bi 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 P1 Ω = 1 X , we haveπ(A|x) = (P 1 Ω )(x) = 1 X (x) = 1, (A.4)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 anyEQUATIONProof. 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⟨T| F f , ψ(x, a)⟩ = (T| F f )(x, a) = F f (x ′ ) τ (dx ′ |x, a) = F ⟨f, φ(x ′ )⟩ τ (dx ′ |x, a) = f, φ(x ′ ) τ (dx ′ |x, a) , from which we obtain ⟨f, (T| F ) * ψ(x, a)⟩ = f, φ(x ′ ) τ (dx ′ |x, a) .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). 
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. 
 
 Appendix: 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.q(P) = (Id -γTP) -1 r. (A.5)• Value function.v(P) = Pq(P).• Cumulative reward. The RL objective functional J(P) = P (Id -γTP) -1 r, ν = ⟨Pq(P), ν⟩ = ⟨v(P), ν⟩ .(A.7)• 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 asd ν (P) = (1 -γ) ∞ t=0 γ t (T * P * ) t ν = (1 -γ) (Id -γPT) - * ν, (A.8)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) impliesJ(P) = P (Id -γTP) -1 r, ν = Pr, (Id -γPT) - * ν = 1 1 -γ ⟨Pr, d ν (P)⟩ ,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 J(P 1 ) -J(P 2 ) = 1 1 -γ ⟨(P 1 -P 2 )q(P 2 ), d ν (P 1 )⟩ . (A.9)Proof. Using the definition of J(P 1 ) and Lemma A.1 one getsJ(P 1 ) -J(P 2 ) = P 1 (Id -γTP 1 ) -1 r, ν -P 2 (Id -γTP 2 ) -1 r, ν = (Id -γP 1 T) -1 P 1 r, ν -P 2 (Id -γTP 2 ) -1 r, ν = (Id -γP 1 T) -1 P 1 (Id -γTP 2 ) (Id -γTP 2 ) -1 r, ν -(Id -γP 1 T) -1 (Id -γP 1 T) P 2 (Id -γTP 2 ) -1 r, ν = (Id -γP 1 T) -1 [P 1 (Id -γTP 2 ) -(Id -γP 1 T) P 2 ] (Id -γTP 2 ) -1 r, ν = (Id -γP 1 T) -1 [P 1 -P 2 ] (Id -γTP 2 ) -1 r, ν = 1 1 -γ ⟨(P 1 -P 2 )q(P 2 ), d ν (P 1 )⟩ .A direct 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.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 lim h→0 J(P 1 + h(P 2 -P 1 )) -J(P 1 ) h = 1 1 -γ ⟨(P 2 -P 1 )q(P 1 ), d ν (P 1 )⟩ . (A.10)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 J(P h ) -J(P 1 ) = ⟨(P h -P 1 )q(P 1 ), d ν (P h )⟩ (A.11) = h ⟨(P 2 -P 1 )q(P 1 ), d ν (P h )⟩ .(A.12) We can therefore divide the above quantity by h and send h → 0. The result follows by observing thatd ν (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:1. For any f ≥ 0 ∈ B b (X ) it holds (Id -γPT) -1 f ≥ f . 2. The operator (1 -γ)(Id -γPT) -1 is a Markov operator. 3. For any positive measure ν ∈ M(B b (X )) it holds (1 -γ) ∥(Id -γPT) - * ν∥ TV = ∥ν∥ TV .4. For any positive measure ν ∈ M(B b (X )) it holds ∥P * ν∥ TV = ∥ν∥ TV . 5. 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 )EQUATION. Further,(Id -γPT) -1 1 X = ∞ t=0 γ t (PT) t 1 X = ∞ t=0 γ t 1 X = 1 X 1 -γ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 [29, Theorem 19.2] .For the last point ( 5), let f ∈ B b (Ω). As P is a conditional expectation operator, it holds that ∥XPf ∥ ∞ ≤ ∥X∥ P(1 Ω ∥f ∥ ∞ ) = ∥X∥ ∥f ∥ ∞ Where the inequality is just the conditional version of Jensen's inequality [44, Chapter 5] applied on the (convex) ∥•∥ ∞ function, while the equality comes from the fact that TP is a Markov operator. Then, we havesup ∥f ∥ ∞ =1 (Id -γXP) -1 f ∞ = sup ∥f ∥ ∞ =1 ∞ t=0 (γXP) t f ∞ ≤ sup ∥f ∥ ∞ =1 ∞ t=0 γ t (XP) t f ∞ ≤ sup ∥f ∥ ∞ =1 ∞ t=0 γ t X t f ∞ (∥f ∥ ∞ = 1) = 1 1 -γ X .Simulation Lemma. We report here the 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 [30] -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) = Pq(P, T) the corresponding value function. Then the following equality holds q(P, T 1 ) -q(P, T 2 ) = γ (Id -γT 1 P) -1 (T 2 -T 1 ) v(P, T 2 ) (A.13)Proof. Using the same technique of the proof of Lemma A.3 one has q(P, T 1 ) -q(P, T 2 ) = (Id -γT 1 P) -1 r -(Id -γT 2 P) -1 r = γ (Id -γT 1 P) -1 (T 2 -T 1 ) P (Id -γT 2 A) -1 r = γ (Id -γT 1 P) -1 (T 2 -T 1 ) v(P, T 2 )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 [30, 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) = Pq(P, T, r) the corresponding value function. Then the following equality holdsq(P, T 1 , r 1 ) -q(P, T 2 , r 2 ) = (Id -γT 1 P) -1 (r 1 -r 2 ) + γ(Id -γT 1 P) -1 (T 2 -T 1 )v(P, T 2 , r 2 ).Proof. The difference between action-value functions can be written as q(P, T 1 , r 1 ) -q(P, T 2 , r 2 ) = (Id -γT 1 P) -1 r 1 -(Id -γT 2 P) -1 r 2 = (Id -γT 1 P) -1 (r 1 -r 2 ) + (Id -γT 1 P) -1 -(Id -γT 2 P) -1 r 2where 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.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.EQUATION)where M π = Z n P π S * n ∈ R n×n is the matrix with entriesEQUATION)Proof. By hypothesis 12) follows by applying Lemma A.1. Eq. ( 13) can be verified by direct calculation. Denote by (e i ) m i=1 the vectors of the canonical basis in R n . Then, for any i, j = 1, . . . , n (M π ) ij = ⟨e i , M π e j ⟩ = ⟨e i , Z n P π S * n e j ⟩ = ⟨Z * n e i , P π S * n e j ⟩ . (A.15)EQUATIONNow, 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 nS * n v = n i=1 v i ψ(x i , a i ) and Z * n v = n i=1 v i φ(x ′ i ). (A.16)This implies that(M π ) ij = ⟨Z * n e i , P π S * n e j ⟩ = ⟨φ(x ′ i ), P π ψ(x j , a j )⟩ (A.17)Since P π is (G, F)-compatible by hypothesis, we can leverage the same reasoning used in Proposition 2 to show that(P π | G ) * φ(x ′ ) = A ψ(x ′ , a) π(da|x ′ ) (A.18)for any x ′ ∈ X . By plugging this equation in the previous characterization for (M π ) ij we have⟨φ(x ′ i ), P π ψ(x j , a j )⟩ = ⟨P * π φ(x ′ i ), ψ(x j , a j )⟩ (A.19) = A ψ(x ′ i , a) π(da|x ′ i ), ψ(x j , a j ) (A.20) = A ⟨ψ(x ′ i , a), ψ(x j , a j )⟩ π(da|x ′ i ), (A.21)as required.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(P π g)(x) = a∈A g(x, a)π(a|x) = ⟨h, ϕ(x)⟩ a∈A ⟨v, De a ⟩ π(a|x) (A.22) = ⟨h, ϕ(x)⟩ a∈A ⟨v, De a ⟩ ⟨p a , ϕ(x)⟩ (A.23) = h ⊗ a∈A ⟨v, De a ⟩ p a , ϕ(x) ⊗ ϕ(x) (A.24) 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.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 [35, 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 ).π t+1 (•|x) ∈ argmin p∈∆(A) -η a∈A q(P t )(x, a)p a + D(p; π t (•|x)) for all x ∈ X . (B.1)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 [29, Theorem 18.19 ]. Let us denote f t : X × ∆(A) → R the functionf t (x, p) := -η a∈A q(P t )(x, a)p a + D(p; π t (x)). (B.2)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 [29, Theorem 18.19] we have that the correspondence of minimizers µ : X ↠ ∆(A) defined asµ(x) := p * ∈ κ(x) : f t (x, p * ) = min p∈κ(x)f t (x, p) 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η [(P t+1 -P)q(P t )] (x) ≥ D(π(x); π t+1 (x)) -D(π(x); π t (x)) (B.3) 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 [35, Corollay 3 .68], a minimizer p * of f t (x, •) satisfies, for all p ∈ ∆(A)⟨∇f t (x, p * ), π(x) -p * ⟩ ≥ 0. (B.4)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η [(P t+1 -P)q(P t )] (x) -⟨∇ψ(π t (x)) -∇ψ(π t+1 (x)), π(x) -π t+1 (x)⟩ ≥ 0 =⇒ η [(P t+1 -P)q(P t )] (x) ≥ D(π(x); π t+1 (x)) -D(π(x); π t (x)) + D(π t+1 (x); π t (x)) =⇒ η [(P t+1 -P)q(P t )] (x) ≥ D(π(x); π t+1 (x)) -D(π(x); π t (x))in the first line we used the definition of Bregman divergence [35, 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 [35, Lemma 9.11] , and the last implication from the positivity of D(π t+1 (x); π t (x)).). 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 [(P t+1 -P t )q(P t )] (x) ≥ 0 (B.5) and J(P t+1 ) -J(P t ) ≥ 0 (B.6)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[(P t+1 -P t )q(P t )] (x) ≥ -2ε t . (B.7) Proof. By setting π(x) = π (x) in (B.3), and recalling that D(p; q) ≥ 0 with equality if and only if p = q, it follows that η [(P t+1 -P t )q(P t )] (x) ≥ D(π t (x); π t+1 (x)) ≥ 0, 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 [(P t+1 -P t )q(P t )] (x) = [(P t+1 -P t )q(P t )] (x) + [(P t+1 -P t )(q(P t ) -q(P t ))] (x)≥ [(P t+1 -P t )(q(P t ) -q(P t ))] (≥ -(P t+1 -P t )(q(P t ) -q(P t )) ∞ ≥ -P t+1 -P t q(P t ) -q(P t ) ∞ ≥ -2ε t .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.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 asEQUATIONwhere π * : 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[(P * -P t+1 )q n (P t )] (x) ≤ 1 η D(π * (x); π t (x)) - 1 η D(π * (x); π t+1 (x)).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[(P * -P t+1 )q(P t )] (x) ≤ 2ε t + 1 η D(π * (x); π t (x)) - 1 η D(π * (x); π t+1 (x)).Adding and subtracting P t q(P t ) on the left side gives [(P * -P t )q(P t )] (x) ≤ [(P t+1 -P t )q(P t )] (x)+ 2ε t + 1 η D(π * (x); π t (x)) - 1 η D(π * (x); π t+1 (x)),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 hasJ(P * ) -J(P t ) ≤ 1 1 -γ ⟨(P t+1 -P t )q(P t ) + 2ε t , d ν (P * )⟩ 1 η(1 -γ) ⟨D(π * ; π t ) -D(π * ; π t+1 ), d ν (P * )⟩ , (B.9)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 ⟨(P t+1 -P t )q(P t ) + 2ε t , d ν (P * )⟩ ≤ (Id -γP t+1 T) -1 [(P t+1 -P t )q(P t ) + 2ε t ], d ν (P * ) = ⟨P t+1 q(P t+1 ), d ν (P * )⟩ -⟨P t q(P t ), d ν (P * )⟩ + 2ε t 1 -γ .Substituting this bound in (B.9) and summing from t = 0 . . . T -1 one gets toT -1 t=0 J(P * ) -J(P t ) ≤ 1 1 -γ (⟨P T q(P T ), d ν (P * )⟩ -⟨P 0 q(P 0 ), d ν (P * )⟩) + 2 (1 -γ) 2 T -1 t=0 ε t 1 η(1 -γ) ⟨D(π * ; π 0 ) -D(π * ; π T ), d ν (P * )⟩ .Using facts ( 3) and ( 4 J(P * ) -J(P t ) ≤ 2 (1 -γ) 2 r ∞ + T -1 t=0 ε t + 1 η(1 -γ) ⟨D(π * ; π 0 ), d ν (P * )⟩ . (B.10)Now notice that for all t < T it holds J(P t ) = ⟨P t q(P t ), ν⟩ = ⟨P t q n (P t ), ν⟩ + ⟨P t (q(P t ) -q n (P t )), ν⟩ (Equation (B.6)) ≤ ⟨P T q n (P T ), ν⟩ + ⟨P t (q(P t ) -q n (P t )), ν⟩ = ⟨P T q n (P T ), ν⟩ + ⟨P t (q(P t ) -q n (P t )), ν⟩ + ⟨P T (q n (P T ) -q(P T )), ν⟩ ≤ J(P T ) + ε t + ε T , so thatJ(P * ) -J(P T ) ≤ ε T + 1 T T -1 t=0 J(P * ) -J(P t ) + ε t .Combining this with (B.10) we obtainJ(P * ) -J(P T ) ≤ ε T + 1 T 1 + 2 (1 -γ) 2 T -1 t=0 ε t + 2 r ∞ (1 -γ) 2 + 1 η(1 -γ) ⟨D(π * ; π 0 ), d ν (P * )⟩ ,leading to the desired bound.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-definedIn 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. [37] ). 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 [37] 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 [45, 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. ,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, q0 obtained according to Proposition 3 is well defined and belongs to G, implying q0 (•, 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.We now show how to 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, γ T n < γ ′ < 1, then, for every (G, F)-compatible policy πqπ -q π ∞ ≤ 1 1 -γ ′ const ψ r n -r G + γ r ∞ 1 -γ T| F -T n HS .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 q n (P) -q(P) = (Id -γT n P) -1 (r n -r) + γ(Id -γT n P) -1 (T -T n )v(P).Then, by Lemma A.5, point 5, we haveq n (P) -q(P) ∞ ≤ 1 1 -γ ′ r n -r ∞ + γ (T -T n )v(P) ∞ ,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 propertyr n -r ∞ = sup (x,a)∈Ω | ⟨ψ(x, a), r n -r⟩ G | ≤ ∥r n -r∥ G sup (x,a)∈Ω ∥ψ(x, a)∥ G = C ψ ∥r n -r∥ Gwhere 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(T -T n )v(P) ∞ = sup (x,a)∈Ω |[(T| F -T n )v(P)](x, a)| = sup (x,a)∈Ω | ⟨ψ(x, a), (T| F -T n )v(P)⟩ G | = sup (x,a)∈Ω Tr [(v(P) ⊗ ψ(x, a))(T| F -T n )] ≤ T| F -T n HS sup (x,a)∈Ω |v(P)(x, a)| ≤ ∥r∥ ∞ 1 -γ T| F -T n HS .Combining the previous two bounds, we get toq n (P) -q(P) ∞ ≤ 1 1 -γ ′ C ψ r n -r G + γ r ∞ 1 -γ T| F -T n HS ,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 [39, 21, 46] . The following lemma connects Assumption 2 with the notation used in [39] which enables us to use the required result. Lemma C.1 (Relation between (A.8) and [39] 's definition). The following two facts are equivalent 1. g ∈ G satisfies the strong source condition Assumption 2 with parameter β on the probability distribution ρ.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 asJ(π * ) -J(π T ) ≤ O 1 T + δ 2 n -αwith probability not less than 1 -4e -δ . Here, α ∈EQUATIONand π * : 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.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 [25] , 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. We compare the performance of our algorithm with several baselines. In particular, we considered A2C [40] , DQN [4] , TRPO [7] and PPO [6] , which we implemented using the stable baselines library [47] . We used the standard hyperparameters in [47] .Paper ChecklistCode available at: github.com/CSML-IIT-UCL/powr 38th Conference on Neural Information Processing Systems (NeurIPS2024).All the discussion in this work can be extended to the case where also the rewards are random and τ : Ω → P(X × R) takes values in the space of joint distributions over states and rewards (Xt+1, Rt) 


Decision: The paper presents a practical implementation of Policy Mirror Descent (PMD) for Reinforcement Learning (RL), which ordinarily requires knowledge of the action value function of the current policy at each iteration. This paper proposes learning estimators of the reward function and the transition model and combining them to approximate the action-value function. 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 and compare it with other RL methods on classic Gym environments.
Overall all reviewers are quite enthusiastic about the paper and I am glad to recommend acceptance.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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: Limitations are discussed in the paper.
paper_summary: paper_summary: (This review has been updated after revert of desk reject)
strength_weakness: strength_weakness: (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: questions: N/A
limitations: limitations: yes
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: Line 88: "Borrowing from the convex optimization literature", please give the exact reference.

Line 97: at the end, lack of a "space"

From 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.

The experimental validation is limited to simple/toy environments in OpenAI Gym.

In 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.

Since 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.
questions: questions: In 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?

In 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.

Does POWR work in continuous action space? Even simple tasks in Gym including Pendulum, MountainCarContinuous, LunarLanderContinuous, etc.
limitations: limitations: The major limitation of this work is the empirical validation, where the authors only consider three toy/simple environments in OpenAI Gym.

Also, 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.
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: 1) 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. 
2) 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. 

Minor: 
Typo on line 335
questions: questions: How well can this algorithm scale to higher-dimensional settings?
limitations: limitations: Authors have adequately addressed limitation and assumptions. 

I see no potential negative societal impacts as a result of this work.

Comments: title: Rebuttal2
Official response
Rebuttal
comment: We thank the reviewer for their feedback.

We appreciate the reviewer’s suggestion to enhance the connection between our experimental results and the theoretical analysis. In the experimental section (lines 367/368), we have already mentioned that our method, POWR, converges to a global optimum (as demonstrated in Theorems 7 and 9), and exhibits smaller sample complexity compared to other baselines, which aligns with the expected behavior of world models. According to the reviewer’s suggestion, we will make these connections clearer.

Additionally, we would like to emphasize that the main contributions of our paper are theoretical. In particular, many highly influential works in this area such as [10, 11, 12] do not include an experimental section at all. Nonetheless, we included experiments to provide additional support for our theoretical claims. While the experimental section serves as a supplementary validation, our primary focus remains on the theoretical advancements presented in the paper.
Thanks authors for the reminder. I updated the review and it's short, but I'm holding (weakly) a positive opinion.
Thanks for your clarification. 
Based on the discussions so far, most of my concerns have been (or will be) addressed. At this stage, I will increase the score. Thanks!
We thank the reviewer for the feedback and for raising their original score. We will update the introduction to accurately reflect the prior work on world models in RL, as the reviewer suggested.
We thank the reviewer for the prompt reply. 

Yes, we will align with the other baselines in the final version of the paper, running experiments up to 10^6 steps, as stated in the attached PDF.

Given the theoretical nature of this work and the limited time available, we have not thoroughly optimized our implementation of POWR, which resulted in longer runtimes. While there is significant room for optimization—and we are excited to advance POWR further— we would like to highlight that the main focus of this work was to prove the novel theoretical contributions of Proposition 3, Corollary 7, and Theorem 9. Our experiments demonstrate POWR's statistical efficiency, achieving success at least an order of magnitude faster than competitors.
Thank you for your response. Please note that the claim "The notion of world models for RL was introduced by Ha and Schmidhuber in [R3]" is incorrect. World models in RL were already introduced in Reference R1 (1990) and in Dyna ([4], 1991).

I trust that the authors will accurately discuss the background of world models. Having read the other reviews and the authors' responses, I am happy to increase the score for this submission.

[4] Sutton, R., Dyna, an integrated architecture for learning, planning, and reacting. (1991)
Thank you for providing the updated Figure 1. It indeed looks more reasonable. 

But is there a reason not to run 10^6 timesteps for MountainCar? I understand that rebuttal has limited time, but MountainCar is a very easy problem, which should not take that long. Will the authors plot 10^6 timesteps in the final version of the paper?
Thank you for your comments.