[RLHF] Preference Transformer: Modeling Human Preferences Using Transformers For RL

Preference Transformer: Modeling Human Preferences Using Transformers For RL

 

Link : https://arxiv.org/pdf/2303.00957.pdf

 

Abstract

 

사람의 선호로도를 베이스로한 강화학습 방법은 에이전트가 두 behavior 간의 사람의 선호도를 학습할 수 있게 해준다. 하지만, Preference-based RL은 사람의 의도에 관한 reward function을 학습하기 위해 엄청나게 많은 양의 human feedback 데이터가 필요하다. 해당 논문은 트랜스포머를 사용해 사람의 선호도를 모델링하는 뉴럴넷 Preference Transformer를 제안한다. 이전 방법들은 사람의 판단이 의사결정을 내리는데 동등하게 기여하는 Markovian reward에 기반한다고 가정하지만, 해당 논문에서는 weighted sum of non-Markovian reward를 사용하는 모델을 소개한다. Preference Transformer(이하 pt)는 이전 연구들이 해결하지 못했던 진짜 사람의 선호도를 사용하는 다양한 control task를 풀 수 있으며 잘 정의된 보상을 유도한다. 또한 사람의 의사결정에 있어 시간적 의존성을 자동으로 포착함으로써 trajectory의 중요한 event를 attend 할 수 있다.

 

 

Introduction

 

RL은 적절한 보상 함수가 있는 다양한 도메인의 sequential decision-making 문제를 푸는데 성공하였다. 그러나 많은 문제들은 보통 여러개의 objective를 가지며 이를 싱글 스칼라 보상으로 표현하기는 어렵고, 사람의 수고와 광범위한 task 지식이 요구된다. 이러한 reward engineering 문제를 피하기위해, Language instruction, Preferences, Expert demonstrations, real-valued feedback 등의 사람 데이터로부터 보상 함수를 학습하는 다양한 방법들이 있다. 특히 상대적으로 판단(e.g. pairwise comparison)은 비교적 제공이 용이하면서도 정보가 풍부하여 최근에 Preference-based RL 연구에 대한 관심이 증가하고 있다.

 

그러나, 현재까지의 접근법들은 여전히 많은 양의 human feedback 데이터가 필요하며, 이는 preference-based rl이 다양하게 적용될 수 있도록 scale up하는 것을 힘들게 한다. 그래서 해당 저자들은 이 어려움이 이전의 대부분 preference modeling 연구의 일반적인 가정 때문이라고 보았다. 이 가정은,

 

(a) 보상 함수는 Markovian이다. 즉, 오직 현재 상태와 행동에만 의존한다.

(b) 사람은 trajectory의 퀄리티(에이전트의 행동)을 같은 weight를 가지는 reward의 합을 통해 평가한다.

이다.

 

그리고 저자들은 아래와 같은 이유로 위 가정들이 틀렸다고 지적한다.

 

(1) 모든 task-relevant 정보를 state로 인코딩하는 것은 어렵기 때문에, 방문했던 state들에 의존하는(non-Markovian) 보상함수가 필요한 다양한 task들이 존재한다. 그리고 이는 특히, 이전의 이벤트가 나중 이벤트의 rating에 영향을 줄 수 있게 하는 사람의 sequentiality가 trajectory segment에 제공되기 때문에, preference-based learning에서는 사실이다.

(2) 게다가 사람은 remarkable moments에 매우 민감하기 때문에 trajectory의 credit assignment가 요구된다. 예를 들어서, 아이 트래커를 사용한 비디오게임에서 사람의 집중에 관한 연구에서 사람은 큰 보상이나 패널티를 받을 수 있는 중요한 state에서 의사결정을 내리기 위해 더 긴 반응 시간과 multiple eye movements를 요구했다.

 

해당 논문에서는 앞에서 설명한 이전 연구들의 일반적인 가정의 한계를 극복할 수 있는 alternative preference model을 제안한다. 여기서 제안하는 preference model은 weighted sum of non-Markovian reward에 근거하며 이는 사람의 결정에서 시간적 종속성을 더 잘 포착할 수 있고, trajectory에서 중요한 이벤트를 추론할 수 있다. Preference Transformer는 trajectory segment를 input으로 받고, task와 연관있는 히스토리컬 정보를 추출한다. bidirectional, causal self attention layer를 쌓음으로써 PT는 non-Markovian reward와 importance weight를 ouput으로 생성한다.

 

저자들이 주장하는 해당 논문의 contribution은 다음과 같다.

• non-Markovian 보상함수에 근거해 사람의 선호도를 모델링할 수 있는 general 프레임워크를 제안
• 제안된 프레임워크를 위해 디자인된 새로운 preference attention을 포함하는 트랜스포머 베이스 모델구조인 Preference Transformer를 소개
• Preference Transformer는 real human preference 로부터 학습한 보상 함수를 통해, RL 에이전트가 복잡한 네비게이션, 로코모션, 로보틱 manipulation task들을 풀 수 있도록 함
• learned reward function과 importance weights 로부터 PT는 잘 정의된 보상을 도출 할 수 있고, trajectory에서 중요한 이벤트를 포착할 수 있다는 것을 분석함

 

 

Related Work

Preference-based reinforcement learning

RL 에이전트가 reward engineering 없이 사람의 선호를 활용할 수 있도록 하는 다양한 방법들이 최근 제안되었다. 해당 연구들은 preference-based RL이 효과적으로 복잡한 control task를 딥 뉴럴넷을 통해 풀 수 있다는 것을 보여주었다. 그리고 feedback-efficiency를 향상시키기 위해서 pre-training, data augmentation, exploration, meta-learning 같은 방법들도 제안되었다. Preference-based RL은 또한 어려운 tasks들에 대해 large-language model을 파인 튜닝하는 것도 성공하였으며, 이러한 접근법에서는 language input의 partial observability 때문에 보상함수를 pre-trained language model을 사용하였다. 하지만 해당 논문에서는 transformer를 control task를 풀기 위해 new preference modeling 하는데 사용하였다.

 

Transformer for reinforcement learning and imitation learning

트랜스포머는 RL에서 다양한 목적으로 연구되었다. 복잡한 환경(e.g. StarCraft, DMLab-30)에서 RL 에이전트에 트랜스포머를 사용하면 sample efficiency, generalization ability를 향상시킬 수 있다는 것을 보여준 연구도 있고, Offline RL에서는 RL 문제를 sequential modeling의 context로 formulating함으로써 트랜스포머를 이용한 연구들도 있다. 게다가 트랜스포머는 imitation learning에도 성공적으로 적용되었는데, one-shot imitation learning에 트랜스포머의 generalization ability를 이용한 것도 있고, multi-domain 또는 multi-task imitation learning을 위해 트랜스포머를 활용한 연구도 있다. 해당 연구에서도 또한 트랜스포머가 preference-based learning에도 효과적임을 보이겠다.

 

Non-Markovian reward learning

Non-Markovian reward는 에피소드 끝에서 보상이 주어지거나, 보상이 딜레이가 있는 경우, 방문했던 상태에 의존한 보상 등의 realistic reward setting을 다루기 위해 연구되었다. 최근 몇 연구들에서는, state-action 시퀀스가 주어졌을 때 trajectory의 return을 예측하고, reward redistribution을 위한 objective에서 학습된 추가적인 모델을 training하는 return decomposition에 초점을 맞추고 있다. LSTM을 reward learning에 사용해 sequential information을 포착하거나, 어떤 사전 지식 없이 에피소딕 리턴은 매 타임스텝마다 유니폼하게 분포도있다는 가정을 하는 간단한 접근법 등이 있다. 해당 논문에서는 non-Markovian reward를 학습하기 위해 transformer-based reward model을 사용해 human decision의 temporal dependencies을 포착한다.

 

 

Preliminaries

먼저, 해당 논문에서는 에이전트가 discrete time에서 상호작용하는 환경에서의 RL framework에 대해 고려한다. 매 타임스텝 t에서, 에이전트는 현재 상태 s_t를 환경으로부터 받고, policy \pi에 의한 행동 a_t를 선택한다. 환경은 reward r(s_t, a_t)와 에이전트 transition로부터 다음 상태 s_t+1을 준다. RL의 목표는 expectied return R_t=\Sum_k=0^\infin(\gamma^k r(s_t+k, a_t+k)를 최대화 하는 policy \pi를 학습하는 것이다. 

많은 애플리케이션에서, 사람의 의도를 포착하는 적절한 보상 함수를 디자인하는 것은 어렵다. Preference-based RL은 사람의 선호로부터 보상 함수를 학습하는 이슈에 대해 설명한다. 이전 연구들과 비슷하게 해당 연구에서도 두개의 trajectory segments가 있는 상황을 고려하며 trajectory segment 길이 H에 대해 trajectory segment는 \sigma={(s_1, a_1), ... (s_H, a_H)}로 표기한다. 세그먼트 페어 (\sigma^0, \sigma^1)이 주어졌을 때, teacher(human)은 어떤 세그먼트를 더 선호하는지 선택하고 이 preference label y는 0, 1, 0.5 값을 가질 수 잇다. y=1이면 세그먼트1 (\sigma^1)을 세그먼트 0보다 선호하는 것을 의미하고 0은 그 반대이며 0.5는 동일하게 선호함을 의미한다. 그래서 각 피드백은 (\sigma^0, \sigma^1, y) 튜플로 저장된다.

이전 연구들과 동일하게 보상함수 \hat{r}은 \psi로 파라미터라이즈된 함수이며, preference predictior는 아래 식과 같이 Bradly-Terry model로 정의되고 Bradly-Terry model은 두 페어의 결과로 전체 랭킹을 예측하는 모델이다.

 

 

위에서 튜플로 정의된 preference dataset D가 주어졌을때, 리워드 함수 \hat{r}은 preference predictor와 실제 human label의 cross entropy term을 줄이는 objective로 업데이트 된다.

이렇게 만들어진 learned reward model을 통해서 expected return을 최대화하도록 모든 RL 알고리즘을 사용해서 policy pi를 업데이트 할 수 있다.

 

 

Preference Transformer

위는 PT의 프레임워크 Figure이다. 두 trajectory segments사이에 preference가 주어지면, PT는 non-Markovian reward \hat{r_t}와, 각 세그먼트에 importance weights w_t를 생성한다. 그리고 weighted sum of non-Markovian reward(i.e. \Sum_t w_t \hat{r_t})에 의해 preference predictor를 모델링하고, 사람의 선호도를 align한다.

 

Preference Modeling

위에서 언급한 것 처럼 대부분의 이전 연구들은 보상을 Markovian으로 가정하고, 사람이 trajectory의 quality를 평가할 때 동등한 weight를 가진다고 가정한다. 이 가정에 근거하면 preference predictor는 위의 eq1처럼 정의된다. 하지만, 이와 같은 formulation은 진짜 사람의 선호를 모델링하는데 몇 몇 한계를 가지고 있다. 많은 경우에 Markovian reward를 사용해 task를 specify하는 것은 어렵고, 사람은 remarkable moments에 민감하기 때문에, trajectory에서 credit assignment이 요구된다.

따라서 해당 논문에서는 segment를 선호할 확률이 weighted sum of non-Markovian reward에 exponentially 의존하는 새로운 preference predictor를 제안한다. 

temporal depencies를 포착하기 위해, 앞 전체의 sub-trajectory를 input으로 non-Markovian reward function \hat{r}를 고려한다. importance weight w는 또한 전체 trajectory segment {(s_i, a_i)}_i=1^H의 함수이고, 이는 preference predictor가 세그먼트 안에서 credit assignment를 할 수 있게 한다. 그리고 이는 conventional design의 generalized version이다. 만약 보상 함수가 오직 현재 state-action pair에 의존하면 importance weight w는 언제나 1이고 그렇게 되면 eq1의 standard version과 동일하다.

 

Architecture

eq3의 Preference predictor를 모델링하기 위해 새로운 트랜스포머 아키텍쳐 Preference Transformer를 제안하하며 PT는 Causal Transformer, Preference attention layer로 이루어져있다.

위 figure는 전체 PT이고, x_t를 causal transformer를 통해 hidden embedding하고, 이는 initial timestep부터 timestep t까지의 context information을 나타낸다. 그리고 preference attention layer는 bidirectional self-attention으로, hidden embedding of x_t로부터 non-Markovian reward \hat{r}과 그들의 convex combinations z_t를 계산한다. 그리고 z_t를 모아서 weighted sum of non-Markovian reward \Sum_t w_t\hat{r_t}를 모델링한다.

 

Causal Transformer

GPT를 backbone으로 사용하고 H length의 trajectory segment \sigma가 주어지면, input은 segment pair가 되므로 2H의 input embedding(for state, action)을 가지며 이 input embedding은 linear layer와 layer normalization을 거치고, 같은 타임스텝의 state, action은 같은 positional embedding을 공유하고 이는 input embedding에 더해진다. input embedding은 causal transformer network에 들어가 output embedding을 만든다.

 

Preference attention layer

preference attention layer는 causal transformer에서 나온 hidden embedding x_t를 받고, reward \hat{r} 과 importance weight w_t를 생성한다. t번째 input x_t는 linear transformation하여 key k_t, query q_t, value \hat{r}으로 매핑되고, hidden embedding x_t는 trajectory segment의 이전 입력들에만 으존하기 때문에 self-attention t 번째의 value, 즉 \hat{r}은 non-Markovian reward를 모델링하는 것으로 여긴다. 특히 t 타임스텝에서, 우리는 t개의 state-action pairs를 reward \hat{r_t}를 approximation하는데 사용한다.  그리고, self-atteion 연산은, i번째 ouput z_i는 i번째 key, query에서 attention weight가 곱해진 value의 covex combination으로 정의된다.

그런 다음, weighted sum of reward는 outputs z_i 들의 평균값으로 계산된다.

w_t는 eq3를 따르는 importance weight로 모델링 된다. 왜냐하면 preference attention은 causally masked되지 않고 또한 full sequence에 의존하지 않기 때문이다.(i.e. bidirectional self-attention). 요약하자면 해당 논문은 트랜스포머 네트워크로부터 나온 평균 output을 받음으로써 weighted sum of non-Markovian reward를 모델링하고, 세그먼트를 선호할 확률은 이에 비례한다고 가정한다.

 

Traning and Inference

Training. eq2의 cross entropy loss를 줄이는 방향으로 PT를 학습하고, PT를 통해 모델된 preference predictor를 human label로 aligning함으로써, 우리는 PT가 적절한 보상 함수와 trajectory segments에서의 중요한 이벤트를 포착할 수 있다.

Inference. RL training에서, 모든 state-action pairs는 learned reward function으로 라벨링되어 있다. non-Markovian 보상 함수를 학습했기 때문에, 우리는 H개의 past transitions을 PT에 제공하고, preference attention layer를 통해 받은 t번째 value \hat{r_t}를 timestep t의 reward로 사용한다.

 

 

Experiements

실험은 아래 내용을 확인하기 위해 디자인 되었다.

• PT는 사람의 선호도를 사용해 복잡한 control task 문제를 풀 수 있는가?
• PT는 well-aligned reward를 도출해내고, 중요한 이벤트를 attend 할 수 있는가?
• PT는 종합적인 선호도를 봐야하는 세팅에서 얼마나 잘 동작하는가?

Setups

실험 환경은 complex control tasks로 D4RL과 Robomimic이 사용되었다. reward learning을 위해 query 즉 trajectory segments는 offline dataset에서 uniformly random select하였고, preference는 real human trainer가 collect하였다. 이를 가지고 reward function을 학습하고, IQL을 사용해 RL agent를 학습하였다. IQL은 특히 D4RL benchmarks에서 강력한 성능을 보였기 때문이라고 설명했다. Markovian policy와 value function은 원래의 IQL implementation을 따랐다. 그리고 baseline model은 Markovian based model과 Non-Markovian based model로 나눴고, MR-based model preference predictor는 

으로 reward function을 학습했고 MLP를 사용, NMR-based model preference predictor는

으로 reward function을 학습하고 LSTM을 사용했다. 두 모델 모두 sum of reward에는 equal weight를 가진다.

 

Benchmark tasks with real human teachers

위 실험 결과는 IQL이 각각 다른 rewad function으로 학습된 성능을 나타낸다. PT는 지속적으로 거의 모든 task의모든 baseline을 능가하는 성능을 보여주며 또한 오직 PT만 IQL이 task reward로 학습된 성능과 비슷한 걸 볼 수 있다. 이는 PT가 real human preference를 가지고 적절한 보상 함수를 도출해낼 수 있으며 의미있는 behavior를 가르칠 수 있다는 것을 암시한다. 그리고 특히 복잡한 task인 AntMaze같은 경우는 baseline성능과 큰 차이로 우세한데, 이는 PT가 에이전트의 behavior의 긴 context와 중요한 이벤트를 포착할 수 있고 이는 특히 복잡한 task에서 중요하기 때문이다.

 

Reward and Weight Analysis

PT가 잘 정의된 보상을 도출하고, 사람의 선호도로부터 중요한 이벤트를 포착하는지 검증한다.

위 figure는, 성공, 실패한 trajectory segment의 learned reward function, importance weight를 나타낸다. 에이전트가 골에 가까울 수록 reward value가 증가하는 반면 에이전트가 코너를 탈출하려고 고군분투 하고 있을 때는 reward가 낮다. 이는 PT가 리워드 함수에서 에이전트의 behavior의 context를 잘 포착함을 나타낸다. 그리고 importance weight는 보상 함수와 비교했을 때 다른 트렌드를 보여준다. 흥미롭게도, 성공, 실패한 trajectory모두 중요한 이벤트(e.g. turning right to reach the goal or fliping)에 높은 importane weight가 할당된다. 

 

Benchmark Tasks with Scripted Teachers

이전 연구 https://arxiv.org/abs/2106.05091 와 비슷하게, PT는 scripted teacher로부터 얻은 synthetic preferences를 사용해 평가 되었다. 여기서 teacher는 deteriministic teacher로, benchmark로부터 나온 task reward로 preference를 생성한다. 즉 task reward가 높은 segment를 선호한다고 평가하는 것이다. 그리고 이 scripted teacher는 eq1에서소개한 preference model의 special case이다.

Scripted teacher와 진짜 사람의 preference로 Training한 MR, NMR, PT에 대한 성능이다. PT말고 baseline 모델들은 locomotion task처럼 비교적 쉬운 task에서는 scripted teacher을 사용한 것이 성능이 더 좋은 것을 볼 수 있다. 

 

Figure 6-b scripted teacher와 human teacher 사이의 agreement rate에 대한 결과이다. 이는 scripted teacher는 real human behavior를 정확하게 모델링하지 않고, sciprted teacher의 evaluation은 오해의 소지가 있는 정보를 생성할 수도 있음을 보여준다. 그리고 6-a는 human, scripted teacher trajectory의 차이점을 보여주는 example이다. human teacher는 context of behavior를 퐉해 올바른 라벨을 제공하는 반면, scripted teacher는 그렇지 못하다. 그리고 또한 이 실험으로 현재 존재하는 preference-based RL benchmark https://arxiv.org/abs/2111.03026 가 충분하지 못하다는 것을 보여준다. preference-based RL을 위한 새로운 benchmark가 디자인 되야한다고 저자들은 시사하고 있다.

 

Learning Complex Novel Behaviors

위 Figure는 로코모션의 hopper task에서 multiple backflip을 하는 피겨이다. 이처럼 복잡한 behavior을 위한 적절한 reward function을 디자인하기 어려울 때 PT를 사용할 수 있다는 것을 보여주는 실험이다. multiple-backflip에서 보상 함수는 rotation 수 같은 non-Markovian context를 포함하고 있어야 하기 때문에 single backflip에 비해 매우 challenging하다. Markovian reward로는 PT에서 학습하는데 성공한 double backflip 같은 것을 할 수 없다

 

 

Discussion

해당 논문에서는 Weighted Non-Markovian reward를 학습할 수 있는 새로운 Preference attention layer를 제안하고, 이는 현재 preference modeling methods에 비해 상당히 outperform하는 것을 해당 논문에서 보여주었다. 그리고 저자들은 Preference Transformer가 Preference-based RL을 scaling하는데 꼭 필수적이라고 이야기한다.