SAGE: Bridging the Feasibility Gap in Diffusion Planning via Self-Supervised Energies

TL;DR

SAGE (Self-supervised Action Gating with Energies) is a plug-and-play inference-time module that makes diffusion planners more robust. By separating "what is good" (Value) from "what is possible" (Feasibility), it uses a self-supervised energy signal to filter out plans that look rewarding on paper but are physically impossible to execute.

Background Positioning: SAGE represents a shift toward selector-side optimization in offline RL, demonstrating that we can significantly improve SOTA planners without retraining the core generative models.

The Problem: The "Value-Feasibility" Tension

In offline RL, we learn from a fixed dataset. Diffusion planners like Diffuser and DV work by generating multiple potential futures (trajectories) and picking the one with the highest predicted reward.

However, there is a hidden trap: Value functions tend to extrapolate. A critic might assign a high score to a trajectory that "cheats"—for example, a robot arm teleporting through a wall or a runner jumping impossible distances—because the critic hasn't seen enough "illegal" examples to know better. When the agent tries to execute the first step of such a plan, it fails catastrophically because the plan is locally inconsistent with real-world dynamics.

Methodology: How SAGE Works

SAGE solves this by treating feasibility as a distinct signal, learned through predictive self-supervision.

1. Learning the Latent Space (JEPA)

The authors first train a Joint-Embedding Predictive Architecture (JEPA). Unlike traditional autoencoders that try to reconstruct pixels, JEPA learns to make future states predictable in a latent space. This creates a representation that is highly sensitive to the temporal logic of the environment.

2. The Action-Conditioned Predictor

Once the latent space is frozen, SAGE trains a Transformer-based predictor. Given a current state $z_t$ and an action $a_t$, it predicts $z_{t+1}$. If a candidate plan suggests a transition $(s_t,a_t,s_{t+1})$ that results in a high prediction error (high energy), SAGE flags it as "dynamically inconsistent."

3. Inference-Time Gating

At test time, the process is simple:

1. Sample $C$ candidate plans from the diffusion model.
2. Compute the Consistency Energy for the first $K$ steps of each plan.
3. Filter out the $P$ fraction of highest-energy candidates.
4. Pick the best remaining plan using the original value critic.

Experiments & Results

SAGE was tested across the D4RL suite, covering locomotion, manipulation, and navigation.

• Robustness: In Maze2D, SAGE effectively suppressed "wall-crossing" trajectories that often plague standard diffusion planners.
• Performance: It achieved a +3.75 point lead over the strong DV baseline in Franka Kitchen tasks and a +2.9 point lead in AntMaze.
• Efficiency: The computational cost of this gating is remarkably low, adding only ~6.8% latency to the inference loop.

Deep Insight: Why Why This Works

The core "insight" of SAGE is factorization. By delegating the task of "physical common sense" to a self-supervised energy model, the value critic is freed from having to learn both "what to do" and "what is possible." Since the energy model doesn't care about rewards—only about whether $s_{t+1}$ follows $s_t$ given $a_t$—it can be trained on the entire offline dataset, making it much more reliable than a reward-sensitive critic.

Conclusion & Limitations

SAGE provides a practical path to more reliable offline planning. Its modular nature means it can be added to existing pipelines like Diffusion Policy or Diffuser without complex re-implementation.

Limitations:

• It depends on the quality of the offline dataset to define "feasibility."
• If the prefix length $K$ is set too long, the energy signal might become noisy due to compounding prediction errors.

Overall, SAGE proves that sometimes, the best way to improve a complex generative model isn't to change how it dreams, but to refine how we choose from those dreams.