Organizations: School of Data Science, Fudan University · Laboratory for Big Data and Decision, National University of Defense Technology · ShanghaiTech University · College of Computer Science and Artificial Intelligence, Fudan University · Shanghai Innovation Institute
Classical on-policy algorithms such as PPO and mirror descent policy optimization provide stable proximal policy updates through tractable action likelihoods, but are typically instantiated with simple Gaussian policies whose expressiveness can be limited in complex continuous-control tasks. Generative policies based on diffusion and flow models provide more expressive action distributions, but they naturally define distributions over multi-step denoising paths whose terminal action density is often intractable, creating a mismatch with likelihood-based on-policy proximal updates. To address this mismatch, we introduce GSB-MDPO (Generalized Schrödinger Bridge Mirror Descent Policy Optimization), which formulates on-policy generative policy optimization as a Generalized Schrödinger Bridge problem over state-conditioned generation paths and instantiates the resulting path-measure update through mirror descent policy optimization. The key insight is that the GSB path-space KL plays the role of the proximal term in MDPO while upper-bounding the terminal action KL, enabling direct control of the executed action distribution without explicit terminal action likelihood evaluation. Experiments on 14 continuous-control tasks across Playground and Gym-MuJoCo demonstrate the empirical effectiveness of GSB-MDPO and support path-space regularization as a principled proximal update for multi-step generative policies.
Figures & tables
Figure 1: Playground results. Curves show mean ± one standard deviation over five seeds.
Figure 2: Gym-MuJoCo results. Panels (A)–(F) show learning curves on six tasks, with mean ± one standard deviation over five seeds. Panel (G) reports representative training time and GPU memory on HalfCheetah-v5.
Figure 3: Ablations on HalfCheetah-v5. Curves show the mean and min–max envelope over two seeds.
Figure 4: Ablation of the number of generation steps for GSB-MDPO on HalfCheetah-v5. Solid curves show the mean over two seeds, and shaded regions show the min–max envelope across seeds.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Hyperparameter
PPO
DPPO
FPO
GenPO
GSB-MDPO
Benchmark and rollout setup
Total environment steps
60M
60M
60M
60M
60M
Number of environments
2048
2048
2048
2048
2048
Episode length
1000
1000
1000
1000
1000
Number of evaluations
10
10
10
10
10
Unroll length
30
30
30
30
30
Appendix
Table 1: Hyperparameters for Playground experiments. A single configuration is used for all Playground tasks within each method.
Hyperparameter
PPO
DPPO
FPO
GenPO
GSB-MDPO
Benchmark and rollout setup
Total environment steps
40M
40M
40M
40M
40M
Evaluation interval
1M
1M
1M
1M
1M
Evaluation episodes
10
10
10
10
10
Number of environments
1024
1024
1024
1024
1024
Rollout length
24
24
24
24
24
Appendix
Table 2: Hyperparameters for Gym-MuJoCo experiments. A single configuration is used for all Gym-MuJoCo tasks within each method.
Figure 5: Toy visualization of the optimal GSB-MDPO update. Left: old policy pold with four approximately balanced modes. Middle: quadrant-wise preference weights exp(A/β) . Right: target exponential-tilted marginal p∗ and the learned distribution pθ . The learned policy closely matches the theoretical tilted target, increasing mass in the preferred quadrant while preserving the old policy’s mode structure.
Figure 6: Fraction of rollout samples whose step log-ratio is clipped during GSB-MDPO training on HalfCheetah-v5. Results are averaged over two random seeds, with shaded regions indicating variability across seeds.
Method
Wall-clock
Rel. time
Peak GPU memory
Mean GPU memory
PPO
6.24 min
1.00×
0.60 GB
0.59 GB
FPO
6.96 min
1.12×
1.00 GB
0.97 GB
DPPO
8.67 min
1.39×
2.50 GB
2.38 GB
GenPO
28.73 min
4.61×
64.98 GB
64.13 GB
GSB-MDPO
7.39 min
1.19×
2.59 GB
2.51 GB
Appendix
Table 3: Representative computational cost on HalfCheetah-v5 for 5 M environment interactions. Wall-clock time is end-to-end training time. GPU memory is measured from the training process tree.
Combining entropy regularisation with mirror descent in online off-policy reinforcement learning requires expressive policies with efficient sampling and explicit entropy control. Gaussian policies struggle to represent multimodal targets, whereas generative policies often require iterative sampling or lack tractable entropy estimates. We introduce Stochastic MeanFlow Policies (SMFP), a one-step generative policy class that combines MeanFlow-based noise-to-action mappings with Gaussian reparameterisation. The latent mapping captures diverse action modes, while a learnable conditional Gaussian scale provides an analytic entropy lower bound. We use this bound to regulate stochasticity alongside direct value guidance and advantage-weighted MeanFlow regression, enabling entropy-regularised mirror-descent learning with one-step policies. Across seven MuJoCo benchmarks, SMFP achieves strong performance against Gaussian and generative baselines with single-step inference. Additional experiments demonstrate multimodal action modelling and effective exploration in sparse-reward navigation.
Zeyuan Wang, Da Li, Yulin Chen +5
Laboratory for Big Data and Decision, National University of Defense Technology, China · Samsung AI Center Cambridge, UK · Queen Mary University of London, UK +3
Expressive generative policies such as diffusion and flow models are appealing for maximum-entropy (MaxEnt) online reinforcement learning because of their ability to model multimodal and highly non-Gaussian action distributions. However, training effective soft generative policies faces two obstacles that often arise together. First, marginal action densities are often unavailable, so existing methods typically rely on entropy bounds, heuristic proxies or approximations. Second, iterative shared-parameter samplers raise inference cost and require backpropagation through time over repeated network evaluations, increasing memory cost and potentially destabilizing policy optimization. These obstacles motivate us to seek a generative actor-critic method that exposes a tractable soft objective while generating each action in just one network evaluation. To this end, we propose soft generative actor-critic (SoftGAC), whose actor is a single stochastic feedforward network with internal Gaussian bridge layers. This explicit bridge structure allows us to lift the soft objective to path space, yielding an analytical path-KL regularizer against a high-entropy reference process for both actor improvement and the soft Bellman bootstrap, without estimating the endpoint entropy. Crucially, the actor generates each action in only one network evaluation under a total parameter budget comparable to strong baselines, substantially reducing inference cost relative to the existing iterative diffusion and flow policies. It also avoids long unrolled backpropagation through time. Across challenging continuous-control benchmarks, our method achieves competitive or superior returns compared with strong diffusion and flow policies, demonstrating a clear advantage in the compute-return tradeoff.
Ke He, Le He, Shunpu Tang +2
Guangzhou University · Zhejiang University · Southeast University
Diffusion policies are a powerful policy class for continuous control, but their iterative denoising process creates a substantial computational bottleneck. Reducing this cost requires adapting the number of denoising steps to the difficulty of each action while preserving task performance. We introduce Prefix-Optimal Generative Policies (POGP), a framework that learns a prefix value function at every intermediate denoising step through a Bellman-style recursion over the denoising chain. The prefix value function serves two purposes: it provides an auxiliary training objective that encourages intermediate outputs to become high-quality actions, and it enables a test-time stopping rule that terminates denoising when additional steps are unlikely to produce meaningful improvement. Across four MuJoCo environments and comparisons with 12 baselines, POGP reduces the required number of denoising iterations by approximately 2.7-fold while retaining near-full task performance. Compared with state-of-the-art dynamic diffusion baselines, prefix training also improves final task performance by approximately 3.5%. These results indicate that supervising intermediate denoising steps is useful not only for adaptive early stopping, but also as an auxiliary objective that improves the learned policy.
Rohit Kumar Salla, Manoj Saravanan, Simon Stepputtis
Virginia Polytechnic Institute and State University (Virginia Tech), Blacksburg, Virginia, USA