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

The gist The proposed VQActFlow policy is a multi-task manipulation policy that tokenizes action chunks and generates code sequences via Variational Flow Matching, maintaining an explicit preference over action modes throughout generation.

VQActFlow Architecture

1. VQActFlow generates manipulation actions as sequences of discrete tokens drawn from a VQ-VAE codebook (Page 4, VQActFlow generates manipulation actions as sequences of discrete tokens drawn from a VQ-VAE codebook (Sec. III-A)).

2. A temporal encoder E maps an action chunk to a sequence of continuous latents z = E(a1:H) in R L×de, where L = H/2 S is the compressed sequence length determined by S temporal downsampling stages and de is the codebook embedding dimension (Page 4, A VQ-VAE Action Tokenization).

3. Each latent is quantized to its nearest codebook entry as in Eq. 1, and a temporal CNN decoder D reconstructs the action chunk from the quantized embeddings, so each output timestep depends on a local neighborhood of codes rather than on a single code in isolation (Page 4, A VQ-VAE Action Tokenization).

Variational Flow Matching Policy

1. The VFM policy trains the network on a cross-entropy over codebook indices rather than an L2 regression toward codebook embeddings, treating the K codes as competing action modes and producing an explicit preference over them at every integration step (Page 4, Variational Flow Matching Policy).

2. Given a demonstration sample, the frozen VQVAE encoder produces target embeddings x1 = zq in R L×de and ground-truth indices k∗ l L l=1. We sample source noise x0 ∼ N (0, I) in R L×de and a timestep t ∼ U(0, 1), and construct the optimal-transport interpolant xt = (1−t) x0+t x1 from Eq. 3 (Page 4, Variational Flow Matching Policy).

3. The policy network fθ processes xt conditioned on t, an observation encoding o, and language conditioning c, producing logits over the codebook at every position: h(l) = fθ(xt, t, o, c) (l) in R K l = 1..., L (Page 4).

4. Training minimizes the cross-entropy between the predicted categorical posterior and the ground-truth codebook indices: LVFM = − X L l=1 (h(l) k∗ l − logX K j=1 exp h(l)) (Page 4).

5. At each integration step, we convert the categorical output to a velocity field via the posterior-weighted mean of the codebook. Let p(l) k = exp(h(l) k)/P j exp(h(l) j). The posterior mean in embedding space is µ1 (l) = X K k=1 p(l)k · ek, (Page 4).

6. This velocity points toward the codebook-weighted average of the action modes currently favored by the policy, with p(l)k controlling the relative weight on each candidate. As the policy concentrates probability on a single mode during integration, this average becomes dominated by the embedding of the favored mode (Page 4).

Inference-Time Guidance Mechanisms

1. The explicit preference is maintained throughout generation, providing a unified interface for inference-time guidance, on which CFG over language conditioning (Sec. IV-C) steers the policy toward the language-instructed action mode and a learned codebook critic (Sec. IV-D) pushes the policy toward modes feasible in the current scene (Page 1).

2. Classifier-Free Guidance amplifies the difference between taskconditioned and unconditional predictions to steer the policy toward the language-instructed action mode, effectively sharpening which action mode the policy commits to under ambiguous observations (Page 4, Classifier-Free Guidance for Language Conditioning).

3. The codebook critic supplies a learned feasibility signal that biases the policy’s mode preference toward actions feasible in the current scene (Page 4, Codebook Critic for Scene-Conditioned Feasibility).

4. During ODE integration, the critic gradient adjusts the policy’s preference over action modes: p˜ = Proj∆[p + λ(t) · ΠT [∇pCψ(o, c, p, t)] (Page 4).

5. The guidance weight adapts to the entropy of the current distribution: λ(t) = λmax 1 − H(p) log K (Page 4).

Experimental Results

1. VQActFlow outperforms both continuous and discrete baselines across three platforms: the LIBERO simulation benchmarks, a Unitree G1 humanoid performing whole-body pick-and-place, and an ALOHA-style bimanual platform performing contactrich tasks (Page 1).

2. On LIBERO-Goal, VQActFlow improves from 73.0% unguided to 81.0% at w=4, then degrades monotonically to 63.0% at w=8 as pertask variance widens (Page 6, CFG Evaluation on LIBERO-Goal).

3. On LIBERO-90, VQActFlow with CFG alone reaches 80.5%, surpassing Continuous Flow Matching DiT (CFM), MT-ACT [1], and all VQ-based baselines when using w=2 and λ=1.0 (Page 6, Multi-task Scaling on LIBERO-90).

4. On the Unitree G1 humanoid, CFG raises aggregate success from 23.8% to 57.5% (Page 7).

5. In bimanual manipulation, combining w=4 and λ=1 yields the best result at 77.5%, which is an improvement over Discrete Policy with w=4 (61.3%) (Page 8, Bimanual Manipulation).

Conclusion

VQActFlow demonstrates that maintaining an explicit preference over discrete action modes during generation provides a more effective interface for inference-time guidance than continuous-embedding alternatives (Page 6). The contrastively-trained codebook critic operates through the same interface and supplies a feasibility signal that CFG alone does not (Page 7). VQActFlow deploys successfully on two distinct hardware platforms, a Unitree G1 humanoid and an ALOHA-style bimanual system, executing contact-rich tasks where continuous flow matching at matched capacity fails (Page 8).

I. INTRODUCTION

Robot policies are increasingly expected to handle multiple manipulation tasks within a single model [1], [2], [3], [4] (Page 1). Demonstration data for this setting is multimodal: the same observation recurs under different instructions and is followed by qualitatively different actions (Page 1). A policy trained on such data must select the instructed action mode from this multimodal distribution, distinguishing modes that differ in target object, contact strategy, or bimanual coordination (Page 1). When the observation is ambiguous, a wrong action mode selection is not a small numerical error but the robot carrying out a coherent behavior for the wrong task or performing the right task in an infeasible manner (Page 1). The policy must therefore commit to the action mode consistent with the instruction and feasible in the scene (Page 1).

II. RELATED WORK

Diffusion and flow-based robot policies. Diffusion models [15] generate samples by iteratively denoising Gaussian noise toward the data distribution (Page 2). Flow matching [6] reframes this as transporting noise along ODE trajectories defined by learned velocity fields (Page 2). Diffusion Policy [5] applies the diffusion paradigm to robot action chunks, modeling multi-step action distributions via DDPM denoising and achieving strong visuomotor control on contact-rich manipulation tasks (Page 2). Discrete formulations of flow matching [18], [19] transport probability over categorical states directly, whereas Variational Flow Matching [13], [14] retains a continuous transport, casting it as variational inference over categorical posteriors (Page 2).

Robot action tokenization. OpenVLA [20] discretizes actions via per-dimension scalar binning to reuse language model vocabularies, a simple scheme that sacrifices interdimension correlations (Page 2). Discrete Policy [12] learns a VQVAE codebook over action sequences and shows that the learned codes cluster by task similarity, suggesting that VQVAE tokenization captures task-discriminative structure in multi-task action distributions (Page 2).

Guidance for generative policies. Classifier guidance [9] steers diffusion sampling with a trained classifier’s gradients; CFG [8] avoids a separate classifier by interpolating conditional and unconditional scores (Page 2). Guided Flows [23] extends guidance to flow matching with faster sampling than diffusion (Page 2). VQActFlow differs from prior robotics guidance in what guidance acts on: rather than steering continuous action distributions or trajectories, both CFG and the contrastively-trained codebook critic act on the policy’s categorical preference over discrete action tokens (Page 2).

III. PRELIMINARIES

A. Vector-Quantized Variational Autoencoder

A VQ-VAE [11] compresses a continuous input into a sequence of indices into a learned codebook C = 1, where each ek ∈ R de is a codebook embedding (Page 3).

1. textbfExplicit Mode Commitment During Generation: Maintain explicit preference over discrete action modes throughout generation by generating logits over codes at every integration step using Variational Flow Matching (VFM). This ensures the mode commitment is represented at every step rather than recovered only at the end, directly addressing the need to resolve ambiguous observations toward the instructed task during sampling, as opposed to earlier methods that abandon discrete structure during generation.

2. textbfEnhanced Inference-Time Steering: Implement guidance through a unified categorical interface by using both Classifier-Free Guidance (CFG) and a contrastively-trained codebook critic simultaneously. CFG steers the policy toward the language-instructed action mode, while the critic pushes the policy toward action modes feasible in the current scene, providing complementary signals at each integration step.

3. textbfContext-Aware Feasibility Signal: Introduce a learned codebook critic that supplies a feasibility signal by training it via contrastive learning against structurally diverse negative examples across temporal coherence, observation grounding, and language grounding. This allows the critic to guide the policy using p˜ = Proj∆[p + λ(t) · ΠT [∇pCψ(o, c, p, t)], amplifying feasibility along axes that CFG alone does not.

4. textbfTask-Specific Guidance Scheduling: Implement a guidance weight schedule where λ(t) is small when p is diffuse and strengthens as p crystallizes, coupling guidance strength to distributional confidence. This ensures the policy only receives meaningful feasibility feedback once it has concentrated probability on likely code selections, preventing premature steering toward arbitrary codes early in the integration process.

5. textbfImproved Trajectory Smoothness: Achieve significantly smoother trajectories compared to continuous baselines by leveraging the discrete bottleneck and VFM transport mechanism instead of continuous flow matching. The paper notes that VQActFlow remains near 2 × 102 rad/s3, roughly two orders of magnitude smoother than CFM, attributing this to the VQ-VAE’s temporal compression structure.