World-Model Policy Arbiter for Goal-Conditioned Reinforcement Learning

2026-10-07

AI Summary

This study addresses the limited generalization of single policies and the difficulty of dynamic multi-policy selection in offline goal-conditioned reinforcement learning by proposing the WMPA framework. This work introduces a novel imagined-trajectory-based arbitration mechanism for policy ensembles. Specifically, it leverages a world model to simulate future trajectories of frozen policies and evaluates them via a shared goal-conditioned value function, enabling real-time arbitration of the optimal policy during test-time inference without retraining. This design facilitates adaptive policy switching across diverse environments. Experimental results demonstrate that WMPA increases the average success rate on the OGBench benchmark from 44% to 58%, significantly overcoming performance bottlenecks in complex manipulation tasks.

Abstract

Offline goal-conditioned reinforcement learning (GCRL) has produced a diverse set of goal-reaching algorithms, yet no single algorithm performs best across environments, goals, and even different phases of the same task. Rather than deploying only the best-performing policy, we ask whether a set of frozen goal-conditioned policies can be used collectively as a portfolio, deciding at every state which policy should act. Choosing a policy at each state is not straightforward. The policies' own value functions cannot be compared directly: they may use different scales, and some policies have no value function. We need to judge each policy by the states it is likely to reach, even though we can execute only one policy at a time. We also need to avoid switching so often that control becomes unstable. To address these challenges, we introduce World-Model Policy Arbiter (WMPA), a test-time framework that, given a set of frozen policies as input, rolls out each frozen policy in a learned state-space world model, evaluates the imagined futures with a shared goal-conditioned value function, and executes the highest-scoring policy for a short commitment interval before the next round of arbitration (policy selection). WMPA assumes access to a bank of frozen goal-conditioned policies and requires neither policy retraining nor privileged task-specific knowledge. Under the official OGBench evaluation protocol on 18 state-based datasets spanning maze navigation as well as cube, scene, and puzzle manipulation, WMPA improves the macro-average success rate from the 44% achieved by the best policy selected per dataset to 58%, with statistically significant gains on 12 datasets. These gains include +33 percentage points on cube-double-play and +36 percentage points on scene-play.

Problem

• goal-conditioned reinforcement learning
• policy selection
• offline RL
• policy portfolio
• value function comparison

Innovation

• Goal-Conditioned Reinforcement Learning
• World Model
• Policy Arbiter
• Offline RL
• Imagined Rollouts