Sum-and-Sample RL: Replay Buffers


Reinforcement learning is thought to work best when on-policy. This means learning from actions sampled from the current policy. But this approach can be very sample-inefficient if the policy is initially crap. In this blog post, I will explore one simple method for learning from off-policy samples.

Note: This is related to previous blog post, where I wrote a bit about one path towards merging SFT and RL. The previous post relied on importance sampling, a method of approximating samples from a target distribution using an importance distribution. But importance sampling is very tricky, as it introduces a tricky hyperparameter: the importance distribution.

Policy gradient

Off-policy policy gradient

Let’s look at the math of off-policy policy gradient (OPPG) a bit to see what can go wrong.

OPPG optimizes the expected reward, but with trajectories from an importance distribution q(x)q(x)

arg max⁡pEp(x)[r(x)]=arg max⁡p∑xp(x)r(x)=arg max⁡p∑xp(x)q(x)q(x)r(x)=arg max⁡pEq(x)[p(x)q(x)r(x)].\begin{align} &\argmax_p E_{p(x)}[r(x)]\\ &= \argmax_p \sum_x p(x)r(x)\\ &= \argmax_p \sum_x p(x) \frac{q(x)}{q(x)} r(x)\\ &= \argmax_p E_{q(x)}\left[\frac{p(x)}{q(x)} r(x)\right]. \end{align}

Sum-and-sample gradient estimator

paper, coin termed here

Past work: https://arxiv.org/abs/2207.00632