cs.ROFeb 25, 2026

Primary-Fine Decoupling for Action Generation in Robotic Imitation

Authors: Xiaohan Lei, Min Wang, Wengang Zhou, Xingyu Lu, Houqiang Li

Organizations: MoE Key Laboratory of Brain-inspired Intelligent Perception and Cognition, University of Science and Technology of China · Institute of Artificial Intelligence, Hefei Comprehensive National Science Center

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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.

Figures & tables

Figure 1: A 2D example illustrating multi-modal expert demonstrations and trajectories predicted by different imitation policies. Behavioral cloning predictions collapse into a single mean. Discrete Policy succeeds but introduces temporal discontinuities. Generative Policy bounces between mode 1 and 2. Our work predicts consistent and fine-grained trajectory.
Figure 2: Overview of our PF-DAG framework. The input observation features are extracted via Observation Feature Extraction and then fed to the Primary Mode Policy π1\pi_{1} . The GT action chunks are compressed into discrete primary modes using VQ-VAE and supervise π1\pi_{1} , which are only used in training stage. The Mode Conditioned MeanFlow Policy π2\pi_{2} takes the selected primary mode mm and observation features as input, generating high-fidelity continuous actions.
 AdroitDexArtMetaWorld 
MethodHammerDoorPenLaptopFaucetToiletBucketMedium (6)Hard (5)Success
IBC0.00±0.000.00\scriptstyle{\pm 0.00}0.00±0.000.00\scriptstyle{\pm 0.00}0.10±0.010.10\scriptstyle{\pm 0.01}0.01±0.010.01\scriptstyle{\pm 0.01}0.07±0.020.07\scriptstyle{\pm 0.02}0.15±0.010.15\scriptstyle{\pm 0.01}0.00±0.000.00\scriptstyle{\pm 0.00}0.11±0.020.11\scriptstyle{\pm 0.02}0.09±0.030.09\scriptstyle{\pm 0.03}0.08
BC-H0.10±0.090.10\scriptstyle{\pm 0.09}0.07±0.050.07\scriptstyle{\pm 0.05}0.16±0.030.16\scriptstyle{\pm 0.03}0.09±0.020.09\scriptstyle{\pm 0.02}0.13±0.040.13\scriptstyle{\pm 0.04}0.21±0.020.21\scriptstyle{\pm 0.02}0.10±0.010.10\scriptstyle{\pm 0.01}0.15±0.030.15\scriptstyle{\pm 0.03}0.18±0.050.18\scriptstyle{\pm 0.05}0.15
DP0.48±0.170.48\scriptstyle{\pm 0.17}0.50±0.050.50\scriptstyle{\pm 0.05}0.25±0.040.25\scriptstyle{\pm 0.04}0.69±0.040.69\scriptstyle{\pm 0.04}0.23±0.080.23\scriptstyle{\pm 0.08}0.58±0.020.58\scriptstyle{\pm 0.02}0.46±0.010.46\scriptstyle{\pm 0.01}0.20±0.050.20\scriptstyle{\pm 0.05}0.19±0.030.19\scriptstyle{\pm 0.03}0.30
DP31.00±0.00\mathbf{1.00\scriptstyle{\pm 0.00}}0.62±0.040.62\scriptstyle{\pm 0.04}0.43±0.060.43\scriptstyle{\pm 0.06}0.83±0.010.83\scriptstyle{\pm 0.01}0.63±0.020.63\scriptstyle{\pm 0.02}0.82±0.04\mathbf{0.82\scriptstyle{\pm 0.04}}0.46±0.020.46\scriptstyle{\pm 0.02}0.45±0.050.45\scriptstyle{\pm 0.05}0.35±0.020.35\scriptstyle{\pm 0.02}0.51
FlowPolicy1.00±0.00\mathbf{1.00\scriptstyle{\pm 0.00}}0.58±0.050.58\scriptstyle{\pm 0.05}0.53±0.120.53\scriptstyle{\pm 0.12}0.85±0.020.85\scriptstyle{\pm 0.02}0.42±0.100.42\scriptstyle{\pm 0.10}0.80±0.050.80\scriptstyle{\pm 0.05}0.39±0.060.39\scriptstyle{\pm 0.06}0.47±0.070.47\scriptstyle{\pm 0.07}0.37±0.070.37\scriptstyle{\pm 0.07}0.51
Table 1: Quantitative comparison of PF-DAG against state-of-the-art baselines on 18 tasks from three simulation benchmarks.
Figure 3: Visual comparison of failure modes in baselines versus PF-DAG. Mode Collapse outputs “average” actions, while Mode Bouncing randomly switches between consecutive time steps.
Task DescriptionTask ParametersSuccess Rate
# Variations# DemosEnd EffectorTactileAction DimVanilla BCDP3PF-DAG (Ours)
Pick Cube350Gripper✗7+17+10.200.600.70\mathbf{0.70}
Place Baymax110Gripper✗7+17+10.400.850.90\mathbf{0.90}
Wipe Table550XHand✓7+127+120.000.550.70\mathbf{0.70}
Place Toy Into Bin450XHand✓7+127+120.000.600.80\mathbf{0.80}
Table 2: Quantitative comparison of success rates of different methods on real world manipulation tasks. The table presents key task parameters alongside the performance of each method.
AblationsBenchmarks 
w. PMw. MFToken.# ModesAdroit (3)DexArt (4)MetaWorld (11)Weighted Success
✓✓VQ-VAE640.77±0.03\mathbf{0.77\scriptstyle{\pm 0.03}}0.72±0.04\mathbf{0.72\scriptstyle{\pm 0.04}}0.70±0.02\mathbf{0.70\scriptstyle{\pm 0.02}}0.72\mathbf{0.72}
✓✗VQ-VAE640.02±0.010.02\scriptstyle{\pm 0.01}0.00±0.000.00\scriptstyle{\pm 0.00}0.02±0.010.02\scriptstyle{\pm 0.01}0.01
✗✓VQ-VAE640.62±0.020.62\scriptstyle{\pm 0.02}0.65±0.010.65\scriptstyle{\pm 0.01}0.51±0.030.51\scriptstyle{\pm 0.03}0.56
✓✓VQ-VAE80.72±0.030.72\scriptstyle{\pm 0.03}0.70±0.020.70\scriptstyle{\pm 0.02}0.55±0.050.55\scriptstyle{\pm 0.05}0.61
✓✓VQ-VAE10240.66±0.020.66\scriptstyle{\pm 0.02}0.68±0.030.68\scriptstyle{\pm 0.03}0.52±0.060.52\scriptstyle{\pm 0.06}0.58
Table 3: Ablation study on the impact of PF-DAG’s key components and hyperparameters. w. PM denotes whether the primary mode policy is included. w. MF indicates whether the mode-conditioned MeanFlow policy is included. Token. means action tokenization method. # Modes represents the number of discrete primary modes.
Figure 4: Hardware and manipulated objects used in real world experiments.
Figure 5: Illustration of critical properties of PF-DAG. (a) Action chunks are projected to 2D via PCA, colored by their assigned primary mode. (b) PF-DAG’s one-step MeanFlow decoder achieves FPS comparable to 1-NFE DP3 while maintaining significantly higher success. (c) PF-DAG preserves high success even with short chunks by avoiding primary mode bouncing.
Figure 6: Visualization of the mode probability distribution predicted by the Primary Mode Policy π1\pi_{1} at three selected frames. The vertical axis of the heatmap represents the mode index.
VariationsODE SolverSuccess
CFM1-NFE Euler0.69±0.030.69\scriptstyle{\pm 0.03}
CFM10-NFE Euler0.69±0.020.69\scriptstyle{\pm 0.02}
CFMRunge-Kutta0.68±0.010.68\scriptstyle{\pm 0.01}
Meanflow1-NFE0.72±0.02\mathbf{0.72\scriptstyle{\pm 0.02}}
Table 4: Comparison of success for MeanFlow (MF) and Conditional Flow Matching (CFM), varying the ODE solver and NFE.
Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

 AdroitDexArtMeta-World (Easy)
Alg \\backslash TaskHammerDoorPenLaptopFaucetToiletBucketButton PressButton Press Topdown
Diffusion Policy45±545\scriptstyle{\pm 5}37±237\scriptstyle{\pm 2}13±213\scriptstyle{\pm 2}69±469\scriptstyle{\pm 4}23±823\scriptstyle{\pm 8}58±258\scriptstyle{\pm 2}46±146\scriptstyle{\pm 1}99±199\scriptstyle{\pm 1}98±198\scriptstyle{\pm 1}
3D Diffusion Policy100±0\mathbf{100\scriptstyle{\pm 0}}62±462\scriptstyle{\pm 4}43±643\scriptstyle{\pm 6}83±183\scriptstyle{\pm 1}63±263\scriptstyle{\pm 2}82±4\mathbf{82\scriptstyle{\pm 4}}46±246\scriptstyle{\pm 2}100±0\mathbf{100\scriptstyle{\pm 0}}100±0\mathbf{100\scriptstyle{\pm 0}}
PF-DAG (Ours)100±0\mathbf{100\scriptstyle{\pm 0}}65±3\mathbf{65\scriptstyle{\pm 3}}65±3\mathbf{65\scriptstyle{\pm 3}}90±2\mathbf{90\scriptstyle{\pm 2}}72±5\mathbf{72\scriptstyle{\pm 5}}82±2\mathbf{82\scriptstyle{\pm 2}}65±3\mathbf{65\scriptstyle{\pm 3}}100±0\mathbf{100\scriptstyle{\pm 0}}100±0\mathbf{100\scriptstyle{\pm 0}}
Appendix
Table 5: Quantitative comparison of PF-DAG against baselines on more tasks from Adroit, DexArt, and MetaWorld benchmarks.
Figure 7: Qualitative visualization of the Primary Mode Policy ( π1\pi_{1} ) at keyframes from four different simulation tasks. Each row corresponds to a single task episode. Within each row, three keyframes show the point cloud observation (left) and the corresponding predicted probability distribution over the discrete primary modes (right) as a heatmap. The vertical axis of the heatmap represents the mode index. The shifting patterns in the heatmaps demonstrate that the policy learns a dynamic, context-dependent mapping from observation to a belief over high-level actions as the task progresses.
HyperparameterDescriptionValue
Prediction Horizon TpT_{p}The total number of timesteps in a predicted action chunk.32 / 16
Execution Horizon TaT_{a}The number of timesteps from the chunk executed before re-planning.16 / 8
Learning RateThe peak learning rate after the warmup phase.1e-4
Weight DecayThe weight decay value for the AdamW optimizer.0.01
Batch SizeThe number of samples processed per training step.128
Codebook Size KKThe number of discrete primary modes in the VQ-VAE codebook.64
Appendix
Table 6: Hyperparameters for the PF-DAG framework.
AblationsBenchmarks 
# ModesAdroit (3)DexArt (4)MetaWorld (11)Weighted Success
640.77±0.030.77\scriptstyle{\pm 0.03}0.72±0.04\mathbf{0.72\scriptstyle{\pm 0.04}}0.70±0.02\mathbf{0.70\scriptstyle{\pm 0.02}}0.72\mathbf{0.72}
80.72±0.030.72\scriptstyle{\pm 0.03}0.70±0.020.70\scriptstyle{\pm 0.02}0.55±0.050.55\scriptstyle{\pm 0.05}0.61
160.79±0.02\mathbf{0.79\scriptstyle{\pm 0.02}}0.71±0.030.71\scriptstyle{\pm 0.03}0.68±0.010.68\scriptstyle{\pm 0.01}0.70
320.76±0.030.76\scriptstyle{\pm 0.03}0.71±0.020.71\scriptstyle{\pm 0.02}0.67±0.050.67\scriptstyle{\pm 0.05}0.70
1280.77±0.010.77\scriptstyle{\pm 0.01}0.69±0.030.69\scriptstyle{\pm 0.03}0.67±0.030.67\scriptstyle{\pm 0.03}0.69
Appendix
Table 7: Ablation study on the mode number KK of PF-DAG.
MethodTotal Jerk ( ↓\downarrow )
DP31.25
PF-DAG (Ours)0.45
Appendix
Table 8: Total end-effector jerk ( ↓\downarrow ) comparison on the real-world ’Wipe Table’ task.

Explore similar work

CardsList

  1. VQActFlow: Vector-Quantized Action Mode Steering for Multi-Task Robot Manipulation

    Jun 19, 2026Zhigen Zhao, Mark Leggiero, Yipu Chen +5Latent Action LearningLanguage-Conditioned Robot Manipulation

    Institute for Robotics and Intelligent Machines (IRIM), Georgia Institute of Technology, Atlanta, GA 30332

  2. CoLA-Flow Policy: Temporally Coherent Imitation Learning via Continuous Latent Action Flow Matching for Robotic Manipulation

    Jan 30, 2026Wu Songwei, Jiang Zhiduo, Sun Wandong +4Flow MatchingLatent Action Learning

    Harbin Institute of Technology, Harbin 150001, China · Honor Device Co., Ltd., Shenzhen, China

  3. Language-Grounded Decoupled Action Representation for Robotic Manipulation

    Mar 13, 2026Wuding Weng, Tongshu Wu, Liucheng Chen +5Robotic ManipulationLanguage-Conditioned Robot Manipulation

    School of Computer Science and Technology, Tongji University, China