HaWMPO: Hallucination-Aware World Model-based Policy Optimization for Generalist Robot Policy

Summary

Generalist robot policies have demonstrated strong generalization across robotic manipulation tasks, yet their success rates remain limited in com- plex long-horizon scenarios.

Abstract Summary

Generalist robot policies have demonstrated strong generalization across robotic manipulation tasks, yet their success rates remain limited in com- plex long-horizon scenarios. Recent methods improve Visual-Language-Action (VLA) policies through online reinforcement learning on real robots, but such training relies on costly physical interactions, suffers from low sample efficiency, and may introduce hardware and safety risks. World models offer a promising alternative by enabling policy optimization with imagined rollouts. However, long-horizon rollouts generated by world models often suffer from prediction hal- lucinations, producing biased state transitions that can mislead policy learning. To address this issue, we propose Hallucination-aware World Model-based Pol- icy Optimization (HaWMPO), a closed-loop reinforcement learning pipeline for VLA policy post-training with world models. Specifically, HaWMPO introduces an action-conditioned hallucination-aware model to estimate the reliability of gen- erated image sequences, and incorporates hallucination scores into group relative policy optimization through a Reward-Soft mechanism, suppressing unreliable ac- tion chunks during training. On the LIBERO benchmark, HaWMPO achieves the best average success rate, with gains of 15.0% over the base model and 2.8% over the strongest baseline; real-world experiments on a G1 robot further validate its effectiveness, raising the average success rate on two manipulation tasks from 67.5% to 80.0%.

Key Points

Generalist robot policies have demonstrated strong generalization across robotic manipulation tasks, yet their success rates remain limited in com- plex long-horizon scenarios.
Recent methods improve Visual-Language-Action (VLA) policies through online reinforcement learning on real robots, but such training relies on costly physical interactions, suffers from low sample efficiency, and may introduce hardware and safety risks.
World models offer a promising alternative by enabling policy optimization with imagined rollouts
However, long-horizon rollouts generated by world models often suffer from prediction hal- lucinations, producing biased state transitions that can mislead policy learning.
To address this issue, we propose Hallucination-aware World Model-based Pol- icy Optimization (HaWMPO), a closed-loop reinforcement learning pipeline for VLA policy post-training with world models.

Methodology

World-model prediction errors can accumulate during imagined rollouts and introduce misleading signals into VLA policy optimization. We introduce an action-conditioned Hallucination-Aware Model (HAM) to estimate the reliability of each generated video chunk.

HAM takes an eight-frame predicted chunk $\hat{\mathbf{V}}_{c}$, its conditioning actions $\mathbf{a}_{c}$, the current observation $I_{c}$, and an anchor image $I_{0}$ as inputs. Here, $I_{0}$ is the initial observation of the rollout, providing a fixed visual reference.

Methodology — HaWMPO: Hallucination-Aware World Model-based Policy Optimization for Generalist Robot Policy — Figure 2: Architecture of HAM. Predicted video frames, the current observation, the initial anchor image, and conditioning actions are encoded and fused to predict a chunk-level hallucination score. Higher scores indicate less reliable predictions.

Experiments

In this experiment, we follow the experimental setup of prior work and adopt OpenVLA-OFT as the base model. Before reinforcement learning, the base model is first supervised fine-tuned using a one-shot SFT setting, where only one expert trajectory is used for each task to obtain an initial policy.

To ensure a fair comparison with existing methods, we continue to use the LIBERO simulation benchmark as the primary evaluation environment.

Experiments — HaWMPO: Hallucination-Aware World Model-based Policy Optimization for Generalist Robot Policy — Figure 5: World-model predictions with higher and lower HAM scores. Higher-scored examples exhibit more evident discrepancies from the corresponding observed frames, whereas lower-scored examples more closely match the observations.

Sources and Demonstrations

Figures are reproduced from Chen et al. See the full paper for experimental details and the project page, when the paper links one, for demonstrations and videos.