💻 computer science

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

This paper introduces KAN-We-Flow, a lightweight and efficient 3D flow-matching policy for robotic manipulation that replaces heavy UNet backbones with a novel RWKV-KAN architecture and Action Consistency Regularization to achieve state-of-the-art performance with significantly fewer parameters.

Original authors: Zhihao Chen, Yiyuan Ge, Ziyang Wang, Youwei Zhang

Published 2026-08-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Zhihao Chen, Yiyuan Ge, Ziyang Wang, Youwei Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Robots have long struggled with the delicate art of manipulation. While they excel at moving heavy objects in straight lines, tasks requiring a human-like touch—such as threading a needle, turning a doorknob, or assembling a laptop—have remained elusive. For years, researchers have tried to teach robots these skills by showing them videos of experts performing the actions, a process known as imitation learning. The most successful recent attempts have relied on complex mathematical models that act like a guessing game. These models start with a random, noisy guess and slowly refine it, step by step, until the guess becomes a clear plan for movement. While this method works well to produce high-quality movements, it is painfully slow. The robot must run through dozens of refinement steps for every single action it takes, creating a delay that makes real-time control impossible. It is like trying to steer a car by guessing the next turn, correcting the guess, guessing again, and correcting again, all before the car has even moved an inch.

To solve this speed problem, scientists have turned to a newer approach called flow matching. Instead of taking many small steps to refine a guess, this method learns a direct path from a random starting point to the correct action, theoretically allowing the robot to decide its next move in a single step. However, a major hurdle remained: the computer brains powering these models were still too heavy and bulky. They relied on massive, energy-hungry architectures that required powerful servers, making them impossible to run on the small, battery-powered computers found on actual robots. The field needed a way to keep the speed of the new method without the weight of the old hardware.

In a new study, a team of researchers has built a robot policy that achieves exactly this balance. They introduced a system called KAN-We-Flow, which replaces the heavy, traditional computer architecture with a much lighter and more efficient design. The researchers combined two recent advances in artificial intelligence to create a new type of processing block. The first part handles the flow of time, allowing the robot to understand how its actions connect over a sequence, much like following a story rather than just looking at a single snapshot. The second part acts as a fine-tuner, adjusting the robot's movements with a high degree of precision using a method that learns flexible, curved patterns rather than rigid, straight lines. By weaving these two capabilities together, the team created a model that is 86.8% smaller than the previous best methods, yet it runs faster and makes fewer mistakes.

The results of this work are striking when tested against standard benchmarks. In a series of simulated environments involving tasks like opening doors, turning pens, and assembling furniture, the new system consistently outperformed its predecessors. On a difficult task involving a robotic hand manipulating a pen, the new model achieved a success rate of 68%, compared to 55% for the previous leading method. In another test involving a door, it reached 83% success, beating the next best system by a significant margin. Perhaps most importantly for real-world use, the system is incredibly fast. It can decide on a new action in about 8 to 11 milliseconds, which is fast enough to allow the robot to control its movements at a rate of 100 times per second. This speed is comparable to the reflexes of a human, a stark contrast to the older methods that took over 100 milliseconds per step, often causing the robot to stumble or miss its target entirely.

To ensure the robot stayed on track, the researchers added a specific training rule that acts as a safety net. During the learning phase, the system is taught not only to predict the next move but also to ensure that if it were to project its current path forward, it would land exactly where an expert human would have ended up. This rule prevents the robot from drifting off course, especially during long and complex tasks. The team found that this simple addition stabilized the training process and improved the final accuracy without requiring any extra time to make decisions. The entire system was tested on three different sets of tasks, ranging from simple object handling to complex assembly, and it consistently delivered the highest success rates while using a fraction of the computer memory required by older models.

The significance of this work lies in its potential to move robot learning out of the laboratory and onto the factory floor or into the home. By proving that a lightweight, efficient model can outperform massive, slow ones, the researchers have removed a major barrier to deployment. The new system does not require a supercomputer to run; it is small enough to fit on the robot itself. This means that robots could eventually learn new skills on the fly, adapting to new objects and environments without needing a constant connection to a powerful cloud server. The study suggests that the future of robotic manipulation does not require building bigger, heavier brains, but rather smarter, more efficient ones that can think quickly and act with the precision of a human hand.

Technical Summary: KAN-We-Flow

Problem Statement

Current generative visuomotor policies for robotic manipulation, particularly those based on diffusion models (e.g., DP3), excel at modeling complex, multi-modal action distributions. However, they suffer from significant inference inefficiencies. These models typically rely on heavy UNet-style backbones and require recursive denoising over many steps to generate actions, leading to high latency that hinders real-time control and deployment on resource-constrained edge robots.

While Flow Matching (FM) paradigms offer a solution by learning a continuous vector field that enables one-step action generation, prior implementations often inherit the same large, parameter-heavy UNet architectures. Consequently, despite faster sampling than diffusion, these models still face bottlenecks in parameter count and computational cost, limiting their practicality for edge deployment. There is a critical need for a policy architecture that maintains the accuracy and inference speed of flow models while drastically reducing computational overhead.

Methodology

The authors propose KAN-We-Flow, a lightweight flow-matching policy designed to replace heavy UNet backbones with a novel architecture combining Receptance Weighted Key Value (RWKV) and Kolmogorov-Arnold Networks (KAN).

1. RWKV-KAN UNet Backbone

Instead of a standard UNet, the policy utilizes a three-stage encoder-decoder composed of RWKV-KAN blocks. This design aims to capture long-horizon temporal dependencies and perform efficient feature calibration with linear complexity.

RWKV Component: The RWKV module handles time and channel mixing. It employs a time-mixing mechanism with exponential decay to propagate task context across the action trajectory with linear complexity, and a channel-mixing branch using gated MLPs with squared-ReLU nonlinearity. This replaces the computationally expensive self-attention mechanisms found in Transformers and the heavy convolutional stacks in UNets.
GroupKAN Component: Following RWKV, a GroupKAN layer applies learnable spline-based, groupwise functional mappings. Unlike traditional MLPs that use fixed activation functions and heavy matrix multiplications, KANs place learnable univariate functions (parameterized by splines) on the edges of the network. The authors partition channels into groups, processing each with an independent KAN operator that shares parameters across time. A lightweight Channel Affinity Modulation (CAM) mechanism is added to adaptively reweight channels based on sequence-aware statistics. This structure enables efficient approximation of multivariate functions with significantly fewer parameters.

2. Consistency Flow Matching (CFM)

The policy learns a velocity-consistent straight-line flow in the action space.

Objective: The model minimizes a consistency loss ( $L_{CFM}$ ) which enforces that decoding from two nearby time steps ( $t$ and $t+\Delta t$ ) yields the same endpoint and that their instantaneous velocities agree.
Multi-segment Extension: To improve expressivity, the time interval $[0, 1]$ is partitioned into $K$ segments (set to 2 in experiments), optimizing segment-wise consistency to balance efficiency and quality.

3. Action Consistency Regularization (ACR)

To stabilize training and improve precision without adding inference steps, the authors introduce Action Consistency Regularization (ACR).

Mechanism: ACR acts as a lightweight auxiliary loss that aligns the policy's predicted action trajectory with expert demonstrations. It performs a single-step Euler extrapolation from the current state to the horizon end ( $t \to 1$ ) and penalizes the deviation between this extrapolated action and the expert action over a specific control window.
Benefit: This provides a global end-point anchor, reducing drift and exposure bias effects during training while incurring negligible computational cost.

Key Contributions

Novel Architecture: The introduction of the RWKV-KAN block, marking the first application of these specific architectures (RWKV and GroupKAN) to robotic visuomotor policy learning. This design achieves drastically reduced model size and faster inference without sacrificing precision.
Action Consistency Regularization (ACR): A new auxiliary objective that anchors one-step decoded actions to expert demonstrations at the horizon, improving stability and precision at negligible cost.
Efficiency-Accuracy Balance: The integration of these components results in a policy that is both highly accurate and computationally efficient, suitable for real-time edge deployment.

Experimental Results

The authors evaluated KAN-We-Flow on three standard benchmarks: Adroit, Meta-World, and DexArt.

Performance: The method achieved state-of-the-art success rates across all domains.
- On Adroit, it improved success rates on the Door task to $83.0 \pm 2.0$ (vs. $62.0 \pm 5.0$ for FlowPolicy) and Pen to $68.0 \pm 2.5$ (vs. $55.0 \pm 13.0$ ).
- On Meta-World, it outperformed baselines across Easy, Medium, Hard, and Very Hard difficulty tiers (e.g., $92.0 \pm 1.0$ on Easy tasks vs. $88.2 \pm 1.1$ for MP1).
Efficiency:
- Parameters: The model uses only 33.6M parameters, representing an 86.8% reduction compared to the 255.1M parameters of the DP3 baseline.
- Inference Speed: The policy achieves real-time inference speeds of 8.0–10.1 ms on Adroit tasks and 6.7–11.8 ms on Meta-World tasks, enabling $\sim$ 100 Hz control. This is significantly faster than diffusion baselines (e.g., DP3 at $\sim$ 130 ms) and improves upon existing flow-matching methods like FlowPolicy.
Ablation Studies: Experiments confirmed that removing either RWKV, GroupKAN, or ACR leads to performance degradation, validating the complementary roles of linear-time sequence mixing, parameter-efficient spline calibration, and horizon-end anchoring.

Significance and Claims

The paper claims that KAN-We-Flow demonstrates a clear path toward accurate, real-time, and edge-deployable robot policies that do not rely on heavy UNet architectures. By replacing traditional backbones with RWKV and KAN, the authors successfully decouple high performance from high computational cost. The work suggests that the combination of efficient sequence modeling (RWKV) and expressive, parameter-efficient function approximation (KAN) is particularly well-suited for robotic manipulation, offering a viable solution for deploying sophisticated visuomotor policies on resource-constrained hardware. The authors note that future work will focus on extending this approach to real-robot evaluations, multi-view sensing, and contact-rich tasks.