https://arxiv.org/abs/2310.12921
Vision-Language Models are Zero-Shot Reward Models for Reinforcement Learning
Reinforcement learning (RL) requires either manually specifying a reward function, which is often infeasible, or learning a reward model from a large amount of human feedback, which is often very expensive. We study a more sample-efficient alternative: usi
arxiv.org
비전 기반 환경에서 RL을 훈련시키는 데 있어 보상 함수를 정의하는 비용이 매우 높다. 실제 환경에서 보상 함수를 수작업으로 설계하기 힘들고 human feedback을 이용하는 경우 비용이 크다. 따라서 VLM 같은 사전학습된 비전-언어 모델을 보상 신호로 사용한다. 기존 연구들은 VLM을 보상 모델로 사용하려면 복잡한 파인튜닝이나 ad-hoc procedure를 필요로 한다.
본 연구는 VLM-RM 방식을 제안한다. CLIP을 사용하여 현재 상태 이미지와 언어 목표의 임베딩 코사인 유사도로 보상을 계산한다. 고전 RL 환경인 CartPole과 MountainCar 벤치마크에서 검증한다. CLIP 기반 보상이 ground truth 보상과 높은 상관관계를 보인다.
Partially observable Markov decision processes
POMDP $(S, A, \theta, R, O, \phi, \gamma, d_0)$ 튜플로 구성된다.
- $S$ 상태 공간 : 환경이 가질 수 있는 모든 상태들의 집합
- $A$ 행동 공간 : 에이전트가 취할 수 있는 행동의 집합
- $θ(s′|s, a)$ 전이 확률 함수 : 현재 상태 $s$에서 행동 $a$를 취했을 때, 다음 상태 $s'$로 전이될 확률
- $R(s,a,s')$ 보상 함수 : 상태 $s$에서 행동 $a$를 하고, 상태 $s'$로 전이되었을 때 얻는 보상 값
- $O$ 관측 공간 : 에이전트가 직접 관찰할 수 있는 정보들의 공간
- $ϕ(o|s′)$ 관측 확률 분포 : 상태 $s'$에 있을 때, 관측 $o$를 얻을 확률
에이전트는 매 시간마다 상태 s ∈ S에 있으며, 행동 a ∈ A를 선택한다. 선택한 행동에 따라 새로운 상태 $s'$로 전이된다. 에이전트는 새 상태 $s'$로부터 관측 $o$를 얻으며, 이 관측은 확률 분포 $ϕ(o|s′)$에 따른다.
에이전트는 할인된 총 보상을 최대화한다.
$$ g(\tau; R) = \sum_{t=0}^{T} \gamma^t R(s_t, a_t, s_{t+1})$$
여기서 $T$는 trajectory의 길이이며, 우리는 유한한 길이의 trajectory만 고려한다. (즉, $|τ| < ∞$)
에이전트의 목표는 정책 $\pi$에 따라 생성된 모든 trjectory에 대해 기대 보상 $G(\pi)$를 최대화하는 정책을 학습하는 것이다.
$$G(\pi) = \mathbb{E}_{\tau(\pi)}[g(\tau; R)]$$
Vision-Language Models AS Reward Models (VLM-RMs)
Using Vision-Language Models as Rewards
$$(S, A, \theta, O, \phi, \gamma, d_0)$$
본 논문은 VLM을 사용해 보상 $R(s)$을 정의한다.
- \( \psi(s) \): 상태 \( s \)를 이미지로 매핑한 관측
- \( l \): 작업의 자연어 설명
- \( c \): 선택적인 컨텍스트(context, 상황 정보), 예: 사용자가 VLM과 상호작용하며 보상을 정의할 수 있도록 도와주는 추가 문맥
CLIP을 VLM으로 사용한다. 언어 벡터와 이미지 벡터 사이의 코사인 유사도가 보상이 된다.
$$R_{\text{CLIP}}(s) = \frac{\text{CLIP}_L(l) \cdot \text{CLIP}_I(\psi(s))}{\|\text{CLIP}_L(l)\| \cdot \|\text{CLIP}_I(\psi(s))\|}$$
Goal-Baseline Regularization to Improve CLIP Reward Models
관측에서 관련 없는 정보를 제거함으로써 보상의 품질을 향상한다. 앞에서 만든 CLIP 보상 함수는 단순히 "상태 이미지"와 "작업 설명" 사이의 코사인 유사도만 계산한다. 즉, 이미지가 "로봇이 무릎 꿇는 것처럼" 보이면 점수를 높게 주는 방식이다.
해당 보상 함수는 CLIP 이미지 임베딩 안에 작업과 관계없는 다른 요소들이 불필요하게 들어가서 잡음이 생긴다. 따라서 "목표 상태"와 "기준 상태"를 비교해서 변화된 부분만 뽑는다.
이를 위해서 2번째 baseline 설명 $b (b∈L≤n)$가 필요하다. $b$는 환경의 기본 상태를 나타내며, 목표나 적합과 무관하게 환경 자체에 대한 일반적인 설명이다.
- 목표 설명 $I$ : “a humanoid robot kneeling” (무릎 꿇은 로봇)
- 기준 설명 $b$ : “a humanoid robot” (일반 로봇)
상태 임베딩을 기준과 목표 임베딩이 이루는 선에 project 시켜, 관측 표현에서 불필요한 성분을 제거한다. 이렇게 정의된 보상 함수를 R_CLIP-Reg라고 한다.
- 목표 임베딩 \( g = \frac{\text{CLIP}_L(l)}{\|\text{CLIP}_L(l)\|} \)
- 기준 임베딩 \( b = \frac{\text{CLIP}_L(b)}{\|\text{CLIP}_L(b)\|} \)
- 상태 임베딩 \( s = \frac{\text{CLIP}_I(\psi(s))}{\|\text{CLIP}_I(\psi(s))\|} \)
이때, 기준과 목표가 이루는 선분을 $L$이라 하고, 보상 함수는 다음과 같이 정의한다.
$$R_{\text{CLIP-Reg}}(s) = 1 - \frac{1}{2} \left\| \alpha \cdot \text{proj}_L(s) + (1 - \alpha) \cdot s - g \right\|^2_2$$
벡터 \( b \rightarrow g \)는 기본 상태 → 목표 상태 방향이다. 현재 상태 임베딩을 $g - b$ 방향으로 투영해서 목표 상태에 얼마나 가까운지를 계산한다.
- α = 0일 경우, 기존의 기본 CLIP 보상 함수 \( R_{\text{CLIP}} \)와 동일
- α = 1일 경우, 상태 벡터 \( s \)에서 \( g - b \)에 수직인 성분이 완전히 제거
RL with CLIP Reward Model
- 환경과 상호작용하면서 얻은 관측 데이터를 replay buffer에 저장한다.
- replay buffer에서 관측 데이터 배치를 샘플링한다.
- 샘플링된 관측을 CLIP 인코더에 입력하여, 해당하는 상태 임베딩을 생성한다.
- 보상 계산
- 코사인 유사도를 계산하여 보상으로 사용
- 작업 임베딩은 학습 내내 한 번만 계산
- 계산된 보상을 바탕으로, 표준 RL 업데이트를 수행
즉, 기존의 reward function만 바뀌었을 뿐, 학습 루프 구조는 일반적인 RL과 동일하다.