# PWM: Policy Learning with Multi-Task World Models

arXiv: https://arxiv.org/abs/2407.02466 · HTML: https://mlsanity.com/paper/2407.02466/

## Key facts

- Cited by 17 papers in the map (data through 14 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.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)
- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 30 of 30)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Segment Anything](https://mlsanity.com/paper/2304.02643/) (2023)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [RT-2: Vision-Language-Action Models Transfer Web Knowledge to Robotic Control](https://mlsanity.com/paper/2307.15818/) (2023)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Octo: An Open-Source Generalist Robot Policy](https://mlsanity.com/paper/2405.12213/) (2024)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Soft Actor-Critic Algorithms and Applications](https://mlsanity.com/paper/1812.05905/) (2018)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [Open X-Embodiment: Robotic Learning Datasets and RT-X Models](https://mlsanity.com/paper/2310.08864/) (2023)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [Deep Reinforcement Learning at the Edge of the Statistical Precipice](https://mlsanity.com/paper/2108.13264/) (2021)
- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [Variational Dropout and the Local Reparameterization Trick](https://mlsanity.com/paper/1506.02557/) (2015)
- [A Brief Survey of Deep Reinforcement Learning](https://mlsanity.com/paper/1708.05866/) (2017)
- [Temporal Difference Learning for Model Predictive Control](https://mlsanity.com/paper/2203.04955/) (2022)
- [dm_control: Software and Tasks for Continuous Control](https://mlsanity.com/paper/2006.12983/) (2020)
- [Monte Carlo Gradient Estimation in Machine Learning](https://mlsanity.com/paper/1906.10652/) (2019)
- [RVT: Robotic View Transformer for 3D Object Manipulation](https://mlsanity.com/paper/2306.14896/) (2023)
- [Foundation Models in Robotics: Applications, Challenges, and the Future](https://mlsanity.com/paper/2312.07843/) (2023)
- [Model-Based Value Estimation for Efficient Model-Free Reinforcement Learning](https://mlsanity.com/paper/1803.00101/) (2018)
- [Gradient Estimation Using Stochastic Computation Graphs](https://mlsanity.com/paper/1506.05254/) (2015)
- [Accelerated Policy Learning with Parallel Differentiable Simulation](https://mlsanity.com/paper/2204.07137/) (2022)
- [Objective Mismatch in Model-based Reinforcement Learning](https://mlsanity.com/paper/2002.04523/) (2020)
- [Do Differentiable Simulators Give Better Policy Gradients?](https://mlsanity.com/paper/2202.00817/) (2022)
- [Gradients are Not All You Need](https://mlsanity.com/paper/2111.05803/) (2021)
- [Adaptive Horizon Actor-Critic for Policy Learning in Contact-Rich Differentiable Simulation](https://mlsanity.com/paper/2405.17784/) (2024)

## Built on by (most cited first, 17 of 17)

- [Understanding World or Predicting Future? A Comprehensive Survey of World Models](https://mlsanity.com/paper/2411.14499/) (2024)
- [World-Env: Leveraging World Model as a Virtual Environment for VLA Post-Training](https://mlsanity.com/paper/2509.24948/) (2025)
- [EgoForge: Goal-Directed Egocentric World Simulator](https://mlsanity.com/paper/2603.20169/) (2026)
- [DMWM: Dual-Mind World Model with Long-Term Imagination](https://mlsanity.com/paper/2502.07591/) (2025)
- [Morphology-Conditioned World Model for Cross-Embodiment Quadrupedal Locomotion](https://mlsanity.com/paper/2604.08780/) (2026)
- [Learning Visual Feature-Based World Models via Residual Latent Action](https://mlsanity.com/paper/2605.07079/) (2026)
- [Prismatic World Model: Learning Compositional Dynamics for Planning in Hybrid Systems](https://mlsanity.com/paper/2512.08411/) (2025)
- [Toward Safe Autonomous Robotic Endovascular Interventions using World Models](https://mlsanity.com/paper/2604.20151/) (2026)
- [HaM-World: Soft-Hamiltonian World Models with Selective Memory for Planning](https://mlsanity.com/paper/2605.05951/) (2026)
- [The SIGReg Objective as Variational Free Energy: A Theoretical Active-Inference Account of JEPA World Models](https://mlsanity.com/paper/2607.13612/) (2026)
- [Coupled Local and Global World Models for Efficient First Order RL](https://mlsanity.com/paper/2602.06219/) (2026)
- [Specification-Driven Generation and Evaluation of Discrete-Event World Models via the DEVS Formalism](https://mlsanity.com/paper/2603.03784/) (2026)
- [Dream-MPC: Gradient-Based Model Predictive Control with Latent Imagination](https://mlsanity.com/paper/2605.04568/) (2026)
- [Predicting Consequences and Reinforcing Navigation Policies with Latent World Models](https://mlsanity.com/paper/2608.26190/) (2026)
- [Representation Learning Enables Scalable Multitask Deep Reinforcement Learning](https://mlsanity.com/paper/2606.05555/) (2026)
- [Progressive Experience Fusion for Multi-Task World Model Control in Endovascular Navigation](https://mlsanity.com/paper/2608.18647/) (2026)
- [ADAPT: Physics-Aware Diffusion-based World Models for Adaptive Predictive Transferable HVAC Control](https://mlsanity.com/paper/2608.19804/) (2026)

## 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/.