Operator World Models for Reinforcement Learning
Authors: Pietro Novelli
Marco Pratticò
Massimiliano Pontil
Carlo Ciliberto
Venue: NeurIPS 2024 poster (2024)
Keywords: Reinforcement Learning
Transfer Operators
World Models
Policy Gradient
Conditional Mean Embeddings
Mirror Descent
Primary area: 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.

Operator World Models for Reinforcement Learning
Pietro Novelli Istituto Italiano di Tecnologia pietro. novelli@iit.it
Marco Pratticò Istituto Italiano Tecnologia marco. pratticooiti..ittt..............
Massimiliano Pontil Istituto Italiano di Tecnologia AI Centre, University College London massimilian iano. pontilliiiittt..
Carlo Ciliberto
Al Centre, University College London cilibertoouucc.c. ac.
Abstract
Policy Mirror Descent (PMD) a powerful and theoretically sound methodology for sequential decision-making. However, it is not directly applicable to Reinforce- ment Learning (RL. due to inaccessibility of explicit aetion-value functions. We address this challenge by introducing novel approach based on learning world model the environment using conditional mean embeddings. Leveraging tools from operator theory derive 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, prove convergence rates to the global optimum. Preliminary experiments in finite and infinite state settings support the effecttivnness of our method'.
RL algorithm for which we
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 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] name 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 global optimum with associated rates [10. 11, 12].
In their original formulation, PMD methods require explicit knowledge of the action-value func- tions for all policies generated during the optimization process. This clearly inaccessible in RL applications. Recently, [12] showed how PMD convergence rates can be extended settings in which inexact estimators of the action-value function are used (see [13] for a similar result from 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 the availability of a (perfect) simulator to queried arbitrary state-action pairs. Moreover, these strategies require repeating this sampling/learning process for any policy generated by the PMD algorithm, which computationally expensive and demands numerous interactions with the
'Code available at: github. www.CCSMLIIT--ULL/AAAAAA
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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 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, 21J. 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 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 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 policy in closed form via matrix operations. wi An (inexact) PMD algorithm combining the learned CMES world models with mirror descent update steps to generate improved policies. iči) Showing that the algorithm is well-defined when learning the world model as an operator between a suitable family of Sobolev spaces. ěv) 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 7. We assume X' and A to be Polish, 7 : Ω + P(2') be Borel measurable function from the joint space Ω = X x A to the space P(X') of Borel probability measures on X. We define a policy be a Borel measurable function ग X - P(A). When A (respectively X) is a finite set. the space P(A) (AA) S R! (respectively P(x)) = A(X) C R|XI) corresponds to the probability simplex. Given discount factor y > 0. an initial state distribution // E P(2') and a Borel measurable bounded and non-negative reward function r : + RR we denote by
DO
( = EE,TTT
Σ
Er'r(X.,A.))
AAll the discussion in this work can be extended to the case where also the rewards are random and T + P(X x R) takes values in the space of joint distributions over states and rewards (Xr+1,

---

the (discounted) expected return of the policy x applied to the MDP, yielding the Markov process A()EN, where Xn distributed according to v and for each & N the action A, distributed according to x (-)x)) and Xt+1 according to r(-[X,, A.).
In sequential decision settings, the goal is to find the optimal policy n, maximizing (1) over the space of all measurable policies. In reinforcement learning, one typically assumes that knowledge of the transition r, 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 (24, at) and corresponding rewards r r ae) and transitions 2(+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 IT = (AA) @ RR of column substochastic matrices, namely matrices M E RlA|x| with non-negative entries and whose columns sum up to one, namely M* = 1x, 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 convex functional over 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], the cost of dimension-dependentt constants.
Starting from an initial policy no. PMD generates a sequence (TT))en according to the update step -17 (9m. (+,2), p) + D(p, (+(x=)).
(-11 [z) = argmin rEA(A)
(2)
for any x E X, with 7) > a step size, D a suitable Bregman divergence [28] and fr Ω - the so-called action-value function of policy æ, see also [12, Sec. 4]. The action-value function
(z, a) E
Σ 1==
(rr(X, A.) Xo z, Ao =
(3)
is the discounted return obtained by taking action A in state x E 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 4r, associated to each iterate of the algorithm. This requires evaluating the expectation (3), which not possible in RL where we do not know the reward r and MDP transition distribution T in advance. While sampling strategies can be adopted to estimate q 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 âz, 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 Sec. 5, where we extend it to infinite state spaces X.
Theorem 1 (Inexact PMD (Sec. 5 in [12] - Informal), In the tabular setting, ler (TT) be a sequence of policies obtained by applying the PMD update in (2) with functions qn, Ω + R in
place of qT, and D a suitable Bregman divergence. For any T E IN and > 0, if [1 - -x. -2.1loo S for all I 1,...,? T, then
max ((T) ((T)- ((TT) s O(E + 1/T)).
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 uniform norm Il While this result supports estimating the action-value function in RL, implementing this strategy in practice poses two main challenges, even tabular settings. First, approximating the expectation in (3) in II-Il,o norm via sampling requires "starting" the MDP from each state X, multiple times. This
3

---

is often not possible in RL, where we do not have control over the starting distribution v. Second, repeating this sampling process to learn a qT. for each policy Ti can become extremely expensive in terms of both the number of computations and interactions with the environment.
In this work, we propose new strategy tackle the problems above. Instead of re-sampling the MDP to estimate qa, at each iteration t, we learn estimators f and f for the reward and transition distribution, respectively. For any policy R, we then leverage the relation between these quantities in (3) to generate an estimator for Øx. This approach tackles the above challenges since 1) it enables us control the approximation error on any valuee function in terms of the approximation error of é and f; 2) it does not require sampling the MDP to learn a new âx, for each na generated by PMD.
3 Operator World Models
In this section, we present an operator- based formulation of the problem introduced in Sec. (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 (X) denote the space of bounded Borel measurable functions on a space X. Formally, T: (x' + (2t) is the linear operator such that, for any (X)
(Tf)
((xxx()) T (dr' (2,a) a) = [f(x') 1x, [
for all (æ, a) E ?
(5)
where X' is sampled according to 1o [-1,,, a). Note that T the Markov operator [30, Ch. 19] encoding the dynamics of the MDP and its conjugate T M(2) - MM(X) is the operator mapping signed
Borel measures js E ((2) to their transition via t as (B) SBxO =(da' [r, a) ()x(dx da) for
any measurable B C X. For any policy æ we define the operator p, (22 + (2y) such that for
all g E (2)
(o) = / ((x, a) r(dalz) = E Lo(X, A) X = x] for all x E X,
(6)
where the expectation taken over the action A sampled according T(-1a). Also P, is Markov operator and its conjugate P -M(X) - ((2) is the operator mapping any v E MMX)) to its joint S2.
measure with π. namely (P)))(C) = Je π (da (d)) for any measurable
Operator Formulation of RL. With these two operators in place, we can characterize the expected re- ward after single interaction between a policy and the MDP as (T,, a) F[r(X', X') x, An = a). This observation can be applied recursively, yielding the operatorial characterization of the action-value function from (3)
xc I (x(2 = ==>==>=TE (E, (,),, A.) x, Ao a (GTP)))r - (ld - TPp)) -1r, (7)
where the last equality follows from T and P. being Markov operators [30, Ch. 19] ( 1), making the Neumann series convergent. Analogously, we can reformulate the RL. objective introduced in (1) as the pairing
P
J(TT) = (Tpm
-
-Pa
(8)
for v E 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 -0--77 -Tp>) This characterization motivates us to learn and from data and then express any action-value function the interaction of these two terms with the policy π as in (7), rather than learning each 9. 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 suitable family

---

of functions. Let : X' + F and 3 9 - G two feature maps with values into the Hilbert spaces FF 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 (X)) and B2( of the form f(x) = (J, xxx))) and g(z,a) = (g, (at, a)) for any f F and g E and any (x. a) E Q. We say that the linear MDP assumption holds with respect to (4,,)) if
Assumption 1 (Linear MDP- - Well-specified CME). The restriction of T to F is a Hilbert-Schmids operator [| E ())
In CME settings, the assumption above is known as requiring the CME of to be well well-specified. The following result, proved in Appendix A.2, clarifies this aspect and establishes the relation of Asm. with the standard definition of linear MDP.
Proposition 2 (Well-specified CME). Under Asm. 1, (T|=)* = (T*)lg and, for any (I, a) E E
(Tl=)" a) / p(x') (<21,,, a) = [(x(()) = x, A A= dj.
Proposition 2 shows that (9) is equivalent to the standard linear MDP assumption [31, Ch. 8] when X is a finite set (taking yp the one-hot encoding) while being weaker in infinite settings. From the CME perspective, the proposition characterizes the action of (Tlz)" 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 learning strategy: (9) characterizes the action of T as evaluating the conditional expectation of a vector (XX'') given (2r, a). Given set of points (II,, az);== and corresponding x' sampled from r(-(2;,, (a),,,)), this can be learned by minimizing the squared loss, yielding the estimator (see [15, Sec 4.2))
T n
argmin TEHS(F.9) ni
lko(x!) TT (a,,,,,,, a:)) + A A = SKK Zn.
(10)
When F and G are finite dimensional, Sn, and Zn matrices with rows, each corresponding respectively to the vectors (I,, a.) and (xx(() for = 1.....,, n. In the infinite setting, they generalize to operators S, G - RD and Zn, F - R". The matrix K₃ S, + ndld, IRexn is the regularized Gram (or kernel) matrix with (i, entry corresponding to
(Kx) = (zb(zz, (2, a.), (2 az)) + ndd:j.
(11)
We conclude our discussion 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, and following the standard practice supervised settings, we can learn an estimator for solving a problem akin to
(10). This yields a function of the form Fn = S b E by (012,, a:) as the linear combination of the embedded training points with the entries of the vector b = K y where y E Rd is the vector with entries y/ = r(z,, az).
Estimating the Action-value Function 4x We now propose our strategy to generate an estimator for the action-valus function 4r of 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 FG.
Definition 1 (GGFF) comppaiibblll A policy π X P(A) is compatible with two subspaces F. C (x) and G C (2) if the restriction Px to G is a bounded linear operator with range C F, that is (Px)lg g) G + F.
Definition is analogous to the linear MDP Asm. in that it requires the restriction of Par to G to take values in the associated space F. However, it 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 wide range of policies (in particular those generated by our POWR method) is (G, F))ooppatbble for our choice of function spaces. Definition is the key condition that enables us to generate an estimator for 9r, as characterized by the following result.
5

---

Algorithm POWR: POLICY MIRROR DESCENT WITH OPERATOR WORLD-MODELS FOR RL
Input: Dataset (I,, ai o4 r:) discount factor E (0, 1), step size > 0, kernel function k(r, x²) = (6(x), with Φ X' + H as in Proposition 4, initial weights Co = E Rax|A
/* World Model Learning */
let E E Rnx|Al with rows E, = OONEHOTIAA( (o)o)).
let X E Rax" such that Kij = k(Ti, I.) ( )a,=a, + nddsj
> Eq. (11)
let E R such that Hij = =k(2', xzg)
compute x and = xxx y with y = (Fi,,..., ra) E E
> Eq. (10)
/* Policy Mirror Descent */
for 1. = 0, 1, . T 1 do:
T++1 = SOFTMAX(1HCI) E RRXXAA
> PMD Step (15)
M = H o E R=x*
Proposition 3, Eq. (13)
Ct+1 = Ci+diig(() with == - (M...)
Proposition 3, Eq. (12)
end for
return कт : X + ((A) such that (2) =SOFTMAX(T HL, with HI, (b(2,,1)) E R".
- Proposition 3. Let - B2, E (S(F, 9) and rn Sgb E 9 for respectively B E Raxn and bER" R". Let be ccompaabll Then,
hn (Id -0,p7
S7 -BM,)
(12)
where M = S x is the matrix with entries
S a) =ddalx!).
(Ma) = (f(x!),, Pa (
(13)
Proposition 3 leverages a kernel trick argument to express the estimator for the action-value function of π as the linear combination 9- i= G (0(2;, az) of the (embedded) training points (, a:) and the entries e; of the vector = (ld BBM,) E RR" We prove the result in Appendix A.4. We note that in (12) both B and M₂ are n x n matrices and therefore the characterization of qnr amounts to solving projection methods such as Nyström approximation to learn Tn and rn [33]. These strategies have been recently shown significantly reduce the computational load of learning while retaining the same empirical and theoretical performance as their non-approximated counterparts [34, 35].
linear system. For settings where n is large, one adopt random
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 Ps 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| < do 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 Ôx, of the action-valle function 4r, 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- Leible divergence DKL(P: p/) = Eae A Pa log(Pa/PL, 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 E X' as
the update of (2). was
Te. (e (2.-)
me+ Te+ (1xx) ΣOEA (e(a(z)e (x,a))
(14)
where we used the estimated action-value function dn from Proposition 3. Additionally, the formula above can be applied recursively, expressing T4+1 the softmax operator applied to the discounted

---

sum of the action-valle functions up to the current iteration
T(41(-1z)) = SOFTMAX
(log(To(11))) + + Σ 9r. (z,
(15)
a=0
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 & to perform world model learning. Specifically, to apply Proposition 3 and obtain proper estimators 9.., we need to guarantee that all policies generated by the PMD update (14) are (G. comopptblll (Definition 1). The following result describes a suitable family of such spaces.
Proposition 4 (Separable Spaces), Let o : H. be a feature map into a Hilbert space H.
Let J = H @ H and 9 = RIAA @ H. with feature maps respectively (p(x) = 6(x) @ ((x) and (ba, a) - ((x) @ ca. with ea E (RIA) the one-hot encoding of action a E A. Let = - ((A) be
policy such that π ((al-) = (pa: (())) with Pa E H for any E A. Then, is (g, -ccoppablle...
Proposition 4 (proof in Appendix A.5) states that for the specific choice of function spaces F = HOH and 9 = RIA &H, we can guarantee F.,.)ppaatabbbiii provided that H. is rich enough to "contain" all : :((al-) for
A. We postpone the discussion on identifying a suitable spaces H for PMD to Sec. 5, since ()) -oooppatiiii is not needed to mechanically apply Proposition 3 and obtain an estimator /.. This is because (12) exploits kernel-trick and rather requires only to be able to evaluate x the training data. The latter is possible for tc+1, thanks to the explicit form of the PMD update in (14). We can therefore present our algorithm.
bypass the need to know p explicitly
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 initial phase, learn the world model Tn S K Zn and the reward On = b fitting the conditional mean embedding described in (10) on a dataset (I;, a, x/. re);-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 E fr, to obtain the policy (-+1(- by (inexact) PMD) via the softmax operator in (15). We use the newly obtained policy to compute the matrix M defined in (13), which is a key component to obtain the estimator a for (8 We note that in the case of the separable spaces of Proposition 4, this matrix reduces to the n n matrix with entriesk (x², z,) ) Te+1 (az) (rr)), where k(21, xz) = (o(x)) o(z;)) is the kernel matrix between initial and evolved states. Finally, we obtain c = (Id K M)- Proposition 3.
bb and model 0 S according to
Clearly, the world model learning and PMD phases can be alternated in POWR, essentially finding 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 [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, and rn 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 Well-defined. To properly apply Proposition 3 to estimate the action- value function of any PMD iterate ni 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.
7

---

Theorem 5. Let X C. Rd be compact set and let - (222)(2)) be the Sobolev space of smoothness > 0 (see e.g. [38]). Let x = and (..(:.a) a) belong to for any E A and i.e [a(a)z) for any E X. Then the policy T1+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-generaie policies. This observation is further supported by the fact that Sobolev spaces of smoothness > 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)),((z)) defines a so-called reproducing kernel, namely a positive definite function that is (usually) efficient evaluate, even if ((x)) is high or infinite dimensional. For example, H (2')(())) with [ has associated kernel k(x; x²) = e-1la-='11 with bandwidth a > 0 [39]. By applying Thm. 5 to the iterates generated by Alg. 1 we have the following result.
k(r, x')
Corollary 6. With the hypothesis of Proposition 4 let H. = w with s > d/2. Let Tn E HS(F, 9) and In E 9 characterized as in Proposition 3. Let To(al-) x e) for qo such that qu(-.a) a) E H any a E A. Then. any E N the PMD iterates ne generated by Alg. I are such that [e(al-) E HH and hence are (G.
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.I.
Inexact PMD Converges. We now present 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). Lel (Te)eeN be a sequence of policies generated by Alg. 1 that are all (g, F)-copattible If the action-value functions qn, are estimated with an error 11gx. gz. ol S Et, the iterates Alg. converge to the optimal policy as
-m.)) S ET +0
(16)
t=
where To X - (AA) is a measurable maximizer of (8).
Thm. 7 shows that inexact PMD behave comparably to its exact version provided that 1) the action value functions âr, are estimated with increasing accuracy, and 2) that the sequence of policies is (g, F) coompatibb,,, for example in the Sobolev-based setting of Cor. 6. Specifically, if l1g=, f1. Il < ((/L)) for any L E N, the convergence rate of inexact PMD is of order Olog T IT), only logarithmic factor slower than exact PMD. This means that we do necessarily need 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 proof in Appendix B.1 since the original result did not allow for 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. 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 Tn. Tn and P,.. The exact qa has analogous closed-form characterization in terms of T, r and Pa 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 Tn and ra the empirical estimators of the transfer operator T and reward function r as defined in Proposition 3, respectively. If satisfies Asm. 1, E G. and TTll < < < 1, then, for every (G.F)-compaaibl policy n
const Ir - 2116 + ?1|l|| 00 ||T
11 - 9r 5 1-7
T Tn
8

---

In the result above, when applied to a function in G. such as I'. the uniform norm is to be interpreted
as the uniform norm of the evaluation of such function, namely || (a.(a. I (rn, (2 a)) I.
and analogously for Tn. The proof, reported in Appendix C.2. follows by first decomposing the difference 9. ब- with the simulation lemma [31, 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,, (and r,,) is made of i.i.d. triplets (æi, di s() with (1;, a;) sampled from a distribution p E P((2) supported on all f (such as the state occupancy measure (see e.g. [11] or Appendix A.3) of the uniform policy <--lar) = Al) and ar sampled from 7(- (z4 0z). 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 E P(L)) and Co the covariance operator EvEA Jx p(dx,, a)b((r @ v0(1,, a). The transition operator and the reward function r are such that Th E HS(F,9) and r E G. Further, [|(T)=))" some 3 > 0.
pppppllll "| < co and "0""rl < oo for
Asm. 2 imposes 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 finite dimensional (e.g. in the tabular RL setting) and imposes additional smoothness on T and when belonging to Sobolev space. Equipped with this assumption, we can now state the convergence theorem for Alg. 1.
Theorem 9. Let (TT, tEN be sequence of policies generated by Alg. 1 in the same setting of Cor. 6. If the action-value functions qs. are estimated from a dataset (I,, ag xf)i with (xi, a4) ~ p E P((2) such that Asm. 2 holds with parameter B, the iterates of Alg. I converge to the optimal policy
[(T.) - J(TT) S 50
( +
with probability not less than 1 - 5 Here, a E 24228 1423 measurable maximizer of (8).
B and T. X' + A(A) is a
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. -oompptible 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 eтror 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 L2(S), μi) and would only require Asm. 1 (Linear MDP) to hold. Rather, we need to impose Asm. to guarantee faster rates in uniform norm.
Experimental results
We empirically evaluated POWR on classical Gym environments [26]. ranging from discrete wrrzzenLakee-,,Txxi---a to continuous state spaces (MountainCar-vO) 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. for a number of steps to generate 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 (rimesteps 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

---

-100
203
-19
199
TR3-
PRD
POWR 10unt
-39
3V'
Timesoap (ogscale)
Timestap Joguce
Timerop: (lopucal))
(a) FrozenLake-vi
(b) Tazi-v3
XountainCar-vo
The plots show the average cumulative reward different environments with respect to the timesteps (i.e. number of interactions with MDP). The dark lines represent the mean of the cumulative reward and the shaded area is the minimum and maximum values reached across 7 independent runs. The horizontal dashed lines represent the reward threshold proposed by the Gym library [26].
Figure 1:
environment, it avoids converging 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 arguably due to 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 novel reinforcement learning (RL) algorithm leveraging these results. Our approach operates in two, possibly alternating, phases: learning world model and planning via PMD. During exploration, we utilize conditional mean embeddings (CMEs) to learn world model operator, showing that this procedure is well-po well-posed when performed over suitable Sobolev spaces. The planning phase involves PMD steps for which guarantee convergence to global optimum at polynomial rate under specific MDP regularities.
Our analysis opens avenues for further exploration. Firstly, extending PMD to infinite action spaces remains 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, whether alternating world model learning with inexact PMD updates benefits the oppooaatio-exxplotatiii
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 -enerreedd 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 0233220003000000 "Future Artificial Intelligence Research (FAIR)", EU grant ELISE (GA no 951847) and EU Project ELIAS (GA no 101 120237).
References
[1] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[2] OpenAl: Marein Andrychowicz, Bowen Baker, Maciek Chociej, Rafal Jozefowicz, Bob Mc- Grew, 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.
10

---

[3] David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schritwieser. Ioannis Antonoglou, Veda Panneershelvam. Mare Lanc- tot, et al. Mastering the game of go with deep neural networks and tree search. Nature, 529(7587)4444449) 2016.
[4] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv1312.5602, 2013.
[5] Hongzi Mao, Mohammad Alizadeh, Ishai Menache, and Srikanth Kandula. Resource man- agement with deep reinforcement learning. In Proceedings of the 15th ACM workshop on hot topics in networks, pages 50-56, 2016.
[6] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv1700.00337, 2017.
[7] John Schulman, Sergey Levine, Philipp Moritz, Michael Jordan, and Pieter Abbeel. Trust region policy optimization, 2017.
[8] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor, 2018.
[9] Dimitri Bertsekas. Dynamic Programming and Optimal Control: Volume I, volume 4. Athena scientific, 2012.
[10] 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 Confer- ence on Artificial Intelligence, volume 34, pages 5668-5675, 2020.
[11] 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.
[12] Lin Xiao. On the convergence rates of policy gradient methods. Journal of Machine Learning Research, 23(282): 2022.
[13] 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.
[14] Kenji Fukumizu, Francis R. Bach, and Michael I. Jordan. Dimensionality reduction for super- vised learning with reproducing kernel Hilbert spaces. Journal of Machine Learning Research, 5:73 99, 2004.
[15] 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.
[16] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. Ad- vances in neural information processing systems, 31, 2018.
[17] Richard S Sutton. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 2(4):160-166 1991.
[18] Jürgen Schmidhuber. Making the world differentiable: on using self supervised recurrent neural networks for dynamic reinforcement learning and planning in stationat environ- ments, volume 126. Inst. für Informatik, 1990.
[19] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301 04 104 2023.
[20] Carlo Ciliberto, Lorenzo Rosasco, and Alessandro Rudi. A general framework for consistent structured prediction with implicit loss embeddings. Journal of Machine Learning Research, 2 21(98):1-67, 2020.
11

---

[21] 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.
[22] 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.
[23] Preston Rozwood, Edward Mehrez, Ludger Paehler, Wen Sun, and Steven I. Brunton, Koopman- assisted reinforcement learning. arXiv preprint urXiv:2403.022900 2024.
[24] Antoine Moulin and Gergely Neu. Optimistic planning by regularized dynamic programming.
In International Conference on Machine Learning, pages 25337-25357. PMLR, 2023.
[25] Alex Ayoub, Zeyu Jia, Csaba Szepesvari, Mengdi Wang, and Lin Yang. Model-based rein- forcement learning with value-targeted regression. International Conference on Machine Learning, pages 463-474. PMLR, 2020.
[26] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arxiv. arXiv preprint 16X1::16000015540, 10, 2016.
[27] Amir Beck and Mare Teboulle. Mirror descent and nonlinear projected subgradient methods for convex optimization. Operations Research Letters, 31(3):167-175, 2003.
[28] Sébastien Bubeck. Convex optimization: Algorithms and complexity. Foundations and Trends@ in Machine Learning, 8(3-4):231-357, 2015.
[29] Sham M. Kakade. A natural policy gradient. Advances in Neural Information Processing Systems, 14, 2001.
[30] C.D. Aliprantis and K.C. Border. Infinite Dimensional Analysis: A Hitchhiker's Guide. Studies Economic Theory. Springer, 1999.
[31] Alekh Agarwal, Nan Jiang, Sham M. Kakade, and Wen Sun. Reinforcement learning: Theory and algorithms. 2021. URL https://tthheersbookkkkk wwt.rrrbboooooooob...... io, 2022.
[32] Steffen Grünewälder, Guy Lever, Luca Baldassarre, Sam Patterson, Arthur Gretton, and Mas- similano Pontil. Conditional mean embeddings regressors. In Proceedings of the 29th International Conference on International Conference on Machine Learning, pages 1803-1810, 2012.
[33] Christopher Williams and Matthias Seeger. Using the Nyström method to speed up kernel machines. Advances in Neural Information Processing Systems, 13, 2000.
[34] Alessandro Rudi, Raffaello Camoriano, and Lorenzo Rosasco. Less is more: Nyström computa- tional regularization. Advances in Neural Information Processing Systems, 28, 2015.
[35] 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.
[36] Amir Beck. First-Order Methods in Optimization. Society for Industrial and Applied Mathe- matics, 2017.
[37] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Bellemare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, al. Human-level control through deep reinforcement learning. Nature, 51887540):5299533, 2015.
[38] Robert A. Adams and John J. F. Fournier. Sobolev Spaces. Elsevier, 2003.
[39] Holger Wendland, Scattered data approximation, volume 17. Cambridge University Press, 2004.
12

---

[40] Simon Fischer and Ingo Steinwart. Sobolev norm learning rates for regularized least-squaree algorithms. Journal of Machine Learning Research, 21 21(20)):1-38, 2020.
[41] Volodymyr Mnih, Adria Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lill- icrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. arXiv1602.01 07716000..17,,.. 2016.
[42] Vladimir R. Kostic, Pietro Novelli, Andreas Maurer, Carlo Ciliberto, Lorenzo Rosasco, and Massimiliano Pontil. Learning dynamical systems via Koopman operator regression in reproduc- ing kernel Hilbert spaces. Advances in Neural Information Processing Systems, 35:4017-4031, 2022.
[43] Giacomo Turri, Vladimir Kostic, Pietro Novelli, and Massimiliano Pontil. randomized algorithm to solve reduced rank operator regression. arXiv preprint arXiv:2312.17348, 2023.
[44] Gene H. Golub and Charles F. Van Loan. Matrix Computations. Johns Hopkins University Press, 2013.
[45] O. Kallenberg. Foundations of Modern Probability. Probability and Its Applications. Springer New York, 2002.
[46] 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.
[47] Vladimir Kostic, Karim Lounici, Pietro Novelli, and Massimiliano Pontil. Sharp spectral rates for koopman operator learning. Advances in Neural Information Processing Systems, 36, 2023.
[48] Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. bable-baselines:: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268):1-8, 2021.
13

---

Appendix
The appendices organized as follows:
Appendix discuss the operatorial formulation of RI and show how to derive the operator- based results in this work.
Appendix focuses on policy mirror descent (PMD) and its convergence rate 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 this work.
A Operatorial Results
A.1 Auxiliary Lemma
We recall here a corollary of the Sherman- Sherman-Woodbury identity [44]. Lemma A.1. Let A and B two conformable linear operators such that (I+ + AB) is invertible.
Then (I + AB) A A(I + BA)
Proof. The result is obvious if A is invertible. More generally, we consider the following two applications of the Sherman-Woodbury Woodbury [44] formula
(I + AB) = A(T + B
(A.1)
and
+I+ + BA)" II (1+ BA) 'BA.
(A.2)
Multiplying the two equations by respectively the right and the left, we obtain the desired result.
A.2 Markoy operators and their properties
We recall here the notion of Markov operators, which is central for 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. bounded linear operator (,B(()) is a Markov operator if positive and maps the unit function to itself. that is:
a. f 2 BB(x) = Pf 0 Bo()).
b. Plx =y,
where 1x X (respectively }).
(respectively ly) denotes the function taking constant value equal to on X
We recall that Markov operators are a convex subset of ((B,(X), B&())). Here we denote this space as LM (X),)), (J))). Direct inspection of (5) and (6) shows that the transfer operator T associated an MDP and the policy operator P, associated to policy æ are both Markov operators.
Markov Operators and Policy Operators. In (6) we defined the policy operator P associated to policy It turns out that the converse is also true, namely that any such Markov operator policy operator.
Proposition A.2. Let e (M(BB(X). (S)) be a Markov operator. Then there exists TP, such that the associated policy operator corresponds to P, namely P = P.
Proof. Define the map Tp X' + M(A) taking value in the space of bounded Borel measures over A such that, for any X' and any B A Borel measurable subset
Tp (B).) = (p1xxB)
(A.3)
14

---

We need to guarantee that for every E X' the function xp is signed measure. To show this, first note that the operation defined by ит well-defined, since for any measurable set B the function 1xxB is also measurable, making xp [3] well defined as well. Moreover, since 1ø(a) = 0 for any a E A, it implies that xxx0 = 0 and therefore Tp (p(())) 2z) = 0 for any x E X. Finally, aadditiiviy follows from the definition of indicator functions, namely U B,
0c i=1 1B. for any family of TP (B.lz) for any x E X.
pair-wise disjoint sets (B,) M= which implies тe (U [. -
We now apply the two properties Markov operators to show that тт takes values in P(A), namely it is non-negative measure that sums to 1. Since Markov operators map non-negative functions in non-negan neganeeetie functions and since xxx5 2 0 for any S X. we have (-|xx) 2 0 as well for any E X. Moreover, since Ω = X' x A and P1g 1x, have
x aA)r=) (P (x) 1x(x)=1 1,
(A.4)
for any E X. Therefore (-(- is probability measure for any E X. Direct application of (6) shows that the associated policy operator corresponds P, namely Pmp P as desired.
Given the correspondence between policies and their Markov operator according to (6) and Proposi- tion 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- -peeiiiee CME), Under Assumption 1, (Tl=)* (T*) g and, for any (æ, a) E Ω
(Ti=)* / p(x)) = A=
Proof. Recall that since they are Hilbert spaces F and G = G* are isometric to their dual and therefore we can interpret any f E as the function f(-) = (-)) with some abuse of notation, where clear from context. By Assumption I we have that T - takes values in G. This means that (T[=f) E 9 or, in other words
(TI=f. = (T)=f)(z,a)
J, f(x') (dx'(zz,a
= /s (F,((()))) (d)
= (5: / (((x) r(dx'|z,, 12)))
from which we obtain
(TI=)) 5, sox²) r(de' (1,)))).)))
Since the above equality holds for any f E F (9) holds, as desired.
We note that the result can be extended to the setting where (Fx C 9, namely the image of T| contained in G, namely sort of (F,9) (F.9)-compatibiliy for the transition operator (see Definition 1).
A.3 The operatorial formulation of RT.
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)
(A.5)
15

---

Value function.
(P) Pq(P).
(A.6)
Cumulative reward. The RL objective functional
J(P) (P (ld - TTP) = (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 we can represent the evolution state distribution Pt at time & to the next state distribution Mt+1 T*P*,, v4. Applying this relation recursively, we recover the state visitation probability associated to the starting state distribution Po V P(X), the MDP with transition T and the policy P as
(P) (1
-1-7) (T*P") = (1 -) 7) (ld -PT) v,
(A.8)
where the 7 convex combination of the probability distributions Vr, hence guaranteeing du (P) to be well-defined (namely it belongs to (x))).
2 is normalizing factor guarantee that the series corresponds
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 few of these results that will be useful in the following.
Remark A.1. Algebraic manipulation of the cumulative expected reward (P) implies
J(P) = (P (Id - -TP) rr,u) = (Pr, (Id -))
1- - (Pr, (P)))
where used Lemma A.I and (P) is the state visitation distribution starting from and following the policy P.
The following result, known Performance Difference Lemma [see e.g. 11, Lemma 1.16), will be instrumental to prove the convergence rates for PMD in Theorem
Lemma A.3 (Performance difference). Let P1, P2 two policy operators. The The following equality holds
J(P.) J(P2) 1 1-y ((P1 d.))())).
(A.9)
Proof.
Using the definition of J(P;) and Lemma A. one gets
J(P,) - J(22) = (Pi (ld -TP,; v, (Pz (ld - TTPp) -1r
(ld - 7P, Pir, v) (P2 (ld (ald~>TPP)) -Pp))_ T.P)
(ld - ~PTT) -1 Pı (ld TP2) TP:)) T.P)
((ld ([ld-7P11 TP (ld
TP2)) rr,))
(Idd 7PiT) - [P] (Id->TPP (Id->TPp) - (ld (ld->PIT) P2] (Id TPp2) r.,) (r,))
(ld 7P.T) 1 [P - P2] (ld -TPP r,,
1 ((P, -
direct consequence of the operator formulation the performance difference lemma is the following operator- based characterization of the differential behavior of the RI objective. The result can be found in [11] for the case of finite state and action spaces, however here the operatorial formulation allows for much more concise proof.
16

---

Corollary A.4 (Directional derivatives). For any two Markov P1.P2 Bb(x)) - BL))),, we have that the directional derivative in P1 towards P2 J(P + h(P2 - P1)) - I(P1) 1-7 ((P
lim h00
(A.10)
h
Proof. The result follows by recalling that the space of Markov operators is convex, namely for any h E [0, 1] the term PΛ = Pi + h(P2 - P1) is still a Markov operator. Therefore, we can apply Lemma A.3 to obtain
-(P))-((PP) ((P) P1)
(A.11)
= - ((P2 - Pi ))(()()))))))))) d,
(A.12)
We can therefore divide the above quantity by and send > The result follows by observing
that de (Px)) = (Id - TTPx) >->>)>>> - (Id TPP)) = (Pi) for h - O. since |P - for any
E [0, 1] and the function M + (I - TM) continuous on the open ball of radius 1/7 > in ((), (x()) with respect to the operator norm.
Properties of (Id - -PT)-1 The quantity (Id - -PT) (note, not (Id - TP) plays a central role in the study of POWR. We prove here few properties that will be useful in the following.
Lemma A.5 (Properties of ld PPT). The following facts are true:
1. For any 2 0 (x)) it holds (Id -PT)-
2. The operator (1 (ld -PT) is Markov operator.
3. For any positive measure .M (MB(() it holds (1 7) [I(ld - -PP)) = ||
4. For any positive measure u E (MBB(A)) it holds [|P* = ||
5. For any bounded linear operator X. policy operator P and discount factor < |X]|, it holds (ldd -i---XPP-111 '1 S 1/(1 1 [X
Proof. Since both T and P are Markov operators by construction, it immediately follows that their composition is Markov operator as well. Using the Neumann series representation of (Id -PT) it follows that for all f 2 0 E (x))
00
(Id (dd-PPT)-1
Σ
7(PT)'s 1+
(PTT'
proving (1). Further,
1x
(Id 7PT) 1x (PT)
Σ
1x
1-7
t=U
1=0
showing that (1- (PPT) and P are Markov operators, (3) and (4) follow from the direct application [30, Theorem 19.2]. For the last point (5), let (22). As P is conditional expectation operator, it holds that || I |X|| P(1e [[flllo) = ||x|| Il
Lx and proving (2). Finally, since (1 - -PT)-
Where the inequality is just the conditional version of Jensen's inequality [45, Chapter 5] applied on the (convex) I| function, while the equality comes from the fact that TP is a Markov operator. Then, we have
sup
(ld 7XP)
sup [I
(1XP)')Il
t=0
(XP()))) 8
sup
|1
t=0
*||||||||||
$ sup
| I,- =1 t=
(11fl 1)
1
17

---

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). Letry> > O and let 11 T2 two linear operators with operator norm strictly less than 7 Let P be policy operator. Denote by q(P: T) (Id r the (generalized) action-value function associated to these terms and u(P, T) = Pq(P.T) the corresponding value function. Then the following equality holds
yTP)
(P, T1) - T2) = -idd-TTy P (T2 - TT)
(A.13)
Proof. Using the same technique the proof Lemma A.3 one has
(P, T1) <PP. T2) (Id >TPP) r (Id TT2P)
= (Id - PTP -1 (T2 T1) (Id -2244 r
Y (Id - T (T Tz) u
(P)TT) T2)
where we have used fact that for any two invertible operators M and P it holds M-1 - p-1 M- - 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 account also for the reward function to vary.
Corollary A.7. Let > O and let 1, To two linear operators with operator norm strictly less than Letri and 72 be two reward functions and a policy operator. Denote by q(P, T.r) = TP) 1 the (generalized) action-valas value function associated to these terms and v(P, T,r) PT( the corresponding value function. Then the following equality holds
qP,, T1. - q(P, T2, = (Id TT,P) -2) r2) + +ld - 7T,P) -T- T2.
Proof. The difference between action-value functions can be written as qP,,T (1. - <P, T2, = (ld - >T, (Id TTpP) (ld T P) (r) r2) + [Ild PTP (ld -T2P)
where we added and removed term (Id > P) 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 to control the approximation error of the estimates dn, appearing in the convergence rates for inexact PMD in Theorem 7.
A.4 Action-value Estimator for -cooppablll Policies
We can leverage the notation introduced this section to prove the following form for the world model-based estimator of the action-value function.
Proposition 3. Ler T, B2, E *S(F,,) 9) and r,, S G for respectively B E Rax= and b E Ra Let π be (G, F) -compatbbll... Then,
fn = (Id - -T,p)) = S,, - -BM,)-
(12)
where M = 2n St E |R is the matrix with entries
(Mn) = (o(x)), Pa,((,, / (1(2), =(dalx)).
(13)
Proof. By hypothesis
in = (ld - 7Tdd- 1dd=
(A.14)
18

---

Eq. (12) follows by applying Lemma A.1. Eq. (13) can be verified by direct calculation. Denote by (e,) the vectors of the canonical basis in R Then, for any i. .2
(Ma) (ei, e2) (ei, Zn S (Z, e,))
(A.15)
Now, we recall that the two operators Sp 9 + R" and Zn F R" the evaluation operators for respectively the points (Ii, (ii and (a/))==- Namely, for any vector E RR
n
and
Z2"
(pp(x
(A.16)
This implies that
(Mx)s (Z = (v(x)), Pa (2),,
(A.17)
Since Px is
by hypothesis, leverage the same reasoning used in Proposi-
tion 2 to show that
(P*la)><('') = - v((',a) =
(A.18)
for any x' E X. By plugging this equation in the previous characterization for (Mx)ss we have
(p(x)), Pax (Z) az)) = (P)))())), ((g,,
(A.19)
/ (((), r(da 24) ((x))))
(A.20)
az)) t(dalz/),
(A.21)
required.
A.5 Separable Spaces
We show here the sufficient condition for commpuiiily 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 2 H. and G = RIAI @ H. with feature maps respectively sp(x) = <(x) @ 6(x) and (x a) <(x) @ with ca E RIAI the one-hot encoding of action E A. Let X - ((A) be
a policy such that N (al-)) (Pa: o(-)) with Pa E foror any E A. Then, (9, F))comaatiblle
Proof. The proposition follows from observing that for any v E RIAI and E H. applying P according to (6) to the function g(z, a) (h, o(x)) (v, Dea) yields
(++))(a)= = Σ ()( = (h, o(x))
Σ
(v, Dea) T(alx)
(A.22)
BEA
aEd
(h, o(x))
Σ
(v, Dea) (pa o(x))
(A.23)
aEd
(v. Dea) Pa; M(x) s obxx)
(A.24)
h Hence (P=9)(z) = (zx)) with h V' E H @ H. = and EeEA (v, Dea) Pa E H. Therefore, the restriction of P. G takes value in J 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 [12] formulated within the notation used in this work.
19

---

Let D : A(A) x (in) - R a Bregman divergence [36. Definition 9.2] over the probability simplex, where (iA) denotes the relative interior of ((A). In the following, for any t we will denote by и. the policy produced at iteration /. by 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 1 > O. We denote P, Pst. the associated operator. We recall here the PMD update from (2), highlighting the dependency on the policy operator Pt via the action-value function q(p,).
q(Pr) + D(ps }
(++1(122) E argmin PEA(A) 7 Σ AEAA
for all x E X.
(B.1)
While this point-wise characterization is sufficient to define the updated policy *++1 : X' - ((A) from the previous Ty and its action-vall value function ((Pt), we need to guarantee that TT41 is measurable. If that were not the case, we would not be able to guarantee the existence of Pr+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(A) x rint (AA) - R be Bregman divergence continuous in its first argument. There exists measurable policy T++1 : X - (AA) that satisfies for all x E X.
Proof. The proof follows from the Measurable Maximum Theorem [30, Theorem 18.19). Let us denote ft X A(A) + R the function
ft(x,)) = Σ
(P))(2, + D(pr aEA
(B.2)
Let also K : - A(A) be the constant correspondaace + ((A) for all x X. clearly has nonempty compact values, and it is also weakly measurable since for any open set G C ((A), its lower inverse x² (G. == {i E X: x(x) n & Ø) = belongs to the Borel sigma-algebra of X. Finally, since ((P,) B, (S2). and by assumption D is continuous in its first argument, then we have that Se a Carathéodory function. Then, by [30, Theorem 18.19] we have that the correspondence of minimizers μ.: X + A(A) defined as
m(x) = {{p- E (px))(x) f((..p)) min fe(2,))
pÉx(z)
admits a measurable selector, which we denote T441 X A(A), proving the statement of the Lemma.
The previous Lemma the key technical step enabling us 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
Lemma B.2 (Three-points lemma). Let met1 X
((A) measurable minimizer of (B.2) and Po+1 its associated operator. For every measurable policy : X + ((A) (alongside its associated
operator P) it holds
n [P (iP=11-))((P]] (x) ((T(())) (()))))))))))))))))))))))mmm))))
D(T(())
(B.3)
Proof. The function fe(x.)) p) in (B.2) is convex and differentiable in as it is a sum of a linear function and (strictly convex) Bregman divergence. By the first-order optimality condition [36, Corollay 3.68], a minimizer of fe(z,)) satisfies, for all p E A(4)
(Vie(z, T<x) p-) 2 0.
(B.4)
Since 1 (a) is a minimizer of fi(x,-) by assumption, letting p. 7141 (x), the first order optimality condition (B.4) becomes
7 [PP++1 - (P)) (x) - (Vv(a(()))
- ( ( # 20 =
n (P PP)(P,)) (z) 2 (=)))
7 (Pt+1 - P))(P:)) (x) 2 (z) D((
20

---

Where in the first line we used the definition of Bregman divergence [36. Definition 9.2] D(p; :=
(p()) - (4)) - (Ve(g)) q) for suitable Legendre function & A(A) the first implication follows from the three-points property of Bregman divergences [36, Lemma 9.11], and the last implication from the positivity of (1)
Corollary B.3 (MD Iterations are monotonically increasing). This Corollary is essentially restate- ment [12, Lemma 7J. Let (Pt)eeN be the sequence of policy operators associated to the measurable minimizers of (B.1)for for all & N. For all a E X it holds
[P1+1 - Pt)q(PP)) (x) 20
(B.5)
and
J(P(1))-(PP))
- 2(P))
(B.6)
Lee. the objective function functioon is always increased by mirror descent iteration, Further, if i(;)) E B₂(
such that [[q(P:) ((pr) St, then (B.5) holds inexactly on @(P.) as
[+1 Pc) -2Et-
(B.7)
Proof. By setting : x(x) = Te (1) in (B.3), and recalling that D(p: q) 2 with equality if and only p it follows that
[(P:+1 Pr)a (x) 2 ((+1 2 0,
giving (B.5). Integrating (B.5) over (ld - v and using the Performance Differ-
ence Lemma A.3 one gets (B.6). Finally, we get (B.7) from
[P - (.))(P.) (x) [(P++1 P.)g(P:)] (1) - ((P)))
z [(P++ Px)(ύ(P.) (P)))
z (Pt+1 Pt)( ((.)))
(B.8)
2 - |P++1 Pr]l]l((P:) ((P))l
25,-
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 Pall s P + |P||| =
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. and using fully operatorial formalism.
Theorem 7 (Convergenge of Inexact PMD). Let (Te)ee be sequence of policies generated by Algorithm that are all ccoopptblle If the action-value functions qa. are estimated with an error |xx, qxc lloo S Es, the iterates of Algorithm I converge to the optimal policy
(
T-1 +
- -(FT) S er +0
-
ct
(16)
1=
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 qn (P,) fr. = (ld - - P,) and gPp)) = 4n. = (Id - TP.) -AP.))-a 'r, respectively. From hypothesis, Algorithm 1 is well-defined since all policies it generates sequence of policies (T) are generated via the update rule (14) on the inexact action-value functions gn. (P.), defined in (12). As the update rule (14) (measurable) minimizer of (B.1) when D equals the Kullback-Leibe Leibler divergence, the three-points Lemma B.2 with ((x) = n. (x) yields
(G. -)copppaiibe The resulting
[PP - Pc+1 (z) 1 D(T - 1 D( (())))))
21

---

Adding and subtracting the term [P. P(+1)q(Pe)) (z), and bounding the remaining difference as [(P. Pt+1)a((P) 9n, (P.))) (z) _ 2e, see the derivation of (B.8) one gets
[(P. Pr+1) (z) - 2s, + - D(T.
- - DTT (T)
1
Adding and subtracting (Pt(P:) on the left side gives
1 Dm (1); ((T. (2);; 1+)1
[PP. - P:)) (a:) S [(P++ P.))(P:)) (1)) 28,
11
"
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
J(P)) (Pt))
((P44 P.)q(P.) + 22 dy (P-))
(B.9)
711 7) (D(T.: - ( (, (P.)),
where used (A.8) on the right-hand-side terms. Since (P++1 (P.) 2s, 2 0 because of (B.7), we can use fact (1) from Lemma A.5 with (Id - -P T) and the performance differ- ence Lemma A.3 get
((P++1 PP)) + 25, (P.))) & ((ld - :P -Idd-pP++:
[(P (P.) + 22r22, dz (P (P.)) - (P) dv +
2e,
1-7
Substituting this bound in (B.9) and summing from t = 0... T one gets to
Σ J(P.) J(P.) - (PPr((Pr), P +
1
2
(1 71- 1=0
711->) 7) (D(T.: D(T.; d (P.))
Using facts (3) and (4) from Lemma A.5 we have that the terms d on the right hand side can be bounded as
dP (P(ld - - (P.))
<(Id-
PPP) dd,
(Duality) < ||rl
(Lemma A.5 and ||l| = 1) ||||
1-7'
while (D(Tri m (P.)) can be dropped due to the positivity of Bregman divergences yielding T-1
T-1
Σ -P..)- - (1 -1-1) 7)2
|1
Σ
Et
7111 (1(1 -7) (D(T=; mo), (P.)). (B.10)
Now notice that for all f < T it holds
((P,) = (P.g(P:),
(Prn((P.): + (Pr(o(P:) 4n (P:)), v)
(Equation (B.6)) < (Prqn(P (Pr),)) + (Pr(g(P:) - qn (P.),,
(Pr (Pr),)) + (Pi(g(P,) (P.)), v) (Pr(un(PT) -((r))) -Pr(u)(P))
< J (PT) t + ET,
so that
1 T
J(P.) - J(Pr) S Ser +
J(P.) - J(P,) t
T 1=0l
Combining this with (B.10) obtain
T-1 Σ 0=
2 2||1||
J(P.) J(Pr) En + T leading to the desired bound.
1+ (1 7
E, +
(D(T.) mo) dx.)))
(1 (1-2)²
n(1 7)
22

---

C POWR Convergence 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 proposed in this work, the resulting PMD iterations 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 m for which we can compute an estimator according to the formula in Proposition 3, we need to guarantee that all such policies are ccompatlll 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 = RIAy & H. The following results provide candidate for choosing such space.
Theorem 5. Let X C RR be compact set and let H (X) be the Sobolev space of smoothness > 0 (see e.g. [38]). Let ne(a)-) and qm, (-, a) belong to for any E A and T((alz) > 0or any X. Then the policy T++1 solution to the PMD update in (14) belongs to H.
Proof. We recall that Sobolev spaces [38] over compact subset of RD 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. H. f + ff', et H and, if f(x) > 0 for all E X. 1/f E 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 in (14) is expressed precisely in terms of these operations and the hypothesis that (a)-) and Gr, (.,a) belong to H for any a E A.
Combining the choice of space H according to the above result and combining with the PMD iterations of Algorithm I we have the following corollary.
Corollary 6. With the hypothesis of Proposition 4 let H. = (22 with d/2. Let Tn E G characterized as in Proposition 3. Let no(al-) ) x en for go such that qu(:,a) a) E H. any a E A. Then. for any t E N the PMD iterates ne generated by Algorithm are such that (al-)) E H. and hence are -)ooppatllb
HS(F G) and
Proof. We proceed by induction. Since q(-, a) E H to guarantee that no(al.) E H for any a E A. Moreover, Ta cot) > for any a E A since is the (normalized) exponential of function. Hence mo is (G, -Cooppaille Therefore, du obtained according to Proposition is well defined and belongs to G, implying go(: a) E H. for any a E A. Now, assume by the inductive hypothesis that the policy n. (al-)) generated by POWR at time and the corresponding estimator dz, a) of the action value function belong to H and that n Te(a]z) > O for any (2r, a) E Ω. Then, by Theorem 5 we have that also _.-1 the solution the PMD update in (14) belongs to H (and is therefore (Compaccall Additionally, since π2+1 can be expressed as the softmax of a (finite) sum of functions in H. we have also (a(x) > 0 for al (z, a) Ω. proving the inductive hypothesis and concluding the proof.
apply the same reasoning in Theorem 5
The above corollary guarantees us that if we are able to learn our estimates for the action-value function in H. 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
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 Tn and ra the empirical estimators of the transfer operator T and reward function r defined in Proposition 3, respectively. If T
23

---

satisfies Assumption 1, r E G, and T||Tx|| < <n' < 1, then, for every (G, FG. -coompatible policy π
1 1-76 constt//nn-r/yy + ?||r|| 1-y || -T Tn
11 97 Irl1xx
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 y+ <(P) and fr yn (P) respectively. We can apply Corollary A.7 to obtain
In (P) (P) (Id -0,,TTT,77 (r r) -
TT, P TP.)
Then, by Lemma A.5, point 5, we have
4n (P) - (P)) S 1
- [|r -1 + -T)((T-T))(())ll))]
where c(P) = P(ld - TTP)) the value function of the MDP, and we used that 7 |
y'. Because of Assumption 1, G, and P being (G. -ccomppalle,l it holds that v(P) E F, while Proposition 3 implies In E 9, and (T (T- Ta)((P) E 9 as well. Therefore, using the reproducing property
I' - -l
sup (1(2,)) - r)g I Ir,, وااs (x,a)ER
[, a)1l = -Cv lr rllg (z,a)E0
sup
where we assumed bounded kernel (1ixx, a), 1(,, a)) s Cp for all (x,0) E Q. Similarly, for the term depending on Tn T we have
TT - Tw)u(P))ls = sup
[[(T], T TT)o(P)((2,a))
(x,a)en
sup
(V(x,,)) (T|x Tn))(P)))
(x,a)e
[T((P(P) @ (b, a) (T)( Tw)]
sup (r,a)e
< || Tallls sup (x;))
(x,a)E
|r|||
||
Tallss
1-7
Combining the previous two bounds, we get to
2111l10 || - T Ta
(P), - (P))ll $ 1-7 Collra r11a + 1 1
H5
as desired.
According to the result above, we can control the approximation enror for the action value function in
terms of the approximation errors ||r, - rilg and T Tn HS This can be done by leveraging state-of-the-arr 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
1. g e satisfies the strong source condition Assumption 2 with parameter on the probability distribution p.
2. E [0]]+20 as in the notation of [40].
Proof. Using the same notations in [40], we have
g 1428
Σ
aip? e, and [ax)inelle s0.
1 +
g
24

---

For 2 we have that that for E |
C Cp g - Σ
asfl Ei
(C.1)
SEN
whose G-norm is ||cp ((as))EN))) 1la 80.
For 2 = 1, let g = EiEn biul e, with (1(b))eEN(l < oo (since). Now, واو8 < oo is
equivalent to [[batt DEENl < co. By letting b; = 1 a, we have that ((a))) lle < oo and that
g Σ 0:14 $
iEN
that is E [9]]-29.
With the connection between [40] and Assumption 2 in place characterize the bound on the approximation error for the world model-based estimation the action-value function.
Proposition C.2. Let Tn and Tn the empirical estimators of the transfer operator T and reward function r as defined in Proposition 3, respectively. When P a FGommacattei policy as in Definition and the strong source condition Assumption 2 is attained with parameter B3, it holds
|| (P) ,(P))l S 1(62naa))
(C.2)
with rates a E
and probability not less than 6
2423 1+28
Proof. We use Lemma C.I to apply Theorem 3.1 (ii) from [40] to show that under Assumption with parameter it holds, with probability not less than 4e
n' r و 8,c,n
(C.3)
The rate ar
2+20 1+28 is determined by the properties of the inclusion g - BB( and the constant > 0 is independent of n and 6. Similarly, point (2.) of [21, Theorem 2] shows that under Assumption
3
|| T FullsS(F.0)) S 82 CTn
(C.4)
again with probability not less than 1- 4e rates at E 2+23 1+25 and with cT independent of and 8. Combining every bound and denoting a (ar,, aT), conclude
8² I-7
>CT ||T||/|l Cect
(P)))
(C.5)
1-7
as required.
C.3 Convergence Rates for POWR
With bound the estimation error of the action-value function by Algorithm 1, are finally ready to state the complexity bounds for POWR.
Theorem 9. Let (Te) be sequence of policies generated by Algorithm I in the same setting of Corollary 6. If the action-value functions qn, are estimated from a dataset (Ii, ag x with (Ii,a,) p P(2)) such that Assumption 2 holds with parameter B, the iterates of Algorithm I converge to the optimal policy as
J(T.) J((T) (0 +
with probability not less than 1 - 5 Here, a E 2+28 1+23 and n. X + A(A) is measurable maximizer of (8).
Proof. Since the setting of Corollary 6 implies that P, are (G, F) compptibbe for all I, and Assump- tion 2 holding, then g(P)) and g, (Px) belong to 9 for (Pt)eeN- This assures that we can use the statistical learning bounds Proposition C.2 into Theorem 7, yielding the final bound.
D
25

---

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 x 10 timesteps during the training. On the other hand, the new lower threshold 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 lower variance compared other techniques. This implies that our approach requires a stable amount of timesteps to learn how to solve a specific environment.
(a) FrozenLake-vi
(b) Taxi-v3
(c) MountainCar-vo
Figure Mean timestep at which various algorithms attain a specified reward threshold during their training. The reward targets are set at 0.8 for FrozenLake-v1, 6 for Taxi-v3. and - -110 for MountainCar-vo. The absence box indicates that the corresponding algorithm was unable to meet the reward threshold within the training process.
D.2 Other methods
We compare the performance of our algorithm with several baselines. 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 [48].
26

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
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 perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long 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 not discussed in the paper.
The authors are encouraged to create 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 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might 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 address problems of privacy and faimess.
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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant 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 complete (and correct) proof?
Answer: [Yes]
27

---

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 cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide 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 supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibiliyy
Does the paper fully disclose all the information needed to reproduce the main ex- perimental 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)?
Question:
Answer: [Yes]
Justification: the paper, we provided the algorithm description and analysis, including the hyperparameters of the baselines and proposed algorithm.
Guidelines:
The answer NA means that the paper does include experiments. If the paper includes experiments, 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 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 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 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 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 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 submis- sions 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 reproduce that algorithm. (b) If the contribution is primarily
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 way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe 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., 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 instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
28

---

reproducible, they 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/
wwwpsaabmssssnoooooooooooooooooo oniioonnpiillll oliccyyyy for more details.
While 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 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.. wtt.cccppppllcccc ptpppppgguides////pppppppppppaep httg/gudde///// Cedeee/mmmmubmiic ieriiisssonpiiiiiie
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 subset of experiments 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 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, hyper- parameters, 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 tested our POWR and the baselines.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to 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, supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
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 accompanied by emror bars, confi- dence 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 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 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 emrr 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
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
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 well 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 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: / htt////eurrp...//
inns?
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 deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is a special consid- eration due to laws or regulations 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: believe that this work doesn't have any societal impact.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact 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.
30

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is 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 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
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, scraped datasets)?
Question:
Answer: [NA]
Justification: This work doesn't 't have such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 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 dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.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, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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.
31

---

If this information is not available online, the authors are encouraged to reach out to the asset's '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/moodl 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 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 crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, 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 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 NeurlPS 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.
32


Reviews:
[{"review_id": "8VL5ExN3lc", "reviewer": "Reviewer_vreF", "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.", "questions": "See above. \n\nThe 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.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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].\n\n[1] Schmidhuber, J. Making the world differentiable: On using supervised learning recurrent neural\nnetworks for dynamic reinforcement learning and planning in non-stationary environments. (1990)\n\n[2] Schmidhuber, J. An on-line algorithm for dynamic reinforcement learning and planning in reactive environments. (1990)\n\n[3] Ha, D., and Schmidhuber, J. World Models (2018)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive review.\n\nIn particular, we thank the reviewer for the additional references. We will add them to the paper with the following expanded discussion in the introductory section of our paper, on line 42 (we denoted with R1,R2,R3 the references suggested by the reviewer):\n\n_The notion of world models for RL has been introduced by Ha and Schmidhuber in [R3] (building on ideas from [R1,R2]), where RNNs are used to learn the transition probability of the MDP. Traditional world model methods such as those proposed in [R3,16] 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._", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fS6Ts0HYvZ", "reviewer": "Reviewer_68rN", "summary": "(This review has been updated after revert of desk reject)", "questions": "N/A", "limitations": "yes", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "(updated after revert of desk reject)\nIt's a solid theoretical paper with nice writing and experimental results on continuous control benchmarks.", "weaknesses": "(updated after revert of desk reject)\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We kindly point out to the reviewer that the decision on desk rejection was reverted by the Program Chairs\n\n(The initial desk rejection had been due to the automated checker failure to detect the NeurIPS checklist in the supplementary material. This happened to several papers this year). \n\nShould the reviewer have any questions, we’d be glad to answer them.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LHedtc8jlO", "reviewer": "Reviewer_41Mf", "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.", "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?\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.", "limitations": "The 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.", "rating": 7, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "To 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", "weaknesses": "Line 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback and kind words. \n\n- __Line 88 - references__: the requested references are provided in the same sentence referred to by the reviewer, line 90. They are references [19,20] in the paper, namely:\n\n    [19] Beck and Teboulle. Mirror descent and nonlinear projected subgradient methods for convex optimization. Operations Research Letters, 2003\n    \n    [20] Sébastien Bubeck, Convex optimization: Algorithms and complexity. Foundations and Trends in Machine Learning, 2015\n\n- __X-axis Range for Plots__: In Fig. 1, we reported the average reward starting from 10^3 because, for most environments, the region $[0,  10^3]$ did not provide much information. However, we agree with the reviewer that the full plot is useful to have a complete picture. Please refer to the pdf attached to our global post here on OpenReview.\n\n- __Variance in Figure 1 (a) and (b)__: Both FrozenLake and Taxi are simple deterministic tabular environments. Once POWR has collected enough evidence about the reward function and the transition operator, there are no fluctuations anymore: the mean reward is constant, and the variance converges to zero. This is precisely the advantage of learning an operator-based world model rather than adopting an implicit one from which we can only sample possible trajectories. We also point out that the behavior is different for MountainCar, since the environment has infinitely many states, and POWR still incurs errors in the approximation of the transition operator.\n\n- __Fig. 1 (c) stops at ~10k steps__: Thanks for pointing this out. Indeed, in the original figure the convergence of POWR for MountainCar was not fully evident from the plot. We have re-run the experiment for a larger number of epochs. See the pdf attached to our global post here on OpenReview. Given the short time for the rebuttal period, we run POWR collecting over 40K samples. The new results show that POWR retains its performance after reaching the “success” threshold, suggesting that it achieved empirical convergence.  \n\n- __POWR in Continuous Action Spaces__: that is a very interesting question. We refer to our global reply here on OpenReview since this was a topic of interest also for other reviewers. \n\n- __Other feedback__: We thank the reviewer for pointing out typos and other issues/suggestions. We will amend the paper accordingly. (Concerning the broken hyperlinks, it was due the way we separated the paper from the appendices/supplementary material. Hyperlinks, indeed, work as intended in the full paper pdf provided in the supplementary material uploaded in the original submission).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sM8vGVBoB9", "reviewer": "Reviewer_Tuc7", "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.", "questions": "How well can this algorithm scale to higher-dimensional settings?", "limitations": "Authors have adequately addressed limitation and assumptions. \n\nI see no potential negative societal impacts as a result of this work.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1) 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.", "weaknesses": "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. \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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive feedback. \n\n- __Experiments on more complicated MDPs__: we agree with the reviewer and we are very excited to test our methods on more complicated environments. As the reviewer pointed out, our main focus in this work was to prove the novel theoretical contributions of Prop. 3, Cor. 7, and Thm. 9. Our experiments in section 4 are meant as empirical support to our theoretical findings. As we discussed in the conclusions of this work, competitive evaluation of POWR on richer MDPs requires additional research on:\nScaling the method to high-dimensional settings (see also our additional comments below)\nCombining POWR with deep representation learning schemes to learn the operator-world model beyond settings suited for Sobolev representations (e.g. images). \n\n- __POWR in Continuous Action Spaces__: that is a very interesting question. We refer to our global reply here on OpenReview since this was a topic of interest also for other reviewers. \n\n- __Scaling to high dimensional settings__: breaking down this question into the computational and statistical perspectives:\n    - _Statistics_: from the statistical perspective (namely Thm. 9), the method is not (directly) affected by the dimensionality of the problem. The learning rates in Thm. 9 mainly depend on the norms of the operator world model and reward function as elements of the feature space where learning is carried out (cf. the final bound reported in Thm. A.11 in the appendix). These norms can be interpreted as capturing “how well suited” the chosen feature space is to represent the operator world model. \n\n        To (partially) answer the reviewer’s question, in the case of the Sobolev feature spaces used in this work, there exists a relation between the dimensionality of the ambient space (namely the state space) and the norm of a function in terms of its smoothness (see [28,29]). In other words - and very loosely speaking - the larger the space dimension, the smoother the target function needs to be to maintain a small norm. This means that if the reward function or the transition operator are not very smooth, we might incur large constants in Thm. 9, hence slower rates. Choosing the hypothesis space is key to fast convergence, and we have discussed how to extend POWR with other hypothesis spaces than Sobolev spaces in our conclusion and future work. \n\n        This behavior is well-understood, albeit rather technical, in the kernel learning literature for traditional Sobolev spaces (for instance, see Sec. 3 and Ex. 3 in [B] or combining Chapter 5 in [28] with the results on interpolation spaces from [C] or [D]. References provided below). It is less clear how this interpretation will extend to less traditional feature spaces. We thank the reviewer for the question. We will add the discussion above as a remark following Thm. 9. \n\n    - _Computations_: by leveraging the kernel trick argument from Prop. 3, the dimensionality of the ambient space or the feature space does not come into play. In contrast, computations are affected by the number of samples observed (as is the case with most kernel methods). However, as observed in the recent kernel literature, this issue can be easily mitigated by adopting random projection approaches such as Nystrom sampling [25,26], significantly reducing computations without compromising model performance. \n\n- __Typos__: Thanks for pointing out typos. \n\n__Additional references__\n\n[B] Cucker and Smale. \"On the mathematical foundations of learning.\" Bulletin of the American mathematical society, 2002.\n\n[C] Steve and Zhou. \"Learning theory estimates via integral operators and their approximations.\" Constructive approximation, 2007.\n\n[D] Caponnetto and De Vito. \"Optimal rates for the regularized least-squares algorithm.\" Foundations of Computational Mathematics, 2007.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Thanking all reviewers for their feedback on our work, in this global reply we share our insights on the extension of our results to continuous action spaces, a question shared among reviewers. In addition, in the attached pdf we report Figure 1 updated according to the suggestions of reviewer __41Mf__.

Extending POWR to action spaces of infinite cardinality could in principle be naively tackled by approximating the normalization integral in the denominator of eq. (13) via sampling methods such as Monte Carlo (MC) sampling. This approach, however, poses two questions: 

- _Computational_: the integral approximation is conditional to the input state, implying that whenever we need to evaluate the policy on a new state, a new approximation is required. Unrolling this computation backward through all PMD iterations is potentially computationally expensive.

- _Approximation_: it is not clear whether this is a good approximation of the ideal solution to the PMD step in equation (2) and how it will impact the convergence rates proved in the theoretical section.
 
This makes the naive approach to extending POWR via MC sampling potentially limited.

An alternative promising approach is to approximate the solution to the PMD step in equation (2) via an iterative optimization strategy. This would still introduce approximation errors, but they could be controlled via the optimization rates for mirror descent. In particular, the recent  work in [A]

[A] _Aubin-Frankowski, Korba, and Léger. "Mirror descent with relative smoothness in measure spaces, with application to sinkhorn and em." Advances in Neural Information Processing Systems 35 (2022): 17263-17275._

shows that it is possible to cast and perform mirror descent also on general probability spaces (a question that was still open) maintaining analogous convergence rates to the finite setting. To our preliminary investigation, adapting the analysis in [A] to Policy Mirror Descent is feasible. However, the open questions in this sense are:
How to concretely perform the optimization of eq. (2) in practice, and
Whether this additional iterative procedure would compromise the computational efficiency of POWR’s pipeline. 

We hope that these comments provide some additional context to the question posed by the reviewers, and why continuous action spaces ended up not being addressed in the paper at this time.

Meta-review: 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.
Decision: Accept (poster)