다목적 에이전트 'JAT' 공개: 단일 모델로 게임과 로봇 제어 동시 수행

JAT(Jack of All Trades)는 강화학습 의사결정, 로봇 제어, 기초적인 시각 및 언어 작업을 단일 네트워크로 처리할 수 있는 오픈소스 범용 트랜스포머 에이전트다.

DeepMind의 Gato 연구를 오픈소스로 재현하고 개선한 프로젝트로, 작업마다 별도의 모델을 학습시키지 않고도 다양한 환경을 하나의 모델로 제어할 수 있다.

지원 기능

JAT는 총 157개 작업에 걸쳐 단일 모델로 작동한다.

Atari 레트로 게임 플레이는 물론, MuJoCo 물리 시뮬레이션 환경에서의 제어, Meta-World 환경에서의 로봇 팔 조작, BabyAI의 지시 기반 내비게이션 환경을 지원한다.

이와 함께 텍스트와 이미지 입력도 지원하여 기본적인 자연어 이해 및 시각 처리 작업까지 하나의 인터페이스에서 처리할 수 있다.

작동 원리와 성능

JAT는 EleutherAI의 GPT-Neo 아키텍처를 기반으로 구축되었으며, 순차적 의사결정 데이터 처리에 최적화된 임베딩 방식을 도입했다.

환경 관측값, 보상, 행동 토큰을 교차 배치하는 인과적 시퀀스 구조를 사용한다. 관측 데이터가 이미지일 때는 CNN, 연속형 벡터나 이산형 값일 때는 선형 계층을 거치며, 텍스트는 GPT-2 토크나이저, 이미지는 ViT 인코더로 처리한다.

단순히 다음 행동만 예측하는 대신 미래 관측값을 함께 예측하는 보조 손실 함수를 적용해 학습 효율을 높였다.

전체 157개 평가 작업에서 JAT는 전문가 기준 평균 65.8%의 성능을 달성했다. BabyAI에서는 전문가의 99.0%, MuJoCo에서는 84.8%, Meta-World에서는 65.5%를 기록했다. Atari 57에서는 전문가의 14.1%, 사람 기준 37.6%의 점수를 냈으며 21개 게임에서는 사람의 점수를 넘어섰다.

이용 방법

모델 가중치, 소스 코드, 각 환경별 전문가 정책은 모두 Hugging Face Hub에 공개되어 있다.

수십만 건의 전문가 궤적 데이터가 포함된 JAT 데이터셋 역시 함께 배포되어 datasets 라이브러리를 통해 즉시 로드해 활용할 수 있다.

왜 중요한가

DeepMind가 발표했던 Gato는 범용 에이전트의 가능성을 보여주었으나 가중치와 데이터셋이 공개되지 않아 외부 검증이 어려웠다.

JAT는 모델뿐 아니라 대규모 의사결정 궤적 데이터셋과 전문가 정책까지 온전히 공개해 오픈소스 범용 에이전트 연구의 기틀을 마련했다.

현재는 기본 행동 복제(Behavioral Cloning) 방식으로 학습되어 전문가의 성능을 넘어서지 못한다는 한계가 있지만, 게임과 로봇 제어를 아우르는 오픈소스 기반 모델로서 연구자들에게 훌륭한 출발점을 제공한다.