mlsanityData to 7 Oct 2026

PWM: Policy Learning with Multi-Task World Models

arXiv 2407.02466 ↗First posted 2 Jul 2024cs.LG cs.AI cs.ROcited by 17

On the map: World models → Latent dynamics models and model-based RL

Abstract (arXiv)

Reinforcement Learning (RL) has made significant strides in complex tasks but struggles in multi-task settings with different embodiments. World model methods offer scalability by learning a simulation of the environment but often rely on inefficient gradient-free optimization methods for policy extraction. In contrast, gradient-based methods exhibit lower variance but fail to handle discontinuities. Our work reveals that well-regularized world models can generate smoother optimization landscapes than the actual dynamics, facilitating more effective first-order optimization. We introduce Policy learning with multi-task World Models (PWM), a novel model-based RL algorithm for continuous control. Initially, the world model is pre-trained on offline data, and then policies are extracted from it using first-order optimization in less than 10 minutes per task. PWM effectively solves tasks with up to 152 action dimensions and outperforms methods that use ground-truth dynamics. Additionally, PWM scales to an 80-task setting, achieving up to 27% higher rewards than existing baselines without relying on costly online planning. Visualizations and code are available at https://www.imgeorgiev.com/pwm/.

Key facts. Cited by 17 papers in the map (data through 7 Oct 2026). Placed in Latent dynamics models and model-based RL, one of 11 directions in World models. Builds on 30 of its 41 references that are in the map. First posted on arXiv 2 Jul 2024.

Builds on

30 of 41 references.

Closest ancestors · 3

Same direction, newest first

  1. 2023TD-MPC2: Scalable, Robust World Models for Continuous Control25 Oct
  2. 2023Mastering Diverse Domains through World Models10 Jan
  3. 2019Dream to Control: Learning Behaviors by Latent Imagination3 Dec

Other references · 27

  1. 2020Language Models are Few-Shot Learnerscited by 32,199
  2. 2017Proximal Policy Optimization Algorithmscited by 14,441
  3. 2023Segment Anythingcited by 8,158
  4. 2016Layer Normalizationcited by 5,579
  5. 2021On the Opportunities and Risks of Foundation Modelscited by 3,961
  6. 2013Playing Atari with Deep Reinforcement Learningcited by 3,211
  7. 2023RT-2: Vision-Language-Action Models Transfer Web Knowledge to Robotic Controlcited by 2,993
  8. 2024Octo: An Open-Source Generalist Robot Policycited by 1,427
  9. 2018Soft Actor-Critic Algorithms and Applicationscited by 1,179
  10. 2019Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learningcited by 1,168
  11. 2023Open X-Embodiment: Robotic Learning Datasets and RT-X Modelscited by 849
  12. 2022A Generalist Agentcited by 802
15 more
  1. 2021Deep Reinforcement Learning at the Edge of the Statistical Precipicecited by 781
  2. 2015Variational Dropout and the Local Reparameterization Trickcited by 489
  3. 2017A Brief Survey of Deep Reinforcement Learningcited by 466
  4. 2022Temporal Difference Learning for Model Predictive Controlcited by 375
  5. 2020dm_control: Software and Tasks for Continuous Controlcited by 368
  6. 2019Monte Carlo Gradient Estimation in Machine Learningcited by 281
  7. 2023RVT: Robotic View Transformer for 3D Object Manipulationcited by 224
  8. 2023Foundation Models in Robotics: Applications, Challenges, and the Futurecited by 212
  9. 2018Model-Based Value Estimation for Efficient Model-Free Reinforcement Learningcited by 152
  10. 2015Gradient Estimation Using Stochastic Computation Graphscited by 120
  11. 2022Accelerated Policy Learning with Parallel Differentiable Simulationcited by 92
  12. 2020Objective Mismatch in Model-based Reinforcement Learningcited by 87
  13. 2022Do Differentiable Simulators Give Better Policy Gradients?cited by 83
  14. 2021Gradients are Not All You Needcited by 82
  15. 2024Adaptive Horizon Actor-Critic for Policy Learning in Contact-Rich Differentiable Simulationcited by 18

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2024: 1; 2025: 3; 2026: 13.

17 papers.

Sort:

By area:

  1. 2024Understanding World or Predicting Future? A Comprehensive Survey of World Modelscited by 161
  2. 2025World-Env: Leveraging World Model as a Virtual Environment for VLA Post-Trainingcited by 57
  3. 2025DMWM: Dual-Mind World Model with Long-Term Imaginationcited by 10
  4. 2026EgoForge: Goal-Directed Egocentric World Simulatorcited by 9
  5. 2026Morphology-Conditioned World Model for Cross-Embodiment Quadrupedal Locomotioncited by 3
  6. 2025Prismatic World Model: Learning Compositional Dynamics for Planning in Hybrid Systemscited by 2
  7. 2026Learning Visual Feature-Based World Models via Residual Latent Actioncited by 2
  8. 2026The SIGReg Objective as Variational Free Energy: A Theoretical Active-Inference Account of JEPA World Modelscited by 2
  9. 2026Toward Safe Autonomous Robotic Endovascular Interventions using World Modelscited by 1
  10. 2026Dream-MPC: Gradient-Based Model Predictive Control with Latent Imaginationcited by 1
  11. 2026HaM-World: Soft-Hamiltonian World Models with Selective Memory for Planningcited by 1
  12. 2026Coupled Local and Global World Models for Efficient First Order RLcited by 0
5 more
  1. 2026Specification-Driven Generation and Evaluation of Discrete-Event World Models via the DEVS Formalismcited by 0
  2. 2026Representation Learning Enables Scalable Multitask Deep Reinforcement Learningcited by 0
  3. 2026Progressive Experience Fusion for Multi-Task World Model Control in Endovascular Navigationcited by 0
  4. 2026ADAPT: Physics-Aware Diffusion-based World Models for Adaptive Predictive Transferable HVAC Controlcited by 0
  5. 2026Predicting Consequences and Reinforcing Navigation Policies with Latent World Modelscited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2025Accelerating Model-Based Reinforcement Learning with State-Space World Models
  2. 2026Scaling World-Model Reinforcement Learning Through Diffusion Policy Optimization
  3. 2026Q-Learning With World Models
  4. 2025World Models as Reference Trajectories for Rapid Motor Adaptation
  5. 2026WAM-RL: World-Action Model Reinforcement Learning with Reconstruction Rewards and Online Video SFT
  6. 2025Safe Planning and Policy Optimization via World Model Learning
  7. 2025Robotic World Model: A Neural Network Simulator for Robust Policy Optimization in Robotics
  8. 2022On the Feasibility of Cross-Task Transfer with Model-Based Reinforcement Learning