Primary-Fine Decoupling for Action Generation in Robotic Imitation

Xiaohan Lei; Min Wang; Wengang Zhou; Xingyu Lu; Houqiang Li

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

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 $\pi_{1}$ . The GT action chunks are compressed into discrete primary modes using VQ-VAE and supervise $\pi_{1}$ , which are only used in training stage. The Mode Conditioned MeanFlow Policy $\pi_{2}$ takes the selected primary mode $m$ and observation features as input, generating high-fidelity continuous actions.

	Adroit			DexArt				MetaWorld
Method	Hammer	Door	Pen	Laptop	Faucet	Toilet	Bucket	Medium (6)	Hard (5)	Success
IBC	$0.00\scriptstyle{\pm 0.00}$	$0.00\scriptstyle{\pm 0.00}$	$0.10\scriptstyle{\pm 0.01}$	$0.01\scriptstyle{\pm 0.01}$	$0.07\scriptstyle{\pm 0.02}$	$0.15\scriptstyle{\pm 0.01}$	$0.00\scriptstyle{\pm 0.00}$	$0.11\scriptstyle{\pm 0.02}$	$0.09\scriptstyle{\pm 0.03}$	0.08
BC-H	$0.10\scriptstyle{\pm 0.09}$	$0.07\scriptstyle{\pm 0.05}$	$0.16\scriptstyle{\pm 0.03}$	$0.09\scriptstyle{\pm 0.02}$	$0.13\scriptstyle{\pm 0.04}$	$0.21\scriptstyle{\pm 0.02}$	$0.10\scriptstyle{\pm 0.01}$	$0.15\scriptstyle{\pm 0.03}$	$0.18\scriptstyle{\pm 0.05}$	0.15
DP	$0.48\scriptstyle{\pm 0.17}$	$0.50\scriptstyle{\pm 0.05}$	$0.25\scriptstyle{\pm 0.04}$	$0.69\scriptstyle{\pm 0.04}$	$0.23\scriptstyle{\pm 0.08}$	$0.58\scriptstyle{\pm 0.02}$	$0.46\scriptstyle{\pm 0.01}$	$0.20\scriptstyle{\pm 0.05}$	$0.19\scriptstyle{\pm 0.03}$	0.30
DP3	$\mathbf{1.00\scriptstyle{\pm 0.00}}$	$0.62\scriptstyle{\pm 0.04}$	$0.43\scriptstyle{\pm 0.06}$	$0.83\scriptstyle{\pm 0.01}$	$0.63\scriptstyle{\pm 0.02}$	$\mathbf{0.82\scriptstyle{\pm 0.04}}$	$0.46\scriptstyle{\pm 0.02}$	$0.45\scriptstyle{\pm 0.05}$	$0.35\scriptstyle{\pm 0.02}$	0.51
FlowPolicy	$\mathbf{1.00\scriptstyle{\pm 0.00}}$	$0.58\scriptstyle{\pm 0.05}$	$0.53\scriptstyle{\pm 0.12}$	$0.85\scriptstyle{\pm 0.02}$	$0.42\scriptstyle{\pm 0.10}$	$0.80\scriptstyle{\pm 0.05}$	$0.39\scriptstyle{\pm 0.06}$	$0.47\scriptstyle{\pm 0.07}$	$0.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 Description	Task Parameters					Success Rate
Task Description	# Variations	# Demos	End Effector	Tactile	Action Dim	Vanilla BC	DP3	PF-DAG (Ours)
Pick Cube	3	50	Gripper	✗	$7+1$	0.20	0.60	$\mathbf{0.70}$
Place Baymax	1	10	Gripper	✗	$7+1$	0.40	0.85	$\mathbf{0.90}$
Wipe Table	5	50	XHand	✓	$7+12$	0.00	0.55	$\mathbf{0.70}$
Place Toy Into Bin	4	50	XHand	✓	$7+12$	0.00	0.60	$\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.

Ablations				Benchmarks
w. PM	w. MF	Token.	# Modes	Adroit (3)	DexArt (4)	MetaWorld (11)	Weighted Success
✓	✓	VQ-VAE	64	$\mathbf{0.77\scriptstyle{\pm 0.03}}$	$\mathbf{0.72\scriptstyle{\pm 0.04}}$	$\mathbf{0.70\scriptstyle{\pm 0.02}}$	$\mathbf{0.72}$
✓	✗	VQ-VAE	64	$0.02\scriptstyle{\pm 0.01}$	$0.00\scriptstyle{\pm 0.00}$	$0.02\scriptstyle{\pm 0.01}$	0.01
✗	✓	VQ-VAE	64	$0.62\scriptstyle{\pm 0.02}$	$0.65\scriptstyle{\pm 0.01}$	$0.51\scriptstyle{\pm 0.03}$	0.56
✓	✓	VQ-VAE	8	$0.72\scriptstyle{\pm 0.03}$	$0.70\scriptstyle{\pm 0.02}$	$0.55\scriptstyle{\pm 0.05}$	0.61
✓	✓	VQ-VAE	1024	$0.66\scriptstyle{\pm 0.02}$	$0.68\scriptstyle{\pm 0.03}$	$0.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 $\pi_{1}$ at three selected frames. The vertical axis of the heatmap represents the mode index.

Variations	ODE Solver	Success
CFM	1-NFE Euler	$0.69\scriptstyle{\pm 0.03}$
CFM	10-NFE Euler	$0.69\scriptstyle{\pm 0.02}$
CFM	Runge-Kutta	$0.68\scriptstyle{\pm 0.01}$
Meanflow	1-NFE	$\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

	Adroit			DexArt				Meta-World (Easy)
Alg $\backslash$ Task	Hammer	Door	Pen	Laptop	Faucet	Toilet	Bucket	Button Press	Button Press Topdown
Diffusion Policy	$45\scriptstyle{\pm 5}$	$37\scriptstyle{\pm 2}$	$13\scriptstyle{\pm 2}$	$69\scriptstyle{\pm 4}$	$23\scriptstyle{\pm 8}$	$58\scriptstyle{\pm 2}$	$46\scriptstyle{\pm 1}$	$99\scriptstyle{\pm 1}$	$98\scriptstyle{\pm 1}$
3D Diffusion Policy	$\mathbf{100\scriptstyle{\pm 0}}$	$62\scriptstyle{\pm 4}$	$43\scriptstyle{\pm 6}$	$83\scriptstyle{\pm 1}$	$63\scriptstyle{\pm 2}$	$\mathbf{82\scriptstyle{\pm 4}}$	$46\scriptstyle{\pm 2}$	$\mathbf{100\scriptstyle{\pm 0}}$	$\mathbf{100\scriptstyle{\pm 0}}$
PF-DAG (Ours)	$\mathbf{100\scriptstyle{\pm 0}}$	$\mathbf{65\scriptstyle{\pm 3}}$	$\mathbf{65\scriptstyle{\pm 3}}$	$\mathbf{90\scriptstyle{\pm 2}}$	$\mathbf{72\scriptstyle{\pm 5}}$	$\mathbf{82\scriptstyle{\pm 2}}$	$\mathbf{65\scriptstyle{\pm 3}}$	$\mathbf{100\scriptstyle{\pm 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 ( $\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.

Hyperparameter	Description	Value
Prediction Horizon $T_{p}$	The total number of timesteps in a predicted action chunk.	32 / 16
Execution Horizon $T_{a}$	The number of timesteps from the chunk executed before re-planning.	16 / 8
Learning Rate	The peak learning rate after the warmup phase.	1e-4
Weight Decay	The weight decay value for the AdamW optimizer.	0.01
Batch Size	The number of samples processed per training step.	128
Codebook Size $K$	The number of discrete primary modes in the VQ-VAE codebook.	64

Appendix

Table 6: Hyperparameters for the PF-DAG framework.

Ablations	Benchmarks
# Modes	Adroit (3)	DexArt (4)	MetaWorld (11)	Weighted Success
64	$0.77\scriptstyle{\pm 0.03}$	$\mathbf{0.72\scriptstyle{\pm 0.04}}$	$\mathbf{0.70\scriptstyle{\pm 0.02}}$	$\mathbf{0.72}$
8	$0.72\scriptstyle{\pm 0.03}$	$0.70\scriptstyle{\pm 0.02}$	$0.55\scriptstyle{\pm 0.05}$	0.61
16	$\mathbf{0.79\scriptstyle{\pm 0.02}}$	$0.71\scriptstyle{\pm 0.03}$	$0.68\scriptstyle{\pm 0.01}$	0.70
32	$0.76\scriptstyle{\pm 0.03}$	$0.71\scriptstyle{\pm 0.02}$	$0.67\scriptstyle{\pm 0.05}$	0.70
128	$0.77\scriptstyle{\pm 0.01}$	$0.69\scriptstyle{\pm 0.03}$	$0.67\scriptstyle{\pm 0.03}$	0.69

Appendix

Table 7: Ablation study on the mode number

K

of PF-DAG.

Method	Total Jerk ( $\downarrow$ )
DP3	1.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

Jun 19, 2026

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

Multi-task robot manipulation policies are challenging to learn from demonstration because traditionally a single network must select among qualitatively different action modes from a multimodal demonstration distribution, conditioned on language and visual context. A wrong mode selection means executing the wrong task or an action infeasible in the scene. Tokenizing continuous actions into a learned discrete codebook separates these modes at the representation level, offering structural advantages for multi-task learning. We propose VQActFlow, a multi-task manipulation policy that tokenizes action chunks and generates code sequences via Variational Flow Matching. VQActFlow maintains an explicit preference over action modes throughout generation. Inference-time guidance acts on this preference to steer mode commitment. We instantiate this with classifier-free guidance over language conditioning, which steers the policy toward the instructed action mode, and a learned codebook critic that supplies a complementary feasibility signal. We evaluate VQActFlow on three platforms: the LIBERO simulation benchmarks, a Unitree G1 humanoid performing whole-body pick-and-place, and an ALOHA-style bimanual platform performing contact-rich tasks. Across these benchmarks, VQActFlow outperforms both continuous and discrete baselines.

Zhigen Zhao, Mark Leggiero, Yipu Chen +5

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

Latent Action Learning Language-Conditioned Robot Manipulation Multi-Task Robotic Manipulation Vector Quantization Robotic Manipulation Robot Imitation Learning

Ask PDF

Jan 30, 2026

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

Learning long-horizon robotic manipulation requires jointly achieving expressive behavior modeling, real-time inference, and stable execution, which remains challenging for existing generative policies. Diffusion-based approaches offer strong modeling capacity but incur high inference latency, while flow matching enables fast, near-single-step generation yet often suffers from unstable execution when operating directly in the raw action space. We propose Continuous Latent Action Flow Policy (CoLA-Flow Policy), a trajectory-level imitation learning framework that performs flow matching in a continuous latent action space. By encoding action sequences into temporally coherent latent trajectories and learning an explicit latent-space flow, CoLA-Flow Policy decouples global motion structure from low-level control noise, enabling smooth and reliable long-horizon execution. The framework further integrates geometry-aware point cloud conditioning and execution-time multimodal modulation, using visual cues as a representative modality to enhance real-world robustness. Experiments in simulation and on real robots show that CoLA-Flow Policy achieves near-single-step inference, improves trajectory smoothness by up to 93.7% and task success by up to 25 percentage points over raw action-space flow baselines, while remaining significantly faster than diffusion-based policies.

Wu Songwei, Jiang Zhiduo, Sun Wandong +4

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

Flow Matching Latent Action Learning Long-Horizon Robotic Manipulation Robot Imitation Learning

Ask PDF

Mar 13, 2026

Language-Grounded Decoupled Action Representation for Robotic Manipulation

The heterogeneity between high-level vision-language understanding and low-level action control remains a fundamental challenge in robotic manipulation. Although recent methods have advanced task-specific action alignment, they often struggle to generate robust and accurate actions for novel or semantically related tasks. To address this, we propose the Language-Grounded Decoupled Action Representation (LaDA) framework, which leverages natural language as a semantic bridge to connect perception and control. LaDA introduces a fine-grained intermediate layer of three interpretable action primitives--translation, rotation, and gripper control--providing explicit semantic structure for low-level actions. It further employs a semantic-guided soft-label contrastive learning objective to align similar action primitives across tasks, enhancing generalization and motion consistency. An adaptive weighting strategy, inspired by curriculum learning, dynamically balances contrastive and imitation objectives for stable and effective training. Extensive experiments on simulated benchmarks (LIBERO and MimicGen) and real-world demonstrations validate that LaDA achieves strong performance and generalizes effectively to unseen or related tasks.

Wuding Weng, Tongshu Wu, Liucheng Chen +5

School of Computer Science and Technology, Tongji University, China

Robotic Manipulation Language-Conditioned Robot Manipulation Generalization in Robotic Manipulation Robot Imitation Learning

Ask PDF