Improving Diffusion Planners by Self-Supervised Action Gating with Energies

Yuan Lu; Dongqi Han; Yansen Wang; Dongsheng Li

cs.LGMar 3, 2026

Improving Diffusion Planners by Self-Supervised Action Gating with Energies

Authors: Yuan Lu, Dongqi Han, Yansen Wang, Dongsheng Li

Organizations: University College London, London, UK · Microsoft Research.

Abstract

Diffusion planners are a strong approach for offline reinforcement learning, but they can fail when value-guided selection favours trajectories that score well yet are locally inconsistent with the environment dynamics, resulting in brittle execution. We propose Self-supervised Action Gating with Energies (SAGE), an inference-time re-ranking method that penalises dynamically inconsistent plans using a latent consistency signal. SAGE trains a Joint-Embedding Predictive Architecture (JEPA) encoder on offline state sequences and an action-conditioned latent predictor for short horizon transitions. At test time, SAGE assigns each sampled candidate an energy given by its latent prediction error and combines this feasibility score with value estimates to select actions. SAGE can integrate into existing diffusion planning pipelines that can sample trajectories and select actions via value scoring; it requires no environment rollouts and no policy re-training. Across locomotion, navigation, and manipulation benchmarks, SAGE improves the performance and robustness of diffusion planners.

Figures & tables

Figure 1: The SAGE framework. (a) Learn a predictive JEPA state representation from masked state windows with an EMA teacher. (b) Train an action-conditioned latent predictor whose prediction error defines a transition energy. (c) At test time, score diffusion-generated candidate plans with this energy and gate value-based selection toward locally feasible actions

Figure 2: SAGE energy localises feasibility violations. Per-step latent-consistency energy on an offline episode with clean and after corrupted action segment; corruption induces a sharp, local spike (shaded interval). Full D4RL results in Appendix D.3 Figure 10 .

Figure 3: 100 trajectories example sampled from Maze2d, MCSS can selects wall-crossing or out-of-bounds trajectories (left), while SAGE’s prefix energy filtering and soft penalty suppresses these failures without collapsing trajectory diversity (right).

Dataset	Env	Diffuser*	RGG	LoMAP	LDCQ	DV*	SAGE (Ours)
M-E	HC	88.9	90.8	91.1	90.2	92.7	95.4
M-E	Hop	103.3	109.6	110.6	113.3	108.8	111.3
M-E	W2d	106.9	107.8	109.2	109.3	108.6	109.4
M	HC	42.8	44.0	45.4	42.8	50.4	51.6
M	Hop	74.3	82.5	93.7	66.2	80.9	83.9
M	W2d	79.6	81.7	79.9	69.4	82.8	84.8

Table 1: MuJoCo locomotion results for Diffusion Planners. Values are means over 500 evaluation seeds; standard errors are omitted here for readability, we report the detailed

\pm

SE in Appendix D.5 Table 13 . * denotes results reproduced by us, others are reported from the literature,

		IL	Non-Diffusion			Diffusion Policies		Diffusion Planners
Dataset	Env	BC	BCQ	CQL	IQL	DQL ∗	IDQL ∗	Diffuser*	RGG	LoMAP	LDCQ	DV ∗	SAGE (Ours)
M	Kitchen	47.5	8.1	51.0	51.0	55.1	66.5	52.5	—	—	62.3	73.6	74.5
P	Kitchen	33.8	18.9	49.8	46.3	65.5	66.7	55.7	—	—	67.8	90.0	96.6
Average		40.7	13.5	50.4	48.7	60.3	66.6	54.1	—	—	65.1	81.8	85.6
Antmaze-L	D	0.0	2.2	61.2	47.5	70.6	67.9	27.3	—	39.3	57.7	76.0	77.0
Antmaze-L	P	0.0	6.7	53.7	39.6	81.3	63.5	17.3	—	20.7	—	76.4	82.1

Table 2: Normalised performance of various offline-RL methods on the D4RL benchmark. Reported values are means over 500 episode seeds; * denotes results reproduced by us; remaining numbers are reported from the literature, — indicates the task was not reported in the original work. Standard errors can be found in the detailed table provided in Appendix D.5 Table 13 .

Figure 4: Average performance as a function of the prefix window length

K

across four evaluation domains.

Figure 5: Average performance as a function of the keep rate

\mathcal{P}

across four evaluation domains.

Figure 6: Average performance as a function of the penalty strength

\lambda

across four evaluation domains.

Domain	#	DV*	SAGE	$\Delta$ (95% CI)	$p$
Overall	18	95.59	97.69	2.10 [1.42,2.78]	$1.1\!\times\!10^{-9}$
MuJoCo	9	82.86	84.44	1.59 [1.16,2.02]	$4.8\!\times\!10^{-13}$
Kitchen	2	81.80	85.55	3.75 [3.21,4.29]	$1.1\!\times\!10^{-42}$
AntMaze	4	81.62	84.53	2.90 [0.46,5.34]	$0.020$
Maze2D	3	161.63	163.10	1.47 [-0.54,3.47]	$0.152$

Table 3: Statistical significance tests of SAGE vs. DV. We test the task-average improvement

\Delta

within each domain using result of 500 evaluation episode seeds per task. We report the conservative unpaired (

\rho{=}0

) sensitivity.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Rendered snapshots of the benchmark tasks used in this work. The benchmark spans MuJoCo locomotion, Franka Kitchen manipulation, and maze-based navigation (Maze2D and AntMaze).

	HalfCheetah	Hopper	Walker2d
State space $\mathcal{S}$	$\mathbb{R}^{17}$	$\mathbb{R}^{11}$	$\mathbb{R}^{17}$
Action space $\mathcal{A}$	$\mathbb{R}^{6}$	$\mathbb{R}^{3}$	$\mathbb{R}^{6}$
Episode length	$1000$	$1000$	$1000$

Appendix

Table 4: Environment details for MuJoCo locomotion experiments.

	Kitchen
State space $\mathcal{S}$	$\mathbb{R}^{60}$
Action space $\mathcal{A}$	$\mathbb{R}^{9}$
Episode length	$280$

Appendix

Table 5: Environment details for Franka Kitchen experiments.

	Medium-Play	Medium-Diverse	Large-Play	Large-Diverse
State space $\mathcal{S}$	$\mathbb{R}^{29}$	$\mathbb{R}^{29}$	$\mathbb{R}^{29}$	$\mathbb{R}^{29}$
Action space $\mathcal{A}$	$\mathbb{R}^{8}$	$\mathbb{R}^{8}$	$\mathbb{R}^{8}$	$\mathbb{R}^{8}$
Episode length	$1000$	$1000$	$1000$	$1000$

Appendix

Table 6: Environment details for AntMaze experiments.

	UMaze	Medium	Large
State space $\mathcal{S}$	$\mathbb{R}^{4}$	$\mathbb{R}^{4}$	$\mathbb{R}^{4}$
Action space $\mathcal{A}$	$\mathbb{R}^{2}$	$\mathbb{R}^{2}$	$\mathbb{R}^{2}$
Goal space $\mathcal{G}$	$\mathbb{R}^{2}$	$\mathbb{R}^{2}$	$\mathbb{R}^{2}$
Episode length	$300$	$600$	$800$

Appendix

Table 7: Environment details for Maze2D experiments.

Block	Hyperparameter	Value
Architecture
Encoder $e_{\theta}$	MLP widths	$\dim(s)\!\to\!512\!\to\!512\!\to\!d_{z}$ ( $d_{z}=256$ )
	Activation / norm	GELU; LayerNorm
Teacher $e_{\bar{\theta}}$	Update rule	EMA of $e_{\theta}$
	Momentum	cosine schedule $0.99\to 0.9999$
Predictor $g_{\phi}$	Module	Transformer encoder

Appendix

Table 8: Predictive State Representation Encoder architecture and optimisation.

Block	Hyperparameter	Value
Inputs and tokenisation
Inputs	Latents / actions	$z_{t:t+W}$ ( $d_{z}=256$ , whitened); $a_{t:t+W-1}$
Tokenisation	Per-step token	bundle $[z_{t},a_{t}]$ with type embedding + time positional embedding
Model
Backbone	Architecture	block-causal Transformer
	Depth / heads	2 layers; 4 heads

Appendix

Table 9: Action-Conditioned Latent Predictor architecture and optimisation.

Setting	Value used
Guidance / selection regime	MCSS (generate $C$ candidates, rank with DV scorer $J$ )
State–action generation	Separate (diffuse states; inverse dynamics produces actions)
Discount / credit assignment	$\gamma=0.997$
Planner horizon / stride	$H=32$ , stride $=1$
Candidates $C$	$50$
Planner sampler / steps	DDIM, 20 denoising steps

Appendix

Table 10: Base diffusion planner configuration.

Setting	Value used
Inverse dynamics type	Diffusion inverse dynamics (conditional DDPM over $a_{t}$ )
Conditioning	$(s_{t},s_{t+1})$ (normalised states)
Backbone	MLP
Hidden size	256
Solver / diffusion steps	DDPM; 10 diffusion steps (train), 10 steps (sample)
Training steps	200,000 gradient steps

Appendix

Table 11: Inverse dynamics / policy configuration

Domain	State forward (ridge)	State forward (MLP)	Random latent (ridge)	SAGE (JEPA+AC)
MuJoCo	0.883 $\pm$ 0.022	0.927 $\pm$ 0.018	0.825 $\pm$ 0.026	0.980 $\pm$ 0.003
Maze2D	0.991 $\pm$ 0.001	0.983 $\pm$ 0.001	0.818 $\pm$ 0.025	0.993 $\pm$ 0.001
Kitchen	0.834 $\pm$ 0.002	0.646 $\pm$ 0.003	0.542 $\pm$ 0.002	0.987 $\pm$ 0.000
AntMaze	0.825 $\pm$ 0.003	0.805 $\pm$ 0.004	0.702 $\pm$ 0.003	0.940 $\pm$ 0.001

Appendix

Table 12: Feasibility discrimination (AUROC

\uparrow

) for classifying real transitions vs. action-shuffled transitions. Mean

\pm

standard error across environments within each domain.

Figure 8: Feasibility discrimination on D4RL (AUROC

\uparrow

). Bars show mean across environments; error bars denote standard error.

Figure 9: Average computational overhead of SAGE against MCSS

(C=50)

across D4RL datasets

Figure 10: Additional feasibility energy distributions across D4RL tasks. Each panel compares clean transitions against corrupted transitions (e.g., action noise/shuffling and/or state corruption); corrupted variants consistently shift to higher energy, supporting that the SAGE feasibility signal detects locally dynamics-inconsistent samples.

Figure 11: Bootstrap distribution of SAGE

-

DV improvement. Violin plots show the nonparametric bootstrap distribution (10,000 resamples) of the task-average improvement

\Delta_{\mathcal{T}}

(SAGE

-

DV) computed from

N=500

per-episode D4RL normalised scores per task. We report results overall and by domain (MuJoCo, Kitchen, AntMaze, Maze2D). The horizontal line marks

\Delta_{\mathcal{T}}=0

Diffusion-Based Planning

Diffusion-based planning generates diverse action sequences or trajectories by progressively denoising candidates conditioned on goals, scenes, and constraints.

See all 139 papers

Diffusion Modulation via Environment Mechanism Modeling for Planning
Feb 23, 2026Hanping Zhang, Yuhong GuoDiffusion-Based Planning Offline RL
School of Computer Science, Carleton University, Ottawa, Canada
ExecDiffuser: Closing the Sample-to-Execution Gap in Diffusion Planning
Oct 27, 2026Shuo Pei, Bosun Liang, Saihui Zha +3Autonomous Driving Planning Diffusion-Based Planning
Shuo Pei, Saihui Zha, Huachun Tan are with the Beijing Institute of Technology, Guangdong, China. · Shuo Pei, Bosun Liang, Chen Sun, and Yong Wang (Corresponding author) are with the Department of Data and Systems Engineering, The University of Hong Kong, Hong Kong, China.
Energy-based Compositional Diffusion Planning
Jun 19, 2026Tao Sun, Utkarsh Aashu Mishra, Jiaxin Lu +2Trajectory Generation Diffusion-Based Planning
Stanford University, California, USA · Georgia Institute of Technology, Georgia, USA · University of Texas at Austin, Texas, USA

Feb 23, 2026J/K move · Enter open · S save

Diffusion Modulation via Environment Mechanism Modeling for Planning

Hanping Zhang, Yuhong Guo

School of Computer Science, Carleton University, Ottawa, Canada

Diffusion-Based Planning Offline RL Model-Based Planning

Open paper AskPDF arXiv

Diffusion models have shown promising capabilities in trajectory generation for planning in offline reinforcement learning (RL). However, conventional diffusion-based planning methods often fail to account for the fact that generating trajectories in RL requires unique consistency between transitions to ensure coherence in real environments. This oversight can result in considerable discrepancies between the generated trajectories and the underlying mechanisms of a real environment. To address this problem, we propose a novel diffusion-based planning method, termed as Diffusion Modulation via Environment Mechanism Modeling (DMEMM). DMEMM modulates diffusion model training by incorporating key RL environment mechanisms, particularly transition dynamics and reward functions. Experimental results demonstrate that DMEMM achieves state-of-the-art performance for planning with offline reinforcement learning.

Abstract

Diffusion-Based Planning

Explore similar work

Diffusion Modulation via Environment Mechanism Modeling for Planning

ExecDiffuser: Closing the Sample-to-Execution Gap in Diffusion Planning

Energy-based Compositional Diffusion Planning