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

WAE-PCN: Wasserstein-autoencoded Pareto Conditioned Networks

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

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

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

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