필요한 사전지식

• MLP
• 트랜스포머
• 토크나이저
• GPT
• MSE
• contrastive learning
• Decision Transformer
• Behavioral Cloning

풀어내는 문제

POMDP를 풀어내는 데 도움 될 표현을 학습하는 문제.

모델 아키텍처

• 이름: Control Transformer
• minGPT 사용
• Input은 관측값과 행동의 수열.
• 관측과 행동을 토큰화한 방법
  • 관측과 행동 각각 동일한 차원의 벡터 공간에 임베딩 함.
  • 이미지 임베딩은 3개의 층을 갖는 CNN으로 학습함.
  • 행동을 토큰화하는 토크나이저
    • 이산적인 경우 룩업 테이블로 토큰화 (코드 근거: nn.Embedding(config.vocab_size, config.n_embd), nn.Tanh())
    • 연속적인 경우 한 개의 층으로 구성된 feedforward 네트워크로 토큰화
  • 각 토크나이저를 control transformer와 같이 학습함.

사전학습 목표

표현을 학습하기 위해 아래의 목표 함수를 최적화함.

• 첫 번째 목표: Forward Dynamics Prediction
  • $o_t$와 $a_t$로 $s_{t+1}$을 예측
  • 물론 $s_{t+1}$은 에이전트가 접근할 수 없음. 그렇기에 실제로는 $\hat{s}_{t+1} = SG(\bar{g}(o_{t+1}))$를 예측함. 여기에서 $\bar{g}$는 관측치를 토큰화하는 토크나이저고 $SG$는 Stop Gradient다.
    • $\bar{g}_{\bar\theta}$를 파라미터 $\bar\theta$를 갖는 토크나이저라 하자. 그러면 $\barθ = τ (\barθ)+(1−τ)θ$
      • 위처럼 moving average를 사용한 이유: 관측치 토큰화 네트워크 또한 학습하기 때문에, 예측 대상인 관측 임베딩 값이 계속 바뀜. 예측 대상이 바뀔 때 moving average가 안정적으로 학습하도록 도움.
  • $L_{fwd} := MSE(f_{fwd}(\phi(a_t), \phi(o_t)), \hat{s}_{t+1})$ $f_{fwd}$는 선형 네트워크임.
  • Stop Gradient를 사용한 이유를 설명하지 않음
    • 추측: 위의 loss를 줄이기 위해서, 예측인 $f_{fwd}(\phi(a_t), \phi(o_t))$를 수정하거나 타겟인 $\hat{s}_{t+1}$을 수정해야 한다. 상식적으로 예측을 수정해서 타겟과 가깝게 만들어야 한다. 만약 $SG(\bar{g}(o_{t+1}))$가 아니라 $\bar{g}(o_{t+1})$을 사용하면, 타겟을 수정해서 예측과 가깝게 만든다. 이를 방지하기 위함으로 보임.
• 두 번째 목표: Inverse Dynamics Prediction
  • $o_t$와 $o_{t+1}$으로 $a_t$를 예측
  • $L_{inv} := MSE(f_{inv}(\phi(o_t), \phi(o_{t+1})), a_t)$ $f_{inv}$는 선형 네트워크임.
  • GPT의 causal mask를 사용하면 $o_{t+1}$이 $a_t$에 attend할 수 있음. 그래서 causal mask를 수정함.
• 세 번째 목표: Random Masked Hindsight Control
  • $(o_t, a_t, ... , o_{t+L}, a_{t_L})$ 수열에서 확률적으로 $k$ 개의 관측치와 $k'$ 개의 행동 토큰을 가림.
  • 가려지지 않은 토큰의 수열로 $k\prime$ 개의 행동 토큰을 예측
  • 가려진 토큰은 $\bold{m}=\{-1, -1, ..., -1\}$로 대체된다.
  • 양방향 예측을 위해 causal mask를 삭제
  • 세 번째 목표가 필요한 이유
    • 위의 두 목표와 다르게 "글로벌" 예측임. 즉 예측 대상 토큰과 예측에 쓰일 토큰이 멀리 떨어질 수 있음.
      • 그렇기에 관측치와 행동 간에 장기적인 관계를 학습할 수 있음.
      • 이에 반해 위의 두 목표는 근처의 토큰으로 근처의 토큰을 예측하는 "로컬 예측"이다. 예시로 두 번째 목표의 경우, $a_t$를 예측하기 위해 $a_t$의 근처 관측치인 $o_t$와 $o_{t+1}$만 사용
  • 처음에 $k$ 개와 $k'$ 개의 관측과 행동을 가리고 행동을 예측하는 것은 어렵다.
    • 그렇기에 마치 curriculum learning을 하듯이, $k$와 $k'$을 초기에는 작게 설정하고 학습을 진행하면서 점점 증가하도록 구현
    • 학습 후반부에는 $k = L$ $k' = \frac{L}{2}$
      • $k'$이 작은 이유: $k'$이 $L$에 가까워지면 $o_t$에서 $o_{t+L}$에 도달하기 위한 행동의 수열을 생성하는 문제와 비슷해진다.
        • 문제는 특정 trajectory를 만드는 행동의 수열은 유일하지 않고 확률적인 환경으로 인해 특정 행동 수열이 동일한 trajectory로 이어지지 않는다.
        • 이러한 어려움을 피하기 위해 $k'$을 작게 설정함.

특이사항

• 연속적인 행동을 이산화하지 않음
  • 이에 더불어 MSE를 사용하기에 multimodal 행동을 제대로 모델링할 수 없음
    • 특히 사람의 행동은 multimodal인 경우가 많음

실험

• 단계: 사전 학습 단계와 미세 조정 단계로 나뉨
  • 위에서 설명한 목표 함수를 최적화하도록 사전 학습한 후, 원하는 작업에 미세 조정함.
  • 백본의 파라미터를 고정시킨 상태에서 미세 조정했을 때, 복잡한 작업을 잘 수행하지 못함.
    • 그렇기에 미세 조정 때 백본까지 같이 학습시킴.
• 환경: DeepMind Control suite를 구성하는 10개의 작업
  • 5 개 작업에 대해서는 pretraining과 fine-tuning 단계에서 학습함.
  • 나머지 5 개 작업에 대해서는 fine-tuning 단계에서만 학습함.
    • 이유: 접하지 못한 작업 수행에도 일반화 가능한지 알아보기 위해.
• 비교 모델. DT를 제외한 모든 모델은 동일한 아키텍처를 가짐.
  • Scratch: 사전 학습없이 CT로 policy 학습한 모델
  • ACL: 토큰을 random하게 가린 후 예측하면서 동시에 contrastive loss를 줄이도록 학습한 모델. SMART와 동일한 데이터셋에 사전 학습함.
  • DT: Decision Transformer. SMART와 동일한 데이터셋에 사전 학습
  • CT-single: 다섯 개의 작업을 사용하지 않고 한 개의 작업만 사용하여 사전 학습한 모델
• 특정 task에서 비교 모델이 우위를 보이는 경우도 있었으나, 전반적으로 SMART 모델이 우월한 성능을 보임.

한계

• DeepMind Control Suite에서만 실험해 봄.
  • DeepMind Control Suite를 구성하는 작업은 실제 효용이 없고 크게 복잡하지 않음.
  • 작업을 수행하는 데 필요한 표현이 복잡할지 의문임
  • 그렇기에 복잡하고 현실적인 작업을 수행할 때, 이 방식이 효과적일지 추측하기 어려움.
• 관측치만 있어도 학습할 수 있는 VC-1과는 다르게, 에이전트가 환경에서 작업을 수행하며 얻어낸 trajectory가 필요함.
  • 그렇기에, 실제 환경에서 동작하는 에이전트를 학습시키고자 한다면 안전하게 trajectory를 얻어낼 방법이 필요함.