English

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

Machine Learning 2026-04-30 v1 Artificial Intelligence

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.

Keywords

Cite

@article{arxiv.2604.26516,
  title  = {Lyapunov-Guided Self-Alignment: Test-Time Adaptation for Offline Safe Reinforcement Learning},
  author = {Seungyub Han and Hyungjin Kim and Jungwoo Lee},
  journal= {arXiv preprint arXiv:2604.26516},
  year   = {2026}
}

Comments

Accepted at AISTATS 2026. First two authors contributed equally. Project page: https://seungyubhan.github.io/sas/. Code: https://github.com/seungyubhan/sas

R2 v1 2026-07-01T12:40:58.919Z