Deep SPI: Safe Policy Improvement via World Models

TL;DR

The paper tackles safe policy improvement in online reinforcement learning with high-dimensional observations by coupling representation learning and world-model learning with controlled policy updates. It introduces a neighborhood operator $\mathcal{N}^{C}$ defined via the importance-ratio bounds $D^{\inf}_{IR}$ and $D^{\sup}_{IR}$ to guarantee monotonic improvements and convergence, and derives a deep SPI theorem linking local world-model losses to policy performance with a bound $\zeta$. Building on this, the authors propose DeepSPI, an on-policy algorithm that minimizes $L_{R}^{\mathcal{B}}$ and $L_{P}^{\mathcal{B}}$ while updating the policy in a clipped PPO-style objective using the utility $U(\cdot)$ that penalizes deviations outside the neighborhood. Empirically, DeepSPI achieves competitive or superior performance to PPO and DeepMDPs on ALE, while providing representation-quality guarantees and addressing issues of out-of-trajectory predictions and confounding policy updates. The work bridges theory and practice by delivering a principled, safe, and scalable approach to deep reinforcement learning in complex environments with learned representations.

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.

Figures (7)

• Figure 1: In $\mathcal{M}$, continuously playing $a_1$ in states from $\mathcal{S}_1$ eventually leads the agent to the region $\mathcal{S}_2$, and playing $a_3$ in $\mathcal{S}_2$ eventually leads the agent to $\mathcal{S}_4$ where a reward of $1$ is incurred at each time step, whatever the action played. Playing $a_2$ in $\mathcal{S}_1$ leads the agent to the region $\mathcal{S}_3$, where all actions incur negative rewards. Here, $\phi\mathopen{}\mathclose{\left(s\right)} = \bar{s}_i$ for any $s \in \mathcal{S}_i$ and $i = \left\{ 1, 2, 4 \right\}$. For $s \in \mathcal{S}_3$, we have either $\phi\mathopen{}\mathclose{\left(s\right)} = \bar{s}_3$ or $\phi\mathopen{}\mathclose{\left(s\right)} = \bar{s}'_3$.
• Figure 2: MDP where the probability of transitioning from $s_1$ to $s_2$ is $1 - \epsilon$, for $0 < \epsilon < 1/4$.
• Figure 3: Relative improvement of DeepSPI compared to PPO over the full stochastic ALE suite (41/61).
• Figure 4: Human normalized score over the stochastic, standard 57 envs. from ALE. Plots per environment available in Appendix \ref{['appendix:full-plots']}.
• Figure 5: Median transition and reward losses during training, aggregated across all the ALE. For the sake of visualization, we cut $L_{P}^{}$ lower values from the plot.

Theorems & Definitions (23)

• Theorem 1
• Theorem 2
• Theorem 3
• Theorem 4
• Remark 1: Occupancy measure
• Definition 1: Neighborhood operator
• Lemma 1
• proof
• Lemma 2
• proof