Primary-Fine Decoupling for Action Generation in Robotic Imitation

Abstract: Multi-modal distribution in robotic manipulation action sequences poses critical challenges for imitation learning. To this end, existing approaches often model the action space as either a discrete set of tokens or a continuous, latent-variable distribution. However, both approaches present trade-offs: some methods discretize actions into tokens and therefore lose fine-grained action variations, while others generate continuous actions in a single stage tend to produce unstable mode transitions. To address these limitations, we propose Primary-Fine Decoupling for Action Generation (PF-DAG), a two-stage framework that decouples coarse action consistency from fine-grained variations. First, we compress action chunks into a small set of discrete modes, enabling a lightweight policy to select consistent coarse modes and avoid mode bouncing. Second, a mode conditioned MeanFlow policy is learned to generate high-fidelity continuous actions. Theoretically, we prove PF-DAG's two-stage design achieves a strictly lower MSE bound than single-stage generative policies. Empirically, PF-DAG outperforms state-of-the-art baselines across 56 tasks from Adroit, DexArt, and MetaWorld benchmarks. It further generalizes to real-world tactile dexterous manipulation tasks. Our work demonstrates that explicit mode-level decoupling enables both robust multi-modal modeling and reactive closed-loop control for robotic manipulation.

• The paper introduces PF-DAG, a two-stage policy that selects a discrete primary behavior before generating continuous, mode-conditioned action residuals with one-step MeanFlow decoding.
• PF-DAG achieves 0.72 aggregate success across 56 simulation tasks versus 0.51 for leading baselines, while real-world tests show gains on contact-rich tasks and lower jerk than DP3.
• A codebook size of 64 performs best, while ablations show that continuous decoding is essential and the framework’s main limitations are mode-classification errors, codebook tuning, and reduced granularity in fast tasks.

PF-DAG (Primary-Fine Decoupling for Action Generation) addresses a persistent failure mode in robotic imitation learning: policies trained on multi-modal demonstration data either collapse modes into an averaged action or bounce erratically between valid modes at test time. The paper proposes a two-stage framework that separates coarse, discrete mode selection from fine-grained continuous action generation, and supports this design with both a theoretical MSE analysis and extensive empirical evaluation (2602.21684).

Motivation and problem setting

The starting observation is that expert demonstrations for manipulation tasks frequently admit several valid behaviors for the same observation — e.g., steering left or right around an obstacle — so the conditional action distribution $p(a \mid o)$ is genuinely multi-modal. The paper situates existing approaches along three failure-prone axes. Behavioral cloning regresses toward the conditional mean, averaging incompatible modes. Discrete-token policies (bin-based tokenizers as in RT-1/RT-2/OpenVLA, frequency-space methods like FAST, and VQ-based tokenizers) represent multiple modes but introduce reconstruction error and temporal discontinuity at fine granularity. Latent-variable generative policies such as ACT and Diffusion Policy can express multiple modes, but independent per-step resampling of the latent produces "mode bouncing": random switches between modes across consecutive timesteps, degrading trajectory smoothness and task accuracy.

The authors' decomposition hypothesis is that many demonstrations factor naturally into a small set of interpretable primary modes ("lift-and-fold", "lift-and-rotate") plus continuous within-mode residuals. PF-DAG operationalizes this by making primary mode selection explicit and consistent, while delegating continuous variation to a conditional generator.

Method

PF-DAG consists of three components operating in a receding-horizon closed-loop scheme: given observation $$\mathbf{o}_t = (\mathbf{p}_t, \mathbf{s}_t, \mathbf{f}_t)$$ — point cloud, proprioception, and optionally 5-fingertip tactile readings ($5 \times 120 \times 3$ taxel force vectors) — the policy predicts a chunk of $T_p$ actions and executes the first $T_a$ before re-planning.

Primary mode learning and selection. A compact MLP-based VQ-VAE compresses ground-truth action chunks into a codebook of $K$ discrete primary modes. A lightweight MLP classifier $\pi_1(m \mid \mathbf{o})$ is trained with cross-entropy against encoder-assigned VQ indices; at inference, greedy argmax selection fixes the mode. Crucially, tokenization is used only for high-level mode choice, not for final action output.

Mode-conditioned MeanFlow decoder. A one-step generative head based on MeanFlow predicts a residual $\Delta\mathbf{a}$ conditioned on $(\mathbf{o}, m)$, added to the VQ-decoder prototype $\hat{\mathbf{a}}^{(m)} = D_\psi(\mathbf{e}_m)$. The average velocity field $$\mathbf{o}_t = (\mathbf{p}_t, \mathbf{s}_t, \mathbf{f}_t)$$0 is trained with the standard MeanFlow identity objective over randomly sampled intervals $$\mathbf{o}_t = (\mathbf{p}_t, \mathbf{s}_t, \mathbf{f}_t)$$1, using a DiT-style transformer backbone. At inference, $$\mathbf{o}_t = (\mathbf{p}_t, \mathbf{s}_t, \mathbf{f}_t)$$2 yields single-function-evaluation generation.

Theoretical analysis

Under unbiasedness assumptions, the paper shows the two-stage predictor's irreducible MSE is $$\mathbf{o}_t = (\mathbf{p}_t, \mathbf{s}_t, \mathbf{f}_t)$$3, which by the law of total variance equals $$\mathbf{o}_t = (\mathbf{p}_t, \mathbf{s}_t, \mathbf{f}_t)$$4. This is strictly smaller than the single-stage bound whenever inter-mode variance is positive — the exact quantity that BC-style averaging is forced to absorb. An appendix sharpens this into an explicit trade-off: PF-DAG exchanges the guaranteed inter-mode error $$\mathbf{o}_t = (\mathbf{p}_t, \mathbf{s}_t, \mathbf{f}_t)$$5 of MSE-optimal averaging for a classification error term $$\mathbf{o}_t = (\mathbf{p}_t, \mathbf{s}_t, \mathbf{f}_t)$$6 from learned mode misselection, arguing that the former is catastrophic for execution while the latter is manageable. The main-text analysis assumes an oracle $$\mathbf{o}_t = (\mathbf{p}_t, \mathbf{s}_t, \mathbf{f}_t)$$7; the appendix version is the more honest account, and it is empirical success rather than theory that carries the claim that $$\mathbf{o}_t = (\mathbf{p}_t, \mathbf{s}_t, \mathbf{f}_t)$$8 is benign.

Empirical results

Across 18 core simulation tasks spanning Adroit (Shadow Hand), DexArt (Allegro Hand), and MetaWorld, with results on 56 tasks total, PF-DAG achieves a 0.72 aggregate success rate versus 0.51 for DP3 and FlowPolicy, 0.30 for Diffusion Policy, and ≤0.15 for IBC and BC-H. Notable per-task gains include Pen (0.65 vs. 0.53), Faucet (0.72 vs. 0.63), and MetaWorld Medium/Hard suites (0.68/0.72 vs. 0.45/0.35). Training uses 10 demos per Adroit/MetaWorld task and 90 for DexArt, with three seeds per configuration.

Real-world experiments cover four tasks on xArm with either a gripper or a 12-DOF tactile XHand. PF-DAG outperforms vanilla BC and DP3 throughout: 0.70 on Pick Cube (vs. 0.60), 0.90 on Place Baymax (vs. 0.85), and 0.70–0.80 on contact-rich Wipe Table and Place Toy Into Bin where vanilla BC scores 0.00. A supplementary jerk metric on Wipe Table reports total end-effector jerk of 0.45 for PF-DAG versus 1.25 for DP3, directly supporting the stability claim.

Ablations

The component ablations are decisive on the side of the MeanFlow decoder: removing it (decoding raw VQ reconstructions) collapses weighted success to 0.01, confirming quantization error alone is destructive at small $$\mathbf{o}_t = (\mathbf{p}_t, \mathbf{s}_t, \mathbf{f}_t)$$9; removing the Primary Mode Policy costs only 0.16 absolute success but still materially hurts performance. The codebook-size sweep identifies $5 \times 120 \times 3$0 as optimal (0.72), with $5 \times 120 \times 3$1 (0.61) underfitting and $5 \times 120 \times 3$2 (0.58) making $5 \times 120 \times 3$3 hard to learn — a real expressivity/learnability tension the authors acknowledge. K-means tokenization performs comparably (0.70), indicating robustness to the specific quantizer. Replacing CFM with MeanFlow improves success (0.72 vs. 0.69), and varying ODE solvers/NFE for CFM yields negligible gains. Two additional findings deserve emphasis: total NFE has surprisingly little effect on simulated-task success within the tested range (which the authors attribute to task tolerance to perturbations, offered explicitly as a hypothesis), and PF-DAG maintains high success at short chunk lengths, approaching reactive closed-loop operation without the mode bouncing that normally penalizes short horizons.

Limitations and open questions

The paper concedes two limitations plainly. First, discretized-chunk mode selection reduces temporal granularity in very high-dynamics, low-latency regimes. Second, the codebook size requires per-task tuning sweeps; the authors propose shared or meta-learned codebooks, end-to-end distillation, and multi-task pretraining as remedies but do not implement them. Additionally, the theoretical strict-improvement result depends on oracle mode selection and unbiasedness assumptions, and the NFE-insensitivity observation is unexplained beyond conjecture. Open questions include whether $5 \times 120 \times 3$4 remains benign as task diversity grows beyond the current benchmarks, and how the framework extends to long-horizon hierarchical control.

Conclusion

PF-DAG demonstrates that explicitly decoupling discrete primary mode selection from mode-conditioned one-step continuous decoding yields consistent gains over diffusion- and flow-based imitation baselines across 56 simulation tasks and tactile dexterous manipulation, backed by a variance-decomposition argument for its MSE advantage. Its practical significance rests on the combination of stability (reduced mode bouncing, lower jerk), sample efficiency, and single-NFE inference, tempered by a codebook-capacity trade-off that remains hand-tuned in the current formulation (2602.21684).