Reward Distribution Matching

• Reward distribution matching is a framework that defines the objective as matching the entire reward distribution—including variance, risk, and multiple high-reward modes—rather than just maximizing expected rewards.
• It employs advanced techniques such as non-Markovian policies and distributional metrics like the Wasserstein distance to accurately replicate expert or target behaviors.
• Empirical results demonstrate that these methods enhance sample efficiency, robustness, and practical performance in complex reinforcement learning, imitation learning, and multi-agent applications.

Reward distribution matching is a paradigm in decision-making, control, and learning systems where the objective extends beyond maximizing the expected reward to shaping, aligning, or imitating the entire reward distribution produced by a system, agent, or set of agents. Instead of focusing solely on mean outcomes, reward distribution matching considers variance, risk sensitivity, diversity, or the richness of multiple high-reward modes, and develops policies or mechanisms to align distributions with a specified target across domains such as reinforcement learning, imitation learning, resource allocation, and economic systems.

1. Fundamental Concepts and Problem Formulation

Reward distribution matching generalizes classical reward maximization by formulating objectives that explicitly target the entire distribution of accumulated rewards or outcomes, not just their mean. A canonical instance is risk-sensitive imitation learning, where the learner aims to match the expert’s return distribution—quantified, for example, by the Wasserstein distance—rather than best-approximate its mean performance (Lazzati et al., 15 Sep 2025). More broadly, the target distribution $\tilde{\pi}(y|x)$ may be derived from scalar rewards by exponentiating and normalizing (i.e., $\tilde{\pi}(y|x) = \exp(\beta r(x,y))/Z(x)$), as in energy-based models or maximum entropy RL (Zhu et al., 18 Sep 2025).

This contrasts with standard occupancy measure matching or BeLLMan optimality criteria, which ensure only expected rewards are aligned. For many reasoning, control, and multi-agent applications, matching higher-order properties (variance, tail behavior, risk) or even the full empirical return distribution can yield more robust, diverse, and interpretable policies.

2. Policy Classes and Expressivity

The expressivity of policy classes is critical for reward distribution matching. Standard Markovian policies, which select actions based solely on the current state and (possibly) time, are insufficient to reproduce arbitrary reward distributions. This is most apparent in matching the expert’s return distribution, where "risk attitude" is encoded nontrivially in the shape of the distribution and may depend on trajectory-level (non-Markovian) features (Lazzati et al., 15 Sep 2025).

The solution is to use non-Markovian policies augmented to access summary statistics of the trajectory, most commonly the cumulative reward up to a given timestep. Formally, for a reward function $r$, the cumulative reward along trajectory $\omega$ up to time $h$ is $$G(\omega; r) = \sum_{h'=1}^{h-1} r_{h'}(s_{h'}, a_{h'})$$. A sufficiently expressive class consists of policies $\pi(a|s, \omega) = \varphi_h(a|s, G(\omega; r))$—that is, policies at each $h$ condition actions on both $s$ and the cumulative past reward. This expansion is essential for the learner to truly match the expert’s risk profile or other higher-order moments of the return distribution.

3. Algorithms and Optimization Principles

Practical instantiations of reward distribution matching adopt different frameworks, typically depending on availability of the environment model and expert information:

• Empirical Return Distribution Matching: In the offline setting (transition model unknown), risk-sensitive behavioral cloning (RS-BC) estimates the expert’s empirical return distribution via discretization and counts, then clones the action probabilities conditioned on state and cumulative reward (Lazzati et al., 15 Sep 2025). In the online or known-model setting (RS-KT), a linear program in an augmented state space (state, cumulative reward) is solved to find the occupancy measure and policy minimizing Wasserstein distance to the empirical expert’s return distribution.
• Normalized Target Distribution and Reverse KL: In LLMs and general RL, scalar rewards $r(x, y)$ are mapped to normalized target distributions $\tilde{\pi}(y|x)$ using a partition function, with policies trained to minimize $D_{KL}(\pi_\theta \| \tilde{\pi})$. This is operationalized through squared error ("trajectory balance") losses and flow-based optimization—directly matching the distribution rather than solely optimizing for high reward (Zhu et al., 18 Sep 2025).
• Imitation Learning by Distribution Matching: Methods such as adversarial imitation learning, occupancy measure matching, and off-policy KL objectives (e.g., ValueDICE, GAIL, and related frameworks) directly minimize divergence metrics (e.g., forward or reverse KL, Jensen–Shannon divergence) between the policy’s visitation distributions and those from expert demonstrations, often employing adversarial training or dual variable reparameterizations (Kostrikov et al., 2019).
• Reward Redistribution via Likelihood Maximization: To address delayed or sparse rewards, likelihood-based frameworks model per-step rewards as random variables and use explicit trajectory-level likelihoods (e.g., via leave-one-out strategies) to distribute global returns across the trajectory, matching resulting per-step reward distributions with uncertainty regularization (Xiao et al., 20 Mar 2025).

4. Theoretical Properties and Sample Complexity

Matching full reward distributions rather than only means introduces higher sample complexity, especially for non-Markovian policies as the necessary class is larger. Theoretical analyses show:

• Sample Complexity: When matching return distributions under known dynamics (with suitable policy class), sample complexity depends polynomially on the time horizon $H^{-2}$ and is independent of the number of states or actions, as in RS-KT (Lazzati et al., 15 Sep 2025). Offline behavioral cloning (RS-BC) incurs higher polynomial costs in $S, A, H$.
• Optimality and Approximation Guarantees: When the expert’s reward is not observed, uniform convergence and empirical process theory (e.g., Dvoretzky–Kiefer–Wolfowitz inequality) are used to show that a finite dataset suffices to approximate all possible return distributions (over a parameterized reward cover) to specified Wasserstein tolerance.
• Necessity of Expressivity: Markovian policies may match the expert’s expected return but cannot reproduce the variance or higher moments. This is shown constructively in both theoretical analysis and simulations.
• Distributional Losses and Metric Selection: Wasserstein distance is especially suited to distribution matching objectives due to its sensitivity to both means and higher-moment discrepancies, controlling both location and shape alignments between return distributions.

5. Empirical Results and Practical Benefits

Empirical evaluations of risk-sensitive and distribution-matching algorithms confirm several key findings:

• Non-Markovian policy architectures substantially reduce Wasserstein errors in return distribution match compared to Markovian baselines such as standard behavioral cloning or occupancy measure matching (Lazzati et al., 15 Sep 2025).
• In both low- and high-dimensional tabular settings, as the number of expert trajectories increases, empirical Wasserstein distance to the expert’s return distribution decreases for RS-BC and RS-KT, while standard baselines stall due to expressivity inadequacy.
• In LLM settings, FlowRL demonstrates $5$–$10$\% improvement in math and code reasoning metrics over reward-maximizing baselines, specifically attributable to the maintenance of solution diversity and coverage of multiple high-reward trajectories (Zhu et al., 18 Sep 2025).
• In imitation learning, distribution matching approaches enable sample-efficient learning, especially off-policy settings, and provide robustness to reward misspecification, supporting richer behavioral diversity and risk attitude transfer.

6. Significance and Future Directions

Reward distribution matching broadens the scope of classical reinforcement and imitation learning, with direct implications for applications where diversity, robustness, or risk sensitivity are essential (e.g., LLM reasoning, dialog, multi-agent coordination, financial AI, safe RL). Key directions include:

• Efficient Policy Classes: Further work on compact, tractable non-Markovian policy representations to support distribution matching in large and continuous spaces.
• Distributional Objectives in LLM and Complex RL: Extending normalization and flow-based target construction to sequencing, reasoning, and "open-ended" generative tasks.
• Distribution Matching in Multi-Agent Systems: Theoretical and algorithmic advances for matching joint or marginal distributions in cooperative or competitive environments.
• Evaluation Metrics and Calibration: Studying appropriate divergences (e.g., Wasserstein, total variation, or energy-based metrics) for practical alignment, as well as behavior under distribution shift scenarios.
• Integration with Model-Based Planning and Bayesian Inference: Using distribution matching as a bridge between Bayesian policy inference, posterior sampling, and risk- and diversity-aware control.

In summary, reward distribution matching represents a substantial generalization of classical reward maximization, grounding policy optimization in an expressivity-rich, risk-informed, and distributionally robust framework. This is realized both through new policy classes and metrics and through algorithms that directly and efficiently align full reward distributions with complex behavioral or theoretical targets.