\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{Residual aggregation and attention residuals.} Standard residual networks update $x_t=x_{t-1}+y_t$, so recursively $x_t=y_0+y_1+\cdots+y_t$. Equivalently, the next layer receives a uniform aggregation of all previous residual increments. A natural generalization is \begin{equation} y_{t+1}=f_{t+1}\left(\sum_{s=0}^t a_{t+1,s}y_s\right),\qquad a_{t+1,s}\ge0,\quad \sum_{s=0}^t a_{t+1,s}=1. \end{equation} AttnRes instantiates the weights with an attention distribution over previous residual outputs, \begin{equation} a_{t+1,s}\propto\exp\left(w_{t+1}^{\top}\mathrm{RMSNorm}(y_s)\right). \end{equation} This preserves the residual pathway while allowing each layer to select useful earlier representations instead of uniformly accumulating all of them~\citep{kimi2026attention}. We use this mechanism in the policy transformer; our ablations suggest that applying AttnRes too broadly inside the vision stack is not automatically beneficial. \section{Method} \label{sec:method} \subsection{Policy formulation} \label{sec:policy_formulation} We consider simulated imitation learning in RoboIMI. At each control step, the policy receives visual observations, robot state, and a task context, and predicts an action chunk of length \texttt{exec}. The main iMF-AttnRes socket and object-transfer policies use ph32/exec16 or related horizon-execution settings. The rollout evaluator records cumulative reward, median reward, mean maximum reward, nonzero-reward episode count, task-specific success-like threshold count, inference FPS, control FPS, and inference latency. The policy contains three conceptual parts. First, an observation encoder maps visual and state inputs into policy tokens. Second, a transformer policy core uses AttnRes in place of fixed residual accumulation. Third, the action generator uses iMF to predict an average flow over action noise-to-data paths, enabling one-to-few inference evaluations. In the strongest socket variants, the generator uses infer2 or infer3; in the strongest object-transfer run, it uses infer1. \subsection{Improved Mean Flow action generation} \label{sec:imf_action_generation} Classical flow matching trains an instantaneous velocity field. 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} If a policy takes very few integration or denoising steps, using only this local instantaneous velocity can cause large path errors when the learned trajectory is curved. Mean Flow addresses this by training an average velocity over an interval. For two times $r