Reinforcement Learning with Discrete Diffusion Policies for Combinatorial Action Spaces

Reinforcement learning (RL) struggles to scale to large, combinatorial action spaces common in many real-world problems. This paper introduces a novel framework for training discrete diffusion models as highly effective policies in these complex settings. Our key innovation is an efficient online training process that ensures stable and effective policy improvement and . By leveraging policy mirror descent (PMD) to define an ideal, regularized target policy distribution, we frame the policy update as a distributional matching problem, training the expressive diffusion model to replicate this stable target. This decoupled approach stabilizes learning and significantly enhances training performance. Our method achieves state-of-the-art results and superior sample efficiency across a diverse set of challenging combinatorial benchmarks, including DNA sequence generation, RL with macro-actions, and multi-agent systems. Experiments demonstrate that our diffusion policies attain comparable or superior performance compared to other baselines. Crucially, our extensive empirical analysis reveals a key trade-off: FKL demonstrates superior sample efficiency and faster initial convergence, whereas RKL ensures stable training and higher asymptotic performance on challenging tasks.

Teaching artificial intelligence to make decisions is usually done one step at a time, but this method breaks down when an AI needs to make many complex choices all at the exact same time-such as coordinating an entire soccer team's strategy or designing a complete DNA sequence. To solve this bottleneck, this research introduces a new framework that allows the AI to build massive, multi-part decisions simultaneously rather than piecing them together one by one. Because training an AI to do this can be highly unstable, our framework uses a two-step process: it first calculates a reliable mathematical blueprint of the current AI behavior, and then simply trains the AI to improve that blueprint. In out work, we also uncovered a key trade-off in how the AI learns: one approach learns incredibly fast with less data, while the other learns more slowly but remains highly stable for the most difficult challenges. When tested on designing synthetic DNA, planning long-term strategies in video games, and managing a simulated 11-player soccer team, this new approach achieved top-tier results, proving it can successfully tackle massive decision-making problems that traditional AI simply cannot handle.