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

TL;DR

This paper introduces a novel online Riemannian optimization framework for safety-critical control using world models, enabling better safety and performance in complex, black-box environments.

Contribution

It develops a sample-based Penalized Predictive Control method that leverages Riemannian geometry of feasibility manifolds for improved safety and convergence guarantees.

Findings

Simulations show the approach outperforms marginal and frozen density models.

The safety bound depends on the curvature and score estimation error, improving with richer context.

The method effectively adapts after environment shifts.

Abstract

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 $\xi_t$. 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 $\hat{p}(u \mid \xi_t)$ that endows the action space with a Riemannian geometry guiding the Planner's gradient descent. The barrier curvature $\kappa(\xi_t)$, the minimum curvature of the conditional log-density $-\ln\hat{p}(\cdot\mid\xi_t)$, 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(\xi_t)$, 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.