Deep SPI: Safe Policy Improvement via World Models
Authors: Florent Delgrange
Raphaël Avalos
Willem Röpke
Venue: ICLR 2026 Poster (2026)
Keywords: reinforcement learning
guarantees
representation learning
model-based
Primary area: reinforcement learning
TL;DR: We consider safe policy improvement for on-policy reinforcement learning  in general-state spaces; we provide safe policy improvement guarantees tailored to (learned) world models and representation learning.

Abstract: Safe policy improvement (SPI) offers theoretical control over policy updates, yet existing guarantees largely concern offline, tabular reinforcement learning (RL). We study SPI in general online settings, when combined with world model and representation learning. We develop a theoretical framework showing that restricting policy updates to a well-defined neighborhood of the current policy ensures monotonic improvement and convergence. This analysis links transition and reward prediction losses to representation quality, yielding online, ''deep'' analogues of classical SPI theorems from the offline RL literature. Building on these results, we introduce DeepSPI, a principled on-policy algorithm that couples local transition and reward losses with regularised policy updates. On the ALE-57 benchmark, DeepSPI matches or exceeds strong baselines, including PPO and DeepMDPs, while retaining theoretical guarantees.

Published as a conference paper at ICLR 2026
DEEP SPI: SAFE POLICY IMPROVEMENT VIA WORLD MODELS
Florent Delgrange AI Lab, Vrije Universiteit Brussel Flanders Make
Raphael Avalos
Al Lab, Vrije Universiteit Brussel Cohere
Willem Röpke
AI Lab, Vrije Universiteit Brussel Cohere
ABSTRACT
Safe policy improvement (SPI) offers theoretical control over policy updates, yet existing guarantees largely concern offline, tabular reinforcement learning (RL). We study SPI in general online settings, when combined with world model and representation learning. We develop a theoretical framework showing that restricting policy updates to well-defined neighborhood of the current policy ensures monotonic improvement and convergence. This analysis links transition and reward prediction losses to representation quality. yielding online, "deep" analogues classical SPI theorems from the offline RL literature. Building on these results, introduce DeepSPI, principled on-policy algorithm that couples local transition and reward losses with regularised policy updates. On the ALE-57 benchmark, DeepSP matches or exceeds strong behaviorals, including PPO and DeepMDPs, while retaining theoretical guarantees.
INTRODUCTION
Reinforcement learning (RL) trains agents to act in complex environments through trial and error (Sutton and Barto, 2018). To scale to high-dimensional domains, modern approaches rely function approximation, making representation learning essential for learning latent spaces where behaviorally similar states are grouped and policies and value functions are easier to estimate (Echchahed and Castro, 2025). complementary approach is model learning, where a predictive model of the environment is trained (Ha and Schmidhuber, 2018). Such models support planning, latent simulation, and improving value estimates (Hafner et al., 2021; Schritwieesr et al., 2020; Xiao et al., 2019).
In the online setting, where the agent updates its policy during interaction, avoiding catastrophic errors is critical. Two key challenges arise: out-of tftjof-trajectory (OOT) world models and confounding policy updates. OOT issues arise when the world model fails to capture rarely visited regions of the state space, leading to unreliable predictions and unsafe updates when the latent policy explores these regions (Suau et al., 2024). Confounding updates occur when both the policy and its underlying representation are updated simultancously: poor representations can lock the agent into suboptimal behavior, while the policy itself prevents corrective updates to the representation. Safe Policy Improvement (SPI) mitigates such risks by ensuring that new policies are not substantially worse than their predecessors (Thomas al., 2015). Classical SPI methods provide rigorous results in tabular MDPs but depend on exhaustive state-action coverage, making them unsuitable for continuous or high-dimensional spaces.
We address this gap by connecting representation and model learning with safe policy improvement in general state spaces. Our contributions are threefold. First, we introduce a neighborhood operator that constrains policy updates and guarantees convergence. Second, we combine this operator with principled model losses to bound the gap between a policy's performance in the world model and in the true environment, thereby enabling safe policy improvement in complex MDPs. This analysis also shows that our scheme enforces representation quality by ensuring that states with similar values remain close in the learned latent space. Third, connect theory PPO (Schulman et al., 2017)

---

Published as a conference paper at ICLR 2026
and propose DeepSPI, practical algorithm that achieves strong empirical performance the Arcade Learning Environment (ALE; Bellemare et al. 2013) while retaining theoretical guarantees.
1.1 RELATED WORK
Regularizing policy improvements. Regularized updates, as in TRPO, PPO, and related analyses, are now standard for stabilizing policy optmiiatiio (Schulman et al., 2015; 2017; Geist et al., 2019; Kuba al., 2022). Our work extends this perspective to the joint training of world model and representation, where we constrain policy updates in a principled neighborhood while controlling model quality through transition and reward losses.
SPI methods provide principled guarantees on policy updates from fixed datasets (offline RL) (Thomas et al., 2015; Ghavamzadeh et al., 2016a; Laroche et al., 2019; Simão et al., 2020; Castellini al., 2023). These methods assume tabular state spaces and offline data, where error bounds must hold globally across all state-action pairs, often via robust MDP formulations (lyengar, 2005; Nilim and Ghaoui. 2005). Our setting fundamentally different: we study online RL with high-dimensional inputs, where such global constraints are intractable. We take inspiration from the SPI literature but introduce local, on-policy losses that make safe improvement feasible in practice. In spirit, other model-based methods share the goal of providing SPI-like guarantees in more general settings, but are again purely offline, omit any form of representation learning, and rely on assumptions that differ substantially from ours (Yu et al., 2020; 2021; Kidambi et al., 2020).
Representation learning and model-based RL. Auxiliary transition and reward prediction losses are central to many model-based methods, from DeepMDP to Dreamer and related world-model approaches (Gelada et al., 2019; Hafner et al., 2021). In particular, the losses we consider for learning transitions and rewards generalize wide range of objectives used across the model-based RL literature (François-Lavet et al.. 2019; van der Pol et al., 2020; Kidambi et al., 2020; Delgrange et al., 2022; Dong et al., 2023; Alegre et al., 2023). Conceptually, our representation guarantees are closely connected to classical notions of state abstraction in MDPs (Li et al., 2006) and to bisimulation (Larsen and Skou, 1991; Desharnais et al., 1998). Building on bisimulation, prior work develops representations that group states into areas in which the agent is guaranteed behave similarly under the current policy (Castro, 2020: Zhang et al., 2021; Castro et al., 2021; Agarwal et al., 2021 a; Avalos al.. 2024). By contrast, we directly link representation quality and model accuracy our safe policy improvement analysis, yielding tractable guarantees in the online setting.
2 BACKGROUND
In the following, given measurable space X, write A (2') for the set of distributions over X. For any distribution E ((Y), we denote by supp(n)) its support.
Markov Decision Processes (MDPs) offer a formalism for sequential decision-making under uncer- tainty. Formally, an MDP is tuple of the form M = (S, A, P. Rt 87, 87,7) consisting of a set of states S, actions A, a transition function P: S x A + A(S), a bounded reward function R A + IR with || RMAX. initial state $1 S, and a discount factor E [O. 1). Unless otherwise stated, we generally assume that S and A are compact. An agent interacting in M produces trajectories, i.e., infinite sequences of states and actions (st, visited along the interaction so that 8p 81 and 8,41 P(- St, ar) for all 2 0.
At each time step I, the agent selects an action according to a (stationary) policy π : S - A(A) mapping states distributions over actions. Running an MDP under æ induces unique probability measure F. over trajectories (Revuz, 1984), with associated expectation operator E,; we write
E. 8p =) when the initial state is fixed to s E S. A policy has full support if supp(r(: x)) = A for all s E S, and we denote the set of all policies by TT. A stationary measure of is distribution over states visited under म, and is defined as a solution of (-) = Egns, Eaw [P s, a)). Such measure is often assumed to exist in continual RL (Sutton and Barto. 2018), is unique in episodic RL (Huang, 2020), and defines the occupancy measure in discounted RL (Metelli et al., 2023).'
Details on the formalization of episodic processes and value functions be found in Appendix A.
2

---

Published as a conference paper at ICLR 2026
Value functions. The performance of the agent executing a policy π II in each single state s E S can be evaluated through the value function (*) = Ex 15ح=0 r R(s,, at) 8u = s). The goal of agent is to maximize the return from the initial state, given by MM)
V (s)). To evaluate the quality of any action a E A, we consider the action value function Q" (s. a) = R(s,a) a) + y P(-laa)) (s'),, being the unique solution of Bellman's equation with V* (1a Q" a). Alternatively, any given action can be evaluated through the advantage function A* a) = (x) - (s), giving the advantage of selecting an action over the current policy,
Representation learning in RL. In realistic environments, the state-action space is too large for tabular policies or value functions. Instead, deep RL employs an encoder Φ: S - S that maps states to tractable latent space S, from which value functions can be approximated. Learning such encoders is referred to as representation learning (Echchahed and Castro, 2025). To improve representations, agents are often trained with additional objectives, commonly auxiliary tasks requiring predictive signals. Policy-based methods then optimize latenl policy n: S - A(A) jointly with d, executed in the environment as 7(- (())). By convention, we write (- 5) for æ ((s) when is clear, and denote the set of all latent policies by IT. For any π E II, the composed policy ι o e belongs to IT.
Model-based RL augments policy learning with a world model M = A5 $1,7), which can improve (i) sample efficiency by generating trajectories (e.g.. Hafner et al. 2021). (ii) value estimation through planning (e.g.. Buckman et al. 2018), 8), and (iii) representation learning by grouping states with similar behavior (e.g., Gelada et al. 2019; Zhang et al. 2021). When 5 = S, the model must replicate environment dynamics, which is often intractable. Instead, we focus on & defined by the learned representation Φ so that M becomes an abstraction of M. Learning transition and reward functions then additionally serves as an auxiliary signal for the representation, encouraging states with similar behavior to map close in S. Since S the latent space, II corresponds to the policies of M. We further assume S is equipped with a metric d: S x 5 + [0, oc) to measure distances.
No WAY HOME: WHEN WORLD MODELS AND POLICIES GO OUT OF TRAJECTORIES
World models are usually learned toward minimizing a reward loss L and/or transition loss I. from experiences r) collected along the agent's trajectories. Those experiences are either gathered in the form of a batch or a replay buffer B. In general, the loss functions take the following form: L = Ey fn (o. R: 1) and Lp EnpB fo (Ø, P: m). where fR (resp. fP) assign a "cost" relative to the error between R and R (resp. P and P) according to the experiences r) and their representation. Henceforth, refer to the policy Th used to insert experiences in B as the behavioral policy.
3.1 OUT-OF-TRAJECTOR WORLD MODEL
One may consider leveraging the model M improve the policy Tb. This can be achieved by directly planning a new policy मे in M or drawing imagined trajectories in the world model to evaluate new actions and improve on sample complexity during RI. However, since the world model is learned from experiences stored in B, we can only be certain of its average accuracy according to this data. This is problematic because some regions of the state space of M may have been rarely, or not at all, visited under Tp. In that case, the predictions made in M might cause the agent to "hallucinate" inaccurate trajectories in the latent space and spoil the policy improvement. This problem, known as the out-of-trajectory (OOT) issue (Suau et al., 2024), arises when a policy in M deviates substantially from mb. which can render the model unreliable.
To illustrate this problem, consider the world model of Figure 1. Assume the model is trained by collecting trajectories produced by Tb in M where ma 8) < e for all s E S1, with e > 0. For sufficiently small the region S3 in the original environment would remain largely unexplored while having almost no impact on the losses LR Lp. Therefore, the representation of states in S3 (5x and 83) may turn completely inaccurate. Here, the model incorrectly assigns reward of 20 to 83. whereas the true reward is strictly negative. Consequently, the optimal policy in M deterministically selects 02 in s When executed in the original environment, this policy drives the agent to Sg thereby degrading the behavioral policy Tb-
3

---

Published as a conference paper at ICLR 2026
S₁
82
45
St
Si
81
84
Sg
8/3
+20
(a) A large MDP M whose state space is divided in four regions& S Ui=11 S..
(b) A simple world model M whose state space = (31, 82, 8a, 8g, 8a).
Figure 1: M. continuously playing di states from S eventually leads agent to the region S2. and playing 0g in S2 eventually leads the agent to Sy where reward of 1 is incurred at each time step, whatever the action played. Playing as in S₁ leads the agent to the region Sg. where all actions incur negative rewards. Here, (ss) 8; for any
S; and i 21 4). For E Ss. have either (bs) 33 o(s) $/3.
3.2 CONFOUNDING POLICY UPDATE
Updating both the representation and the policy solely from experience collected under behavioral policy can degrade performance rather than improve it. In the same spirit as policy confounding (Suau et al., 2024), call this phenomenon confounding policy update. The MDP in Figure illustrates the issue.
The agent maps the states 82 and $3 to the same latent state 5, i.c. ((s) = 5 iff s E {82,88) States 51 and each have their own latent state. We consider the behavioral policy Tb == n6 Φ where Tg is a stochastic policy with a small exploration rate Ç:
d22
$2
SI
SA
Q1
(a1
11-5. S.
(a 3) S.
83
for 0 <<<<<<< A good representation would ideally group states from which the agent behaves similarly. Because trajectories that reach 83 and pick d2 are unlikely, the two states appear identical - W (sa)[=00 0. Therefore, this justifies using
Figure 2: MDP where the probabil- ity transitioning from 81 to 82 is 1 for <<< 1/4.
under Tb: as representation for The because the values of 82 and $3 are nearly identical: the agent exhibits close behaviors under from those states.
Suppose exploration under n, eventually discovers that playing 02 in ø sometimes yields the +2 reward. Based on exploration data, an RL agent might therefore be tempted to change the latent policy to (a2 8) = 1 without modifying the representation é. With the representation still grouping 82 and $3, the new policy would now deterministically pick a2 in both concrete states. Whenever the agent actually reaches $3, it would receive the large negative reward -2/c, which turns the overall return (from 8 negative, thus worse than under m, even though d2 is indeed optimal in 82.
solution to this problem would have been to split the representation of 82 and 83 in distinet latent states. In general, representation and policy learning must be coupled since any change in the policy that alters the distribution over states can invalidate previously adequate representation. However, in this example, the agent has no incentive to do so based on the experiences collected under Ty. As we will show below, updating both the policy and the representation jointly should be handled carefully to ensure policy improvement.
Our goal is to establish sufficient conditions to guarantee safe policy improvement during the RL process, either based on world models, state representations, or both, thus alleviating OOT world model and confounding policy update issues. Notice that, in the examples, both problems occur when performing aggressive updates from Th new policy π (the mode of the distributions drastically shifts). Intuitively, smooth updates indeed ensure to alleviate those issues: constraining the policy search to policies "close" to (i) prevents hallucinations in parts of the world model that have been underexplored; (ii) reduces the risk of significantly degrading the return when updating the policy. While the benefits of regularizing policy improvements have already been both theoretically and practically justified (e.g., Geist et al. 2019; Kuba et al. 2022), their implications when mixing model-based and representation learning in RL have been underexplored.
4

---

Published as a conference paper at ICLR 2026
Roadmap. To rigorously address the OOT and confounding- issues, the next sections develop the theoretical foundations of our approach, showing how controlled policy updates, local model losses, and representation stability interact. We briefly summarize how the main results connect.
Our analysis combines neighborthod-rrsstiieedd policy updates, model-quality bounds, and representation guarantees. Sect. 4 introduces the neighborhood operator defining trust region around the behavioral policy; restricting updates to this region ensures monotonic improvement and convergence (Thm. 1). Sect. 5 then links the reward and transition losses to value discrepancies: Thm. 2 shows that, when these losses are small, and updates remain in the neighborhood. the world model stays accurate under the learned representation. Combining these ingredients yields our first SPI result (Thm. 3), guaranteeing that direct policy updates in the world model translate to improvement under controlled error. Finally, Thm. shows that the same loss-based control stabilizes the encoder, ensuring that value-distinet states remain separated in the latent space.
4 YOUR FRIENDLY NEIGHBORHOOD POLICY
Motivated by the intuition that constraining policy updates can mitigate OOT and confounding policy issues, we consider measuring the update as the importance ratio (IR) of the policies. This measure provides guarantees for constraining policy and representation updates, and with an appropriate optimisation scheme, ensures both policy improvement and convergence. Section 5, we will further show that properly constraining the IR allows for safe policy improvements in world models while providing representation guarantees.
Let n. E II. the extremal importance ratios are defined Di (π, ext { (als)//(aaa)) s E S, a E supp(T(. s))} where ext E {inf,sup}. We define a neighbor- hood operator based the IR, .CC II 2 for some constant 1 C < 2, establishing
trust region for policies updates that constraints the IR between - and C::
2 C s DiR (TT, r²) I DiR (e,((zz (x, S SC and supp(T(-) {)) supp(='((1) (-1s)) E
}
No (T)
ा E
E
A critical question is whether agent that restricts its policy updates to defined neighborhood is truly following a sound policy improvement scheme. The following theorem shows that does and further guarantees convergence.
Theorem 1. (Policy improvement and convergence guarantees) Assume S and A are finite spaces. Let No II be a policy with full support and (In) be a sequence of policy updates defined as
Tn+1 == arg E
E
A (8,
(TE)) smpten awn' (-18)
where Pr, is a sampling distribution with supp((>=.) = for each 2 0. Then, the value function V is monotonically improving, converges to V", and so is the return pl TT, M).
The proof consists in showing the resulting policy update scheme is an instance of mirror learning (Kuba et al., 2022), which yields the guarantees. Notice that since To has full support, all the subsequent policies n. have full support as well. To maintain the guarantees, considering a stationary measure & as the sampling distribution only possible when supp(E..)) S. Note that this is always the case in episodic tasks (as the policy itself has full support). This is more generally true in ergodic MDPs (Puterman, 1994).
WITH GREAT WORLD MODELS COMES GREAT REPRESENTATION
This section explains how the neighborhood operator of Eq. enables safe policy improvement during world-model planning and representation updates in complex environments. Standard SPI methods ignore representation learning and require exhaustive state-action coverage B to obtain guarantees, making them unsuitable for general state-action spaces. Even in finite domains, bounding the count
There are clear similarities between the IR, our neighborhood operator, and the PPO loss function (Schulman al., 2017). We discuss this connection in Section 6.
5

---

Published as a conference paper at ICLR 2026
of each cotte-action pair does not scale. Laroche et al. (2019) proposed baseline bootstrapping for under-sampled pairs, but their approach remains impractical in large-scatd scale settings despite conceptual similarities to our operator. Further discussion of SPI limitations is provided in Appendix D.
Learning an accurate world model. SPI typically relies on optimizing a policy with respect to latent model learned from the data stored in B. In contrast to previous methods, our approach scales to high-dimensional feature spaces by (i) learning representation & and (ii) considering local error measures as opposed to global measures across the whole state-action space. We formalize them as tractable loss functions. Their local nature makes them compliant with stochastic gradient descent methods. Formally, given distribution B E A(S x A), we define the reward loss L and the transition loss Lp as
IF 8,0=B
R(s,a) R(s,a)),
=
E (o (8,18880) 8, P(: ((s),
(4)
8,0~
where by P is the pushforward measure of P by Φ. and W the Wasserstein distance (Vaserstein, 1969).
W between Jr, v E (S)) is defined as W (n, inf xn[xAtee)))) ((a,z")~ d(3, ('), where A(u, v) is the set of all couplings of μ and v. While the Wasserstein operator may seem scary at first glance, generalizes over transition losses that can be found in the literature (cf. Sect. 1.1). In particular, when the latent space is discrete, this distance boils down to the total variation distance. Another notable case is when the transition dynamics are deterministic, in which case the transition loss reduces to LE Es,,,''BB d(o(5', P(}(s), a)). Finally, in general, a tractable upper bound can be obtained as p.p S Es Ey-P(- (()).a) d()(s"), 8') (proof Appendix C).
Lipschitz constants. To provide the guarantees, for any particular policy म E II, we assume the world model equipped with Lipschitz constants TT KT. defined follows: for all $1, 32 E,,
R²'
E (1
R1)
E
R(52, a2)
S K d(51.52).
W
E P(. 81, 01);
E (2)x((112)
P(. 82, 02)
KT (511 82).
(131))
Intuitively, the Lipschitzness of the latent reward and transition functions guarantees that the latent space is well-structured, so that nearby latent states exhibit similar latent dynamics. Gelada et al. (2019) control those bounds by adding gradient penalty term the loss and enforce Lipschitzness (Gulrajani et al., 2017). One can also obtain constrained Lipchitz constants as a side effect by enforcing the metric d to match the bisimulation distance in the latent space (Zhang et al., 2021). Interestingly, when the latent space is discrete, Lipschitz constants can be trivially inferred since K = 2feAx and K = 1 (Delgrange et al., 2022). Note also that as the spaces are assumed compact, restricting to continuous functions ensures Lipschitz continuity.
For the sake of presentation, we restrict our attention the following assumption for Thms. 2 and 3: Assumption 1. We assume that the agent operates in the episodic RL setting, i.e., we consider the standard RI. framework where the environment is eventually resel with probability one.
Our results extend to general settings where stationary distribution is accessible (c.f. Remark 3).
World model quality. Before introducing our safe policy improvement theorem, we first show that the local losses effectively measure the world model's quality with respect to the original environment. Namely, their difference in return obtained under any latent policy in well-defined neighborhood is bounded by the local losses derived from the reference, behavioral policy's state-action distribution. This is formalized in the following theorem.
Theorem 2. Suppose > 1/2 and K < 1/7. Let C E (1,1/), T E II be the base policy, (āt o d) (T))) where Il is latent policy and o: S S a state representation. Then,
1 + Kv L 1/D (m).5)
M) - (TT,M))= & AEL((T))
where AEL(Tb) denotes the average episode length when M runs under Th, Kv K (x and T. Ems 1 are the local losses of Eq. 4 over the stationary distribution Ex,, induced by Th. I.

---

Published as a conference paper at ICLR 2026
In simpler terms, if the deviation (supremum IR, or SIR for short) between the behavioral policy and any new policy т stays strictly lower than 1/7, the gap in return between the environment and the world model for this new policy can be bounded using data collected via The Minimizing local losses from ny's 's data ensures that refining the representation & for и improves model quality: when these losses vanish, .M and M are almost surely equivalent under म. The bound depends the Average Episode Length (AEL), but even a loose upper bound is sufficient to preserve guarantees. It is also strongly influenced by the discount factor 7, which defines implicit horizon. Smaller values permit larger deviations from Th and relax the accuracy required of the world model.
Safe policy improvement. We consider the setting where the world model is used to improve the behavioral policy Ta = T6 o Ø, with R, E and the representation ø is fixed during each update. Restricting updates
well-defined neighborhood guarantees that p( @, MM p(mp, MM) > M p( (T, M) S. where ( is defined as the cumulative modeling error from the local losses. Theorem 3. (Deep, Safe Policy Improvement) Under the same preamble as in Thm. 2, assume that ó fixed during the policy update and the behavioral a latent policy with Tp Tb & and Th E III Then, the improvement of the return of M under Hh can be guaranteed on as
p( o o M) - p(Tb, M) 2 p( M)
M) -s:
where Ç = AEL(TT) (LR*/1
Ky KyLpp"
(m),))
Theorem 3 addresses the OOT issue (Section 1): if the SIR of the behavioral remains strictly below 1/4, then minimizing the local losses reduces the error Ç. ensuring safe policy improvement when the world model is used to enhance the policy. While our focus is not on offline SPI, Appendix E (Thm. 5) additionally provides a PAC variant of the result, following the standard use of confidence bounds in the SPI literature.
Representation learning. Finally, we analyze how learning a world model using our loss functions as an auxiliary task facilitates the learning of a useful representation. A good representation should ensure that environment states that are close in the representation also have close values, directly supporting policy learning. Specifically, we seek "almost" Lipschitz continuity (Vanderbei, 1991) of the form :К: V81, $2 E S. (V"
V* (#2) K ((()uu())) (sw()))) + Lm, (oou) where Lx. an auxiliary loss depending on the data collected by Th- Notably, a critical question is whether updating the policy and its representation, respectively to and Ø, maintains Lipschitz continuity. Crucially, as the behavioral m, is updated to o é with respect to the experience collected under Tb, the bound must hold for L.,- The following theorem is a probabilistic version of this statement, formalized as a concentration inequality:
Theorem 4. (Deep SPI for representation learning) Under the same preamble in Thm. 2, let 1. 1. +7KV Then, with probability at least 5 under En, we have for all 5 (Voog"
8 > 0 and 8 == 4.
82 E S that
(*'ss) V Kv d(((s))() ((22)) +
Theorem addresses confounding policy updates (Section 3.2): minimizing the losses increases the probability that learned representations remain almost Lipschitz under controlled policy changes (with an SIR below 1/7). This prevents distinct states from collapsing into identical latent representations that degrade performance. We note that Gelada al. (2019) proved similar bound when Tb (the policy update was disregarded), which in contrast to ours, surely holds with
1 RKv
1
(81))
(=(2))
However, in general spaces, for any specific E S, (s) (x(xss might simply equal zero, making the bound undefined. In particular, in the continuous setting, widely assumed to be endowed with Borel sigma-algebra, where the probability of every single point is indeed zero.
6 ACROSS THE SPI-VERSE: PPO COMES INTO PLAY
These theorems inspire a practical RL algorithm that combines policy improvement and guarantees with solid empirical performance. The critical part of approach is make updates
7

---

Published as a conference paper at ICLR 2026
restricted to the policy neighborhood while minimizing the auxiliary losses Ln= Lp. In fact, neighborhood operator has close connections to PPO (Schulman et al., 2017), where the policy update given by*
#n+1 = arg sup E
E
(s, - Dr. (α' |s)
(5)
a n'Ell
with Dr, (T' () = Ba ReLu [='(ale//an(a)
clip( (hip(='(a=)))) (=l)))/)())))) 1 e)] ATM (s)
for some 5 > O. By fixing C 1, instead of strictly constraining the updates to the neighborhood. the regularization TT,, (T' corrects the utility (sls'()ss) A (s, a) (compare Eq. 3 and Eq. 5), so that there is no incentive for x' to deviate from nr with an IR outside the range [2 C, C?. Under the same assumption as in Theorem I, PPO is also an instance of mirror learning (Kuba et al., 2022), meaning it also benefits from the same convergence guarantees.
Strictly restricting the IR a neighborhood is much harder in practice, considering a PPO objective thus appealing alternative. However, it is sufficient to add the auxiliary losses L,p, La the objective of Eq. to maintain the guarantees. Indeed, updating the representation e by minimizing the additional losses may push the the policy ó outside the neighborhood. As solution propose to incorporate the local losses by replacing all occurrences of TT in Eq. 5 by the utility
U (s, s') == T (s, aR Eu(s, a) ap Ep(ss,,s'))
(6)
where a)
[R(s, a) (()(s),a)) ap's,a,s') = Egr-P(io(s) 5()(s)),5'. o'
P(: s, a). and aR, ap E (0. Intuitively, fR. fp are transition auxiliary losses that al- low retrieving L fTn and -f." in expectation w.r.t. the current policy nn. When optimized, since they are clipped in PPO-fashion, U allows restricting the policy updates to the neighborhood.
Algorithm 1: DeepSP1
From this loss, propose DeepSP a prin- cipled algorithm leveraging the policy im- provement and representation learning pabilities developed in our theory. As our losses rely distributions defined over the current policy, we focus on the on-policy set- ting. While model-based approaches are not standard in this setting, we stress that highly parallelized collection of data (e.g.. via vec- torized environments) enables a wide cover- age of the state space (cf. Mayor et al., 2025; Gallici et al., 2025), which is suitable to opti- mize the latent model. DeepSP1 updates the world model, the encoder, and the policy si- multancously while guarantecing the represen- tation is suited to perform safe policy updates.
Inputs: Horizon T, batch size B, vectorized environment env, parameters
Initialize vectors
()1 aa E AT , R
repeat
for L <- to T' do
Draw actions from the current policy:
(g. ~ (( ((+t.)) V1 < is
Perform single parallelized (B) step:
S141 wnv.enstteee cewv.stepsss..) a.)
Update # by descending
Vo DeepSPI loss(s. (8. a, U 0) b change in Eq. 5 by U from Eq. 6
8141
until convergence
return /
6.1 ILLUSTRATIVE EXAMPLE
To illustrate the representation learning capabilities of DeepSPI. we consider the toy grid-world shown in Fig. 3. This environment mirrors the confounding policy update discussed in Sect. 3.2, instantiated earlier in Fig. 2.
11
The agent starts in the cell labeled I. Upon leaving the orange cell immediately to its right, it is sent to the top e and to the bottom branch It must then traverse corridor of n blue cells (here n=5). Moving one cell to the right yields reward of +1, and the agent cannot move backwards.
branch with probability
with probability
Toy maze environment illustrating the confounding policy update problem.
Figure
At the final corridor cell, marked with a *, moving right yields a reward +1 regardless of whether the agent is in the top or bottom branch, and the episode terminates. The difference is when the agent
we give the formulation of Kuba et al. (2022), which is equal the of Schulman al. (2017).
8

---

Published as a conference paper at ICLR 2026
IQM
Optimality Gap
Median
Mean
PPO
DeepMDP (PPO)
DeepSPI
0.60 0.63 0.66
0.48 0.50
0.32
0.40 0.44 0.48 0.52 1.8
1.9
2.0
Human Normalized Score
Figure Aggregate results on stochastic versions of the standard environments from ALE, with 95% confidence intervals (CIs). Higher values for the mean, median, and interquartile mean (IQM) indicate better performance, while lower optimality gap is preferable (cf. Agarwal et al. 2021b). Cls are obtained through percentile bootstrapping with stratified resampling. Plots per environment available in Appendix H.3.
moves up from the cell: in the top branch, it receives a reward of +n/"", whereas in the bottom branch it receives -(2-r)n(ten") before termination. As in Sect. 3.2, this construction ensures that both * states are merged in the latent space, choosing "right" remains acceptable (their values coincide), but choosing "up" produces negative expected return from the initial state I (details in Appendix G). To improve upon the policy that always chooses "right," the agent must learn to assign distinct representations the
cells, select "up" in the top one, and "right" in the bottom one.
We compare the behaviour of PPO and DeepSPT in this environment. Since our goal to high- light the agent's representation-leamiii capabilities, each observation is provided as raw pixels. The agent must therefore learn
both policy and an encoder mapping pixels to a structured latent space. Further details on the environment and observation scheme are given in Appendix G.
Value from initial state
Resresentation distance
0.4
6-
0.3
DeepSPI
0.2
As shown in Fig. 4, the repre- sentation learned by PPO col- lapses the top and bottom cells into a single latent state. With such representation, the best policy PPO can learn to always choose "right," which leads return of 4.8. In contrast, DeepSPI benefits from the representation quality guarantees of Thm. 4. which ensure that states with different values remain separated in the latent space for all policies in suitable neighborhood. This is exactly what we observe: the learned representation distinguishes the two * cells. As result, the agent learns to choose "up" in the top * cell and "right" in the bottom one, achieving return of 8.
3-
DeapSPI
PPO
Step
Step
Figure 4: Value from cell in the (left) and distance between the representation
cell from the top and bottom branches (right).
IQM Human Normalized Score
EXPERIMENTS
0.6
In this section, evaluate DeepSPI the Atari Ar- cade Learning Environment (ALE; Bellemare et al. 2013),
0.4
where each frame four stacked frames. While ALE domains feature a wide range of dynamics, to further high- light the robustness of our approach probabilistic set- tings, we inject stochasticity by following Machado et al. (2018) with sticky actions (repeat the previous action with probability Pa) and random initialisation (games (re)start after riooPp frames), using Pa, = 0.3 and riooOP 60.
0.2
0.0
0.00
0.25
0.50
0.75
1.00 107
Global step
Figure 6: Sample efficiency w.r.t. IQM nor- malized scores on the stochastic ALE-57. Shaded regions give pointwise 95% CIs ob- tained via percentile stratified bootstrap.
As baselines, we use PPO (vectorized cleanRI; Huang et al., 2022) and DeepMDPs (Gelada et al., 2019). DeepMDPs auxiliary transition and reward losses (Sect. 5) that can be plugged into RL algorithms to improve representations, with guarantees. Unlike DeepSPI, these representation updates unconstrained and push policies out of the neigh- borhood, so none of the SPI guarantees presented this paper apply to DeepMDPs. For a fair comparison, plugged the DeepMDP losses to (vectorized) PPO, and use the architecture for

---

Published as a conference paper at ICLR 2026
Asterix
Asteroids
MarioBros
RoadRunner
Robotank
VideoPinball
1330
2913
30002
2020
200
1933-
290
25002
161
107
lef
161
1e9
163
DeepSPI
DeepMDD
PPO
DreamSPI
Figure 8: Sample environments from ALE where DreamSPI learns meaningful behaviors.
DeepSPI. We default cleanRL hyperparameters for all three algorithms, except data collection (128 environments, horizon steps).
As latent space, we the raw 3D representation obtained after the convolution layers (as recom- mended and used by Gelada et al., 2019). For the transition function, we use a mixture of multivariate normal distributions (the transition network outputs 5 means/diagonal matrices). To enforce Lips- chitz constraints on reward and transition functions, we model R. P with Lipschitz networks, using norm-constrained GroupSort architectures to enforce 1-Lipschitzness Anil et 2019.
As shown in Fig. 5 and 6, DeepSPI delivers strong performance, improving on both PPO and DeepMDP. Notably, these results are obtained while preserving SPI- valuable combination, the expense of performance and substantial data requirements. Beyond perfor- mance, we assess whether the DeepSPI world model exhibits accurate dynamics. Fig. 7 reports Lp. Ln during training. Note that DeepSPI consistently achieves lower tran- sition losses, indicating more accurate predic- tions. We discuss the statistical significance of that statement in Appendix H.3. In contrast to the off- polipoliy setting of Gelada et al. (2019), did not observe competing transition and reward losses in our parallel on-policy setting, as our losses are always computed under the current policy rather than eplay-bbuffer data.
style properties;
such theoretical control typically comes
transition loss
reward loss
10
0.010
0.008
0.006
0.004
0.002
0.000
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
1=7
107
DeepSPI
DeepMDP
Figure 7: Median transition and reward losses during training, aggregated across all the ALE. For the sake of visualization, we cut Lp lower values from the plot.
To probe the predictive quality of the latent model and illustrate Thm. 3, we introduced DreamSPI, naïve variant where DeepSPI learns the world model and representation, and PPO updates the policy from imagined trajectories (Appendix F). Unlike off aff-policy approaches that exploit replay buffers, on-policy setting updates the world model only from fresh interaction data, making model learning and planning harder to combine. Even so, DreamSPI learns several environments and exhibits coherent behaviours (cf. Fig. 8 & Appendix H.3). Although its aggregate median score remains below the baselines, this is expected given stricter data requirements than standard model-based approaches. Importantly, the ability to maintain model offers benefits that extend far beyond raw scores, enabling future applications in safety, verification, and reactive synthesis.
8 CONCLUSION AND FUTURE WORK
We developed a theoretical framework for safe policy improvement (SPI) that combines world-model and representation learning nontrivial settings. Our results show that constraining policy updates within a well-defined neighborhood yields monotonic improvement and convergence, while auxiliary transition and reward losses ensure that the latent space remains suitable for policy optimisation. We further provided model-quality guarantees in the form of "deep" SPI theorem, which jointly accounts for the learned representation and the reward/transition losses. These results directly address two critical issues in model- based RL: out-of-trajectory errors and confounding policy updates. Building on this analysis, we proposed DeepSPT, I, prineipled algorithm that integrates the thcoretical ingredients with PPO. On ALE, DeepSPI I is competitive with and often improves upon PPO and DeepMDPs s, while providing SPI guarantees.
This work opens several directions. A first avenue is to make pure deep SPI model-based planning practical. Our experiments with DreamSPI suggest that this is feasible but requires improved sample efficiency. Another direction goes beyond return optimization: a principled world model, grounded in our theory, can support safe reinforcement learning via formal methods, through synthesis (Delgrange al., 2025; Lechner et al., 2022), shielding (Jansen al., 2020).
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
All theoretical results are stated with explicit assumptions, and complete proofs are included in the appendix. The experimental setup is described in detail in the main text and supplementary material, including environments, hyperparameters, and training procedures. We provide the full source code as supplementary material enable reproduction of our results. Datasets used the experiments are publicly available (we use envpool Atari).
ACKNOWLEDGMENTS
We thank Marnix Suilen and Guillermo A. Pérez for their valuable feedback during the preparation of this manuscript. This research was realized under F. Delgrange's VUB OZR mandate (VUB- OZR4417) and was supported by the "DESCARTES" iBOF project. W. Röpke and R. Avalos are supported by the Research Foundation Flanders (FWO), with respective grant numbers 1197622N and 11F5721N.
REFERENCES
Rishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, and Marc G. Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. In 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021a.
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C. Courville, and Mare G. Belle- mare. Deep reinforcement learning at the edge of the statistical precipice. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances Neural Information Processing Systems 34: Annual Conference on Neu- ral Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pages 29304 29304-29320, 2021b. URL http:::tps/.ooo httptttp:://rrreeeeee /ttpc///ooccoo////eoooooooocoocc ings. ngs neurips iis.pps.oo... ppssccc/pppee//22//1/1 426edd55 www.oaaattbbttatt....aooao html.
Lucas Nunes Alegre, Ana L. C. Bazzan, Diederik M. Roijers, Ann Nowé, and Bruno C. da Silva. Sample-efficient multi-objective learning via generalized policy improvement prioritization. In Noa Agmon, Bo An, Alessandro Ricci, and William Yeoh, editors, Proceedings of the 2023 International Conference on Autonomous Agents and Multiagent Systems, AAMAS 2023, London, United Kingdom, 29 May 2023 2 June 2023, pages 2003-2012. ACM, 2023, doi: 10.5555/ 3545946.3598877 URL https: /dl...... acm. org hww.doiooocoo 5555 354 6.. 3598872.
Cem Anil, James Lucas, and Roger B. Grosse Sorting out lipschitz function approximation. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 291-301. PMLR, 2019. URL http: httpp/ppcceeng...eeee .ll..rr....... press / v97 aniil9a.nan.. html.
Raphaël Avalos, Florent Delgrange, Ann Nowe, Guillermo Perez, and Diederik M Roijers. The wasserstein believer: Learning belief updates for partially observable environments through reliable latent space models. In The Twelfth International Conference on Learning Representations, 2024.
URL https:tpp//...c http::::///ppeeeeeeeeee
net wttp:/iiee.eeeetttrreeeee forum? witeeiiiekktttteeeee
M. G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279, jun 2013.
Richard Bellman. Dynamic Programming. Dover Publications, 1957. ISBN 9780486428093.
Jacob Buckman, Danijar Hafner, G. Tucker, Fugene Brevdo, and Honglak Lee. Sample-efficient reinforcement learning with stochastic ensemble value expansion. ArXiv, abs/1807.00 076.006655.00 2018
URL https: / //api..../// semant i icscholar org/ C/c/pooppssi : 49577255.
Alberto Castellini, Federico Bianchi, Edoardo Zorzi, Thiago D. Simão, Alessandro Farinelli, and Matthijs T. J. Spaan. Scalable safe policy improvement via Monte Carlo tree search. In Andreas
11

---

Published as a conference paper at ICLR 2026
Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning. volume 202 of Proceedings of Machine Learning Research, pages 3732-3756. PMLR, 23-29 Jul 2023. URL https: http:///oooceccceee wgw.oeeeeiigg........oo. m.s.mml.. pressppp/////... wtw.ssg...sssssseessa.eea.e....e
v202/ castell wttps//ss/////caaasllcsssaa/ in miaa.. html.
Pablo Samuel Castro. Scalable methods for computing state similarity in deterministic Markov Decision Processes. In Proceedings of the Thirty-Fourth AAAI Conference on Artificial Intelligence (AAA1-20), 2020.
Pablo Samuel Castro, Tyler Kastner, Prakash Panangaden, and Mark Rowland. MICo: Improved representations via sampling-basse state similarity for Markov decision processes. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 301 13-30126. Curran Associates, Inc., 2021.
Florent Delgrange, Ann Nowé, and Guillermo A. Pérez. Distillation of rl policies with formal guarantees via variational abstraction of markov decision processes. Proceedings of the AAAI Conference Artificial Intelligence, 36(6):6977-6555, Jun. 2022. doi: 00.1609/aai..v36i v36.6.20622..
Florent Delgrange, Guy Avni, Anna Lukina, Christian Schilling, Ann Nowe, and Guillermo Perez. Composing reinforcement learning policies, with formal guarantees. In Proceedings of the 24th International Conference on Autonomous Agents and Multiagent Systems, AAMAS 2025, Detroit, Michigan, USA, May 19-23, IFAAMAS, 2025.
Josée Desharnais, Abbas Edalat, and Prakash Panangaden. A logical characterization of bisimulation for labeled markov processes. In Thirteenth Annual IEEE Symposium on Logic in Computer Science, Indianapolis, Indiana. USA, June 21-24, 1998, pages 478 487. IEEE Computer Society, 1998. doi: 10.1 10.100//ICS.. 1998. 705681. URL httpss:: doi orgg C 11 109 LLCS. 1 1998. . 705 18668.
Kefan Dong. Yannis Flet-Berlic Berliarr Allen Nie, and Emma Brunskill. Model- based offline reinforcement learning with local misspecification. In Brian Williams, Yiling Chen, and Jennifer Neville, editors, Thirty-Sevetth AAAI Conference on Artificial Inielligence, AAAI 2023, Thirty- Conference on Innovative Applications of Artificial Intelligence, IAAI 2023, Thirteenth. Symposium on Educational Advances in Artificial Intelligence, EAAI 2023, Washington, DC, USA, February 7-14, 2023, pages 7423-7431. AAAI Press, 2023, doi: 10. N0.1699/AAALL3376655550000 URL https: / /doi. org/ 10 10.1609/4 aaaii. w3776.25900
Ayoub Echchahed and Pablo Samuel Castro. A survey of state representation learning for deep reinforcement learning. Trans. Mach. Learn. Res., 2025. URL https:/ wwpppreeiee.neettt..eeeeeeeeeeee
forum?id=g0k3 4vUHt2..
Vincent François-Lavet, Yoshua Bengio, Doina Precup. and Joelle Pineau. Combined reinforcement learning via abstract representations. In The Thirty-Thing AAAI Conference on Artificial Intelligence,
AAAI 2019, The Thirty-First Innovative Applications of Artificial Intelligence Conference, IAAI 2019, The Ninth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI
2019, Honolulu, Hawaii, USA, January 27- - February 1, 2019, pages 3582-3589, AAAI Press, 333033358.. URL https: /ttp///ii.o/// http://d....oo..////// 1609//00..0///aaa...
2019, doi:
v33i01. 33013582.
Matteo Gallici, Mattie Fellows, Benjamin Ellis, Bartomeu Pou, Ivan Masmitja, Jakob Nicolaus Foerster, and Mario Martin. Simplifying deep temporal difference learning. The Thirteenth International Conference on Learning Representations, 2025. URL https:///oot /ttp:///peee.ie.eeem htt://ooommiddiieeeeeee
Matthieu Geist, Bruno Scherrer, and Olivier Pietquin. A theory of regularized markov decision processes. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, 1CML 2019, 9-15 June 2019, Long Beach, Cali- fornia, USA, volume 97 of Proceedings of Machine Learning Research, pages 2160-2169, PMLR. 2019. URL http: / / wwwppooooennnnge ngs. ml wrwss/rssssseeppe / //ggggisttt 1t 9a html.
Carles Gelada. Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare G. Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine
12

---

Published as a conference paper at ICLR 2026
Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 2170-2179. PMLR, 2019.
Mohammad Ghavamzadeh, Marek Petrik, and Yinlam Chow. Safe policy improvement by minimizing robust baseline regret. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Roman Garnett, editors, Advances in Neural Information Processing Systems 29: Annual Conference on Neural Information Processing Systems 2016, December 5-100, 2016, Barcelona, Spain, pages 2298-2306, 20 URL https: http::/:s///ppppeeeeeeeeee neuri wttpuuripscccccooooo.eooooo...ee cttppper///////// 201 6
paper/ html.
Mohammad Ghavamzadeh, Marek Petrik, and Yinlam Chow. Safe policy improvement by mini- mizing robust baseline regret. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29, Curran Associates, Inc., neurips ca/paper....eeeeeeeeeeeeeeeeeeee files/ pateee//pppee//1////// wttp//pper_fiieeeeppppeee//////e Paaela-pppppppppe pdf..
2016b. URL htttpp::/poocceeeeiititet file/ 00330555222000006633333333332232
Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville. Improved training of wasserstein gans. I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https: http::::s/pprooeednngg wewrips.cc//ppper.. fwweffilepppa https///appsepa/111///// #tle//922cbbceeecc233327777722f9
David Ha and Jürgen Schmidhuber. World models. CoRR, abs/1803.10112. 2018. URL http: ///axxiv.rggaa////////.
10122.
Danijar Hafner, Timothy P. Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL net/fose net//form?2e Ooabmidddddddodoool00
Wassily Hoeffding. Probability inequalities for sums of bounded random variables. Journal of the American Statistical Association, 58(301):13-30, 1963. ISSN 01621459. URL http:
/ /www. jstor. .
Bojun Huang. Steady state analysis of episodic reinforcement learning. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tier Lin, editors, Ad- vances Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
Shengyi Huang, Rousslan Fernand Julien Dossa, Chang Ye, Jeff Braga, Dipam Chakraborty, Kinal Mehta, and João G.M. Araújo. Cleanrl: High-quality single-fil implementations of deep rein- forcement learning algorithms. Journal of Machine Learning Research, 23(274):1-18, 2022. URL papers hvtrr//v23/21 223224421212.... html.
http: / /jmrr .rgg hawml..oor//pprrr////
Garud N. Iyengar. Robust dynamic programming. Math. Oper. Res., 30(2):257-280, 2005, doi: 10.1287/MOOR.I 1040.0129, URL httpt:h/ddo....... orooiioorgg11 100 1287 /moor.ooo...../ooo/ 1040. 0129.
Nils Jansen, Bettina Könighofer, Sebastian Junges, Alex Serban, and Roderick Bloem. Safe rein- forcement learning using probabilistic shields (invited paper). In Igor Konnov and Laura Kovács, editors, 31st International Conference on Concurrency Theory, CONCUR 2020, September 1- 2020, Vienna, Austria (Virtual Conference). volume 171 of LIPIcs, pages 3:1-3:16. Schloss Dagstuhl- Leiiniiz-eeetrure für Informatik, 2020. doi: 1044300LIPICS.CNCCUU 2020.3. URL https: /ww..ii.o///// org 100 4230 LLIP Ics CONCUR 2020.3.
J. L. W. V. Jensen. Sur les fonctions convexes et les inégalités entre les valeurs moyennes. Acta Mathematica, 30:175-192 1906.
Leonid Kantorovich and Gennady S. Rubinstein, On a space of totally additive functions. Vestnik Leningrad. Univ, 13:52-59, 1958.
13

---

Published as a conference paper at ICLR 2026
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020. URL https::tp//.t..........t /ttttt::://prrceeeee tttp:///ooceeedeeeoeooooo ngs neurips wwwppppp.ccppppppe Abstract htmll.
2020/hase f7efa4f86 ae Saw.Ea/A666..9999442222222
Jakub Grudzien Kuba, Christian A. Schröder de Witt, and Jakob N. Foerster. Mirror learning: A unifying framework of policy optimisation. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato, editors, International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pages 7825-7844 PMLR, 2022.
Romain Laroche, Paul Trichelair, and Remi Tachet Des Combes. Safe policy improvement with baseline bootstrapping. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 3652-3661. PMLR, 09-15 Jun 2019. URL https: pptppeeeeeee...... mlr. prwss//sssss...p.. perress/v91e v97/1aroche htt.ehllll.................m
Kim Guldstrand Larsen and Ame Skou, Bisimulation through probabilistic testing. Inf. Compul., 94 (1):1-28, 1991. doi: 10. 100111000990-54009199900000 0890-5400 (91) 90030-6.
URL https: / http://oooooooooooo oorg/1..111/////
Mathias Lechner, Dorde Zikelic, Krishnendu Chatterjee, and Thomas A. Henzinger. Stability ver- ification in stochastic control systems via neural network supermartingales. In Thirty-Sixth AAAI Conference on Artificial Intelligence, AAAI 2022, Thirty-Fourth Conference on Inno- vative Applications of Artificial Intelligence, IAAI 2022, The Twelveth Symposium Edu- cational Advances in Artificial Intelligence, EAAI 2022 Virtual Event. February 22 March 2022, pages 7326-7336, AAAI Press, 2022. doi: 10.1609/AAALV3317.2006 https: /oi......./// or /10 609660//aa v3617.2063
URL
Lihong Li, Thomas J. Walsh, and Michael L. Littman. Towards a unified theory of state abstraction for mdps. In International Symposium on Artificial Intelligence and Mathematics, Al&h 2006, Fort Lauderdale, Florida, USA, January 4-6, 2006, 2006. URL http: amwtiietttii....m..mmm.. cs umass. pdf.
edu/ att.ttmmmtteeee 6 wewpaattts/pccceeinnenee pttppppttcgggccoeooopppoopp
Marlos C. Machado, Marc G. Bellemare, Erik Talvitie, Joel Veness, Matthew Hausknecht, and Michael Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. Journal of Artificial Intelligence Research, 61:523-500 2018.
Walter Mayor, Johan Obando-Ceron Aaron Courville, and Pablo Samuel Castro. The impact of on-policy parallelized data collection on deep reinforcement learning networks. In Forty-second International Conference Machine Learning, 2025. URL https: http///pppeeeeeeerreeee/ forum? dininnyyddccyyyc zužhSo.
Alberto Maria Metelli, Mirco Mutti, and Marcello Restelli. A tale of sampling and estimation in discounted reinforcement learning. In Francisco J. R. Ruiz, Jennifer G. Dy, and Jan-Willem de Meent, editors, International Conference on Artificial Intelligence and Statistics, 25-27 April 2023, Palau de Congressos, Valencia, Spain, volume 206 of Proceedings of Machine Learning Research, pages 4575-4601 PMLR, 2023. URL https: http://////c/cceeeeeeeeee mlr mirssr v206 10.0mmeteeli23a...
Gaspard Monge. Mémoire sur la théorie des déblais des remblais. In Histoire de l'Académie royale des sciences avec les mémoires de mathématique et de physique tirés des registres de cette Académie, pages 666 605.
705. 1781.
Alfred Müller.
Integral probability metrics and their generating classes of functions. Advances in Applied Probability, 29(2):429 1997. ISSN 00018678. URL ht.p: www.jst.o........ooo.......o. org/ stable/ maable/142801.. 428011.
Arna Nilim and Laurent El Ghaoui. Robust control of markov decision processes with uncertain transition matrices. Oper. Res., 53(5):780-798, 2005. doi: 10.128770PRE.100000222 URL
fttps://dos .org/10r0000 287 opre. 1000. 0216.
14

---

Published as a conference paper at ICLR 2026
Efe A. Ok. Real Analysis with Economic Applications. Princeton University Press, 2007. ISBN 9780691117683.
Martin L. Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming. Wiley Series in Probability and Statistics. Wiley, 1994. ISBN 978-0-471611777-2 doi: 10.1002/ 9780470316887. URL www.h:ttps/.mm http:/oiiiooooooo http://:///o...r///.//// 1002/ 00000/0077470311680 /t00////11.00001111771111666
D. Revuz. Markov Chains. North-Holland mathematical library. Elsevier Science Publishers B. second (revised) edition, 1984. ISBN 9780444864000.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy Lillicrap, and David Silver. Mastering Atari, Go, chess and shogi by planning with a learned model. Nature, 588
(7839):604 609, 2020. doi: 10. 00008/44586-020-00333
John Schulman, Sergey Levine, Pieter Abbeel, Michael I. Jordan, and Philipp Moritz. Trust region policy optimization. Francis R. Bach and David M. Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, volume 37 JMLR Workshop and Conference Proceedings, pages 1889-1897. 1889.......rrg. 2015. mllr ppress//v37/cchumaaaa URL http: http:////.p.eeeennn....eeeee//// htt.mmmannn..........1.mm html.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. CoRR, abs/ (0007,06677,222 2017.
R. Serfozo. Basics of Applied Stochastic Processes. Probability and Its Applications. Springer Berlin Heidelberg, 2009. ISBN 9783540893325.
Thiago D. Simão, Romain Laroche, and Rémi Tachet des Combes. Safe policy improvement with estimated baseline policy. In Amal EI Fallah Seghrouchni, Gita Sukthankar, Bo An, and Neil Yorke-Smith, editors, Proceedings of the 19th International Conference on Autonomous Agents and Multiagent Systems, AAMAS 200 Auckland, New Zealand, May 9-13, 2020, pages 1269-1277, International Foundation for Autonomous Agents and Multiagent Systems, 2020. doi: 10.5555/ $398761.3398999 URL https:::: : /dl acm orgg doi / 0.550 0331//39987611333999009 3398908. /ot
Elias M. Stein and Rami Shakarchi. Real Analysis: Measure Theory, Integration, and Hilbert Spaces. Princeton University Press, 2005. ISBN
113869.
Miguel Suau, Matthijs T. Spaan, and Frans Olichoek. Bad Habits: Policy Confounding and Trajef-Traecctrry Generalization in RL. RIJ, 4:1711-1732, 2024.
Marnix Suilen, Thom S. Badings, Eline M. Bovy, David Parker, and Nils Jansen. Robust markov decision processes: A place where AI and formal methods meet. In Nils Jansen, Sebastian Junges, Benjamin Lucien Kaminski, Christoph Matheja, Thomas Noll, Tim Quatmann, Mariëlle Stoelinga, and Matthias Volk, editors, Principles of Verification: Cycling the Probabilistic Landscape Essays Dedicated to Joost-Pieter Katoen on the Occasion of His 60th Birthday, Part III, volume 15262 of Lecture Notes Computer Science, pages 126-154. Springer, 2024. doi: 10.1007/ URL https: www.ddi..or...oooooo / 10.
Richard S. Sutton and Andrew G. Barto. Reinforcement learning an introduction, 2nd Edition. MIT Press, 2018
Philip S. Thomas, Georgios Theocharous, and Mohammad Ghavamzadeh. High confidence policy improvement. In Francis R. Bach and David M. Blei, editors, Proceedings of the 32nd International Conference Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Conference Proceedings, pages 2380-2355 JMLR.org, 2015. URL htth:: / //proceeding igg...n........ ml ml prrsss / v37 stooassh 5 html..
Elise van der Pol, Thomas Kipf, Frans A. Oliehoek, and Max Welling. Plannable approximations to MDP homomorphisms: Equivariance under actions. In Amal El Fallah Seghrouchni, Gita Suk- thankar, Bo An, and Neil Yorke-Smith, editors, Proceedings of the 19th International Conference on Autonomous Agents and Multiagent Systems, AAMAS '20, Auckland, New Zealand, May 9-13, 2020, pages 1431 1-1--1 1439. International Foundation for Autonomous Agents and Multiagent Sys- tems, 2020. doi: 10. (0.5555/339876113399999 URL https: / /dl. acm. org/ hetg//o//100 /t.ii/1..55555 5555/ 3398761. 3398926.
15

---

Published as a conference paper at ICLR 2026
Robert J. Vanderbei. Uniform continuity is almost Lipschitz continuity. Technical Report SOR-91-11,
Statistics and Operations Research Series, Princeton University, 1991.
Leonid Nisonovich Vaserstein. Markov processes over denumerable products of spaces, describing large systems of automata. Problemy Peredaci Informacii, 5:64-72, 1969.
Cédric Villani. The Wasserstein distances, pages 93-111. Springer Berlin Heidelherg, Berlin, Heidelberg. 2009. ISBN 978-3-540-711050-99 doi: 10. 00.10077978-3-54--700000 //doi. org.or//////oo..
https:
Patrick Wienhöft, Marnix Suilen, Thiago D. Simão, Clemens Dubslaff, Christel Baier, and Nils Jansen.
More for less: Safe policy improvement with stronger performance guarantees. In Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence, IJCAI 2023, 19th-25th August 2023, Macao, SAR, China, pages 4406 4415. ijcai.org, 2023. doi: 10.24963/JCCAL22033 490. URL https::///// : /doi.. org/r0 2 4344933340 i jcai. /42.3333caa..22333//1 2023/490.
Wolfram Wiesemann, Daniel Kuhn, and Berę Rustem. Robust markov decision processes. Math. Oper. Res., 38(1): (1:1)11533188 2013. doi: 10.1287/MOOR. 120.0566. URL /sttps://di... org/ 10 10.1187/11 moor 11.11....555
Frank Wilcoxon. Individual Comparisons by Ranking Methods, pages 196-202. Springer New York, New York, NY, 1992. ISBN 978-1-4612-4380-9, doi: 10. 01177997781-611-4338---1111 URL
https: /ttps://do/ /ww../oii.////////////../ orooorggg/// / 1007 191//9988111----888888
Chenjun Xiao, Yifan Wu, Chen Ma, Dale Schuurmans, and Martin Müller. Learning combat compounding-errov error in model-base reinforcement learning. CoRR, abs/1912.11200 2019.
Tianh Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y. Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. MOPO: model-based offline policy optimization. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Tien Lin, editors, Advances in Neural Information Processing Systems 33; Annual Con- ference Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL httpthttaa//pooceeeeeeee wtw.rocceedinng................. httppps..cpppppppe htcpc//ape//////////a/ html.
Tianhe Yu. Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. COMBO: conservative offline model-based policy optimization. In Marc' Aurelio Ran- zato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, ed- itors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pages 28954 2021. URL https: http:///pcceeenn..eeeeeo/o neweeddiiinnnnnnn.....ee.. www.inuur.r.ppppeeeeeeeeeeeeeeee wctppappr///2//////.... / #29a17974690222
Abstrrat.. html.
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Representations, 2021. URL https: /ttp///pnweee..eeeeeeeeo net /forrm? 1F==DFCWDKREED.
httpps://penrrevee..
16

---

Published as a conference paper at ICLR 2026
Appendix
A REMARK ON VALUE FUNCTIONS AND EPISODIC PROCESSES
An episodic process is formally defined as an MDP M = A. 87, 7) where:
there a special state Syeser E S, intuitively indicating the termination of any episode;
the reset state does not incur any reward: Rreee a) = 0 for all actions A:
(1i1) Srser is almost surely visited under any policy: for all policies E II, Pr ( (s2, 3i: 8: Sreset }) 1;
and
(iv) M restarts from the initial state once reset: P([st} Sreset, a) 1 for all A.
Note that by items (iii) and (iv), Sreser is almost surely infinitely often visited: we have for all x E that
Pa (1(0
2 O, 3j i: 8j 8yeset })= =
Alternatielly and equivalently, episodic process may also be defined without unique reset state by the means of several terminal states, which go back to the initial state with probability one.
An episode of M is thus the prefix so: 00..... 1 St of trajectory where 8t Shewe and for all I, 8; # Sreeet- Notice that our formulation embeds (but is not limited to) finite-horizon tasks, where an upper bound on the length of the episodes fixed. The average episode length (AEL) is then formally defined AEL(T) E,, [T] with
T(r)
Σ 1=0
11+1).1f8. = Syeser and Vj < i, 83 7 Sresei)
for any trajectory r = (8t; (t)
Often, when considering episodic tasks, RL algorithms stops accumulating rewards upon the termination of every episode. In practical implementations, this corresponds to discarding rewards when flag done, indicating episode termination, is set to rue. In such case, we may slightly adapt our value functions as:
E
II 1s + sreser . R(se. az) 80 = if s of Syeset
(*(s)
otherwise;
or, when formalized as Bellman's equation:
(s, = R(s,a) + EewP((ea)) V*(8') if & Syeset
[o
otherwise; and
(*(() E a-1a(-1s)
(s)
All our results extend to this formulation (cf. Remark 2).
Remark (Occupancy measure). In RL. theory, the discounted occupancy measure
(=) Σ 'P (f(s) = =s})
is often considered as the default marginal distribution over states the agent visit along the interaction, mostly because of its suitable theoretical properties. In fact, for any arbitrary MDP, nr the stationary distribution of the episodic process obtained by considering reset probability of from every state of the original MDP (Puterman, 1994; Metelli et al., 2023). Again, we contend that all our results can be extended to the occupancy measure with little effort.
17

---

Published as a conference paper at ICLR 2026
В POLICY IMPROVEMENTS THROUGH MIRROR LEARNING AND CONVERGENCE GUARANTEES
In this section, we prove that NC (Eq. 2) is a proper mirror learning neighborhood operator. As consequence, appropriately updating the policy according to N is guaranteed to be an instance of mirror learning, yielding the convergence guarantees of Theorem
For completeness, we recall the definition of neighborhood operator from Kuba et al. (2022).
Definition 1 (Neighborhood operator). The mapping N: II 2 (mirror learning) neighborhood operator, if
(continuity) It is a continuous map;
2. (compactness) Every ((T) is a compact set; and
3. (closed ball) There exists metric d: 11 x 11 + 00), such that for all policies E 11, there exists > 0, such that d( x²) & c implies n' (T))
The trivial neighborhood operator is N(T) = 1l. Lemma No a neighborhood operator.
Proof. Henceforth, fix policy и E TT. When taking the supremum, infimum, maximum, or minimum value over states and actions, always consider actions be taken from the support of the behavioral policy (in the denominator of the quotient).
Item 2 (compactness) is trivial due to Dir (T, 2 2 - C and DiR (TT, x²) S C for any n' E E (T). This (T,,
means N (x)) contains its extrema, i.e., all the policies x' satisfying Dirt (TT, x²) ) = 2 C' and DiR (T,, x') S C, or DiR' (T, x') 2 2 C' and DiR (x, C.
In the following, for any II and sequence (Tn) we write nn - for the convergence of the sequence to π with respect the metric
|| Telloo if supp(ar s)) supp(T2((1 vs E S, and otherwise.
d(T1, T2)
(7)
In other words, means that n. converges to in supremum norm n - oo when the support of the converging policy stabilizes and becomes the same as the limit policy.
Let prove item I (continuity). We show that Nr is continuous correspondence by showing it is upper and lower hemicontinuous (Ok, 2007).
No is upper hemicontinuous (uhc) if it is compact-valued (item 1) and, for all policies E II and every sequences (In)no and (##)220 with T², E -'T(TTT)) for all n 2 O, Tn - x and T. + n implies n' E (N) Let (Σn)m20 and (Tz) be sequences of policies with x/ e (mT) for all n 2 0.
Fix S and A. Consider the mapping
a s)
faa: [1T x') II x II a supp(=T( s))} + [0, oo), ( (π,
a(a
It is clear faa is continuous since the application of ı to a[a s) is continuous and the division of two continuous functions also continuous (when considering actions from the support of x(- s)). Importantly, for ext E (sup, inf). D x²) ext (J, rT): S,a supp(ar() [s]]} is also continuous: since S and A are finite, the supremum (resp. infimum) boils down to taking the maximum (resp. minimum) of finitely many many continuous functions, which is continuous operation.
Now, assume that Tr + π and ay + n'. The continuity of D means that D (Txy =/2) + D (π, T'). Since Ti AC (TT), we have DiR' (Tny 7/) 2 2 C and Dik (mn, m/) S for all n 2 0. By the fact that D (may m's) converges Dext (π,, x²) for
(inf, sup). we also have that Dinf (IT, x') )
and D (T, T') C.
Then, is uhc.
NoC is lower hemicontimuous (Ihe) if, for every policy æ, sequence (Ta) n20 with म, - ग,, and policy e eNO()), there exists a sequence (Tz)a20 with oú, + n' and such that there is a 2 0 from which, for all n 2 no. Th EVO (m,).
18

---

Published as a conference paper at ICLR 2026
Therefore, let (In)n2o be sequence of policies that nn and r' E <N(x). Since Tn n, have
Vg 3mg E N: Vn 2 no: || Tn Tllo0 S and supp( 8)) = supp(=(-) 8)) Vs E. In particular, this holds for < Kmin/22 where Tmin min {T 8): 8 S, a supp(#(- [])] Let no 2 the step associated with 5 < #mix//2 and n 2 no- Write 6, = |m x 20 and let
2Con
E (0,
min (C 2C',
Construct sequence (z/)noo that, for all E S A, and n 2 no.
Th 8) = (1 - en) . T'(a 8) + en (a 8). Intuitively, TV, is a mixture of distributions x'( |s) and (- 8). Consequently, T s) is a well-defined distribution. Finally, note that Th - x'' because 6, + and does Cn.
Now, we restrict our attention to a E supp(Tn.( s)). Note that since n. stably converges to π with its support, т has the same support n.. Furthermore, since x' E (C), n/ has also the same support kn. In consequence, nh has the same support as Ta-
Having that said, start by showing the upper bound:
T (a
8)
x (a
(1 nn)
En
x
Tn (a
C. 1 (a
8)
I (1 En)
+ En
(because x'aa| S C- :Ta s))
(a 8)
C. TT 8)
$ (1 - Gn)
+ Cn
(because (a 2 x(a 8) - 6,))
T(a
6n
(1 en)
1 6n/x(a10)
S (1 6n)
6n/
Cn-
Note that for all x E [0, 1/2],
1 1 2x because 1 1+ 2x
1 x
(1+2r)( -x) -1
x(1 -xx)
20 =
20 +
2
1 x
1-x
1 1- I
Then, since 0 < 80/maain < 1/2, we have
x s)
<1 (n)) C. (1 + 26m//miiu) + En
Tn ( 8)
Let In = 1 1+ 26n,, and note that
208n
20 6 Teain
-20 . ón//min Amin
C(1 xn)
En Tmin (C - 1) + +20h
+ 20 5n/ 6n/smin
- --20 6n//mii
- Jn C
Then,
T' (a 8)
< (1 -
(1-6n))r Cn) C Cn
(a s)
En C En . En C + En
C(1 In)
In C+ C(1 En.)
In C
1 - In C
1 1- In C
Fn C(( En C)
En C2 - Fn,) ) C(1 - Fa)
- In C
In C x22 C2 2
Fn 2 + x
2 2 C - C
1 xn C
19

---

Published as a conference paper at ICLR 2026
In C2 +C
- In C
1 In C
C
1 In C
C,
which means that DiR (m m²) S C.
We now show the lower bound:
T (a] s)
π r'(a 8)
TT (a (1 kn)
+ En
Tn (a 8)
(2 C) ((
8)
2 (1 - nx)
+
(because J' 2 (2 -
(2-C)- C.- [Ta s))
(a
(2 C* (a 5)
2 (1 6x)
+ +
(because (a(a S <(a s) + 6,))
<(a 8) + 8n
(2- C)
= (1 (x)
+ Cn
1 bn/T(a 8)
C)
2 (1 (..)
+ En
1+ 6n/mtwin
2 (1- c.) (2 - C) (1 11--6y 6n/ + En
(because for all x E R, Itx 2 1 x)
= (1 - (x) (2 C - -2. -2.6n/ 6n//Tmiu + C (.nn/Tuin) + Cn
(1 [1-6,)) (2 C (C 2) 5n//miia) + cn.
= 2 C C (C'(22 6n//xmin - (n - C (C'- 2) 6n//mmin) + On
2 C++CC - 2) 6n/ Amin + <n (C - 1 + (2 ())_8//Fmia)
= 2 C + (C - 2) /n//Tmia
+ en (C - 1) +en .(2 - C) 6
2C8, (C - 1)
2C8, (2 CC)
= 2 - C + (C - 2) 6n/
Mmin (C -C- 1) + 2Con Mmiu (C - 2C'nn
2 C + +nn ( C -
2C' (C-
2C6,
-1 Tmiin (2 C)
Tiin
Tmin
(C (Cnin(C-11 1) 2C8, Miin (C 2C
22-C. C.
To see how we obtain the last line, note that it suffices to show the content of the parenthesis multiplied by 8, is greater than zero, i.e.,
C 2
20 (C-(1-)
2C6, - (2 C)
2 0
Tinn
Tuin (C (C-1)+2CC
Tuin (C' - + 208
2C 1C) Amin (2 -))
C
Tminn (C - 1) + 208,
Tmin
20Tmin (C (C- 1) 2Con
(2 C) 2 (2 C) (Tmin (C 1) + 2C8m)
2C min
(C - 1) 2 (2 - C) (Amin (C - 1) + 2C5, 2C8
1
20 mmin
(C 1) 2 (2 C): (Tmin (C- 1])
20 2 2 C,
which is always satisfied because C > 1. Therefore, since this holds for any s S and both T/, and n, have the same support, we have that DiR! (M,, x') 2 2 2 C.
Thus, we have Disl (T, x²) 2 2 - C' and D (T,, n') S C, Th E (Tx)..). Therefore, is lhe.
Since NCC is uhe and Ihe, it is continuous. This concludes the proof of item 1.
It remains to show item 3. Lete (C'- 1).min T(a 8), with taken from supp(T s)). Assume d(z, S (cf. Eq. 7). For all s S,a a E supp(T((- s)), we have
n (a] 8)
20

---

Published as a conference paper at ICLR 2026
< 8) +c
< 8) + (C - 1) min T(a ss
s,
<a(a 8) +C 1) a(a 8)
= = 8) ()+ 1)
x(a s.
or equivalently:
<a 8)
x(a 8) S C.
It remains to show the lower bound:
z' 5) 2 x(a 8)
= x(a 8) (C - 1)
min (a s)
2 x (a 8) (C 1)
T(a 8)
<a 8) (1 C + 1)
x(a S)
2 x(a s)(2- - C),
or equivalently:
a'aa s)
2 C. x(a s)
This concludes the proof of item 3.
Then, Theorem 1 is obtained as corollary of Lemma 1, and the fact that the update process
Tn+11 arg sup IE E [ATH (8, x' (*.) S~ErT avri((1a)
is an instance mirror learning (Kuba et al., 2022).
C CRUDE WASSERSTEIN UPPER BOUND
Lemma 2. LetsE E S and a E A. the following upper bound holds:
(( P(. (()) S Rg EZ 2(((8))
Proof.
W(opP(- 8,a), P(- P(-1)(s),a))
sup 11f
(I)
S (g-p(()sa))
sup ((((') E (F
11 W(Dasr)) P(. o(s),a)
EvaP(1ssa)
min xEAlEarr P((d(e)))) EyrP(()s,a) EgrNP(Ho)())))) a(o(s'),5))
Ee-P(1sa))
E E d(31, 32)
(2)
Here, (1) corresponds to the dual Kantorovich Suntorovich-Rubinstinn formulation (Kantorovich and Rubinstein, 1958) where Ipp corresponds to the Lipschitz norm, while (2) follows from the primal Monge formulation (Monge, 1781), with trivial coupling induced by )o()))) the Dirac measure with impulse 6(8').
21

---

Published as a conference paper at ICLR 2026
D REMARK ON SAFE POLICY IMPROVEMENT METHODS
Standard principled safe policy improvement methods (SPI; Thomas et al., 2015; Ghavamzadeh et al., 2016a; Laroche et al., 2019; Simão al., 2020; Castellini et al., 2023; Wienhöft al., 2023) do not consider representation learning. Instead, SPI methods assume S == S and learn R, P by maximum likelihood estimation with respect to the experience stored in B collected by the behavioral Tb. Then, the policy improvement relies on finding the best policy in M that is (probably approximately correctly) guaranteed to improves on the behavioral policy (up to an error term Ç > 0) against a set of all admissible MDPs, called robust MDPs (lyengar, 2005; Nilim and Ghaoui, 2005; Wiesemann et al., 2013; Ghavamzadeh et al., 2016b; Suilen et al., 2024):
sup p M) xell
such that
arg M') 2 M') - -,, where (M.
R(s, a) R(s,a)) S Baux . c(s,a) E(M.e) = {MM= M= (S. drv (P(- s,,), P(- s.)) S e( Vs V
and
e(s, a) being an error term depending on the number of times each state and action are present in the dataset B, and drv being the total variation distance (Müller, 1997) which boils down the L1 distance when the state-action space is finite. To provide probably approximately correct (PAC) guarantees, the state-action pairs need to be visited sufficient amount of time, depending on the size of the state-action space, to ensure e sufficiently small.
Note that the reward and total variation constraints are very related to our local losses L and Lp: the representation corresponds here to the identity and drv coincides with Wasserstein as the state space is discrete (Villani, 2009). The major difference here is that the bounds need to hold globally, i.e., for all state-action pairs, which make their computation typically intractable in complex settings (c.g., high-dimensional feature spaces).
We argue this objective is ill-suited to complex settings. First, classic SPI does not apply to general spaces. Second, assuming we deal with finite, high-dimensional feature spaces (c.g., visual inputs or the RAM of video game). it is simply unlikely that B contains all state-action pairs. SPI with baseline bootstrapping (Laroche al., 2019) allows bypassing this requirement by updating To only in state-action pairs where a sufficient number of samples are present in B. Nevertheless, this number is gigantic and is linear in the itte-actton space while being exponential in the size of the encoding of y and the desired error Ç. This deems the policy update intractable. Finally, mentioned, standard SPI does not consider representation learning. This is further obstacle to its application in complex settings.
E SAFE POLICY IMPROVEMENTS: PROOFS
Notations Henceforth, we denote by vR the value function of the world model M obtained under any latent policy R E IT. When is clear from the context that ø is the representation used jointly with a latent policy n. we may simply write V instead of (=xod) for the value function of executing क in M. In the following, we may also write as ~ Err a shorthand for first drawing
&x and then T(- s) for any policy
II.
We start by recalling a result from Gelada et al. (2019) that will be useful in the subsequent proofs. Lemma 3 (Lipschitzness of the latent value function). Let M be latent MDP and be a policy for M. Assume that M has reward and transition constants K and p with K R
< 1/x. Then, the latent value function is
K (1/(1->k))
(iipss-k))--ipschtt,, LLipschitz i.e., for all 31, 82 S.
K d(5
VV((5)) - V*)
-V61)
1 1- R
Note that the bound is straightforward when the latent space is discrete and the discrete metric 1l (+) is chosen for d: the largest possible difference in values is 2Rws/117
We also consider bounding expected value difference between the original MDP and the latent MDP by the local losses evaluated with respect to behavioral policy Tb Importantly, the expectation is measured over states and actions generated according to To. whereas the values correspond to those evaluated under another latent policy म. The following Lemma states that the value difference yielded by a latent policy can be measured according to another behavioral policy, provided that the latent policy lies within a well-defined neighborhood of the behavioral policy.
22

---

Published as a conference paper at ICLR 2026
Lemma 4 (Average value difference bound). Let To II be the behavioral policy, ( o 6) E // (Tb) that R I1 and 6: S + S is state representation. Assume M is equipped by the Lipschitz constants K and K and Kv = (1-,k)) Assume that K3, is strictly lower than 1/7. Then, the average difference of value of M and p
M under गt is bounded by
"R Er +7Kv L.p IIp 1/D/ (m..)) -
(T)
S
Proof. The proof follows by adapting the proof of (Gelada et al., 2019, Lemma 3) by taking extra care of the behavioral policy. Namely, want to evaluate the value difference bound for the latent policy я, assuming states and actions are/have been produced by executing the behavioral policy лb. The idea is to incorporate the divergence from Th to म in the bound, formalized the supremum IR between the underlying distribution of the two policies.
E SNIETD
V* (s) (d(s))
E S~EET
E a><((-((())
R(s, a) +
E
[V'(69]
R(o(s), +7
E PmP( (()).a)
(8))
8(-P(11),))
(NT(116(s))
E
E a~* (-10(())
R(ssa)- - R(((())))) +7
amT((1d())
P((aaa) 8 p(-((-)())..)
E
E
[B(ssa)- - H()(s), +y
Exp
(-H(-10())
amT(1
5' s(-l(a)))))
(')((-(())..)
E
E
(R(s,a) - R()(s),a))
(-(()-)())
E am{(11o(s))
E
(0((5)) +
E
[V)((())
stuP(-la,a)
s(-saa)))) ;~~((1(()))))
< E
E Smtn, a~# (-(()((
[R(s,a)-((dh)) - R((((), +7 E
((66)-'((0(55)) +Y
E of P(lsa 3 (()(()))) (Jensen's inequality)
(V(((s)))
< E
E smiEr, a((((1)))))
R(s. R()(s).)) +7 Sn amF(1
E s'aP(a)a,)) ('~T(((((s).a))
-
? E
E
E s' w(((ss).))
[V*(s') (ò(s'))
S a)){(((()))
(Triangle inequality)
{a ((s))
+ E E Sn a>)))
E P()aa)) P(-P(-)())))))
[D*(6())) - V7
mb(a (8)
T[a x(s))
E
[V*(8') - Va (
sn (-
ab(a 5) s'MP())a.)
(because supp (- ((s))) = supp( s)) for all E S)
23

---

Published as a conference paper at ICLR 2026
D (Tb; R)
E 5,amEn
[R(s,a) - R()(s),,)) a)] + +7. DiR (Tb; π)
E s,anEm m.
E
(7('')
E 3)-((10()),.))
(=(8')
8 Flonda P(l
+ +
Dik (Tb: x))
E s,an
E s(11()))))))
[V - (=(6())1
(because DiR (Tb; R)
sups.e
[(al((3)) ma(1l))
-DiR (Tb: .) LR 7. DIR (Tb. π)
E 8,0~5m
E
V (5')
E P(-((())).
V*(3")
F'mwby ((1ss,))
+ DiR
(Tb: π
E aaan
[V*(s') ( (d({s)
s'~1(-11,,))
(by definition of L )
SDR (Th, π) -. + RKv. DJR (mb: 1))
E s.antm
W3 (or P(. 88 88,)), P(. (())
+ y
IR (Th, π
E Sent of P(-lea))
(by Theorem and the dual formulation of Wasserstein)
=DiR (Tn. π) (Th.
(I E + tgKv E + yDiR (Th, R)
E a
[V*(5') (=
P(-saa)
(by definition of T
<DiR (Th; π) (I +7Kv 1.(00) + y DDDi (TT, .).
E s (-1saa)
E
(*( - V* (b(s)))
(Jensen's 's inequality)
=DiR (Tb; т) (I +yKv -b + (Tb,
.)
E ("))
(as En, is stationary measure)
To summarize, have:
DouD DiuP(TTT (Tb, R). (L ErD + 7Kv 18x0) + 1ip
V
a
Or equivalently,
-1- 7D sup (Tb; 7)) E
- (ss < sun (Tb, 7) (Tb,
(I Lm + +KV L
2K 1.p00
(s) (d((s)) < D (Tb; 1)
1- sup (Tb, T))
L + +k I
1/D (m,, 7
which well-defined because D (Ab, 7) is assumed strictly lower than 1/7.
In the main text, we made the assumption the environment is episodic Let us formally restate this assumption: Assumption 2. The environment M and the world model M are episodic.
Assumption 3. Vs E S, dos)) = Brese if and only if
S
Note that, as mentioned in Section 2, Assumption 2 ensures the existence of stationary distribution & and the ergodicity of both the original environment and the latent model. Assumption 3 guarantees that the reset states are aligned in the original and latent MDPs
We are now ready to prove Theorem 2.
24

---

Published as a conference paper at ICLR 2026
Theorem 2. Suppose y 1/2 and K < 1/7. Let C E (1,1/), Th II be the base policy, (7 (Tb) where π E II is a latent policy and Φ: - S state representation. Then,
1 / Kv Lp p (p(ToO, M| S ABL(TT)
Sma
3/D (1b
where (T)( denotes the average episode length when M runs under nb. Ky = K (1/()))))), and S L. are the local losses of Eq. over the stationary distribution Ex, induced by nb.
Proof. The first part of the proof follows by the expected value difference bound of Lemma 4. The second part of the proof follows by adapting of the one of Delgrange et al., 2025, Theorem 1, where the authors considered discrete latent MDPs and reach-avoid objectives (rewards were disregarded).
Our goal is to get rid of the expectation. First, note that for any measurable state so that Em, (f+}) > 0, we have V (s) V7 (d(s)) & ((+), Ex'-E,, (s*) - pd (+(s))) For simplicity, we write (8) as shorthand for E., (1+}) when considering such states. Second, note that as Sreee is almost surely visited episodically (Assump-
tion 2), restarting the MDP (i.e., visiting Sreses) is a measurable event, meaning that Syese has a non-zero probability Sx, (Sreee)) E (0, 1). Then,
p(F o Φ, - - (((T,
(8)
(*((s) - [313)]]
(9)
1 V (5(()) - V (8)
(10)
V (Sreser) - v7 (5((rese))) )
(by Assumptions and 3)
E
(11)
7 Em (Srewet) 1ith // + Ky 5.
(12)
Em, (Breset) (1 (VD52R (RoiT)
Finally, the result follows from the fact that (/mm)(mmm) (smm) corresponds to the AEL. Indeed, when M is episodic, it is irreducible and recurrent (Huang, 2020); thus, given the random variable
TT(T = So. Co, 01 .1. = Σ T T = and # for all 0 <<<<<< T},
we have (s) 1/8. (e,TTlls-00 [=0=s] for any s S and stationary policy n, where E, [T Sg s] is the mean recurrence time of under x (Serfozo, 2009, Chapter 1, Theorem 54). In particular, this means that (/Em)( = [T Syww 800 = Sres)] F., [T] is the AEL of M under mh. which yields
E Kb
I. /7 + Ky Lp
p(Too, -ac.MM)-pFFRMM. (R] s E,, (T)
1/D2 (mi:)) -7
Remark 2 (Extension to episodic value functions). Lemma 4 and Theorem 2, we considered the standard definition of value function. One may wonder whether the results hold when considering episodic value functions, as defined in Appendix A. It turns out that it is the case, as one can easily adapt the proofs for those particular value functions.
We start by adapting the proof of Lemma
E. Sr/Er
(*)
V (o(s))
25

---

Published as a conference paper at ICLR 2026
B amF((lo(s))
R(s.a) ty
E
[v'
E s
1 1s # Sreset)
R(d(s), +7
E
[+
am)(-10()))
E S.- an<7(-lo(a))
E a(s,a) +y
[v*(})]
E amT(11d())
R()(s),)) +7
V (3()')
s(-P(-)),a)
8'mP((1(()).))
The remaining of the proof is identical.
Concerning Theorem 2, we take a detour by defining new value function as
U (+)
E
R(s.a) +7
U (5')
{8' o Speae]]
4s E S
a~T((1o(s)
The latent counterpart 0 is defined similarly. By definition of the episodic value function (Appendix A) and since V (Sreset) = 0, it is clear that
(*((s) (s) (s) 1 (s # Syeser)
if + Snewe otherwise; and
UU"( (8) (8*())
if f ((8mse)) otherwise.
(²(s)
V
1 1 + (Syese)))
(13)
Therefore,
E errEr,
(s) U (ó(s))
[aa d(s)).
<
E 6,0
(Rs,s a) R(d(s).a))
mb(a 8)
F(a (())
+
E
E
(*' {s' * S }1
E (T'('') 1 (3' * ((
sTarET
Tb(a
P(
8'~(((()).a)
(Triangle inequality and importance sampling)
=D "TTT R)
E saa
R(s, - R(((s), + 7
sup (Th, T)
E s
E
(=('')
E
V=(3")
o P(-s.a)
3'~P(-l(s))a)
(by Eq. 13 and definition the SIR)
< Duup (Tb, (ii ())
E stanie
[((sa) (o(0))(((sss)))))) 7 -DD (To.
D F)
F s,amen P(-saa)
(V*('') - -V=)))(((())]
+ + D
(Tb, πt)
E saamEe,
E
(*(33)
B ((-ssa)
(TP(((())
(Triangle inequality)
<DD 7) Em + ( .)
1R (Tb,
sup (Tb.
E'
(*())) (d(s))) + T (Tb, #) Kv L.p
sruEe
(by the same developments in the proof of Lemma 4)
Ler KKy I6
<Dit (Tb,
- + Dik suu (Tb, (TbT)((TTTT)) R)- R.- (T
+ 7 D
sup Duuup(TT (To. 7)
Kv L
(Lemma 4)
3/D (TT. (m.... -
-Daup (Tb;(TTT Rt
15n 1
7
7Kv L
1+
Duupp(Tb 7) -
(Tbr πτ) 111 IR
=Doup -DiupP(Tb (Tb;
1 .KV 1,p00)
1+
?
Dou (Tb, T) 1 -
Len + +Kkv Li
1/D (m..F))))))) -
26

---

Published as a conference paper at ICLR 2026
Now, in the proof of Theorem 2, it suffices to replace Equation 9 by observing that, in the episodic case, have O - p p M) = (s)) - 47 (01) = (s:) - 0 (s1)
(again, by Equation 13)
11 - U (51)
1 (8reset) 0* (<((S
Modulo this change, the remaining of the proof remains identical; one just needs replace the occurrences of EswEF, (s) - V* (d(()) by ES-E-, (U* (
Since the subsequent results all rely on Lemma 4 and Theorem 2, they all extend to episodic value functions.
Theorem 3. (Deep, Safe Policy Improvement) Under the same preamble as in Thm, 2, assume that & if fixed during the policy update and the behavioral is a latent policy with To == T o Φ and T6 E II. Then, the improvement of the return of M under it can be guaranteed on Th as p( R M) p( 2 p(T,, M) - p(Tb: M) S.
where Ç = (EL(TT) (L /1 + KvL
1 -
1
(1-7)
Proof. First, note that
p o M) MM
=p( o M, - M(F, + ((R pp(mb M).
(14)
By Theorem 2, we have with Dik (Tb, mb) = that
Lg / + Kv L.
Esb
[p(T,, - M)] S R.M [T]
1 -
which implies that
/7 + Kv I.pp
M) - M) E.M T
1 1-7
1 /7 + Kv
M)
p( + Way [T]
(15)
1-7
On the other hand, we have
L /- Kv Lp Sop
( M(F, E. [T]
1/ D (mbi
which implies that
ET / + Ky Lp 1/Dgeu (m..
( - p( M) 2 mp
(
(16)
By plugging Equations 15 and 16 into Equation 14, we get the desired result: p Φ, M) pl M) = p( Φ, M) MT..M)
+p(n,M)
p (T M)
LR /r + Kv 1.5T Em 1/ 1/Di (ma
Erh /r + Ky L
E
M) B.M [T]
1-7
2 -- J.M [T]
1. + Kv 16r + M( p (To, M) BM [T]
La
Emy
/r + Kv Lp
1/ uop (5)n, -
1 1-7
=P(T,M M) ((Fo, - E.M [T]
Em / + KvL 5m
1/Dae (m),)) (1 -
1
+
1-
27

---

Published as a conference paper at ICLR 2026
In the following, we provide probabilistic version of Theorem 3, as is standard in the SPI literature. Essentially, we derive probably approximately correct estimations from interaction data of LR. Lp. Then, we use those estimations to get an approximation of the enror term of the safe policy improvement inequality of Theorem 3.
Those PAC guarantees rely on a discrete latent space. While it may seem restrictive, learning discrete latent spaces turns to be beneficial not only theoretically (e.g., it yields trivial Lipschitz bounds the latent reward and transition functions), but also in practice (see, e.g., Hafner et al., 2021).
Finally, note that provide two versions of the theorem: (1) where we have access to upper bound of the AEL (which is mild in practice), and (2) another one where this bound cannot be derived. The latter case yields additional challenge as need estimate the AEL from sample states drawn according to the stationary distribution. In this case, the bound yields a probabilistic algorithm that is guaranteed to almost surely terminate without any predefined endpoint, depends on the current approximation of the losses.
Theorem 5 (Probabilisticc Deep SPI with confidence bound), Under the same preamble as in Theorem assume now S is discrete. Let (/,, de. Tt, st) 1tt S T} be set of transitions drawn from Ex,, by simulating Ma i.e,, 8t ~ Eoy Og ~ Tb(. 8r), Ft. = R(se, ar), and 8/ ~ P(- 8t, as) for all 1 S. T. Let E,
and define
r
T
T 1
Lp =
Σ
((((s)) (())), in = T t=1
+
Σ
r - R(O()) Exeas ==
= =8resct) 7
T
T 1=
R ==
I- and R = max {1, 11,4 Then, the policy can be safely improved as
p( o Φ. M) p(Tb, M) 2 M(a. p( -c.
(17)
with probability at least 1 - 8 under the following conditions:
(1) one has access to upper bound L AEL(TT)), the number of collected transitions is lower-bounded by T 1 -R" log $ (V)++ and Ć L (I<//> + Kv Lp
or
(2) without access to such bound, we take
a/Emeer (La/7 + Kv Lp) +E+ n.. 01/7 + Kv)
R² log(%/3)
T
max
2
EEreset
and { ==
1. (La/7 + KvLp te.
Ener
Proof. Let E,8 > O. First, note that need T R log(3/2)
to satisfy both (a) LR + Sr and
(b) Lip+e> + 1fr0 with probability 1 8 and T
It log(s/s)
satisfy simultaneously (a), (b), and (c)
ga
Enser - - < Em, (Sneet) with probability - 6. This statement is proven by Delgrange et al. (2022) and Delgrange et al. (2025). The result essentially due raw application of Hoeffding's inequality and the fact that Wasserstein boils down to total variation when the state space is discrete (Villani, 2009).
Let E' 0.
Case 1. Assume we have an upper bound on ABL(T)), say L. Then follows that
S L. L Kv
(C is the safe policy improvement error term of Theorem 3)
L (iRte' + kv(Lp +
with probability least 1 whenever
R log(5/2)
T>
e²
To ensure error of at most E, choose e' such that
L (Lute' + kv +6))=+ 5 L. (4e + KulL) +e.
28

---

Published as a conference paper at ICLR 2026
Equivalently,
L E L + Kvve) EE
- E's
LK (1/7 Kv)
Thus, it suffices that
R log(6/2)
- -R" log(//))
T
(Ix
e²
g²
to satisfy& < S { with probability at least 1 8.
Case 2. Suppose we not have an upper bound on AEL(TT). From the proof of Theorem 2, we know that (T)( = 1/E,, (Sreve). In this case we include an estimate Ereser in the bound and use the high-probability deviations
La + e' LR
Smp
Ee+E's Lp, Ereee - E' Ero (Sreset).
We have
5
LR + KvVLLL)
(18)
ETS (Sreset)
(uute' + (v(bp + +
(19)
Sreset
with probability at least whenever
R log(6/3)
T2
2
To guarantee an error at most we require
K LB +
Srevet - o¹ (Lutz' + (v(lp++e))
(20)
Erewet
Assuming e' Ereset, we multiply both sides of (20) by (Erescr - E') and expand:
( La + Kv Kv Lr 16
(Ir
e'
+ Ereser EE'
Ema
La KvLp)) ( + Kv) NEE.
Cancel the common term
(LR + Ky LyLp and group the s' terms:
Ereses
K6
+ + ( + (*KK)KK
Ereeet
Therefore a sufficient condition is the explicit upper bound
E Ereset
<
min
Exeset:
(21)
La + Kv KvLp) + Kv )
Ereset
Together with the concentration requirement on T, the choice (21) ensures at least
on Ç of at most with probability
Finally, the safe policy improvement bound follows from the fact that Ç is greater than < with probability 1 6. Then, due to the SPI bound of Theorem 3, the improvement guaranteed to be even larger when using Ç instead of Ç as error term. This guarantees the improvement when Ç is small enough.
Remark (Episodic assumption). For the sake of presentation, we have considered and proved the bounds for episodic processes (cf. Appendix A), One could extend them to more general cases under the assumption that one has access to a stationary distribution &n. of M. As mentioned in Section 2, the existence distribution is often assumed in continual RL (Sutton and Barto, 2018) and guaranteed unique in the episodic case (Huang, 2020). Then, replacing the difference of returns in Theorem by an expectation (similar to Theorem 2 with Lemma 4) would allow to remove the AEI term and obtain similar results.
stationary
29

---

Published as a conference paper at ICLR 2026
Theorem 4. (Deep SPI for representation learning) Under the same preamble in Thm. 2, let
0 and
== 4. 1. -. Then, with probability at least under Eov we have for all 81, 82 E S that (m,.))-7)
- (52) Kv d((((1). tE.
Proof. First, let us consider bounding the following absolute value difference for every possible state 8 E S, i.e.,
| V (é(s)) To that aim, we consider Markov's inequality:"
( (s) (*((ss) s/2))
(=(((s) 2/22) S Em ((sE S: (
2 Bonbn (*((s) - (o(s))
(Markov's inequality)
3 S2 LR +JKV L.
(by Lemma 4)
6 (1/DDP((m..)) - 7)
Consider any joint distribution A(E,, Cm,), i.e., any joint distribution over S x S whose marginals both match Sa, Then, by the union bound, we have
8(((51.2)) e exx:: (T1)s (- > >/22 (*2)) - (" > f/2))
S A(((81, E x S: (*(ss)
(d(c))) (
< ((())1)22) ES x S: [V(ss)) ('((ss))) 2//2 +((((5))52) ES xxS: (V* - ( > s/2))
(union bound)
= En E 8: [ (+ (/2)) + E. (152 E: V7((2) - (=(((>2)) s/2))
(A has Er, marginal distributions)
< 4. Erp + Lpm 8. (2/D 7)
Therefore, since this holds for any such A. we have with least probability 8 that for all $1, 82 E S, V - V (o(s1)) < 5/2 and V²((2) - V (6(+2)) < s/2. In consequence, with same probability, we have
(=1( v*((s))
(=1 7 (Fd())) + 7766 (52) V'' (*(((s)))
< (+1)s) V (o(s:)) P
+ V² (b(s))))
V7
(triangle inequality)
< (o((s)))) - (0(22)) +
< Kr ((6(s))))) tc.
(by Lemma 3)
"also referred as Chebyshev's inequality (Stein and Shakarchi, 2005).
30

---

Published as a conference paper at ICLR 2026
F
DREAM SPI
Algorithm 2: DreamSPI
Input:
(others) world model and encoder parameters 8, actor/critic parameters 6, imagination horizon H
Init. E (T++))xB
AT rERTxB SE E (
E A
repeat
for t +---0 1 to do
a1 ~ ā(. (())))
Tí, 81+1 env step(s;, d.)
U Update 0 by descending V DeepSPT_ U .0)
Only &. P, and R are updated here
world model world_model
(S.A.P.R)
Set latent states: 81 - s 1st
Perform latent imagination:
for <- H do
a ~ | 82) Pr, 8c41 f worldd
Update by descending V.ppo_te (opo__oss(8, a. 8, A
Perform a standard PPO update of the actor/critic w.r.t. the imagined trajectories
81 - 81+1 until convergence return /
We report in Algorithm 2 the algorithm we used in our experiments evaluate the quality of the world model's predictions. Note that the algorithm is on-policy; we leverage parallelized environments to make sure data coming from the interaction covers sufficiently the state space (Mayor et al., 2025). Empirically, we found most beneficial to use discrete latent spaces, and model the transition function with categorical distributions (32 classes of 32 categories, as in Dreamer; r; Hafner et al., 2021). This observation agrees with the observation made by Hafner et al. (2021) on the benefits of categorical latent spaces in world models.
G ADDITIONAL DETAILS ON THE ILLUSTRATIVE EXAMPLE
s uin
a
+n/ym
a+
a+
8
S
S
win
-
aang
S
SI
S fail
at
(2
a
S
2
urin
Figure 9: Underlying MDP of the grid world of Fig. Actions leading to self-loops are omitted for clarity.
In this section, we expand on the illustrative example introduced in Sect. 6.1. The underlying MDP for the grid world is shown in Fig.9. Formally, the MDP has four actions Gai with dir E {t, -,, +), and 2n + states:
the initial state 87, which transitions 8 whenever a.,, is played;
31

---

Published as a conference paper at ICLR 2026
the hazardous state st, sending the agent si with probability 1-c and s with probability independently of the action played;
the 2n corridor states io and of for i E nn, forming the top and bottom branches; and
the terminal states Swin Swin win and 8fail-
We focus on the value of the initial state (s1) for policies π E II. We highlight three policies of particular interest:
(a) The good policy "good: the policy moves right everywhere except in s,, where it chooses dt:
V (81) = =1(11 e)
n
Σ
xr)]
+
gm
mm1
Σ
/ + (1 +
=9
(1-y + +11-c)m c)n + c²")
Learning Tgood requires that the representation distinguishes the two branches and assigns distinct latent states to 8 and s With n = 5, this corresponds to a return of = 8.01, which is the value reported in Fig. 4 for DeepSPI. This highlight the representation learning capabilities of our algorithm.
(b) The bad policy Toad: the policy moves right everywhere except in s, and s where it chooses a+:
Viaal (51) = 7 [11 E)
n-1
+
7"
Σ
7 + 7"h
-(2 c)m
7m
7"6
n- 2
7 n
1=0
7 7m
n
1 7
0.
Such a policy may arise due to the policy confounding update described in Sect. 3.2, where the representation incorrectly merges Si and s With 5, this corresponds to a return of = -1.99. (c) and the ugly "always right" policy u. this policy deterministically selects a every state. Its value is
72 22--70 -2
v (81)
Σ
2
With n = 5, this corresponds to a return of 2J 4.8. This coincides with the values reported in Fig. 4 for PPO, indicating that PPO alone fails address the confounding policy update in this example.
As mentioned in the main text, we want to highlight the representation learning capabilities of DreamSPI. For this reason, we provide a view of the grid in raw pixels to the agent (cf. Fig. 10). In our experiments, we choose == 0.2. To evaluate PPO and DeepSP in this environment, we use the default parameters from cleanRL (Huang et al., 2022), both for PPO and DeepSPI. I. However, discrete representation with a limited capacity of 256 latent states (precisely, we use 4 categories of 4 classes, with the same latent representation as the used by Hafner et al., 2021). For DeepSPI, 1,7-1, which leads to a neighborhood constant C' < 1/7, as the theory suggests. Fig. 4 reports the median of 10 independent runs/seeds per algorithm, well the interquartile range (25-75%). Note that the values (*r(s) reported in Fig. 4 are computed analytically.
enforce for both algorithms compact, small
restrict the ratio to
H EXPERIMENTS: EVALUATION ON THE ATARI LEARNING ENVIRONMENTS
H.1 IMPLEMENTATION
Our implementation can be found at https::pp/..... http:///////ttttooo////////ttttt
32

---

Published as a conference paper at ICLR 2026
E
Figure 10: Observations of the grid perceived by the agent. The white-most cell corresponds to the agent's 's location in the grid. Each observation has size 84 x 84.
Percent improvement (relative)
109.
-10".
-10
Robotank
gattiezore
Bramiidor
Lisertaine
Stardunnee
Anider
Contipde
Enduro
Spaceinvazte
PrvataEye
desfccoon
tutarhhem.
Freeway
meningDetty
lietoccley
Pntall
Sking
Ternis
Beeing
Obert
Craycllmma
hong
Defender
Allen
Frostbite
dreekkout
Phoenix
te Kruil
Kunafulkestee Astortx
Zanoon
Boa Reaciturner tarshoeengee faboom
Mar catras
Alonts
VideoFinball -inball-
Grevitar-
Uplloonn
Et
Gepnee
Venture
Chopper Cmd.
Hero
Namerhisc fimelhor
Kanganoo
Bowling
Figure 11: Relative improvement DeepSF1 compared to PPO over the full stochastic ALE suite (41/61).
H.2 SETTING
Each presented experiment on the environments from ALE has been conducted across seeds for each algorithm. Each run requires (mean + std) 16.75 + 1.7 min for PPO, 60.24 + 20.16 min for DeepMDP, 62.49 + 20.71 min for DeepSPI, I, and 80.8 + 1.35 min for DreamSPI on an NVIDIA A40. This corresponds to a 3.6x overhead when using DeepSP instead of PPO, which we consider modest cost given the guarantees we obtain. Because our method is on- on-policy and fully parallelizable, the wall-clock time remains well below that of off-policy approaches that do not exploit vectorized environments. For comparison, SAC requires roughly 40 hours for the same number of collected frames on an NVIDIA 1100 (Huang et al., 2022), and Dreamer-ve needs about two days on NVIDIA 1000 (Hafner et al., 2021).
33

---

Published as a conference paper at ICLR 2026
H.3 ADDITIONAL PLOTS
In this section, we present additional figures to highlight statistics and the performance of our algorithm, DeepSPI. Fig. 11 presents the relative improvement of DeepSPI w.r.t. PPO (Fig. 11). We formally compute the relative improvement as
score
scorépaseline
ScOTebaseline
and we use the maximum median human normalized score as the metric to compare in each environment. See next page for comparison of each of the algorithms (average episodic return) per environment, across training steps. We report the median and interquartile range (25-75%) for each environment. Recall that one training step corresponds gathering four Atari frames in the environment.
Transition and reward losses. In the main text, we stated that the transition loss achieved by DeepSPI I is, in general, lower than for DeepMDPs. s. We elaborate here on the statistical significance of this claim. First, as for the human normalized score, we provide in Fig. 12 aggregate metrics for the transition and reward losses. This analysis already reveals that there is no statistically significant difference between the capacity to predict rewards between the two algorithms. We take closer look at the transition loss.
For each environment i, we summarize the transition loss of DeepSPI and DeepMDP by scalars (SP) and EDDP, and form paired differences d, = =SPP- EMDe, The reported mean difference d = 7. Σ, d; = -0. 131881 therefore means that, on average across environments, DeepSPI's 's transition loss is about 0.14 units lower than DeepMDP's s. To quantify uncertainty on this average effect, we use a paired bootstrap: we resample the n environments with replacement, recompute d for each bootstrap sample b 1,..., B, and form the 95% confidence interval as the 2.5th and 97.5th percentiles of (d(b) 6=1- The resulting interval [-0. 0.2226, 0.05907] lies entirely below zero, which under the usual frequentist interpretation provides strong evidence that the true mean gap in transition loss is negative (DeepSPI better) rather than consequence sampling noise.
The paired Wilcoxon signed-rank test (Wilcoxon, 1992) further supports this conclusion without invoking normality of the d,: it ranks the absolute differences |dg]. assigns each rank the sign of d;, and uses the signed rank sum test statistic for the null hypothesis Hu median d;) O. We obtain a very small two-sided p-value = 6.6 x 10" indicating that observing differences this systematically negative would be extremely unlikely if DeepsP I and DeepMDP had the same typical transition loss.
Finally, the aggregates of Fig. 12 provide a complementary robust view: the interquartile mean (IQM) of transition loss is lower for DeepSPT than for DeepMDP, indicating that DeepSPI improves not only the mean performance but also the performance on the central bulk of environments. Taken together, the negative mean difference with 95% confidence interval that excludes zero, the significant Wilcoxon test, and the lower IQM all consistently indicate that DeepSPI I achieves statistically significantly lower transition loss than DeepMDP across Atari.
H.4 HYPERPARAMETERS
As mentioned in the main text, we use the same parameters for PPO as the default cleanRL's parameters. We list the DeepSPI parameters in Table and those of DreamSP I in Table 2. We used the same pa- rameters as DeepSPI for DeepMDPs. For DeepSP I, we performed grid search for the transition density in [IndependentNorrall, Mixturelndependenttoomei search revealed that the mixture of independent normal distributions (i.e., with diagonal covariance matrices) worked best for DeepSPI. We also found that using Lipschitz networks to enforce the Lipschitzness of the latent space (cf. Sect. 5) was faster than enforcing gradient penalty (as used by Gelada et al. 2019) since, in contrast to gradient penalties, enforcing Lipschitz condition through the architecture does not require additional sampling from the latent transition function (which might turn out costly, especially with mixture distributions). Furthermore, constconstaiieed GroupSort architectures ensure Lipschitzness by construction. For the reward and transition coefficients, we performed a grid search in aR, ap E {10-2,5 x 10-3 10-3, 30-3, 5,5 x 10 10-4). We found the
= 5), Categorical(n_cat 32, n_cls 32)}. The grid
best performance at ag 0.01 and ap = 5 x 10 10-4
34

---

Published as a conference paper at ICLR 2026
Adventure
Amidar
Assault
Asterix
Asteroids
0.933
2000
2000
2500
-0.05
1000
-5
097
0.0
5.5
0/2
2.5
Atlantis
BankHeist
BattieZone
BearRRider
Berzerk
Bowling
1200
1000
10
200
1aT
367
347
Boxing
Breakout
Centipede
ChopperCommand
CrazyClimber
Defender
4000
53
12000
100
menn
1000
25000
20000
Iet
3x7
DemonAttack
DoubleDunk
Enduro
Et
FishingDerby
Freeway
200
150
9333
10
100
0.5
1.0
Frostbite
Gopher
Gravitar
HauntedHouse
IceHockey
1000
300
20000
11000
-10-
-12
0.0
65
Lo
0.5
0.
4.5
0.0
10
lamesbond
Kaboom
Kangaroo
Krull
KungFuMaster
LaserGates
3500
1000
10000
2.5
L0
2.5
10r
301
167
167
307
307
MarioBros
MontezuraReeveng
MsPacman
NammTThisGame
Phoenix
Pitfall
2000
15000
0000
-100
X03
2000
5000
167
367
Pong
PrivateEye
Qbert
Riverraid
RoadRunner
Robotank
0000
GCo00
10000
6000
10-
5000
20000
Seaquest
Skiing
Solaris
Spacelnvaders
StarGunner
Superman
800
200
40000
002
-1900
-20000
200
0.0
1.0
65
10
25
15
0.
25
1.
00
0.5
0.0
3.0
Surround
Tennis
TimePilot
Tutankham
UpNDown
Venture
200
40
150
-3.5
52600
20
-100
0.0
2.5
2.5
0.0
VideoPinball
WizardOrWor
YarsRevenge
Zaxxon
0003
2000
10000
1000
100
25
1n7
1n7
180
DeepSP
DeepMDP
PPO
DreamSPI
35

---

Published as a conference paper at ICLR 2026
Hyperparameter
Value
Learning rate
2.5 10-4
Number of envs
128
Number of rollout steps
8
LR annealing
True
Activation functiin
ReLU
Discount factor
0.99
GAE
0.95
Number of minibatches
4
Update epochs
4
Advantage normalization
True
Clipping coefficient
0.1
Entropy coefficient
0.01
Value loss coefficient
0.5
Max gradient norm
0.5
Transition loss coefficient (ap)
5 10
Reward loss coefficient (axg)
0.01
Transition density
Mixture of Normal (diagonal covariance matrix)
Number of distributions
5
Lipschitz networks
True
Table Summary of D hyperparameters.
Hyperparameter
Value
Imagination horizon
actor/critic update epochs
actor/eritic number of minibatches
x 8 32
Discount factor ? Encoder learning rate Actor Iearning rate Critic lcarning rate
0,995
10 10-4
2.75 10
2.75 x 10-5
World model learning rate
10-4
Global LR annealing
False
Weight decay (Adamm)
True; with decay 10
Transition density
Categorical (32 categories of 32 classes, Hafner et al., 2021)
Transition loss coefficient (ap)
0.01
Reward loss coefficient (ag)
0.01
Lipschitz networks
False (unnecessary with discrete random variables)
Other parameters
Same as DeepsPI
Table 2: Summary of DreamSP1 hyperparamcters.
36

---

Published as a conference paper at ICLR 2026
Median
IQM
Mean
Optimality Gap
DeepMDP
DeepSPI
5.20
6.32
Transition better)
Median
IQM
Mean
Optimality
DeepMDP
DeepSPI
0.0036 0.0039 0.0042
0.00488 0.00496 0.00504
0,00828 0.00834 0.00840
0.00028 0.00034 0.00040
Reward loss Dewer is better)
Figure 12: Aggregate median, IQR, Mean, and optimality gap for the reported transition and reward losses over all the Atari environments considered in experiments, with 95% confidence intervals. The confidence intervals are obtained via percentile bootstrapping with stratified resampling, For information, refer Agarwal et al., 2021b.
37


Reviews:
[{"review_id": "zpN3fjHCfJ", "reviewer": "Reviewer_q5uR", "summary": "The paper presents a method for safe policy improvement (SPI) for an online learning setting in combination with world models and representation learning. The authors offer proof that restricting the update of the policy gives guarantees of convergence. They also introduce auxiliary loss functions on the latent space ensure its stability (i.e. prevents the collapse of states in the representation). Afterwards experiments on the Atari benchmark show the performance improvement over PPO/DeepMDPs. The authors also present an investigation into the quality of the resulting world model.", "questions": "- In section 4 when bounding the trust region, the choice of setting 1 < C < 2 seems to be arbitrary, why specifically bound it up to 2?\n- What happens in the evaluations when changing the stochasticity of the system, i.e., how do DeepSPI compare to PPO when p_a and n_{NOOP} are closer to zero, and similarly when they are even higher? This would further motivate the presented results as it shows that the system is not susceptible to small parameter changes.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The problem addressed in the paper for online on-policy safe improvements is very relevant.\n- The proofs and formulations presented in the paper appear correct/rigorous and the analysis done so far forms also a good basis for future works.\n- The evaluations are well done over the full Atari suite, and include meaningful benchmarks, albeit all of them designed by the authors themselves.", "weaknesses": "- A lot of the background in section 2 lacks references for some of the very general statements the authors make about the field. E.g., L125-126 or L132-134. Similarly other mentions of \"in literature\" like L263 don't have references.\n- Numeric results seem to only show minor improvement over PPO (however this is not directly comparable since the authors' approach provides SPI guarantees).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "EYUAJv29ZW", "reviewer": "Reviewer_Y6jS", "summary": "This paper provides a strong theoretical contribution that generalizes safe policy improvement to deep, on-policy RL with learned world models. The integration of SPI principles, Wasserstein-based transition losses, and Lipschitz-continuous latent dynamics is mathematically interesting. The theoretical development is rigorous, connecting mirror descent formulations (Kuba et al., 2022) to practical PPO-style updates, and the proofs (Appendices B–E) seems to be technically comprehensive (although I did not check every detial). The empirical section, though limited to ALE benchmarks, is well-executed and demonstrates the feasibility of theory-grounded improvements.\n\nHowever, some aspects could be clarified or strengthened: (i) practical relevance of the γ < 1/ C assumption, (ii) the sensitivity of DeepSPI to neighborhood size and Lipschitz constants, (iii) scalability to continuous control tasks, and (iv) the limited ablation on auxiliary losses.", "questions": "The choice of neighborhood constant C and coefficients αₚ, αᵣ (Eq. 5) critically affects guarantees and performance, but no ablation or stability analysis is reported.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The work extends classical SPI (Thomas et al., 2015; Laroche et al., 2019) to high-dimensional, online RL. The neighborhood operator $\\mathcal{N}^{C}(\\pi)$ (Eq. 2) defines a trust region via importance ratios, bridging SPI and policy regularization in a provably convergent way (Thm. 1).\n\n2. The paper formalizes how transition/reward prediction losses ($L_{P}$, $L_{R}$) act as local regularizers ensuring Lipschitz-consistent latent dynamics (Thm. 2–4). This explicitly ties representation smoothness to policy safety, a novel link not addressed by DeepMDPs or Dreamer.\n\n3. The identification of PPO as a special case of the proposed mirror-learning-based neighborhood update (Eq. 4 vs Eq. 3) is conceptually valuable. DeepSPI’s modified utility (Eq. 5) yields a principled way to blend safety and empirical efficiency.\n\n4. On the ALE-57 suite, DeepSPI consistently outperforms PPO in 43 of 61 environments (Fig. 3), while reducing transition loss and ensuring smoother learning (Fig. 5). The implementation is reproducible and uses standard open frameworks (CleanRL, EnvPool).", "weaknesses": "1. Evaluations are restricted to Atari; no tests on continuous-control or partially observable tasks (e.g., DMControl, Procgen). It remains unclear whether Lipschitz constraints and local losses scale to these settings.\n\n2. Although DeepMDP losses are included, the study omits other model-based SPI or regularized RL baselines (e.g., DreamerV3, SAC + KL). Statistical significance and variance across seeds are not discussed (e.g., arxiv.org/abs/1904.06979).\n\n3. While claimed “on-policy,” DeepSPI introduces additional transition modeling networks. The paper lacks runtime, sample-efficiency, or wall-clock comparisons versus PPO, given the relatively limited improvement versus the increased algorithmic complexity.\n\n4. This paper does noticeably lack a dedicated “proof roadmap” or overview section, which would guide the reader through the logical structure and dependencies of its multiple theorems (e.g., Theorems 1–5).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "gfOqWTrsst", "reviewer": "Reviewer_9StS", "summary": "The paper proposes DeepSPI, a safe policy improvement algorithm that combines mirror learning updates with a learned world model. The work provides theoretical results that jointly consider the impacts of policy updates and world model updates, which motivates an extension of PPO that incorporates a world model for representation learning. Experiments across the Atari benchmark demonstrate improved performance compared to PPO and DeepMDP.", "questions": "**[Q1]**: Are the theoretical results compatible with model-based RL using imaginary rollouts as in the Dreamer series? The experiments suggest that the use of an on-policy approach without imaginary rollouts cannot achieve performance close to state-of-the-art algorithms like DreamerV3, which limits the practical impact of the work.\n\n**[Q2]**: Are the theoretical results compatible with off-policy RL methods? \n\n**[Q3]**: How does the approach relate to conservative model-based methods in offline RL (e.g., [3])?\n\n**[Q4]**: Please provide additional implementation details, such as network architectures and the computation time of DeepSPI relative to PPO.\n\n\n**References:**\n\n[3] Kidambi et al., “MOReL: Model-Based Offline Reinforcement Learning.” In NeurIPS 2020.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**[S1]**: The paper is clear, well-written, and easy to follow.\n\n**[S2]**: The paper introduces interesting theoretical analysis that combines mirror learning policy updates with world model representation learning. This analysis extends results from DeepMDP [1] to consider the interplay between policy updates and representation learning, and results in bounds with less restrictive assumptions. \n\n**References:**\n\n[1] Gelada et al., “DeepMDP: Learning Continuous Latent Space Models for Representation Learning.” In ICML 2019.", "weaknesses": "**[W1]**: Experimental results do not demonstrate convincing practical benefits of the proposed approach. The experiments show modest improvements for DeepSPI compared to PPO and DeepMDP, and the performance of DeepSPI is significantly worse than other results that have been reported in the literature on the Atari benchmark. DreamerV3 [2], for example, is also an actor-critic algorithm that leverages a learned world model, and reports better results after only 100k steps (Atari100k results in [2]) compared to the performance of DeepSPI after 10M steps.\n\n**[W2]**: It seems like the learned world model is under-utilized in DeepSPI, only being used for representation learning but not for model-based rollouts as in the Dreamer series. The attempt to consider a model-based implementation using imaginary rollouts (DreamSPI) does not perform well, which the authors suggest may be due to its on-policy nature. It is not clear if the performance guarantees can be extended to account for model-based training with imaginary rollouts.\n\n**References:**\n\n[2] Hafner et al., “Mastering Diverse Domains through World Models.” arXiv, 2023. arXiv:2301.04104.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "EOQSrhCTDC", "reviewer": "Reviewer_2xGk", "summary": "This paper presents theoretical bounds on expected return degeneration (or values) as a consequence of representation learning composed with policy updates.\nThe authors theoretically show that making controlled (regularized) updates reduces out-of-trajectory misgeneralization during learning.\nSecondly, they show that including the world model losses as part of the returns (aka an auxiliary reward/ objective) reduces the policy confounding problem.\nFinally, an extension to PPO with auxiliary losses was implemented to validate the authors theory, which showed marginal performance increase on the ALE benchmark.", "questions": "1. I don't understand part of the claim on the support for the policy sequence $(\\pi_n)_{n\\ge 0}$. Why does $\\pi_n, n \\rightarrow \\infty$ guarantee full support?Theorem 1 claims that the sequence converges to an optimal value, and I can understand this yields a (stochastically) optimal policy. However in lines 234-239 it is stated that since $\\pi_0$ has full support, all $\\pi_n$ have full support as well, which again guarantees full support over the full state space. \nBut if $\\pi_n \\rightarrow \\pi^* $ as $n \\rightarrow \\infty$, then $ \\pi^* $ should not have full support. The optimal policy should put zero density/ probability on suboptimal actions. If it doesn't, then $\\pi_n$ should not converge such that $V^{\\pi_n} = V^*$.\n   \n2. Could you shortly reflect on the theory and practical impact of your answer to my Q1? How would you revise your claims/ theorem?\n\n3. Can you add equation numbering to the reward-transition loss of lines 257-259, and then refer to this Eq. again when discussing Theorem 2. It took me a moment to see that $L_R^\\xi, L_P^\\xi$ referred to these losses over the state-occupancy instead of the buffer $\\mathcal{B}$.\n\n4. Could the authors think of (and run) a didactic experiment that validates the newly proposed algorithm. This could perhaps be built on the toy MDPs from figure 1 or figure 2, at the moment, it is not clear experimentally whether the authors proposed method", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "- Strong mathematical clarity, the background is compact yet comprehensive, and I could understand all nomenclature from this point on. The authors are precise without impeding reading flow. \n\n- Clear communication and examples of model learning issues and confounding in section 3.\n\n- Overall excellent discussion of theory. Almost every theoretical result is followed by a simplifying explanation that discusses its impact and how it will be used later.\n\n- Results are linked to the claims proposed in the introduction.\n\n- Simple and intuitive modification of the PPO algorithm in Section 6 that is decently motivated from the preceding theory. \n\n- Humorous Spiderman references, although reducing the section-title fontsize in section 3 to fit on one line, is probably not allowed.", "weaknesses": "- I find the use of the sensitive word \"safe\" quite imprudent and unfitting. While we can endlessly debate the concept of safety, and the authors are probably not to blame (but prior work is), as far as I understand the paper is not dealing with cost constraints or risk aversion. Instead, \"safe\" refers to not \"destroying\" the world model due to policy updates during learning. Considering that there exist actual risk-sensitive RL applications (e.g., in finance or healthcare), I believe the authors should advertise their method in a different way. How about a word play on \"conservative-world-model policy improvement\" or \"unconfounding regularization\"? Which is more accurate in my opinion to the fix presented in Eq.4 \\& Eq.5.\n\n- Although section 3.2 presents a clear example, as far I understand this idea lies extremely close to well known prior work that should be cited. Specifically, the $Q^\\pi$-irrelevance abstractions. See, Li L. et al. (2006). *Towards a Unified Theory of State Abstraction for MDPs.*\n\n- The after-discussion of Theorem 2 should more strongly state the limitation that the result cannot give guarantees for large discount factors $\\gamma \\rightarrow 1$. This is ok, and aligns with intuition. The last 2 sentences of lines 301-308 already sort-of discusses this, but I think it should be even more explicit.\n\n- The experimental section is weak and I'm not able to gather much from it, other than confirming that \"the authors' method does not break PPO\". Performance is not marginally improved or decreased, and varies from environment to environment. So, although the setup is decent, the problem is that we do not learn much from these experiments.\n\t- I am missing a more didactic experiment that compares the authors' method to a baseline method to validate that they \"fix\" the OOT and confounding policy update issues that they claim to solve. Instead the experiments go straight into performance and loss curves. Although useful to know, they are not that meaningful in the papers' context.\n\t- Furthermore, I miss a stronger connection from the theory to e.g., parameter choices for say $\\epsilon$ (or $C$) that provides some approximate guarantee of modulating the permitted errors derived in the previous sections. Could a more TRPO-like algorithm be derived from the authors' theory, that explicitly tries to satisfy constrained updates? Although the authors decently build up to the PPO extension, I find the auxiliary loss regularization to read somewhat like an afterthought that undermines the interesting preceding results.\n\t- The authors claim in paragraph 448-458 that their transition loss is improved, but looking at figure 5 this is not statistically significant compared to the DeepMDP baseline.\n\n- Conclusion line 479-481 overstates the result that DeepSPI improves upon PPO, the performance benefit is quite marginal and more noisy.\n\n *Minor comments:*\n- line 148, the \"baseline policy\" is known in RL as the behavior policy. I don't see a reason to ignore convention here.\n- Appendix page 30, figure is too large, the label overlaps with the page number.\n- Line 286-288, can you put this paragraph into an \"assumption\" environment and clearly state that you're working with this \"from this point on\". This detail is now buried within the surrounding text, and it's not clear whether the assumption belongs to the later results, or to all results.\n\n\n*My overall verdict:* \nThis is an extremely clearly written document with theoretical results that relatively nicely align with intuition. \nThe authors carefully guide the reader through the theory and explain the impact of their result without obfuscating their result with unnecessary detail.\nThat said, I am not a pure theoretical person in the RL field, so I cannot accurately judge the impact or novelty of the proposed contributions over prior work. \nMy main problem with the paper at the moment is the weak experiment section, I want to see a setting where the authors can validate their theoretical result, even if it is simple. I think the PPO extension does not do the rest of the paper justice.\nOverall, I would be leaning towards acceptance if most of my points can be addressed and questions answered.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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:

Decision: Accept (Poster)