# Safety-Critical Contextual Control via Online Riemannian Optimization with World Models

arXiv: https://arxiv.org/abs/2604.19639 · HTML: https://mlsanity.com/paper/2604.19639/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Latent dynamics models and model-based RL, one of 11 directions in World models.
- Builds on 22 of its 38 references that are in the map.
- First posted on arXiv 21 Apr 2026.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)

## Builds on (most cited first, 22 of 22)

- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Learning to Simulate Complex Physics with Graph Networks](https://mlsanity.com/paper/2002.09405/) (2020)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [Genie: Generative Interactive Environments](https://mlsanity.com/paper/2402.15391/) (2024)
- [Control Barrier Functions: Theory and Applications](https://mlsanity.com/paper/1903.11199/) (2019)
- [End-to-End Safe Reinforcement Learning through Barrier Functions for Safety-Critical Continuous Control Tasks](https://mlsanity.com/paper/1903.08792/) (2019)
- [Control Barrier Function Based Quadratic Programs for Safety Critical Systems](https://mlsanity.com/paper/1609.06408/) (2016)
- [PRECOG: PREdiction Conditioned On Goals in Visual Multi-Agent Settings](https://mlsanity.com/paper/1905.01296/) (2019)
- [Ctrl-World: A Controllable Generative World Model for Robot Manipulation](https://mlsanity.com/paper/2510.10125/) (2025)
- [A predictive safety filter for learning-based control of constrained nonlinear dynamical systems](https://mlsanity.com/paper/1812.05506/) (2018)
- [Sample-efficient Cross-Entropy Method for Real-time Planning](https://mlsanity.com/paper/2008.06389/) (2020)
- [Video Prediction Models as Rewards for Reinforcement Learning](https://mlsanity.com/paper/2305.14343/) (2023)
- [Learning for Safety-Critical Control with Control Barrier Functions](https://mlsanity.com/paper/1912.10099/) (2019)
- [Data-Driven Model Predictive Control with Stability and Robustness Guarantees](https://mlsanity.com/paper/1906.04679/) (2019)
- [Maximum Entropy Reinforcement Learning via Energy-Based Normalizing Flow](https://mlsanity.com/paper/2405.13629/) (2024)
- [Landing with the Score: Riemannian Optimization through Denoising](https://mlsanity.com/paper/2509.23357/) (2025)
- [Learning-Based Predictive Control via Real-Time Aggregate Flexibility](https://mlsanity.com/paper/2012.11261/) (2020)
- [Data-driven MPC with stability guarantees using extended dynamic mode decomposition](https://mlsanity.com/paper/2308.00296/) (2023)
- [Data-Driven Input-Output Control Barrier Functions](https://mlsanity.com/paper/2502.17688/) (2025)
- [Learning State-Action Control Barrier Functions for Model-Free Control under Constraints](https://mlsanity.com/paper/2505.15515/) (2025)
- [Robust adaptive NMPC using ellipsoidal tubes](https://mlsanity.com/paper/2603.05029/) (2026)

## Abstract (arXiv)

Modern world models are becoming too complex to admit explicit dynamical descriptions. We study safety-critical contextual control, where a Planner must optimize a task objective using only feasibility samples from a black-box Simulator, conditioned on a context signal xit. We develop a sample-based Penalized Predictive Control (PPC) framework grounded in online Riemannian optimization, in which the Simulator compresses the feasibility manifold into a score-based density hatp(u mid xit) that endows the action space with a Riemannian geometry guiding the Planner's gradient descent. The barrier curvature kappa(xit), the minimum curvature of the conditional log-density -lnhatp(·midxit), governs both convergence rate and safety margin, replacing the Lipschitz constant of the unknown dynamics. Our main result is a contextual safety bound showing that the distance from the true feasibility manifold is controlled by the score estimation error and a ratio that depends on kappa(xit), both of which improve with richer context. Simulations on a dynamic navigation task confirm that contextual PPC substantially outperforms marginal and frozen density models, with the advantage growing after environment shifts.