Amortizing Trajectory Diffusion with Keyed Drift Fields

Puthumanaillam, Gokul; Ornik, Melkior

Amortizing Trajectory Diffusion
with Keyed Drift Fields

Gokul PuthumanaillamMelkior Ornik

University of Illinois Urbana-Champaign

KDP: retrieve condition-matched dataset windows, train a generator using attraction and repulsion, then generate and execute trajectories in one step — **Keyed Drifting Policies.** Condition-matched retrieval and drift-based refinement during training; a full trajectory window in one forward pass at inference.

01

Abstract

Diffusion-based trajectory planners can synthesize rich, multimodal action sequences for offline reinforcement learning, but their iterative denoising incurs substantial inference-time cost, making closed-loop planning slow under tight compute budgets. We study the problem of achieving diffusion-like trajectory planning behavior with one-step inference, while retaining the ability to sample diverse candidate plans and condition on the current state in a receding-horizon control loop. Our key observation is that conditional trajectory generation fails under naïve distribution-matching objectives when the similarity measure used to align generated trajectories with the dataset is dominated by unconstrained future dimensions. In practice, this causes attraction toward average trajectories, collapses action diversity, and yields near-static behavior. Our key insight is that conditional generative planning requires a conditioning-aware notion of neighborhood: trajectory updates should be computed using distances in a compact key space that reflects the condition, while still applying updates in the full trajectory space. Building on this, we introduce Keyed Drifting Policies (KDP), a one-step trajectory generator trained with a drift-field objective that attracts generated trajectories toward condition-matched dataset windows and repels them from nearby generated samples, using a stop-gradient drifted target to amortize iterative refinement into training. At inference, the resulting policy produces a full trajectory window in a single forward pass. Across standard RL benchmarks and real-time hardware deployments, KDP achieves strong performance with one-step inference and substantially lower planning latency than diffusion sampling.

02

Conditioning-aware drift fields

Comparison of full-window neighborhoods, which average trajectories, and key-space neighborhoods, which preserve conditional diversity — Full-window distances obscure the current condition. Key-space distances recover condition-matched neighborhoods.

Conditional generation needs a conditioning-aware neighborhood.

Diffusion planners produce expressive, multimodal trajectories, but each control step requires sequential denoising. Naïve one-step distribution matching can instead pull samples toward average, near-static trajectories.

KDP compares samples in a compact key space aligned with the condition, while applying updates in the full trajectory space. Attraction toward matched dataset windows and repulsion from generated neighbors define the drift field.

A stop-gradient drifted target amortizes this refinement into training. At inference, the generator samples a trajectory window in one forward pass, executes the first action, and replans.

Algorithm 1: Keyed drifting policy training (one SGD step)

Input: Dataset minibatch $\{\ypos_i\}_{i=1}^B$ ; key map $\kkey$ ; generator $g_\psi$ ; temperatures $\{\tau_m\}_{m=1}^M$ ; clamp operator $\mathcal{C}$ ; clamp mask $M$ ; weights $\lambda_a,\lambda_s$ .

$c_i \gets \kkey(\ypos_i)$ , sample $z_i \sim \mathcal{N}(0,I)$ .
$\xgen_i \gets \mathcal{C}_{c_i}\!\big(g_\psi(z_i,c_i)\big)$ .
$\yneg_j \gets \xgen_j$ and mask self-negatives ( $j=i$ ) in repulsion weights.
$V_i \gets 0$ for all $i$ .
for $m=1$ to $M$ do
$\Wpos_{ij} \gets \mathrm{softmax}\!\big(-\dkey(\xgen_i,\ypos_j)/\tau_m\big)$ .
$\Wneg_{ij} \gets \mathrm{softmax}\!\big(-\dkey(\xgen_i,\yneg_j)/\tau_m\big)$ .
$\mupos_i \gets \sum_j \Wpos_{ij}\ypos_j,\quad \muneg_i \gets \sum_j \Wneg_{ij}\yneg_j$ .
$V_i \gets V_i + \frac{1}{M}\big(\mupos_i - \muneg_i\big)$ .
end for
$V_i \gets M \odot V_i$ . // constraint-aware drift
$\widetilde{x}_i \gets \mathcal{C}_{c_i}\!\big(\mathrm{sg}(\xgen_i + V_i)\big)$ .
$\Delta^s_{i,t} \gets \xgen_i[t,0{:}d_s]-\widetilde{x}_i[t,0{:}d_s]$
$\Delta^a_{i,t} \gets \xgen_i[t,d_s{:}D]-\widetilde{x}_i[t,d_s{:}D]$
$\mathcal{L} \gets \sum_{i,t}\Big(\lambda_s\lVert \Delta^s_{i,t}\rVert_2^2 + \lambda_a\lVert \Delta^a_{i,t}\rVert_2^2\Big)$
Update $\psi$ by SGD on $\mathcal{L}$ .

Algorithm 2: Receding-horizon planning with KDP

Input: Current context $c$ ; generator $g_\psi$ ; number of candidates $K$ ; optional scorer $J_\phi$ .

Sample $z_1,\dots,z_K \sim \mathcal{N}(0,I)$ .
Generate candidates: $\hat{x}^{(k)} \gets \mathcal{C}_c(g_\psi(z_k,c))$ for $k=1,\dots,K$ .
if $J_\phi$ is available then
Select $\hat{x}^{*} \gets \arg\max_k J_\phi(\hat{x}^{(k)},c)$ .
else
Select $\hat{x}^{*} \gets \hat{x}^{(1)}$ .
end if
Execute $a \gets \hat{x}^{*}[0,d_s{:}D]$ and repeat at the next control cycle.

03

Quantitative results

KDP highlighted in pink

Scroll within a table to see all columns.

Short-horizon continuous control

Table 1. Short-horizon control on D4RL locomotion. ME: medium-expert; M: medium. Higher scores are better; fewer sequential network evaluations (NFE) are better.
Method	HalfCheetah		Hopper		Walker2D		Avg↑	NFE↓
Method	ME	M	ME	M	ME	M	Avg↑	NFE↓
Offline selection
BC	55.2	42.6	52.5	52.9	107.5	75.3	64.3	–
CQL	91.6	44.0	105.4	58.5	108.8	72.5	80.1	–
IQL	86.7	47.4	91.5	66.3	109.6	78.3	80.0	–
DT	86.8	42.6	107.6	67.6	108.1	74.0	81.1	–
TT	95.0	46.9	110.0	61.1	101.9	79.0	82.3	–
MOPO	63.3	42.3	23.7	28.0	44.6	17.8	36.6	–
MOReL	53.3	42.1	108.7	95.4	95.6	77.8	78.8	–
MBOP	105.9	44.6	55.1	48.8	70.2	41.0	60.9	–
Generative trajectory planners
DQL	96.8	69.1	111.1	90.5	110.1	87.0	91.1	5
CAC	84.3	69.1	100.4	80.7	110.4	83.1	88.0	2
CBC	32.7	31.0	90.6	71.7	110.4	83.1	69.9	2
Diffuser	88.9	42.8	103.3	74.3	106.9	79.6	82.6	20
KDP	92.5	62.1	103.6	90.3	108.5	87.2	90.7	1

Table 2. Planning cost on one NVIDIA A100. D: Diffuser; UR: unranked; R: ranked. PL and E2E latency are in milliseconds per step.
Method	NFE↓	BEF↓	PL	E2E
BC	1	1	1.01	1.22
D-UR(K=1, T=20)	20	20	149.32	149.51
D-R(K=16, T=20)	20	320	172.16	172.34
D-R(K=64, T=20)	20	1280	279.98	280.03
KDP-UR(K=1)	1	1	1.98	2.11
KDP-R(K=16)	1	32	3.97	4.10
KDP-R(K=64)	1	128	4.76	4.96

NFE Sequential network evaluations
BEF Batch-equivalent forwards
PL / E2E Planning / end-to-end latency
K / T Candidate trajectories / denoising steps

Long-horizon planning & objective ablations

Table 3. Long-horizon, goal-conditioned planning. UM: U-Maze; Med.: medium; Lar.: large. Planning latency (PL) is in seconds per step.
Method	Maze2D				AntMaze
Method	UM	Med.	Lar.	PL	UM	Med.	PL
CQL	5.7	5.0	12.5	–	84.0	53.7	–
IQL	47.4	34.9	58.6	–	62.2	70.0	–
Diffuser	113.9	121.5	123.0	3.848	76.0	31.9	6.153
KDP	122.3	130.4	133.0	0.031	81.3	67.6	0.052

Table 4. Training objective ablations. Locom.: locomotion; Act. div.: action diversity.
Ablation	Locom.↑	Maze2D↑	Act. div.↑
Full method (KDP)	90.7	128.2	0.93
No keying	2.91	9.6	0.14
Include self-negatives	78.08	91.3	0.93
Attraction-only (no repulsion)	1.05	3.7	2.2e6
No drift normalization	86.93	131.5	0.93
Single τ	92.37	102.1	0.93

Dexterous manipulation & action chunking

Table 5. Dexterous manipulation on Adroit. T is planning time in seconds.
Task	Diffuser		KDP
Task	Score	T(s)	Score	T(s)
pen-clone	10.7	1.634	53.4	0.020
door-clone	56.7	1.598	61.8	0.024
hammer-clone	53.1	1.532	79.6	0.034
relocate-clone	56.2	1.685	62.8	0.036

Table 6. Action chunking. L is the number of actions executed before replanning.
Task	L	Diff.	KDP
adroit-pen	1	38%	57%
adroit-pen	8	41%	73%
adroit-hammer	1	47%	62%
adroit-hammer	8	41%	63%
adroit-relocate	1	35%	65%
adroit-relocate	8	64%	71%

Real-world closed-loop control

Table 7. Real-world closed-loop results on a Crazyflie and an SO-100 arm. Diffuser: K = 64, T = 20; KDP: K = 64. Planner and E2E p50 latency are in milliseconds.
Domain	Method	Succ. (%)	TTS (s)	Replan (Hz)	Planner p50	E2E p50
Navigation	Diffuser	92	12.5	4.1	242	312
Navigation	KDP	94	12.0	38	8	26
Manipulation	Diffuser	88	35	3	311	410
Manipulation	KDP	90	32	17	9	47

Succ. Success rate
TTS Median time-to-success over successful trials
Replan Achieved replanning frequency
p50 Median latency

04

Experiments

Select a video to play

Locomotion

HalfCheetah

Hopper

Walker2D

Navigation & dexterous manipulation

Maze navigation

Hammer

Pen

Door

Real-world hardware

Navigation · 01Crazyflie

Navigation · 02Crazyflie

ManipulationSO-100 arm

05

Citation

@misc{puthumanaillam_keyed_drift_fields,
  title = {Amortizing Trajectory Diffusion with Keyed Drift Fields},
  author = {Puthumanaillam, Gokul and Ornik, Melkior},
  howpublished = {\url{https://keyed-drifting.github.io/}}
}