Self-supervised Action Gating with Energies (SAGE)

• The paper introduces SAGE, which separates feasibility from value by using a latent consistency energy to rerank trajectories in offline RL.
• It employs a self-supervised JEPA encoder and an action-conditioned predictor to capture local dynamics and identify infeasible action prefixes.
• The approach integrates with existing diffusion planning pipelines with minimal overhead, improving performance across locomotion, navigation, and manipulation benchmarks.

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 (Lu et al., 3 Mar 2026).

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 (Lu et al., 3 Mar 2026).

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 (Lu et al., 3 Mar 2026).

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:

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

This stage is purely self-supervised and uses only offline state sequences; no rewards or rollouts are needed. The teacher $$\hat\tau_t^* \in \arg\max_{i \in [C]} J_\psi(\hat\tau_t^{(i)}), \qquad a_t \leftarrow \hat a_t^{(*)}.$$1 is frozen after training and becomes the representation used by the next stage.

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

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

and an action-conditioned predictor $$\hat\tau_t^* \in \arg\max_{i \in [C]} J_\psi(\hat\tau_t^{(i)}), \qquad a_t \leftarrow \hat a_t^{(*)}.$$3 is trained as a block-causal Transformer over latent-action tokens:

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

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

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

The short-horizon rollout consistency loss is

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

The action-usage hinge uses mismatched action sequences $$\hat\tau_t^* \in \arg\max_{i \in [C]} J_\psi(\hat\tau_t^{(i)}), \qquad a_t \leftarrow \hat a_t^{(*)}.$$7, obtained by permuting actions within the batch, to discourage action-agnostic prediction:

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

The final objective is

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

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 $s_{\mathrm{ctx}}=(s_t,\dots,s_{t+W-1}),$0 steps of each candidate plan, with $s_{\mathrm{ctx}}=(s_t,\dots,s_{t+W-1}),$1. For candidate $s_{\mathrm{ctx}}=(s_t,\dots,s_{t+W-1}),$2, with frozen latent encoder $s_{\mathrm{ctx}}=(s_t,\dots,s_{t+W-1}),$3, the latent consistency energy is

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

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 (Lu et al., 3 Mar 2026).

SAGE does not replace the planner’s score $s_{\mathrm{ctx}}=(s_t,\dots,s_{t+W-1}),$5. Instead, it augments it through a two-stage reranker:

1. Sample $s_{\mathrm{ctx}}=(s_t,\dots,s_{t+W-1}),$6 candidate trajectories from the diffusion planner:

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

1. Compute each candidate’s energy $s_{\mathrm{ctx}}=(s_t,\dots,s_{t+W-1}),$8 on the first $s_{\mathrm{ctx}}=(s_t,\dots,s_{t+W-1}),$9 steps.
2. Keep only the lowest-energy fraction $$s_{\mathrm{tgt}}^{(k)} \triangleq s_{t+W-1+k}, \qquad k \in \mathcal K.$$0 of candidates.
3. Among those survivors, choose the one with the best value minus energy penalty:

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

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 (Lu et al., 3 Mar 2026).

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 $$s_{\mathrm{tgt}}^{(k)} \triangleq s_{t+W-1+k}, \qquad k \in \mathcal K.$$2, that the planner has an existing scoring function $$s_{\mathrm{tgt}}^{(k)} \triangleq s_{t+W-1+k}, \qquad k \in \mathcal K.$$3 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 (Lu et al., 3 Mar 2026).

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 $$s_{\mathrm{tgt}}^{(k)} \triangleq s_{t+W-1+k}, \qquad k \in \mathcal K.$$4, keep rate $$s_{\mathrm{tgt}}^{(k)} \triangleq s_{t+W-1+k}, \qquad k \in \mathcal K.$$5, and penalty weight $$s_{\mathrm{tgt}}^{(k)} \triangleq s_{t+W-1+k}, \qquad k \in \mathcal K.$$6. The reported findings are: too-small $$s_{\mathrm{tgt}}^{(k)} \triangleq s_{t+W-1+k}, \qquad k \in \mathcal K.$$7 gives a weak feasibility signal; too-large $$s_{\mathrm{tgt}}^{(k)} \triangleq s_{t+W-1+k}, \qquad k \in \mathcal K.$$8 causes prediction errors to compound and the energy becomes overly conservative; too-small $$s_{\mathrm{tgt}}^{(k)} \triangleq s_{t+W-1+k}, \qquad k \in \mathcal K.$$9 over-prunes and hurts diversity; too-large $$\tilde s_{\mathrm{ctx}}^{(1)}, \tilde s_{\mathrm{ctx}}^{(2)}.$$0 weakens gating; and too-large $$\tilde s_{\mathrm{ctx}}^{(1)}, \tilde s_{\mathrm{ctx}}^{(2)}.$$1 makes energy dominate so selection becomes too conservative. The main configuration used in experiments is

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

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 (Lu et al., 3 Mar 2026). 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” (Geng et al., 2023) 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.