KAN We Flow? Advancing Robotic Manipulation with 3D Flow Matching via KAN & RWKV

TL;DR

KAN-We-Flow introduces a lightweight flow-matching visuomotor policy built on a RWKV–KAN backbone to achieve real-time manipulation with far fewer parameters than UNet-heavy baselines. The RWKV–KAN architecture captures long-horizon context and per-channel calibration using spline-based functions, complemented by Action Consistency Regularization to stabilize horizon-end behavior without extra inference steps. Across Adroit, Meta-World, and DexArt, the method achieves state-of-the-art success rates while delivering substantial efficiency gains (e.g., ~86.8% fewer parameters) and real-time control (~100 Hz). This work demonstrates a practical path toward edge-deployable robotic policies that maintain high precision and responsiveness.

Abstract

Diffusion-based visuomotor policies excel at modeling action distributions but are inference-inefficient, since recursively denoising from noise to policy requires many steps and heavy UNet backbones, which hinders deployment on resource-constrained robots. Flow matching alleviates the sampling burden by learning a one-step vector field, yet prior implementations still inherit large UNet-style architectures. In this work, we present KAN-We-Flow, a flow-matching policy that draws on recent advances in Receptance Weighted Key Value (RWKV) and Kolmogorov-Arnold Networks (KAN) from vision to build a lightweight and highly expressive backbone for 3D manipulation. Concretely, we introduce an RWKV-KAN block: an RWKV first performs efficient time/channel mixing to propagate task context, and a subsequent GroupKAN layer applies learnable spline-based, groupwise functional mappings to perform feature-wise nonlinear calibration of the action mapping on RWKV outputs. Moreover, we introduce an Action Consistency Regularization (ACR), a lightweight auxiliary loss that enforces alignment between predicted action trajectories and expert demonstrations via Euler extrapolation, providing additional supervision to stabilize training and improve policy precision. Without resorting to large UNets, our design reduces parameters by 86.8\%, maintains fast runtime, and achieves state-of-the-art success rates on Adroit, Meta-World, and DexArt benchmarks. Our project page can be viewed in \href{https://zhihaochen-2003.github.io/KAN-We-Flow.github.io/}{\textcolor{red}{link}}

Figures (6)

• Figure 1: Comparison of KAN-We-Flow with the state-of-the-art methods FlowPolicy and DP3 regarding accuracy, parameter, and inference time. (a) KAN-We-Flow achieves superior success rates among different benchmarks' tough tasks; (b) Our approach obtains an 86.8% parameter reduction, compared with FlowPolicy and DP3; (c) Compared with DP3, our KAN-We-Flow achieves 92.6% inference time decrease in the Adroit–Pen task, enabling real-time control.
• Figure 2: Overview of KAN-We-Flow. The policy receives a noised action and a condition that comprises three encoded parts, a point-cloud perception embedding, a robot-state embedding, and a time embedding. The concatenated representation is processed by a lightweight RWKV-KAN U-shaped backbone instead of a large UNet-style backbone; RWKV mixes long-range time/channel context with linear complexity, while KAN performs learnable spline-based feature calibration. Then, a straight-line flow is learned with conditional consistency flow matching to produce a one-step velocity field, and the resulting actions are generated at real-time inference speed; an additional action consistency regularization aligns Euler-extrapolated trajectories with demonstrations to stabilize training.
• Figure 3: The architecture of the GroupKAN.
• Figure 4: Visualization of manipulation results. We evaluate on three datasets: Meta-World, Adroit, and DexArt. We illustrate representative rollouts from Meta-World (Assembly, Disassemble), Adroit (Hammer), and DexArt (Laptop). During interaction, the KAN-We-Flow policy predicts future action sequences and continues issuing actions until the task is successfully completed.
• Figure 5: Visualization of success rates. We plot top-1, top-3, and top-5 success rates (SR1, SR3, SR5), computed as the average of the highest 1, 3, and 5 trial outcomes per task; KAN-We-Flow achieves superior performance across tasks.