World-Model Policy Arbiter for Goal-Conditioned Reinforcement Learning

The gist 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 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.

How it works

The World-Model Policy Arbiter (WMPA) is a test-time framework that selects among a set of frozen goal-conditioned policies by rolling each one out in a learned state-space world model, evaluating the imagined futures with a shared goal-conditioned value function, and executing the highest-scoring policy for a short commitment interval before selecting again. This process addresses challenges such as comparing policies that may use different scales or have no value function by scoring them based on their predicted consequences.

The WMPA mechanism involves three main components:

1. Learned State-Space World Model: An ensemble of independently parameterized MLP dynamics models predicts the next state given the current state and action, trained to match non-normalized environment states. This model is decoupled from the policies, allowing it to be learned without modifying the policy bank.

2. Goal-Conditioned Value Function: A shared value function scores every candidate’s imagined states on a common scale, ensuring that policies trained with different objectives are compared appropriately. The paper describes two parameterizations for this function, one using a latent distance for geometric progress and another using a direct MLP for combinatorial progress.

3. Model-Based Policy Rollout and Ranking: At each arbitration point, every frozen policy is rolled out in closed loop for k model steps to generate imagined states, which are then scored by the shared value function to yield a score Si(st, g). The policy index with the highest score i⋆ = arg maxi Si(st, g) is selected and executed for a commitment interval of length c.

Key Contributions and Findings

The paper makes several key contributions to the field of offline GCRL. First, it introduces the problem of policy arbitration for offline GCRL, arguing that this setting is natural because the policy bank incurs no additional data collection cost and goal-conditioning multiplies situations where different members excel. Second, it details the World Model Policy Arbiter (WMPA) method, which compares frozen policies by scoring their closed-loop imagined futures with a shared goal-conditioned value function without retraining or privileged signals. Third, it derives a conditional guarantee over the best fixed bank member, showing that arbitration improves success when locally useful alternatives are available and true-dynamics scores separate inferior choices.

Empirical evaluation on 18 state-based OGBench datasets shows significant performance gains for WMPA. WMPA achieves a mean success rate of 58%, compared to 44% for the best policy and 39% for random switching with the same interval and policy bank. The largest improvements occur on multiobject manipulation tasks, including +41 percentage points on scene-noisy and +35 on cubedouble-noisy. Furthermore, WMPA outperforms action-level model-predictive control (MPC) applied to the best policy with the same horizon k and value function in several datasets.

Analysis of Arbitration Benefits

The study investigates whether the gains from policy switching arise solely from selecting the best policy for each episode or if they also rely on within-episode composition. The per-episode oracle, which selects a successful bank policy for each episode whenever one exists, is significantly outperformed by WMPA on 7 datasets. This indicates a benefit of composing policies within an episode. The analysis shows that the headroom between the best and second-best policy precludes gains from switching alone.

Further ablation studies clarify the source of performance gains. A random-switch baseline, which draws a policy uniformly from the bank at every arbitration point, is outperformed by WMPA by 19 points on average, with significant gains on 13 datasets. This suggests that WMPA’s major gains come neither from the commitment schedule nor from cycling through policies, but from model-based candidate ranking. The stall-restart controller, which switches only when the agent stalls rather than on a fixed schedule, does not recover the gains of WMPA on manipulation datasets.

Value Head and Model Influence

The choice of value head significantly impacts performance. Table 8 shows that Q-select, which uses a twin-Q value function to rank actions without imagined states, is significantly above the best policy on 10 of 18 datasets at c=1. WMPA scores imagined states with either the metric value or the direct value. The results indicate that scoring every rollout with a single value function is essential for a meaningful comparison among policies, as scoring each rollout with its own value function fails to yield gains over the shared value function on 12 and 14 datasets.

The study also examines the effect of dynamics-model accuracy, finding that replacing learned dynamics with true-simulator branches raises the success rate by +2 points on average, with significant gains on 2 of 18 datasets. Finally, Table 7 demonstrates that when pooling policies trained on play and noisy data into a single twelve-member bank, WMPA outperforms the original six-member run whenever the added members contribute new competencies.

Conclusion

WMPA successfully selects among frozen goal-conditioned policies at test time by rolling each forward in a learned world model, scoring the imagined futures with a shared value function, and executing the highest-scoring policy for a short commitment interval before selecting again. The results suggest viewing world models as test-time evaluators of existing behaviors, not only as generators of new actions. The method requires three requirements for useful arbitration: the bank must offer locally advantageous alternatives, the value head must separate substantially inferior choices under true dynamics, and score errors must be small relative to that separation.

--- Page 18 ---

Preprint. Under review.

Table 13 reports, for each dataset, the bank statistics introduced in Section 4.3. Here, Bmax denotes the success rate of the best policy, and Dgap denotes the performance gap between the best and second-best policies.

--- Page 20 ---

Table 14 reports how the portfolio is used.

--- Page 21 ---

Table 5 reports results for the stall-restart controller.

--- Page 23 ---

Table 6 is discussed in Section 4.1 and Table 7 is discussed in Section 4.2.

--- Page 24 ---

Table 8 reports comparison of value heads and selectors at a fixed cell on all the datasets.

--- Page 25 ---

Table 9 tests whether the value functions of the policies themselves can replace the shared value function in the setting most favorable to them.

--- Page 17 ---

Table 10 presents a post hoc diagnostic analysis of value-head ranking quality.

--- Page 22 ---

Table 13 reports bank statistics and the static per-episode oracle.

--- Page 23 ---

Table 6 reports WM action search vs. arbitration with the same horizon k and the same value function on every dataset.

--- Page 24 ---

Table 8 compares value heads and selectors at a fixed cell on all the datasets.

--- Page 15 ---

Table 9 tests whether the value functions of the policies themselves can replace the shared value function in the setting most favorable to them.

--- Page 16 ---

Table 7 reports true-simulator branches inside WMPA at the selected cell on every dataset.

--- Page 19 ---

Table 12 reports pooled play+noisy banks.

--- Page 20 ---

Table 8 compares value heads and selectors at a fixed cell on all the datasets.

--- Page 24 ---

Table 13 reports bank statistics and the static per-episode oracle.

--- Page 18 ---

Table 6 compares WM action search vs. arbitration with the same horizon k and the same value function on every dataset.

--- Page 21 ---

Table 9 tests whether the value functions of the policies themselves can replace the shared value function in the setting most favorable to them.

1. Bold header: World-Model Policy Arbiter for Test-Time Control

This system can select among a set of frozen goal-conditioned policies at every state by rolling each forward in a learned world model and scoring the imagined futures with a shared goal-conditioned value function, ensuring that the highest-scoring policy is executed for a short commitment interval before the next round of arbitration.

2. Bold header: Robust Performance Gains on Manipulation Tasks

The WMPA system can achieve significant performance improvements, such as +36 percentage points on cube-double-play and +35 percentage points on scene-play, by leveraging compositional planning within an episode, which is identified as a key source of gain over simply selecting the single best policy.

3. Bold header: Robustness to Model Errors

The framework is designed to be resilient to inaccuracies in the learned dynamics, as Table 7 replaces the learned dynamics with true-simulator branches on every dataset and results show that this substitution raises the success rate by +2 points on average, with significant gains on 2 of 18 datasets.

4. Bold header: Separation of Inferior Choices

The system relies on a strong selection criterion to avoid poor choices, as Assumption I.1 (Separation of inferior choices) requires that the value head, when applied to true-dynamics rollouts, to separate any member whose advantage is more than δ below the best by a score margin of at least ∆.

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.