Jack of All Trades (JAT) 다목적 트랜스포머 에이전트

Hugging Face는 단일 네트워크를 사용해 다양한 비전‑언어 및 의사결정 작업을 수행하도록 설계된 일반형 에이전트 프로젝트인 Jack of All Trades (JAT)를 발표했습니다. 이 프로젝트는 Gato 아키텍처를 오픈 재현한 것으로, 포괄적인 전문가 RL 에이전트 세트, 특화된 학습 데이터셋, 그리고 멀티모달 트랜스포머 모델을 제공합니다.

JAT 데이터셋 및 전문가 정책

JAT 프로젝트는 일반형 에이전트 학습을 위해 특별히 설계된 최초의 데이터셋을 제공하며, 수십만 개의 전문가 궤적을 포함합니다. 이러한 궤적은 여러 다양한 환경에서 최첨단 전문가 에이전트를 훈련시켜 생성되었습니다:

  • Atari: 고전 아케이드 게임.
  • BabyAI: 간단한 네비게이션 환경 (BabyAI 봇 활용).
  • Meta-World: 로봇 조작 작업.
  • MuJoCo: 물리 기반 연속 제어.

통합 사용자 인터페이스를 제공하기 위해 JAT 데이터셋에는 Wikipedia, Oscar, OK‑VQA, Conceptual‑Captions의 텍스트 데이터도 포함됩니다.

JAT 에이전트 아키텍처

JAT는 EleutherAI의 GPT‑Neo 구현을 기반으로 구축되었습니다. 이 아키텍처는 관찰 임베딩, 행동 임베딩, 그리고 해당 보상을 교차 삽입(interleaving)하여 순차적 의사결정 작업을 처리하도록 설계되었습니다.

임베딩 및 인코딩 메커니즘

JAT는 데이터 모달리티에 따라 서로 다른 인코딩 전략을 사용합니다:

  • 이미지 (예: Atari): Convolutional Neural Network (CNN)으로 처리.
  • 연속 벡터: 선형 레이어로 처리.
  • 이산 값: 선형 프로젝션 레이어로 처리.
  • 텍스트: GPT‑2 방식으로 토크나이즈.
  • 일반 이미지: ViT‑type 인코더 사용.

예측 및 손실 함수

모델은 인과 마스크(causal mask)를 사용해 다음 임베딩을 자동회귀적으로 예측하며, 관찰을 한 타임스텝씩 이동시켜 에이전트가 이전 관찰 및 행동을 기반으로 다음 행동을 예측하도록 합니다. 손실은 각 모달리티별로 별도로 계산됩니다:

  • MSE 손실: 이미지와 연속 값에 사용.
  • Cross‑Entropy 손실: 이산 값에 사용.

최종 손실은 시퀀스 내 각 요소에 대한 손실들의 평균입니다.

성능 결과

JAT는 157개의 학습 작업에 대해 평가되었으며, 네 가지 주요 도메인에서 전문가 에이전트 대비 평균 65.8%의 성능을 달성했습니다.

도메인별 성능

  • BabyAI: 전문가 점수의 99.0%, 단 하나의 작업에서만 50% 이하.
  • MuJoCo: 전문가 점수의 84.8%.
  • Meta‑World: 전문가 점수의 65.5%.
  • Atari 57: 전문가 점수의 14.1% (인간 성능의 37.6%에 해당), 21개의 게임에서 인간 성능 초과.

모델은 NLP와 CV 작업에서도 기본적인 능력을 보이지만, 다양한 RL 도메인에서 단일 네트워크로 전문가 수준의 성능을 모방할 수 있는 것이 주요 강점입니다.

관찰 예측의 영향

JAT 프로젝트 동안 수행된 연구에서는 에이전트에게 보상을 최대화하는 것 외에 미래 관찰을 예측하도록 요구하면 학습이 향상되는지를 조사했습니다. 관찰 손실과 행동 손실을 균형 맞추는 가중치 파라미터 $\kappa$를 사용했으며, $\kappa = 0.005$에서 "최적점"을 발견했습니다.

이 특정 가중치에서는 관찰을 예측하는 학습이 에이전트의 학습 효율을 높였습니다. 그러나 $\kappa$가 너무 높을 경우(예: 0.5) 관찰 예측 목표가 학습을 방해했습니다. 이는 보조 목표가 적절히 균형 잡히면 일반형 에이전트에 도움이 될 수 있음을 시사합니다.

향후 연구 방향

JAT 에이전트를 더욱 발전시키기 위해 Hugging Face는 다음 세 가지 주요 개선 영역을 제시합니다:

  1. 데이터 품질: 더 다양한 전문가 에이전트로부터 궤적을 추가 수집해 데이터셋을 확장하고 편향을 감소.
  2. 오프라인 RL: 기본 행동 복제(Behavioral Cloning)를 넘어, 최적이 아닌 궤적을 활용해 전문가를 능가할 가능성을 탐색.
  3. 샘플링 전략: 균등 샘플링 대신 더 어려운 작업에 집중할 수 있는 동적 멀티‑태스크 샘플링 전략 구현.

Sources