\documentclass{article} \usepackage{corl_2026} \usepackage[T1]{fontenc} \usepackage{microtype} \usepackage{booktabs} \usepackage{multirow} \usepackage{graphicx} \usepackage{amsmath,amssymb} \usepackage{enumitem} \usepackage{tcolorbox} \usepackage[capitalize]{cleveref} \usepackage{url} \title{iMF-AttnRes: Fast Mean-Flow Action Generation with Attention Residuals for Simulated Robot Manipulation} \author{Anonymous Authors} \begin{document} \maketitle \begin{abstract} Diffusion-style vision-language-action policies provide expressive action distributions, but iterative inference can be too slow for closed-loop manipulation. We study whether improved Mean Flow (iMF), originally motivated by fast-forward generative modeling, can be adapted to action generation so that a policy uses one to a few flow evaluations rather than long denoising chains. We combine iMF with Attention Residuals (AttnRes), which replace fixed residual accumulation in a deep policy transformer with learned depth-wise aggregation over previous layer outputs. In RoboIMI socket peg insertion, the best iMF-AttnRes policy reaches an average reward of 1513.56 and median reward of 1901.5 with 15.122 ms average inference time, compared with diffusion-style infer100 baselines at 861.53/338.291 ms and 1019.39/397.912 ms. In RoboIMI object transfer, the best iMF-AttnRes policy reaches average reward 526.22 and 44/100 success-like episodes, compared with a native diffusion-policy baseline at 319.2 and 29/100. These results suggest that average-flow action generation is a promising route to low-latency robot policies, while also revealing sensitivity to horizon choices, vision-token design, and simulation-to-real validation. \end{abstract} \keywords{Robot learning, imitation learning, flow matching, vision-language-action models} \section{Introduction} \label{sec:introduction} Learning-based robot manipulation policies increasingly use expressive generative action models. Diffusion Policy showed that action diffusion can be an effective visuomotor policy class for manipulation~\citep{black2024pi0,chi2023diffusion}, inheriting the representational flexibility of denoising diffusion models~\citep{ho2020denoising} and transformer-based diffusion backbones~\citep{peebles2023scalable}. However, this expressivity often comes with a deployment cost: the policy must be evaluated repeatedly during sampling. In the RoboIMI socket peg experiments studied here, two diffusion-style VLA baselines use infer100 sampling and require 338.291 ms and 397.912 ms per inference on their recorded rollouts. Such latencies are undesirable when the robot must repeatedly close the perception-action loop. This paper asks a direct question: can a robot action generator retain the quality benefits of generative policies while reducing inference to one or a few evaluations? We explore this question by adapting improved Mean Flow (iMF) to VLA-style imitation learning. Flow matching and rectified flow formulate generation through continuous vector fields~\citep{lipman2023flow,liu2022flow}. Mean Flow reframes generation around average velocity, enabling one-step generation~\citep{geng2025mean}; iMF further addresses instability caused by substituting conditional velocities into nonlinear mean-flow relations~\citep{geng2025improved}. For robot control, this distinction matters because fewer evaluations directly increase control responsiveness. We pair iMF with Attention Residuals (AttnRes). Residual connections are essential for deep visual and transformer networks~\citep{he2016deep,vaswani2017attention}, but standard residual accumulation adds all layer outputs with fixed unit weight. The motivation in our notes is to rewrite residual networks as a sum over layer increments, then replace the uniform sum with a learned softmax aggregation. AttnRes implements this view by letting each layer attend over preceding residual outputs~\citep{kimi2026attention}. We use this mechanism in the policy transformer to reduce depth-wise feature dilution while keeping the action generator fast. Our contributions are: \begin{enumerate}[leftmargin=*] \item We formulate an iMF-AttnRes VLA policy for simulated robot imitation learning, combining one-to-few-step average-flow action generation with learned residual aggregation. \item We provide 100-rollout evaluations on two RoboIMI manipulation environments, socket peg insertion and object transfer, comparing against diffusion-style VLA baselines, a native Diffusion Policy baseline, ACT-style action chunking, and SmolVLA-style compact VLA inference~\citep{zhao2023learning,shukor2025smolvla}. \item We report both task quality and deployment speed. On socket insertion, the strongest iMF-AttnRes run improves average reward over the ph16 diffusion-style baseline by 1.76$\times$ and inference FPS by 22.97$\times$; on object transfer, the strongest iMF-AttnRes run improves average reward over the native diffusion-policy baseline by 1.65$\times$. \end{enumerate} We deliberately keep the claims scoped to simulation: hardware varies across runs, and real-robot transfer remains future work. \section{Related Work} \label{sec:related_work} \paragraph{Diffusion and flow policies for robot action generation.} Diffusion Policy introduced action diffusion as a visuomotor policy learning framework~\citep{black2024pi0,chi2023diffusion}, building on denoising diffusion objectives~\citep{ho2020denoising}. Diffusion transformers further show that transformer backbones can scale generative modeling~\citep{peebles2023scalable}. These models are expressive because they iteratively refine samples, but iterative refinement is also a latency source. Flow matching provides an alternative continuous generative formulation by learning vector fields that transport probability paths~\citep{lipman2023flow}. Rectified flow emphasizes straighter transport paths and faster sampling~\citep{liu2022flow}. Our method follows this fast-generation lineage but uses the improved mean-flow relation to target one-to-few-step action generation. \paragraph{Vision-language-action models and action chunking.} Transformer robot policies such as RT-1 and RT-2 show how large sequence models can map observations and task context to robot actions~\citep{brohan2022rt,brohan2023rt}. OpenVLA studies open VLA modeling at scale~\citep{kim2024openvla}, while compact or efficient VLA models such as SmolVLA and FAST reduce deployment cost through smaller backbones or efficient action tokenization~\citep{shukor2025smolvla,pertsch2025fast}. Action Chunking with Transformers (ACT) predicts chunks of future actions to improve temporal consistency and execution efficiency~\citep{zhao2023learning}. The iMF-AttnRes policy is complementary: it retains chunked execution but changes the generative action sampler so that each chunk can be produced with one to a few evaluations. \paragraph{Mean-flow training for one-step generation.} Let $x_t$ denote a sample at time $t$ and let the instantaneous flow satisfy \begin{equation} \frac{d x_t}{dt} = v(x_t,t). \end{equation} For two times $r 4}; for sim\_transfer, it is \texttt{max\_reward >= 4}. Hardware differs across runs, so speed comparisons are practical deployment measurements rather than perfectly normalized throughput benchmarks. \subsection{Socket peg insertion} \label{sec:socket_results} \Cref{tab:socket_results} summarizes the socket peg insertion results. The best iMF-AttnRes policy by average reward is the infer3 model, with average reward 1513.56 and median reward 1901.5. It exceeds both diffusion-style infer100 baselines: the ph16 baseline obtains average reward 861.53 and median reward 668.0, while the ph32/exec16 baseline obtains average reward 1019.39 and median reward 804.0. The latency gap is large. The iMF-AttnRes infer2 and infer3 policies require 14.409 ms and 15.122 ms per inference, while the two infer100 baselines require 338.291 ms and 397.912 ms. \begin{table*}[t] \centering \small \setlength{\tabcolsep}{3.5pt} \caption{Socket peg insertion over 100 rollouts. Success-like uses the recorded \texttt{max\_reward > 4} count.} \label{tab:socket_results} \begin{tabular}{llrrrrrr} \toprule Method & Setting & Avg. reward & Median & Avg. max & Nonzero & Success-like & Latency ms \\ \midrule iMF-AttnRes & infer1, ph32, exec16, 50k & 1275.22 & 1490.5 & 3.12 & 84/100 & 0/100 & 16.186 \\ Diffusion-style VLA & infer100, ph16, 150k & 861.53 & 668.0 & 2.58 & 97/100 & 2/100 & 338.291 \\ Diffusion-style VLA & infer100, ph32, exec16, 150k & 1019.39 & 804.0 & 2.47 & 92/100 & 4/100 & 397.912 \\ iMF-AttnRes & infer2, ph32, exec16, 150k & 1472.62 & 1825.0 & 3.29 & 83/100 & 5/100 & 14.409 \\ iMF-AttnRes & infer3, ph32, exec16, 150k & \textbf{1513.56} & \textbf{1901.5} & 3.28 & 83/100 & 3/100 & 15.122 \\ ACT & action chunking & 289.60 & 13.0 & 1.29 & 55/100 & 1/100 & 100.212 \\ SmolVLA & 100k & 466.16 & 106.0 & 1.72 & 89/100 & 0/100 & \textbf{2.614} \\ \bottomrule \end{tabular} \end{table*} The nonzero-reward count reveals an important caveat. The ph16 diffusion-style baseline has 97/100 nonzero episodes, higher than the iMF-AttnRes infer2 and infer3 policies at 83/100. Yet its average and median rewards are much lower. Thus, in this task, nonzero contact or partial progress is not sufficient; the iMF-AttnRes policies more often produce high-reward trajectories when they engage the task successfully. SmolVLA is the fastest method in raw latency and control FPS, but its reward is lower than iMF-AttnRes. ACT underperforms both iMF-AttnRes and the diffusion-style VLA baselines in this setup. \begin{figure}[t] \centering \begin{tcolorbox}[width=0.96\linewidth,colback=white,colframe=black!35,title={Placeholder for \texttt{fig-socket-reward-latency}}] \small Paper Banana prompt: create a 4:3 publication-quality reward-latency comparison for socket peg insertion. Plot methods from the socket table with average reward on the y-axis and average inference time or inverse latency on the x-axis. Highlight iMF-AttnRes infer2/infer3 at 1472.62/1513.56 average reward and 14.409/15.122 ms, diffusion-style infer100 baselines at 861.53/1019.39 average reward and 338.291/397.912 ms, ACT at 289.6 and 100.2116 ms, and SmolVLA at 466.16 and 2.614 ms. Use log-scale latency if helpful and label the speed-quality Pareto frontier. \end{tcolorbox} \caption{Planned socket reward-latency figure. The current draft contains the Paper Banana prompt placeholder instead of a rendered image.} \label{fig:socket_reward_latency} \end{figure} \subsection{Object transfer and ablations} \label{sec:sim_transfer_results} \Cref{tab:sim_transfer_results} reports the object-transfer results. The strongest iMF-AttnRes run reaches average reward 526.22 and 44/100 success-like episodes, exceeding the native diffusion-policy DiT/DDPM/ResNet baseline at average reward 319.2 and 29/100 success-like episodes. This is the clearest sim\_transfer evidence that iMF-AttnRes can improve both reward and success-like threshold count. \begin{table*}[t] \centering \small \setlength{\tabcolsep}{3.2pt} \caption{Object transfer / sim\_transfer over 100 rollouts. Success-like uses the recorded \texttt{max\_reward >= 4} count.} \label{tab:sim_transfer_results} \begin{tabular}{llrrrrrr} \toprule Method & Setting & Avg. reward & Median & Avg. max & Nonzero & Success-like & Inf. FPS \\ \midrule Native Diffusion Policy & DiT + DDPM + ResNet & 319.20 & 6.0 & 1.68 & 55/100 & 29/100 & 32.09 \\ Diffusion-style VLA & emb384, layer18 & 233.52 & 0.0 & 1.12 & 39/100 & 17/100 & 1.859 \\ iMF multi-token ResNet18 & step34999, ph16, exec08 & 260.66 & 0.0 & 1.12 & 33/100 & 23/100 & \textbf{441.80} \\ iMF full AttnRes vision & ph16, exec08, 50k & 228.42 & 0.0 & 0.94 & 31/100 & 16/100 & 55.997 \\ iMF-AttnRes DiT only & ph16, exec16, 50k & 240.64 & 0.0 & 1.02 & 32/100 & 19/100 & 137.996 \\ iMF-AttnRes DiT only & ph32, exec08, 50k & 163.28 & 0.0 & 0.76 & 26/100 & 12/100 & 85.638 \\ iMF-AttnRes DiT only & ph32, exec32, 50k & 260.72 & 0.0 & 1.18 & 38/100 & 21/100 & 9.557 \\ iMF-AttnRes DiT only & ph16, exec08, 50k & 229.56 & 0.0 & 1.18 & 41/100 & 18/100 & 69.984 \\ iMF-AttnRes DiT only & ph08, exec08, 50k & 237.02 & 0.0 & 1.32 & 48/100 & 18/100 & 69.192 \\ iMF-AttnRes DiT only & ph32, exec16, 50k & 49.88 & 0.0 & 0.32 & 13/100 & 4/100 & 138.903 \\ iMF-AttnRes & infer1, ph32, exec16, 50k & \textbf{526.22} & \textbf{86.0} & \textbf{2.14} & \textbf{63/100} & \textbf{44/100} & 8.911 \\ \bottomrule \end{tabular} \end{table*} The ablations are mixed and therefore useful. The ResNet18 multi-token iMF model is extremely fast at 441.80 inference FPS, but its average reward is 260.66, below the native diffusion-policy baseline. The full-AttnRes vision model reaches 228.42 average reward, suggesting that replacing residuals inside the vision encoder is not automatically helpful. Horizon and execution length are also sensitive: the ph32/exec16 DiT-only iMF variant reaches only 49.88 average reward despite high inference FPS. The strongest object-transfer run therefore combines iMF-AttnRes with the right horizon and execution setting rather than showing a universally dominant architectural change. \begin{figure}[t] \centering \begin{tcolorbox}[width=0.96\linewidth,colback=white,colframe=black!35,title={Placeholder for \texttt{fig-sim-transfer-ablation}}] \small Paper Banana prompt: create a 4:3 grouped bar chart for sim\_transfer ablations. Show average reward and success-like episode count for native diffusion policy, best sim-transfer iMF-AttnRes infer1, ResNet18 multi-token iMF, full-AttnRes vision, and selected horizon/execution variants. Emphasize that best iMF-AttnRes reaches 526.22 average reward and 44/100 success-like episodes versus native diffusion policy at 319.2 and 29/100, while several ablations underperform. \end{tcolorbox} \caption{Planned sim\_transfer ablation figure. The prompt is retained for later Paper Banana rendering.} \label{fig:sim_transfer_ablation} \end{figure} \subsection{Derived speed-quality comparisons} \label{sec:derived_comparisons} \Cref{tab:derived_comparisons} lists the derived comparisons used to summarize the tradeoff. On socket insertion, iMF-AttnRes infer3 improves average reward by 1.76$\times$ over the ph16 diffusion-style baseline and by 1.48$\times$ over the ph32 diffusion-style baseline. The same infer3 run improves inference FPS by 22.97$\times$ over the ph16 diffusion-style baseline. The infer2 run is 23.48$\times$ lower latency than the ph16 diffusion-style baseline and 28.45$\times$ higher inference FPS than the ph32 diffusion-style baseline. On object transfer, best iMF-AttnRes improves average reward by 1.65$\times$ and success-like episodes by +15 over native Diffusion Policy. \begin{table}[t] \centering \small \caption{Derived comparisons from the 100-rollout logs.} \label{tab:derived_comparisons} \begin{tabular}{llr} \toprule Comparison & Metric & Result \\ \midrule Socket iMF infer3 vs ph16 diffusion & Avg. reward & 1.76$\times$ \\ Socket iMF infer3 vs ph32 diffusion & Avg. reward & 1.48$\times$ \\ Socket iMF infer3 vs ACT & Avg. reward & 5.23$\times$ \\ Socket iMF infer3 vs SmolVLA & Avg. reward & 3.25$\times$ \\ Socket iMF infer2 vs ph16 diffusion & Latency reduction & 23.48$\times$ \\ Socket iMF infer3 vs ph16 diffusion & Inference FPS & 22.97$\times$ \\ Socket iMF infer2 vs ph32 diffusion & Inference FPS & 28.45$\times$ \\ Sim transfer iMF vs native diffusion & Avg. reward & 1.65$\times$ \\ Sim transfer iMF vs native diffusion & Success-like episodes & +15 \\ ResNet18 iMF vs layer18 baseline & Inference FPS & 237.65$\times$ \\ \bottomrule \end{tabular} \end{table} \section{Limitations} \label{sec:limitations} All experiments are simulation rollouts. The data do not establish real-robot transfer, robustness to sensing changes, or safety under hardware execution. CoRL submissions should provide convincing robotics evidence; the present draft should therefore be viewed as a simulation-first manuscript that still needs real-robot or stronger transfer evidence before formal submission. The speed measurements are also not perfectly hardware-normalized. Runs were collected on RTX 5880 Ada, L20, and RTX 5090 machines as encoded by their run names. Large latency differences between infer100 diffusion-style policies and infer1--3 iMF-AttnRes policies are meaningful because they reflect algorithmic sampling cost, but exact FPS ratios may include hardware and implementation effects. Future experiments should rerun all main policies on the same GPU, with identical rollout parallelism and identical observation preprocessing. Finally, iMF-AttnRes is sensitive to design choices. In sim\_transfer, several iMF variants underperform the native diffusion-policy baseline, and full AttnRes in the vision encoder performs worse than more conservative policy-transformer AttnRes. This suggests that average-flow training is not a plug-in guarantee; horizon, execution length, visual tokenization, and residual placement must be tuned carefully. \section{Conclusion} \label{sec:conclusion} We presented a first simulation study of iMF-AttnRes for fast robot action generation. The method adapts improved Mean Flow to VLA imitation learning and uses Attention Residuals to replace fixed residual accumulation in the policy transformer. In socket peg insertion, iMF-AttnRes achieves higher average and median reward than diffusion-style infer100 baselines while reducing inference latency from hundreds of milliseconds to roughly 15 ms. In object transfer, the best iMF-AttnRes run improves average reward and success-like episode count over the native diffusion-policy baseline. At the same time, ablations show that the method is sensitive to horizon/execution settings and residual placement. The next step is controlled, hardware-normalized evaluation with real-robot transfer. \clearpage \bibliography{refs} \end{document}