⚡ electrical engineering

Amortizing Trajectory Diffusion with Keyed Drift Fields

이 논문은 반복적 디노이징의 계산 비용을 줄이면서도 다중 모드 행동 생성과 조건부 계획 능력을 유지하기 위해, 조건 인식형 거리 기반의 '키드 드리프트 필드'를 학습하여 단일 단계 추론으로 고품질 궤적을 생성하는 '키드 드리프트 정책 (KDP)'을 제안합니다.

원저자: Gokul Puthumanaillam, Melkior Ornik

게시일 2026-03-17

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Gokul Puthumanaillam, Melkior Ornik

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

🚗 비유: "내비게이션"의 두 가지 방식

로봇이 길을 찾을 때 (행동을 계획할 때) 기존 방식과 이 논문이 제안하는 방식의 차이를 상상해 보세요.

1. 기존 방식 (확산 모델, Diffusion): "완벽한 지도를 그리는 화가"

기존의 AI 는 길을 찾을 때 마치 정교한 그림을 그리는 화가와 같습니다.

과정: 처음엔 흐릿한 점들만 있는데, "이게 길일까?", "아니면 저길까?"라고 수십 번을 반복해서 점들을 지우고 다시 그립니다.
장점: 아주 세밀하고 다양한 길을 그릴 수 있습니다.
단점: 그림이 완성되기까지 시간이 너무 오래 걸립니다. 로봇이 실시간으로 움직이려면 "다음 걸음"을 결정해야 하는데, 화가가 그림을 다 그리기 전에 로봇은 이미 넘어져버립니다.

2. 새로운 방식 (KDP): "순간포착 전문가"

이 논문이 제안하는 KDP 는 순간포착 전문가입니다.

과정: 지금 상황을 보고 한 번에 가장 적합한 길을 그려냅니다. 반복해서 수정하지 않습니다.
핵심 아이디어: "지금 내가 어디에 있는지 (조건)"를 정확히 파악하면, 미래의 모든 가능성을 다 그릴 필요 없이 가장 유력한 후보들만 빠르게 골라낼 수 있다는 것입니다.

🔑 핵심 비밀: "열쇠 (Key)"로 문을 여는 방법

그런데 여기서 큰 문제가 있었습니다. "한 번에 그리는 것"이 너무 쉬워서, AI 가 중요한 건 무시하고 평균적인, 밍밍한 길만 그리는 경향이 있었습니다.

문제 상황: AI 가 "지금 내가 A 지점에 있다"라고 했을 때, 미래의 모든 가능한 경로 (미래의 상태) 를 다 비교하다 보니, 미래의 불확실한 부분에 너무 신경을 써서 "그냥 중간쯤 있는 평범한 길"로 결론을 내렸습니다. (다양성이 사라짐)
KDP 의 해결책: "열쇠 (Key)" 개념 도입
저자들은 **"미래를 비교할 때, '지금의 상황 (조건)'과 맞는 것끼리만 비교하라"**는 규칙을 만들었습니다.
- 비유: 길을 찾을 때, "지금 내가 있는 곳 (현재 상태)"이라는 열쇠로 문을 엽니다.
- 이 열쇠로 **현재 상황과 딱 맞는 과거의 성공 사례들 (데이터)**만 골라냅니다.
- 그리고 그 사례들을 바탕으로 미래의 전체 경로를 한 번에 그려냅니다.

이렇게 하면 AI 는 "미래가 어떻게 될지"를 헷갈려 하지 않고, **"지금 이 상황에서 가장 잘 맞는 행동"**에 집중하게 되어, 다양하면서도 정확한 행동을 빠르게 결정할 수 있게 됩니다.

🏋️‍♂️ 훈련 방법: "수업 중 연습" vs "시험 시간"

기존 방식은 시험 시간 (실제 작동) 에 "수십 번의 수정"을 하느라 느렸습니다.
KDP 는 **훈련 시간 (학습)**에 모든 수정 작업을 끝내버립니다.

훈련 중: AI 가 길을 그릴 때, "이건 좀 더 오른쪽으로 가야 해", "저건 왼쪽으로 가야 해"라고 **수백 번의 수정 (드리프트)**을 가합니다. 이때 AI 는 이 수정 과정을 한 번의 계산으로 기억해 냅니다.
시험 (실제 작동) 중: 이미 훈련을 통해 "수정된 결과"를 외워두었기 때문에, 한 번의 계산으로 바로 완벽한 길을 그려냅니다.

이를 **"학습 단계에서 수정 작업을 미리 해치우는 것 (Amortization)"**이라고 합니다.

🌍 실제 성과: 로봇이 얼마나 빨라졌을까?

이 기술을 실제 로봇 (드론, 로봇 팔) 에 적용해 보니 놀라운 결과가 나왔습니다.

속도: 기존 방식보다 수십 배에서 수백 배 빨라졌습니다. (예: 242ms → 8ms)
정확도: 속도가 빨라졌다고 해서 로봇이 넘어지거나 길을 잃지 않았습니다. 오히려 더 정확하게 움직였습니다.
실시간 대응: 로봇이 갑자기 장애물을 만나도, 매우 빠르게 새로운 길을 찾아서 즉시 피할 수 있게 되었습니다. (기존 방식은 길을 찾느라 너무 늦어서 충돌하는 경우가 많았습니다.)

💡 한 줄 요약

"복잡한 미래를 예측할 때, 과거의 데이터를 무작정 다 비교하지 말고, '지금의 상황'이라는 열쇠로 딱 맞는 사례들만 골라내서 한 번에 결정하자!"

이 방법은 로봇이 더 빠르고, 더 똑똑하게, 그리고 더 안전하게 우리와 함께 살 수 있게 해주는 중요한 기술입니다.

1. 연구 배경 및 문제 제기 (Problem)

배경: 확산 모델 (Diffusion Models) 은 오프라인 강화학습 (Offline RL) 에서 복잡한 다중 모드 (multimodal) 행동 분포를 학습하여 풍부한 궤적 계획을 가능하게 합니다.
문제점: 기존 확산 기반 계획기는 추론 시 반복적인 탈노이즈 (denoising) 과정을 거치기 때문에 계산 비용이 매우 높습니다.
- 지연 시간 (Latency): 제어 루프 내에서 실시간으로 작동하기 어렵습니다.
- 조건부 생성의 실패: 고차원 상태 - 행동 시퀀스에서 단순한 전체 윈도우 (full-window) 거리 기반의 분포 매칭을 시도할 경우, 제약이 없는 미래 차원 (unconstrained future dimensions) 이 거리 계산을 지배하게 됩니다. 이로 인해 생성된 궤적이 데이터의 평균으로 수렴하고 (mode collapse), 행동 다양성이 사라지며, 정적 (static) 인 제어 동작을 유발합니다.
목표: 확산 모델과 유사한 다중 모드 샘플링 및 조건부 제어 능력을 유지하면서, 단일 단계 추론 (one-step inference) 으로 계획이 가능한 방법을 개발하여 실시간 제어에 적용하는 것입니다.

2. 제안 방법론: Keyed Drifting Policies (KDP)

저자들은 반복적인 탈노이즈 과정을 훈련 단계로 이전 (amortize) 하여 단일 단계 생성기를 학습하는 Keyed Drifting Policies (KDP) 를 제안합니다.

핵심 아이디어: 조건 인식형 이웃 (Conditioning-aware Neighborhood)

관찰: 조건 (현재 상태) 에 맞는 궤적만 선택적으로 학습 신호로 사용해야 합니다. 전체 윈도우 공간에서의 거리는 조건과 무관한 미래 정보에 의해 왜곡됩니다.
해결:
1. Key Space (키 공간): 이웃을 선정할 때는 전체 궤적이 아닌, 조건과 정렬된 키 공간 (예: 초기 상태 $s_0$ ) 에서의 거리를 사용합니다.
2. Full Trajectory Space (전체 궤적 공간): 실제 업데이트 (드리프트) 는 전체 궤적 공간에서 수행됩니다.
- 이를 통해 조건에 맞는 데이터 윈도우로 끌어당기면서 (Attraction), 생성된 샘플 간의 다양성을 유지하기 위해 서로 밀어내는 (Repulsion) 메커니즘이 조건에 민감하게 작동합니다.

구체적인 알고리즘 구성

드리프트 필드 (Drift Field) 구성:
- 인력 (Attraction): 조건 (Key) 이 일치하는 데이터셋 윈도우로 샘플을 끌어당깁니다.
- 반발력 (Repulsion): 현재 모델이 생성한 다른 샘플들로부터 밀어내어 모드 붕괴를 방지합니다.
- 키 거리 기반 가중치: softmax 를 사용하여 키 공간 ( $d_\kappa$ ) 의 거리에 기반한 가중치를 계산합니다.
Stop-Gradient Drifted Target:
- 반복적인 드리프트 업데이트를 모델 파라미터 학습에 반영하기 위해, 생성된 샘플에 드리프트 벡터를 더한 후 Stop-Gradient를 적용하여 타겟 ( $\tilde{x}_i$ ) 을 생성합니다.
- 생성기는 이 타겟으로 회귀 (regression) 하도록 훈련됩니다.
제약 조건 처리 (Hard Constraints):
- 현재 상태 ( $s_0$ ) 나 목표 ( $g$ ) 와 같은 조건은 '클램핑 (clamping)' 연산자를 통해 고정하며, 드리프트 업데이트 시 이러한 좌표는 변경되지 않도록 마스킹합니다.
추론 (Inference):
- 훈련된 생성기는 단일 순전파 (forward pass) 로 전체 계획 윈도우를 생성합니다.
- 여러 후보 ( $K$ 개) 를 병렬로 생성한 후, 학습된 스코어러 (scorer) 를 통해 최적의 행동을 선택합니다.

3. 주요 기여 (Key Contributions)

조건부 궤적 생성을 위한 실용적인 키 드리프트 필드:
- 하드 조건 제약 하에서 붕괴 (collapse) 를 방지하고, 조건에 민감한 이웃 구조를 도입하여 학습 신호의 안정성을 확보했습니다.
단일 단계 계획 레시피:
- 확산 계획기의 인터페이스를 유지하면서 $T$ 단계의 순차적 탈노이즈 루프를 제거하고, 단일 단계 생성기로 대체하여 추론 속도를 획기적으로 개선했습니다.
종합적인 평가:
- 표준 시뮬레이션 벤치마크 (D4RL) 와 실시간 하드웨어 배포 (드론, 매니퓰레이터) 를 통해 낮은 지연 시간 예산 하에서도 강력한 성능을 입증했습니다.

4. 실험 결과 (Results)

시뮬레이션 벤치마크 (D4RL)

locomotion (이동): KDP 는 Diffuser 와 같은 확산 기반 모델과 유사하거나 더 나은 성능을 내면서, NFE(순차적 정책 평가 횟수) 를 1 로 줄였습니다.
Maze2D/AntMaze (장기 계획): 조건부 계획이 필요한 장기 계획 작업에서 KDP 는 Diffuser 보다 훨씬 낮은 지연 시간 (약 30~60 배 빠름) 을 유지하면서도 경쟁력 있는 성공률을 보였습니다. 특히, 후보 샘플 수 ( $K$ ) 를 늘려도 지연 시간이 선형적으로 급증하지 않아 Best-of-K 선택이 용이했습니다.
Adroit (정교한 조작): 고차원 제어 및 접촉이 필요한 작업에서 KDP 는 Diffuser 보다 훨씬 빠른 추론 속도 (수십 ms 대 수 초) 와 더 높은 성공률을 기록했습니다. 이는 전체 윈도우를 반복적으로 정제하는 것보다 실행 시 중요한 행동 블록에 집중하는 훈련 방식이 효과적이었음을 시사합니다.

실제 하드웨어 배포

Crazyflie (드론) 및 SO-100 (매니퓰레이터):
- KDP 는 Diffuser 대비 재계획 (replanning) 주파수를 크게 높였습니다 (드론: 4.1Hz → 38Hz, 매니퓰레이터: 3Hz → 17Hz).
- 실시간 환경에서 계획이 낡아지기 전에 빠르게 갱신할 수 있어, 실제 작업 성공률과 안정성이 향상되었습니다.

Ablation Study (성분 분석)

Keying 제거: 전체 윈도우 거리를 사용하면 성능이 급격히 저하되고 행동 다양성이 사라짐 (collapse).
Repulsion 제거: 반발력이 없으면 행동 다양성이 무한히 커지거나 (divergence) 평균화로 인해 성능이 떨어짐.
Self-negative masking: 자기 자신을 반발 대상에서 제외하지 않으면 드리프트가 무효화됨.

5. 의의 및 결론 (Significance)

실시간 생성적 계획의 실현: 확산 모델의 강력한 생성 능력 (다중 모드, 조건부 제어) 을 유지하면서, 추론 시의 반복적 계산을 훈련 단계로 이전함으로써 실시간 제어 (Real-time Control) 가 가능해졌습니다.
시스템 효율성: 순차적 심도 (Sequential depth) 를 줄이고 병렬성 (Parallelism) 을 활용하여 현대 가속기 (GPU/TPU) 환경에서 최적의 성능을 발휘합니다.
향후 방향: 희소 보상 (sparse reward) 이나 복잡한 목표 조건부 계획 영역에서는 여전히 개선의 여지가 있으나, KDP 는 실시간 로봇 제어에 있어 생성적 계획의 새로운 표준을 제시합니다.

요약하자면, 이 논문은 확산 모델의 반복적 추론 비용을 훈련 비용으로 전환하고, 조건에 민감한 거리 측정 (Keyed Drift) 을 도입하여 단일 단계로 고품질의 다중 모드 궤적을 생성하는 KDP 를 제안함으로써, 로봇의 실시간 의사결정 속도와 성능을 동시에 비약적으로 향상시켰습니다.