확산 기반 정책에서의 숨겨진 보상 복구

기존의 확산 기반 로봇 제어 방식은 전문가의 행동을 단순히 흉내 내는 데 그쳐 새로운 환경에 적응하는 능력이 부족했습니다. 이 논문은 확산 모델 내부에 숨겨진 보상 체계를 수학적으로 증명하고 추출함으로써, 로봇이 단순 모방을 넘어 스스로 학습하고 낯선 상황에서도 더 똑똑하게 대처할 수 있는 길을 열었습니다.

기존의 Diffusion Policy는 노이즈가 섞인 행동에서 노이즈를 제거하여 전문가의 행동을 복원하는 Score Matching 방식을 사용한다. 이는 전문가의 행동 분포를 잘 따라 하지만, 왜 그 행동이 좋은지에 대한 '의도'나 '보상' 개념이 없어 학습 데이터에 없는 상황이 오면 대처하기 어렵다는 한계가 있다.

ENERGYFLOW는 이 문제를 해결하기 위해 물리적인 에너지 개념을 도입한다. 모델이 직접 노이즈를 예측하는 대신, 특정 상태와 행동에 대한 '에너지 값'을 출력하는 함수를 먼저 정의한다. 그리고 이 에너지의 경사면(Gradient)을 따라가는 방향이 곧 노이즈를 제거하는 방향이 되도록 설계한다. 이렇게 하면 모델은 자연스럽게 전문가가 선호하는 행동일수록 낮은 에너지를 갖도록 학습하게 된다.

이 과정에서 핵심은 '보존장(Conservative Field)' 제약이다. 에너지가 경로에 상관없이 일정하게 유지되어야 한다는 물리적 원칙을 모델에 강제함으로써, 모델이 복잡한 노이즈 패턴을 억지로 외우는 대신 논리적이고 매끄러운 에너지 지형을 그리도록 유도한다. 결과적으로 로봇은 학습하지 않은 위치에서도 에너지가 낮은 방향(정답에 가까운 방향)을 더 정확하게 찾아낼 수 있게 된다.

ENERGYFLOW는 스칼라 에너지 함수 E_phi(o, a)를 파라미터화하며, Score Function S_phi를 -grad_a E_phi로 정의한다. 이를 통해 학습된 필드가 항상 보존적(Conservative)임을 보장하며, Denoising Score Matching 목적 함수를 최소화하여 에너지 지형을 학습한다.

학습 과정에서는 분산 폭발(Variance-Exploding) 설정을 따르며, 노이즈 스케줄 sigma(t)에 따라 오염된 행동 a_t를 입력받아 grad_a_t E_phi(a_t, s, t)를 계산한다. [상태 s와 노이즈 섞인 행동 a_t를 입력으로] -> [에너지 함수의 행동에 대한 편미분을 수행해] -> [Score 벡터를 얻고] -> [이 값이 실제 주입된 노이즈와 일치하도록 가중치를 갱신한다].

추출된 에너지를 강화학습 보상으로 쓰기 위해 Centered Shaping 기법을 적용한다. r_phi(a, s) = -[E_phi(a, s, gamma) - E_a'~N(0,I)[E_phi(a', s, gamma)]] 식을 사용한다. [특정 행동의 에너지에서] -> [가우시안 분포에서 샘플링된 참조 행동들의 평균 에너지를 빼서] -> [상태별 상대적 선호도 점수를 얻고] -> [이를 강화학습의 보상 신호로 활용한다].

RoboMimic과 Meta-World 벤치마크의 10개 작업에서 실험한 결과, ENERGYFLOW는 평균 성공률 93.8%(RoboMimic)와 92.5%(Meta-World)를 기록하며 기존 Diffusion Policy(91.2%, 90.7%)를 능가했다. 특히 ToolHang과 같은 복잡한 조작 작업에서 성능 향상이 두드러졌다.

분포 외(OOD) 일반화 실험에서 초기 위치 섭동(Perturbation)이 커질수록 기존 모델들은 성능이 급격히 하락했으나, ENERGYFLOW는 보존장 제약 덕분에 훨씬 완만한 성능 저하를 보이며 강건함을 입증했다. 또한 추출된 에너지 보상을 SAC 알고리즘에 적용했을 때, Sparse Reward나 Raw Energy를 사용한 경우보다 훨씬 빠르고 안정적인 학습 곡선을 보여주며 Oracle 보상에 근접한 성능을 달성했다.

ENERGYFLOW는 Diffusion Policy의 1D Conditional U-Net 백본을 수정하여 사용한다. 기존의 벡터 출력 헤드를 Global Average Pooling과 3층 MLP로 구성된 스칼라 출력 헤드로 교체하여 단일 에너지 값을 산출하도록 설계했다. 또한 에너지의 2차 미분이 포함되는 Score Matching 학습을 위해 모든 활성화 함수를 Mish와 같은 C2 연속 함수로 대체했다.

이론적으로는 Rademacher Complexity 분석을 통해 보존장 제약이 가설 공간의 복잡도를 출력 차원 d의 제곱근에 비례하던 것에서 표현의 매끄러움(Smoothness) L에 비례하도록 줄여줌을 증명했다. 이는 고차원 행동 공간에서 모델의 일반화 능력을 수학적으로 보장하는 핵심 장치이다. 구현 측면에서는 PyTorch의 autograd.grad를 사용하여 순방향 패스에서 Score를 계산하고 이를 통해 역전파를 수행하는 Differentiable Training 인프라를 구축했다.

본 논문은 상태 의존적 통합 상수 c(s)를 완전히 제거하지 못해 서로 다른 상태 간의 절대적인 보상 값을 비교하는 데 한계가 있음을 명시했다. 또한 Score Matching 오차가 존재할 경우 보상 추출의 정확도가 Lipschitz 연속성에 따라 제한될 수 있음을 언급했다.