MambaAct: Cut π0.5 VLA Denoising Latency by 64.8%

What problem does MambaAct solve?

A robot asked to place a bowl inside a drawer needs to interpret images, understand the instruction and generate motion. A Vision-Language-Action model, or VLA, learns these capabilities through a policy. However, interpreting a scene efficiently does not guarantee fast action generation. In a flow-matching policy, the action expert runs repeatedly to transform noise into a continuous sequence of actions.

MambaAct: Efficient Mamba Action Experts for Flow-Matching Vision-Language-Action Models, an October 2026 arXiv preprint by Khuyen Pham, Tran Nguyen Le, Edoardo Fazzari and Cesare Stefanini, targets this repeated computation. It preserves the pretrained visual-language backbone of π0.5, replaces the Transformer action expert with a compact Mamba3 expert and fine-tunes the resulting policy on downstream demonstrations. The official implementation is giakhuyendihoc/MambaAct.

The headline result is a 64.8% reduction in ten-step denoising latency on H100, from 160.26 ms to 56.40 ms. Complete policy-query latency drops from 228.6 ms to 107.6 ms, a reduction of 52.9%. These percentages describe different measurement boundaries. This tutorial follows the public implementation and reproduction documentation; the benchmark numbers are author-reported, rather than measurements independently produced by vnrobo.

Three concepts to understand first

An action chunk contains several future actions generated by one policy query. For example, the policy can predict ten steps, execute the first five and then observe the environment again. This reduces query frequency while allowing new observations to correct the remaining trajectory.

Flow matching learns a velocity field that transforms noisy actions into valid actions. Here, velocity means a direction in the model's action space; it is not automatically the physical velocity of a robot joint. During inference, the model starts from Gaussian noise and follows repeated Euler updates toward a structured action chunk.

Mamba is a sequence mixer based on a selective state-space model. Instead of explicitly forming every pairwise token interaction through self-attention, it propagates information through a selective state along an ordered sequence. In this project, that sequence consists of action tokens within a chunk. It should not automatically be interpreted as persistent memory across entire episodes.

For an introduction to how images, instructions and demonstrations enter a VLA policy, see our SmolVLA policy-training guide. Understanding the data interface makes the reproduction steps below much easier to debug.

Architecture: preserve the backbone, replace action generation

Source: Figure 2 in the original paper. The diagram shows the actual conditioning and action-generation pathways.

Compute the visual-language prefix once

Two RGB observations, a wrist view and a third-person view, enter the π0.5 PaliGemma backbone together with the instruction. SigLIP encodes the images, and Gemma processes the visual-language prefix. The main MambaAct variant executes all 18 VLM layers.

Preserving this backbone refers to its architecture and pretrained initialization. The published recipe still fine-tunes the backbone, including the vision encoder and executed VLM layers. It is not a frozen-backbone experiment in which only a small replacement head learns.

Compress the prefix into 64 context tokens

The Context Compressor first projects prefix representations to the action expert's hidden width. Visual tokens are pooled over contiguous segments using validity masks, and language tokens are compressed similarly. Four learned queries extract additional global summaries using attention.

Modality embeddings distinguish these groups. Following refinement, masked mean pooling produces a condition vector called z. Each action block has its own MLP that maps z to FiLM scale, shift and residual-gate parameters.

This distinction matters: the architecture does not remove every attention operation. The global queries still use attention during context construction. The saving comes from avoiding repeated token-level cross-attention inside the action expert at every denoising step.

Mix actions with four Mamba3 blocks

The main expert uses four SISO Mamba3 blocks with hidden width 512, state size 64, expansion factor 1, head dimension 64, chunk size 64 and FFN width 2048. Its core contains 29.7 million parameters, compared with 311.5 million in the original π0.5 expert. The complete VLA policy is much larger than this core.

Actions are padded to an internal dimension of 32 and projected into exactly H action tokens. The paper specifies no additional robot-state token in this action sequence. Flow time is embedded separately and injected into each block. Internal padding does not imply that the physical robot has 32 joints or 32 executable action channels.

The context and FiLM parameters are computed once per query and reused across all ten denoising steps. Action representations and time conditioning are recomputed as the noisy chunk changes. A final RMSNorm and linear head predict flow velocity.

What does the training objective predict?

The following is an algorithm sketch, not a runnable repository API. Let A0 be the clean demonstration action chunk and eps be Gaussian noise:

u ~ Beta(1.5, 1.0)
t = 0.001 + 0.999 * u
Xt = t * eps + (1 - t) * A0
velocity_target = eps - A0
loss = mean_squared_error(predicted_velocity, velocity_target)

Inference, 10 steps:
X(t - 0.1) = X(t) - 0.1 * predicted_velocity(X(t), t, observation)

The expert learns how the noisy sample should change, rather than directly classifying an action token. The sign is important: using A0 - eps with the displayed inference update would reverse the intended direction. Larger flow time corresponds to the noisy end of the trajectory.

Initialization also receives special treatment. Time projections and FiLM scale/shift outputs start at zero, while the initial residual multiplier is 0.2. This keeps the newly initialized Mamba branch active but initially attenuated. The authors describe 0.2 as an initialization choice, not a comprehensively optimized constant.

Step 1: prepare the policy environment

The Reproducibility Guide specifies Linux, Python 3.11, PyTorch 2.7.1, CUDA 12.8, bfloat16 and one NVIDIA H100 for the π0.5 experiments. Before starting, install a working NVIDIA driver and uv, and reserve storage for model weights, training data and caches.

# Run these commands on your experiment machine.
git clone https://github.com/giakhuyendihoc/MambaAct.git
cd MambaAct
bash scripts/setup.sh

export MAMBAACT_OUTPUT_ROOT="$PWD/artifacts"
export MAMBAACT_PI05_CHECKPOINT="$MAMBAACT_OUTPUT_ROOT/checkpoints/pi05_base_pytorch"

nvidia-smi
cd openpi
uv run python -c 'import torch; print(torch.__version__, torch.cuda.is_available())'
cd ..

The setup script follows openpi/uv.lock, installs the vendored Mamba implementation and applies the repository's Transformers replacements. Use the project's dedicated environment rather than modifying an environment serving another model. Installing an unrelated Mamba package with a similar name is not equivalent to using the evaluated Mamba3 source.

H100 is the reference hardware, not evidence that the same recipe fits every consumer GPU. Reducing batch size or introducing gradient accumulation changes the experiment. Such adaptations can be useful, but record them explicitly instead of presenting the result as an unchanged reproduction.

Step 2: convert the π0.5 initializer

From the MambaAct root directory, convert the public base checkpoint into the PyTorch layout expected by this implementation:

cd openpi
uv run python examples/convert_jax_model_to_pytorch.py \
  --checkpoint-dir gs://openpi-assets/checkpoints/pi05_base \
  --config-name pi05_libero_baseline \
  --output-path "$MAMBAACT_PI05_CHECKPOINT" \
  --precision bfloat16
cd ..

test -s "$MAMBAACT_PI05_CHECKPOINT/model.safetensors"

Both MambaAct and the controlled baseline must start from the same base checkpoint. Initializing only one model from an already LIBERO-fine-tuned policy would change the comparison. If the launcher reports a missing model.safetensors, fix conversion or the environment variable rather than bypassing the file check.

The LIBERO training configuration points to the converted LeRobot dataset physical-intelligence/libero. Training data and the evaluation simulator serve different purposes. Having the demonstration dataset available does not mean that LIBERO assets, task definitions and initial states are installed for rollout evaluation.

Step 3: install LIBERO separately

Follow the installation instructions in the official LIBERO repository using a separate environment. Its dependency stack differs from the Python 3.11 policy environment. Complete the benchmark's asset and path configuration before attempting a long training run.

The example below assumes that you have already installed the benchmark dependencies in a Conda environment named libero. Run it from the MambaAct root after installation:

export MAMBAACT_LIBERO_ROOT="$PWD/third_party/LIBERO"
export MAMBAACT_LIBERO_PYTHON="$(conda run -n libero python -c 'import sys; print(sys.executable)')"

"$MAMBAACT_LIBERO_PYTHON" -c 'from libero.libero import benchmark; print(sorted(benchmark.get_benchmark_dict()))'

Set the root variable to your actual LIBERO checkout location if you installed it elsewhere. The Python variable must identify an executable in the simulator environment, not the policy environment. This separation allows the OpenPI server and simulator client to use their own dependencies.

A successful import is an initial check, not proof that rendering works. Run the official offscreen environment example as well. On a headless machine, failures involving EGL or MuJoCo should be diagnosed before training, because an otherwise valid checkpoint cannot be evaluated without a working simulator.

Step 4: fine-tune both models with the same recipe

The training launcher computes normalization statistics before invoking the trainer. The published settings are:

bash scripts/train_libero.sh mambaact
bash scripts/train_libero.sh pi05

Run these sequentially on a shared single-GPU machine to avoid competing for accelerator memory. With the default output root, MambaAct writes its final checkpoint under artifacts/checkpoints/mambaact_libero/mambaact_l18/30000. The baseline uses pi05_libero_baseline/pi05_l18/30000 beneath the same checkpoint root.

The seed-42 results report approximately 13.78 training hours for full MambaAct on the reference system. Treat that as a planning reference rather than a runtime guarantee. Also remember that final model weights are not the same as a resumable training snapshot containing optimizer state. An interrupted run may require more recovery work than simply pointing the trainer at the last weights directory.

Step 5: run inference through the evaluator

The supplied evaluator is the most direct route to benchmark-compatible inference. It starts a policy server in the OpenPI environment, then invokes the simulator client using the LIBERO Python executable:

bash scripts/eval_libero.sh mambaact \
  "$MAMBAACT_OUTPUT_ROOT/checkpoints/mambaact_libero/mambaact_l18/30000"

bash scripts/eval_libero.sh pi05 \
  "$MAMBAACT_OUTPUT_ROOT/checkpoints/pi05_libero_baseline/pi05_l18/30000"

The standard configuration uses 224×224 images, action horizon 10, ten denoising steps and replanning after five environment actions. Spatial, Object, Goal and LIBERO-10 contribute ten tasks each. At 50 trials per task, one checkpoint requires 2,000 evaluation rollouts.

Read the suite JSON files and summary.md under artifacts/reports/libero/. A decreasing training loss does not demonstrate that the policy completes tasks in closed-loop simulation. Check that every suite finished before comparing the overall average.

For a standalone policy service, pass the configuration name, rather than the shorter model alias:

bash scripts/serve_checkpoint.sh mambaact_libero \
  "$MAMBAACT_OUTPUT_ROOT/checkpoints/mambaact_libero/mambaact_l18/30000"

The default server port is 8000. Any custom client must preserve image mapping, instruction handling, normalization and action representation. Starting this service alone does not connect a UR5 or translate actions into your robot controller's protocol.

Reading the LIBERO results correctly

The public results document reports independent training seeds 42, 17 and 27, all evaluated with seed 7. The uncertainty values below are population standard deviations across those training runs, not confidence intervals.

The average improvement is 1.08 percentage points. MambaAct's seed-42 run reaches 98.40%, but that best-run number must not replace the three-seed average of 97.80%. Its larger L-10 variation also shows why reporting only the strongest checkpoint is incomplete.

For multi-seed reproduction, change TRAIN_SEED and RUN_NAME during training and give evaluation a distinct RUN_ID. For example:

TRAIN_SEED=17 RUN_NAME=mambaact_trainseed17 \
  bash scripts/train_libero.sh mambaact

RUN_ID=mambaact_trainseed17_eval50_seed7 EVAL_SEED=7 TRIALS=50 \
  bash scripts/eval_libero.sh mambaact \
  "$MAMBAACT_OUTPUT_ROOT/checkpoints/mambaact_libero/mambaact_trainseed17/30000"

Repeat the documented procedure for both models and all three seeds. Calculate each run's average across the four suites before computing the mean and standard deviation across runs. Averaging the four suite standard deviations does not produce the standard deviation of the overall score. Three seeds provide useful descriptive evidence, but they do not by themselves establish a broad statistical-significance claim.

Profiling latency without mixing measurement boundaries

The compressor adds some prefix overhead, while the repeatedly evaluated pathway becomes substantially cheaper. Full-query timing also includes I/O, CPU–GPU transfer, serialization, communication and scheduling. The two component timings therefore do not describe every cost included in the complete query.

bash scripts/profile_checkpoint.sh mambaact \
  "$MAMBAACT_OUTPUT_ROOT/checkpoints/mambaact_libero/mambaact_l18/30000"

The profiler uses three warm-up queries and ten measured queries on a canonical observation. Keep hardware, software, denoising count and background load matched when comparing models. The ratio 228.6 / 107.6 gives approximately 2.12× query speedup. It does not directly specify the robot's control frequency, because chunk execution and controller scheduling are separate parts of the system.

Generalization and physical-robot evidence

Source: Figure 3 in the MambaAct paper. Wrist and third-person cameras supply the two observation views.

On LIBERO-Plus, evaluated without additional fine-tuning on that benchmark, MambaAct scores 79.7% overall against 78.7% for the matched baseline. The gain is not uniform: Noise drops from 77.4% to 73.5%, while Camera remains almost unchanged. A higher aggregate score should not be presented as improved robustness to every perturbation.

For the UR5 experiments, each task uses 112 demonstrations from 28 training configurations. Inference runs on an RTX 5090 with horizon 50, executing the first 30 actions at 30 Hz before replanning. Nut pouring achieves a final score of 96.9% against 91.3%; tube placing reaches 88.8% against 85.0%.

These physical-task results use normalized scoring with partial credit. An intermediate accomplishment can receive 0.5 points. They should not be described as binary full-task completion rates or compared directly with a benchmark that only awards complete success.

Ablations further qualify the architecture claim. The compact Transformer with gated conditioning reaches 74.22 ms E2E, faster than MambaAct's 107.61 ms, but scores only 85.7% against 98.4% in the seed-42 experiment. The central benefit is preserving task performance under aggressive expert compression. It is not a universal statement that every Mamba model is faster than every Transformer.

MambaAct also composes with prefix-depth reduction. Its L12 variant reaches 97.5% and 102.91 ms on LIBERO, offering another accuracy–latency trade-off. This supports exploring complementary optimizations, while leaving their combined performance to be measured rather than assumed.

Troubleshooting the first reproduction

If setup fails while compiling Mamba, inspect the CUDA toolkit, compiler and PyTorch compatibility before changing model code. The vendored implementation and locked environment are part of the reproduction contract. A successful import of a different Mamba package is not a substitute for building the required kernels.

If evaluation cannot start, first check the simulator Python path and model.safetensors, then inspect the report directory's server.log. The launcher uses port 8000 by default, so another service can prevent startup. Rendering errors belong to the simulator environment; model-loading errors belong to the policy environment. Keeping these logs separate makes the cause much easier to identify.

If a trained policy runs but performs poorly, verify camera order, instruction fields, action normalization and replanning settings before attributing the result to the sequence mixer. A correct architecture with an incorrect data interface can generate valid-shaped actions that are behaviorally wrong.

For your own robot, validate action units, delta-joint conventions, workspace limits and emergency-stop behavior before sending commands. LIBERO evaluation does not establish those deployment properties. MambaAct provides a concrete optimization path for flow-matching VLAs: make each repeated expert evaluation cheaper while retaining the pretrained backbone and the standard ten-step generation procedure.

Related Posts

• Fine-tuning SmolVLA and preparing policy training
• Training π0-FAST in the LeRobot ecosystem
• DexVerse: interpreting OpenVLA and π0.5 benchmarks