FlashVLA: tăng tốc π0.5/SmolVLA

FlashVLA là một hướng tiếp cận mới để làm cho các mô hình VLA flow-matching chạy nhanh hơn trong robot manipulation. Thay vì xem policy như một hàm nặng: nhận ảnh, language instruction, proprioception rồi phải hoàn tất toàn bộ quá trình denoising trước khi robot nhúc nhích, FlashVLA biến action decoding thành một pipeline streaming. Robot vừa thực thi chunk action hiện tại, vừa để GPU chuẩn bị chunk tiếp theo. Sau giai đoạn warm-up, mỗi forward pass không còn “giải” một trajectory từ đầu đến cuối, mà cập nhật một action buffer gồm nhiều slot đang ở các mức denoising khác nhau.

Điều này đặc biệt quan trọng với π0.5 và SmolVLA. Hai model này thuộc nhóm policy dùng vision-language backbone mạnh, nhưng inference có thể chậm khi cần nhiều denoising step, nhiều camera view, hoặc control loop có deadline chặt. Trong benchmark chính thức, FlashVLA giúp π0.5 trên LIBERO giảm time/step từ 53.8 ms xuống 22.1 ms với async d=1, và xuống 20.6 ms với d=2. Trên RoboTwin 2.0, success rate trung bình của π0.5 tăng từ 86.0% lên khoảng 90.6%. Với SmolVLA, FlashVLA giữ success rate gần như nguyên vẹn nhưng giảm per-step time từ 41.2 ms xuống 29.7 ms hoặc 28.7 ms tùy chế độ overlap.

Nếu bạn đã đọc LeRobot hands-on, OpenVLA deep dive, hoặc bài EXPO-FT cho π0.5, FlashVLA là mảnh ghép tiếp theo: nó không tập trung vào dataset hay reward, mà tập trung vào latency engineering cho policy VLA khi đưa vào robot thật.

Nguồn chính của bài này là paper FlashVLA: Streaming Action Decoding for Fast and Asynchronous VLA Inference trên arXiv, repo chính thức z-lab/flashvla, và model collection z-lab/flashvla trên Hugging Face. Tại thời điểm viết bài, repo đã có code training, benchmark latency, eval LIBERO, adapter RoboTwin 2.0, policy π0/π0.5/SmolVLA/LingBot-VLA, và checkpoint release cho π0.5 trên LIBERO/RoboTwin. Bài này đi theo góc nhìn beginner-friendly: hiểu ý tưởng trước, rồi mới tới lệnh cài đặt, training và inference.

Vấn đề FlashVLA giải quyết

Một VLA policy cho manipulation thường nhận ba nhóm input. Nhóm đầu tiên là quan sát hình ảnh: camera ngoài, wrist camera, hoặc nhiều view cùng lúc. Nhóm thứ hai là ngôn ngữ: “pick up the cup”, “put the block in the bowl”, “clean the table”. Nhóm thứ ba là robot state: joint position, gripper state, end-effector pose, hoặc proprioception đã normalize. Output là một chuỗi action cho vài bước điều khiển kế tiếp.

Trong policy kiểu flow matching, output action ban đầu được biểu diễn như noise. Model học một velocity field để biến noise thành action sạch theo thời gian giả t. Khi inference, ta chạy nhiều denoising step. Cách truyền thống là đồng bộ: model dự đoán xong một action chunk, robot thực thi chunk đó, rồi chờ model dự đoán chunk tiếp theo. Nếu inference mất 50 ms nhưng control step cần 20 ms, robot hoặc phải chạy chậm, hoặc phải chấp nhận jitter.

FlashVLA nhìn vấn đề theo hướng khác. Trong một episode, robot không cần toàn bộ tương lai dài vô hạn. Robot chỉ cần chunk gần nhất đủ tốt để thực thi, trong khi những chunk xa hơn có thể tiếp tục được denoise dần. Vì vậy FlashVLA giữ một buffer có N slot, mỗi slot là một chunk dài C action. Slot đầu gần nhất sắp được robot dùng. Slot sau ở xa hơn trong tương lai. Mỗi forward pass cập nhật toàn bộ buffer một bước denoising, sau đó khi buffer đã đầy, policy lấy slot đầu làm action to execute, shift buffer sang trái, và thêm noise mới vào slot cuối.

Nói đơn giản: baseline dùng nhiều forward pass để tạo một chunk; FlashVLA dùng nhiều slot trong buffer để pipeline nhiều chunk cùng lúc. Sau cold start, mỗi forward pass trả được một chunk đã trải qua đủ số denoising step. Đây là lý do paper gọi phương pháp là streaming action decoding.

Kiến trúc: shared observation, multi-buffer action

FlashVLA không yêu cầu bỏ π0.5 hay SmolVLA để xây model mới từ đầu. Repo triển khai các policy type như pi05-flashvla, smolvla-flashvla, pi0-flashvla, và lingbot-flashvla. Ý tưởng chung là giữ phần vision-language prefix, nhưng thay phần action suffix thành cấu trúc có nhiều buffer slot.

Trong π0.5, input prefix gồm image tokens, language tokens và state embedding. Action suffix gồm N * C action tokens. Với cấu hình LIBERO trong repo, num_buffer_slots: 4 và chunk_size: 10, tức buffer chứa 40 action token. Trong training, cùng một observation được dùng cho nhiều cấu hình buffer song song. Mỗi slot được gán một mức thời gian t khác nhau: slot gần hiện tại đã gần sạch hơn, slot xa hơn còn nhiều noise hơn. Model học velocity cho mọi slot trong một pass.

Trong inference, logic nằm ở hàm sample_actions và các nhánh _cold_start, _steady_streaming. Khi episode mới bắt đầu, buffer chưa đủ lịch sử denoising, nên policy trải qua cold start. Trong giai đoạn này, predict_action_chunk có thể trả None; wrapper sẽ dùng action giữ pose hiện tại hoặc action lạnh phù hợp với config. Khi step_counter đạt N - 1, buffer đã đủ “chín”. Từ đó trở đi, mỗi lần gọi policy sẽ:

1. Build prefix từ ảnh, language và state.
2. Denoise toàn bộ buffer một bước với dt = 1 / N.
3. Lấy slot đầu buffer[:, :C, :] làm chunk để thực thi.
4. Shift các slot còn lại sang trái.
5. Sinh noise mới cho slot cuối.

Hai chi tiết kỹ thuật đáng chú ý. Thứ nhất, FlashVLA dùng block-causal attention cho action tokens: token trong cùng một slot được attention đầy đủ, nhưng các slot chỉ nhìn theo hướng causal để không rò rỉ tương lai. Thứ hai, với SmolVLA, repo có tuỳ chọn use_adarms_time_cond: true, dùng time conditioning qua AdaRMS-style modulation cho action expert. Điều này giúp model biết token nào đang ở mức denoising nào mà không cần làm rối prefix observation.

Cài đặt môi trường

Repo FlashVLA dùng Python 3.12, PyTorch CUDA 12.8, TorchVision và TorchCodec. Lệnh cài đặt cơ bản:

git clone https://github.com/z-lab/flashvla.git
cd flashvla
conda env create -f environment.yml
conda activate flashvla

File environment.yml cài torch==2.9.1+cu128, torchvision==0.24.1+cu128, torchcodec==0.9.1+cu128, rồi install package local bằng -e .. Nếu bạn dùng workstation RTX 4090/5090 hoặc server CUDA mới, đây là đường đi sạch nhất. Nếu dùng Jetson, laptop không có CUDA 12.8, hoặc cluster có module CUDA riêng, đừng copy môi trường này một cách máy móc. Hãy giữ version logic của FlashVLA nhưng chỉnh PyTorch wheel theo driver thực tế.

Để chạy LIBERO, repo yêu cầu cài thêm simulator dependency:

conda install -y cxx-compiler make
CMAKE_POLICY_VERSION_MINIMUM=3.5 pip install --no-build-isolation -e ".[libero]"
export MUJOCO_GL=egl
export __EGL_VENDOR_LIBRARY_FILENAMES=/usr/share/glvnd/egl_vendor.d/10_nvidia.json

--no-build-isolation là chi tiết nhỏ nhưng quan trọng. Một số dependency kéo cmake qua pip dưới dạng Python entry script; khi build isolation bật, build env có thể không thấy compiler/cmake đúng cách. Với headless server, MUJOCO_GL=egl giúp MuJoCo render qua GPU mà không cần desktop.

Training π0.5 FlashVLA trên LIBERO

Repo cung cấp config sẵn:

bash train/train.sh train/configs/pi05/libero/pi05_flashvla.yaml

Các tham số quan trọng trong config:

policy:
  type: pi05-flashvla
  pretrained_path: lerobot/pi05_base
  dtype: bfloat16
  device: cuda
  state_cond: false
  num_buffer_slots: 4
  chunk_size: 10
  freeze_vision_encoder: false
  normalization_mapping:
    VISUAL: IDENTITY
    STATE: MEAN_STD
    ACTION: MEAN_STD

dataset:
  repo_id: HuggingfaceVLA/libero
  video_backend: torchcodec

batch_size: 8
grad_accum_steps: 4
steps: 50000

Với beginner, nên đọc config theo bốn lớp. Lớp model cho biết backbone là lerobot/pi05_base, policy type là pi05-flashvla, chạy bfloat16, và không freeze vision encoder. Lớp buffer cho biết policy học streaming với 4 slot, mỗi slot dài 10 action. Lớp normalization giữ ảnh ở identity, nhưng state/action dùng mean-std, phù hợp với robot state có scale khác nhau. Lớp optimization dùng AdamW, learning rate đỉnh 1e-4, warm-up 1000 step, cosine decay tới 2.5e-6.

Nếu GPU memory hạn chế, giảm batch_size trước, giữ grad_accum_steps để effective batch không đổi. Nếu muốn fine-tune nhanh để debug pipeline, giảm steps xuống vài trăm hoặc vài nghìn, nhưng đừng dùng checkpoint đó để kết luận benchmark. FlashVLA cần học pattern multi-buffer; nếu train quá ngắn, lỗi thường không phải ở thuật toán mà ở checkpoint chưa hội tụ.

Training SmolVLA FlashVLA

SmolVLA có config riêng:

bash train/train.sh train/configs/smolvla/libero/smolvla_flashvla.yaml

Điểm khác biệt:

policy:
  type: smolvla-flashvla
  pretrained_path: lerobot/smolvla_base
  num_buffer_slots: 5
  chunk_size: 10
  use_adarms_time_cond: true
  attention_mode: cross_attn

batch_size: 64
steps: 20000

SmolVLA nhỏ hơn π0.5 nên config dùng batch lớn hơn và số step thấp hơn. num_buffer_slots: 5 nghĩa là buffer có 50 action token. use_adarms_time_cond giúp đưa thời gian denoising vào expert branch bằng modulation, còn attention_mode: cross_attn giữ kiến trúc phù hợp với SmolVLA. Kết quả cross-architecture trong README cho thấy FlashVLA không chỉ là hack riêng cho π0.5: với SmolVLA, success rate trung bình trên benchmark được báo cáo là 80.1% ở baseline và 80.1% với FlashVLA sync, trong khi latency inference giảm từ 19.7 ms xuống 10.1 ms.

Inference đồng bộ và bất đồng bộ

Khi deploy, phần đáng hiểu nhất là AsyncStreamingActionManager. Wrapper này giữ current_chunk, next_chunk, chunk_index và overlap_steps. Nếu overlap_steps = 0, robot chạy sync: dùng hết chunk hiện tại rồi mới block để lấy chunk mới. Nếu overlap_steps = 1 hoặc 2, wrapper launch inference sớm trước khi chunk hiện tại kết thúc. GPU tính chunk kế tiếp trong khi CPU/simulator/robot vẫn đang execute action cũ.

Pseudocode control loop:

manager = AsyncStreamingActionManager(policy, overlap_steps=1)
manager.reset()

obs = env.reset()
while not done:
    if manager.needs_observation():
        batch = preprocessor(obs).to(policy.device)
        action = manager.act(batch)
    else:
        action = manager.pop_cached_action()

    action = postprocessor({"action": action})["action"]
    obs, reward, done, info = env.step(action.numpy())

needs_observation() là tối ưu thực tế: nếu action đã nằm trong CPU cache, simulator có thể bỏ qua render ảnh mới cho policy ở step đó. Với robot thật, điều này tương đương việc control thread không nhất thiết phải đợi camera pipeline ở mọi tick. Tuy nhiên, bạn vẫn cần quyết định policy update rate và control rate rõ ràng. FlashVLA không thay thế low-level controller; nó chỉ làm cho high-level action chunk đến đều hơn.

Một chi tiết dễ bị bỏ qua là warm-up CUDA graph. Code có hàm warmup(processed_obs) để chạy trước num_buffer_slots + 1 inference, capture graph rồi reset state. Nếu không warm-up, lần đầu policy chuyển sang _steady_streaming có thể trigger torch.compile và CUDA graph capture, gây block 10-30 giây. Trong benchmark latency, điều này sẽ làm số liệu vô nghĩa; trong robot thật, nó có thể khiến episode đầu bị đứng.

Chạy evaluation

Với LIBERO, repo cung cấp script:

bash sim_eval/libero/eval.sh
GPUS="0 1 2 3" bash sim_eval/libero/eval.sh

Để debug một suite:

python sim_eval/libero/eval.py --policy.path=z-lab/flashvla-pi05-libero \
  --env.type=libero --env.task=libero_spatial --eval.n_episodes=50 \
  --policy.n_action_steps=5 --inference_overlap_steps=1 \
  --policy.compile_model=true

Với RoboTwin 2.0, môi trường phức tạp hơn vì SAPIEN/RoboTwin có dependency riêng. FlashVLA chạy policy server trong env flashvla, còn RoboTwin chạy client trong env simulator. Sau khi copy adapter vào tree RoboTwin, bạn chạy hai terminal:

# terminal 1 - flashvla env
bash eval_server.sh /path/to/flashvla_robotwin_ckpt 9999 0 current_state 1 16 true

# terminal 2 - RoboTwin env
EVAL_STEP_LIM_OFFSET=48 ROBOTWIN_VENV=/path/to/robotwin_venv \
  bash eval_client.sh beat_block_hammer demo_clean

current_state quan trọng với RoboTwin vì action là absolute qpos; zero action có thể làm arm crash. EVAL_STEP_LIM_OFFSET=48 bù cold start cho cấu hình 4 buffer slot, 16 action mỗi call: (4 - 1) * 16 = 48.

Kết quả benchmark nên đọc như thế nào?

Trên LIBERO, π0.5 baseline đạt trung bình 96.9% success với 53.8 ms mỗi step. FlashVLA sync d=0 đạt 97.9%. FlashVLA async d=1 đạt 97.8% và 22.1 ms mỗi step, tức nhanh hơn 2.43 lần. Với d=2, success đạt 98.3% và time/step 20.6 ms, nhanh hơn 2.62 lần. Điều đáng chú ý là tốc độ tăng mà success không giảm; thậm chí Long suite tăng từ 92.4% lên 97.4% ở d=2.

Trên RoboTwin 2.0, baseline π0.5 đạt 86.0% trung bình trên 50 task. FlashVLA sync đạt 90.5%, async d=1 và d=2 cùng đạt khoảng 90.6%. Điều này gợi ý rằng multi-buffer training không chỉ tối ưu latency mà còn có thể regularize policy, vì model học dự đoán nhiều tương lai gần ở các mức noise khác nhau.

Benchmark latency riêng trên RTX 4090/5090 cũng đáng chú ý. Với hai camera view trên RTX 4090, π0.5 là 45.8 ms, Realtime-VLA là 29.2 ms, FlashVLA là 26.7 ms. Với ba view trên RTX 5090, π0.5 là 44.8 ms, Realtime-VLA là 34.2 ms, FlashVLA là 27.1 ms. Paper nhấn mạnh π0.5 và FlashVLA đều dùng cùng CUDA Graph và kernel-fusion optimization trong phép đo này, nên phần chênh lệch đến từ cách decode action, không phải chỉ từ trick implementation.

Khi nào nên dùng FlashVLA?

FlashVLA hợp nhất khi bạn có VLA flow-matching đang bị nghẽn ở inference latency: robot cần loop nhanh hơn, nhiều camera làm forward pass nặng hơn, hoặc bạn muốn chạy async để action không bị khựng ở ranh giới chunk. Nó cũng hợp khi bạn đã có pipeline LeRobot/π0.5/SmolVLA và muốn thử checkpoint/config có sẵn thay vì viết scheduler mới.

Không nên kỳ vọng FlashVLA giải quyết mọi vấn đề. Nếu dataset demonstration kém, camera calibration sai, action normalization lệch, hoặc low-level controller không ổn định, streaming decoding sẽ chỉ làm policy sai nhanh hơn. Trước khi deploy robot thật, hãy test theo thứ tự: latency benchmark offline, LIBERO/RoboTwin hoặc sim riêng, replay action trong safety cage, rồi mới cho robot tương tác object thật.

Một checklist thực dụng:

Kết luận

FlashVLA là một bài học hay cho robotics engineer: đôi khi cải tiến lớn không nằm ở việc tăng model size, mà ở cách tổ chức inference quanh constraint của robot. Với streaming action decoding, π0.5 và SmolVLA có thể tận dụng cùng observation prefix, duy trì action buffer nhiều slot, overlap GPU inference với robot execution, và giảm latency đáng kể mà không hy sinh success rate.

Nếu bạn đang xây VLA manipulation stack, hãy xem FlashVLA như một lớp deployment/training recipe cho action decoder. Bắt đầu bằng LIBERO để hiểu buffer và cold start, sau đó chuyển sang dataset LeRobot của bạn. Khi mọi thứ ổn trong sim, bật async overlap từng bước: d=0, rồi d=1, rồi mới d=2. Đừng tối ưu tốc độ trước khi đo stability; với robot thật, một chunk nhanh nhưng sai vẫn là một chunk nguy hiểm.

Bài viết liên quan

• FineVLA: hướng dẫn VLA hai tay
• DexVerse: Benchmark OpenVLA và π0.5
• EXPO-FT: fine-tune π0.5 bằng online RL