MambaAct: Efficient Mamba Action Experts for Flow-Matching Vision-Language-Action Models

AI summary

Overview

Research area: Robotics — efficient vision-language-action (VLA) policies for robotic manipulation, specifically flow-matching VLA architectures.

Technical level: Advanced. The paper assumes familiarity with flow matching, transformer action experts, state-space models (Mamba), and VLA benchmarks such as LIBERO.

Scope: The paper proposes MambaAct, a compact Mamba-based action expert that replaces the transformer action expert of π0.5 while keeping its pretrained visual-language backbone, and measures the resulting success-rate and latency trade-off in simulation and on two real UR5 tasks.

What This Paper Is About

Flow-matching VLA policies such as π0 and π0.5 generate actions by iteratively denoising an action chunk, which means the action expert is evaluated many times per policy query while the visual-language prefix is computed only once. The paper reports that 10-step denoising accounts for 70.1% of end-to-end latency in π0.5 and 86.4% in SmolVLA, making the repeatedly evaluated expert the dominant cost. The goal is to accelerate that specific component by swapping the large transformer action expert for a much smaller Mamba-based expert, without changing the pretrained VLM, the flow-matching formulation, or adding distillation.

Key Contributions

1. Action-expert replacement as an efficiency axis: MambaAct replaces the repeatedly evaluated transformer action expert of π0.5 with a compact Mamba expert while retaining the pretrained visual-language backbone and standard flow-matching inference, with no teacher model, distillation stage, or additional pretraining.
2. Aggressive expert compression: the new expert has 29.7M parameters versus 311.5M for the original π0.5 action expert, and total model size drops from 3.62B to 2.96B parameters, while mean LIBERO success improves by 1.08 percentage points across three fine-tuning seeds.
3. A conditioning interface that avoids per-step cross-attention: a lightweight Context Compressor maps VLM prefix tokens into K = 64 context tokens (44 visual, 16 language, 4 learned-query global summaries), pooled into a single vector that block-specific FiLM MLPs convert into scale, shift, and residual-gate parameters for each Mamba block.
4. Controlled ablations and composability checks: a 2×2 study of Mamba versus transformer mixers and gated versus cross-attention conditioning, a composition test with VLM-depth reduction (12 of 18 layers), and a claim that the approach transfers to SmolVLA (the specific transfer numbers are not present in the provided text).

Main Findings

• LIBERO success (three seeds {17, 27, 42}, 50 trials per task, evaluation seed 7): average success rises from 96.72 ± 0.14% (π0.5) to 97.80 ± 0.43% (MambaAct), a gain of 1.08 percentage points. Per-suite: Spatial 98.00 ± 0.43 → 98.73 ± 0.34, Object 98.33 ± 0.66 → 99.60 ± 0.43, Goal 97.73 ± 0.84 → 98.60 ± 0.28, L-10 92.80 ± 0.43 → 94.27 ± 1.52.
• Latency (seed-42 checkpoints, single H100): prefix latency rises from 40.76 ms to 46.61 ms due to the Context Compressor, denoising falls from 160.26 ms to 56.40 ms (64.8% reduction), and end-to-end latency falls from 228.6 ms to 107.6 ms (52.9% reduction), corresponding to a 2.12× speedup over the matched baseline.
• Best single run: seed 42 reaches 99.2%, 99.8%, 98.8%, and 95.8% on Spatial, Object, Goal, and L-10, averaging 98.4% — close to SnapFlow's reported 98.8%, though the authors note SnapFlow uses 10 trials per task, reports no seed variability, retains the full π0.5 architecture, and relies on progressive self-distillation.
• Zero-shot robustness (LIBERO-Plus, 10,030 test-only task instances, seven perturbation categories): MambaAct reaches 79.7% overall versus 78.7% for the matched π0.5, with the largest gain on robot-pose variation (70.9% → 77.6%), the highest listed lighting score at 95.2%, and a matching best layout score of 86.4%. The main regression is under observation noise, 77.4% → 73.5%, while camera robustness is essentially unchanged (61.6% → 61.5%).
• Real-world UR5 tasks (RTX 5090): both models reach 100% in-distribution success; MambaAct improves out-of-distribution success from 86.5% to 95.2% on nut pouring and from 76.9% to 82.7% on tube placing, with final scores of 96.9% versus 91.3% and 88.8% versus 85.0%. Denoising latency drops 58.8% on both tasks, giving roughly 33.9% and 33.3% lower end-to-end latency; training time falls from 14.3 h to 14.0 h and GFLOPs from 5001.2 to 4488.8 (10.2% reduction).
• Ablation, architectural source of gain (seed 42): with gated conditioning, Mamba averages 98.4% versus 85.7% for a size-matched transformer; with cross-attention, 96.7% versus 75.6%. Gated conditioning beats cross-attention in both cases (Mamba: 96.7% → 98.4%, latency 119.66 ms → 107.61 ms), but the gated experts are larger (30M vs. 17–18M parameters), so the authors state the accuracy contribution of the interface remains partially confounded by expert size.
• Composability with VLM-depth reduction (L12, 12 of 18 layers): on LIBERO, training time drops from 13.78 h to 10.81 h and end-to-end latency from 107.6 ms to 102.91 ms while average success is 97.5% versus 98.4% for the full-depth model. In the real-world setting, model size falls from 2.96B to 2.30B parameters, GFLOPs per query from 4488.84 to 3209.70, and latency from 83 ms to 74 ms, with 100% ID success and final scores of 95.0% (nut pouring) and 90.0% (tube placing); the provided text is truncated at this point.

Methodology in Plain English

The starting point is π0.5, a policy that pairs a pretrained vision-language model with an action expert. Given camera images and a language instruction, the VLM produces a prefix representation once. The action expert then takes a noisy action chunk and, over 10 denoising steps, repeatedly predicts a velocity that is used to push the chunk toward a clean action sequence. Because the expert runs at every step, its per-step cost is multiplied by 10.

MambaAct keeps the VLM, including the SigLIP vision encoder and Gemma prefix encoder, and jointly fine-tunes it. Only the action expert is swapped out. Two problems follow from that swap: the new expert is randomly initialized, and it needs to know about the current scene and instruction without re-running attention over the full prefix at every denoising step.

The paper solves the first conditioning problem with a Context Compressor. Prefix tokens are projected to the expert's width, visual tokens are pooled into 44 contiguous segments and language tokens into 16, and four learned queries extract global summaries via multi-head attention. The resulting tokens are averaged into a single vector, which small per-block networks turn into a scale, a shift, and a residual gate — the FiLM interface. That vector is computed once and reused across all denoising steps, so the expert never performs token-level cross-attention to the prefix.

The second problem is handled by the expert itself. The noisy action chunk is projected into exactly H action tokens (no separate robot-state token; actions are padded to dimension 32), and a 4-block Mamba-3 stack mixes information along that ordered action sequence. Flow time is injected separately into each block through a small projection. Initialization is deliberately conservative: time projections and FiLM scale/shift outputs are zero-initialized, gate logits are set to log(0.2/0.8), so the new Mamba branch starts as a small residual (multiplier 0.2) on top of the action representation, and the gates are learned during fine-tuning. Training uses the standard flow-matching regression loss against the target velocity ε − A₀, with 30,000 optimizer steps, batch size 16, AdamW, gradient clipping at 1.0, weight decay 10⁻¹⁰, a cosine schedule from 5×10⁻⁵ to 5×10⁻⁶, and bfloat16 on a single H100. The expert uses width 512, state size 64, expansion factor 1, and head dimension 64; all 18 VLM layers are executed in the default configuration.

Evaluation covers the four LIBERO suites (Spatial, Object, Goal, LIBERO-10), zero-shot LIBERO-Plus perturbation categories, and two UR5 tasks — tube placing and nut pouring — with 112 demonstrations per task from 28 configurations, four episodes each, and 28 in-distribution plus 52 out-of-distribution trials. Real-world policies predict chunks of horizon 50 and execute the first 30 actions at 30 Hz.

Why This Matters

The paper reframes VLA efficiency: instead of shortening the denoising schedule, pruning the VLM, or shrinking the whole model, it targets the module that is called repeatedly. Because the acceleration multiplies across denoising steps, a modest per-step saving becomes a large end-to-end saving — 64.8% of denoising latency and 52.9% of end-to-end latency on LIBERO, with a smaller total model (3.62B to 2.96B parameters). The result that a 10× smaller action expert can also improve success, rather than merely preserve it, is the distinctive claim and suggests that the standard transformer expert may be over-parameterized for the action-mixing role.

Real-world applications:

• Industrial pick-and-place and assembly, where the reported tube-placing task is a stand-in for transferring parts between fixtures under tight cycle-time budgets.
• Table-top manipulation with containers and granular material, matching the nut-pouring task, where partial-credit scoring reflects the multi-stage nature of the skill.
• Deployment on cost- or power-constrained hardware, where the reduced parameter count (2.96B, or 2.30B with depth reduction) and lower FLOPs per query (4488.8, or 3209.7) matter alongside latency.
• Robust operation under environmental variation — lighting, layout, and robot-pose shifts — where the LIBERO-Plus gains (79.7% overall) indicate the compressed expert is not purely overfitting to benchmark conditions.

Industry relevance: the gains are reported as composition-compatible with other efficiency techniques (VLM-depth reduction, and by extension step compression or visual-token pruning), which matters for teams that have already optimized one axis of a deployed policy. Releasing code and weights at the linked repository lowers the barrier to reproducing and adopting the expert swap.

Future Directions

• Disentangling conditioning from capacity. The authors state that the accuracy contribution of gated FiLM conditioning is partially confounded by expert size (30M gated versus 17–18M cross-attention variants), which calls for a strictly size-matched comparison of the two interfaces.
• Diagnosing the noise regression. The only clear LIBERO-Plus regression for MambaAct is under observation noise (77.4% → 73.5%), and understanding whether this stems from the selective state-space mixing or the compressed context is an open question.
• Combining axes of acceleration more broadly. Only VLM-depth reduction is tested as a prefix-side complement; combining the Mamba expert with denoising-step compression or visual-token pruning is left open, and the authors argue the mechanisms are orthogonal.
• Broadening backbone transfer and real-world scope. The paper claims transfer to SmolVLA but the specific results are absent from the provided text; extending validation to more backbones, more tasks, and unseen object configurations beyond the two UR5 setups would test generality.

Target Audience

Robotics and embodied-AI researchers working on VLA policies; efficiency and deployment engineers who need lower inference latency for real-time control; practitioners of state-space models looking for a control application; and readers already familiar with flow matching and π0.5-style architectures who want a concrete, ablation-backed case study in replacing a repeatedly evaluated module rather than pruning an entire network.

Authors’ abstract

Vision-language-action (VLA) models have shown strong promise for generalist robotic manipulation, but their inference cost remains a major bottleneck for real-time control. This limitation is especially pronounced in flow-matching VLAs such as $π_{0.5}$, where the action expert is evaluated repeatedly during iterative denoising. We introduce MambaAct, an efficient Mamba action expert for flow-matching VLA policies. MambaAct retains the pretrained visual-language prefix encoder of $π_{0.5}$, while replacing its transformer-based action expert with a compact Mamba-based action expert. Within this expert, Mamba serves as the action mixer that models dependencies across the noisy continuous-action sequence. To preserve visual-language conditioning without repeated token-level cross-attention, MambaAct uses a lightweight Context Compressor to map Vision-language model (VLM) prefix tokens into a compact conditioning signal that gates and modulates each Mamba action block. MambaAct uses a $10\times$ smaller action expert and reduces flow-matching denoising latency by $64.8\%$ on LIBERO and $58.8\%$ in real-world deployment, while improving mean LIBERO success by $1.08$ percentage points across three independent fine-tuning seeds and maintaining gains on LIBERO-Plus and real-world benchmarks. Our code and model weights are publicly available at https://github.com/giakhuyendihoc/MambaAct.