RA-VLA: Retrieval-Augmented VLA for Test-Time Adaptation 리뷰

Background

VLA는 학습 데이터에 있는 행동만 잘하고 본 적 없는 manipulation task에서는 잘하지 못한다. 또, 익숙하지 않은 명령을 내릴 때, 학습 과정에서 경험한 익숙한 behavior로 돌아가거나 이상한 행동들을 하는 경향이 있다.

이런 현상이 발생하는 이유는 인공지능 모델이라는 것 자체가 학습된 데이터에 대한 bias을 가지기 때문이다. 이때 보통 VLA를 새로운 환경 / task에 적응시키기 위해서는 추가적인 데이터를 모으고 학습을 해야 한다.

이런 문제를 해결하기 위해서 보통 In-Context Imitation Learning을 통해 test time adaptation을 한다. 재학습 과정 없이 in-context에 expert demonstrations을 넣어서 policy가 weight update 없이도 올바른 행동을 하도록 유도하는 것이다.

하지만 In-Context Imitation도 bottleneck가 존재한다.

1. 우선, in context guidance을 어떻게 precise actions으로 변환할 것인지가 문제이다. 일반적인 retriever는 task을 해결하기 위해 필요한 behavior을 바탕으로 retrieve하지 않고 시각적 similarity을 바탕으로 가져오기 때문이다.
2. In context guidance을 주더라도 VLA가 학습된 behavior에 갇혀서 guidance을 잘 활용하지 못하는 behavioral intertia 문제이다.
3. In context guidance으로 인해 처리해야 하는 내용이 많아져서 생기는 latency 문제가 있다.

이런 문제를 해결하기 위해서 RAG을 도입한다. 우선 Retrieval 자체를 visual similiarity 말고 action aware하게 하여 action 생성에 더 helpful한 정보를 가져오고자 했다

또, 가져온 guidance을 반강제로 사용하도록 하기 위해서 behavioral intertia가 제거되도록 contextual adherence loss을 사용한다.

RA-VLA

Policy $\pi_{\theta}$ 는 timestep t에서 vision input $V_t$, language input $L$, 그리고 robot proprioceptive state $s_t$ 을 받아서 action chunk $A_t \in \mathbb{R^{h \times c}}$ 을 출력한다.

\[A_t = [a_t, a_{t+1}, \ldots, a_{t+h-1}] = \pi_\theta(V_t, L, s_t)\]

VLM $f_{\phi}$ 은 $V_t$ 와 $L$ 을 처리해서 multimodal feature $H_t = f_{\phi} (V_t, L)$ 을 얻는다. Action head는 diffusion transformer으로 noise distribution을 target action chunk으로 refine하는 역할을 한다. 그러기 위해서 vector field $F_t$ = $($$A_t - A_t^{(0)}$) 을 맞추도록 학습된다.

\(A_t^{(k+\Delta k)} = A_t^{(k)} + \Delta k \cdot g_\psi(H_t, s_t, A_t^{(k)}, k)\)

Method

Context Buffer Construction

Retrieval 대상의 expert demonstrations을 저장하기 위해 buffer B을 도입한다. RA-VLA는 long horizon expert demonstrations을 여러 segment ($V’, L’, s’, A’$) 으로 나눠서 buffer에 저장한다. 이 segment들은 단순히 길이 h, stride s의 sliding window을 통해 나뉘어졌다.

이 segment들을 guidance으로 사용하면서 input sequence는 늘리지 않기 위해서 RAG에서 document들을 embedding으로 저장하는 것처럼 각 segment을 VLM으로 미리 처리하여 $H’ = f_{\phi} (V’, L’)$ 을 만들어 buffer 내용은 key-value pair 형태로 저장해놓았다.

Key는 retriever가 만든 $V_t$ 와 $L$ 의 average pooled tokens $z’$이고 Value는 $(H’, s’, A’)$ 으로 구성된다.

Expert Segment Retrieval

Query는 $V_t$ 와 $L$ 정보으로 구성된 것으로 이를 바탕으로 retrieval encoder R이 k개의 가장 연관성 높은 segments $H_{\mathrm{ret}}$을 가져온다. 이때 연관성은 query와 key 간의 cosine similarity으로 정의된다.

\(H_{\mathrm{ret}} = \left\{ H_{\mathrm{ret}}^{(1)}, \ldots, H_{\mathrm{ret}}^{(K)} \right\}\) Retrieval encoder는 two layer Transformer encoder으로 구성된다.

Grounded Action Generation

Action head $g_{\psi}$ 는 retrieve한 segment을 사용해서 augmented vector field $F_t$ 을 생성한다.

\(F_t = g_\psi(H_t, H_{\mathrm{ret}}, s_t, A_t^{(k)}, k)\) Action 생성의 attention 과정에서는 현재 observation feature $H_t$에 expert demonstration의 $H_{ret}^{(i)}$ 을 concatenate한다. 이때 모든 retrieved context을 한번에 넣는 것이 아니라 layer마다 하나씩 추가하였다. 또, score가 높은 context 일수록 더 나중 cross atttention layer에 추가함으로써 더 fine-grained action 생성을 유도하였다.

Action-Aware Retrieval

이때 retrieval이 단순히 시각 정보 + 언어 정보를 바탕으로 되도록 하지 않도록 constrative learning 기반 behavior alignment loss을 사용한다.

이 loss의 objective는 VL context의 embedding space가 behavior pattern이 비슷한 것들끼리 놓이도록 하는 것이다.

Positive pair, 즉 비슷한 behavior pair을 구성하기 위해서 동일한 task의 random한 두 expert demonstration을 사용한다. 이때 두 demonstration은 길이가 다르기도 하고 segment으로 나눴을 때 다른 두 demonstration 간의 segment가 대응이 안 될 수 있다. 따라서 이 둘의 action sequence에 Dynamic Time Warping을 적용하여 둘의 길이가 달라도 대응/비교가 가능하도록 해서 positive pair을 구성한다.

이런 positive pair을 바탕으로 contrastive learning을 하면 retriever가 학습을 할 때 visual observation인 z만 가지고도 이 둘은 비슷한 behavior을 가지고 있다는 것을 유추할 수 있게 되는 것이다. 어떻게 보면 positive pair의 구성으로 인한 강제적인 action aware contrastive learning으로 보인다.

\[\mathcal{L}_{\mathrm{align}} = -\sum_{(i,j)\in P} \log \frac{ \exp\left(\operatorname{sim}(z_i,z_j)/\tau\right) }{ \sum_{k\in\{j\}\cup N_i} \exp\left(\operatorname{sim}(z_i,z_k)/\tau\right) }\]

Learning Context-Grounded Action Generation

Behavior aware guidance을 잘 retrieve해서 제공한다고 하더라도 behavioral inertia 문제 (pretrained 된 대로 행동하는 경향성)로 인해 이 guidance을 잘 사용하지 못 할 가능성이 있다.

Action head가 retrieved expert segments을 더 잘 사용하게 만들기 위해서 contextual adherence loss을 추가한다.

각 training iteration마다 정답 vector field $F_t$ 와의 두 가지 MSE을 계산한다. 하나는 Retrieve된 내용을 가지고 만든 vector field와의 MSE와 다른 하나는 random 내용을 가지고 만든 vector field와의 MSE이다.

\[\mathcal{L}_{\mathrm{rel}} = \left\| F_t - g_{\psi} \left( H_t, H_{\mathrm{ret}}, s_t, A_t^{(k)}, k \right) \right\|_2^2, \qquad \mathcal{L}_{\mathrm{irrel}} = \left\| F_t - g_{\psi} \left( H_t, H_{\mathrm{rand}}, s_t, A_t^{(k)}, k \right) \right\|_2^2\]

이 두 MSE을 가지고 contextual adherence loss을 구성한다.

\(\mathcal{L}_{\mathrm{adhere}} = \max\left( 0,\, m-\left( \mathcal{L}_{\mathrm{irrel}} - \mathcal{L}_{\mathrm{rel}} \right) \right)\) 이 loss는 retrieved된 guidance을 쓰도록 강제한다. Policy가 retrieved된 / random context을 사용하지 않고 현재 observation만 사용한다면 추가된 것은 사용하지 않는다는 것이기 때문에 결국 $\mathcal{L}{\text{rel}}$ 와 $\mathcal{L}{\text{irrel}}$ 은 동일해지기 때문이다.

Policy $\pi_{\psi}$ 을 학습하는 최종 objective는 다음과 같다.

\[\mathcal{L}_{\mathrm{overall}} = \mathcal{L}_{\mathrm{rel}} + \lambda \mathcal{L}_{\mathrm{adhere}}\]

Retrieved된 내용을 사용해 정답 action을 더 정확하게 만드는 loss와 Retrieved된 내용을 사용하도록 강제하는 loss이다.

이때, retrieved 된 내용 중에서 action 값은 flow matching 과정에서 denoising 과정의 noise initialization에 사용된다. 첫 Gaussian noise initialization을 retrieved된 top k actions의 평균으로 함으로써 fast / stable training을 가능하게 해준다.