Reward-Weighted Regression (RWR)

• Reward-Weighted Regression (RWR) is an EM-based algorithm for policy optimization that uses reward-weighted likelihood to drive monotonic improvement in reinforcement learning.
• It utilizes a two-step process where trajectories are sampled and weighted by return, followed by a closed-form policy update via weighted log-likelihood maximization.
• RWR guarantees global convergence under ideal conditions while highlighting challenges such as function approximation issues and distributional drift in practical settings.

Reward-Weighted Regression (RWR) is an Expectation-Maximization (EM)–based algorithm family for policy optimization in reinforcement learning (RL), designed to produce monotonic policy improvement by fitting new policies to maximize a return-weighted likelihood of observed trajectories. RWR's formalization, convergence guarantees, and applications span domains such as MDPs, LLM post-training, and flow-based generative models. It provides a tractable surrogate for RL objectives, endowing EM-type structure with sample-based and variational interpretations that facilitate both theoretical analysis and efficient large-scale training.

1. Mathematical Definition and EM Formulation

In a Markov Decision Process (MDP) $$\mathcal{M}=(\mathcal{S}, \mathcal{A}, p_T, R, \gamma, \mu_0)$$ with compact state and action spaces, continuous strictly positive reward $R(s, a)$, and continuous transition kernel $p_T$, RWR operates by performing the following update at every iteration $k$ with policy parameters $\theta_k$:

• E-step: Sample trajectories $\{\tau_i\}_{i=1}^N$ using current policy $\pi_{\theta_k}$; for each, compute accumulated discounted return $R(\tau_i)=\sum_{t=0}^\infty \gamma^t R(s_t, a_t)$.
• Weighting: Assign each trajectory a "reward weight," typically $w(\tau_i) \propto \exp(R(\tau_i)/\beta)$ (with temperature $\beta > 0$), or $w(\tau_i)=R(\tau_i)$ in some theoretical results.
• M-step: Fit the next policy by maximizing the weighted log-likelihood:

$$\theta_{k+1} = \arg\max_\theta \sum_{i=1}^N w(\tau_i) \log \pi_\theta(\tau_i)$$

or equivalently, the sum over all state-action pairs within sampled trajectories.

In the infinite-data limit and with exact representations,

$$\pi_{k+1}(a|s) = \frac{Q^{\pi_k}(s,a)\, \pi_k(a|s)}{V^{\pi_k}(s)}$$

where $Q^{\pi_k}$ and $V^{\pi_k}$ are the policy's action-value and value functions, and this defines the RWR operator $B$:

$$(B\pi)(a|s) = \frac{Q^\pi(s,a) \pi(a|s)}{\int Q^\pi(s,a')\pi(a'|s) da'}$$

(Štrupl et al., 2021).

2. Convergence Theory and Optimality Guarantees

RWR exhibits global convergence to the unique optimal policy under certain regularity and compactness assumptions:

• Compact $\mathcal{S},\mathcal{A}\subset\mathbb{R}^n$;
• Strict positivity and continuity of $R(s,a)$ and $p_T$;
• Exact (nonparametric) policy representations and strictly-positive initial policy $\pi_0$.

The convergence analysis (Štrupl et al., 2021) establishes:

• Monotonic Policy Improvement: $V^{\pi_{n+1}}(s) \ge V^{\pi_n}(s)$ pointwise, with strict increase whenever the policy places nonzero mass on suboptimal actions.
• Limit Behavior: The sequence $\pi_n$ weakly concentrates on maximally-rewarded (greedy) actions; limiting policy $\pi_L$ is supported on the set of global maximizers $M(s)=\arg\max_a Q^*(s,a)$ in every state.
• Global Optimum: The limit $(V_L, Q_L, \pi_L)$ satisfies Bellman's optimality equations; $\pi_L$ is the unique optimal policy of the MDP.
• Finite MDPs—R-linear (Geometric) Convergence: In discrete state/action spaces, $V_n \to V^*$ with $\|V^* - V_n\|_\infty = O(\alpha_m^n)$ for an explicit $\alpha_m<1$ depending on the suboptimal action values (Štrupl et al., 2021).

3. Algorithmic Variants and Application Domains

Trajectory Weighting Variants

• Exponential Weighting: $w(\tau) \propto \exp(R(\tau)/\beta)$ concentrates updates on high-return trajectories for small $\beta$, accelerating greedy improvement but risking instability or premature convergence.
• Linear Weighting: $w(\tau) = R(\tau)$ preserves the optimum and underlies theoretical convergence analysis (Štrupl et al., 2021).

Modern Extensions

• LLM Fine-Tuning: RWR bounds underlie methods such as Dynamic Fine-Tuning (DFT) and Anchored Supervised Fine-Tuning (ASFT), where the loss is formulated as a reward-weighted negative log-likelihood over human-preferred trajectories, sometimes augmented with KL-anchoring to ensure distributional stability (Zhu et al., 28 Sep 2025).
• Generative Modeling with Human Feedback: In video generation, RWR is instantiated as Flow-RWR, training rectified-flow models using reward-weighted MSE on velocity or noise predictions, with per-sample weights obtained from learned reward models (Liu et al., 23 Jan 2025).

4. Theoretical Connections and Distinctions

RWR is formally distinct from other EM-like or information-theoretic policy optimization schemes:

• Closed-Form M-step: Unlike Cross-Entropy Method (CEM), Relative Entropy Policy Search (REPS), or Maximum A Posteriori Policy Optimization (MPO), RWR provides a closed-form policy update maximizing the weighted log-likelihood, obviating explicit KL constraints or dual optimization (Štrupl et al., 2021).
• Lower-Bound Structure: When applied to supervised post-training (e.g., SFT, DFT), RWR can be derived as optimizing a surrogate lower bound on the true RL objective, with tightness controlled via the choice of auxiliary weighting distribution $q$ (Zhu et al., 28 Sep 2025).
• Variational and Importance-Sampling Interpretation: Given limited or demonstration-only data, RWR formalizes policy improvement as importance-sampled EM, where improved lower bound tightness trades off with increased variance and risk of distributional drift unless regularized (Zhu et al., 28 Sep 2025).

5. Limitations and Open Problems

• Function Approximation: The global convergence guarantee is lost when parameterized function approximators (e.g., neural networks) are employed for policies or value functions; RWR then exhibits classic EM pathologies such as local suboptimality (Štrupl et al., 2021).
• Distributional Drift in Auxiliary Weighting: In applied settings (e.g., DFT for LMs), iterative reweighting without anchoring induces policy drift and collapsing effective sample size, necessitating KL-based regularization (as in ASFT) for variance control and stable learning (Zhu et al., 28 Sep 2025).
• Tradeoffs in Weighting Schedules: Aggressive weighting (small $\beta$) may speed optimization but can cause mode collapse or suboptimal convergence due to vanishing variance. More tempered, smooth weighting increases robustness and effective exploration (Štrupl et al., 2021, Liu et al., 23 Jan 2025).
• Alignment vs. Sample Efficiency: In reward-modulated generative modeling (Flow-RWR), the single-sample weighting simplifies implementation but often underperforms pairwise preference-optimization approaches such as Flow-DPO, particularly on high-precision alignment objectives (Liu et al., 23 Jan 2025).

6. Implementation Protocols and Empirical Results

RWR variants typically require:

• Computation or estimation of trajectory or sample-level rewards, potentially using learned reward models (e.g., VideoReward for video, binary indicator for demonstrations in LMs).
• Generation and normalization of sample weights, typically $\exp(r/\beta)$ or variants thereof, with stabilization via mean/variance normalization or batch-level scaling.
• Weighted regression or loss minimization—e.g., weighted MSE for generative models, reward-weighted log-likelihood for sequence models.

Empirical findings indicate:

• For LLMs, ASFT (an RWR-derived, KL-regularized method) outperforms both standard SFT and unanchored DFT on mathematical reasoning, medical QA, and code generation, with documented improvements in accuracy (e.g., +17.89 pp on math reasoning at 100k scale) and KL divergence stabilization (Zhu et al., 28 Sep 2025).
• For rectified-flow models in video generation, Flow-RWR facilitates reward alignment but is typically superseded by pairwise-based methods (Flow-DPO) in achieving fine-grained text-video consistency and overall alignment, though the former is simpler to implement (Liu et al., 23 Jan 2025).

7. Comparative Table: RWR in Key Domains

Domain	RWR Instantiation	Key Update Objective
MDP/Classic RL	Weighted log-likelihood	$\mathbb{E}[w(\tau)\log \pi_\theta(\tau)]$
LLM Fine-Tuning	DFT, ASFT	$-\mathbb{E}[w(\tau)\log\pi_\theta(\tau)]$ + KL anchoring
Flow-based Video Generation	Flow-RWR	$$\mathbb{E}[w_i\|\mathbf v_i - \mathbf{\hat v}_i\|^2]$$

RWR unifies several contemporary optimization routines under a broadly-applicable EM-based framework. It admits rigorous convergence proofs in exact settings, underpins recent developments in post-training LLMs and feedback-aligned generation, and highlights principled design choices for reward modulation and stability. Open research questions remain regarding its behavior and guarantees when combined with high-capacity function approximators and in high-variance, limited-data regimes (Štrupl et al., 2021, Zhu et al., 28 Sep 2025, Liu et al., 23 Jan 2025).