Simulus: Combining Improvements in Sample-Efficient World Model Agents

기술 요약: Simulus – 샘플 효율적 세계 모델 에이전트의 개선 사항 통합

문제 제기
세계 모델 (WM) 은 샘플 효율적 강화 학습 (RL) 의 최전선에 있지만, 유망한 개선 사항들을 통합하는 복잡성으로 인해 그 잠재력이 종종 실현되지 못하고 있습니다. 개별적인 진전은 존재하지만, 이를 결합하려면 상당한 전문성과 엔지니어링 노력이 필요합니다. 또한, 기존 토큰 기반 세계 모델 (TBWM) 은 주로 이미지 관찰과 이산적 행동에 제한되어 있어, 연속 제어 및 다중 모달 설정에서의 적용성이 제한적입니다. furthermore, 내재적 동기부여 (인지적 불확실성 감소를 위한) 와 우선순위 리플레이와 같은 구성 요소는 통합 비용이 높고, 내재적 탐색이 작업과 무관한 경험에 귀중한 상호작용 예산을 낭비할 수 있다는 위험으로 인해 샘플 효율적 에이전트에서 자주 간과됩니다. 이러한 상황은 개별적으로 검증된 개선 사항들이 체계적으로 통합되지 않았던 딥 RL 의 레인보우 (Rainbow) 이전 시대를 반영합니다.

방법론
저자들은 일련의 상호 보완적인 개선 사항을 통합하고 검증하도록 설계된 모듈형 토큰 기반 세계 모델 에이전트인 Simulus를 소개합니다. 이 에이전트는 REM 아키텍처를 기반으로 하며, 표현 모듈 ($V$), 세계 모델 ($M$), 컨트롤러 ($C$), 그리고 리플레이 버퍼로 구성된 네 가지 고유한 모듈로 이루어져 있습니다.

1. 유연한 토큰화 프레임워크 ($V$): Simulus 는 관찰 및 행동 모달리티의 임의의 조합을 지원하는 모듈형 토큰화 시스템을 사용합니다. 이는 이미지, 연속 벡터, 범주형 변수, 그리고 "2D 범주형"(범주형 변수의 이미지와 유사한 그리드) 을 네이티브로 처리합니다. 인코더는 고정 길이의 토큰 시퀀스를 생성하여 통합된 표현으로 연결합니다. 이를 통해 에이전트는 연속 제어 및 상징적 다중 모달 작업을 포함한 다양한 환경에서 작동할 수 있습니다.
2. 세계 모델 ($M$): 동역학 모델은 이산 토큰에 대한 시퀀스 모델링으로 환경 동역학을 모델링하기 위해 병렬 관찰 예측 (POP) 메커니즘이 보강된 Retentive Network(RetNet) 를 활용합니다.
   • 내재적 동기부여: 인지적 불확실성을 줄이기 위해 Simulus 는 다음 관찰 예측 헤드의 앙상블을 사용합니다. 이러한 헤드 간의 불일치는 Jensen-Shannon Divergence(JSD) 를 통해 측정되어 내재적 보상 ($r_{int}$) 을 생성합니다. 이러한 보상은 외부 보수와 결합되어 "상상"(모델 내 계획) 중 컨트롤러가 고불확실성 영역으로 이동하도록 유도합니다.
   • 우선순위 리플레이: 리플레이 버퍼는 저장된 예제에 대한 세계 모델 손실 값을 유지합니다. 학습은 단일 매개변수 $\alpha$로 제어되는 균일 샘플링과 우선순위 샘플링의 혼합을 사용하며, 여기서 샘플은 소프트맥스 변환된 손실에 비례하여 선택됩니다.
3. 컨트롤러 ($C$): 컨트롤러는 세계 모델이 생성한 시뮬레이션 경험 전체에서 작동하는 확장된 액터-크리틱 아키텍처입니다.
   • 회귀로서의 분류 (RaC): 무제한 보상 및 반환 신호를 견고하게 처리하기 위해 Simulus 는 가치 예측을 위해 지수형 버킷을 갖춘 분류 네트워크를 채택합니다. 전통적인 평균 제곱 오차 대신 HL-Gauss 손실 방법 (교차 엔트로피 기반) 을 사용하여 가치 추정 안정성을 향상시킵니다.
4. 학습 주기: 에이전트는 데이터 수집, 표현 학습, 세계 모델 학습, 그리고 상상 속에서의 제어 학습의 주기를 따릅니다. 컨트롤러는 $\lambda$-returns 과 가치 기준선을 갖춘 REINFORCE 목적 함수를 사용하여 학습됩니다.

주요 기여

• 간과된 구성 요소의 통합: 이 논문은 내재적 동기부여, 우선순위 세계 모델 리플레이, 하이브리드 RaC 접근 방식을 식별하고 통합합니다. 이러한 구성 요소들이 개별적으로 효과적일 뿐만 아니라 결합되었을 때 시너지 효과를 낸다는 것을 입증합니다.
• 다용도 토큰 기반 에이전트: 다중 모달 토큰화를 샘플 효율적 체제에 적응시킴으로써, Simulus 는 시각, 연속, 그리고 다중 모달 입력을 처리할 수 있는 다용도 에이전트로서 TBWM 을 확립하여 이미지 전용 이산 행동 모델의 한계를 넘어섭니다.
• 모델 유도 탐색의 검증: 이 연구는 세계 모델 유도 탐색이 엄격한 상호작용 예산 하에서도 성능을 크게 향상시킨다는 것을 보여주며, 내재적 동기부여가 샘플 효율적 설정에는 비용이 너무 많이 든다는 관념에 도전합니다.
• 오픈소스 기반: 저자들은 코드와 모델 가중치를 오픈소스로 제공함으로써 TBWM 연구에 대한 강력한 출발점을 마련하고, 경쟁력 있는 세계 모델 에이전트 개발의 진입 장벽을 낮춥니다.

결과
Simulus 는 엄격한 상호작용 예산을 가진 세 가지 다양한 벤치마크에서 평가되었습니다.

• 시각 Atari 100K: Simulus 는 계획 없는 방법들 사이에서 최첨단 (SOTA) 샘플 효율성을 달성하여 DreamerV3, TWM, STORM, DIAMOND 와 같은 베이스라인을 능가했습니다. 특히, 이는 인간 수준의 인터쿼타일 평균 (IQM) 과 중앙값 점수에 도달한 최초의 계획 없는 세계 모델이며, 26 개 게임 중 13 개에서 초인적인 성능을 달성했습니다.
• 연속 제어 (DMC Proprioception 500K): 연속 제어 작업에서 Simulus 는 DreamerV3 와 일관되게 맞먹거나 약간 더 나은 성능을 보여주어, 컴팩트한 어휘를 가진 연속 도메인에서 토큰 기반 아키텍처의 효과성을 입증했습니다.
• 상징적 다중 모달 (Craftax-1M): 상징적 그리드와 연속 특징을 결합한 부분 관측 가능한 2D 생존 환경에서 Simulus 는 동시 세계 모델 방법과 모든 모델 프리 베이스라인 (RND, PPO-RNN, E3B 포함) 을 능가하여 긴 토큰 시퀀스를 효율적으로 처리했습니다.

Ablation 연구는 각 구성 요소가 개별적으로 기여하며, 내재적 보상이 특히 중요하다는 것을 확인했습니다. 모든 구성 요소의 결합은 가장 강력한 성능을 발휘하여 상당한 시너지를 나타냈습니다.

의의
이 논문은 세계 모델 에이전트의 실현되지 않은 잠재력이 모듈형 개선 사항을 결합하는 어려움에서 비롯된다고 주장합니다. DQN 의 Rainbow 프레임워크에서 영감을 얻어, Simulus 는 토큰화 유연성, 내재적 동기부여, 우선순위 리플레이, 그리고 고급 가치 예측 기술의 체계적인 결합이 최첨단 샘플 효율성으로 이어진다는 것을 보여줍니다. 이 연구는 구성 요소가 모놀리식 아키텍처에서 흔히 보이는 간섭을 피하도록 신중하게 통합된다면, 토큰 기반 세계 모델이 다양한 RL 도메인을 위한 다용도이고 확장 가능한 접근 방식임을 확립합니다. 저자들은 토큰 기반 에이전트가 현재 일부 베이스라인에 비해 학습 속도 제한에 직면해 있지만, 모듈형 설계는 정책 추론과 향후 연구 확장성에 있어 뚜렷한 장점을 제공한다고 주장합니다.