Flow Matching for Offline Reinforcement Learning with Discrete Actions
About
Generative policies based on diffusion models and flow matching have shown strong promise for offline reinforcement learning (RL), but their applicability remains largely confined to continuous action spaces. To address a broader range of offline RL settings, we extend flow matching to a general framework that supports discrete action spaces with multiple objectives. Specifically, we replace continuous flows with continuous-time Markov chains, trained using a Q-weighted flow matching objective. We then extend our design to multi-agent settings, mitigating the exponential growth of joint action spaces via a factorized conditional path. We theoretically show that, under idealized conditions, optimizing this objective recovers the optimal policy. Extensive experiments further demonstrate that our method performs robustly in practical scenarios, including high-dimensional control, multi-modal decision-making, and dynamically changing preferences over multiple objectives. Our discrete framework can also be applied to continuous-control problems through action quantization, providing a flexible trade-off between representational complexity and performance.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Offline Reinforcement Learning | hopper medium | Normalized Score3.73e+3 | 52 | |
| Offline Reinforcement Learning | walker2d medium | Normalized Score1.25e+3 | 51 | |
| Offline Reinforcement Learning | halfcheetah medium | Normalized Score4.92 | 43 | |
| Offline Reinforcement Learning | Hopper expert | Normalized Score15 | 19 | |
| Reinforcement Learning | Acrobot v1 | Mean Return-147.2 | 14 | |
| Offline Reinforcement Learning | Walker2d Expert | Episodic Return211.3 | 12 | |
| Offline Reinforcement Learning | HalfCheetah Expert | Episodic Return1.26 | 12 | |
| Multi-objective offline reinforcement learning | Deep Sea Treasure | ND15 | 2 | |
| Multi-objective offline reinforcement learning | Resource Gathering | ND4 | 2 |