SQAM이 흐름 정책 미세조정의 야코비안 곱 비용을 없앤다

무엇인가

흐름 정책(flow policy)은 다양한 행동 분포를 표현할 수 있어 오프라인 및 오프라인-투-온라인 강화학습의 표준 정책 클래스가 됐고, 최근 대형 사전학습 로봇 정책의 백본으로도 쓰인다. 문제는 이런 정책을 학습된 가치 함수에 대해 오프폴리시 RL로 미세조정하는 일이 간단치 않다는 점이다. 흐름 정책은 여러 흐름 단계(flow step)를 거쳐 행동을 생성하는데, 크리틱은 마지막 행동만 평가한다. 기존 방법 대부분은 정책을 얼려두고 그 바깥에 별도의 잔차 정책을 붙이거나 입력 노이즈를 학습하는 방식으로 우회한다. 이 경우 개선의 여지가 얼어붙은 정책의 표현력에 묶이거나, 흐름 정책의 표현력을 스스로 포기하게 된다. 어드조인트 매칭(adjoint matching)은 가치 정보를 마지막 행동에서 각 흐름 단계로 역전파해 흐름 모델 자체를 업데이트하는 이론적 대안이지만, 매 단계마다 정책을 통과하는 벡터-야코비안 곱(VJP)을 요구하므로 비용이 흐름 단계 수와 정책 크기에 비례해 커진다.

어떻게 동작하나

이 논문의 출발점은 경험적 관찰이다. 사전학습된 흐름 정책의 배치 평균 속도 야코비안 J_τ = ∇v_base가 대각 성분에 집중된다는 것이다(논문 Figure 2, 10개 도메인 전체는 Section F.4). 저자들은 이를 근거로 J_τ = c_τ·I 라는 등방성 근사를 취하고, 그 아래에서 린 어드조인트 ODE를 닫힌 형태로 푼다. 그 해는 ã_τ = -τ·exp(2∫c_s ds)·∇Q(s, X_1)로, 모든 흐름 시점에서 최종 행동의 크리틱 기울기와 같은 방향을 가리킨다(Proposition 1). 실무적으로 c_τ를 추정하면 애초에 피하려던 야코비안이 되돌아오므로, 지수 인자를 버리고 â_τ = -τ·∇Q(s, X_1)만 남긴다. 이것이 스칼라 어드조인트다. 샘플링된 행동 X_1에서 크리틱 기울기 한 번만 계산하면 되고, 어떤 흐름 단계에서도 정책을 통과하는 VJP가 필요 없다. 학습 목표는 기존 어드조인트 매칭 손실의 타깃을 스칼라 어드조인트로 바꾼 형태이며, TRQAM의 경로 공간 KL 신뢰영역을 선택적으로 함께 쓸 수 있다.

무엇과 다른가

스칼라 어드조인트만으로는 특히 조작(manipulation) 도메인에서 성능이 잘 나오지 않는다. 정확한 어드조인트는 최종 행동의 기울기를 매 단계 속도 야코비안을 통해 통과시키는 반면, 스칼라 어드조인트는 그 기울기를 훨씬 직접적으로 사용하기 때문에 정책이 생성한 행동에서의 크리틱 오차에 그대로 노출된다. 그래서 저자들은 정책이 생성한 행동 a_π에서의 크리틱 값을 데이터셋 행동 a_data의 값과 비교해 벌점을 주는 항을 추가한다: L_critic += c·(Q(s, sg(a_π)) - Q(s, a_data)). 이 '밸류 페널티'는 액터 업데이트가 이미 뽑아둔 단일 행동에서 작동하므로 추가 샘플링 비용이 없다. 데이터셋 행동에만 크리틱을 맞추는 IQL류의 인샘플 크리틱이나, 행동 간 거리를 벌점으로 주는 ReBRAC식 정규화와 달리, 정책 업데이트가 실제로 기울기를 가져다 쓰는 지점의 크리틱 값을 직접 통제한다는 것이 차이다.

어떻게 쓰나

실험은 OGBench의 보상 기반 단일 과제 50개(10개 도메인 × 5개 과제)에서 진행됐다. 모든 방법이 동일한 흐름 정책에서 출발해(행동 복제 300K 스텝 사전학습) 오프라인 1M 스텝, 온라인은 1.5M 스텝까지 이어서 학습했고, 8개 시드 평균 성공률을 보고한다. 비교 대상은 FQL, CGQL-L, DSRL, IFQL, QAM, QAM-E, TRQAM 일곱 가지다. 결과적으로 SQAM은 50개 과제 전체에서 75%를 기록해 최강 베이스라인 TRQAM의 68%를 앞섰고, 개선은 가장 어려운 네 개 도메인에 집중됐다. antmaze-giant에서 최고 베이스라인 대비 21%포인트, humanoidmaze-large 18%포인트, cube-triple 22%포인트, cube-quadruple 35%포인트를 더했다. cube-quadruple에서는 최강 베이스라인이 19%에 그친 반면 SQAM은 54%에 도달했다. 이 경향은 온라인 단계에서도 유지된다.

전제와 한계

스칼라 어드조인트가 왜 정확한 어드조인트를 대신할 수 있는지는 두 가지 성질로 설명된다. 크기 면에서 두 어드조인트 모두 흐름 시점이 0에 가까워질수록 같은 방식으로 수축해 회귀 타깃의 크기가 비슷하고, 방향 면에서 스칼라 어드조인트는 모든 시점에서 정확한 어드조인트와 양의 정렬을 유지한다. 다만 둘이 동일하지는 않으며, 두 어드조인트 사이의 각도는 조작 도메인에서 이동(locomotion) 도메인보다 크고, 바로 그 도메인에서 스칼라 어드조인트 단독 성능이 떨어진다. 밸류 페널티의 효과도 여기서 갈린다. cube-triple과 cube-quadruple에서 밸류 페널티가 가장 좋았고, 인샘플 크리틱과 ReBRAC은 정규화 없는 스칼라 어드조인트와 비슷한 수준에 머물렀다. 크리틱 기울기 노름에 벌점을 주는 방식도 개선을 보였는데, 밸류 페널티를 a_π 주변에서 1차 전개하면 Q(s,a_π) - Q(s,a_data) ≈ ∇Q(s,a_π)ᵀ(a_π - a_data)가 되어 두 벌점이 같은 기울기를 통해 연결되기 때문이다. 정확한 어드조인트를 쓰는 TRQAM에도 밸류 페널티를 적용하면 개선되지만 폭은 SQAM보다 훨씬 작다. 계수 c는 cube-double을 제외한 모든 도메인에서 0.3을 쓰고, cube-double은 큰 c가 오프라인 성능을 떨어뜨려 페널티를 끈다. KL 예산은 대부분의 도메인에서 좁을수록 좋거나 평탄했고, 상태 공간이 큰 puzzle-4x4만 더 넓은 예산을 선호했다.

대형 사전학습 정책으로의 확장성도 확인했다. 사전학습된 비전-언어-행동(VLA) 정책 RLDX-1을, 20자유도 손 두 개를 장착한 실제 양팔 로봇에서 행동 청크 16으로 미세조정했다. 과제 시연으로 지도학습 미세조정(SFT)을 한 뒤 SQAM으로 오프라인 15K 스텝, 온라인 10K 스텝을 추가로 돌렸다. 비교 대상은 사전학습 정책을 얼려두고 수정 정책을 따로 학습하는 잔차 미세조정 기법 EXPO-FT다. 과제당 30 에피소드의 홀드아웃 성공률에서 SQAM은 세 과제 모두 SFT를 상회했고, EXPO-FT는 같은 세팅에서 잘 작동하지 않았다. 저자들은 SQAM이 흐름 정책을 직접 업데이트하는 반면 EXPO-FT는 처음부터 학습해야 하는 잔차 정책에 의존해 사전학습 정책의 표현력을 활용하지 못하는 점을 원인으로 추정한다.

개발자 관점에서 이 논문이 주는 실무적 시사점은 두 가지다. 첫째, 흐름 기반 정책을 RL로 개선할 때 정책 업데이트 자체를 건드리는 경로가 계산적으로 현실적이 됐다. 단계별 VJP를 없애면 비용이 흐름 단계 수와 정책 크기에서 분리되므로, 큰 사전학습 정책을 소량의 실데이터로 개선하는 파이프라인에 넣기 쉬워진다. 둘째, 병목이 정책 업데이트가 아니라 크리틱일 수 있다는 지적이다. 스칼라 어드조인트를 쓰면 정책이 생성한 행동에서의 크리틱 값이 성능을 좌우하므로, 밸류 페널티 계수 c와 KL 예산 ε_KL을 도메인별로 스윕해 확인해야 한다. 논문도 cube-double처럼 페널티를 꺼야 하는 예외가 있음을 명시한다.

저자들이 밝힌 전제와 한계는 다음과 같다. 스칼라 어드조인트는 정확한 어드조인트의 근사이며, 두 방향 사이의 각도가 조작 도메인에서 더 크다는 사실을 논문 스스로 인정한다. 지수 인자를 버리는 것도 Figure 4(좌)의 경험적 근거에 기반한 선택이다. 밸류 페널티는 TRQAM에도 도움이 되지만 SQAM만큼은 아니다. 실제 로봇 실험은 통제된 실험실 환경에서 운영자가 감독하는 조건으로 수행됐고, 사람 대상이나 개인 데이터는 포함하지 않는다. 논문은 결론에서 어드조인트의 거친 근사로 충분한 반면 성능은 가치 학습 방식이 크게 좌우한다는 점을 강조하며, 흐름 정책을 쓰는 오프폴리시 RL의 병목이 정책 업데이트보다 크리틱에 있을 수 있다는 해석을 남긴다.