Autoregressive Diffusion World Models for Off-Policy Evaluation of LLM Agents

K. Liu, G. Xiong, W. Zhang, and S. Tang, "Autoregressive Diffusion World Models for Off-Policy Evaluation of LLM Agents," arXiv:2606.05558, 2026.

새 에이전트를 배포 전에 평가하려면 보통 실제 환경에 굴려봐야 합니다. 그런데 멀티턴 환경에서 에이전트를 라이브로 돌리는 건 비쌉니다. API 예산을 태우고, 되돌릴 수 없는 부작용을 낼 수도 있습니다. 그래서 등장한 것이 오프폴리시 평가(off-policy evaluation, OPE)입니다. 과거에 모아둔 오프라인 데이터만으로 새 정책의 가치를 추정하고, 환경과 더는 상호작용하지 않는 방식입니다.

문제는 LLM 에이전트에서 기존 OPE가 줄줄이 무너진다는 점입니다. 이 논문은 월드 모델 자체를 디퓨전 과정으로 세워 그 난점을 정면 돌파합니다. 네 개의 멀티턴 에이전트 벤치마크에서, 다섯 개 고전 OPE 추정기가 최소 세 개 셀에서 실패하는 동안 ADWM만 모든 셀에서 양의 순위 상관을 냈습니다.

저자

네 명이 에모리대와 상하이교통대에 걸쳐 있습니다. 제1저자 Kaixuan Liu와 시니어 저자 Shengpu Tang은 에모리대 컴퓨터과학과 소속입니다. Shengpu Tang은 오프라인 강화학습과 OPE를 의료 의사결정에 적용해온 연구자로, "검증 데이터로 정책을 고르려면 OPE가 믿을 만해야 한다"는 문제의식을 오래 다뤄 왔습니다. 이 논문의 출발점이 정확히 거기입니다.

교신저자 Guojun Xiong은 상하이교통대 존 홉크로프트 센터의 부교수로 강화학습과 온라인 순차 의사결정이 전공입니다. 정책 유도 궤적 법칙을 단일 스텝 조건부로 정확히 분해하는 이 논문의 핵심 정리가 그의 이론적 배경과 맞닿아 있습니다. 함께한 Weinan Zhang은 SJTU의 강화학습·에이전트 AI 대가로, 피인용 3만 회를 넘는 연구 축적이 월드 모델 설계의 무게를 더합니다.

OPE를 깊이 아는 사람들과 월드 모델·강화학습 이론에 강한 사람들이 만난 조합입니다. "왜 기존 OPE가 LLM 에이전트에서 깨지는가"를 정확히 짚을 수 있었던 배경입니다.

배경

OPE의 목표는 평가 정책 \(\pi_e\)의 가치 \(J(\pi_e) = \mathbb{E}_{\tau \sim p_{\pi_e}}[R(\tau)]\)를, 행동 정책 \(\pi_b\)가 모은 오프라인 데이터만으로 추정하는 것입니다. 핵심 난관은 분포 이동(distribution shift)입니다. \(\pi_b\)가 만든 궤적은 \(\pi_e\)가 만들 궤적과 다르므로, 데이터에 그대로 몬테카를로를 돌리면 심하게 편향됩니다.

고전적 처방마다 약점이 분명합니다. 중요도 샘플링(importance sampling)은 오프라인 궤적을 정책 가능도 비율 \(\prod_t \frac{\pi_e(a_t \mid h_t)}{\pi_b(a_t \mid h_t)}\)로 재가중하는데, 이 가중치가 궤적 길이에 따라 지수적으로 폭발해 멀티턴에서 못 씁니다. 값 기반 직접 추정은 평가 정책이 행동 정책과 크게 다르면 편향이 누적됩니다.

더 유망한 길은 월드 모델입니다. 오프라인 데이터에서 환경 동역학을 배우고 평가 정책으로 롤아웃을 시뮬레이션하는 방식입니다. 그런데 에이전트 설정에서는 이 시뮬레이션이 자기회귀적이어야 합니다. 각 LLM 행동이 방금 받은 관측에 의존하므로, 월드 모델과 에이전트가 인과 순서로 한 스텝씩 번갈아야 합니다. 자기회귀 트랜스포머 월드 모델은 이 구조를 자연스럽게 따르지만 관측을 토큰 단위로 생성하느라 오차가 스텝 안에서 쌓이고 스텝을 건너 더 불어납니다. 긴 시야 과제에서 치명적인 실패 모드입니다.

디퓨전 모델은 원리적 해법을 줍니다. 각 전이를 독립적인 디노이징 과정으로 모델링하면 오차가 스텝을 건너 전파되지 않습니다. 그런데 기존 디퓨전 OPE는 연속 제어용으로 설계됐습니다. 상태와 행동이 둘 다 실숫값 벡터라 하나의 텐서로 함께 노이즈를 넣고 빼는 걸 전제합니다. LLM 에이전트에서는 이 전제가 무너집니다. 행동이 이산 텍스트이고, 환경을 관측한 뒤에야 정책에서 샘플링됩니다. 전체 궤적을 한 번에 생성하려면 관측을 만들기 전에 모든 행동을 미리 알아야 한다는 순환 의존이 생기고, 이건 에이전트 평가가 요구하는 한 스텝씩 롤아웃과 양립할 수 없습니다.

어떻게 풀었나

위 그림(Figure 1)이 세 패러다임을 나란히 놓습니다. 온폴리시 평가는 에이전트를 실제 환경에 굴려야 해서 비싸고 위험합니다. 전통적 오프폴리시 평가는 오프라인 데이터로 모델 기반 시뮬레이터를 배우지만 분포 이동과 누적 오차라는 두 근본 문제를 안습니다. ADWM은 둘을 동시에 칩니다. 정책 유도(policy guidance)가 평가 정책 \(\pi_e\)를 매 디노이징 스텝에 주입해 분포 이동을 누르고, 디퓨전 사전분포에 행동 인식 디노이징을 결합해 누적 오차를 막습니다.

핵심은 정리 1(Theorem 1)입니다. 정책 유도 궤적 법칙이 단일 스텝 조건부의 곱으로 정확히 분해된다는 것입니다.

\[\tilde{q}_{\omega,\eta}(o_{t+1} \mid h_t, a_t; \pi_e) \propto \underbrace{P_\theta(o_{t+1} \mid h_t)}_{\text{사전분포}} \cdot \underbrace{P_\theta(a_t \mid o_{t+1}, h_t)^{\omega}}_{\text{행동 사후분포}} \cdot \underbrace{C_\alpha(o_{t+1}; h_t, a_t, \pi_e)^{\eta}}_{\text{정책 연속}}\]

세 인자가 각자 다른 일을 합니다. 사전분포 \(P_\theta(o_{t+1} \mid h_t)\)는 생성된 관측을 오프라인 데이터의 지지집합에 묶어 둡니다. 행동 사후분포는 별도 분류기 없이 현재 스텝 행동 정보를 주입하는데, classifier-free guidance에서 바로 복원됩니다. 연속 인자 \(C_\alpha\)는 후보 관측을 \(\pi_e\)와의 장기 양립성으로 재가중해, \(\pi_e\)가 효과적으로 행동할 수 있는 미래를 여는 관측을 고르게 합니다. 이 곱의 로그는 합으로 풀리므로 score가 더해지는 구조가 되고, 각 항을 월드 모델 재학습 없이 역과정에 독립적으로 주입할 수 있습니다. 핵심은 이게 단일 스텝 조건부라는 점입니다. \(\pi_e\)가 관측이 생성되기 전에 행동을 내놓을 필요가 없어, 기존 디퓨전 OPE를 막던 순환 의존이 완전히 풀립니다.

구조(Figure 2)는 이렇게 돕니다. 인코더 \(E\)가 텍스트 관측을 잠재 상태 \(z_t\)로 보내고, 디퓨전 월드 모델 \(p_\theta\)가 \(K\)스텝 디노이징으로 다음 잠재를 만듭니다. 프로젝터 \(G_\psi\)가 잠재를 소프트 토큰으로 바꿔 평가 정책 \(\pi_e\)가 자기 임베딩 공간에서 읽게 합니다. \(\pi_e\)는 두 역할을 합니다. 디노이징을 조종하고(정책 유도), 그 관측을 보고 행동 \(a_t\)를 샘플링해 자기회귀 롤아웃을 끕니다. 보상 헤드가 보상을 예측하고, 이를 몬테카를로로 평균해 가치 추정 \(\hat{J}(\pi_e)\)를 냅니다. 환경 상호작용은 한 번도 필요 없습니다.

의미 인코더를 그냥 가져다 쓰지 않은 점도 중요합니다. 부정어 하나만 다른 두 관측이 의미 공간에서는 가깝지만 동역학은 완전히 다를 수 있어, 세 유도 항을 동시에 망가뜨립니다. 그래서 인코더를 표면 의미가 아니라 환경 동역학이 빚도록 데이터 위에서 끝까지 함께 학습시킵니다.

결과

평가는 보상과 후학습 영역을 넓게 덮는 네 벤치마크에서 합니다. HotpotQA(조밀한 스텝별 F1 보상), ScienceWorld(형태화된 부분 보상), WebShop(연속 부분 보상), ALFWorld(희소한 0/1 성공)입니다. 정책 품질을 단계적으로 흔들기 위해 \(\varepsilon\)-그리디 혼합으로 \(\varepsilon \in \{0, 0.25, 0.5, 0.75, 1.0\}\)의 평가 정책 군을 만들고, 실제 환경에서 얻은 정답 곡선과 ADWM이 월드 모델 롤아웃만으로 낸 \(\hat{J}\) 곡선 사이의 스피어만 순위 상관 \(\rho\)를 봅니다. 모든 셀에서 월드 모델을 학습시키는 행동 정책 \(\pi_b\)는 평가 정책 \(\pi_e\)와 엄격히 다릅니다. 월드 모델은 \(\pi_e\)를 한 번도 본 적 없이 학습됩니다.

ADWM은 여섯 개 구성 전부에서 \(\rho > 0\)인 유일한 추정기입니다. 평균 +0.82, 최소 +0.67, 최대 +0.90입니다. 고전 베이스라인은 하나같이 최소 세 셀에서 실패합니다. 중요도 샘플링(IS)과 이중 강건(DR)은 LLM이 유발하는 중요도 비율이 스무 자릿수를 넘나들며 폭발해, 상관이 −0.90까지 떨어집니다. 가중 중요도 샘플링(WIS)은 폭발은 막지만 긴 시야에서 지배적 가중치 하나로 퇴화해 ALFWorld 두 셀에서 \(\rho=0\)을 내고, 다른 곳의 양의 상관은 부동소수점 언더플로 동률이 만든 허상입니다. 직접법(DM)은 평가 정책의 행동 분포와 무관한 행동 정책 상태값을 추정해 모든 \(\varepsilon\) 스윕에서 \(\rho=0\)을 내고 교차 정책 분할에서는 부호가 뒤집힙니다(\(\rho=-0.45\)). FQE만 부분적 예외인데, COBS 프로토콜의 \(\varepsilon\)-선형성 보정을 거치면 ALFWorld-iter1에서 +0.82를 내지만 나머지에서는 무너지거나 뒤집혀 평균 +0.10에 그칩니다. 표에서 IS, WIS 같은 베이스라인에는 사실 행동 LLM의 정확한 토큰별 로그확률까지 특권적으로 줬는데도 그렇습니다. 현실 배포(폐쇄형 API, 삭제된 체크포인트, 숨은 토크나이저)에서는 그것조차 못 얻습니다. ADWM은 중요도 가중치도, 선형 귀납 편향도 없이 다음으로 좋은 베이스라인을 평균 상관에서 +0.34 앞섭니다.

무엇이 중요한가

세 유도 요소가 각각 다른 보상 영역을 책임집니다. 어느 하나도 혼자서는 전 영역을 못 덮습니다.

로컬 CFG(행동 사후분포 유도)는 희소 보상에서 가장 결정적입니다. 빼면 WebShop이 +0.90에서 +0.10으로 무너지고 ScienceWorld가 +1.00에서 +0.60으로 떨어집니다. 연속 유도(continuation guidance)는 스텝별 보상이 없을 때 궤적을 성공 경로 위에 붙드는 주 메커니즘으로, 제거하면 WebShop이 +0.30, ScienceWorld가 +0.90으로 내려갑니다. \(\psi\) 어댑터(소프트 토큰 프로젝터)는 언어적으로 풍부한 환경에서 지배적입니다. 빼면 HotpotQA가 +0.40으로 떨어집니다. ADWM의 견고함은 세 요소의 결합에서 나온다는 결론입니다. 학습 손실도 네 벤치마크 모두에서 50 에폭 안에 안정적으로 수렴합니다.

회고

저자들이 그은 경계는 분명합니다. IS, WIS, DR 같은 고전 베이스라인은 원래 저차원 연속 제어용으로 설계됐고, 무한한 LLM 행동 공간과 긴 시야 롤아웃의 제한된 분포 겹침은 이들이 기대는 절대 연속성 가정을 위반합니다. 그래서 저자들은 이 베이스라인을 경쟁 상대가 아니라 진단용 참조로 포함한다고 못 박았습니다. 셀당 시드는 다섯 개이고, 닫힌 형태의 신뢰구간은 2000회 부트스트랩으로 냈습니다. 평가 정책은 모두 선행 연구가 공개한 체크포인트라, 재현 가능한 범위 안에서 검증했습니다.

이 논문이 보여주는 건 "디퓨전을 OPE에 쓴다"는 막연한 아이디어가 아니라, 이산 텍스트 행동이라는 LLM 특유의 제약 때문에 기존 디퓨전 OPE가 왜 못 도는지, 그리고 그 순환 의존을 단일 스텝 분해로 어떻게 끊는지의 구체적 메커니즘입니다.

정리

• ADWM은 새 LLM 에이전트를 실제 환경에 굴리지 않고 과거 로그만으로 가치를 추정하는 오프폴리시 평가 프레임워크입니다. 월드 모델 자체를 디퓨전 과정으로 세운 것이 핵심입니다.
• 정리 1이 정책 유도 궤적 법칙을 사전분포 곱하기 행동 사후분포 곱하기 정책 연속의 단일 스텝 조건부로 정확히 분해합니다. 평가 정책이 매 디노이징 스텝을 조종하므로 관측보다 행동이 먼저 필요한 순환 의존이 풀리고, 각 전이가 독립 디노이징이라 오차가 스텝을 건너 누적되지 않습니다.
• 네 개 멀티턴 벤치마크 여섯 구성에서 ADWM만 모든 셀에서 양의 순위 상관(평균 +0.82, 최소 +0.67)을 내며, 다음으로 좋은 베이스라인을 평균 +0.34 앞섭니다. 에이전트 A/B 테스트 비용을 줄이는 실전 도구로 읽을 수 있습니다.