💻 computer science

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

이 논문은 무거운 UNet 백본을 새로운 RWKV-KAN 아키텍처와 행동 일관성 정규화(Action Consistency Regularization)로 대체하여, 훨씬 적은 파라미터 수로 최첨단 성능을 달성하는 로봇 조작을 위한 경량화되고 효율적인 3D 플로우 매칭(flow-matching) 정책인 KAN-We-Flow를 소개한다.

원저자: Zhihao Chen, Yiyuan Ge, Ziyang Wang, Youwei Zhang

게시일 2026-08-19

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Zhihao Chen, Yiyuan Ge, Ziyang Wang, Youwei Zhang

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

로봇은 오랫동안 정교한 조작이라는 예술적 과제에 어려움을 겪어왔습니다. 로봇은 무거운 물체를 직선으로 이동시키는 데는 뛰어나지만, 바늘에 실을 꿰거나, 문손잡이를 돌리거나, 노트북을 조립하는 것과 같이 인간과 같은 섬세한 손길이 필요한 작업은 여전히 난제로 남아 있습니다. 수년 동안 연구자들은 전문가가 동작을 수행하는 영상을 로봇에게 보여주는 방식인 모방 학습(imitation learning)을 통해 이러한 기술을 가르치려 노력해 왔습니다. 가장 성공적인 최근의 시도들은 일종의 추측 게임처럼 작동하는 복잡한 수학적 모델에 의존했습니다. 이 모델들은 무작위의 노이즈가 섞인 추측에서 시작하여, 추측이 명확한 움직임 계획이 될 때까지 단계별로 조금씩 정교하게 다듬어 나갑니다. 이 방법은 고품질의 움직임을 만들어내는 데는 효과적이지만, 매우 느리다는 단점이 있습니다. 로봇은 단 하나의 동작을 취할 때마다 수십 번의 정교화 단계를 거쳐야 하며, 이는 실시간 제어를 불가능하게 만드는 지연을 발생시킵니다. 이는 마치 자동차가 움직이기도 전에 다음 회전 구간을 추측하고, 추측을 수정하고, 다시 추측하고, 다시 수정하는 과정을 반복하며 운전하는 것과 같습니다.

이 속도 문제를 해결하기 위해 과학자들은 '플로우 매칭(flow matching)'이라 불리는 더 새로운 접근 방식에 주목했습니다. 이 방법은 추측을 다듬기 위해 여러 번의 작은 단계를 거치는 대신, 무작위 시작점에서 올바른 동작으로 이어지는 직접적인 경로를 학습하여, 이론적으로 로봇이 단 한 번의 단계만으로 다음 동작을 결정할 수 있게 합니다. 그러나 큰 장애물이 남아 있었습니다. 이 모델을 구동하는 컴퓨터의 두뇌가 여전히 너무 무겁고 비대하다는 점이었습니다. 이 모델들은 강력한 서버를 필요로 하는 거대하고 에너지를 많이 소비하는 구조에 의존했기에, 실제 로봇에 탑재된 소형 배터리 구동 컴퓨터에서는 실행이 불가능했습니다. 분야 전체에는 새로운 방식의 속도는 유지하면서도 기존 하드웨어의 무게를 덜어낼 방법이 필요했습니다.

새로운 연구에서 연구진은 정확히 이 균형을 달성하는 로봇 정책을 구축했습니다. 그들은 기존의 무거운 전통적 컴퓨터 구조를 훨씬 가볍고 효율적인 설계로 대체하는 'KAN-We-Flow'라는 시스템을 도입했습니다. 연구진은 두 가지 최신 인공지능 기술을 결합하여 새로운 유형의 처리 블록을 만들었습니다. 첫 번째 부분은 시간의 흐름을 처리하여, 로봇이 단순히 단일 스냅샷을 보는 것이 아니라 마치 이야기를 따라가는 것처럼 자신의 동작이 일련의 과정 속에서 어떻게 연결되는지 이해하도록 돕습니다. 두 번째 부분은 미세 조정 역할을 수행하며, 경직된 직선이 아닌 유연하고 곡선적인 패턴을 학습하는 방식을 통해 로봇의 움직임을 높은 정밀도로 조정합니다. 이 두 가지 역량을 하나로 엮음으로써, 연구팀은 이전의 최고 방식보다 86.8% 더 작으면서도 더 빠르게 실행되고 오류는 더 적은 모델을 만들어냈습니다.

이 연구의 결과는 표준 벤치마크 테스트에서 놀라운 성과를 보여주었습니다. 문 열기, 펜 돌리기, 가구 조립 등을 포함한 일련의 시뮬레이션 환경에서, 이 새로운 시스템은 기존의 방식들을 일관되게 능가했습니다. 로봇 손이 펜을 조작하는 어려운 과제에서 새 모델은 68%의 성공률을 기록했는데, 이는 이전의 선도적인 방식이 기록한 55%와 대조됩니다. 문과 관련된 또 다른 테스트에서는 83%의 성공률에 도달하여 차세대 시스템을 상당한 차이로 앞질렀습니다. 무엇보다 실제 사용 측면에서 중요한 점은 이 시스템이 믿기지 않을 정도로 빠르다는 것입니다. 이 시스템은 약 8~11밀리초 만에 새로운 동작을 결정할 수 있으며, 이는 초당 100회의 속도로 로봇의 움직임을 제어할 수 있을 만큼 충분히 빠른 속도입니다. 이 속도는 인간의 반사 신경과 맞먹는 수준으로, 로봇이 비틀거리거나 목표를 놓치게 만들었던 기존 방식들의 100밀리초 이상의 단계와 극명한 대조를 이룹니다.

로봇이 경로를 이탈하지 않도록 연구진은 안전망 역할을 하는 특정 학습 규칙을 추가했습니다. 학습 단계 동안, 시스템은 단순히 다음 동작을 예측하는 것뿐만 아니라, 현재의 경로를 앞으로 투영했을 때 정확히 전문가인 인간이 도달했을 지점에 착륙하게 된다는 것을 보장하도록 교육받습니다. 이 규칙은 특히 길고 복적인 작업 중에 로봇이 경로를 벗어나는 것을 방지합니다. 연구팀은 이 간단한 추가 사항이 의사 결정에 추가적인 시간을 소요하지 않으면서도 학습 과정을 안정시키고 최종 정확도를 높였다는 것을 발견했습니다. 전체 시스템은 단순한 물체 핸들링부터 복잡한 조립에 이르기까지 세 가지 서로 다른 과제 세트에서 테스트되었으며, 기존 모델에 필요한 컴퓨터 메모리의 극히 일부만을 사용하면서도 일관되게 가장 높은 성공률을 기록했습니다.

이 연구의 의의는 로봇 학습을 실험실에서 공장 현장이나 가정으로 옮길 수 있는 잠재력에 있습니다. 가볍고 효율적인 모델이 거대하고 느린 모델보다 뛰어날 수 있음을 증명함으로써, 연구진은 배포의 주요 장벽을 제거했습니다. 이 새로운 시스템은 실행을 위해 슈퍼컴퓨터를 필요로 하지 않습니다. 로봇 자체에 탑재될 수 있을 만큼 작기 때문입니다. 이는 로봇이 강력한 클라우드 서버에 지속적으로 연결되지 않고도, 새로운 물체와 환경에 맞춰 실시간으로 새로운 기술을 습득할 수 있음을 의미합니다. 본 연구는 로봇 조작의 미래가 더 크고 무거운 두뇌를 만드는 것이 아니라, 인간의 손처럼 빠르게 생각하고 정밀하게 행동할 수 있는 더 스마트하고 효율적인 두뇌를 만드는 데 있다는 것을 시사합니다.

기술 요약: KAN-We-Flow

문제 정의

현재의 생성적 시각-운동 정책(generative visuomotor policies), 특히 확산 모델(diffusion models, 예: DP3)에 기반한 모델들은 복잡하고 다중 양상적인(multi-modal) 행동 분포를 모델링하는 데 탁월합니다. 그러나 이들은 심각한 추론 비효율성 문제를 겪고 있습니다. 이러한 모델들은 일반적으로 무거운 UNet 스타일의 백본에 의존하며, 행동을 생성하기 위해 여러 단계에 걸친 재귀적 디노이징(denoising) 과정을 필요로 합니다. 이는 높은 지연 시간(latency)을 초래하여 실시간 제어 및 자원이 제한된 에지 로봇(edge robots)으로의 배포를 저해합니다.

Flow Matching(FM) 패러다임은 하나의 단계만으로 행동 생성을 가능하게 하는 연속적인 벡터장을 학습함으로써 해결책을 제시하지만, 기존 구현체들은 여전히 거대하고 파라미터가 많은 UNet 아키텍처를 그대로 상속받는 경우가 많습니다. 결과적으로, 확산 모델보다 샘플링 속도는 빠르지만, 여전히 파라미터 수와 계산 비용 측면에서 병목 현상이 발생하여 실제적인 에지 배포에 한계가 있습니다. 따라서 Flow 모델의 정확도와 추론 속도를 유지하면서도 계산 오버헤드를 획기적으로 줄일 수 있는 정책 아키텍처가 절실히 필요합니다.

방법론

저자들은 무거운 UNet 백본을 대체하기 위해 **RWKV(Receptance Weighted Key Value)**와 **KAN(Kolmogorov-Arnold Networks)**을 결합한 새로운 아로키텍처인 KAN-We-Flow라는 경량 Flow-matching 정책을 제안합니다.

1. RWKV-KAN UNet 백본

표준 UNet 대신, 이 정책은 RWKV-KAN 블록으로 구성된 3단계 인코더-디코더를 활용합니다. 이 설계는 장기적인 시간적 의존성(long-horizon temporal dependencies)을 포착하고 선형 복잡도로 효율적인 특징 교정(feature calibration)을 수행하는 것을 목표로 합니다.

RWKV 구성 요소: RWKV 모듈은 시간 및 채널 믹싱을 담당합니다. 이 모듈은 지수적 감쇠(exponential decay)를 사용하는 시간 믹싱 메커니즘을 사용하여 선형 복잡도로 작업 컨텍스트를 행동 궤적 전체에 전파하며, squared-ReLU 비선형성을 가진 게이트형 MLP를 사용하는 채널 믹싱 브랜치를 사용합니다. 이는 트랜스포머(Transformer)에서 발견되는 값비싼 셀프 어텐션(self-attention) 메커니즘과 UNet의 무거운 합성곱 스택을 대체합니다.
GroupKAN 구성 요소: RWKV에 이어, GroupKAN 레이어는 학습 가능한 스플라인 기반(spline-based)의 그룹별 함수 매핑을 적용합니다. 고정된 활성화 함수와 무거운 행렬 곱셈을 사용하는 전통적인 MLP와 달리, KAN은 네트워크의 엣지(edge)에 학습 가능한 단변량 함수(스플라인으로 매개변수화됨)를 배치합니다. 저자들은 채널을 그룹으로 분할하고, 각 그룹을 독립적인 KAN 연산자로 처리하되 시간(time)에 따라 파라미터를 공유하도록 합니다. 여기에 시퀀스 인식 통계에 따라 채널 가중치를 적응적으로 재조정하는 경량 채널 친화도 변조(Channel Affinity Modulation, CAM) 메커니즘이 추가되었습니다. 이 구조는 훨씬 적은 파라미터로 다변량 함수를 효율적으로 근사할 수 있게 합니다.

2. 일관성 흐름 매칭 (Consistency Flow Matching, CFM)

정책은 행동 공간 내에서 속도가 일관된 직선형 흐름(velocity-consistent straight-line flow)을 학습합니다.

목적 함수: 모델은 두 인접한 타임 스텝( $t$ 와 $t+\Delta t$ )으로부터 디코딩했을 때 동일한 종단점(endpoint)을 생성하고 그 순간의 속도가 일치하도록 강제하는 일관성 손실( $L_{CFM}$ )을 최소화합니다.
다중 세그먼트 확장: 표현력을 높이기 위해 시간 간격 $[0, 1]$ 을 $K$ 개의 세그먼트(실험에서는 2로 설정)로 분할하여, 효율성과 품질의 균형을 맞추기 위해 세그먼트별 일관성을 최적화합니다.

3. 행동 일관성 정규화 (Action Consistency Regularization, ACR)

추론 단계를 추가하지 않고도 학습을 안정화하고 정밀도를 높이기 위해, 저자들은 **행동 일관성 정규화(ACR)**를 도입했습니다.

메커니즘: ACR은 정책이 예측한 행동 궤적을 전문가의 시연(expert demonstrations)과 정렬시키는 경량 보조 손실(auxiliary loss) 역할을 합니다. 현재 상태에서 호라이즌 끝( $t \to 1$ )까지 단일 단계 오일러 외삽(Euler extrapolation)을 수행하고, 이 외삽된 행동과 특정 제어 윈도우에서의 전문가 행동 사이의 편차를 처벌합니다.
이점: 이는 전역적인 종단점 앵커(global end-point anchor)를 제공하여, 훈련 중 드리프트(drift)와 노출 편향(exposure bias) 효과를 줄여주면서도 계산 비용은 거의 들지 않습니다.

주요 기여

새로운 아키텍처: RWKV-KAN 블록의 도입으로, 이러한 특정 아키텍처(RWKV 및 GroupKAN)를 로봇 시각-운동 정책 학습에 최초로 적용했습니다. 이 설계는 정밀도를 희생하지 않으면서도 모델 크기를 획기적으로 줄이고 추론 속도를 높였습니다.
행동 일관성 정규화 (ACR): 호라이즌 끝에서 단일 단계 디코딩된 행동을 전문가의 시연에 고정시켜, 매우 적은 비용으로 안정성과 정밀도를 향상시키는 새로운 보조 목적 함수입니다.
효율성-정확도 균형: 이러한 구성 요소들의 통합은 높은 정확도와 높은 계산 효율성을 동시에 갖춘 정책을 만들어내어, 실시간 에지 배포에 적합하도록 합니다.

실험 결과

저자들은 Adroit, Meta-World, DexArt의 세 가지 표준 벤치마크에서 KAN-We-Flow를 평가했습니다.

성능: 이 방법은 모든 도메인에서 최첨단(SOTA) 성공률을 달s성했습니다.
- Adroit에서, Door 작업의 성공률을 $83.0 \pm 2.0$ 으로 개선했습니다 (FlowPolicy의 $62.0 \pm 5.0$ 대비). Pen 작업은 $68.0 \pm 2.5$ 를 기록했습니다 (vs. $55.0 \pm 13.0$ ).
- Meta-World에서, Easy, Medium, Hard, Very Hard 난이도 계층 전반에서 베이스라인을 능가했습니다 (예: Easy 작업에서 MP1의 $88.2 \pm 1.1$ 대비 $92.0 \pm 1.0$ ).
효율성:
- 파라미터: 모델은 단 33.6M 파라미터를 사용하며, 이는 DP3 베이스라인의 255.1M 파라미터와 비교했을 때 86.8% 감소한 수치입니다.
- 추론 속도: 이 정책은 Adroit 작업에서 8.0–10.1 ms, Meta-World 작업에서 6.7–11.8 ms의 실시간 추론 속도를 달성하여 약 100 Hz 제어가 가능합니다. 이는 확산 베이스라인(예: DP3의 약 130 ms)보다 훨씬 빠르며, 기존의 Flow-matching 방식인 FlowPolicy보다도 개선된 성능입니다.
절제 연구 (Ablation Studies): RWKV, GroupKAN, 또는 ACR 중 하나라도 제거하면 성능 저하가 발생함을 확인하여, 선형 시간 시퀀스 믹싱, 파라미터 효율적인 스플라인 교정, 그리고 호라이즌 엔드 앵커링의 상호 보완적인 역할을 검증했습니다.

의의 및 주장

본 논문은 KAN-We-Flow가 무거운 UNet 아키텍처에 의존하지 않고도 정확하고 실시간 대응이 가능하며 에지에 배포 가능한 로봇 정책을 향한 명확한 경로를 보여준다고 주장합니다. 전통적인 백본을 RWKV와 KAN으로 교체함으로써, 저자들은 높은 성능과 높은 계산 비용 사이의 관계를 성공적으로 분리했습니다. 이 연구는 효율적인 시퀀스 모델링(RWKV)과 표현력이 풍부하고 파라미터 효율적인 함수 근사(KAN)의 결러가 로봇 조작에 특히 적합하며, 정교한 시각-운동 정책을 자원이 제한된 하드웨어에 배포하기 위한 유망한 솔루션을 제공한다는 점을 시사합니다. 저자들은 향후 연구 과제로 실제 로봇 평가, 멀티 뷰 센싱, 그리고 접촉이 빈번한(contact-rich) 작업으로의 확장을 언급했습니다.