Reinforcement Learning | Florent Delgrange

It's a matter of timescale: non-linear utility in successor features and multi-objective planning and learning

Tue, 14 Jul 2026 00:00:00 +0000

Deep SPI: Safe Policy Improvement via World Models

Thu, 23 Apr 2026 00:00:00 +0000

CTDE2: Continuous Training Discrete Execution

Wed, 22 Apr 2026 00:00:00 +0000

Deep SPI

Tue, 03 Mar 2026 00:00:00 +0000

Implementation of the techniques presented in our paper Deep SPI: Safe Policy Improvement via World Models.

Code: GitHub repository
Paper: ICLR 2026
Blog: Deep SPI explainer

The Core Problem

When you train a deep RL policy with auxiliary losses to improve the representation (the observation encoder before your value/critic heads), you face a critical timing problem: The representation optimized under your behavioral policy may not be reliable for the next policy. Add auxiliary losses to regularize the latent space, improve the policy, and suddenly the encoder has shifted, invalidating the very representation you were relying on.

Deep SPI solves this by coupling world-model learning with controlled policy updates: you improve the policy step by step, in a neighborhood that keeps it close to regions where the world model is well-calibrated. This way, updates that look good in the model actually translate to improvements in the real environment.

How It Works: The Algorithm

Deep SPI operates on a deep RL agent where the learned representation $\phi$ (the encoder) maps raw observations into a compact latent space. This representation is shared by both the policy and the world model—the policy predicts actions from $\phi(s)$, and the world model predicts rewards and transitions from $\phi(s)$. The core insight is that as the policy improves, this shared representation can shift, breaking the world model’s calibration. Deep SPI solves this by constraining policy updates to stay in a neighborhood where the world model remains reliable.

Learning a Reliable World Model

Learn two predictors $\overline{R}$ (reward) and $\overline{P}$ (transition) on behavioral data, operating on the learned representation $\phi(s)$:

$$ L_R := \mathbb{E}_{(s,a,r)\sim \pi_b}\left|r - \bar R(\phi(s), a)\right|. $$ $$ L_P := \mathbb{E}_{(s,a,s')\sim \pi_b}\left[\mathbb{E}_{\overline{s}^+ \sim \bar P(\cdot \mid \phi(s), a)}\|\phi(s') - \overline{s}^+\|\right]. $$

These measure calibration exactly where you have data, telling you whether the model is reliable for the next policy update.

Enforcing a Policy Neighborhood

Restrict new policies to stay close to the current policy: action probabilities can’t deviate more than a fixed ratio (similar to PPO’s clipping). This keeps updates close to the data regime, ensuring the world model remains trustworthy.

The Core Guarantee

With constrained updates and small prediction losses, improvements in the learned model transfer back to the real environment. The gap between model and reality shrinks as the prediction losses shrink — so if you’re predicting rewards and transitions well, your planned improvements will actually happen in the real world.

Representation Quality

The encoder remains useful for the next policy: states with very different values can’t collapse to the same latent point without breaking the guarantee. This prevents representation degradation and keeps the learned state abstraction useful as you improve the policy.

Implementation: PPO + World Models

The practical question: how do you enforce the neighborhood in a modern deep RL pipeline?

The answer is PPO, because its clipped objective already acts like a trust region. We build upon Clean RL implementations for robust JAX-based training. But there’s a critical subtlety: auxiliary losses (like $L_R$ and $L_P$) update the encoder, so they can indirectly push the composed policy $\pi := \bar{\pi} \circ \phi$ outside the neighborhood even if the policy head looks well-behaved.

The fix: fold the auxiliary terms into the advantage function $A(s, a)$:

$$ U(s, a) := A(s, a) - \alpha_R \cdot L_R - \alpha_P \cdot L_P $$

Then optimize a PPO-like objective. This way, the encoder and policy head move together in a way that respects the neighborhood constraint.

The Data Coverage Problem

There’s a critical practical detail: PPO is on-policy, which means world-model learning can suffer from narrow, correlated coverage if you collect long rollouts from only a few environments. That’s exactly the wrong regime for learning a well-calibrated world model.

The solution: use many vectorized JAX environments with shorter rollouts instead of a few environments with long rollouts. This gives much broader state coverage while staying on-policy—crucial because a local guarantee is only helpful if your local data actually cover enough of the neighborhood to learn the model well.

Implementation Details

Framework: JAX-based, for efficient jitted computation and GPU scaling
Vectorized environments: many parallel JAX-jitted environments with shorter rollouts for better state coverage
Reproducibility: checked-in uv.lock for exact dependency reproducibility
Benchmarks: ALE-57 (57 Atari games), with results matching or exceeding PPO and DeepMDPs
Environments: uses envpool for high-throughput Atari simulation with GPU support
Code structure: modular design separating encoder, policy, world model, and PPO training logic

Why It Matters

Most auxiliary-loss approaches (even VAE-MDPs and related work) derive their guarantees on-policy, meaning they are valid only for the policy that collected the data. Deep SPI extends this to show that with controlled updates, the guarantees survive the policy change. This is especially valuable because:

Monotonic improvement: You can reason about whether each update actually helps, not just whether it looks good in hindsight
Representation stability: The encoder doesn’t degrade as you improve the policy—it stays useful for downstream learning
No offline requirement: Deep SPI works in online settings; you don’t need a pre-collected dataset

Results

On ALE-57, Deep SPI matches or exceeds PPO and DeepMDPs while maintaining explicit monotonic improvement bounds. The algorithm demonstrates that principled methods with guarantees and practical performance need not be in tension.

ALE-57 Results: Deep SPI is competitive across multiple metrics. Higher is better for IQM (interquartile mean), median, and mean returns; lower is better for optimality gap.

Sample efficiency: Deep SPI’s wall-clock learning curves match strong baselines, showing the guarantees don’t impose significant computational overhead.

Learn More

Paper: Deep SPI: Safe Policy Improvement via World Models (ICLR 2026)
Blog post: Deep SPI explainer with visualizations and detailed worked examples
Code: GitHub repository with reproducible setup and Atari experiments

Composing Reinforcement Learning Policies, with Formal Guarantees

Fri, 23 May 2025 00:00:00 +0000

This work received the Best Poster Award at BeNeRL 2025.

Composing RL policies, with formal guarantees

Mon, 05 May 2025 18:05:41 +0200

Implementation of the techniques presented in our paper Composing Reinforcement Learning Policies, with Formal Guarantees. The source code can be found on GitHub.

This project aims at combining reinforcement learning (RL) and reactive synthesis to equip learning agents with reliable control policies in large, complex environments.

While deep RL is very effective for allowing learning agents to solve complex tasks, (1) it requires (extensive) reward engineering to align the user’s intentions with the learned agent’s behaviors, and (2) the learned agent’s policies are not reliable (no guarantee). Reactive synthesis on the other hand produces reliable control policies that are guaranteed to meet specifications provided in intuitive formal language. However, this requires access to an explicit environment’s model, which is rarely possible in complex environments; even when it is, synthesis does not scale to high-stakes scenarios.

This proposed framework tackles those issues through a hierarchical decomposition of the environment into sub-regions/sections that we call “rooms”. Given a graph (e.g, “map”, “skill graph”) that represent this decomposition, we apply RL in each room to get low-level policies satisfying low-level specifications. Then, we apply synthesis to produce a high-level planner that selects which policy to apply in each room. The resulting method provides guarantees on the resulting agent’s controller, allows for a separation of concerns, and mitigates reward engineering.

The approach relies on learning world models and discrete latent spaces, which enables the formal verification of the low-level policies (via model-checking). Reactive synthesis then composes with the low-level RL policies and the world models to produce high-level planners with guarantees.

New RL algorithm and environments

The project also includes WAE-DQN, an RL algorithm learning a discrete and verifiable world model along with its policy, and new “two-level” environments:

A large, parameterizable grid world with moving obstacles
A 8-room A ViZDoom scenario with ennemies randomly spawning on the map at regular interval.

The two environments come with low- and high-level variants. In the low-level variant, the agent is placed in a room of the two-level environment and its goal is to reach the exit safely, by avoiding moving obstacles. In the high-level variant, the goal of the agent is to navigate safely through the rooms composing the environment to reach a target location.

Integrating RL and Planning through Optimal Transport World Models

Tue, 01 Apr 2025 00:00:00 +0000

Activating Formal Verification of Deep Reinforcement Learning Policies by Model Checking Bisimilar Latent Space Models

Tue, 20 Aug 2024 00:00:00 +0000

The Wasserstein Believer: Learning Belief Updates for Partially Observable Environments through Reliable Latent Space Models

Wed, 29 May 2024 00:00:00 +0000

WBU

Fri, 10 May 2024 00:00:00 +0000

ICLR 2024 paper: https://openreview.net/forum?id=KrtGfTGaGe

The source code is available on GitHub.

Wasserstein Belief Updater (WBU) is an RNN free RL algorithm for POMDPs that learns a representation of the history via an approximation of the belief update in a reliable latent space model, providing theoretical guarantees for learning the optimal value.

This work concerns agents learning how to behave, i.e., their control policy, through reinforcement learning (RL). In real-world scenarios, the environment’s state is very often perceived either through noisy sensors, cameras, or more geneally imperfect observations (e.g., visual observation vs. exact coordinates on a map). In that case, the observation is non-Markovian and the environment is partially observable. This usually leads to complications compared to theoretical perfect-observation RL (i.e., with Markovian observation). For optimal decision making, the agent must in that case base its decision either on (a) the full observation-action history, or (b) the distribution over the possible real states of the environment in which the agent could be at each time step. The latter is called the belief of the agent and is a sufficient statistic to optimize the agent’s return.

The easiest method to tackle partial observability is to process the full history through an RNN to obain a compressed hidden state that can be fed to the policy of the learning agent. While appealing, RNNs don’t yield any guarantee that the representation learned is actually useful (a sufficient statistic) to optimize the agent’s return.

With WBU, we rather propose to learn a representation of the belief. Belief learning is difficult in RL because (1) the dynamics of the environments must be known to exactly compute the belief, and (2) it does not scale as it requires to integrate over the full state space (usually intractable). To tackle those challenges, WBU

learns a world model, through Wasserstein auto-encoded MDPs. This model comes with theoretical abstraction quality guarantees. It is learned through discrete latent spaces which eases the computation of the belief through the latent space.
minimizes the discrepancy between the theoretical belief update rule and the latent belief computed. This yields theoretical representation quality guarantees: close points in the representation space of the beliefs are guaranteed to yield close expected returns (Lipschitz continuity). This guarantees to support policy learning.

WAE-PCN: Wasserstein-autoencoded Pareto Conditioned Networks

Mon, 29 May 2023 00:00:00 +0000

Wasserstein Auto-encoded MDPs @ ICLR 2023

Mon, 01 May 2023 10:42:42 +0100

WAE-MDPs

Wed, 01 Mar 2023 00:00:00 +0000

Official implementation of the Wasserstein Auto-encoded MDP (WAE-MDP) framework. WAE-MDPs enable the distilation of (any kind of) reinforcement learning policies into simpler controllers, paired with a discrete, tractable model of the environment (a latent space model). The two are provided with bisimulation guarantees, which allow formally verifying the behaviors of the agent operating under the simplified policy. The source code provided allows for replicating the experiments of the paper Wasserstein Auto-encoded MDPs: Formal Verification of Efficiently Distilled RL Policies with Many-sided Guarantees.

The source code is available on GitHub.

Deep RL generally lacks guarantees but are very effective to solve complex tasks in intricate environments. On the other hand, model-checking techniques allow to verify agent’s policies, but need the access to an environment’s model. Since model-checking usually relies on an exhaustive exploration of the input model, the latter must in general be finite (discrete spaces) and tractable.

WAE-MDPs distill RL policies learned via any RL technique (in particular, through deep neural networks) to learn a discrete, latent representation of the intractable input space. The distillation yields both a simpler policy and a world model that are amenable to model-checking.

One of the main features of WAE-MDPs is that the quality of the abstraction can be formally checked through PAC bounds on bisimulation pseudometrics. In short, the world model is guaranteed to be bisimilarly close to the real environment; it replicates closely the agent-environment interaction behaviors through its latent space. In addition, WAE-MDPs are equipped with representation guarantees: the representation learned ensures that input states that are grouped to the same latent state are bisimilarly close (the agent behaves the same way from those states).

Beyond those PAC-verifiable guarantees, the bisimulation guarantees are also maintained during learning. This is in contrast to VAE-MDPs, which learn a surrogate, variational objective to avoid the intractable computation of the Wasserstein metric, in the fixpoint definition of bisimulation. Instead, WAE-MDPs soundly approximate this computation by learning a discriminator (“max”) that distinguishes between states that are likely to be produced in the learned world model, and those of the real environment. This yields a $2$-player game objective (min/max) where the goal of the minimizer is to learn both a representation and a world model that makes the latent states generated by the encoder and the world model indistinguishable.

By doing so, WAE-MDPs avoid posterior collapse issues, which drastically speeds up learning, and yield distilled policies and world models of better quality.

Wasserstein Auto-encoded MDPs: Formal Verification of Efficiently Distilled RL Policies with Many-sided Guarantees

Sun, 01 Jan 2023 00:00:00 +0000

Formal Verification of Efficiently Distilled RL Policies with Many-sided Guarantees @ BNAIC/BeNeLearn 2022

Mon, 07 Nov 2022 10:42:42 +0100

Distillation of RL Policies with Formal Guarantees via Variational Abstraction of Markov Decision Processes

Tue, 28 Jun 2022 00:00:00 +0000

A Framework for Flexibly Guiding Learning Agents

Wed, 01 Jun 2022 00:00:00 +0000

VAE-MDPs

Fri, 17 Dec 2021 00:00:00 +0000

A TensorFlow 2 implementation of Variational Markov Decision Processes, a framework allowing to (i) distill policies learned through (deep) reinforcement learning and (ii) learn discrete abstractions of continuous environments, the two with bisimulation guarantees.

The source code provided allows replicating the experiments presented in the paper Distillation of RL Policies with Formal Guarantees via Variational Abstraction of Markov Decision Processes.

The source code is available on GitHub.

Safe Reinforcement Learning

Thu, 18 Oct 2018 00:00:00 +0000