🤖 machine learning

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

This paper introduces SAS, a transformer-based framework that enables test-time adaptation for offline safe reinforcement learning by generating and selecting Lyapunov-compliant imagined trajectories as in-context prompts to realign agent behavior toward safety without retraining.

Original authors: Seungyub Han, Hyungjin Kim, Jungwoo Lee

Published 2026-04-30

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Seungyub Han, Hyungjin Kim, Jungwoo Lee

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you hire a highly skilled robot driver who learned to drive by watching thousands of hours of expert driving videos. This robot is great at following the road, but it has a problem: it was trained in a perfect, sunny simulation. When you finally put it on a real road with unexpected rain, potholes, or a sudden detour, the robot panics. It tries to apply its old rules to a new situation and might drive straight into a wall or a ditch.

This is the core problem the paper addresses: Offline Reinforcement Learning (RL). It's like training a robot on a massive dataset of past experiences without letting it practice in the real world. When the real world changes slightly, the robot becomes unsafe.

The authors propose a solution called SAS (Self-Alignment for Safety). Think of SAS not as a new teacher, but as a safety coach that talks to the robot right before it starts driving.

Here is how SAS works, broken down into simple concepts:

1. The "Imagination" Phase

Before the robot takes a single real step, SAS asks it to imagine several different ways it could drive the next few seconds.

The Metaphor: Imagine you are about to cross a busy street. Before you step out, you quickly imagine three scenarios: "If I walk left, I might hit a car," "If I walk right, I might trip," and "If I walk straight, I'm safe."
The Tech: The robot uses its internal "world model" (a mental map of how the world works) to generate these imaginary paths, or "rollouts."

2. The "Lyapunov" Safety Check

How does the robot know which imagined path is safe? It uses a mathematical tool called a Lyapunov function.

The Metaphor: Think of the robot's safety as a ball rolling down a hill. A "Lyapunov" check is like a sensor that ensures the ball is always rolling down toward a safe valley (the goal) and never rolling up toward a cliff (danger).
The Tech: The paper defines a "safety score" based on how often the robot has seen similar situations in its training data. If an imagined path leads to a place the robot has never seen before (a "low-density" area), the safety score drops, and the path is flagged as dangerous. The robot checks: "Does this path keep the safety score going down (or staying safe)?"

3. The "Self-Alignment" (The Magic Trick)

This is the most unique part. Usually, to fix a robot, you have to retrain it from scratch, which takes a lot of time and data. SAS does something smarter: It uses the robot's own imagination to correct itself.

The Metaphor: Imagine the robot is about to drive. Instead of retraining it, SAS says: "Hey, look at these three imaginary paths you just made. Two of them hit a wall. One of them is safe. Let's pretend that safe path is a hint or a prompt."
The robot then takes that safe imaginary path and feeds it back into its own brain as a "demonstration." It's like the robot saying, "Okay, I see the safe way now. I will follow this specific example."
The Result: The robot aligns its behavior to be safe without changing its underlying code or weights. It's a "self-correction" using a prompt, similar to how you might ask a smart AI, "Here is a safe example, now do the same," without needing to retrain the AI.

4. The "Hierarchical" Brain

The paper explains that the robot's brain (a Transformer model) works like a two-level manager.

The Metaphor: There is a Boss (high-level policy) who decides the general strategy (e.g., "Go to the goal"), and a Worker (low-level policy) who actually moves the wheels.
The Tech: SAS acts as the Boss. By feeding the safe "imagined" path as a prompt, SAS is essentially whispering a new instruction to the Boss: "For this specific situation, the strategy should be 'follow this safe path'." This allows the robot to adapt instantly to new hazards.

What Did They Find?

The authors tested this on various robot simulations (like moving a car, a point, or a drone through obstacles).

The Outcome: Robots using SAS made significantly fewer mistakes and crashed less often than robots that just used their original training.
The Bonus: They didn't sacrifice performance. The robots were still fast and reached their goals, but they did so much more safely.
The Key Takeaway: SAS allows a robot trained on old data to adapt to new, dangerous situations instantly by using its own "imagination" to find a safe path and then following that path as a rule.

Summary

SAS is like a safety net made of the robot's own thoughts. Instead of forcing the robot to learn new rules (which is slow and hard), SAS asks the robot to imagine the future, pick the safest version of that future, and then use that safe version as a guide for what to do right now. It turns "what if" into "what to do," keeping the robot safe without needing a single second of retraining.

1. Problem Statement

Offline Reinforcement Learning (RL) agents, trained on static datasets (e.g., D4RL, DSRL), often fail when deployed in real-world environments due to distribution shifts. Even if an agent performs well on training data, slight differences in dynamics or state transitions can lead to unsafe behaviors (e.g., collisions, exceeding cost thresholds).

Existing solutions face significant limitations:

Belief-based adaptation: Requires retraining from scratch to model uncertainty.
Constrained/Conservative RL: Often sacrifices performance (return) or requires fine-tuning with additional cost signals.
Test-time adaptation: Most current methods cannot adapt a pretrained offline agent to safety constraints at deployment without parameter updates or human supervision.

The core question addressed is: Can an offline RL agent be self-aligned for safety at test time using only its pretrained distribution, without retraining?

2. Methodology: SAS (Self-Alignment for Safety)

The authors propose SAS, a transformer-based framework that enables test-time adaptation through Lyapunov-guided self-alignment. The method operates without updating model parameters, relying instead on in-context learning and Bayesian inference.

Core Components

Lyapunov Density Model (LDM) for Offline RL:
- The authors reformulate Lyapunov stability using occupancy measures derived from offline data.
- They define an energy function $E(s, a) = -\log \hat{\rho}_\pi(s, a)$ , where $\hat{\rho}_\pi$ is the model-based occupancy measure estimated from imagined trajectories.
- A Lyapunov surrogate $G_{SAS}(s, a)$ is constructed to ensure that the agent stays within a control-invariant set ( $R_{SAS}^G$ ). This set guarantees that trajectories do not enter low-density (unsafe/out-of-distribution) regions.
- $G_{SAS}$ is defined as the minimum energy encountered along any trajectory, ensuring the policy avoids collapsing into unsafe, low-probability regions.
Probabilistic Inference View:
- The paper interprets the Transformer-based RL agent as a Hierarchical RL system where the high-level policy selects latent skills ( $z_t$ ) and the low-level policy executes actions.
- In-context learning is framed as Bayesian inference over these latent skills. By providing a prompt (a sequence of state-action pairs), the model infers the most likely high-level policy consistent with that prompt.
- Safety conditions are encoded as observable indicators ( $U_t$ $U_{t}$ and $V_t$ $V_{t}$ ):
  - $U_t$ : Indicates if the current state-action pair is within the safe density region ( $G_{SAS} > 0$ ).
  - $V_t$ : Indicates if the Lyapunov function is non-increasing ( $G_{SAS}(s_t, a_t) \geq G_{SAS}(s_{t+1}, a_{t+1})$ ).
The SAS Algorithm (Algorithm 1):
The process involves two loops of "imagination" and selection:
- Step 1 (Condition $U_t$ ): From a fixed initial state, the transformer imagines $N$ trajectories. It selects the trajectory segment that minimizes the maximum energy (worst-case risk) to form an initial prompt.
- Step 2 (Condition $V_t$ ): Using the initial prompt, the transformer imagines $M$ new trajectories. It selects the trajectory that maximizes the cumulative Lyapunov descent (sum of $V_t$ indicators).
- Step 3 (Deployment): The selected safe segment is used as an in-context prompt ( $p_{1:K}$ ) for the transformer. The agent then conditions its actions on this prompt and the current state, effectively "steering" the pretrained policy toward safe behavior without gradient updates.

3. Key Contributions

New Formulation: A novel definition of Lyapunov stability for offline RL based on occupancy measures, allowing safety constraints to be enforced via data density rather than explicit cost functions.
Theoretical Framework: A Hierarchical RL interpretation of Transformer-based RL via Bayesian inference, where prompting corresponds to inferring latent high-level skills. This provides a theoretical basis for why in-context learning works for safety.
Test-Time Adaptation Method (SAS): A practical algorithm that generates Lyapunov-guided prompts to realign agent behavior. It guarantees safety by restricting the agent to control-invariant sets derived from offline data.
Probabilistic Safety Guarantee: The authors provide a formal bound (Theorem 4) showing that the probability of the agent leaving the safe set decreases exponentially with the number of imagined rollouts ( $N$ ) and re-imagination trials ( $M$ ).

4. Experimental Results

The method was evaluated on Safety Gymnasium, MuJoCo, and Bullet-Safety-Gym benchmarks.

Performance vs. Safety: SAS consistently reduced cost and failure rates (up to 2x improvement) while maintaining or improving return (reward) compared to baselines like Decision Transformer (DT), Constrained Decision Transformer (CDT), and various safe RL baselines (BC-Safe, COptiDICE, etc.).
Generalization: SAS improved safety even for backbones (DT) that were not trained with explicit cost signals, demonstrating that the density-based Lyapunov guidance acts as an implicit safety proxy.
Ablation Studies:
- Lyapunov Descent ( $V_t$ ): Removing the descent condition significantly increased costs, proving its necessity for long-horizon stability.
- Prompt Selection: Random prompting or max-energy prompting performed worse than the Lyapunov-guided selection, confirming the efficacy of the specific selection criteria.
- Model Robustness: The method showed robustness across different prompt lengths, numbers of imagined trajectories, and backbone model sizes (from DT to GPT-2).

5. Significance and Impact

Bridging LLM and RL: The paper successfully adapts the concept of self-alignment (popular in Large Language Models) to the domain of Safe RL, enabling agents to "reason" about safety using their own internal world models.
Deployment Viability: By eliminating the need for retraining or fine-tuning, SAS offers a scalable solution for deploying offline RL agents in dynamic, uncertain environments where safety is critical.
Theoretical Rigor: It moves beyond heuristic safety checks by providing a probabilistic guarantee that the agent remains within a control-invariant set, grounded in Lyapunov stability theory and Bayesian inference.
Efficiency: While it requires inference-time computation for imagination (rollouts), it avoids the computational cost of full retraining, making it suitable for scenarios where data is static but the environment is not.

In summary, SAS represents a significant step forward in making offline RL agents robust to distribution shifts by leveraging Lyapunov-guided imagination to self-correct behavior at test time through in-context learning.