Hugging Face 팀이 Gato 재현 프로젝트에서 시작해 Transformer 기반 범용 에이전트 JAT를 개발해 4개 도메인에서 평균 65.8% 전문가 성능 달성

Context

기존 강화학습은 단일 환경에서만 정책을 학습했으므로, 게임, 로봇 제어, 자연어 이해 등 다양한 작업을 수행할 수 있는 범용 에이전트 개발이 필요했다. Gato(Reed et al., 2022) 논문은 비전-언어-결정 작업을 모두 수행하는 Transformer 학습을 제안했으나, 오픈소스 재현과 개선이 없었다.

Technical Solution

• Atari, BabyAI, Meta-World, MuJoCo 등 다양한 환경에서 전문가 정책 학습: 각 환경별로 최첨단 성능에 도달한 에이전트 157개 개발
• JAT 데이터셋 구축: 전문가 에이전트에서 수집한 수십만 개의 전문가 궤적 포함
• Transformer 기반 아키텍처 설계: 관찰 임베딩과 행동 임베딩, 보상을 interleave하는 인코딩 메커니즘 도입
• 멀티모달 데이터 처리: 이미지는 CNN, 연속 벡터는 선형 레이어, 이산 값은 선형 투영층 사용
• 손실 함수 모달리티별 분리: 이미지와 연속 값은 MSE 손실, 이산 값은 교차 엔트로피 손실 적용
• 인과적 예측: 행동을 다음 시점으로 1단계 시프트해 이전 관찰과 행동으로부터 다음 행동 예측

Impact

• Atari 57개 게임에서 전문가 성능의 14.1% 달성(인간 성능의 37.6%에 해당)
• 21개 Atari 게임에서 인간 성능 초과
• 4개 도메인(Atari, BabyAI, Meta-World, MuJoCo) 전체에서 평균 65.8% 전문가 정규화 점수 달성

Key Takeaway

단일 Transformer 아키텍처로 다중 도메인의 이질적인 데이터(이미지, 연속값, 이산값)를 통합 처리하려면, 모달리티별 인코더 분화와 손실 함수 분리를 통해 각 데이터 타입에 최적화된 학습이 필수다. 범용 에이전트 개발은 충분한 전문가 궤적 수집과 오프라인 강화학습, 다중 작업 샘플링 전략 개선으로 성능 향상 가능성이 남아있다.