아지마는 상상한 미래 관측으로 행동을 고르도록 훈련된다

무엇인가

이 논문은 되돌릴 수 없는 행동과 희소한 피드백 속에서 실행 전에 결과를 예상해야 하는 에이전트 문제를 다룬다. 소코반에서는 겉보기에 합리적인 밀기 한 번이 퍼즐을 풀 수 없게 만들 수 있고, 실제 환경에서 대안 행동을 직접 시험하는 것은 비용이 크거나 불가능할 수 있다. 저자들은 행동 전에 미래 관측을 상상하고 그 상상 위에서 추론해 실행을 결정하는 Agentic Imagination(Agima) 훈련 프레임워크를 제안한다.

어떻게 동작하나

Agima의 핵심은 하나의 옴니 모델이 정책과 세계 모델을 겸한다는 점이다. BAGEL 백본은 전체 14B, 활성 7B 파라미터를 가진 이미지·텍스트 생성 모델로, 현재 관측과 목표를 받아 행동 후보를 내고 그 결과 이미지를 생성한다. 행동 묶음 K개에 대해 행동과 예상 관측을 교차 배치해 τ=(a_t, ô_{t+1}, a_{t+1}, ô_{t+2}, …) 같은 궤적을 만들고, 각 예상 상태를 다음 계획의 내부 관측으로 쓴다. 선택된 행동을 실행한 뒤에는 실제 관측이 다음 상상 라운드의 출발점이 되며, 상상 관측과 실제 관측은 같은 인터페이스로 모델 입력에 쓰인다.

무엇과 다른가

SFT 단계는 텍스트와 이미지가 교차하는 궤적으로 정책·세계 모델·결과 평가를 한 모델에 묶는다. 텍스트에는 다음 토큰 교차엔트로피, 이미지에는 MSE 손실을 적용하고 λ_img로 가중해 합친다. 훈련 궤적은 세 가지다. 선형 롤아웃은 한 행동 시퀀스를 따라 각 단계의 상태를 예측한다. 반사실 분기는 주 궤적을 유지하면서 특정 상상 상태에서 다른 행동을 시도하고, 그 즉시 결과가 나쁘면 대안을 기각한 뒤 원래 계획을 계속한다. 전방 사고는 더 긴 후보 계획 여러 개를 롤아웃해 비교한 뒤 실행할 계획을 고른다.

어떻게 쓰나

RL 단계는 SFT가 정해 둔 추론 형식을 넘어 상상을 계획에 더 효과적으로 쓰도록 만든다. GRPO를 사용해 각 문제에서 모델이 실제 환경 관측 없이 완전한 계획 궤적과 최종 실행 행동을 생성하게 하고, 실행 후 얻은 이진 성공 보상으로 언어 토큰만 최적화한다. 생성된 이미지는 문맥에 남지만 직접 RL 손실을 받지 않는다. 그룹 크기 8, 배치 16으로 200스텝 훈련했고, 두 상자 소코반이 SFT 후 거의 포화되어 세 상자 소코반에서 RL을 수행했다. 선형·전방·분기 체크포인트는 약 50% 선형 롤아웃과 50% 해당 추론 궤적을 섞어 초기화해, RL이 언제 어떤 형식을 쓸지 학습하게 했다.

전제와 한계

실험은 Maze2D, Sokoban, Pacman, PushT 네 환경에서 단순한 설정으로 학습하고 학습에 없는 더 어려운 설정으로 평가했다. K는 다음 실제 관측 전에 계획·실행하는 행동 수다. K=5에서 Agima는 이미지 없이 행동 묶음만 예측하는 동일 백본 기반 Agima-A7B(No Imagine)보다 Pacman 36점, Maze2D 17점, Sokoban 38점, PushT 17점 높았다. 두 상자 소코반만 SFT한 모델이 네 상자 문제에서 40% 성공률을 기록했고, 행동만 예측하는 모델은 2%였다. 다만 K=∞의 매우 긴 개루프 계획에서는 여러 환경에서 절대 성능과 행동 전용 모델 대비 우위가 줄었다.

상상 관측의 내용이 실제로 쓰인다는 통제 실험도 있다. 예상 관측을 빈 화면이나 오래된 프레임으로 바꾸면 네 환경 모두 성공률이 급감했고, 예상 프레임을 정답 프레임으로 바꾸면 Pacman·Maze2D·PushT에서 29~44%포인트, Sokoban에서 3%포인트 성공률이 올랐다. 반사실 분기는 소코반 Structural Fork에서 선형 롤아웃 대비 세 상자 성공률을 34%에서 39%로, 네 상자 성공률을 22%에서 32%로 높였다. 전방 사고는 소코반 Deadlock 두 상자에서 22%를 60%로, 세 상자에서 39%를 48%로 올렸고, Pacman 경로 순서 함정에서는 G8/F4 22%를 54%로, G10/F8 5%를 12%로 높였다.

RL은 상상 사용 방식 자체를 바꿨다. SFT 데이터에는 선형 롤아웃과 처음부터 두 후보를 비교하는 전방 사고라는 두 형식만 있었는데, RL 100스텝 후 첫 계획을 롤아웃해 성공이 예상되면 추론을 끝내고 실패가 예상될 때만 다른 후보를 생성하는 새 형식이 나타났다. 두 번째 후보가 생성된 경우 첫 계획이 실패한 비율은 RL 전 32%에서 RL 후 100%가 됐고, 세 상자 설정에서는 처음 두 계획이 실패하면 세 번째 후보까지 탐색했다. 분기에서는 궤적당 평균 분기 수가 RL 전 1.70에서 150스텝 후 1.95로 늘고, 퍼즐을 풀 수 없게 만드는 대안이 있는 중요한 결정 상태에서 분기하는 비율인 Critical Branch Rate가 상승했다. 전체 상상 프레임 수는 늘지 않아, 더 많은 상상 예산을 쓰기보다 분기를 더 효과적으로 배치한 것으로 해석된다.

개발자에게 이 논문은 피드백이 비싸거나 행동이 되돌릴 수 없는 게임·로봇·계획 에이전트에서 실행 전 시뮬레이션을 훈련 루프에 넣는 구체적 패턴을 보여준다. 정책과 세계 모델을 한 모델로 묶고, 생성 이미지를 계획 문맥으로 재사용하며, 결과 성공 보상만으로 언제 더 상상할지 학습시키는 구조가 핵심이다. 다만 원문은 별도의 한계 절을 두지 않는다. 본문에서 확인되는 전제와 한계는 긴 개루프 계획에서 성능이 줄어든다는 점, RL이 소코반 세 상자에서만 수행되고 이진 성공 보상에 의존한다는 점, SFT 궤적이 미리 정한 경험적 형식을 따른다는 점, 생성 이미지가 직접 RL 최적화되지 않는다는 점이다. 실제 도입 전에는 상상 이미지의 정확도와 생성 비용, 보상 설계, 학습 분포 밖 일반화 성능을 별도로 검증해야 한다.