Lyapunov-Guided Self-Alignment: Test-Time Adaptation for Offline Safe Reinforcement Learning

TL;DR

This paper introduces SAS, a transformer-based test-time adaptation method for offline safe reinforcement learning that improves safety and performance without retraining by using Lyapunov-guided imagined trajectories as prompts.

Contribution

SAS is a novel framework enabling test-time safety adaptation in offline RL through self-alignment and Lyapunov-guided imagination without retraining.

Findings

SAS reduces safety violations across benchmarks.

SAS maintains or improves task return.

SAS demonstrates effective test-time safety adaptation.

Abstract

Offline reinforcement learning (RL) agents often fail when deployed, as the gap between training datasets and real environments leads to unsafe behavior. To address this, we present SAS (Self-Alignment for Safety), a transformer-based framework that enables test-time adaptation in offline safe RL without retraining. In SAS, the main mechanism is self-alignment: at test time, the pretrained agent generates several imagined trajectories and selects those satisfying the Lyapunov condition. These feasible segments are then recycled as in-context prompts, allowing the agent to realign its behavior toward safety while avoiding parameter updates. In effect, SAS turns Lyapunov-guided imagination into control-invariant prompts, and its transformer architecture admits a hierarchical RL interpretation where prompting functions as Bayesian inference over latent skills. Across Safety Gymnasium and MuJoCo benchmarks, SAS consistently reduces cost and failure while maintaining or improving return.