4 minute read

Background

Robot policy의 generalization 능력을 높이기 위해서 Pretrained VLM의 visual,linguistic representation을 action decoder에 condition 하여 지식을 전달한다.

여기에 더해서 action policy의 input을 action 생성과 align시키기 위해서 다른 연구들도 존재한다.

추론 과정 condition

joowan1108

VLM 내부의 LLM으로 추론 과정(sub task으로 분해)을 guidance으로 action policy에 전달한다.

World Model 사용

joowan1108

World model을 사용하여 미래를 예측한 이미지를 guidance으로 action policy에 전달한다.

하지만 이 방법들의 한계점은 action policy가 모두 vision-language 공간에서 생각을 한다는 것이다. 즉, VLM이 주는 정보는 low level action execution와 gap이 존재한다는 것이다. VLM은 web data로부터 pretrain을 하기 때문에 dynamics에 대한 이해보다는 이미지, 언어 이해에 지식이 치우쳐져 있다.

이 갭을 semantic-kinematic gap이라고 정의한다.

World model도 dynamics을 바탕으로 한 guidance을 주긴 하지만 실제로 action policy에 전달되는 것은 이미지이다.

ACoT-VLA

joowan1108

본 논문은 이런 언어적/시각적 정보가 action sequence guidance에는 최선이 아니라고 주장한다. 그래서 본 논문은 thought라는 것을 action policy에게 전달할 때 기존의 linguistic 정보가 아니라 action intent으로 전달하고자 한다. 이를 통해서 더 직접적인 action 힌트를 주는 것이다.

하지만 이것이 가능하기 위해서는 VLM으로 오는 linguistic,visual 정보를 action intent으로 변환해야 한다. 본 논문은 motion trajectory 정보가 명확히 들어있는 explicit 정보와 linguistic 정보로 얻을 수 있는 함축적인 정보를 사용해서 action space 기반 guidance (action intent)을 생성한다.

문제 정의

자연어 명령 $l$ 과 현재 visual observation $o_t$ 가 있을 때 일반적인 robot policy $\pi_{\theta}$ 는 action sequence $a_{t:t+H-1}$ 을 생성한다.

\[a_{t:t+H-1} = \pi_{\theta} (o_t, l)\]

H는 action horizon이다

최근 연구들은 추가적인 guidance 정보 $g$ 를 통해 policy의 미래 예측 능력을 갖게 하도록 하여 더 정확한 action sequence을 생성하도록 한다. 이때 $g \in {g_{\text{lang}}, g_{\text{vis}}}$이다.

\[\pi_{\theta} (a_{t:t+H-1}, g \mid o_t, l) = \pi_{\theta} (a_{t:t+H-1} \mid o_t, l, g) \pi_{\theta} (g \mid o_t, l)\]

본 논문은 g을 linguistic-visual 정보로 제한하지 않고 action space 기반의 직접적인 정보 $g_{action}$ 까지 추가하여 action policy의 성능을 높이고자 하는 것이다.

Method

joowan1108

Explicit Action Reasoner

생성형 모델에서 prior에 대한 정보를 주면 quality가 높아진다는 원리에 기반하여 kinematic plausible reference action을 action policy에 주입하려고 한다. 이 reference action을 action policy에 주입 가능한 형태를 본 논문은 explicit guidance $g^{ex}_{action}$ 으로 정의하는 것이다.

추론 과정에서 kinematic plausible reference action을 생성하고 action policy에 전달하기 위해서 explicit action reasoner (EAR)을 light weight transformer로 구현한다.

Visual observation $o_t$ 와 언어 명령 $l$ 이 있을 때 pretrained VLM은 contextual key value cache을 생성한다.

\[(\mathbf{K}_{1:N}^{\text{VLM}}, \mathbf{V}_{1:N}^{\text{VLM}}) = \text{VLM}(o_t, l)\]

N은 VLM의 layer 개수이다.

EAR은 이 KV cache와 noisy action sequence $\tilde a_{t:t+H^{ref}-1}$ 을 input으로 받아서 kinematic plausible reference action 생성하는 것이다.

이때 $\tilde a_{t:t+H^{ref}-1}$ 는 transformer의 input으로 들어가기 위해 hidden representation $H^{ref}$ 으로 embedding 된 뒤에 각 transformer의 layer i에서 self attention ($H^{\text{ref}}_{\text{i-1}}$ 끼리)와 cross attention (Contextual KV cache와 $H^{\text{ref}}_{\text{i-1}}$ 간에) 가 진행된다.

\[\tilde{h}_i^{\text{ref}} = \text{Self-Attn}(h_{i-1}^{\text{ref}}) + \text{CrossAttn}(h_{i-1}^{\text{ref}}, \mathbf{K}_i^{\text{VLM}}, \mathbf{V}_i^{\text{VLM}})\]

직관적으로 이해하면 self attention 모듈은 action sequence 간의 시간적 의존성 정보를 capture하고 cross attention은 좋은 action이 되기 위해 어떤 context에 집중해야 하는가에 대한 정보를 capture하는 것이다.

그 다음, $\tilde{h}_i^{\text{ref}}$ 은 residual FFN으로 EAR의 i번째 representation $\tilde h^{ref}_i$ 을 update한다.

\[h_i^{\text{ref}} = h_{i-1}^{\text{ref}} + \text{FFN}(\tilde{h}_i^{\text{ref}})\]

그 다음 EAR $\pi^{\text{ref}}_{\theta}$ 는 flow matching을 통해서 denoised action sequence을 생성하도록 학습된다.

\[\mathbf{a}_{t:t+H_{\text{ref}}-1}^{\text{ref}} = \pi_\theta^{\text{ref}}\left(\tilde{\mathbf{a}}_{t:t+H_{\text{ref}}-1}, \mathbf{K}_{1:N}^{\text{VLM}}, \mathbf{V}_{1:N}^{\text{VLM}}\right)\]

이 생성된 reference action sequence는 MLP projector을 통해서 action embedding $Z^{ex}$ 가 되어 action policy 학습을 위한 explicit action-space guidance $g^{ex}_{action}$ 이 된다.

Implicit Action Reasoner

VLM의 multimodal latent space도 visual 정보나 action과 관련있는 semantics처럼 implicit motion cue를 생성하는 역할을 하는게 implicit action reasoner (IAR)이다.

본 논문은 VLM이 주던 condition이 필요없다는 것이 아니다. VLM의 KV cache의 rich한 정보는 좋지만 action relevant 정보만 추출해서 사용해야 한다는 것이다. 따라서 IAR은 VLM의 KV cache을 직접적으로 사용하여 action 생성에 유용한 정보를 추출한다.

우선 VLM의 각 layer $i \in [1,N]$ 에서 action 생성에 도움이 되는 정보를 추출하기 위해서 query 학습 가능한 matrix $Q_i \in \mathbf{R}^{M \times d}$ (M은 hyperparameter, d는 VLM의 hidden dimension 크기) 를 정의한다.

또, KV cache의 정보는 어떻게 보면 방대하기 때문에 projection matrix $W_Q, W_K, W_V$ 을 사용해서 low dimension으로 축소를 하여 VLM의 KV cache을 응축(?)시킨다.

\[\mathbf{Q}_i' = \mathbf{Q}_i \mathbf{W}_Q^{(i)}, \quad \mathbf{K}_i' = \mathbf{K}_i^{\text{VLM}} \mathbf{W}_K^{(i)}, \quad \mathbf{V}_i' = \mathbf{V}_i^{\text{VLM}} \mathbf{W}_V^{(i)}\]

이때 $\mathbf{W}_Q^{(i)}, \mathbf{W}_K^{(i)}, \mathbf{W}_V^{(i)} \in \mathbb{R}^{d \times d’}$ 이고 $d’ « d$ 이다.

그 다음 cross attention을 통해서 $K’_i, V’_i$ 에서 action relevant 정보를 뽑아온다. 그 다음 average pooling과 MLP projector을 통해서 각 VLM layer에 존재하는 semantics을 compact한 implicit action semantics $z^{im}_i$ 로 변환한다.

\[\mathbf{z}_i^{\text{im}} = \text{MLP}\left(\text{Pool}\left(\text{CrossAttn}\left(\mathbf{Q}_i', \mathbf{K}_i', \mathbf{V}_i'\right)\right)\right)\]

각 layer에서 얻는 representation들을 aggregating해서 implicit action related feature $Z^{im}$ , 즉 implicit action-space guidance $g^{im}_{action}$ 이 된다.

Action Guided Prediction

joowan1108

이렇게 EAR와 IAR으로 얻은 $Z^{ex}$ 와 $Z^{im}$ 으로 효과적인 action guidance을 구성하고 policy learning에 condition 하는 방법은 action guided prediction이라고 정의한다.

본 논문의 action policy $\pi_{\theta}^{\text{head}}$ 추론 과정은 기존의 연구 방법들과 다르다.

정답 action segment + noise가 섞인 action segment $\tilde a_{t:t+H-1}$ 이 있을 때, 이 action을 MLP projector로 embedding 한 뒤에 action query $Q_{\text{action}}$ 로 사용한다. 이 action query로 $Z^{ex}$ 와 $Z^{im}$ 을 참조하여 action에 필요한 정보를 dual cross attention으로 추출하는 것이다.

\[\mathbf{S}^{\text{ex}} = \text{CrossAttn}(\mathbf{Q}_{\text{action}}, \mathbf{Z}^{\text{ex}}, \mathbf{Z}^{\text{ex}})\] \[\mathbf{S}^{\text{im}} = \text{CrossAttn}(\mathbf{Q}_{\text{action}}, \mathbf{Z}^{\text{im}}, \mathbf{Z}^{\text{im}})\]

$\mathbf{S}^{\text{ex}}$ 은 kinematic 정보를 알려주는 것이고 $\mathbf{S}^{\text{im}}$ 은 context로부터의 행동의 방향성 정보를 알려주는 것이다.

이 복합적인 guidance을 합쳐서 사용하기 위해서 얻은 feature들을 concatenate하고 self attention fusion block을 통해 하나의 representaation $\bar{\mathbf{h}}$ 을 얻는다.

\[\bar{\mathbf{h}} = \text{Self-Attn}\left(\left[\mathbf{S}^{\text{ex}} ; \mathbf{S}^{\text{im}}\right]\right)\]

이 $\bar{\mathbf{h}}$ 는 이제 $\pi_{\theta}^{\text{head}}$ 의 입력으로 들어가서 denoised action sequence $a_{\text{t:t+H-1}}$ 을 예측하게 된다.

Training Objective

전체 framework는 flow matching MSE objective으로 학습된다. Training loss는 EAR $\pi_{\theta}^{\text{ref}}$ 와 Action head $\pi_{\theta}^{\text{head}}$ 의 flow matching loss으로 구성된다.

\[\mathcal{L}_{\text{total}} = \lambda_1 \mathcal{L}_{\pi_\theta^{\text{ref}}} + \lambda_2 \mathcal{L}_{\pi_\theta^{\text{head}}}\]

정리하자면 $\pi_{\theta}^{\text{ref}}$ 의 output을 reference action으로 사용하여 explicit guidance을 생성하고 이 explicit guidance을 바탕으로 action head가 더 효과적인 action을 생성하는 것이다. 하지만 $\pi_{\theta}^{\text{ref}}$ 의 처음 성능은 좋지 않기 때문에 초기의 $\pi_{\theta}^{\text{ref}}$ 을 바탕으로 explicit guidance을 얻는다면 이 guidance의 quality가 떨어질 것이기 때문에 action head의 성능도 떨어질 것이다. 따라서 본 논문은 Teacher Forcing Stabilization을 사용한다. 학습할 때는 reference action을 $\pi_{\theta}^{\text{ref}}$ 의 output이 아니라 정답 action으로 두고 explicit guidance을 구성한다. 추론할 때는 $\pi_{\theta}^{\text{ref}}$ 의 성능이 향상되었을 것이므로 이때는 $\pi_{\theta}^{\text{ref}}$ 가 explicit guidance을 구성한다.

You May Also Enjoy