PaperlayerView exact revision on arXiv

Source-linked AI summary

Training Diffusion Policies via Prior-Mapping Co-Evolution

Chubin Zhang, Zhenglin Wan, Feng Chen, Fuchao Yang, Lang Feng, Yaxin Zhou, Xingrui Yu, Yang You, Ivor Tsang, Bo An

arXiv:2512.02581v3cs.LG

TL;DR

Stable unimodal policies can struggle to represent multimodal actions, while expressive generative policies are difficult to optimize online. GoRL separates latent-policy optimization from generative action synthesis and alternates between improving the latent policy and refining the decoder. It outperforms unimodal and generative baselines across continuous-control tasks, with particularly large reported gains on HopperStand and humanoid tasks.

  • Problem

    Stable unimodal policies often lack the expressiveness to represent multimodal actions, while generative policies are difficult to optimize in online RL.

    Source: §1 · pp. 1–2

  • Method

    GoRL optimizes a tractable latent policy separately from a conditional generative decoder, refining the decoder with a fixed prior between latent-policy updates.

    Source: §1 · pp. 2–3, §3 · p. 4

  • Results

    GoRL improves stability, final performance, and mode coverage over unimodal and coupled generative baselines; its HopperStand returns exceed 870, more than 3× the strongest baseline.

    Source: §1 · p. 3, §6 · p. 12

  • Takeaways & Limitations

    Across the reported continuous-control tasks, decoupling optimization from generation combines stable learning with increasingly flexible action modeling.

    Source: §6 · p. 12

    Periodic decoder refinement adds wall-clock cost, and a priori stage lengths may weaken refinement signals or waste interaction budget after returns plateau.

    Source: §6 · p. 12

Abstract

from arXiv · show

Reinforcement learning (RL) faces a persistent tension: policies that are stable to optimize (e.g., Gaussians) are often too simple to represent the multimodal action distributions required for complex control. Conversely, expressive generative policies -- such as diffusion and flow matching -- can be difficult to optimize in online RL due to intractable likelihoods and gradients propagating through long sampling chains. We address this tension with a key structural principle: decoupling optimization from generation. Building on this, we introduce GoRL (Generative Online Reinforcement Learning), an algorithm-agnostic framework that trains expressive policies from scratch by confining policy optimization to a tractable latent space while delegating action synthesis to a conditional generative decoder. Viewed as prior-mapping co-evolution, each stage first improves a tractable latent prior through RL and then consolidates the resulting behavior into a more expressive prior-to-action mapping. This two-timescale schedule, anchored by fixed-prior decoder refinement, enables stable optimization while continuously expanding expressiveness. Empirically, \textsc{GoRL} consistently outperforms unimodal and generative baselines across diverse continuous-control tasks. Notably, GoRL achieves returns exceeding 870 on HopperStand, more than 3* the strongest baseline; on high-dimensional humanoid tasks, it further outperforms the strongest non-GoRL baseline by over an order of magnitude.

1. Introduction

GoRL addresses the tension between stable but limited unimodal policies and expressive generative policies that are difficult to optimize online. It separates latent-policy optimization from action generation and reports gains over unimodal and generative baselines.

  • Stable unimodal policies can miss multimodal high-reward actions, while online distribution shift and long generative sampling chains make expressive policies harder to optimize.Source: §1 · pp. 1–2
  • GoRL separates a tractable latent policy optimized with standard RL from a conditional decoder that synthesizes expressive actions.Source: §1 · p. 2
  • Fixed-Gaussian-prior decoder refinement avoids near-self-reconstruction and consolidates the optimized latent policy’s exploration into a stronger generator.The alternating schedule improves the latent policy with the decoder fixed, then refines the decoder with the latent policy frozen.Source: §1 · pp. 2–3
  • GoRL variants exceed 870 on HopperStand, more than 3× the strongest non-GoRL baseline, and improve over that baseline by more than an order of magnitude on two humanoid tasks.Source: §1 · p. 3

2. Background and Preliminaries

The background contrasts tractable but often unimodal policies with expressive generative policies, whose likelihoods and long sampling chains complicate online optimization. These challenges motivate a decoupled approach.

2. Background and Preliminaries

  • The section frames the contrast between likelihood-based policy optimization and expressive generative policies as motivation for GoRL’s decoupled design.Source: §2 · p. 3

2.1. Likelihood-Based Policy Optimization

  • Common continuous-control methods use tractable policy densities, often diagonal Gaussians, whose unimodal form cannot represent many multimodal action distributions.Source: §2.1 · p. 3

2.2. Generative Models as Expressive Policies

  • Diffusion policies iteratively denoise Gaussian noise, while flow-matching policies integrate learned transport dynamics to generate actions.Both are presented as conditional generative policies capable of representing rich action distributions.Source: §2.2 · pp. 3–4

2.3. Why Generative Policies Are Hard to Optimize Online

  • Generative policies can have expensive or intractable likelihoods, and backpropagation through long sampling chains can amplify critic errors into brittle learning or collapse.Source: §2.3 · p. 4

3. Methodology

GoRL confines policy optimization to a tractable latent space and delegates action synthesis to a conditional generative decoder. Its alternating updates and theoretical results support this separation while allowing the decoder to evolve.

3. Methodology

  • GoRL optimizes the encoder in latent space while the decoder is fixed, then refines the decoder with supervised generative training while the encoder is frozen.Source: §3 · p. 4, Figure 2 · p. 5

3.1. Latent–Generative Factorization

  • The latent–generative factorization treats the encoder as a tractable state-conditioned latent policy and the decoder as a transport from latents to actions.The decoder may define actions implicitly without a closed-form likelihood.Source: §3.1 · p. 5

3.2. Two-Timescale Alternating Optimization

  • Alternating optimization resets the encoder between stages and trains the decoder on rollout actions using fresh Gaussian latent inputs, preventing refinement from becoming near-self-reconstruction.The initial decoder is approximately identity-like; refinement consolidates the improved behavior into a stronger prior-to-action map.Source: §3.2 · pp. 5–6

3.3. Instantiation

  • GoRL is modular: it can pair on- or off-policy latent optimizers with conditional generative decoders, with experiments instantiating PPO and SAC alongside diffusion or flow matching.Source: §3.3 · pp. 6–7

3.4. Latent Optimization Guarantees

  • With a fixed deterministic decoder, latent updates give unbiased gradients for the composite policy, and bounded latent divergence yields a return-change lower bound.The bound implies improvement when expected advantage dominates the divergence penalty; otherwise it quantifies worst-case degradation.Source: §3.4 · p. 7

4. Experiments

The experiments evaluate GoRL’s performance, stability, mechanisms, and action-distribution expressiveness across continuous-control benchmarks. GoRL outperforms the baselines, with particularly large gains on HopperStand and high-dimensional humanoid tasks.

4.2. Main Results: On-Policy Performance

  • Above 870 on HopperStand—more than 3× the strongest non-GoRL baseline—GoRL variants also achieve higher final returns and more stable learning across six DMControl tasks.GoRL also often establishes a lead earlier on FishSwim and FingerSpin.Source: Figure 3 · p. 8, §4.2 · p. 9

4.1.1. Environments

  • The benchmark uses six DMControl tasks, 180M online steps per task, and Gaussian PPO, FPO, and DPPO baselines with comparable architectures and interaction budgets.GoRL uses PPO with either a flow-matching or diffusion decoder; training is divided into four stages, beginning with a fixed approximate-identity decoder and using fixed-prior refinement at later boundaries.Source: §4.1.1 · p. 8, §4.1.2 · p. 8, §4.1.3 · p. 8

4.2. Main Results: On-Policy Performance

  • FPO shows mid-to-late training drops on several tasks, while GoRL’s latent-space optimization allows direct PPO entropy regularization to help sustain exploration.The authors attribute FPO’s instability to surrogate-ratio misalignment under distribution shift and the lack of an easily controlled explicit entropy regularizer.Source: §4.2 · p. 9
  • GoRL also works with off-policy SAC on standard OpenAI Gym benchmarks, suggesting the framework is not specific to PPO.Source: §4.2 · p. 10
  • On both 21-dimensional humanoid tasks, GoRL variants outperform the strongest non-GoRL baseline by more than an order of magnitude.Stage-wise GoRL-Diff returns rise sharply after the first decoder refinement and continue improving through later stages.Source: Table 2 · p. 10, §4.2 · p. 10

4.3.1. Fixed-Prior vs. Evolving Latents

  • On FingerSpin, refining the decoder on evolving latents causes repeated performance drops, whereas fixed-prior anchoring avoids that feedback loop and consolidates exploration into an improved generator.Source: §4.3.1 · p. 10

4.3.2. Necessity of Stage-Wise Re-initialization

  • After decoder updates, re-initializing the latent policy to N (0, I) avoids degradation by restarting policy search from the new decoder’s high-density support region.Without the reset, the latent policy remains aligned to the previous decoder and can degrade immediately, especially after the third update.Source: §4.3.2 · p. 10

4.3.3. Impact of Staged Decoder Refinement

  • Fresh encoders achieve monotonically higher capability ceilings with later frozen decoders, while HopperStand policies evolve from unimodal at 60M steps to bimodal at 180M.On HopperStand, the diffusion decoder’s returns rise from 168 at Stage 0 to 869 at Stage 3, while the Gaussian decoder rises from 134 to 282; gains saturate by Stage 3.Source: §4.3.3 · p. 11, Figure 5 · p. 11, §4.4 · p. 12

5. Related Work

Prior latent-space methods either rely on static offline data or optimize against a frozen pretrained generator. GoRL instead combines online latent-policy optimization with separate decoder refinement.

  • Offline latent-space methods search latents from generators trained on static datasets, while DSRL optimizes a latent controller over a frozen pretrained diffusion backbone whose fixed manifold restricts action support.Source: §5 · p. 12
  • Other online generative-policy methods use specialized objectives or estimators, whereas GoRL keeps policy optimization in a tractable latent space and trains its decoder separately via supervised generative learning.Source: §5 · p. 12

6. Conclusion

GoRL separates stable latent-space optimization from expressive action generation, improving stability, final performance, and mode coverage across continuous-control tasks. Its added refinement cost and experimental design choices motivate further work.

6. Conclusion

  • Across diverse continuous-control tasks, GoRL improves stability, final performance, and mode coverage relative to unimodal and coupled generative baselines.The framework confines policy-gradient updates to a tractable latent space and refines the expressive decoder separately.Source: §6 · p. 12
  • Periodic decoder refinement adds wall-clock cost, and stage boundaries are set a priori in the experiments.Stages that are too short may provide weak decoder-refinement signals, while stages that are too long may waste interaction budget after returns plateau.Source: §6 · p. 12
  • Adaptive refinement based on returns, latent entropy, or value-loss plateaus is proposed as a next step.Source: §6 · p. 13
  • Experiments leave higher-degree-of-freedom systems and visual observations open for future work.The default latent dimension matches the action dimension; other latent sizes may trade off stability, exploration, and expressiveness, and may require projection layers.Source: §6 · p. 13

Impact Statement

This material frames GoRL’s motivation and approach, reviews why standard policy gradients are difficult to apply directly to generative policies, and gives RL background. It also notes potential impacts and that real-world deployment is not evaluated.

Impact Statement

  • The paper identifies potential benefits for complex continuous-control applications, while noting possible misuse in safety-critical settings and that experiments cover simulated benchmarks only.It calls for safety constraints and rigorous evaluation before physical or high-stakes deployment.Source: Impact Statement · p. 13

A. Why Generative Policies Are Hard to Update with Standard Policy Gradients

A.1. Three Classical Routes to Policy Gradients

  • Stable policy gradients rely on tractable likelihoods, efficient differentiable samplers, or accessible cumulative distribution functions.Source: §A.1 · p. 17

A.2. Diffusion Policies Violate All Three Routes

  • Diffusion policies have costly likelihoods, delicate gradients through long denoising chains, and no tractable CDF for implicit gradients.Source: §A.2 · p. 18

A.3. Flow-Matching Policies Violate All Three Routes

  • For flow-matching policies, likelihood evaluation is costly or unavailable, ODE-based pathwise gradients are expensive and fragile, and CDF-based gradients are unavailable.Source: §A.3 · p. 18

A.4. Implications for Online RL and the GoRL Solution

  • GoRL confines policy optimization to a tractable latent policy and trains the expressive decoder separately with supervised objectives anchored to a fixed prior.This avoids placing intractable generative likelihoods inside the RL optimization loop.Source: §A.4 · p. 19

B. Reinforcement Learning Background

B.1. Markov Decision Processes

  • The continuous-control problem is modeled as an MDP with continuous state and action spaces, transition dynamics, rewards, discount factor, and initial-state distribution.The policy objective is expected discounted cumulative return over trajectories sampled under the policy and dynamics.Source: §B.1 · p. 19

B.2. Policy Gradient Theorem

The appendix reviews policy-gradient optimization and specifies the conditional diffusion and flow-matching decoders used in GoRL. It also sets up the induced action policy analyzed in the theoretical proofs.

B.2. Policy Gradient Theorem

  • The policy-gradient theorem supplies the standard gradient estimator, which is approximated from collected trajectories and optimized with PPO’s clipped surrogate objective.Source: §B.2 · p. 19

C. Generative Policy Details

C.1. Diffusion-Based Policies

  • The diffusion decoder iteratively denoises latent noise, using deterministic DDIM sampling at inference and rollout collection and training its noise predictor on on-policy rollout pairs.Source: §C.1 · p. 20

C.2. Flow Matching Policies

  • The flow-matching decoder uses an optimal-transport path that linearly interpolates latent noise and actions, then solves a neural ODE from t=0 to t=1.Its vector field is trained to regress the conditional target field.Source: §C.2 · pp. 20–21

D. Proofs of Theoretical Guarantees

D.1. Notation and Induced Action Policy

  • The induced action policy is obtained by sampling a latent from the state-conditioned encoder and decoding it into an action; its advantage is defined from action values and state values.Source: §D.1 · p. 21

D.2. Proof of Lemma 3.1 (Unbiased Latent Policy Gradient)

With a fixed decoder, the proof expresses the action-policy gradient using the latent-policy score and the advantage of the decoded action. This yields the stated unbiased latent-space estimator.

D.2. Proof of Lemma 3.1 (Unbiased Latent Policy Gradient)

  • Fixing the decoder gives the induced action policy the unbiased latent policy-gradient estimator stated in Lemma D.1.Source: §D.2 · p. 21
  • Sampling from the induced action policy is equivalent to sampling a latent and setting the action to its decoder output.Source: §D.2 · p. 21
  • Because the fixed decoder does not depend on the latent-policy parameters, their effect on the induced policy enters through the latent policy’s log-probability.Source: §D.2 · p. 21

D.3. Proof of Lemma 3.2 (Performance under Small Latent Divergence)

The proof bounds performance change using latent-policy divergence, while the implementation details describe matched baselines and a staged GoRL training schedule.

D.3. Proof of Lemma 3.2 (Performance under Small Latent Divergence)

  • With bounded advantages, latent-policy divergence bounds action and state-distribution shifts, yielding a bound on performance change.The decoder’s deterministic mapping contracts divergence from latent space to action space, and the discussion states that performance degradation is bounded by O(δ).Source: §D.3 · p. 22, §D.4 · p. 23

D.4. Discussion: Stability and Regularization

  • Reinitializing the encoder to the prior anchors latent updates to the decoder’s training support, so sufficiently positive expected advantage yields improved return.The discussion states that performance degradation is bounded by O(δ) when successive encoder policies remain close in latent space.Source: §D.4 · p. 23

E. Training and Implementation Details

E.1. Environments and Preprocessing

  • The six DMControl tasks span varied control demands; agents train from scratch for 180M steps using 2048 parallel environments.Observations are normalized online, rewards are scaled by 10.0, actions are tanh-squashed to [−1, 1], and advantages are normalized per batch.Source: §E.1 · pp. 23–24, Figure 7 · p. 23

E.2. Shared PPO Hyperparameters

  • GoRL’s encoder and Gaussian PPO, FPO, and DPPO share the PPO optimization pipeline to isolate policy parameterization.Shared PPO hyperparameters are listed in Table 4.Source: §E.2 · p. 24, Table 4 · p. 24

E.3. Baselines

  • Baseline actor and decoder capacity is matched, while FPO and DPPO use unregularized formulations because exact entropy requires costly ODE/SDE integration.The text notes that this limits efficient entropy regularization online and may contribute to instability on harder tasks relative to GoRL.Source: §E.3 · pp. 24–25

E.4. GoRL Specifics

  • The GoRL latent policy is a diagonal Gaussian with dimension equal to the task’s action dimension, and its decoder is initialized near an identity mapping.The decoder can be flow matching or diffusion; rollout sampling uses 10 steps, and decoder training samples latents from the fixed prior N(0, I).Source: §E.4 · p. 25
  • GoRL uses four stages of 60M, 60M, 30M, and 30M steps, with fresh rollout data collected for decoder refinement at stage boundaries.Each refinement stage uses 512 episodes; refinement interactions add about 1.5M steps, less than 1% of the main budget.Source: §E.4 · p. 25

E.5. Evaluation and Visualization Details

Evaluation averages performance across five seeds, and the appendix also describes how policy-action density plots are constructed.

E.5. Evaluation and Visualization Details

  • Performance is averaged over five seeds, with evaluations every 6M steps using 128 parallel environments for a full episode.Learning-curve shading denotes one standard deviation.Source: §E.5 · p. 26
  • Density plots use 10,000 actions from a representative stable state, projected onto PC1 and smoothed with KDE at bandwidth 0.8.Source: §E.5 · p. 26

F. Additional Experimental Results

F.1. Computational Cost Analysis

Across three tasks, GoRL’s additional training time accompanies higher average return, unlike FPO and DPPO. An SAC-based GoRL variant also matches or exceeds Gaussian SAC on the tested Gym environments.

F.1. Computational Cost Analysis

  • ≈1.63× return accompanies GoRL’s ≈1.87× wall-clock time relative to PPO; FPO and DPPO reach 0.43× and 0.81× return despite 1.47× and 1.58× overhead.Figure 8 reports return and training time normalized to Gaussian PPO across CheetahRun, FingerSpin, and HopperStand, under a matched 180M-step budget.Source: Figure 8 · p. 26, §F.1 · p. 26

F.2. Off-Policy Compatibility with SAC

  • GoRL(SAC+Diffusion) matches or exceeds Gaussian SAC across Hopper-v2, Walker2D-v2, and HalfCheetah-v2, with later-stage leads on Walker2D and HalfCheetah.The passage characterizes this as an initial compatibility check for off-policy settings.Source: Figure 9 · p. 27, §F.2 · p. 27
  • GoRL with SAC optimizes latent actions through a latent critic, while the frozen decoder synthesizes environment actions only for rollout execution.Thus, actor updates do not backpropagate RL gradients through the diffusion sampling chain.Source: §F.2 · p. 27
Loading 2512.02581v3…