기술 요약: PLUME

문제 정의

숙련된 다지(multi-finger) 조작은 물체의 형상, 포즈, 마찰 계수와 같은 물리적 파라미터에 매우 민encitive합니다. 이러한 파라미터들은 배포 시점에 불확실한 경우가 많으며, 이는 미지의 시스템 역학으로 이어집니다. 시뮬레이션은 알려진 파라미터를 통해 대규모 데이터 수집을 가능하게 하지만, 표준적인 도메인 무작위화(domain randomization)를 통해 훈련된 정책은 정밀한 작업(예: 드라이버 돌리기)에서 실패하는 경우가 많습니다. 이는 최적의 조작 전략이 특정 파라미터 값에 따라 변해야 할 수도 있기 때문입니다. 또한, 복잡한 다지 작업을 위한 고품질 데이터를 수집하는 것은 어려우며, 종종 하위 최적(sub-optimal) 궤적이나 실패를 포함하는 이질적인 데이터셋을 생성하게 됩니다. 기존의 월드 모델링 방법들은 역학적 변화를 암시적으로 처리하기 위해 고정된 호라이즌의 히스토리에 의존하지만, 이는 정밀한 작업에는 불충분합니다. 또한, 오프라인 강화 학습 방법들은 일반적으로 시스템 식별과 정책 학습의 결합된 문제를 다루지 않으며, 온라인 시스템 식별 방법들은 복잡한 보상 엔지니어링을 요구하고 플래닝(planning)을 지원하지 못하는 경우가 많습니다.

방법론

저자들은 월드 모델을 학습하는 동시에 잠재적 물리 파라미터에 대한 신념(belief)을 추정하는 프레임워크인 PLUME(Probabilistic Latent Unified world Modeling and parameter Estimation)을 제안합니다.

핵심 아키텍처

PLUME은 엔드 투 엔드로 훈련되는 세 가지 주요 구성 요소로 이루어져 있습니다:

1. 파라미터 표현 학습 (Parameter Representation Learning): 트랜스포머 백본을 가진 Wasserstein Auto-Encoder(WAE)가 물리적 파라미터(예: 마찰력, 형상)와 보상을 포함한 이질적인 변수들의 통합된 잠재 표현($z_t$)을 학습합니다. 인코더는 이러한 입력들을 토큰화하고 교차 주의(cross-attention) 풀링을 사용하여 잠재 임베딩을 생성합니다. 합리적인 파라미터 신념의 초기화를 용이하게 하기 위해 표준 정규 사전 분포(prior)가 부과됩니다.
2. 히스토리 모델링 (History Modeling): Mamba-2 인코더는 궤적 히스토리(과거 행동 및 관측값)를 컨텍스트 벡터($c_t$)로 압축합니다. 이는 노이즈가 있는 관측값이나 제한된 입자(particle) 상황에서도 단순한 마르코프 가정이 놓칠 수 있는 유용한 역사적 맥락을 유지합니다.
3. 결합된 플로우 매칭 (Joint Flow Matching): 플로우 매칭 모델이 통합된 월드 모델 역할을 수행합니다. 이 모델은 잠재 파라미터를 추정하고 해당 파라미터에 조건화된 시스템 역학을 모델링하도록 학습됩니다. 모델은 다양한 조건부 분포를 처리하기 위해 트랜스포머 백본과 보조 마스크를 사용합니다:
   • 파라미터 추정: $p(Z_{t+1} | Z_t, y_t, u_{t+1}, y_{t+1}, c_t)$
   • 롤아웃 생성: $p(u_{t+1}, y_{t+1}, Z_{t+1} | Z_t, y_t, c_t)$
   • 역 역학 (Inverse Dynamics): $p(u_{t+1} | Z_t, y_t, Z_{t+1}, y_{t+1}, c_t)$

훈련 및 추론

• 훈련: 모델은 다양한 파라미터 값을 가진 궤적들이 포함된 오프라인 데이터셋으로 훈련됩니다. 손실 함수는 파라미터에 대한 재구성 손실, 잠재 공간을 표준 정규 분포와 정렬하기 위한 MMD(Maximum Mean Discrepancy) 규제, 그리고 플로우 매칭 손실을 결합합니다. 결정적으로, 모델이 표현 능력을 희생하면서 스스로의 타겟을 최적화하는 것을 방지하기 위해 플로우 매칭 손실로부터의 그래디언트는 파라미터 인코더로 역전파되지 않습니다.
• 플래닝 (배포): 배포 시, PLUME은 입자 집합($Z_t$)으로 표현되는 잠재 파라미터에 대한 신념을 유지합니다. 이는 다음과 같은 후퇴 수평(receding-horizon) 궤적 샘플링 접근 방식을 채택합니다:
  1. 롤아웃 (Rollout): 모델은 현재 관측값, 히스토리, 그리고 입자 신념에 조건화되어 여러 궤적을 자기회귀적으로 샘플링합니다.
  2. 스코어링 (Scoring): 궤적은 디코딩된 보상과 추정된 전이 가능성(transition likelihood)의 조합을 바탕으로 점수가 매겨집니다. 가능성은 역 역학 모델을 사용하여 추정되며, 낮은 가능성은 비현실적인 역학을 보이는 고보상 궤적에 페널티를 부여합니다(Out-of-Distribution 탐지).
  3. 실행 및 업데이트: 가장 높은 점수를 받은 궤적의 첫 번째 행동을 실행합니다. 잠재 신념은 새로운 관측값을 사용하여 온라인으로 업데이트되며, 이를 통해 월드 모델은 재학습이나 미세 조정 없이 실제 역학에 적응할 수 있습니다.

주요 기여

1. 통합 플로우 매칭 월드 모델: 시스템 파라미터를 추정하여 월드 모델 롤아웃을 조건화함으로써, 단일 정책이 다양한 역학에 적응할 수 있게 합니다.
2. 엔드 투 엔드 잠재 표현: 물리적 파라미터(마찰력, 형상)와 보상과 같이 질적으로 다른 파라미터들을 위한 통합된 잠재 공간을 학습하며, 이를 역학 학습에 유용하도록 월드 모델과 공동으로 최적화합니다.
3. 가능성 가중 플래닝 (Likelihood-Weighted Planning): 월드 모델 롤아웃과 추정된 전이 가능성을 활용하여 로봇 행동을 최적화하고, 비현실적인 고보상 계획을 효과적으로 필터링하는 플래닝 전략을 제공합니다.

실험 결과

저자들은 네 가지 시뮬레이션 작업(드라이버 돌리기, 밸브 돌리기, 디스크 튕기기, 버킷 들어올리기)과 한 가지 하드웨어 작업(드라이버 돌리기)에서 PLUME을 평가했습니다.

• 베이스라인: 비교 대상은 Unified World Model (UWM), Flow Q-Learning (FQL), Decision Diffuser (DD)입니다.
• 성능:
  • 시뮬레이션 드라이버: PLUME은 **81.3%의 유효율(validity rate)**을 달성하여 UWM(70.7%), FQL(21.3%), DD(19.3%)를 크게 앞질렀습니다. 또한 최상의 베이스라인 대비 목표 거리를 31% 단축했습니다.
  • 하드웨어 드라이버: PLUME은 시뮬레이션으로부터의 제로샷 전이(zero-shot transfer)를 통해 93.3%의 유효율(30회 중 28회 성공)을 달성하며 UWM(86.7%)을 능가했습니다.
  • 기타 작업: PLUME은 밸브 돌리기 및 디스크 튕기기 작업에서 모든 베이스라인을 크게 앞질렀으며(목표 거리로 측정), DexArt 버킷 들어올리기 작업에서 일관되게 성공한 유일한 방법이었습니다.
• 절제 연구 (Ablations): 논문은 PLUME이 물리적 파라미터 외에도 부분적으로 관측 가능한 상태 변수(위치, 속도, 방향)를 추론하는 능력이 보상 추정의 정확성에 기여하여 성공의 밑바탕이 되었다고 언급합니다.

의의 및 주장

본 논문은 정밀하고 민감한 다지 조작 작업의 경우, 시스템 파라미터의 표현을 명시적으로 추정하고 조건화하는 것이 파라미터에 무관한 정책을 학습하거나 변화를 암시적으로 모델링하는 방법보다 우수한 성능을 보인다고 주장합니다.

주요 주장은 다음과 같습니다:

• 제로샷 전이 (Zero-Shot Transfer): PLUME은 온라인 파라미터 추론을 통해 실제 역학에 적응함으로써, 미세 조정 없이 시뮬레이션에서 훈련된 정책을 하드웨어 작업(드라이버 돌리기)에 배포할 수 있습니다.
• 데이터 품질에 대한 강건성: 이 방법은 하위 최적 궤적과 실패를 포함하는 이질적인 데이터셋으로부터 유용한 모델을 학습할 수 있으며, 보상 레이블을 활용해 실패 데이터를 역학 학습에 사용하면서도 플래닝을 통해 문제가 되는 행동을 피할 수 있습니다.
• 적응성: 재학습이나 미세 조정이 필요한 방법들과 달리, PLUME은 파라미터에 대한 신념을 업데이트함으로써 배포 시점에 월드 모델을 실제 역학에 정렬시켜, 단일 모델이 다양한 물리적 조건에 대응할 수 있게 합니다.

저자들은 한계점도 인정했는데, 구체적으로 훈련 중 정답(ground-truth) 파라미터 값이 필요하다는 점(시뮬레이션에서는 가능하지만 실제 세계에서는 항상 가능하지 않음)과 트랜스포머의 입자 수에 따른 이차적 메모리 스케일링 문제를 꼽았습니다. 이들은 향arly한 연구 방향으로, 정답 레이블 없이도 사전 분포를 정렬하기 위해 사전 훈련된 모델을 실제 데이터로 미세 조정하는 것을 제안했습니다.