PaperlayerView exact revision on arXiv

Source-linked AI summary

World-Model Policy Arbiter for Goal-Conditioned Reinforcement Learning

Junwei Quan, Evgenii Opryshko, Nicholas Rhinehart, Igor Gilitschenski

arXiv:2610.10932v1cs.LG

TL;DR

Offline goal-conditioned reinforcement learning lacks a single method that dominates across environments, goals, and task phases, while frozen policies may be difficult to compare directly. WMPA evaluates their imagined futures with a shared value function and commits to the highest-ranked policy. Across 18 OGBench datasets, it raises mean success from 44% for the best policy to 58%.

  • Problem

    No single goal-conditioned method dominates across environments, goals, and task horizons, and practitioners may discard policies that could help in other task phases.

    Source: §1 · pp. 1–2

  • Method

    WMPA rolls frozen policies forward in a learned world model, scores their imagined futures with a shared goal-conditioned value function, and executes the highest-scoring policy for a commitment interval.

    Source: §7 · p. 9

  • Results

    Across 18 OGBench datasets, WMPA achieves 58% mean success versus 44% for the best policy, with significantly positive gains on 12 datasets.

    Source: §4.1 · p. 6

  • Takeaways & Limitations

    The results suggest using world models as test-time evaluators of existing behaviors, not only as generators of new actions.

    Source: §7 · p. 10

    The study uses state observations and per-environment transition models, leaving visual environments to future work; WMPA also cannot create behaviors absent from its policy bank.

    Source: §6 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

Offline goal-conditioned policies can excel at different environments, goals, and phases, motivating test-time selection among a frozen portfolio. WMPA compares their imagined futures and can combine complementary policies within an episode.

1 INTRODUCTION

  • On cube-double-play, WMPA succeeds in 69% of episodes, versus 36% for GCIQL and 5% for HIQL, the strongest and weakest of six standard methods, respectively.The example illustrates policies contributing different task phases rather than one global winner.Source: §1 · p. 1
  • Offline GCRL practitioners typically keep the candidate with the highest average success and discard the rest, potentially losing policies useful for particular states or phases.The two-cube example retains GCIQL but discards HIQL, which moves the gripper between cubes.Source: §1 · pp. 1–2
  • Policy values cannot reliably rank candidates because algorithms use different scales or have no value function, alternatives' consequences are unobserved, and frequent switching can disrupt control.The cited examples include temporal-difference values, quasimetric distances, contrastive scores, and GCBC without a value function.Source: §1 · p. 2
  • WMPA rolls each frozen policy forward in a learned world model, scores imagined futures with a shared goal-conditioned value, then executes the highest-ranked policy for a commitment interval.It evaluates existing behaviors rather than optimizing new action sequences, and needs no retraining, real-environment rollouts, hand-designed progress signal, or privileged task knowledge.Source: §1 · p. 2
  • The paper frames offline GCRL as policy arbitration: test-time control over independently trained goal-conditioned policies, distinct from choosing one global winner or retargeting one controller's subgoal.It argues offline GCRL incurs no additional data-collection cost for the policy bank, and goal-conditioning multiplies situations where different members excel.Source: §1 · p. 2

2 BACKGROUND ON GOAL-CONDITIONED REINFORCEMENT LEARNING

The setup defines goal-conditioned policies and evaluates them by goal-reaching success over initial-state and goal pairs. The paper studies switching among a frozen portfolio without modifying its members.

2 BACKGROUND ON GOAL-CONDITIONED REINFORCEMENT LEARNING

  • A goal-conditioned policy maps the current state and goal to a distribution over actions in a controlled Markov process.The process specifies state, action, and goal spaces, an initial-state distribution, and transition dynamics.Source: §2 · p. 3
  • Performance is expected goal-reaching success for trajectories induced by initial-state and goal pairs drawn from an evaluation distribution.Success is an evaluation-only criterion; training supervision instead uses hindsight goal relabeling from logged trajectory states.Source: §2 · p. 3
  • The experiments train the auxiliary world model and value function on the offline dataset, but the general formulation does not require policies to share a dataset or be trained offline.In the experiments, the dataset is the OGBench training split, also used to train all policies offline.Source: §2 · p. 3
  • The portfolio contains frozen policies with a common state, action, and goal interface; the arbiter needs only black-box action distributions, not comparable internal values or representations.At each arbitration point, one member controls the agent for a commitment interval; the objective is expected goal-reaching success without modifying any member.Source: §2 · p. 3

3 METHOD: WORLD-MODEL POLICY ARBITER

WMPA repeatedly imagines the closed-loop futures of frozen policies, scores them with a shared goal-conditioned value function, and commits to the highest-ranked policy. Its learned state-space model, task-family-specific scorer, and rollout ranking implement this arbitration loop.

3 METHOD: WORLD-MODEL POLICY ARBITER

  • At each arbitration point, WMPA imagines k steps for every frozen policy, scores the predicted states with one goal-conditioned value function, and executes the winner for c environment steps.The shared scorer puts candidate futures on a common scale; commitment and value-based selection limit excessive switching.Source: §3 · p. 3, Figure 2 · p. 4

3.1 LEARNED STATE-SPACE WORLD MODEL

  • The state-space world model uses an ensemble of independently parameterized MLP dynamics models trained on environment states and multi-step prediction windows.Its curriculum increases the effective horizon; state-space predictions remain valid policy inputs and decouple model learning from changes to the policy bank.Source: §3.1 · pp. 3–4

3.2 GOAL-CONDITIONED VALUE FUNCTION

  • The shared scorer uses latent metric distance for geometric progress and a direct value for puzzle tasks, where embedding distance need not track progress.Metric distance orders imagined state–goal pairs, including when rollouts leave logged-data support; the direct value is reused from the bank's GCIQL member and estimates the best dataset-supported action value.Source: §3.2 · pp. 4–5

3.3 MODEL-BASED POLICY ROLLOUT AND RANKING

  • Each candidate is rolled out closed-loop for k steps using deterministic policy actions, and its score is the highest value of any imagined state along the rollout.The maximum gives credit for reaching the goal before the horizon ends; WMPA executes the highest-scoring policy for c steps and repeats arbitration.Source: §3.3 · p. 5

4 EXPERIMENTS

Across 18 state-based OGBench datasets, WMPA improves success over individual and switching baselines. The experiments test its performance, scoring and rollout choices, and whether gains depend on composing policies within episodes.

4 EXPERIMENTS

  • The study evaluates six frozen OGBench reference learners across 18 state-based datasets, using three bank seeds and paired episode comparisons.Reported intervals are 95% hierarchical bootstraps over bank seeds and then episodes.Source: §4 · p. 6

4.1 POLICY ARBITER IMPROVES GOAL REACHING

  • 58% mean success versus 44% for the best policy and 39% for random switching; WMPA has the highest mean on 17 datasets and significant gains on 12.Its gain over the best policy is not significant on 5 datasets and is significantly negative on 1.Source: §4.1 · p. 6
  • 53% average success for action-level MPC versus 58% for WMPA, despite MPC receiving k times the model budget; MPC is significantly worse on 7 datasets and better on 1.The comparison supports arbitration among policies beyond locally refining actions from a single best policy.Source: §4.1 · p. 6

4.2 DOES THE WORLD MODEL IMPROVE PERFORMANCE? USUALLY.

  • 19 points higher average success than RANDOM-SWITCH, with significant gains on 13 datasets, attributes WMPA’s gains to model-based candidate ranking rather than commitment scheduling or policy cycling.RANDOM-SWITCH removes both world-model rollouts and the value head while retaining the same commitment schedule.Source: §4.2 · p. 7
  • A displacement-triggered random switch failed to recover manipulation gains, while true-simulator branches added +2 points on average, so model error was not the main performance limit.The displacement threshold was fixed across datasets without tuning.Source: §4.2 · p. 7
  • Q-select exceeds the best policy on 10 of 18 datasets, but WMPA exceeds Q-select by +12 points averaged across datasets; at matched commitment, direct-value WMPA gains +7 points.At the matched interval, direct-value WMPA has significant gains on 11 datasets and no significant losses.Source: §4.2 · p. 7
  • Own-value rollout scoring is significantly worse than shared scoring on 14 of 18 datasets, even among policies whose value functions estimate the same quantity.The selector collapses onto the most optimistic value function.Source: §4.2 · p. 7

4.3 ARBITRATION BENEFITS FROM COMPLEMENTARY BEHAVIORS WITHIN EPISODES

  • WMPA significantly exceeds the per-episode oracle on 7 datasets, including cube-double-play at 69% versus 57% and scene-play at 87% versus 74%.Because the oracle picks a successful fixed policy per episode when one exists, these gains indicate a benefit from composing policies within episodes.Source: §4.3 · p. 7, Figure 3 · p. 8

5 RELATED WORK

Prior work uses models to generate or plan behavior, or selects among policies with comparable values. WMPA instead uses a learned model to arbitrate among frozen, heterogeneous goal-conditioned policies while keeping the final goal fixed.

Goal-conditioned reinforcement learning

  • Model-based GCRL methods modify or create policies, whereas WMPA leaves every policy frozen and uses the model only to choose among them.The cited approaches include planning to fine-tune offline policies (Wang et al., 2024) and planning exploration goals (Hu et al., 2023).Source: §5 · p. 8

Test-time planning with frozen policies and world models

  • Test-time planning methods often choose subgoals for a controller, while WMPA keeps the final goal fixed and chooses which controller acts now.CompPlan also sequences pretrained policies with a learned model, but searches over subgoals for one policy family rather than arbitrating among heterogeneous policies.Source: §5 · p. 9

Policy switching, options, and learned hierarchies

  • Classical policy switching simulates candidate policies, while generalized policy improvement requires comparable per-policy action values; WMPA uses learned-model rollouts and a shared value function instead.WMPA learns neither a high-level policy nor a low-level policy; its dynamics model is trained on the dataset without reference to any policy.Source: §5 · p. 9

6 LIMITATIONS AND FUTURE WORK

The study uses state observations and per-environment transition models, and WMPA is limited to behaviors already present in its policy bank.

6 LIMITATIONS AND FUTURE WORK

  • The evaluation uses state observations and per-environment transition models; visual environments are left to future work.Pixel observations would require a shared visual encoder across policies or a world model that produces realistic images.Source: §6 · p. 9
  • Arbitration cannot create new behaviors and typically does no better than the best policy when bank members share failures, are at task floor, or one member is already saturated.WMPA selects among behaviors that already exist in the bank.Source: §6 · p. 9
  • Arbitration adds a few milliseconds of computation per step; proposed future work includes distilling its choices or generating corrected training data.The paper describes this computation as modest overhead.Source: §6 · p. 9

7 CONCLUSION

WMPA selects among frozen goal-conditioned policies by scoring their imagined futures with a shared value function. It improves on a per-episode hindsight oracle on 7 datasets, supporting a view of world models as evaluators of existing behaviors.

7 CONCLUSION

  • WMPA improves on its bank’s best policy by 13 points on average and by 33–41 points on several multi-object manipulation tasks, without retraining policies.It rolls frozen policies forward in a learned world model and executes the highest-scoring policy for a short commitment interval.Source: §7 · p. 9
  • Exceeding a per-episode hindsight oracle on 7 datasets indicates a benefit from composing policies within an episode; the paper frames world models as evaluators of existing behaviors.The model evaluates behaviors represented in the policy bank rather than generating new actions.Source: §7 · p. 10

REPRODUCIBILITY STATEMENT

The release includes code for training the policy bank and world model, running evaluations, and aggregating success-rate intervals. It also records dataset-specific checkpoints and evaluation setup for reproducibility.

REPRODUCIBILITY STATEMENT

  • The code includes entry points for training the policy bank and world model, plus scripts to run WMPA, controls, and ablations.It also aggregates per-episode logs into success rates and paired hierarchical-bootstrap intervals.Source: REPRODUCIBILITY STATEMENT · p. 10
  • A per-dataset registry records the bank, checkpoints, and reported intervals for all 18 datasets; software versions are also listed.The trained checkpoints and raw evaluation logs are omitted because of their size.Source: REPRODUCIBILITY STATEMENT · p. 10
  • Reset seeds are deterministic functions of dataset, goal, and episode index, allowing methods to be re-evaluated on the same initial states.The datasets are public OGBench releases.Source: REPRODUCIBILITY STATEMENT · p. 10

A PER-DATASET PLANNING SETTINGS

The appendix reports settings and sensitivity for WMPA’s planning horizon and commitment interval, then compares learned arbitration with random switching and stall-triggered restarts.

A PER-DATASET PLANNING SETTINGS

  • WMPA imagines each policy for horizon k, scores its imagined futures, then executes the highest-scoring policy for up to c steps before arbitrating again.Table 2 sets k=c as both WMPA’s horizon and commitment and RANDOM-SWITCH’s redraw interval.Source: A PER-DATASET PLANNING SETTINGS · p. 12, Table 2 · p. 12
  • Across all 18 imagination-horizon and commitment cells in Table 3, every cell retains a gain of at least 26 points over the best policy on three datasets with large gains.The interval is not extensively tuned per dataset; it is kept the same within each dataset family.Source: A PER-DATASET PLANNING SETTINGS · p. 12, Table 3 · p. 13

B NO-MODEL SWITCHING CONTROLS

  • At matched commitment intervals, WMPA averages 58% success versus 39% for RANDOM-SWITCH and is significantly better on 13 datasets.The methods differ in whether the next policy is chosen by learned score or uniform draw.Source: Table 4 · p. 13, B NO-MODEL SWITCHING CONTROLS · p. 13
  • Across manipulation datasets, stall-restart does not recover WMPA’s gains, although it can reach 94–99% on pointmaze-medium-navigate versus WMPA’s 79%.The controller restarts with a randomly selected other policy when recent displacement falls below a threshold.Source: B NO-MODEL SWITCHING CONTROLS · pp. 13–14, Table 5 · p. 14

C MODEL-BASED CONTROLS AND SCORER COMPARISONS

The controls compare policy arbitration with action-level search and alternative selectors or scorers. The analyses find that shared-value scoring matters, while tested score reductions have little effect on macro-average success.

C MODEL-BASED CONTROLS AND SCORER COMPARISONS

  • Direct-value WMPA exceeds Q-select by +7 points on average at matched commitment, significantly on 11 of 18 datasets; Q-select is significantly above the best policy on 10 datasets at c=1 and 5 at c=k.Q-select uses a twin-Q value function to rank proposed actions without imagined states.Source: C MODEL-BASED CONTROLS AND SCORER… · p. 14
  • Own-value scoring averages 4 percentage points below the best of the three members, while shared scoring improves on that member by 11 percentage points.With own values, selection collapses toward the most optimistic member; shared scoring gains significantly over the best member on 12 datasets and own-value scoring on 14.Source: C MODEL-BASED CONTROLS AND SCORER… · p. 14
  • The post hoc concordance criterion selects the metric head on cube-double-play and scene-play, and the direct head on puzzle-4x4-play and cube-triple-play.On cube-triple-play, the criterion favors the direct head, which is also better in online evaluation.Source: C MODEL-BASED CONTROLS AND SCORER… · p. 15
  • Action-level MPC replans each step over perturbed actions from the best policy, using the same horizon and value but 18k model steps per environment step versus WMPA’s 18.Its model budget is 5× larger on cube datasets and 10× larger on scene and puzzle datasets.Source: C MODEL-BASED CONTROLS AND SCORER… · p. 14, Table 6 · p. 15

D SCORE ABLATIONS

  • Changing horizon aggregation shifts macro-average success by at most one point (final state +1, mean -1); changing ensemble reduction does likewise (min +0, mean−std +0).These reductions do not eliminate the low-success failure cases.Source: D SCORE ABLATIONS · p. 15

E BANK COMPOSITION

The bank-composition analysis tests pooled policies trained on different data, while the appendix describes model architecture and training. Gains depend on whether added policies contribute competencies not already represented.

E BANK COMPOSITION

  • The appendix defines controls that replace learned model branches with true-simulator branches, compare Q-select with imagined-state scoring, and contrast member-specific with shared value scoring.The table captions specify the control configurations and matched evaluation settings.Source: Table 7 · p. 16, Table 8 · p. 16, Table 9 · p. 17
  • The pooled-bank test adds six policies trained on sibling-dataset data to a six-member bank, without retraining, while keeping the base dataset’s model, value function, interval, and episodes.The policies share the same environment and observation, action, and goal interfaces.Source: E BANK COMPOSITION · pp. 15–16, Table 12 · p. 18
  • The appendix describes a post hoc head-selection diagnostic based on whether value rankings agree with observed time-to-goal ordering in offline trajectories.Pairs with arrival times differing by fewer than 20 steps are discarded, and confidence intervals are bootstrapped over trajectories.Source: Table 10 · p. 17
  • On puzzle-4x4-play, WMPA gains +18 [+11, +26]* points over the pooled bank’s best policy, which is itself a foreign, noisy-trained GCIQL policy.That policy succeeds on 34% of play episodes versus 24% for the native policy.Source: E BANK COMPOSITION · p. 17
  • Where added policies contribute new competencies, the pooled bank gains +10 [+6, +15]* on cube-double-play; where they duplicate competencies, it matches the original on scene-play and puzzle-4x4-play.The corresponding contrasts are +1 [-3, +5] for scene-play and +2 [-4, +8] for puzzle-4x4-play.Source: E BANK COMPOSITION · p. 17, Table 12 · p. 18

F WORLD-MODEL ARCHITECTURE AND TRAINING DETAILS

  • The dynamics model uses three independently initialized ensemble members, each a three-layer, width-512 MLP trained to predict normalized observation deltas.Training unrolls five steps for cube, scene, and puzzle tasks and ten for the two mazes.Source: F WORLD-MODEL ARCHITECTURE AND… · p. 19
  • The metric value head scores states by negative latent distance to the encoded goal; training uses relabeled goals and, for mazes and scene, a local-smoothness regularizer.The puzzle-family direct value head follows the GCIQL design and is loaded from a pretrained checkpoint, remaining frozen during world-model training.Source: F WORLD-MODEL ARCHITECTURE AND… · p. 19

G EVALUATION PROTOCOL, STATISTICS, AND THE STATIC PER-EPISODE ORACLE

The evaluation compares WMPA with a static per-episode hindsight oracle under paired episodes, then reports portfolio use and runtime alongside conditions for a theoretical improvement guarantee.

G EVALUATION PROTOCOL, STATISTICS, AND THE STATIC PER-EPISODE ORACLE

  • WMPA exceeds the static per-episode oracle by at least one percentage point on 10 datasets; paired intervals favor WMPA on 7, while the oracle is hindsight-only.Evaluation uses 5 goals × 50 episodes × 3 bank seeds per dataset, with paired episodes and hierarchical bootstrap over bank seeds and episodes (104 resamples, 95% intervals).Source: G EVALUATION PROTOCOL, STATISTICS, AND… · pp. 19–20, Table 13 · p. 20

H SELECTION BEHAVIOR AND RUNTIME

  • WMPA switches less often than RANDOM-SWITCH at the same interval, and policy-usage entropy is high on datasets with the largest gains.Test-time latency grows linearly with bank size; twelve-member pooled banks run at 5–6 ms per step.Source: H SELECTION BEHAVIOR AND RUNTIME · p. 20

I WHEN DO IMAGINED ROLLOUTS SUPPORT USEFUL ARBITRATION?

  • Under the stated separation assumption, a positive theorem bound guarantees improvement over every fixed bank member.The bound weighs available improvement against tolerated chunk-advantage loss and score errors large enough to overturn the separation margin.Source: I WHEN DO IMAGINED ROLLOUTS SUPPORT… · p. 22
  • Useful arbitration requires locally advantageous alternatives, true-dynamics score separation, and errors small relative to that margin; accurate dynamics alone are insufficient.Errors below half the margin cannot select a policy more than δ below the best chunk advantage.Source: I WHEN DO IMAGINED ROLLOUTS SUPPORT… · p. 22
Loading 2610.10932v1…