Research map / Reinforcement learning / Policy gradient and actor-critic
PWM: Policy Learning with Multi-Task World Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes PWM, which pre-trains a multi-task world model offline and extracts per-task policies via first-order gradient optimization through the smooth, well-regularized model, achieving higher reward than TD-MPC2 without online planning.
Problem
Multi-task reinforcement learning across different embodiments is hard; world-model methods rely on inefficient zeroth-order optimization or online planning, while first-order gradients struggle with contact discontinuities and chaotic dynamics.
Approach
Well-regularized (smooth) learned world models produce smoother optimization landscapes than the true discontinuous dynamics, so first-order gradient actor-critic training through them is efficient; the framework trains one multi-task world model offline and then extracts a separate policy per task in minutes.
Key results
- PWM achieves 9672 ± 2012 reward on Ant vs 7662 for SHAC and 7080 for TD-MPC2 in single-task dflex experiments.
- In the MT80 setting PWM achieves up to 27% higher reward than TD-MPC2 without online planning.
- Multi-task world model PWM policies are trained in 9.3 minutes per task (10k gradient steps on an RTX6000 GPU).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use a regularized (SimNorm, Mish) world model for FoG policy learning instead of ReLU model (ablation) | PWM on Ant · Ant world model regularization ablation | policy reward | Helped |
| Increase contact stiffness on Hopper (ablation) | SHAC (ground-truth differentiable sim) · Hopper contact stiffness ablation | normalized reward: SHAC performance decreases by 48% | Hurt |
| Increase contact stiffness on Hopper (ablation) | PWM · Hopper contact stiffness ablation | normalized reward | No clear effect |
| Use larger policy batch size for FoG policy training (ablation) | PWM · Ant batch size ablation | reward within fixed wall-clock time | Hurt |
| Pre-train world model with H=16 instead of H=3 (policy trained at H=16) (ablation) | PWM multi-task · MT30/MT80 horizon ablation | score: marginally higher scores | Helped |
| Replace PWM actor FoG objective with TD(lambda) actor objective (ablation) | PWM · Hopper, Ant, Anymal TD(lambda) ablation | reward | Hurt |
| Train a single multi-task policy instead of per-task policies extracted from the multi-task world model (ablation) | PWM · MT30 | reward | Hurt |
| Increase world model pre-training gradient steps (50k to 250k) for PWM policy (ablation) | PWM policy components · dm_control 5-task sample efficiency ablation | reward | Helped |
| Use PWM FoG policy (on world model) instead of TD-MPC2 DDPG-style policy with planning disabled (ablation) | World model policy extraction · dm_control sample efficiency | reward | Helped |
| Train PWM with FoG policy on pre-trained world model instead of online planning | PWM vs TD-MPC2 on MT30/MT80 · MT30 and MT80 multi-task | reward: up to 27% higher reward than TD-MPC2 | Helped |
| Use FoG through learned world model instead of ground-truth differentiable simulator | PWM vs SHAC · Ant single-task | PPO-normalised reward: 1.46 vs 1.16 | Helped |
| Use FoG through learned world model instead of ground-truth simulator | PWM vs SHAC · SNU Humanoid (152 actions) | PPO-normalised reward: 1.36 vs 1.44 | Hurt |
Surprising findings
- More accurate world models do not yield better policies; less regularized (lower-error) models gave suboptimal final policies, indicating an inverse correlation between model accuracy and policy reward.
- A single multi-task PWM policy and TD-MPC2 without planning both fail on MT30, so per-task policy extraction is needed.
Limitations
- Performance depends heavily on availability of substantial offline data to train the world model.
- Requires re-training of the policy for each new task, limiting rapid adaptation.
- TD-MPC2 world models are difficult to train at scale due to their auto-regressive formulation; PWM does not scale well to the 152-action SNU Humanoid.
Experimental setup
Datasets: dm_control MetaWorld TD-MPC2 multi-task offline data SHAC-collected trajectories for dflex tasks
Benchmarks: dflex Hopper dflex Ant dflex Anymal dflex Humanoid dflex SNU Humanoid MT30 MT80
Metrics: 50% IQM reward mean reward 95% CI normalized reward inference time
Baselines: TD-MPC2 SHAC DreamerV3 PPO SAC
Builds on: TD-MPC2 SHAC DreamerV3 TD(lambda)
Closest papers by meaning
- Overcoming Slow Decision Frequencies in Continuous Control: Model-Based Sequence Reinforcement Learning for Model-Free Control (ICLR 2025 · similarity 0.90)
- World Models as Reference Trajectories for Rapid Motor Adaptation (NeurIPS 2025 · similarity 0.89)
- Do Transformer World Models Give Better Policy Gradients? (ICML 2024 · similarity 0.89)
- Towards General-Purpose Model-Free Reinforcement Learning (ICLR 2025 · similarity 0.89)
- AdaWM: Adaptive World Model based Planning for Autonomous Driving (ICLR 2025 · similarity 0.88)
- RLVR-World: Training World Models with Reinforcement Learning (NeurIPS 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.