---
title: Keyed Drifting Policies in Offline RL
url: https://www.emergentmind.com/topics/keyed-drifting-policies-kdp
type: topic
---

# Keyed Drifting Policies in Offline RL

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" [2603.14056], 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$ [2603.14056]. 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 [2603.14056]. At inference time, KDP draws $K$ independent latent samples $z_k$, evaluates $g_\psi$ once per candidate, and optionally scores the resulting trajectories. The practical consequence is that candidate generation requires only $K$ parallel network calls rather than $KT$ 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 [2603.14056]. 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 $H \times D$ trajectory space, KDP defines a compact key space and uses it only for neighborhood selection, while applying updates in the full trajectory space [2603.14056].

For a minibatch of dataset windows $\{x_i\}$, the clamped model output is
\[
\hat x_i = \mathcal C_{c_i}\bigl(g_\psi(z_i,c_i)\bigr),
\quad c_i = \mathrm{state}(x_i)[0].
\]
The key map is
\[
k(x)=\mathrm{state}(x)[0]\in\mathbb R^{d_s},
\]
so under clamping, $k(\hat x_i)=c_i$. Similarity is then measured only in key space:
\[
d^+_{ij} = \|\,k(x_i)-k(x_j)\|_2,\qquad
d^-_{ij} = \|\,k(x_i)-k(\hat x_j)\|_2.
\]
This construction ensures that attraction pulls a generated sample toward dataset trajectories sharing the same initial key $c_i$, rather than averaging across arbitrary futures [2603.14056].

Attraction and repulsion are parameterized with softmax weights at temperature $\tau$:
\[
w^+_{ij}
= \frac{\exp\bigl(-d^+_{ij}/\tau\bigr)}{\sum_{k=1}^B\exp\bigl(-d^+_{ik}/\tau\bigr)},
\]
and
\[
w^-_{ij}
= \begin{cases}
0,& j=i,\\
\displaystyle
\frac{\exp\bigl(-d^-_{ij}/\tau\bigr)}
{\sum_{k\neq i}\exp\bigl(-d^-_{ik}/\tau\bigr)},& j\neq i.
\end{cases}
\]
The self-negative mask $w^-_{ii}=0$ is included so that the repulsion term does not collapse. These weights define full-window attraction and repulsion means,
\[
\mu^+_i=\sum_{j=1}^B w^+_{ij}\,x_j,\qquad
\mu^-_i=\sum_{j=1}^B w^-_{ij}\,\hat x_j,
\]
and hence an instantaneous drift field in trajectory space,
\[
V_i^{(\tau)}=\mu^+_i-\mu^-_i.
\]

To reduce sensitivity to any single bandwidth, KDP averages drift fields across multiple temperatures $\{\tau_m\}_{m=1}^M$:
\[
V_i=\frac1M\sum_{m=1}^M V_i^{(\tau_m)}.
\]
A binary mask $M\in\{0,1\}^{H\times D}$ zeros out entries corresponding to clamped coordinates, and a length normalization
\[
V_i\leftarrow \frac{V_i}{\sqrt{\tfrac1{H\,D}\|V_i\|_2^2}+\epsilon}
\]
stabilizes drift magnitudes [2603.14056].

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 [2603.14056].

## 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 [2603.14056]. The amortized refinement target is
\[
\widetilde x_i
= \mathcal C_{c_i}\Bigl(\mathrm{sg}\bigl(\hat x_i + V_i\bigr)\Bigr),
\]
where $\mathrm{sg}$ denotes stop-gradient.

The training loss regresses the current prediction toward this drifted target using a weighted squared error that upweights actions:
\[
\mathcal L(\psi)
= \mathbb E_{i,t}\Bigl[
\lambda_s\|\hat x_i[t,\,0:d_s]-\widetilde x_i[t,\,0:d_s]\|^2
+ \lambda_a\|\hat x_i[t,\,d_s:D]-\widetilde x_i[t,\,d_s:D]\|^2
\Bigr].
\]
The paper characterizes this as amortizing the mean-shift-style update $\hat x_i \mapsto \hat x_i + V_i$ into the model parameters $\psi$ [2603.14056].

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 [2603.14056]:

1. Sample a minibatch $\{x_i\}_{i=1}^B$ from the dataset; for each sample set $c_i=k(x_i)$ and draw $z_i\sim\mathcal N(0,I)$.
2. Generate clamped predictions $\hat x_i=\mathcal C_{c_i}(g_\psi(z_i,c_i))$.
3. Compute key-space distances $d^+_{ij}, d^-_{ij}$ and form softmax weights $w^+_{ij}, w^-_{ij}$ with $w^-_{ii}=0$.
4. Compute attraction $\mu^+_i$ and repulsion $\mu^-_i$.
5. Average across $M$ temperatures, apply the clamping mask $M$, and normalize each $V_i$.
6. Build the target $\widetilde x_i=\mathcal C_{c_i}(\mathrm{sg}(\hat x_i+V_i))$.
7. Evaluate the weighted mean-squared error $\|\hat x_i-\widetilde x_i\|_W^2$ and update $\psi$ with SGD.

At inference time, the control loop is correspondingly simple. Given the current state $c=s_0$, KDP draws $K$ independent noises $z_k\sim\mathcal N(0,I)$, computes $\hat x^{(k)}=\mathcal C_c(g_\psi(z_k,c))$ in parallel, optionally scores each candidate with a learned scorer $J_\phi$, executes the first action block $a=\hat x^{*}[0,d_s{:}D]$, and repeats [2603.14056]. 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 [2603.14056].

| Setting | Diffuser | KDP |
|---|---|---|
| D4RL Locomotion | Avg score $\approx82.6$, $T=20$, 149 ms | Avg score $\approx90.7$, 1 NFE, 1.98 ms |
| Maze2D (Large) | $\sim123$ score, 3.85 s/step | $\sim133$ score, 0.031 s/step |
| AntMaze (Medium) | $\sim31.9$, 6.15 s/step | $\sim67.6$, 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 $\approx94\%$, 4.1 Hz / 312 ms | success $\approx94\%$, 38 Hz / $\approx26$ ms |
| SO-100 manipulation | 3 Hz / 410 ms | success $\approx90\%$, 17 Hz / $\approx47$ ms |

For D4RL locomotion tasks including HalfCheetah, Hopper, and Walker2D, KDP is reported at average score $\approx90.7$ with 1 NFE and 1.98 ms planning latency, compared with Diffuser at average score $\approx82.6$ using $T=20$ denoising steps and 149 ms planning latency [2603.14056]. On long-horizon goal tasks, Maze2D (Large) is reported at $\sim133$ score and 0.031 s/step for KDP versus $\sim123$ and 3.85 s/step for Diffuser, while AntMaze (Medium) is reported at $\sim67.6$ and 0.052 s/step for KDP versus $\sim31.9$ 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 [2603.14056]. On real-world hardware, Crazyflie navigation is reported at success $\approx94\%$ with replanning at 38 Hz and end-to-end $\approx26$ ms/step, versus Diffuser at 4.1 Hz and 312 ms; SO-100 manipulation is reported at success $\approx90\%$ with replanning at 17 Hz and end-to-end $\approx47$ 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 [2603.14056].

## 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 [2603.14056]. Replacing key distances with full-window distances, labeled *No keying*, produces catastrophic collapse with score $\approx3$ in locomotion. Allowing self-negatives, $w^-_{ii}\neq0$, reduces performance to $\sim78$ 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 [2603.14056].

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 [2603.14056]. 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 [2603.14056]. 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.

Source: https://www.emergentmind.com/topics/keyed-drifting-policies-kdp