Research map / Reinforcement learning / Offline RL
ATraDiff: Accelerating Online Reinforcement Learning with Imaginary Trajectories
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 ATraDiff, a diffusion model trained on offline data that synthesizes full state-conditioned trajectories (as image sequences) to augment the replay buffer and accelerate online, offline-to-online and offline RL.
Problem
Sample inefficiency of online RL with sparse rewards; prior uses of offline data are fixed and limited, and diffusion-based augmentation (SynthER) generates only transitions and cannot handle varying trajectory lengths or offline-to-online distribution shift.
Approach
Generate complete trajectories conditioned on the current state and task with a set of diffusion models of different fixed lengths (coarse-to-precise length selection plus non-parametric pruning), operating on image sequences via a fine-tuned Stable Diffusion, and adapt the diffuser online using importance-weighted replay samples while keeping a frozen copy against forgetting.
Key results
- On D4RL Locomotion offline RL, IQL+ATraDiff reaches e.g. 53.1 on halfcheetah-medium vs 48.3 for IQL and 49.6 for IQL+SynthER.
- ATraDiff improves SAC and REDQ online learning across three Locomotion environments over 250K steps, with largest gains on halfcheetah and walker2d.
- On pixel-based V-D4RL, BC+ATraDiff improves cheetah-expert from 67.4 to 85.3 and walker-medium from 40.9 to 42.2 (normalized return).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| augment offline RL replay data with ATraDiff trajectories | TD3+BC and IQL · D4RL Locomotion offline datasets | normalized return | Mixed |
| replace transition generation with trajectory generation (ablation) | RLPD augmentation · Locomotion | normalized return: 93.5 -> 97.1 | Helped |
| condition generated trajectories on the current state (ablation) | RLPD augmentation · Locomotion | normalized return: 97.1 -> 101.3 | Helped |
| add importance sampling during online adaptation (ablation) | RLPD+ATraDiff · Locomotion | normalized return: 101.3 -> 102.3 | Helped |
| replace transition generation with trajectory generation (ablation) | RLPD augmentation · AntMaze | normalized return: 97.2 -> 97.6 | Helped |
| condition generated trajectories on the current state (ablation) | RLPD augmentation · AntMaze | normalized return: 97.6 -> 98.2 | Helped |
| add importance sampling during online adaptation (ablation) | RLPD+ATraDiff · AntMaze | normalized return: 98.2 -> 98.7 | Helped |
| replace transition generation with trajectory generation (ablation) | RLPD augmentation · Meta-World Multitask | return: 0.60 -> 0.72 | Helped |
| condition generated trajectories on the current state (ablation) | RLPD augmentation · Meta-World Multitask | return: 0.72 -> 0.75 | Helped |
| add importance sampling during online adaptation (ablation) | RLPD+ATraDiff · Meta-World Multitask | return: 0.75 -> 0.77 | Helped |
| replace transition generation with trajectory generation (ablation) | RLPD augmentation · Meta-World Sequencetask | return: 0.64 -> 0.71 | Helped |
| condition generated trajectories on the current state (ablation) | RLPD augmentation · Meta-World Sequencetask | return: 0.71 -> 0.79 | Helped |
Surprising findings
- Training multiple diffusion models of varying lengths is computationally expensive, with ATraDiff roughly doubling training time for SAC and REDQ (e.g., 6.5h to 15h).
- Image-level generation outperforms state-level generation in complicated tasks, which the authors attribute to pretrained Stable Diffusion.
Limitations
- Training multiple diffusion models of different lengths is computationally intensive, and cost grows with longer generated trajectories.
- Online adaptation is needed for complex tasks; fixed diffuser suffers from distribution shift.
- Reported improvements on several tasks are modest (e.g., some offline results within noise of baselines).
Experimental setup
Datasets: D4RL Locomotion D4RL AntMaze D4RL Kitchen Meta-World V-D4RL
Benchmarks: D4RL Locomotion D4RL AntMaze D4RL Kitchen Meta-World V-D4RL
Metrics: normalized return return learning curves (median over 5 seeds, 25-75% percentiles)
Baselines: SAC REDQ SynthER S4RL SPiRL PARROT RLPD TD3+BC IQL BC
Builds on: SynthER Diffuser Stable Diffusion SAC REDQ
Closest papers by meaning
- Planning, Fast and Slow: Online Reinforcement Learning with Action-Free Offline Data via Multiscale Planners (ICML 2024 · similarity 0.89)
- A Snapshot of Influence: A Local Data Attribution Framework for Online Reinforcement Learning (NeurIPS 2025 · similarity 0.88)
- MADiff: Offline Multi-agent Learning with Diffusion Models (NeurIPS 2024 · similarity 0.88)
- Offline-to-Online Reinforcement Learning with Classifier-Free Diffusion Generation (ICML 2025 · similarity 0.88)
- DiffStitch: Boosting Offline Reinforcement Learning with Diffusion-based Trajectory Stitching (ICML 2024 · similarity 0.88)
- Learning from Reward-Free Offline Data: A Case for Planning with Latent Dynamics Models (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.