GR00T N1: An Open Foundation Model for Generalist Humanoid Robots 리뷰

Background

LLM을 학습할 때는 양이 방대한 web data으로 인해 generalization 능력이 강했고 downstream task에 적응하는 속도도 빨랐다.

하지만 humanoid처럼 로봇 분야에서는 web data처럼 방대한 데이터가 없기 때문에 large scale pre-training이 힘들다.

Open X-Embodiment에서는 이 문제를 해결하기 위해 여러 로봇에서 정보를 모아서 cross-embodied learning이 가능하도록 많은 양의 데이터를 모았지만 아무래도 로봇 종류, 사용하는 센서, DoF 등 너무 다양하기 때문에 이런 데이터를 바로 다른 로봇에 적용하기에는 어렵다.

GROOT N1

GROOT N1은 일반적인 VLA처럼 image와 언어로 된 지시사항이 들어오면 그에 맞는 action을 출력하는 모델이지만 cross embodiment가 가능한 모델이다.

Cross embodiment가 가능한 이유는 dual system compositional 구조를 사용하기 때문이다. 우선 System 2 추론 모듈은 pre-trained VLM이고 카메라가 보는 정보와 지시사항을 이해하는 역할을 한다. System 1 행동 모듈은 Diffusion Transformer으로 일반적인 action expert처럼 action flow matching objective으로 학습이 되고 VLM의 output token과 cross attention한다.

Cross Embodiment 가능

이때 다른 점은 embodiment마다 자기만의 encoder와 decoder을 갖게 해서 embodiment에 맞지 않는 데이터가 들어와도 자기만의 encoder으로 처리하고 embodiment에 맞는 action을 decoder로 생성하는 것이다.

Data 구성

또 다른 점은 data의 구성이다. VLA 학습 데이터를 구축할 때 data 피라미드처럼 구성을 하였다.

Web data와 사람 영상들을 제일 많이 사용하고, 그 다음으로 simulation/augmented data, 그 다음으로 실제 환경에서 모은 로봇 데이터를 가장 낮은 비중으로 사용한다. Web data는 VLA의 prior을 구성하고 simulation, 실제 환경 로봇 데이터들은 real embodiment execution의 강화에 도움을 준다.

학습 방법

그런데 피라미드의 하위층을 구성하는 일반적인 web 영상에서는 visual data만 존재하고 action data는 없다. Web 영상에서는 관절이 얼마나 움직였는지에 대한 action에 대한 data가 없다. 즉, VLA를 학습하기 위해서는 $(\text{state}, \text{observation}, \text{instruction}) \rightarrow \text{action}$과 같은 데이터 구성이 필요한데 없다는 것이다.

이런 action-less data에도 action을 학습하기 위해서 latent action codebook과 inverse dynamics model을 사용한다.

Latent action codebook 영상에서 특정 구간 동안의 프레임 변화를 encoder로 압축해서 action vector를 만들어 flow matching의 action label로 사용하는 것 (LAPA: Latent Action Pretraining from Videos)

Inverse dynamics란 상태 변화를 통해 action을 유추하는 것.

이렇게 하나의 구조로 GROOT N1은 다양한 embodiments에 사용될 수 있다.

GROOT N1 Architecture

이미지와 지시사항을 처리하는 VLM 부분과 DiT block은 모든 embodiment가 다 공유한다. 하지만 robot state와 noised action encoder 그리고 DiT output으로 action을 출력하는 action decoder는 embodiment마다 고유한 모듈이다. 이를 통해서 다양한 embodiment에 적용이 가능한 것이다.

VLM (System 2)

VLM은 SmolLM2와 SigLIP-2 image encoder을 finetune한 모델인 Eagle-2을 기반으로 한다. Encoding된 image와 text token는 VLM의 LLM 모듈에서 같이 encoder되어 입력된다.

학습을 할 때는 task의 text 설명과 여러 이미지를 입력으로 받고 vision language feature을 출력한다. 이때, 빠른 추론 속도와 학습 성공률을 위해 중간 layer의 embedding만 사용한다.

Diffusion Transformer Module (System 1)

Action을 모델링하기 위해서 GROOT N1은 DiT Transformer $V_{\theta}$를 사용한다. $V_{\theta}$는 cross-attention과 self-attention이 interleave된 구조를 가지고 있다. Self-attention에서는 noised action token embedding $A^{\tau}_t$와 state embedding $q_t$가 서로를 attend하고 cross-attention block에서는 VLM의 output embedding $\phi_t$가 action, state embedding과 서로 attend한다.

DiT 마지막 block 뒤에는 embodiment별 action decoder가 붙어서 embodiment에 맞는 action token들을 출력한다.

정답 action $A_t$가 주어졌을 때 flow matching timestep $\tau \in [0,1]$, $\epsilon \sim \mathcal{N}(0, I)$라고 하면 noised action chunk는 다음과 같이 계산된다.

\[\mathbf{A}^{\tau}_{t} = \tau \mathbf{A}_{t} + (1-\tau)\epsilon.\]

최종 모델의 action 예측값은 다음과 같다.

\[V_{\theta} \left( \phi_t, \mathbf{A}^{\tau}_{t}, q_t \right).\]

이 예측값이 정답 action vector field인 $\epsilon - \mathbf{A}_{t}$를 최대한 denoising하는 방향으로 가도록 objective를 설계하였다. 이것이 flow matching이다.

\[\mathcal{L}_{\mathrm{fm}}(\theta) = \mathbb{E}_{\tau} \left[ \left\| V_{\theta}(\phi_t, \mathbf{A}^{\tau}_{t}, q_t) - (\epsilon-\mathbf{A}_{t}) \right\|^2 \right].\]

학습을 다 하고 추론할 때는 action chunk를 반복적으로 생성하기 위해서 forward Euler integration을 사용한다. 처음에는 random noised action $\mathbf{A}^0_{t} \sim \mathcal{N}(0, I)$으로 시작해서 다음 timestep마다 학습된 action vector field 방향으로 점차 추론하는 것이다.

\[\mathbf{A}_{t}^{\tau+\Delta\tau} = \mathbf{A}_{t}^{\tau} + \Delta\tau V_{\theta} \left( \phi_t,\mathbf{A}_{t}^{\tau},q_t \right), \quad \Delta\tau=\frac{1}{K}\]

학습 데이터 생성

Latent Actions

정확하게 web scale dataset으로부터 action labels을 어떻게 얻을까? 즉, video로부터 latent action label을 어떻게 얻을까?

LAPA는 VQ-VAE 모델이 연속적인 image 프레임들로부터 feature을 추출하도록 학습시키는 것이다. Encoder는 현재 frame $x_t$와 미래 시점 frame $x_{t+H}$ (window 크기 H) 로부터 latent action $z_t$을 출력한다. 그럼 Decoder는 $z_t$ 와 $x_t$ 로부터 미래 시점 frame $x_{t+H}$ 을 예측하는 것이다.

이 방법으로 학습을 하고 나면 Encoder는 두 frame을 통해서 두 frame 사이에 일어난 latent action $z_t$ 을 알아낼 수 있고 이 $z_t$는 추론할 때 codebook에서 이미 정의된 action embedding 중 가장 가까운 embedding으로 mapping되어 action token으로 변하게 될 수 있는 것이다.

즉, 이 방법으로 학습된 encoder는 두 시점 frame으로부터 어떤 행동이 일어났는지 알아낼 수 있기 때문에 inverse dynamics model이 되는 것이다.

Web scale로 만든 dataset에서 action label을 만들 때 embodiment을 따로 “LAPA” embodiment라는 가상의 embodiment로 정의해서 넣었다.

이 방법으로 얻은 latent action은 대용량의 web scale 데이터로 얻었기 때문에 generalization 능력이 뛰어나다.

다음 이미지들은 latent action의 generalization 능력을 보여주는 그림이다. 왼쪽 이미지들은 다양한 embodiment에 대해서 오른쪽 팔을 왼쪽으로 움직이라는 지시를 내렸을 때이고 오른쪽 이미지들은 오른쪽 팔을 오른쪽으로 움직이라는 지시를 내렸을 때이다.

Neural Trajectories

Robot data는 사람이 직접 정답 trajectory을 만들어야 하기 때문에 web data처럼 많이 생길 수가 없다고 했는데 video generation model은 이 문제를 완화할 수 있다.

88 시간동안 실험 환경에서 직접 모은 teleoperation data로 video generation model을 finetune했고 그 model의 output으로 827 시간의 데이터를 추가하였다. 이렇게 하면 실제로 teleoperation 하는 것보다 더 다양한 데이터를 모을 수 있다.

이것보다 더 많은 diversity을 얻기 위해서 multimodal LLM object을 사용하여 다양한 조합의 “pick up {object} from {location A} to {location B}” 영상들을 얻었다.

Simulation Trajectories

Simulation data을 만들기 위해서는 적은 human demonstration을 분리하고 이 조각들을 subtask라고 가정하고 이 subtask들을 새 environment들에 적용하여 더 많은 demonstration data을 생성하는 것이다.

Training Details

나중에