Policy Distillation & Generalist Policies

Chapter 4까지 우리는 embodiment마다 별도의 residual RL specialist를 얻었다. Carter(wheeled), H1·G1(humanoid), Spot(quadruped) 각각에 대해 X-Mobility가 뽑는 base action 위에 residual을 얹어 morphology에 맞춘 정책을 학습한 것이다. 문제는 배포다. 로봇 fleet마다 서로 다른 네트워크 가중치를 관리하고, 새 embodiment가 추가될 때마다 독립 파이프라인을 유지하는 것은 시스템적으로 지속 불가능하다. 우리가 진짜 원하는 것은 하나의 정책이 embodiment id만 바꿔주면 여러 몸을 모두 몰 수 있는 generalist다.

COMPASS의 3단계 중 마지막인 policy distillation(정책 증류)은 바로 이 통합을 supervised learning 문제로 바꾼다. 각 specialist를 teacher로 두고, 그들이 남긴 입출력 분포 — (world model policy state \(p\), one-hot embodiment id \(e\), PPO Gaussian action의 mean·variance) — 를 데이터셋으로 기록한다. 그리고 단 하나의 conditional policy \(\pi^{dist}_\theta\)가 \((p, e)\)를 받아 teacher의 action 분포를 재현하도록 KL divergence를 최소화한다. RL의 불안정성과 exploration은 이미 specialist 단계에서 끝났으므로, distillation은 안정적인 회귀 문제로 축소된다.

하지만 이 단계는 동시에 이 강의 후반부(Chapter 8)를 관통하는 averaging problem의 씨앗을 심는다. 서로 다른 몸을 위한 여러 정책을 하나의 조건부 네트워크로 압축할 때, 조건화가 충분히 강하지 않으면 정책은 embodiment 간 행동을 평균 내버린다. 왜 COMPASS가 MSE가 아니라 KL을 쓰는지, 왜 one-hot이 가장 단순하면서도 안전한 출발점인지 — 이 모든 설계 결정이 averaging을 억제하기 위한 방어선이라는 관점에서 이 챕터를 읽어야 한다.

Specialist에서 Generalist로 — Distillation의 데이터 셋업

Policy distillation(정책 증류) = teacher policy(들)의 행동을 하나의 student policy가 supervised로 모방하도록 학습하는 절차. 원래 Hinton의 knowledge distillation을 정책 학습으로 옮긴 것으로, RL의 online 탐색을 offline 회귀로 바꾸는 것이 핵심 이득이다.

COMPASS에서 teacher는 Chapter 4에서 얻은 embodiment별 specialist \(\pi^{(i)}_\phi\)다. 각 specialist는 world model이 뽑은 policy state \(p\)를 입력받아 Gaussian action \(\mathcal{N}(\mu^{(i)}(p), \sigma^2)\)를 출력한다. 여기서 \(\sigma\)는 PPO가 학습 과정에서 유지하는 exploration variance다. Distillation을 위해 우리는 각 specialist를 환경에서 굴리면서(rollout) 매 스텝의 튜플 \((p, e^{(i)}, \mu^{(i)}(p), \sigma^2)\)를 기록한다. 즉 teacher의 '무엇을 봤고, 어떤 몸이었고, 어떤 분포로 행동했는가'를 통째로 저장하는 것이다.

데이터 규모는 embodiment당 320 trajectory × 128 step ≈ 40k frame이다. 이 정도 규모는 4x H100에서 다룰 수 있는 실용적 크기이며, 각 embodiment가 균등하게 대표되도록 설계된다. 중요한 미묘함은 어떤 rollout을 기록하는가인데, COMPASS는 성공한 궤적만 남기는 good-only 필터링이 오히려 marginal한 개선만 준다는 것을 발견했다. 실패 케이스가 담고 있는 corner case — 충돌 직전의 회피, 좁은 통로에서의 재조정 — 가 오히려 학습에 유용하기 때문이다(이 점은 Chapter 4의 covariate shift 논의와 정확히 맞닿는다).

student \(\pi^{dist}_\theta\)의 구조는 teacher와 latent pipeline을 공유하되, 결정적 차이는 입력에 embodiment embedding을 조건으로 받는다는 점이다. 같은 latent pipeline 뒤에 embodiment embedding conditioning을 붙이고, 그 위에 MLP가 mean \(\mu_\theta(p, e)\)를 내고 별도의 global variance \(\sigma_\theta^2\)를 학습한다. 여러 teacher의 지식이 이 하나의 조건부 헤드로 흘러 들어가는 것이 distillation의 물리적 실체다.

KL Distillation 목적함수 — 분포를 통째로 옮긴다

COMPASS의 distillation 목적함수는 모든 embodiment \(E\)와 각 embodiment의 데이터셋 \(D^i\)에 대해 teacher 분포와 student 분포 사이의 KL을 최소화하는 것이다.

\[\min_\theta \sum_{i\in E}\sum_{d^i\in D^i} \text{KL}\big(\mathcal{N}(\mu^{(i)}(p),\sigma^2)\,\big\|\,\mathcal{N}(\mu_\theta(p,e^{(i)}),\sigma_\theta^2)\big)\]

여기서 teacher 분포 \(\mathcal{N}(\mu^{(i)}(p),\sigma^2)\)는 고정된 target이고, student는 자신의 mean \(\mu_\theta(p, e^{(i)})\)와 variance \(\sigma_\theta^2\)를 조정해 이 target에 맞춘다. 조건 \(e^{(i)}\)가 결정적이다 — 같은 policy state \(p\)라도 embodiment id가 다르면 다른 target 분포를 맞춰야 하므로, 네트워크는 embodiment마다 행동을 분기시키도록 강제된다.

두 Gaussian 사이의 KL은 닫힌 형태(closed form)로 계산된다. 1차원 성분에 대해:

\[\text{KL}\big(\mathcal{N}(\mu_1,\sigma_1^2)\,\|\,\mathcal{N}(\mu_2,\sigma_2^2)\big) = \log\frac{\sigma_2}{\sigma_1} + \frac{\sigma_1^2 + (\mu_1-\mu_2)^2}{2\sigma_2^2} - \frac{1}{2}\]

이 식을 뜯어보면 distillation이 무엇을 하는지가 드러난다. 우변의 \((\mu_1-\mu_2)^2\) 항은 student mean이 teacher mean에서 벗어난 만큼 벌점을 주는데, 그 벌점의 크기는 \(2\sigma_2^2\)로 나뉜다. 즉 student가 확신하는(variance가 작은) action 차원일수록 mean 오차에 더 민감하고, uncertain한 차원은 관대해진다. 여기에 \(\log(\sigma_2/\sigma_1)\)와 \(\sigma_1^2/2\sigma_2^2\) 항이 더해져 student의 variance가 teacher의 variance를 따라가도록 만든다.

실무적으로 이 KL 항의 gradient는 안정적이다. teacher가 고정되어 있으므로 target이 움직이지 않고, exploration이나 reward 노이즈가 개입하지 않는다. 이것이 distillation을 'RL의 껍데기를 쓴 회귀'로 만드는 이유이며, 4x H100에서 40k frame/embodiment를 무리 없이 수렴시킬 수 있는 배경이다.

Embodiment Embedding — one-hot의 단순함과 한계

student가 어떤 몸을 몰고 있는지 아는 유일한 통로가 embodiment embedding이다. COMPASS의 가장 단순한 버전은 길이 \(N\)의 one-hot 벡터(\(N\) = embodiment 수)로 각 몸을 표현한다. Carter=[1,0,0,0], H1=[0,1,0,0]처럼 완전히 직교하는 표현이다.

one-hot의 장점은 간섭이 없다는 것이다. 각 embodiment가 임베딩 공간에서 서로 최대로 떨어져 있으므로, 한 몸의 gradient가 다른 몸의 표현을 오염시키지 않는다. 이는 averaging problem에 대한 가장 보수적이고 강력한 방어다. 조건 신호가 명확할수록 네트워크가 embodiment 간 행동을 뭉갤 여지가 줄어든다. distillation처럼 여러 teacher를 하나로 합치는 상황에서, one-hot은 '평균 내지 마라'는 신호를 가장 선명하게 전달한다.

한계는 명백하다. one-hot은 unseen embodiment에 대해 무력하다. 벡터 공간에 여유 축이 없으므로, 학습 시 본 적 없는 새 로봇은 표현할 자리 자체가 없다. 새 embodiment를 추가하려면 임베딩 차원을 늘리고 재학습해야 한다. 이는 Chapter 1에서 언급한 '구조적 분리' 패러다임의 대가로, 확장성을 유연성과 맞바꾼 셈이다.

COMPASS 논문은 대안으로 learnable embedding을 future work으로 지목한다. 각 embodiment를 연속적인 \(d\)차원 벡터로 학습하면, embedding space에서의 보간(interpolation)을 통해 unseen embodiment로 일반화할 가능성이 열린다. 예컨대 Spot과 H1 사이 어딘가에 위치하는 새 4족-2족 하이브리드를 두 임베딩의 보간으로 근사하는 식이다. 하지만 이 유연성은 averaging risk를 다시 불러온다 — 임베딩이 가까운 embodiment들의 행동이 서로 새어 나갈(leak) 수 있기 때문이다. one-hot에서 learned로 가는 것은 단순한 구현 디테일이 아니라, 확장성과 averaging 억제 사이의 근본적 trade-off를 다시 협상하는 일이다.

KL vs MSE — variance를 왜 보존해야 하는가

가장 자연스러운 질문. mean만 맞추면 되지 않나? student mean \(\mu_\theta\)가 teacher mean \(\mu^{(i)}\)를 재현하도록 단순 MSE로 회귀하면 훨씬 간단하다. COMPASS는 KL/All(분포 전체)과 MSE/All(mean만)을 ablation으로 비교했고, MSE가 약간 열등하다 — 특히 H1(humanoid)에서 는 결과를 얻었다.

이유는 variance에 담긴 정보에 있다. teacher specialist의 \(\sigma^2\)는 PPO가 학습하며 남긴 exploration의 흔적이자, 각 action 차원의 민감도 지도다. variance가 큰 차원은 정확한 값이 덜 중요한(로봇이 좀 틀려도 괜찮은) 차원이고, variance가 작은 차원은 정밀해야 하는 차원이다. MSE는 모든 차원을 동등하게 취급해 이 지도를 통째로 버린다. 반면 KL은 앞서 본 \((\mu_1-\mu_2)^2 / 2\sigma_2^2\) 항을 통해 민감한 차원에 자동으로 더 큰 가중치를 준다.

variance 보존이 결정적인 두 번째 이유는 training noise와 architectural 차이의 관리다. distillation은 여러 teacher(구조·학습 노이즈가 제각각인)를 하나의 student 구조로 압축한다. teacher마다 다른 variance를 student가 그대로 흡수하면, 네트워크는 '이 embodiment의 이 차원은 원래 불확실하다'는 사실을 인지한 채 학습한다. 이것이 서로 다른 teacher를 무리하게 하나의 sharp한 점 추정으로 뭉개는 것을 막아준다 — 다시 말해 KL은 averaging의 자연스러운 완충재다.

H1에서 격차가 특히 큰 것은 humanoid의 action 분포가 wheeled보다 훨씬 복잡하고 multimodal에 가깝기 때문으로 해석할 수 있다. 복잡한 분포일수록 mean만으로는 재현이 어렵고, variance 정보가 있어야 그 분포의 '폭'을 올바로 복원한다. 이 관찰은 Chapter 8의 generative policy(diffusion·flow) 논의로 직결된다 — Gaussian variance조차 이렇게 중요하다면, 진짜 multimodal 분포를 다루려면 단일 Gaussian을 넘어서야 한다는 복선이다.

대규모 Generalist 계보 — Open-X/RT-X와의 연결, 그리고 averaging의 씨앗

COMPASS의 generalist 정책은 진공에서 나온 것이 아니다. '하나의 정책으로 여러 로봇을' 이라는 문법에는 대규모 선례가 있다. Open X-Embodiment / RT-X는 22개 로봇, 100만+ trajectory, 60+ dataset, 527개 skill을 하나의 데이터셋으로 통합해 positive transfer — 여러 embodiment를 함께 학습하면 개별 학습보다 나아지는 현상 — 를 대규모로 입증했다. CrossFormer는 900K trajectory·20 embodiment를 단일 transformer로, HPT는 이질적 데이터를 공유 표현으로 묶는 접근을 보였다.

이들과 COMPASS의 근본적 차이는 데이터 스케일링 vs 구조적 분리라는 Chapter 1의 두 패러다임 축에 있다. Open-X/RT-X는 방대한 실제 데이터를 하나의 transformer에 부어 embodiment 구분을 데이터가 알아서 배우게 한다. COMPASS는 반대로 embodiment마다 명시적으로 specialist를 분리해 학습한 뒤 distillation으로 신중하게 합친다. 전자는 스케일에, 후자는 구조에 베팅한다.

이 대비가 중요한 이유는 두 접근이 averaging problem을 서로 다른 방식으로 만난다는 데 있다. 대규모 generalist는 embodiment 조건이 약하거나 데이터 불균형이 있으면 흔한 embodiment의 행동으로 평균이 쏠린다. COMPASS의 distillation은 specialist를 먼저 완성한 뒤 합치므로 각 몸의 '정답'이 이미 확보돼 있지만, 그 정답들을 하나의 조건부 네트워크로 압축하는 순간 — 조건화(one-hot의 강도)와 손실 함수(KL의 variance 보존)가 부실하면 — 여전히 평균으로 무너질 수 있다.

따라서 이 챕터의 모든 설계 결정(KL, one-hot, 실패 케이스 포함)은 하나의 목표를 공유한다. 여러 teacher를 하나로 합치되 각자의 정체성을 지키는 것이다. distillation은 통합의 수단이지만 동시에 averaging의 발원지다. Chapter 8에서 우리는 이 씨앗이 어떻게 자라나며, generative policy가 왜 단일 Gaussian distillation의 한계를 넘는 대안으로 등장하는지를 보게 될 것이다. 지금 기억할 것은 하나다 — generalist를 만드는 순간, 당신은 averaging과의 싸움을 시작한 것이다.