Keyed Drifting Policies in Offline RL

• Keyed Drifting Policies (KDP) are one-step conditional trajectory generators that replace iterative denoising with an amortized planning approach.
• They employ conditioning-aware drift fields and keyed neighborhoods to align generated trajectories with the initial state, ensuring diverse and non-averaged outputs.
• Empirical results show KDP significantly reduces inference latency while improving performance on offline RL benchmarks and real-time hardware deployments.

Searching arXiv for the specified paper and closely related references. Keyed Drifting Policies (KDP) are a class of one-step conditional trajectory generators for offline reinforcement learning that seek to reproduce diffusion-style trajectory planning behavior without iterative denoising at inference time. In "Amortizing Trajectory Diffusion with Keyed Drift Fields" (Puthumanaillam et al., 14 Mar 2026), KDP is defined as an approach that amortizes iterative trajectory refinement into the training objective, enabling generation of full $H$-step trajectory windows in a single forward pass while preserving conditioning on the current state and the ability to sample diverse candidate plans in a receding-horizon loop.

1. Definition and distinction from diffusion planners

In standard diffusion-based trajectory planning, a planner samples noise $x_T \sim \mathcal N(0, I)$ and then applies $T$ sequential reverse-denoising steps $p_\theta(x_{t-1}\!\mid x_t, c)$ to recover a trajectory window $x_0$ (Puthumanaillam et al., 14 Mar 2026). Because each denoising step requires a full neural-network evaluation, best-of-$K$ sampling incurs $\mathcal O(KT)$ network calls per control cycle.

KDP replaces this with a one-step conditional generator

$$g_\psi\colon z\sim\mathcal N(0,I),\;c\;\longmapsto\;\hat x=g_\psi(z,c)\in\mathbb R^{H\times D},$$

together with a hard clamp $\mathcal C_c(\hat x)$ that forces the first state of the generated window to match the current observation $c=s_0$ exactly (Puthumanaillam et al., 14 Mar 2026). At inference time, KDP draws $x_T \sim \mathcal N(0, I)$0 independent latent samples $x_T \sim \mathcal N(0, I)$1, evaluates $x_T \sim \mathcal N(0, I)$2 once per candidate, and optionally scores the resulting trajectories. The practical consequence is that candidate generation requires only $x_T \sim \mathcal N(0, I)$3 parallel network calls rather than $x_T \sim \mathcal N(0, I)$4 sequential calls.

The central motivation is a failure mode of naïve conditional distribution matching. The paper states that conditional trajectory generation fails when the similarity measure used to align generated trajectories with the dataset is dominated by unconstrained future dimensions. In that regime, generation is attracted toward average trajectories, action diversity collapses, and behavior becomes near-static (Puthumanaillam et al., 14 Mar 2026). KDP is designed specifically to avoid that failure mode.

2. Conditioning-aware drift fields and keyed neighborhoods

The defining technical device in KDP is a conditioning-aware notion of neighborhood. Instead of computing similarity in the full $x_T \sim \mathcal N(0, I)$5 trajectory space, KDP defines a compact key space and uses it only for neighborhood selection, while applying updates in the full trajectory space (Puthumanaillam et al., 14 Mar 2026).

For a minibatch of dataset windows $x_T \sim \mathcal N(0, I)$6, the clamped model output is

$x_T \sim \mathcal N(0, I)$7

The key map is

$x_T \sim \mathcal N(0, I)$8

so under clamping, $x_T \sim \mathcal N(0, I)$9. Similarity is then measured only in key space: $T$0 This construction ensures that attraction pulls a generated sample toward dataset trajectories sharing the same initial key $T$1, rather than averaging across arbitrary futures (Puthumanaillam et al., 14 Mar 2026).

Attraction and repulsion are parameterized with softmax weights at temperature $T$2: $T$3 and

$T$4

The self-negative mask $T$5 is included so that the repulsion term does not collapse. These weights define full-window attraction and repulsion means,

$T$6

and hence an instantaneous drift field in trajectory space,

$T$7

To reduce sensitivity to any single bandwidth, KDP averages drift fields across multiple temperatures $T$8: $T$9 A binary mask $p_\theta(x_{t-1}\!\mid x_t, c)$0 zeros out entries corresponding to clamped coordinates, and a length normalization

$p_\theta(x_{t-1}\!\mid x_t, c)$1

stabilizes drift magnitudes (Puthumanaillam et al., 14 Mar 2026).

A common misunderstanding is that one-step trajectory generation can be obtained by direct distribution matching alone. The KDP formulation rejects that premise: the paper’s claim is that, in high-dimensional trajectory spaces, unconstrained future state and action blocks dominate full-window distances, so the resulting neighborhood structure is misaligned with the conditioning variable and produces trajectory averaging and diversity loss (Puthumanaillam et al., 14 Mar 2026).

3. Amortized refinement objective

KDP is trained with a stop-gradient drifted target, following the drifting models framework cited in the paper, but with the keyed neighborhood metric as the crucial conditioning-aware modification (Puthumanaillam et al., 14 Mar 2026). The amortized refinement target is

$p_\theta(x_{t-1}\!\mid x_t, c)$2

where $p_\theta(x_{t-1}\!\mid x_t, c)$3 denotes stop-gradient.

The training loss regresses the current prediction toward this drifted target using a weighted squared error that upweights actions: $p_\theta(x_{t-1}\!\mid x_t, c)$4 The paper characterizes this as amortizing the mean-shift-style update $p_\theta(x_{t-1}\!\mid x_t, c)$5 into the model parameters $p_\theta(x_{t-1}\!\mid x_t, c)$6 (Puthumanaillam et al., 14 Mar 2026).

The resulting objective has a specific division of labor. The key-space distances determine which examples count as condition-matched neighbors; the attraction and repulsion are computed as means in full trajectory space; and the stop-gradient target converts what would otherwise be an iterative refinement step into a single supervised regression target. This suggests that KDP is not merely a compressed sampler, but an amortized planner whose training objective embeds a structured neighborhood update.

4. Training and one-step control loop

The paper describes one stochastic gradient descent step for KDP as a seven-stage procedure (Puthumanaillam et al., 14 Mar 2026):

1. Sample a minibatch $p_\theta(x_{t-1}\!\mid x_t, c)$7 from the dataset; for each sample set $p_\theta(x_{t-1}\!\mid x_t, c)$8 and draw $p_\theta(x_{t-1}\!\mid x_t, c)$9.
2. Generate clamped predictions $x_0$0.
3. Compute key-space distances $x_0$1 and form softmax weights $x_0$2 with $x_0$3.
4. Compute attraction $x_0$4 and repulsion $x_0$5.
5. Average across $x_0$6 temperatures, apply the clamping mask $x_0$7, and normalize each $x_0$8.
6. Build the target $x_0$9.
7. Evaluate the weighted mean-squared error $K$0 and update $K$1 with SGD.

At inference time, the control loop is correspondingly simple. Given the current state $K$2, KDP draws $K$3 independent noises $K$4, computes $K$5 in parallel, optionally scores each candidate with a learned scorer $K$6, executes the first action block $K$7, and repeats (Puthumanaillam et al., 14 Mar 2026). The first action block is therefore extracted from a complete predicted trajectory window rather than from a stepwise rollout.

The distinction between training and inference is central. Training performs neighborhood-based drift computation within a minibatch, whereas inference uses only the learned generator and optional scorer. A plausible implication is that the computational burden of iterative refinement is shifted almost entirely offline into optimization, which is the paper’s stated sense of amortization.

5. Empirical performance and hardware deployment

The paper reports results across standard reinforcement learning benchmarks and real-time hardware deployments, emphasizing one-step inference, planning latency, and retention of diverse candidate generation (Puthumanaillam et al., 14 Mar 2026).

Setting	Diffuser	KDP
D4RL Locomotion	Avg score $K$8, $K$9, 149 ms	Avg score $\mathcal O(KT)$0, 1 NFE, 1.98 ms
Maze2D (Large)	$\mathcal O(KT)$1 score, 3.85 s/step	$\mathcal O(KT)$2 score, 0.031 s/step
AntMaze (Medium)	$\mathcal O(KT)$3, 6.15 s/step	$\mathcal O(KT)$4, 0.052 s/step
Pen-clone	10.7 @1.63 s	53.4 @0.02 s
Hammer-clone	53.1 @1.53 s	79.6 @0.03 s
Crazyflie navigation	success $\mathcal O(KT)$5, 4.1 Hz / 312 ms	success $\mathcal O(KT)$6, 38 Hz / $\mathcal O(KT)$7 ms
SO-100 manipulation	3 Hz / 410 ms	success $\mathcal O(KT)$8, 17 Hz / $\mathcal O(KT)$9 ms

For D4RL locomotion tasks including HalfCheetah, Hopper, and Walker2D, KDP is reported at average score $$g_\psi\colon z\sim\mathcal N(0,I),\;c\;\longmapsto\;\hat x=g_\psi(z,c)\in\mathbb R^{H\times D},$$0 with 1 NFE and 1.98 ms planning latency, compared with Diffuser at average score $$g_\psi\colon z\sim\mathcal N(0,I),\;c\;\longmapsto\;\hat x=g_\psi(z,c)\in\mathbb R^{H\times D},$$1 using $$g_\psi\colon z\sim\mathcal N(0,I),\;c\;\longmapsto\;\hat x=g_\psi(z,c)\in\mathbb R^{H\times D},$$2 denoising steps and 149 ms planning latency (Puthumanaillam et al., 14 Mar 2026). On long-horizon goal tasks, Maze2D (Large) is reported at $$g_\psi\colon z\sim\mathcal N(0,I),\;c\;\longmapsto\;\hat x=g_\psi(z,c)\in\mathbb R^{H\times D},$$3 score and 0.031 s/step for KDP versus $$g_\psi\colon z\sim\mathcal N(0,I),\;c\;\longmapsto\;\hat x=g_\psi(z,c)\in\mathbb R^{H\times D},$$4 and 3.85 s/step for Diffuser, while AntMaze (Medium) is reported at $$g_\psi\colon z\sim\mathcal N(0,I),\;c\;\longmapsto\;\hat x=g_\psi(z,c)\in\mathbb R^{H\times D},$$5 and 0.052 s/step for KDP versus $$g_\psi\colon z\sim\mathcal N(0,I),\;c\;\longmapsto\;\hat x=g_\psi(z,c)\in\mathbb R^{H\times D},$$6 and 6.15 s/step for Diffuser.

On dexterous manipulation benchmarks in Adroit, the reported numbers are 53.4 @0.02 s for KDP versus 10.7 @1.63 s for Diffuser on Pen-clone, and 79.6 @0.03 s for KDP versus 53.1 @1.53 s for Diffuser on Hammer-clone (Puthumanaillam et al., 14 Mar 2026). On real-world hardware, Crazyflie navigation is reported at success $$g_\psi\colon z\sim\mathcal N(0,I),\;c\;\longmapsto\;\hat x=g_\psi(z,c)\in\mathbb R^{H\times D},$$7 with replanning at 38 Hz and end-to-end $$g_\psi\colon z\sim\mathcal N(0,I),\;c\;\longmapsto\;\hat x=g_\psi(z,c)\in\mathbb R^{H\times D},$$8 ms/step, versus Diffuser at 4.1 Hz and 312 ms; SO-100 manipulation is reported at success $$g_\psi\colon z\sim\mathcal N(0,I),\;c\;\longmapsto\;\hat x=g_\psi(z,c)\in\mathbb R^{H\times D},$$9 with replanning at 17 Hz and end-to-end $\mathcal C_c(\hat x)$0 ms/step, versus Diffuser at 3 Hz and 410 ms.

Across these settings, the paper states that KDP’s one-step inference retains or improves success rate and action diversity while reducing planning latency by one to two orders of magnitude relative to iterative diffusion (Puthumanaillam et al., 14 Mar 2026).

6. Ablations, failure modes, and stated limitations

The ablation studies reported for locomotion and Maze2D are used to isolate the contribution of each design choice (Puthumanaillam et al., 14 Mar 2026). Replacing key distances with full-window distances, labeled No keying, produces catastrophic collapse with score $\mathcal C_c(\hat x)$1 in locomotion. Allowing self-negatives, $\mathcal C_c(\hat x)$2, reduces performance to $\mathcal C_c(\hat x)$3 locomotion. Removing repulsion, labeled Attraction only, yields pure averaging and zero diversity. Removing drift normalization still works but is less stable across tasks. Using a single temperature produces a small degradation on long-horizon tasks.

These ablations are directly aligned with the paper’s diagnosis of conditional generation failure. The key-space metric is not treated as an implementation convenience, but as the mechanism that makes neighborhood construction condition-aware. The repulsion term is not ancillary; according to the reported ablation, omitting it causes diversity to vanish. The normalization and multi-temperature averaging are presented as stabilizing design choices rather than as the core source of performance (Puthumanaillam et al., 14 Mar 2026).

The paper also states a limitation: very sparse, long-horizon goals, exemplified by AntMaze, still strain one-step generators, and richer goal embeddings or hierarchical extensions may be required. That statement bounds the scope of the method. Although KDP improves substantially over iterative diffusion on AntMaze in the reported experiments, the long-horizon sparse-reward regime remains identified as a difficult setting rather than a solved one.

7. Significance within trajectory generative control

KDP occupies a specific position within trajectory-generative planning. It preserves several properties associated with diffusion-based planners—multimodal trajectory synthesis, candidate sampling, and state-conditioned receding-horizon control—while replacing iterative denoising with a one-step generator trained against a drift-field objective (Puthumanaillam et al., 14 Mar 2026). The method’s defining claim is therefore not simply faster sampling, but the preservation of diffusion-like planning behavior under a different training-inference tradeoff.

The paper’s summary identifies three components as decisive: using a conditioning-aware key metric to form meaningful neighborhoods, defining attraction-repulsion drifts in full trajectory space, and amortizing an iterative mean-shift update into a single regression step via stop-gradient targets (Puthumanaillam et al., 14 Mar 2026). This suggests a broader interpretation of KDP as an amortized trajectory planner whose quality depends on the geometry of conditional neighborhoods rather than on the fidelity of a reverse diffusion chain.

Within that interpretation, the principal contribution of KDP is methodological. It makes explicit that conditional trajectory generation can fail because the similarity metric is misaligned with the conditioning variable, and it proposes keyed drift fields as the corrective structure. The empirical results indicate that this structure is compatible with real-time closed-loop control on both benchmark tasks and hardware systems, where inference latency is a primary constraint.