마스크드 확산 언어 모델은 에이전트 강화학습을 위한 강력하고 조정 가능한 텍스트 기반 세계 모델이다

초록

최근 강화학습(RL)의 성장은 다양하고 특화된 훈련 환경에 대한 필요성을 부각시켰다. 고정된 작업과 보상 난이도를 가진 수작업 환경은 모델 성능이 향상됨에 따라 효과적인 신호를 제공하지 못하게 되며, 긴 시간 지평에 걸친 희소 보상은 특정 작업 흐름이나 도구 구조에 대한 모드 붕괴를 유발한다. 환경 상태를 시뮬레이션하는 세계 모델은 순수 롤아웃 성능과 일치하여, 필요에 따라 다양성을 확장하는 데 유망하다. 그러나 자기회귀(AR) 세계 모델은 좌에서 우로의 편향으로 인해 도구 스키마, 이전 회차, 예상 결과와 같은 전역적으로 상호의존적인 상태 앵커에 조건화할 수 없다. 본 연구에서는 (i) 텍스트 기반 세계 모델링을 초기 상태, 작업 맥락, 도구 스키마, 도메인 규칙, 방향 지시로 분해된 조정 가능한 전이 동역학 문제로 정형화하고, (ii) 9개의 오픈소스 환경과 12개의 프론티어 모델 패밀리를 포괄하는 239,403개의 근거 기반 상태-행동 궤적을 큐레이션했다. AR 언어 모델과 마스크 확산 언어 모델(MDLM)을 비교한 결과, MDLM은 양방향 앵커 인식 잡음 제거를 통해 LLM 대비 4배 큰 파라미터 크기와 유사한 추론 지연 시간에서 더 나은 일관성, 근거 기반성, 실증적으로 검증된 롤아웃 다양성을 달성했다. 또한 결정적 상태 검사를 포함한 플러그 앤 플레이 GRPO 훈련 프레임워크를 도입하고, 3개의 OOD 환경(ScienceWorld, ALFWorld, AppWorld)에서 3개의 1.2B-7B 에이전트 백본(LFM2.5, Qwen3, Mistral)에 대해 제로샷 전이 절제 실험을 수행하여 환경별 미세 조정 없이 기준선 대비 최대 47%의 절대적 성능 향상을 달성했다. 더 나아가, 적대적 시나리오에서의 실패 모드에 대한 행동 분석과 현실성, 결과 정확성, 훈련 유용성에 대한 인간 평가를 수행했다. 본 연구는 이 방향의 연구를 장려하기 위해 오픈소스로 공개한다.

English

Recent growth in reinforcement learning (RL) has surfaced a need for diverse, specialized training environments. Hand-curated environments with fixed task and reward difficulties become ineffective signals as model performance improves, and sparse rewards over long horizons induce mode collapse on specific workflows or tool structures. World models that simulate environment states have matched pure rollout performance, making them promising for scaling diversity on-demand. However, autoregressive (AR) world models suffer from a left-to-right bias preventing conditioning on globally interdependent state anchors such as tool schemas, prior turns, and expected outcomes. We (i) formalize text-based world modeling as a steerable transition-dynamics problem decomposed into initial state, task context, tool schemas, domain rules, and steering directives, and (ii) curate 239,403 grounded state-action trajectories spanning nine open-source environments and twelve frontier model families. We compare AR LMs and masked diffusion language models (MDLMs), showing MDLMs, via bidirectional anchor-aware denoising, achieve better coherence, groundedness, and empirically validated rollout diversity than LLMs over 4x their parameter size, at comparable inference latency. We introduce a plug-and-play GRPO training framework with deterministic state checks, and perform zero-shot transfer ablations on three OOD environments (ScienceWorld, ALFWorld, AppWorld) across three 1.2B-7B agent backbones (LFM2.5, Qwen3, Mistral), achieving up to 47% absolute gains over baselines without environment-specific fine-tuning. We further conduct behavioral analysis of failure modes under adversarial scenarios and human evaluation on realism, outcome correctness, and training utility. We open-source our work to encourage research in this direction.