Enhancing Policy Learning with World-Action Model

• URL: http://arxiv.org/abs/2603.28955v1
• Date: Mon, 30 Mar 2026 19:56:05 GMT
• Title: Enhancing Policy Learning with World-Action Model
• Abstract summary: World-Action Model (WAM) is an action-regularized world model that jointly reasons over future visual observations and the actions that drive state transitions.<n>WAM incorporates an inverse dynamics objective into DreamerV2 that predicts actions from latent state transitions.<n>We evaluate WAM on enhancing policy learning across eight manipulation tasks from the CALVIN benchmark.
• Score: 4.941630596191806
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: This paper presents the World-Action Model (WAM), an action-regularized world model that jointly reasons over future visual observations and the actions that drive state transitions. Unlike conventional world models trained solely via image prediction, WAM incorporates an inverse dynamics objective into DreamerV2 that predicts actions from latent state transitions, encouraging the learned representations to capture action-relevant structure critical for downstream control. We evaluate WAM on enhancing policy learning across eight manipulation tasks from the CALVIN benchmark. We first pretrain a diffusion policy via behavioral cloning on world model latents, then refine it with model-based PPO inside the frozen world model. Without modifying the policy architecture or training procedure, WAM improves average behavioral cloning success from 59.4% to 71.2% over DreamerV2 and DiWA baselines. After PPO fine-tuning, WAM achieves 92.8% average success versus 79.8% for the baseline, with two tasks reaching 100%, using 8.7x fewer training steps.

Related papers

• PAVXploreRL: Physical-Action-Visual World Model Reinforcement Learning with Action Exploration [49.59412060345528]
  Action-conditioned world models are a key component of embodied AI, serving as scalable policy evaluators.<n>We propose PAVXploreRL, a reinforcement learning framework built on a pretrained latent world model.<n>Experiments show that PAVXploreRL consistently outperforms pretrained baselines.
  arXiv  Detail & Related papers  (2026-07-18T02:42:27Z)
• DREAMSTEER: Latent World Models Can Steer VLA Policies During Deployment Without Any Finetuning [19.23819338884138]
  DREAMSTEER is a deployment-time steering framework for pretrained vision-language-action (VLA) policies.<n>Itleverages a latent world model and a value model to steer pretrained VLA policies.<n>It improves task success rate from 23.75% to 66.25% and instruction-following accuracy from 38.75% to 56.25% over the base VLA policy.
  arXiv  Detail & Related papers  (2026-07-03T02:08:50Z)
• Stealthy World Model Manipulation via Data Poisoning [6.257560917160816]
  We propose SWAAP, the first two-stage data poisoning framework for learned world models.<n>In the first stage, SWAAP identifies a harmful target world model that induces low-return behavior under planning.<n>In the second stage, SWAAP realizes this target through stealth-constrained gradient matching.
  arXiv  Detail & Related papers  (2026-06-17T05:24:18Z)
• FlowMPC: Improving Flow Matching policies with World Models [0.0]
  Flow Matching (FM) is a powerful approach for behavior cloning in multimodal action spaces.<n>Because it is not trained to directly maximize expected return, there is still room to improve how FM policies act at test time.<n>This work investigates whether a learned world model can improve FM policies by enabling Model Predictive Path Integral (MPPI) planning over candidate action sequences proposed by the policy.
  arXiv  Detail & Related papers  (2026-06-15T06:50:11Z)
• Feat2Go: Visual Feature-Grounded Value Estimation for Embodied Reinforcement Learning [14.672595536029503]
  Feat2Go is a fine-grained value estimation framework for embodied reinforcement learning.<n>We show that Feat2Go consistently improves the performance of existing VLA models under both single-arm and bimanual manipulation settings.
  arXiv  Detail & Related papers  (2026-05-29T03:36:33Z)
• Hi-WM: Human-in-the-World-Model for Scalable Robot Post-Training [54.896907620476675]
  We propose a post-training framework that uses a learned world model as a reusable corrective substrate for failure-targeted policy improvement.<n>Hi-WM caches intermediate states and supports rollback and branching, allowing a single failure state to be reused for multiple corrective continuations.<n>We evaluate Hi-WM on three real-world manipulation tasks spanning both rigid and deformable object interaction, and on two policy backbones.
  arXiv  Detail & Related papers  (2026-04-23T14:42:54Z)
• Grounded World Model for Semantically Generalizable Planning [94.53923128709965]
  We learn a Grounded World Model (GWM) in a vision-language-aligned latent space.<n>Each proposed action is scored based on how close its future outcome is to the task instruction.<n>On the proposed WISER benchmark, GWM-MPC achieves a 87% success rate on the test set.
  arXiv  Detail & Related papers  (2026-04-13T17:25:41Z)
• Better World Models Can Lead to Better Post-Training Performance [9.713688760042544]
  We study how explicit world-modeling objectives affect the internal representations and downstream capability of Transformers.<n>We compare standard next-token prediction to two explicit world-modeling strategies.<n>We find that explicit world-modeling yields more linearly decodable and causally steerable state representations.
  arXiv  Detail & Related papers  (2025-12-03T03:13:20Z)
• Dual-Stream Diffusion for World-Model Augmented Vision-Language-Action Model [62.889356203346985]
  We propose DUal-STream diffusion (DUST), a world-model augmented VLA framework that handles the modality conflict.<n>DUST achieves up to 6% gains over a standard VLA baseline and implicit world-modeling methods.<n>On real-world tasks with the Franka Research 3, DUST outperforms baselines in success rate by 13%.
  arXiv  Detail & Related papers  (2025-10-31T16:32:12Z)
• Internalizing World Models via Self-Play Finetuning for Agentic RL [65.96875390986655]
  Large Language Models (LLMs) as agents often struggle in out-of-distribution (OOD) scenarios.<n>We show how to encode this world model by decomposing it into two components: state representation and transition modeling.<n>We introduce SPA, a simple reinforcement learning framework that cold-starts the policy via a Self-Play supervised finetuning stage to learn the world model.
  arXiv  Detail & Related papers  (2025-10-16T18:03:39Z)
• CO-RFT: Efficient Fine-Tuning of Vision-Language-Action Models through Chunked Offline Reinforcement Learning [7.780242426487376]
  We propose Chunked RL, a novel reinforcement learning framework for Vision-Language-Action (VLA) models.<n>Within this framework, we extend temporal difference (TD) learning to incorporate action chunking, a prominent characteristic of VLA models.<n>We then propose CO-RFT, an algorithm aimed at fine-tuning VLA models using a limited set of demonstrations.
  arXiv  Detail & Related papers  (2025-08-04T09:11:48Z)
• Pre-Trained Policy Discriminators are General Reward Models [81.3974586561645]
  We propose a scalable pre-training method named Policy Discriminative Learning (POLAR)<n>POLAR trains a reward model (RM) to discern identical policies and discriminate different ones.<n> Empirical results show that POLAR substantially outperforms traditional non-pre-trained methods.
  arXiv  Detail & Related papers  (2025-07-07T16:56:31Z)
• ConRFT: A Reinforced Fine-tuning Method for VLA Models via Consistency Policy [10.596344084789434]
  Vision-Language-Action (VLA) models have shown substantial potential in real-world robotic manipulation.<n>Fine-tuning these models through supervised learning struggles to achieve robust performance due to limited, inconsistent demonstrations.<n>We propose a reinforced fine-tuning approach for VLA models, named ConRFT, which consists of offline and online fine-tuning.
  arXiv  Detail & Related papers  (2025-02-08T05:01:17Z)
• Pre-trained Model Guided Fine-Tuning for Zero-Shot Adversarial Robustness [52.9493817508055]
  We propose Pre-trained Model Guided Adversarial Fine-Tuning (PMG-AFT) to enhance the model's zero-shot adversarial robustness. Our approach consistently improves clean accuracy by an average of 8.72%.
  arXiv  Detail & Related papers  (2024-01-09T04:33:03Z)
• Live in the Moment: Learning Dynamics Model Adapted to Evolving Policy [13.819070455425075]
  We learn a dynamics model that fits under the empirical state-action visitation distribution for all historical policies. We then propose a novel dynamics model learning method, named textitPolicy-adapted Dynamics Model Learning (PDML). Experiments on a range of continuous control environments in MuJoCo show that PDML achieves significant improvement in sample efficiency and higher performance combined with the state-of-the-art model-based RL methods.
  arXiv  Detail & Related papers  (2022-07-25T12:45:58Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.