Research map / Reinforcement learning / Policy gradient and actor-critic
Do Transformer World Models Give Better Policy Gradients?
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.
Shows that transformer world models conditioned on full state-action history yield circuitous gradient paths, and proposes Actions World Models (conditioned only on initial state and actions) that give better long-horizon policy gradients than simulators, Markovian/RNN models, SAC and Online-DT.
Problem
Backpropagation through unrolled learned world models for policy optimization becomes unstable over long horizons because of circuitous gradient paths created by autoregressive state prediction, even when transformers are used.
Approach
Condition the world model only on the initial state and the sequence of actions (an Actions World Model), so gradients flow from rewards to actions directly through the network rather than through autoregressively predicted states; theoretically, gradient norm bounds are inherited from the underlying architecture (RNN exponential, self-attention polynomial O(H^3)).
Key results
- On the chaotic double-pendulum task, transformer AWMs outperform all BPO baselines, including BPO through the true simulator, which fails to produce successful policies.
- On eight Myriad tasks with horizons up to 500, BPO with transformer AWMs is the only world model that performs reliably well as horizon grows beyond 100 steps.
- BPO with transformer AWMs is more sample efficient than SAC and Online-DT on Myriad, with superior performance at long horizons.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use a transformer AWM instead of a Markovian model or HWM | BPO · Double-pendulum | final return | Helped |
| Condition the transformer world model on actions rather than states | BPO with a transformer world model · Myriad | final performance | Helped |
| Use a transformer AWM instead of SAC or Online-DT | BPO · Myriad | sample efficiency | Helped |
| Remove stop-gradient on the policy input (ablation) | BPO with transformer AWM or transformer HWM · Myriad | final performance | No clear effect |
| Use an LSTM backbone for the AWM instead of the transformer backbone (ablation) | BPO with an AWM · Myriad | performance | Helped |
Surprising findings
- Conditioning transformer world models on full state history offers no policy-gradient benefit over Markovian models, despite better prediction.
- Backpropagation through the true simulator can fail on chaotic double-pendulum dynamics where a learned AWM succeeds.
Limitations
- Theoretical analysis is developed for deterministic MDPs.
- Experiments are on low-dimensional continuous domains (Myriad, double pendulum), not large-scale real-world tasks.
- Only worst-case bounds are given; scaling to larger systems is left to future work.
Experimental setup
Datasets: Myriad environments (cancer treatment, bioreactor, mould fungicide, bacteria, harvest, invasive plant, HIV treatment, ti double-pendulum environment data from random trajectories
Benchmarks: Myriad double-pendulum
Metrics: normalized return IQM gradient norm
Baselines: Markovian world model BPO history world model BPO RNN AWM Soft Actor-Critic Online Decision Transformer BPO through simulator
Builds on: Soft Actor-Critic Decision Transformer GPT-2 teacher forcing
Closest papers by meaning
- PWM: Policy Learning with Multi-Task World Models (ICLR 2025 · similarity 0.89)
- Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvement (NeurIPS 2026 · similarity 0.88)
- WPT: World-to-Policy Transfer via Online World Model Distillation (CVPR 2026 · similarity 0.86)
- World Action Verifier: Self-Improving World Models via Forward-Inverse Asymmetry (NeurIPS 2026 · similarity 0.86)
- Revisiting Multi-Agent World Modeling from a Diffusion-Inspired Perspective (NeurIPS 2025 · similarity 0.86)
- ACT: Empowering Decision Transformer with Dynamic Programming via Advantage Conditioning (AAAI 2024 · similarity 0.86)
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.