FlowMPC: Improving Flow Matching policies with World Models

기술 요약: FlowMPC

문제 정의

Flow Matching (FM)은 다중 모드(multimodal) 액션 공간에서 강력한 행동 복제(behavior cloning) 방법으로 부상하였으며, 특히 여러 개의 서로 다른 액션 궤적이 동일한 과업을 수행할 수 있는 로봇 조작 분야에서 유용합니다. 그러나 FM은 근본적으로 기대 보상을 직접 최대화하기보다는 전문가의 행동을 모방하도록 훈련된 모방 학습 방법입니다. 따라서 FM 정책은 데이타의 품질과 커버리지에 의해 제한됩니다. 이들은 작은 오류를 복구하거나, 전문가 데이터셋에 잘 나타나지 않은 상태에서 발생하는 누적된 실수(compounding mistakes)를 해결하는 데 어려움을 겪는 경우가 많습니다.

최근 연구들은 RL을 사용하여 FM 정책을 정교화함으로써 이를 해결하려는 시도를 해왔으나, 이러한 접근 방식은 FM 훈련 목적 함수 자체를 수정해야 합니다. 본 논문은 다른 질문을 조사합니다: 학습된 월드 모델(world model)이 FM 훈련 목적 함수를 변경하지 않고도 FM 정책을 개선할 수 있는가? 목표는 모델 기반 계획(model-based planning)이 테스트 시점에 flow 기반 모방 정책을 보완할 수 있는지 확인하는 것입니다.

방법론: FlowMPC 프레임워크

제안된 프레임워크인 FlowMPC는 모방 학습된 FM 정책을 Model Predictive Path Integral (MPPI) 계획 루프 내의 학습된 월드 모델과 통합합니다. 핵심 아이디어는 FM 정책을 후보 액션 시퀀스를 생성하기 위한 구조화된 제안 메커니즘(proposal mechanism)으로 사용하여, 생성된 시퀀스를 월드 모델이 평가하고 정교화하도록 하는 것입니다.

1. 훈련 패러다임

훈련 파이프라인은 다음의 5단계로 구성됩니다:

1. 전문가 데이터 수집: 특권 정보(privileged state information)를 가진 SAC 에이전트가 전문가 데모를 생성합니다.
2. 궤적 수집: 전체 궤적(RGB, 고유 수용 감각, 액션, 보상)을 수집합니다.
3. FM 정책 훈련: 이러한 전문가 궤적에 대해 FM 정책을 훈련하여 다중 모드 액션 사전 분포(prior)를 학습합니다.
4. 시각적 전처리: 전문가 궤적의 RGB 관측치를 FM 정책의 시각 인코더로 전처리하여 압축된 시각적 잠재 표현($z_{viz,t}$)을 생성하며 재수집합니다.
5. 월드 모델 훈련: 전문가 데이터가 프리로드된 리플레이 버퍼를 사용하여 온라인으로 월드 모델을 훈련하며, 전처리된 시각적 잠재 표현을 고유 수용 감각 상태와 결합하여 사용합니다($x_t = [s_t, z_{viz,t}]$).

2. TD-MPC2에 대한 구조적 적응

이 프레임워크는 TD-MPC2를 기반으로 구축되었으나, 학습된 액터(actor)를 모방 학습된 FM 정책으로 대체하였습니다. 이 하이브리드 구조를 수용하기 위해 몇 가지 중요한 수정이 이루어졌습니다:

• 행동 복제 액션 예측기 ($\pi_{BC}$): 현재 모델 상태를 액션으로 매핑하는 전용 네트워크입니다. 이는 계획 롤아웃 중 터미널 가치 부트스트래핑(terminal value bootstrapping)을 위해 사용되어, 가치 함수가 분포 내(in-distribution) 액션을 평가하도록 함으로써 계획 과정을 안정화합니다.
• 통합 표현 인터페이스: 월드 모델은 별도의 인코더를 학습하는 대신, FM 시각적 잠재 표현을 고유 수용 감각 상태와 직접 결합하여 소비합니다. 이를 통해 월드 모델이 동결된 FM 정책의 시각적 특징을 활용할 수 있게 합니다.
• 단순화된 목적 함수: 과업의 이진(binary) 보상 구조를 고려하여, 표준 TD-MPC2의 복잡한 이산화된 symlog 변환 보상 대신 보상을 위한 이진 교차 엔트로피(binary cross-entropy)와 Q-값 회귀를 위한 표준 평균 제곱 오차(MSE)를 사용했습니다.
• 이중 월드 모델 변체:
  • 상태 공간 변체 (State-Space Variant): 결합된 상태와 시각적 잠재 표현의 공간에서 직접 예측합니다. PickCube 과업($H=1$)에 사용됩니다.
  • 잠재 공간 변체 (Latent-Space Variant): 학습된 잠재 표현과 다단계 롤아웃을 따르는 표준 TD-MPC2 방식을 따릅니다. 더 복잡한 PickSingleYCB 과업($H=3$)에 사용됩니다.

3. 추론 및 계획

테스트 시, 플래너는 다음과 같이 작동합니다:

1. 제안 (Proposal): 현재 관측치를 바탕으로 플래너는 $K$개의 후보 액션 궤적을 샘플링합니다. 이 중 일부($K_\pi$)는 FM 정책에 의해 제안되며, 나머지는 무작위로 샘플링됩니다.
2. 롤아웃 및 점수 산정: 각 궤적은 학습된 월드 모델 내에서 롤아웃됩니다. 궤적은 예측된 보상의 합과 행동 복제 예측기에서 유도된 터미널 가치 추정치($\hat{Q}$)에 의해 점수가 매겨집니다.
3. 정교화 (Refinement): MPPI는 이러한 점수를 기반으로 궤적의 가중치를 재설정하고 정교화합니다. 최적화된 시퀀스의 첫 번째 액션이 실행됩니다.
4. 효율성: FM 시각적 인코딩은 관측당 한 번만 계산되며, 모든 샘플 및 통합 단계에서 재사용됩니다.

주요 결과

이 접근 방식은 두 가지 ManiSkill 조작 과업인 PickCube-v1과 PickSingleYCB-v1에서 평가되었습니다.

• 에피소드 종료 성공률 (End-of-Episode Success): 가장 유의미한 이득은 에피소드 종료 시까지 성공을 유지하는 데서 관찰되었습니다.
  • PickCube: 93.14% (FM 단독)에서 97.44% (FlowMPC)로 향상되었습니다.
  • PickSingleYCB: 56.81% (FM 단독)에서 66.41% (FlowMPC)로 향상되었습니다.
• Anytime 성공률 (Anytime Success): 임의의 시점에 성공 상태에 도달하는 측면에서도 개선이 있었으나, 그 이득은 더 작았습니다.
  • PickCube: 95.78% $\to$ 98.68%.
  • PickSingleYCB: 68.77% $\to$ 69.78%.

결과는 월드 모델이 단순히 일시적인 성공 상태에 도달하는 것을 돕기보다는, 작은 액션 오류를 수정하고 과업 완료 근처에서 안정적인 제어를 유지하는 데 특히 효과적임을 나타냅니다.

의의 및 주장

본 논문은 월드 모델 기반 계획이 새로운 FM 훈련 목적 함수를 요구하지 않고도 flow 기반 모방 정책을 효과적으로 보완할 수 있다고 주장합니다.

• 하이브리드 접근 방식: 이 연구는 모방 학습과 모델 기반 RL 사이의 실행 가능한 하이브리드 방식을 보여줍니다. 모방을 통해 강력한 다중 모드 액션 사전 분포를 학습하고, 계획을 통해 추론 시 결정을 정교화합니다.
• 재훈련 불필 এটি 않음: 개선은 새로운 RL 목적 함수로 정책을 재훈련하는 것이 아니라, 테스트 시점에 정책을 더 잘 활용함으로써 발생합니다.
• 제한된 범위: 저자들은 결과가 두 가지 과업에 국한되어 있으며, 최적의 월드 모델 설계(상태 공간 vs 잠재 공간, 계획 지평)가 과업마다 달랐음을 명시적으로 밝힙니다. 그들은 보편적으로 최적화된 구성을 주장하는 것이 아니라, 표현력이 풍부한 생성 정책과 룩어헤드(lookahead) 계획의 결합이 견고한 정책 개선을 위한 유망한 방향임을 시사합니다.
• 성능 베이스라인: 이 방법은 (이미 한 과업에서 TD-MPC2 성능에 근접했던) 강력한 FM 베이스라인을 능가하였으며, 이는 플래너가 단순히 약한 정책을 구제하는 것이 아니라 강력한 정책을 "날카롭게(sharpen)" 만든다는 것을 시사합니다.