Diffusion Modulation via Environment Mechanism Modeling for Planning

Hanping Zhang; Yuhong Guo

cs.AIFeb 23, 2026

Diffusion Modulation via Environment Mechanism Modeling for Planning

Authors: Hanping Zhang, Yuhong Guo

Organizations: School of Computer Science, Carleton University, Ottawa, Canada

Abstract

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.

Figures & tables

Gym Tasks	BC	CQL	IQL	DT	TT	MOReL	Diffuser	HDMI	HD-DA	DMEMM (Ours)
Med-Expert HalfCheetah	55.2	91.6	86.7	86.8	95.0	53.3	88.9 $\pm$ 0.3	92.1 $\pm$ 1.4	92.5 $\pm$ 0.3	94.6 $\pm$ 1.2
Med-Expert Hopper	52.5	105.4	91.5	107.6	110.0	108.7	103.3 $\pm$ 1.3	113.5 $\pm$ 0.9	115.3 $\pm$ 1.1	115.9 $\pm$ 1.6
Med-Expert Walker2d	107.5	108.8	109.6	108.1	101.9	95.6	106.9 $\pm$ 0.2	107.9 $\pm$ 1.2	107.1 $\pm$ 0.1	111.6 $\pm$ 1.1
Medium HalfCheetah	42.6	44.0	47.4	42.6	46.9	42.1	42.8 $\pm$ 0.3	48.0 $\pm$ 0.9	46.7 $\pm$ 0.2	49.2 $\pm$ 0.8
Medium Hopper	52.9	58.5	66.3	67.6	61.1	95.4	74.3 $\pm$ 1.4	76.4 $\pm$ 2.6	99.3 $\pm$ 0.3	101.2 $\pm$ 1.4
Medium Walker2d	75.3	72.5	78.3	74.0	79.0	77.8	79.6 $\pm$ 0.6	79.9 $\pm$ 1.8	84.0 $\pm$ 0.6	86.5 $\pm$ 1.5

Table 1: This table presents the scores on D4RL locomotion suites for various comparison methods. Results are averaged over 5 seeds.

Environment	MPPI	IQL	Diffuser	HDMI	HD-DA	DMEMM (Ours)
Maze2D U-Maze	33.2	47.4	113.9 $\pm$ 3.1	120.1 $\pm$ 2.5	128.4 $\pm$ 3.6	132.4 $\pm$ 3.0
Maze2D Medium	10.2	34.9	121.5 $\pm$ 2.7	121.8 $\pm$ 1.6	135.6 $\pm$ 3.0	138.2 $\pm$ 2.2
Maze2D Large	5.1	58.6	123.0 $\pm$ 6.4	128.6 $\pm$ 2.9	155.8 $\pm$ 2.5	153.2 $\pm$ 3.3
Multi2D U-Maze	41.2	24.8	128.9 $\pm$ 1.8	131.3 $\pm$ 1.8	144.1 $\pm$ 1.2	145.6 $\pm$ 2.6
Multi2D Medium	15.4	12.1	127.2 $\pm$ 3.4	131.6 $\pm$ 1.9	140.2 $\pm$ 1.6	140.8 $\pm$ 2.2
Multi2D Large	8.0	13.9	132.1 $\pm$ 5.8	135.4 $\pm$ 2.5	165.5 $\pm$ 0.6	159.6 $\pm$ 3.8

Table 2: This table presents the scores on Maze2D navigation tasks for various comparison methods. Results are averaged over 5 seeds.

Gym Tasks	DMEMM	DMEMM-w/o-weighting	DMEMM-w/o- $\lambda_{tr}$	DMEMM-w/o- $\lambda_{rd}$	DMEMM-w/o-tr-guide
Med-Expert Hopper	115.9 $\pm$ 1.6	115.2 $\pm$ 0.4	114.4 $\pm$ 0.8	115.0 $\pm$ 0.4	114.8 $\pm$ 0.2
Med-Expert Walker2d	111.6 $\pm$ 1.1	110.4 $\pm$ 0.8	108.4 $\pm$ 1.2	110.4 $\pm$ 0.6	109.9 $\pm$ 1.0
Medium Hopper	101.2 $\pm$ 1.4	100.4 $\pm$ 1.2	98.6 $\pm$ 1.8	100.1 $\pm$ 1.1	99.8 $\pm$ 1.6
Medium Walker2d	86.5 $\pm$ 1.5	85.6 $\pm$ 1.2	82.8 $\pm$ 1.4	84.4 $\pm$ 0.9	83.0 $\pm$ 1.8
Med-Replay Hopper	100.6 $\pm$ 0.9	98.8 $\pm$ 1.2	97.0 $\pm$ 0.9	98.2 $\pm$ 0.6	96.2 $\pm$ 1.2
Med-Replay Walker2d	85.8 $\pm$ 2.6	84.6 $\pm$ 2.2	82.2 $\pm$ 1.7	83.7 $\pm$ 2.5	82.6 $\pm$ 3.2

Table 3: This table presents the scores on D4RL locomotion suites for all four ablation variants. Results are averaged over 5 seeds.

Panel of Figure 1 — Figure 1: Hyperparameter sensitivity analysis of the tradeoff parameters for transition-based diffusion modulation loss ( $\lambda_{tr}$ ) and reward-based diffusion modulation loss ( $\lambda_{rd}$ ) on Hopper-Medium-Expert and Walker2D-Medium-Expert environments.

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

Shortcut Trajectory Planning for Efficient Offline Reinforcement Learning
Jul 10, 2026Guanquan Wang, Yoshimasa TsuruokaTrajectory Generation Offline RL
Department of Information and Communication Engineering · The University of Tokyo
Amortizing Trajectory Diffusion with Keyed Drift Fields
Mar 14, 2026Gokul Puthumanaillam, Melkior OrnikTrajectory Generation Diffusion-Based Planning
University of Illinois Urbana-Champaign
Muninn: Your Trajectory Diffusion Model But Faster
May 11, 2026Gokul Puthumanaillam, Hao Jiang, Ruben Hernandez +4Trajectory Generation Diffusion-Based Planning
University of Illinois Urbana-Champaign. · Florida International University. · Providence College.

Jul 10, 2026J/K move · Enter open · S save

Shortcut Trajectory Planning for Efficient Offline Reinforcement Learning

Guanquan Wang, Yoshimasa Tsuruoka

Department of Information and Communication Engineering · The University of Tokyo

Trajectory Generation Offline RL Motion Planning Shortcut Learning Model-Based RL

Open paper AskPDF arXiv

Diffusion-based trajectory planners have shown strong performance in offline reinforcement learning, but their iterative denoising process often incurs high inference cost. Consistency-based planners reduce the number of sampling steps, yet they typically rely on a two-stage teacher--student distillation pipeline that increases training cost and may introduce instability. We propose Shortcut Trajectory Planning (STP), an offline model-based reinforcement learning framework that incorporates shortcut models as efficient trajectory generators. STP trains a conditional shortcut trajectory model in a single stage, supports adjustable one-step and few-step inference through step-size conditioning, and selects candidate plans using a critic augmented with feasibility-aware correction. Across standard D4RL benchmarks, including locomotion, navigation, manipulation, and dexterous control tasks, STP achieves strong performance while simplifying the training pipeline for fast generative planning.

Abstract

Diffusion-Based Planning

Explore similar work

Shortcut Trajectory Planning for Efficient Offline Reinforcement Learning

Amortizing Trajectory Diffusion with Keyed Drift Fields

Muninn: Your Trajectory Diffusion Model But Faster