Drift, Then Spread: Learning One-Step Generative Policies for Online Reinforcement Learning
summary
The gist
The gist: DRIFTSPREAD proposes a one-step generative policy based on Drifting for online maximum-entropy (MaxEnt) reinforcement learning, achieving superior aggregate normalized return and efficiency
In short
DRIFTSPREAD proposes a one-step generative policy for online reinforcement learning by combining critic-guided Drift and entropy-induced Spread updates derived from a Maximum Entropy objective. This method achieves superior performance over diffusion and flow baselines on MuJoCo, offering higher return and lower training costs while maintaining fast inference.
Key concepts
- Drift Learning
- This update uses the critic's gradients to move the generated actions in the direction that maximizes expected reward. It's like a guided regression where the actor is adjusted based on how good its current actions are according to the critic, pushing it towards better outcomes.
- Spread Learning
- This update captures entropy by training the actor to match a distribution derived from a higher-temperature target. It uses Maximum Mean Discrepancy (MMD) loss to ensure the generated action distribution spreads out appropriately, capturing more diverse and potentially better actions.
- Drift-Spread Cycles
- The method learns by sequentially applying two updates: Drift moves the policy towards high reward, and Spread adjusts the action distribution's spread. These two updates are cycled repeatedly to refine the policy into a single, expressive one-step generative model.
- Maximum Entropy (MaxEnt) Objective
- This is the underlying goal guiding the learning process. It encourages policies that not only maximize reward but also maintain high entropy, meaning they explore a wider range of actions while still performing well. This helps in finding more robust and flexible policies.
Terminology used across episodes
This episode discusses
- Drift, Then Spread: Learning One-Step Generative Policies for Online Reinforcement Learning · Paper Radio
- One-Step Flow Policy Mirror Descent
- Generative Modeling via Drifting
- Mean Flow Policy Optimization
- On the Wasserstein Gradient Flow Interpretation of Drifting Models
- One-Step Generative Modeling via Wasserstein Gradient Flows
- Sinkhorn-Drifting Generative Models
- Drift Q-Learning
- Score-Based One-step MeanFlow Policy Optimization
- Drifting Field Policy: A One-Step Generative Policy via Wasserstein Gradient Flow
- Reverse Flow Matching: A Unified Framework for Online Reinforcement Learning with Diffusion and Flow Policies
- Wasserstein Proximal Policy Gradient
The paper
Drift, Then Spread: Learning One-Step Generative Policies for Online Reinforcement Learning · Read on arXiv
Chunyu Luo, Hongyu Zhao, Guangyi Ni, Yue Song, Lingkai Kong
The University of Hong Kong · Beihang University · Harbin Institute of Technology, Shenzhen 4 · Tsinghua University
Generative policies offer expressive action distributions for online reinforcement learning (RL), but many diffusion- and flow-based policies require costly iterative sampling. We propose DriftSpread, a one-step generative policy based on Drifting for online maximum-entropy (MaxEnt) RL. To incorporate entropy without a tractable policy density, we build on an energy-based view of Drifting and derive critic-guided Drift and entropy-induced Spread from the negative MaxEnt objective. Our method learns these updates sequentially in the same actor, using action regression for value improvement and distribution matching to capture the effect of entropy. The actor updates require neither direct samples from the critic-defined Boltzmann target nor policy-density or score estimation. Under exact fitting with a fixed critic, we prove convergence to the Boltzmann target in the small-step limit followed by the long-time limit. On MuJoCo, DriftSpread achieves the highest aggregate normalized return among all evaluated methods, while outperforming the diffusion- and flow-based baselines in both training and inference efficiency.
Transcript
Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.
Tom: Today's paper: "Drift, Then Spread".
Jane: The gist: DRIFTSPREAD proposes a one-step generative policy based on Drifting for online maximum-entropy (MaxEnt) reinforcement learning,
Tom: First, who's behind it and why it matters.
Title and authors: Tom: So we're looking at this paper today called "Drift, Then Spread: Learning One-Step Generative Policies for Online Reinforcement Learning" and it’s from Chunyu Luo and his team. It tackles the problem of learning expressive action distributions for online reinforcement learning in a single forward pass, which cuts down on the costly iterative sampling you usually have to do with diffusion or flow-based policies.
Jane: It sounds like they're proposing a one-step generative policy based on Drifting for online maximum-entropy reinforcement learning, and it aims to be more efficient than those other methods we’ve seen. What exactly is the core idea behind this approach?
Lu: The core idea is that they learn this one-step generative policy through two sequential updates applied to the same actor: first, critic-guided Drift, and second, entropy-induced Spread. They are building these two updates from an energy-based view of Drifting and deriving them from the negative maximum-entropy objective.
Meng: So they’re not using one big model for everything; they're doing two separate training steps on the same actor to get there. What does that sequential learning actually look like in practice?
Tom: In each cycle, Drift uses critic gradients to move generated actions toward higher values, and Spread captures the entropy contribution through diffusion of the action distribution. It’s a loop where they fit these two different objectives one after the other.
Jane: That means Drift is focused on improving the expected return based on what the critic tells it, while Spread handles making sure that exploration or entropy is still accounted for in that single forward pass generation. How do they make sure these two updates work together to create a stable policy?
Lu: They implement this by defining an empirical Drift loss, L Drift(theta; theta k) = one/two E s,z
f theta(s, z) - e alpha k+one/two(s, z): , and then they learn Spread through distribution matching using a squared maximum mean discrepancy loss: L Spread(theta; theta k+one/two) = E s MMD squared h(pi theta(times s), pi theta k+one/two(times s)) <ref:2610.15662#pg1>.
Meng: MMD, that's maximum mean discrepancy. So they're training the actor to match the distribution of a plus epsilon eta, where a is derived from the current policy and critic estimate. That’s a bit more complex than just fitting one loss function.
Title and authors: Jane: It seems like they are using this sequential Drift-Spread cycle repeatedly to refine the actor, which is what allows them to derive two distinct updates from the MaxEnt objective without needing explicit policy densities or scores. Is that the main advantage over diffusion methods?
Tom: Exactly, because it avoids those costly iterative sampling steps that diffusion and flow-based policies often require during both training and inference. The paper suggests this method learns one-step generators by repeatedly moving generated samples and training the generator to reproduce their updated positions.
Jane: If we think about the theoretical side, what guarantees do they have that this sequential learning process actually converges to the correct target distribution? Is it stable over time?
Lu: Under exact fitting with a fixed critic, they prove convergence to the Boltzmann target distribution in equation (twelve) in the small-step limit followed by the long-time limit. They show that drift and diffusion terms balance at pi*Q, making it a stationary distribution of equation (thirty-six).
Meng: Convergence to a Boltzmann target is solid, but what about deployment? If we were to deploy this on something like MuJoCo, how does the actual training procedure look when you're trying to get it running in production?
Tom: The training procedure combines data collection and critic learning with sequential Drift-Spread actor updates using an off-policy actor-critic algorithm. The critic update minimizes the squared error to the detached target y = r + gamma(one - delta) i,two Q(s', a') on page seven.
Jane: So, for a practical deployment scenario, it involves collecting data, training a critic to estimate that value function target based on the Bellman equation minus some discount factor delta, and then iteratively updating the actor with the Drift and Spread losses. What does that mean for inference speed?
Tom: They report that this approach reduces average inference time to zero point one five ms per action, which is significantly faster than the one point two four ms they saw for DIME on page eight. That's a real win for real-time applications.
Title and authors: Jane: It’s impressive that it achieves that speed while also showing better performance than other methods when compared on MuJoCo, which is often a benchmark for these kinds of control tasks. What about the trade-offs? Are there any caveats we should keep in mind about this approach?
Lu: The choice of spreading noise definitely affects policy learning; they found that Gaussian noise achieves the highest or near-highest late-training mean returns across all three tasks. Also, the frequency distribution you use for that spreading noise determines the MMD kernel they use to compare the target with the actor on page nine.
Meng: So, if we want a reliable baseline for our own systems, sticking with Gaussian noise seems like a safe bet based on these results. What about training costs? Is it cheaper than something that requires more samples?
Tom: Compared to diffusion and flow baselines, DRIFTSPREAD shows lower training costs on MuJoCo page eight. It’s also worth mentioning that the sequential fitting achieves higher mean returns than joint fitting on Humanoid-v3, which is another complex task.
Jane: So, to summarize this paper "Drift, Then Spread: Learning One-Step Generative Policies for Online Reinforcement Learning," they propose a way to learn expressive action distributions in one step by sequentially applying critic-guided Drift and entropy-induced Spread from the negative MaxEnt objective. It converges under exact fitting to the Boltzmann target distribution, and on MuJoCo, it beats thirteen baselines while being faster than diffusion methods.
Lu: The method incorporates the effect of entropy without explicitly computing policy densities or scores, allowing each action sample to be generated in one forward pass on page seven.
Meng: For us in production, the key is that it's one forward pass generation at inference time, which keeps things lean and fast.
Tom: It’s an interesting mechanism because it manages to incorporate entropy without needing a tractable policy density or score estimation for the model to work.
Jane: This paper gives us a solid framework for learning policies that balance return and exploration in online reinforcement learning in a single step, which is definitely something worth keeping on our radar as we look at next-generation generative models.
The paper's summary: Tom: So, to recap, we’re looking at this paper by Luo and his team called "Drift, Then Spread," which proposes learning one-step generative policies for online reinforcement learning using a specific two-stage update process called critic-guided Drift and entropy-induced Spread.
Jane: Right. The main idea is that they can generate an action in just one forward pass, which cuts down on all that slow iterative sampling you usually have to do with diffusion or flow methods.
Lu: Exactly. They derive these two updates from the negative maximum-entropy objective, so they are fitting both the return and the spread aspect of the policy simultaneously but sequentially.
Meng: It’s interesting because it avoids computing explicit policy densities or scores, which is usually a big part of those other generative models.
Tom: That’s what excites me—if you can get high-quality action distributions in one go, that drastically lowers the computational overhead for online tasks.
Jane: But how do they prove it actually works? What are the theoretical guarantees behind this sequential fitting process?
Lu: Under exact fitting with a fixed critic, they prove convergence to the Boltzmann target distribution in equation (twelve) in the small-step limit followed by the long-time limit.
Tom: Convergence to that target distribution is key—it shows that after enough time, you actually get the right kind of policy behavior.
Jane: They show that drift and diffusion terms balance out at pi*Q, which means it settles into a stable state described by equation (thirty-six).
Meng: That stability is important for deployment because it suggests the learning process isn't going to keep oscillating forever.
Tom: And on the practical side, they show that this method outperforms thirteen other baselines on MuJoCo in aggregate normalized return.
Jane: It’s not just a win on one benchmark; they also claim lower training costs compared to diffusion, flow, and MeanFlow baselines.
Lu: That’s because they manage to incorporate the entropy effect without needing those heavy density calculations, which usually drive up the cost for those other methods.
Tom: And for inference speed, they got a very low average inference time of zero point one five milliseconds per action, compared to twelve point four four milliseconds for DIME. That’s a big gap when you’re dealing with real-time control.
Jane: So what does this mean for someone who is just listening to the show? It means we might be able to move toward more efficient generative models that can actually run online without needing massive iterative sampling loops during operation.
Lu: The main implication is that you get expressive one-step policies by cleverly combining critic guidance and distribution matching sequentially.
Tom: Exactly, it’s a Drifting-based approach to learning those policies that seems really solid because it handles the return and the spread aspect of the objective in a structured way.
Jane: So we’re looking at a method where you derive two updates from energy and fit them through repeated Drift-Spread cycles, achieving efficiency without sacrificing return on tasks like MuJoCo.
Meng: I still have to watch how they handle scaling that MMD loss as the state space gets bigger, because if that MMD becomes too slow, the one-step generation advantage might vanish.
Lu: That’s a fair concern, Meng. The paper does discuss the choice of spreading noise—Gaussian noise seems to give them their best late-training mean returns across all three tasks they tested.
Tom: It sounds like Gaussian noise is a good default setting for this method, which simplifies things for those trying to implement it quickly.
Jane: So, while the theory shows convergence and the numbers on MuJoCo are strong, we still need to see how robust this holds up when you take it out of simulation and onto a real-world control task.
The paper's improvements: Tom: We’re now looking at what the authors suggest as improvements to their Drift, Then Spread method for learning those one-step policies.
Jane: So, beyond just doing a two-step update cycle, what specific tweaks are they proposing to make the policy learning even better?
Lu: They focus on making sure the spreading noise isn't just a fixed Gaussian thing; they look at how the frequency distribution of that noise affects what happens.
Tom: So, it’s not just using any random noise to capture entropy; how you choose that distribution matters for the final performance.
Jane: They found that using Gaussian noise achieves the highest or near-highest late-training mean returns across all three tasks they tested, which is a pretty solid default setting.
Meng: That’s good for deployment because it means we don't have to spend extra time fine-tuning the noise parameters if we just pick Gaussian.
Tom: The paper suggests that the frequency distribution you use for that spreading noise is what actually determines the MMD kernel they use to compare the target with the actor.
Jane: So, if you change your noise distribution, it changes how they measure how well their policy matches the target distribution during training.
Lu: That’s because that choice of frequency distribution directly impacts the MMD loss function, which is what drives that second update step.
Tom: So they are suggesting a degree of flexibility there—you can tune the noise to get better results, but Gaussian is their current best bet for simplicity.
Jane: It shows they aren't just sticking to one fixed setup; they’re allowing the noise choice to be part of the learning process, which is an improvement over just picking one thing and never looking back.
Meng: From a production standpoint, flexibility is helpful because it means we can iterate on deployment settings without having to rebuild the whole training pipeline every time.
Tom: It’s about making the system more adaptable rather than just finding a single perfect configuration for one task.
Lu: The paper also points toward future work where they plan to test this method on real-world control tasks, moving beyond MuJoCo simulations.
Jane: That makes sense, because showing it works in a controlled environment is one thing, but showing it handles the messiness of the actual physical world is where the real test is.
Tom: So they are looking ahead to taking this Drifting framework and applying it where things aren't perfectly smooth or defined by a simple MDP.
Jane: It’s interesting because online reinforcement learning in the real world isn't as clean as simulations, so testing it there will really show its practical power.
Lu: They are also looking at how to adapt that spreading strength dynamically as the learning process goes on, which would mean the noise level changes during training.
Tom: That sounds like a sophisticated control mechanism—adjusting the exploration factor based on how well things are fitting the target distribution in real-time.
Jane: So, they’re not just doing one set of updates; they’re suggesting a way to dynamically adjust the spreading component as the model learns.
Conclusion: Tom: So we’ve covered how DRIFTSPREAD uses critic-guided Drift and entropy-induced Spread to create one-step generative policies for online reinforcement learning and its strong performance on MuJoCo.
Jane: Exactly, and we saw that it achieves high aggregate normalized return while being faster than diffusion methods during inference.
Lu: The main idea is deriving these two updates from the negative maximum-entropy objective, which allows for a single forward pass action generation.
Meng: It’s good that it’s efficient at scale, but I still want to see how stable this training process is when we move past simple MuJoCo benchmarks and onto more complex control problems.
Lalam: I'm confident in the paper's answer; the convergence proof under exact fitting shows that it settles into the Boltzmann target distribution as time goes to infinity.
Tom: That convergence guarantee is important because it tells us that theoretically, if we give it enough time, we should get what we want.
Jane: And for people listening, it means an AI agent could generate a move in one step without having to run a long simulation just to decide what the best move is.
Lu: The paper also points toward future work where they plan to test this method on real-world control tasks, which is the next logical step.
Tom: I’m curious about that because applying these math concepts to actual physical systems is always where the real engineering happens.
Meng: From my side, I wonder what happens if we have to introduce some kind of dynamic adjustment for the spreading strength during training, like they hinted at in their improvements section.
Jane: It means the model could adjust how much it explores based on how confident it is in its current estimate of the return.
Lalam: I’m sure that dynamic adjustment will help keep the policy from getting stuck in local optima during long training runs.
Tom: So, to wrap up, this paper, "Drift, Then Spread: Learning One-Step Generative Policies for Online Reinforcement Learning," gives us a concrete method for generating expressive actions quickly by balancing drift and spread updates.
Jane: It’s a solid foundation for future work on faster online learning systems that can incorporate entropy naturally without the heavy computational cost of density estimation.
Lu: I think the approach is ambitious because it tries to capture both the return and exploration aspects of policy learning simultaneously in this sequential way.
Meng: It's definitely worth keeping an eye on, especially for reducing inference latency, but we’ll need to see how well it handles real-world noise and uncertainty when we test it.
Lalam: I'm sure this paper is a significant step forward in how AI learns to generate actions directly from the reward structure.
More episodes
- 2610.15564-UNISON: A Unified Framework for Learning from Scored Bags
- 2610.15659-Testing Central Structure for Complex Networks
- 2610.15696-DRanker: A Transformer-Based Multi-Task Ranking Model for the Disney+ Homepage
- 2610.15623-Adaptive Role Assignment and Controlled Feature Reconstruction for Federated Image Semantic Communication
- 2610.15830-How Shared Representations Become Specialized: Formation, Maintenance, and Latent Susceptibility
- 2610.15640-Context-Aware Precision Scheduling for Small Language Models
- 2610.15808-Grounding Deep Learning Framework Testing in Data Flow
- 2610.15818-Benchmarking Deep Learning Models for Anomaly Detection in Real-World Cloud Telemetry
- 2610.15899-Lightweight Scheduling Policies for Continuously Batched LLM Serving Across GPU Architectures
- 2610.15871-Entropy-Regularized Mean-Field Games with Average-Reward Criterion: Span-Seminorm Contractions and Model-Free Learning