---
title: Self-supervised Action Gating with Energies (SAGE)
url: https://www.emergentmind.com/topics/self-supervised-action-gating-with-energies-sage
type: topic
---

# Self-supervised Action Gating with Energies (SAGE)

Self-supervised Action Gating with Energies (SAGE) is an inference-time add-on for diffusion planners in offline reinforcement learning that addresses a specific failure mode of value-based trajectory selection: a planner may rank a candidate trajectory as high-value even when its early steps are not actually executable from the current state. SAGE separates **value** from **feasibility** by learning a self-supervised, latent-space “consistency energy” from offline trajectories and using that energy to rerank sampled candidates at test time. It requires no environment rollouts and no policy re-training, and it can integrate into existing diffusion planning pipelines that can sample trajectories and select actions via value scoring [2603.02650].

## 1. Problem setting and motivating failure mode

Diffusion planners work by sampling many candidate trajectories and then selecting the best one using a learned scorer \(J\), often a value or return estimate:

\[
\hat\tau_t^* \in \arg\max_{i \in [C]} J_\psi(\hat\tau_t^{(i)}), \qquad a_t \leftarrow \hat a_t^{(*)}.
\]

This is effective, but the scorer only measures long-horizon desirability. It does **not explicitly check whether the candidate’s early prefix is locally consistent with the environment dynamics**. The resulting failure mode is that a candidate can look good under the value model, yet its first few actions or transitions are not physically or dynamically realizable, so execution fails immediately or becomes brittle under replanning [2603.02650].

The formulation is especially relevant in offline RL, where value estimators may prefer out-of-distribution or overly optimistic futures. SAGE’s central idea is that **feasibility should be scored separately from value**, rather than being implicitly absorbed into one critic. This suggests a division of labor between a long-horizon desirability signal and a short-horizon executability signal.

## 2. Self-supervised training pipeline

SAGE is built from two self-supervised components trained only on offline data: a JEPA encoder that learns a predictive latent representation of states from offline state sequences, and an action-conditioned latent predictor that learns to predict the next latent state from the current latent and action [2603.02650].

In Stage I, an offline trajectory provides a state window

\[
s_{\mathrm{ctx}}=(s_t,\dots,s_{t+W-1}),
\]

and future targets

\[
s_{\mathrm{tgt}}^{(k)} \triangleq s_{t+W-1+k}, \qquad k \in \mathcal K.
\]

Two masked versions of the context are created by randomly masking state features and timesteps,

\[
\tilde s_{\mathrm{ctx}}^{(1)}, \tilde s_{\mathrm{ctx}}^{(2)}.
\]

Let \(e_\theta\) denote the online state encoder, \(e_{\bar\theta}\) the EMA teacher encoder, and \(g_\phi\) a predictor from context latents to future latent targets. The latent embeddings are written as

\[
z_{\mathrm{ctx}}^{(i)} = e_\theta(\tilde s_{\mathrm{ctx}}^{(i)}), \quad
\bar z_{\mathrm{tgt}}^{(k)} = e_{\bar\theta}(s_{\mathrm{tgt}}^{(k)}), \quad
\hat z_{\mathrm{tgt}}^{(i,k)} = g_\phi(z_{\mathrm{ctx}}^{(i)}, k).
\]

The core JEPA loss aligns predicted target latents to teacher target latents:

\[
\mathcal L_{\mathrm{JEPA}} = \big\| g_\phi(e_\theta(s_{\mathrm{ctx}})) - \mathrm{sg}(e_{\bar\theta}(s_{\mathrm{tgt}})) \big\|_2^2.
\]

In the appendix, the full objective includes VICReg regularization to avoid collapse:

\[
\mathcal L_{\mathrm{JEPA}} = \mathcal L_{\mathrm{sim}} + \lambda_{\mathrm{var}}\mathcal L_{\mathrm{var}} + \lambda_{\mathrm{cov}}\mathcal L_{\mathrm{cov}}.
\]

This stage is purely self-supervised and uses only offline state sequences; no rewards or rollouts are needed. The teacher \(e_{\bar\theta}\) is frozen after training and becomes the representation used by the next stage.

In Stage II, the frozen JEPA encoder maps states to latents,

\[
z_{t:t+W} = e_{\bar\theta}(s_{t:t+W}),
\]

and an action-conditioned predictor \(f_\eta\) is trained as a block-causal Transformer over latent-action tokens:

\[
\hat z_{t+1} = f_\eta(z_t, a_t).
\]

The predictor is trained with three losses. The teacher-forced one-step prediction loss is

\[
\mathcal L_{\mathrm{tf}} = \sum_{j=0}^{W-1}\big\| \hat z_{t+1+j} - z_{t+1+j}\big\|_1 .
\]

The short-horizon rollout consistency loss is

\[
\mathcal L_{\mathrm{ro}} = \big\| \hat z_{t+H_{\mathrm{ro}}} - z_{t+H_{\mathrm{ro}}}\big\|_1 .
\]

The action-usage hinge uses mismatched action sequences \(a'\), obtained by permuting actions within the batch, to discourage action-agnostic prediction:

\[
E_{\mathrm{neg}}=\sum_j\|\hat z'_{t+1+j}-z_{t+1+j}\|_1,
\qquad
\mathcal L_{\mathrm{neg}} = [m - E_{\mathrm{neg}}]_+ .
\]

The final objective is

\[
\mathcal L_{\mathrm{AC}} = \mathcal L_{\mathrm{tf}} + \lambda_{\mathrm{ro}}\mathcal L_{\mathrm{ro}} + \lambda_{\mathrm{neg}}\mathcal L_{\mathrm{neg}}.
\]

The paper’s interpretation is explicit: this predictor is trained to be sensitive to the action sequence, so mismatched actions cause larger latent prediction errors. That makes the error a meaningful feasibility signal.

## 3. Latent consistency energy and inference-time action gating

At inference, SAGE evaluates only the first \(K\) steps of each candidate plan, with \(K \ll H\). For candidate \(i\), with frozen latent encoder \(z=e_{\bar\theta}(s)\), the **latent consistency energy** is

\[
E(\hat\tau_t^{(i)}) \;=\; \frac{1}{K}\sum_{k=0}^{K-1} \Big\| f_\eta\!\big(z_{t+k}^{(i)}, a_{t+k}^{(i)}\big) - z_{t+k+1}^{(i)} \Big\|_1 .
\]

Low energy means the candidate’s prefix is predictable under the learned offline dynamics, so it is locally feasible. High energy means the prefix is dynamically inconsistent, likely infeasible, or off-support. The paper explicitly treats this as an **energy-based feasibility measure**, not a value estimate [2603.02650].

SAGE does not replace the planner’s score \(J\). Instead, it augments it through a **two-stage reranker**:

1. Sample \(C\) candidate trajectories from the diffusion planner:
   \[
   \hat\tau_t^{(i)} \sim p_\theta(\tau \mid s_t), \qquad i=1,\dots,C.
   \]
2. Compute each candidate’s energy \(E(\hat\tau_t^{(i)})\) on the first \(K\) steps.
3. Keep only the lowest-energy fraction \(\mathcal P\) of candidates.
4. Among those survivors, choose the one with the best value minus energy penalty:
   \[
   i^* \in \arg\max_{i \in \mathcal I_t} \Big( J(\hat\tau_t^{(i)}) - \lambda\,E(\hat\tau_t^{(i)}) \Big), \qquad a_t \leftarrow \hat a_t^{(i^*)}.
   \]

This design preserves the value planner’s long-horizon objective while correcting its tendency to select implausible prefixes. A plausible implication is that SAGE is most useful when the base generator already produces a diverse candidate set but the selector is prone to optimistic ranking errors.

## 4. Integration into diffusion-planning pipelines

SAGE is explicitly designed to be **planner-agnostic** and **selector-side only**. It works with any diffusion planner that can sample candidate trajectories, it only changes inference-time reranking, it requires **no environment interaction**, it requires **no policy retraining**, and it does not modify the diffusion generator or the critic [2603.02650].

The paper uses a DV-style generate-and-rank planner as the base system because it is exactly the regime where SAGE is most useful: many sampled candidates, ranked by a critic or value score. If the generator already produces a diverse set of candidate trajectories, SAGE can simply re-score them using feasibility. If the base planner outputs only states, actions can be recovered using an inverse-dynamics model.

The method assumes that the planner can sample a pool of candidate trajectories \(\{\hat\tau_t^{(i)}\}_{i=1}^C\), that the planner has an existing scoring function \(J(\tau)\) or value estimate, that short prefixes of candidates are available for feasibility checking, and that offline data is sufficient to learn a latent dynamics-consistency model. It does **not** assume environment rollouts, online calibration, or labels for infeasible trajectories. The “negative” signal for the action predictor comes from in-batch action permutation, not from interacting with the environment.

## 5. Empirical results, robustness, and diagnostics

The evaluation covers D4RL benchmarks spanning **locomotion** (HalfCheetah, Hopper, Walker2d), **navigation** (AntMaze and Maze2D), and **manipulation** (Kitchen). Across these settings, SAGE improves the strongest diffusion-planner baseline DV* and improves the performance and robustness of diffusion planners [2603.02650].

| Domain | Benchmark family | Average score |
|---|---|---|
| Locomotion | MuJoCo locomotion | DV*: **82.9**; SAGE: **84.4** |
| Navigation | AntMaze | DV*: **81.6**; SAGE: **84.5** |
| Navigation | Maze2D | DV*: **161.6**; SAGE: **163.1** |
| Manipulation | Kitchen | DV*: **81.8**; SAGE: **85.6** |

On MuJoCo locomotion, the gain is described as modest but consistent. On sparse-reward navigation, the gains are larger, with SAGE improving all four AntMaze tasks. On Maze2D, the improvement is smaller because performance is near ceiling on some tasks, but SAGE still gives consistent gains, especially on the harder Large maze. On Kitchen, the improvement is one of the clearest reported effects.

Several analyses support the interpretation of latent prediction error as a feasibility signal. In the corruption test, short action windows in offline trajectories are corrupted while the rest is left intact. SAGE’s energy produces a sharp local spike exactly on the corrupted interval, and stays low elsewhere. This shows the energy is localized and sensitive to action-conditioned inconsistency.

| Setting | Reported AUROC |
|---|---|
| MuJoCo | **0.98** |
| Kitchen | **0.98** |
| Maze2D | **0.99** |
| AntMaze | **0.94** |

The action-conditioned predictor is also compared against state-space ridge dynamics, state-space MLP dynamics, and a random-latent ridge model. SAGE’s JEPA+AC setup has the best discrimination, with AUROC around **0.98 overall**, outperforming the alternatives. This supports the claim that the JEPA representation and action-conditioned prediction are both important.

The hyperparameter ablations study prefix length \(K\), keep rate \(\mathcal P\), and penalty weight \(\lambda\). The reported findings are: too-small \(K\) gives a weak feasibility signal; too-large \(K\) causes prediction errors to compound and the energy becomes overly conservative; too-small \(\mathcal P\) over-prunes and hurts diversity; too-large \(\mathcal P\) weakens gating; and too-large \(\lambda\) makes energy dominate so selection becomes too conservative. The main configuration used in experiments is

\[
K=10,\quad \mathcal P=0.8,\quad \lambda=0.1.
\]

The method adds only a small inference cost, about **6.8% wall-clock overhead** relative to the MCSS baseline.

## 6. Interpretation, scope, and acronym disambiguation

SAGE is best understood as a **feasibility-aware reranking layer** for diffusion planners. Its novelty is not in generating trajectories, but in learning a **self-supervised latent consistency energy** that captures whether a candidate’s early prefix is dynamically plausible under the offline data distribution [2603.02650]. A common misconception is to treat the energy as another value estimator; the paper instead defines it as an energy-based feasibility measure and combines it with, rather than substitutes it for, value.

Another possible source of confusion is acronym overlap. The 2023 paper titled “SAGE: Bridging Semantic and Actionable Parts for GEneralizable Manipulation of Articulated Objects” [2312.01307] is **not** the same method. That work does **not** present an action-gating system, nor an energy-based policy, nor a self-supervised learning framework in the sense implied by “Self-supervised Action Gating with Energies.” Its SAGE stands for a language-guided articulated-object manipulation framework that bridges semantic parts and Generalizable Actionable Parts. The overlap is therefore only partial and accidental.

Within offline RL and diffusion planning, the specific contribution of Self-supervised Action Gating with Energies is the separation of feasibility from value at inference time: JEPA learns a predictive latent space from offline state sequences, the action-conditioned predictor learns local latent dynamics, latent prediction error becomes a feasibility energy, and that energy is combined with value during test-time candidate reranking. This suggests a narrow but well-defined role for the method: correcting the selection of high-value yet unrealizable trajectories without modifying the base diffusion planner.

Source: https://www.emergentmind.com/topics/self-supervised-action-gating-with-energies-sage