Research map / Reinforcement learning / MDPs and planning
Diffusion Spectral Representation for Reinforcement Learning
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 Diff-SR, which uses diffusion-model score learning on an energy-based transition model to learn spectral representations for Q-functions in MDPs/POMDPs, avoiding diffusion sampling at inference.
Problem
Diffusion models are expressive for RL but sampling is slow, and it is unclear how to use them for planning and exploration with value-function representations in MDPs and POMDPs.
Approach
Connects diffusion models to energy-based models to get a random-Fourier-feature spectral factorization of the transition operator; the diffusion denoising loss (via Tweedie's identity) learns the representation directly, so Q-functions are linear in the learned features and no sampling from the diffusion model is needed.
Key results
- On MBBL state-based MDP tasks after 200K steps, Diff-SR reaches 7223.5 on HalfCheetah versus 5557.6 for LV-Rep and is best on 9 of 10 tasks among representation baselines shown.
- Diff-SR is reported to be about 4x faster in wall-clock time than PolyGRAD, which samples from a diffusion model.
- On visual Meta-World tasks, Diff-SR achieves over 90% success on seven tasks, four more than the second-best baseline muLV-Rep.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Diff-SR instead of LV-Rep | SAC-based representation RL · Ant-ET | return: 2511. 8 ± 460.0 -> 4788.6 ± 623.1 | Helped |
| Use Diff-SR instead of LV-Rep | SAC-based representation RL · Walker-ET | return: 2523.5 ± 333.9 -> 3722.2 ± 406.1 | Helped |
| Use Diff-SR instead of the strongest baseline | SAC-based representation RL · Humanoid-ET | return: 1794.4 ± 458.3 -> 821.3 ± 118.9 | Hurt |
| Use Diff-SR instead of the strongest baseline | SAC-based representation RL · Reacher | return: -4.1 ± 0.1 -> -6.5 ± 2.0 | Hurt |
| Use Diff-SR instead of the strongest baseline | SAC-based representation RL · Pendulum | return: 177.3 ± 1.9 -> 162.3 ± 3.0 | Hurt |
| Use Diff-SR instead of the strongest baseline | SAC-based representation RL · S-Humanoid-ET | return: 1110.1 ± 181.8 -> 1160.3 ± 150.1 | Helped |
| Use Diff-SR instead of the strongest baseline | SAC-based representation RL · CartPole | return: 200.7 ± 0.2 -> 199.9 ± 1.1 | No clear effect |
| Use Diff-SR instead of PolyGRAD | Diffusion-based RL · Gym-MuJoCo tasks | runtime speed: 4 × faster | Helped |
| Use Diff-SR instead of μLV-Rep | Partially observable continuous-control RL · Walker (partial observation) | return: 1298.1± 276.3 -> 1860.5 ± 912.1 | Helped |
| Use Diff-SR instead of μLV-Rep | Partially observable continuous-control RL · Ant (partial observation) | return: 1621.4 ± 472.3 -> 3189.7 ± 720.7 | Helped |
| Use Diff-SR instead of μLV-Rep | Partially observable continuous-control RL · Humanoid (partial observation) | return: 806.7 ± 120.7 -> 650.1 ± 57.4 | Hurt |
| Use Diff-SR instead of Dreamer-v2 | Partially observable continuous-control RL · Pendulum (partial observation) | return: 172.3 ± 8.0 -> 167.4 ± 4.4 | Hurt |
Surprising findings
- The vanilla latent-variable reduction of diffusion models to spectral representations is shown to be impossible because the decoder depends on (s,a).
Limitations
- Not evaluated on real-world or multi-task data.
- Image-based experiments rely on a VAE latent space rather than raw pixels.
Experimental setup
Datasets: MuJoCo locomotion (Gym) Meta-World
Benchmarks: Gym-MuJoCo (MBBL tasks) Meta-World visual tasks partially observable MuJoCo variants
Metrics: episode return success rate wall-clock runtime
Baselines: ME-TRPO PETS MBBL best PolyGRAD PPO TRPO SAC DeepSF SPEDE LV-Rep Dreamer-v2 SLAC
Builds on: Latent variable representation (LV-Rep) Soft Actor-Critic DrQ-V2 PolyGRAD Random Fourier features
Closest papers by meaning
- Designing Reinforcement Learning for Diffusion Models: A Unified Path-Space View (NeurIPS 2026 · similarity 0.91)
- Score Regularized Policy Optimization through Diffusion Behavior (ICLR 2024 · similarity 0.89)
- Adding Conditional Control to Diffusion Models with Reinforcement Learning (ICLR 2025 · similarity 0.89)
- Online Laplacian-Based Representation Learning in Reinforcement Learning (ICML 2025 · similarity 0.89)
- Prior-Guided Diffusion Planning for Offline Reinforcement Learning (NeurIPS 2025 · similarity 0.89)
- Diffusion-based Reinforcement Learning via Q-weighted Variational Policy Optimization (NeurIPS 2024 · similarity 0.89)
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.