Replacing the final layer of a flow policy with a Gaussian doubles robot speed

During inference, a flow policy has to integrate the action chunk out of noise, usually requiring multiple network passes. Youssef Saied et al. examined the action distributions of Xiaomi-Robotics-1 and GR00T N1.5 and found that continuous action chunks are very close to a full-covariance Gaussian. For Xiaomi's 16-action chunk, the relative error is 0.8% in translation and 1.8% in rotation. This number is the median relative sliced Wasserstein error over 300 observations, computed with full covariance. In the figure, the orange control group is the error when the samples themselves come from a Gaussian, and Xiaomi's curve stays close to it.

GAR's approach is very economical: keep the pretrained network, freeze it, and only replace the final layer with a Gaussian readout. The readout predicts the Gaussian mean and covariance, as well as the binary command, from the frozen features. Fitting uses teacher moments from 6000 contexts; after caching features and moments, Xiaomi's readout trains for about 2 minutes and uses only 0.02% of the available states. At inference, a single action-network pass samples directly, and the likelihood can also be computed in closed form.

The mean is hidden in the initial velocity. At t=0, the input noise plus the velocity is approximately equal to the mean action. The variance is not visible at t=0; two Gaussians with different variances can have the same initial velocity and only diverge later. Therefore GAR does not guess the variance from the mean at t=0; instead, it trains a small head to predict it from the features before the final layer.

On 50 RoboCasa365 tasks, with 2500 episodes per policy. The Xiaomi teacher has a 58.24% success rate; GAR has 57.64% and is 1.91 times faster. The GR00T teacher has 23.08%; GAR has 21.92% and is 1.57 times faster. GAR uses one action-network pass, replacing Xiaomi's five flow steps and GR00T's four. Speed is measured for the entire policy call, including the vision-language backbone, with batch size 1.

Having only one Gaussian per decision does not mean there is only one behavior. Sampled actions change the state, the state changes the next Gaussian, and in the closed loop this creates diffusion over physical time. In the same RoboCasa lunch-packing scene, GAR can produce two different packing orders. The project page has a 2D example: 128 trajectories start from the same point, avoid a circular region, and then reach the goal, without a predefined route.

Some commenters asked: if this result holds, does it mean most teleoperation demonstrations are actually quite unimodal? Does it still hold for tasks where the operator switches between left and right regrasping? Others asked whether the Gaussian can hold up when friction starts causing trouble on real hardware, and someone said the posterior still needs a clock. The paper does not answer these questions. What can be said is that in the control experiments, a Gaussian fitted to the teacher mean and full covariance achieves a success rate close to the teacher, indicating that the Gaussian shape itself does not hurt the task much. GR00T is more sensitive to the mean estimate; mean-only control will be limited by mean error.

Paper: Gaussian Distillation of Flow Policies

Project page: GAR. Code will be released.