Drifting Field Policy: A One-Step Generative Policy via Wasserstein Gradient Flow
About
We propose Drifting Field Policy (DFP), a non-ODE one-step generative policy built on the drifting model paradigm. We frame the policy update as a reverse-KL Wasserstein-2 gradient flow toward a soft target policy, so that each DFP update corresponds to a gradient step in probability space. By construction, this gradient is decomposed into an ascent toward higher action-value regions and a score matching with the anchor policy as a trust region. We further derive a simple, tractable surrogate of the otherwise intractable update loss, akin to behavior cloning on top-K critic-selected actions. We find empirically that this mechanism uniquely benefits the drifting backbone owing to its non-ODE parameterization. With one-step inference, DFP achieves state-of-the-art performance on several manipulation tasks across Robomimic and OGBench, outperforming ODE-based policies.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Robotic Manipulation | OGBench Cube-double-task2 | Success Rate41 | 15 | |
| Square | RoboMimic | -- | 14 | |
| scene-task2 | OGBench | Success Rate93 | 11 | |
| puzzle-3x3-task4 | OGBench | Success Rate3 | 11 | |
| cube-single-task2 | OGBench | Success Rate95 | 11 | |
| puzzle-4x4-task4 | OGBench | Success Rate20 | 11 | |
| task2 | Cube-triple | Offline Success Rate3.2 | 8 | |
| task2 | Cube-quad | Offline Success Rate20 | 8 | |
| task3 | Cube-triple | Offline Success Rate7.6 | 8 | |
| task3 | Cube-quad | Offline Success Rate5.6 | 8 |