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

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

• The paper introduces a novel safety-critical contextual control framework that leverages online Riemannian optimization with black-box, sample-driven feasibility densities.
• It establishes strong theoretical guarantees linking penalty parameters, barrier curvature, and dynamic regret, while demonstrating fast adaptation in nonconvex dynamic tasks.
• The method scales efficiently to complex scenarios, ensuring improved safety with sharper phase transitions and robust performance in robot navigation experiments.

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

Problem Setting and Theoretical Framework

This paper addresses the challenge of guaranteeing safety in a contextual control setting where the system dynamics are encapsulated within a high-fidelity, black-box Simulator rather than an explicit model. Concretely, at each time step, the Planner must select controls restricted to an unknown, potentially non-convex manifold $\mathcal{M}_t$ dictated by the environment and operational safety constraints. The Planner observes only a low-dimensional context vector $\xi_t$, and receives feasibility samples from the Simulator, encapsulated as a conditional density estimate $\hat{p}(u|\xi_t)$.

The core contribution is the formalization of this architecture as a safety-critical contextual control problem with three principal properties:

1. Black-box constraint access: The Planner never observes the dynamics, nor the constraint Jacobians, but instead consumes a density and its score estimated from Simulator samples.
2. Riemannian optimization interface: The action space is endowed with a Riemannian geometry induced by the Fisher information of the learned density, guiding efficient gradient descent.
3. Online adaptation and safety guarantees: The method tracks the time-varying feasibility manifold using online learning, with explicit non-asymptotic guarantees relating safety to score estimation error and a geometric curvature quantity.

Control Design via Penalized Predictive Control

The Planner solves, at each step, a penalized objective of the form: $\min_u \ c(u) - \beta_t \ln \hat{p}_t(u|\xi_t)$ where $c(u)$ is the immediate task cost, $\hat{p}_t$ is the conditional feasibility density estimated from Simulator samples, and $\beta_t$ is a time-varying penalty (stiffness) parameter. The gradient, $$\nabla c(u) - \beta_t \nabla_u \log \hat{p}_t(u|\xi_t)$$, optimizes task objective while ensuring actions lie close to the high-probability regions of the feasibility manifold.

The Riemannian structure arises as the Hessian of the objective combines the cost smoothness and the Fisher information matrix of the density; the penalty weight $\beta_t$ effectively scales the “barrier” curvature $\kappa(\xi_t)$, which in turn governs contraction rates and safety residuals.

Theoretical Results

The authors obtain a hierarchy of structural and quantitative guarantees:

• Strong Convexity and Contraction: On any convex subset of the superlevel set of the density, the penalized predictive control (PPC) objective exhibits strong convexity with modulus proportional to the barrier curvature (Figure 1). Gradient descent contracts geometrically with a rate determined by the Fisher information and penalty parameter.

  Figure 1: Stiffness ablation validates the phase transition in safety rate as predicted by the critical stiffness $\xi_t$0; safety improves markedly once the penalty crosses the threshold determined by curvature.

• Critical Stiffness Threshold: The minimum penalty ensuring the optimizer resides interior to the learned feasibility set (not “just on” the boundary) is derived in closed form, depending only on the maximum cost gradient, barrier curvature, and superlevel-set geometry.
• Dynamic Regret Bound: The authors quantify tracking performance (dynamic regret) in a time-varying setting, revealing that regret scales with cumulative score variation, and crucially, is damped by the square of the local curvature $\xi_t$1; higher curvature implies robustness to manifold drift.

  Figure 2: Dynamic regret in the cost tracks the predicted scaling with manifold path length, confirming the theoretical $\xi_t$2 dependence on score variation.

• Statistical Rate for Score Estimation: Using kernel density estimators on Simulator feasibility samples, they show that the integrated squared error of the estimated score (the input for Riemannian optimization) decays at the nonparametric minimax rate $\xi_t$3 as a function of sample budget and control dimension $\xi_t$4 (Figure 3).

  Figure 3: Safety rate and integrated squared score error improve sharply with increased sample budget, exhibiting faster convergence than the theoretical lower bound.

• Set Approximation and Safety Certificates: The deviation between the estimated and true feasibility sets is controlled by the score error and density level-set regularity. The main safety bound shows that the distance from selected actions to the true feasibility set consists of a set-approximation residual and a cost-to-safety term, with the latter proportional to $\xi_t$5.
• Contextual Gain Quantification: Through a variational-inference Hessian decomposition, the paper establishes a quantifiable advantage for contextual (conditional) estimation versus naive marginalization. In regimes where context strongly modulates feasible actions, the additional curvature due to context (“score covariance”) ensures that contextual Planners enjoy provably tighter safety guarantees.

Empirical Evaluation

A rigorous experimental protocol evaluates the theoretical framework on a series of challenging dynamic robot navigation tasks, manifesting non-convex and switching feasible regions.

• Comparison with Baselines: The proposed online PPC method outperforms strong baselines including offline generative models, sampling-based CEM, learned and oracle CBF-QP, and conservative ellipsoidal approximations in both safety rate and cost (Figure 4, Table 1).

  Figure 4: Across all metrics, online PPC yields the highest safety with lowest normalized cost among black-box methods.

• Sample Efficiency and Scalability: The method demonstrates sharp phase transitions in safety as a function of penalty parameter (Figure 1), and exhibits monotonic improvement in both safety and score estimation as the sample budget increases. Importantly, PPC scales gracefully as the number of obstacles grows, in contrast to CBF-QP and CEM, whose safety degrades rapidly (Figure 5).

  Figure 5: Scalability results indicate graceful degradation for PPC, whereas feasibility for CEM and CBF-QP is increasingly compromised.

• Fast Adaptation to Manifold Shifts: Under abrupt changes to the environment (“reshuffling obstacles”), contextual PPC using the conditional density rapidly adapts, yielding near-oracle safety, while methods relying on marginal or frozen densities fail to recover (Figure 6).

  Figure 6: Trajectories before and after obstacle reshuffle show that PPC-Context recovers safety, while offline and marginal approaches suffer persistent violations.

• Context Kernel and Conditional Advantage: The contextual observation pipeline (Figure 7) and mode-switching experiments (Figure 8, Table 3) demonstrate that the main empirical effect size in safety (over 5% increase in critical mode-switch regimes) validates the predicted gap derived from the contextual curvature increase. The empirical context-safety gain closely matches the theoretical lower bound, with the gap largest where posterior score variance is high.

  Figure 7: Conditional densities assign substantially more mass to the current feasible set than marginals, consistent with large posterior score covariance and the theoretical contextual safety gap.

  

  Figure 8: Contextual PPC achieves a marked reduction in safety violations at mode switches, empirically quantifying the curvature-driven advantage over marginalization.

Practical and Theoretical Implications

The proposed architecture decouples the task of safety enforcement from explicit modeling of dynamics, leveraging black-box or learned world models without sacrificing safety guarantees. All core quantities required for control and safety certificate—barrier curvature, set approximation parameters—are data-driven and computable online from feasibility samples and context signals.

From a control theory perspective, the replacement of traditional Lipschitz-based bounds with curvature of the learned density marks a step toward integrating probabilistic inference and manifold geometry with safety-critical control, bridging gaps between control-as-inference, online optimization, and distributional robustness. The contextual results lay groundwork for quantifying the value of situational awareness and multi-modal observations.

On the practical side, the Riemannian optimization design, bypassing the need for explicit model access or constraint Jacobians, generalizes to arbitrary black-box or foundation-model-based simulators and scales to high-dimensional and non-convex feasible sets. The sample-based nature of all guarantees positions the method as directly compatible with learned world models and generative simulators common in robot learning, autonomous driving, and distributed energy systems.

Conclusion

This work presents a comprehensive framework for safety-critical contextual control via online score-based Riemannian optimization, with all guarantees grounded in learnable, sample-driven quantities. The integration of manifold geometry, online statistical learning, and contextual inference yields both strong theoretical safety bounds and state-of-the-art empirical results in high-variability environments.

Future research directions include integrating learned perception modules for high-dimensional sensory streams, tightening statistical rates for non-Euclidean action spaces, and dynamically tuning context representations to maximize curvature advantages predicted by the mixture Hessian identity. The architecture and analysis set a template for safe, adaptive Planner–Simulator interfaces suitable for next-generation complex world models.