[{"content":"Main problem: in supervised learning, you have a loss function to differentiate. In RL the reward is a black box you can only sample from. Everything else in this chapter is to handle that fact.\nSetting Contextual bandit: see context \\(\\to\\) policy take action \\(\\to\\) get reward.\ncontext: \\(x\\sim \\mu\\) policy: \\(\\pi_{\\theta}(a\\mid x)\\) reward: \\(r\\sim R(x,a)\\) objective: \\(J(\\theta)=\\mathbb{E}[r]\\) Estimation The objective \\(J(\\theta)\\) is an expectation over unknown distributions. We use samples to estimate the expectation. As long as the samples are independently sampled from the correct distribution this gives us an unbiased estimate, known as the Monte Carlo estimate.\nThe Monte Carlo estimate is unbiased and its variance scales as \\(\\sigma^2_{\\theta}/n\\), for \\(n\\) samples.\nUnbiased: \\(\\mathbb{E}[\\hat{J}_{\\theta}]=J(\\theta)\\) \\(\\hat{J}_{\\theta}=\\frac{1}{n}\\sum_{i}r_{i}\\): random variable estimator, calculated from \\(n\\) samples. If the estimator is unbiased we would have \\(\\mathbb{E}[\\hat{J}_{\\theta}]=J(\\theta)\\). \\[ \\begin{aligned} \\mathbb{E}[\\hat{J}_{\\theta}] \u0026amp;= \\mathbb{E}\\left[ \\frac{1}{n}\\sum_{i} r_{i} \\right] \u0026amp;\u0026amp; \\text{definition of } \\hat{J}_{\\theta} \\\\[6pt] \u0026amp;= \\frac{1}{n}\\sum_{i}\\mathbb{E}[r_{i}] \u0026amp;\u0026amp; \\text{linearity of } \\mathbb{E} \\\\[6pt] \u0026amp;= \\frac{1}{n}\\cdot n\\cdot J(\\theta) \u0026amp;\u0026amp; \\text{each } r_{i} \\text{ drawn from the process defining } J(\\theta) \\\\[6pt] \u0026amp;= J(\\theta) \\end{aligned} \\] In expectation, the estimator is the true objective.\nVariance: \\(\\mathrm{Var}(\\hat{J}_{\\theta})=\\sigma^2_{\\theta}/n\\) Since our estimator is unbiased, \\(\\mathbb{E}[\\hat{J}_{\\theta}] = J(\\theta)\\) and we can swap one for the other:\n\\[ \\mathrm{Var}(\\hat{J}_{\\theta}) = \\mathbb{E}\\big[(\\hat{J}_{\\theta}-\\mathbb{E}[\\hat{J}_{\\theta}])^2\\big] = \\mathbb{E}\\big[(\\hat{J}_{\\theta}-J(\\theta))^2\\big] \\]\\[ \\begin{aligned} \\mathrm{Var}(\\hat{J}_{\\theta}) \u0026amp;= \\mathbb{E}\\left[ \\left( \\frac{1}{n}\\sum_{i}r_{i} - J(\\theta) \\right)^{2} \\right] \u0026amp;\u0026amp; \\text{plug in } \\hat{J}_{\\theta} \\\\[6pt] \u0026amp;= \\mathbb{E}\\left[ \\left( \\frac{1}{n}\\sum_{i}r_{i} - \\frac{1}{n}\\sum_{i}J(\\theta) \\right)^{2} \\right] \u0026amp;\u0026amp; J(\\theta)=\\tfrac{1}{n}\\textstyle\\sum_{i}J(\\theta),\\ \\text{const.} \\\\[6pt] \u0026amp;= \\mathbb{E}\\left[ \\left( \\frac{1}{n}\\sum_{i}\\big(r_{i}-J(\\theta)\\big) \\right)^{2} \\right] \u0026amp;\u0026amp; \\text{merge under one sum} \\\\[6pt] \u0026amp;\\overset{\\star}{=} \\frac{1}{n^{2}}\\sum_{i}\\mathbb{E}\\big[(r_{i}-J(\\theta))^{2}\\big] \u0026amp;\u0026amp; \\text{cross terms vanish} \\\\[6pt] \u0026amp;= \\frac{1}{n^{2}}\\cdot n\\,\\sigma^{2}_{\\theta} \u0026amp;\u0026amp; \\sigma^{2}_{\\theta}=\\mathbb{E}[(r_{i}-J(\\theta))^{2}],\\ n \\text{ of them} \\\\[6pt] \u0026amp;= \\frac{\\sigma^{2}_{\\theta}}{n} \\end{aligned} \\]\\(\\star\\) Squaring the sum gives \\(n\\) diagonal terms plus cross terms \\((r_{i}-J)(r_{j}-J)\\) for \\(i\\neq j\\). Independence splits the product, and unbiasedness makes each factor zero, so every cross term is \\(0\\cdot 0=0\\). The diagonal terms survive because squaring kills the sign cancellations.\nThis verifies the intuition that more samples are better as \\(n\\) samples reduces the variance by a factor of \\(n\\). Since samples are expensive in RL, we can also reduce \\(\\sigma^2_{\\theta}\\) itself which is what the next section is about.\nVariance Reduction Variance in reward comes from three sources:\nNoise inherent to the reward function Randomness in the policy\u0026rsquo;s actions Variation due to different contexts We can reduce variance by accounting for term 3, variation across contexts. Given a context, we want to know the relative increase in reward an action gives, not the raw reward for that context, because the nature of each context can lead to different scales of reward values.\nMotivated by this observation, instead of raw reward, let\u0026rsquo;s use a relative-to-average reward \\(r_{i}-b(x_{i})\\). Intuitively this term represents whether that reward is better or worse than average for that context, when \\(b(x)\\approx \\mathbb{E}_{a\\sim \\pi_{\\theta}(x),\\, r\\sim R(x,a)}[r]\\).\nSubtracting a context-dependent baseline shifts the objective by a \\(\\theta\\)-independent constant, so \\(\\nabla J_b = \\nabla J\\) and the argmax is unchanged.\nBaseline is free: \\(J_b(\\theta) = J(\\theta) - \\mathbb{E}_\\mu[b(x)]\\) \\[ \\begin{aligned} J_b(\\theta) \u0026amp;= \\mathbb{E}_{x\\sim\\mu}\\Big[\\mathbb{E}_{a\\sim\\pi_\\theta(x)}\\big[\\mathbb{E}_{r\\sim R(x,a)}[\\,r-b(x)\\,]\\big]\\Big] \u0026amp;\u0026amp; \\text{definition of } J_b \\\\[6pt] \u0026amp;= \\mathbb{E}_{x}\\Big[\\mathbb{E}_{a}\\big[\\mathbb{E}_{r}[r]-\\mathbb{E}_{r}[b(x)]\\big]\\Big] \u0026amp;\u0026amp; \\text{linearity of } \\mathbb{E} \\\\[6pt] \u0026amp;= \\mathbb{E}_{x}\\Big[\\mathbb{E}_{a}\\big[\\mathbb{E}_{r}[r]\\big]-b(x)\\Big] \u0026amp;\u0026amp; b(x) \\text{ const. w.r.t. } r,a \\\\[6pt] \u0026amp;= \\underbrace{\\mathbb{E}_{x}\\Big[\\mathbb{E}_{a}\\big[\\mathbb{E}_{r}[r]\\big]\\Big]}_{=\\,J(\\theta)} -\\;\\mathbb{E}_{x\\sim\\mu}\\big[b(x)\\big] \u0026amp;\u0026amp; \\text{linearity, outer layer} \\\\[6pt] \u0026amp;= J(\\theta)-\\underbrace{\\mathbb{E}_{x\\sim\\mu}[b(x)]}_{\\text{const., no }\\theta} \\end{aligned} \\]The objective is shifted by a constant value so the optimization problem isn\u0026rsquo;t changed.\nSo how should we pick \\(b(x)\\)? A function of \\(x\\) alone can only ever touch the between-context term, and it zeroes that term exactly when it matches the per-context mean \\(\\mathbb{E}[r\\mid x]\\). This is why the target is \\(b(x)\\approx\\mathbb{E}_{a\\sim\\pi_{\\theta}(x),\\,r\\sim R(x,a)}[r]\\).\n\\(b(x)\\) only reaches the between-context term Law of total variance, conditioning on the context:\n\\[ \\mathrm{Var}(r)=\\underbrace{\\mathbb{E}_{x}\\big[\\mathrm{Var}(r\\mid x)\\big]}_{\\text{within context}}\u0026#43;\\underbrace{\\mathrm{Var}_{x}\\big(\\mathbb{E}[r\\mid x]\\big)}_{\\text{between context}} \\]The second term is the context driven variance coming from the fact that different contexts have different typical rewards.\nNow redo the decomposition on the baseline-corrected reward:\n\\[ \\begin{aligned} \\mathrm{Var}\\big(r-b(x)\\big) \u0026amp;= \\mathbb{E}_{x}\\big[\\mathrm{Var}(r-b(x)\\mid x)\\big]\u0026#43;\\mathrm{Var}_{x}\\big(\\mathbb{E}[r-b(x)\\mid x]\\big) \u0026amp;\u0026amp; \\text{law of total variance} \\\\[6pt] \u0026amp;= \\underbrace{\\mathbb{E}_{x}\\big[\\mathrm{Var}(r\\mid x)\\big]}_{\\text{unchanged}}\u0026#43;\\mathrm{Var}_{x}\\big(\\mathbb{E}[r\\mid x]-b(x)\\big) \u0026amp;\u0026amp; b(x) \\text{ const. given } x \\end{aligned} \\]when conditioning on \\(x\\), \\(b(x)\\) becomes a constant and a constant never changes the variance. We can still change the value of the second term by adjusting \\(b(x)\\) so we have \\(b(x)\\) track \\(\\mathbb{E}[r\\mid x]\\) so that variance term goes to zero.\nOff-Policy Estimation Everything so far assumes that the data was collected by the same \\(\\pi_{\\theta}\\) that we are trying to train (on-policy). If instead the data was generated by a different policy \\(\\pi_{\\theta\u0026#39;}\\), the average \\(\\frac{1}{n}\\sum_{i}r_{i}\\) would estimate \\(J(\\theta\u0026#39;)\\) not \\(J(\\theta)\\). The fix is to reweight each sample by how much more or less likely \\(\\pi_{\\theta}\\) would have been to take that action.\nThe reweighting trick Fact 1.1 importance reweighting Let \\(\\rho,\\rho\u0026#39;\\) be distributions over \\(u\\), with coverage \\(\\rho\u0026#39;(u)\u0026gt;0\\) whenever \\(\\rho(u)\u0026gt;0\\). Then \\[\\mathbb{E}_{u\\sim\\rho\u0026#39;}\\!\\left[\\frac{\\rho(u)}{\\rho\u0026#39;(u)}\\,u\\right] = \\mathbb{E}_{u\\sim\\rho}[u]\\] Proof of Fact 1.1 \\[ \\begin{aligned} \\mathbb{E}_{u\\sim\\rho\u0026#39;}\\!\\left[\\frac{\\rho(u)}{\\rho\u0026#39;(u)}u\\right] \u0026amp;= \\int_u \\frac{\\rho(u)}{\\rho\u0026#39;(u)}\\,u\\;\\rho\u0026#39;(u)\\,du \u0026amp;\u0026amp; \\text{definition of }\\mathbb{E}_{u\\sim\\rho\u0026#39;} \\\\[6pt] \u0026amp;= \\int_u u\\,\\rho(u)\\,du \u0026amp;\u0026amp; \\rho\u0026#39;(u)\\text{ cancels} \\\\[6pt] \u0026amp;= \\mathbb{E}_{u\\sim\\rho}[u] \u0026amp;\u0026amp; \\text{definition of }\\mathbb{E}_{u\\sim\\rho} \\end{aligned} \\] Applying this to \\(J(\\theta)\\) with data from \\(\\pi_{\\theta\u0026#39;}\\), we get the importance-weighted estimator \\[ \\hat{J}_{\\theta}=\\frac{1}{n}\\sum_{i} \\frac{\\pi_{\\theta}(a_{i}\\mid x_{i})}{\\pi_{\\theta\u0026#39;}(a_{i}\\mid x_{i})} r_{i}=\\frac{1}{n}\\sum_{i}w_{i}r_{i},\\qquad \\mathbb{E}_{\\theta\u0026#39;}[\\hat{J}_{\\theta}]=J(\\theta) \\]Cost Though we have an unbiased estimate of \\(J(\\theta)\\) using data from a different policy, the cost is an increased variance.\nVariance of a reweighted reward \\[\\mathrm{Var}_{\\theta\u0026#39;}(w_i r_i)=\\underbrace{\\mathbb{E}_\\theta\\big[(w_i-1)r_i^2\\big]}_{\\text{off-policy penalty}}\u0026#43;\\underbrace{\\mathrm{Var}_\\theta(r_i)}_{\\text{on-policy variance}}\\] The second term is what you would have had on-policy. The first is the price of reweighting, and in RL it is usually positive.\nVariance decomposition \\[ \\begin{aligned} \\mathrm{Var}_{\\theta\u0026#39;}(w_i r_i) \u0026amp;= \\mathbb{E}_{\\theta\u0026#39;}\\big[w_i^2 r_i^2\\big]-\\big(\\mathbb{E}_{\\theta\u0026#39;}[w_ir_i]\\big)^2 \\\\[4pt] \u0026amp;= \\mathbb{E}_{\\theta\u0026#39;}\\big[w_i^2 r_i^2\\big]-J(\\theta)^2 \u0026amp;\u0026amp; \\mathbb{E}_{\\theta\u0026#39;}[w_ir_i]=J(\\theta) \\\\[6pt] \u0026amp;= \\mathbb{E}_{\\theta}\\big[w_i r_i^2\\big]-J(\\theta)^2 \u0026amp;\u0026amp; \\text{Fact 1.1 on }f=w_ir_i^2 \\\\[6pt] \u0026amp;= \\mathbb{E}_{\\theta}\\big[w_i r_i^2\\big]-\\mathbb{E}_\\theta\\big[r_i^2\\big]\u0026#43;\\mathbb{E}_\\theta\\big[r_i^2\\big]-J(\\theta)^2 \\\\[6pt] \u0026amp;= \\mathbb{E}_{\\theta}\\big[(w_i-1)r_i^2\\big]\u0026#43;\\mathrm{Var}_\\theta(r_i) \\end{aligned} \\] Policy Gradient Estimation tells us how good a policy is, now we will focus on how to update and improve policies. Like with supervised learning, we will use gradient steps to improve our policy, except we cannot differentiate through the black-box reward, so we will estimate the gradient with samples.\nImprovement with Gradients Gradients are useful because they give a local improvement direction: a small enough step along \\(\\nabla J(\\theta)\\) increases \\(J\\). This comes from Taylor\u0026rsquo;s theorem, which says that for a small step \\(u\\) the function is approximately linear \\[ J(\\theta\u0026#43;u)=J(\\theta)\u0026#43;\\nabla J(\\theta)^{\\top}u\u0026#43;\\underbrace{\\tfrac{1}{2}u^{\\top}\\nabla^{2}J(\\theta_{0})u\u0026#43;\\dots}_{\\text{negligible when }|u|\\le\\varepsilon} \\]so choosing the best small step is the constrained problem \\(\\max_{|u|\\le\\varepsilon}\\ J(\\theta)\u0026#43;\\nabla J(\\theta)^{\\top}u\\).\nThe solution to this is \\(u^\\star=\\alpha \\nabla J(\\theta)\\) with \\(\\alpha=\\frac{\\varepsilon}{\\mid \\nabla J(\\theta)\\mid}\\), this gives us the update \\[ \\theta \\leftarrow \\theta \u0026#43; \\alpha \\nabla J(\\theta) \\] Small step solves to \\(u^{\\star}=\\alpha\\nabla J(\\theta)\\) \\(J(\\theta)\\) is a constant with respect to \\(u\\), so we only need to maximize the linear term. Write \\(v=\\nabla J(\\theta)\\). \\[ \\begin{aligned} \\max_{\\Vert u\\Vert\\le\\varepsilon}\\; v^{\\top}u \u0026amp;= \\max_{\\Vert u\\Vert\\le\\varepsilon}\\; \\Vert v\\Vert\\,\\Vert u\\Vert\\cos\\angle(u,v) \u0026amp;\u0026amp; \\text{definition of dot product} \\\\[6pt] \u0026amp;= \\Vert v\\Vert\\cdot\\varepsilon\\cdot 1 \u0026amp;\u0026amp; \\text{take } \\Vert u\\Vert=\\varepsilon,\\ \\cos=1 \\\\[6pt] \u0026amp;\\Rightarrow\\; u^{\\star}=\\varepsilon\\,\\frac{v}{\\Vert v\\Vert} \u0026amp;\u0026amp; \\text{parallel to } v,\\ \\text{length } \\varepsilon \\end{aligned} \\]Any step of length \\(\\varepsilon\\) pointing along the gradient is the best first-order improvement. As long as \\(\\alpha\\) is small enough that the higher order terms are negligible, this is an improvement step.\nIn supervised learning the loss is known, so the gradient \\(\\nabla_{\\theta}\\ell(f_{\\theta}(x_{i}),y_{i})\\) is computed directly per data point and averaged. For a probabilistic model under cross-entropy loss \\(\\ell=-\\log f_{\\theta}(y\\mid x)\\) the update is \\[ \\theta\\leftarrow\\theta\u0026#43;\\alpha\\,\\frac{1}{n}\\sum_{i}\\nabla\\log f_{\\theta}(y_{i}\\mid x_{i}) \\] which pushes the model to make the observed labels more likely. Keep this in mind, REINFORCE will look almost identical.\nGradient Estimation We cannot compute \\(\\nabla J(\\theta)\\) directly because \\(R(x,a)\\) is unknown and we only see samples of it. We use two tricks to solve this, the reweighting trick from Fact 1.1, and the log derivative trick.\nLog-derivative trick \\[\\nabla f(\\theta)=f(\\theta)\\nabla\\log f(\\theta)\\] Follows from \\(\\frac{d}{dx}\\log x=1/x\\) and the chain rule.\nPolicy gradient Let \\(\\rho_{\\theta}(x,a)=\\mu(x)\\pi_{\\theta}(a\\mid x)\\) be the joint distribution over context and action. For any data collection policy \\(\\pi_{\\theta\u0026#39;}\\) with coverage, \\[\\nabla J(\\theta)=\\mathbb{E}_{x,a\\sim\\rho_{\\theta\u0026#39;}}\\left[\\frac{\\pi_{\\theta}(a\\mid x)}{\\pi_{\\theta\u0026#39;}(a\\mid x)}\\nabla\\log\\pi_{\\theta}(a\\mid x)\\mathbb{E}_{r\\sim R(x,a)}[r]\\right]\\] The gradient is an expectation over things we can sample, so a Monte Carlo estimate is immediate.\nPolicy gradient The only place \\(\\theta\\) appears in \\(J\\) is in the sampling distribution \\(\\rho_{\\theta}\\), which is exactly what makes it hard to differentiate. Reweight so the expectation is over a fixed \\(\\rho_{\\theta\u0026#39;}\\), then \\(\\theta\\) only appears inside the integrand. \\[ \\begin{aligned} \\nabla J(\\theta) \u0026amp;= \\nabla\\mathbb{E}_{x,a\\sim\\rho_{\\theta}}\\big[\\mathbb{E}_{r\\sim R(x,a)}[r]\\big] \u0026amp;\u0026amp; \\text{definition of } J \\\\[6pt] \u0026amp;= \\nabla\\mathbb{E}_{x,a\\sim\\rho_{\\theta\u0026#39;}}\\left[\\frac{\\rho_{\\theta}(x,a)}{\\rho_{\\theta\u0026#39;}(x,a)}\\mathbb{E}_{r}[r]\\right] \u0026amp;\u0026amp; \\text{Fact 1.1, needs } \\rho_{\\theta\u0026#39;}\u0026gt;0 \\text{ where } \\rho_{\\theta}\u0026gt;0 \\\\[6pt] \u0026amp;= \\mathbb{E}_{x,a\\sim\\rho_{\\theta\u0026#39;}}\\left[\\frac{\\nabla\\rho_{\\theta}(x,a)}{\\rho_{\\theta\u0026#39;}(x,a)}\\mathbb{E}_{r}[r]\\right] \u0026amp;\u0026amp; \\text{expectation no longer depends on } \\theta,\\ \\text{linearity} \\\\[6pt] \u0026amp;= \\mathbb{E}_{x,a\\sim\\rho_{\\theta\u0026#39;}}\\left[\\frac{\\rho_{\\theta}(x,a)}{\\rho_{\\theta\u0026#39;}(x,a)}\\nabla\\log\\rho_{\\theta}(x,a)\\mathbb{E}_{r}[r]\\right] \u0026amp;\u0026amp; \\text{log-derivative trick} \\\\[6pt] \u0026amp;= \\mathbb{E}_{x,a\\sim\\rho_{\\theta\u0026#39;}}\\left[\\frac{\\pi_{\\theta}(a\\mid x)}{\\pi_{\\theta\u0026#39;}(a\\mid x)}\\nabla\\log\\pi_{\\theta}(a\\mid x)\\mathbb{E}_{r}[r]\\right] \u0026amp;\u0026amp; \\star \\end{aligned} \\]\\(\\star\\) Since \\(\\rho_{\\theta}(x,a)=\\mu(x)\\pi_{\\theta}(a\\mid x)\\), the \\(\\mu(x)\\) cancels in the ratio, and \\(\\nabla\\log\\rho_{\\theta}=\\nabla\\log\\mu(x)\u0026#43;\\nabla\\log\\pi_{\\theta}(a\\mid x)=\\nabla\\log\\pi_{\\theta}(a\\mid x)\\) because \\(\\mu\\) does not depend on \\(\\theta\\). This is the whole point of the log-derivative trick: the unknown context distribution disappears.\nSetting \\(\\theta\u0026#39;=\\theta\\) (on-policy data) makes every importance weight equal to one.\nREINFORCE With samples \\((x_{i},a_{i},r_{i})\\) collected by \\(\\pi_{\\theta}\\), \\[\\theta\\leftarrow\\theta\u0026#43;\\alpha\\frac{1}{n}\\sum_{i}\\nabla\\log\\pi_{\\theta}(a_{i}\\mid x_{i})r_{i}\\] Compare with the cross-entropy update from the previous section: it is the same form, except each data point is weighted by its reward. The policy is pushed to make high-reward actions more likely, and (for negative rewards) low-reward actions less likely.\nIn the general off-policy case the estimate keeps the weights, \\[ \\widehat{\\nabla J}\u0026#39;_{\\theta}=\\frac{1}{n}\\sum_{i}\\underbrace{\\frac{\\pi_{\\theta}(a_{i}\\mid x_{i})}{\\pi_{\\theta\u0026#39;}(a_{i}\\mid x_{i})}}_{w_{i}}\\nabla\\log\\pi_{\\theta}(a_{i}\\mid x_{i})r_{i} \\] which increases the likelihood of actions that achieved high reward and were comparatively unlikely under \\(\\pi_{\\theta\u0026#39;}\\) (large \\(w_{i}\\)). Both contribute to high variance, especially over multiple updates: weights become positively correlated with reward, and the policy gets pushed into regions poorly covered by \\(\\pi_{\\theta\u0026#39;}\\). This is why taking many off-policy gradient steps on one batch does not work well, and it motivates the next section.\nBaselines carry over for free. In every estimate above we may replace \\(r_{i}\\) with \\(r_{i}-b(x_{i})\\) without introducing bias, since \\(\\nabla J_{b}=\\nabla J\\).\nTrust Regions The gradient ascent step was defined by \\(\\Vert u\\Vert\\le\\varepsilon\\), a Euclidean ball in parameter space. This is an arbitrary choice: two parameter vectors that are close in \\(\\ell_{2}\\) can induce very different action distributions, and vice versa. A sigmoid policy near saturation is the canonical example, a huge step in \\(\\theta\\) barely moves the probability, so ordinary policy gradient crawls exactly when the policy has found the right answer and wants to commit. We want \u0026ldquo;small step\u0026rdquo; to mean small in policy space.\nKL divergence \\[\\mathrm{KL}(\\rho\u0026#39;\\Vert\\rho)=\\mathbb{E}_{u\\sim\\rho\u0026#39;}\\left[\\log\\frac{\\rho\u0026#39;(u)}{\\rho(u)}\\right]\\] The surprise at seeing data from \\(\\rho\u0026#39;\\) when expecting \\(\\rho\\). Not a true distance (asymmetric, no triangle inequality) but always non-negative and zero only when the distributions coincide.\nWe measure the distance between two policies by KL averaged over contexts: \\[ d_{\\mathrm{KL}}(\\theta\u0026#39;,\\theta)=\\mathbb{E}_{x\\sim\\mu}\\big[\\mathrm{KL}(\\pi_{\\theta\u0026#39;}(\\cdot\\mid x)\\Vert\\pi_{\\theta}(\\cdot\\mid x))\\big]=\\mathbb{E}_{x,a\\sim\\rho_{\\theta\u0026#39;}}\\left[\\log\\frac{\\pi_{\\theta\u0026#39;}(a\\mid x)}{\\pi_{\\theta}(a\\mid x)}\\right] \\] This is zero at \\(\\theta\u0026#39;=\\theta\\), grows as the policies diverge, and does not care how the policies are parameterized.\nNatural Policy Gradient Same recipe as before: first order approximation to \\(J\\), but now the constraint is on \\(d_{\\mathrm{KL}}\\) instead of \\(\\Vert u\\Vert\\). Since \\(d_{\\mathrm{KL}}\\) is not quadratic in \\(u\\), we take its second order Taylor approximation. The zeroth and first order terms vanish, and the surviving quadratic term is defined by a matrix called the Fisher information.\nFisher information matrix \\[F(\\theta)=\\mathbb{E}_{x,a\\sim\\rho_{\\theta}}\\big[\\nabla\\log\\pi_{\\theta}(a\\mid x)\\nabla\\log\\pi_{\\theta}(a\\mid x)^{\\top}\\big]\\] and for small \\(u\\), \\(d_{\\mathrm{KL}}(\\theta,\\theta\u0026#43;u)\\approx u^{\\top}F(\\theta)u\\). Read \\(F\\) as a sensitivity meter: large in directions where wiggling \\(\\theta\\) changes the action distribution a lot, small where it barely matters.\nExpected score is zero: \\(\\mathbb{E}_{a\\sim\\pi_{\\theta}(x)}[\\nabla\\log\\pi_{\\theta}(a\\mid x)]=0\\) This identity is used twice (first-order term of the KL expansion, and the PPO gradient).\n\\[ \\begin{aligned} \\mathbb{E}_{a\\sim\\pi_{\\theta}(x)}\\big[\\nabla\\log\\pi_{\\theta}(a\\mid x)\\big] \u0026amp;= \\sum_{a}\\pi_{\\theta}(a\\mid x)\\nabla\\log\\pi_{\\theta}(a\\mid x) \u0026amp;\u0026amp; \\text{definition of } \\mathbb{E} \\\\[6pt] \u0026amp;= \\sum_{a}\\nabla\\pi_{\\theta}(a\\mid x) \u0026amp;\u0026amp; \\text{log-derivative trick in reverse} \\\\[6pt] \u0026amp;= \\nabla\\sum_{a}\\pi_{\\theta}(a\\mid x) \u0026amp;\u0026amp; \\text{linearity} \\\\[6pt] \u0026amp;= \\nabla 1 = 0 \u0026amp;\u0026amp; \\text{probabilities sum to one} \\end{aligned} \\] \\(d_{\\mathrm{KL}}(\\theta,\\theta\u0026#43;u)\\approx u^{\\top}F(\\theta)u\\) Expand \\(d_{\\mathrm{KL}}(\\theta,\\theta\u0026#43;u)=\\mathbb{E}_{x,a\\sim\\rho_{\\theta}}[\\log\\pi_{\\theta}(a\\mid x)-\\log\\pi_{\\theta\u0026#43;u}(a\\mid x)]\\) in \\(u\\) around \\(u=0\\). Note the sampling distribution \\(\\rho_{\\theta}\\) does not depend on \\(u\\), so gradients in \\(u\\) pass straight through the expectation. Zeroth order. \\(d_{\\mathrm{KL}}(\\theta,\\theta)=0\\).\nFirst order. \\[ \\begin{aligned} \\nabla_{u}d_{\\mathrm{KL}}(\\theta,\\theta\u0026#43;u) \u0026amp;= -\\mathbb{E}_{x,a\\sim\\rho_{\\theta}}\\big[\\nabla_{u}\\log\\pi_{\\theta\u0026#43;u}(a\\mid x)\\big] \u0026amp;\u0026amp; \\text{only the second log depends on } u \\\\[6pt] \u0026amp;\\overset{u=0}{=} -\\mathbb{E}_{x,a\\sim\\rho_{\\theta}}\\big[\\nabla_{\\theta}\\log\\pi_{\\theta}(a\\mid x)\\big] = 0 \u0026amp;\u0026amp; \\text{expected score is zero} \\end{aligned} \\] Alternatively: \\(d_{\\mathrm{KL}}\\ge 0\\) with equality at \\(u=0\\), so \\(u=0\\) is a minimum, hence a critical point.\nSecond order. Using \\(\\nabla^{2}\\log\\pi=\\dfrac{\\nabla^{2}\\pi}{\\pi}-\\nabla\\log\\pi\\nabla\\log\\pi^{\\top}\\), \\[ \\begin{aligned} \\nabla^{2}_{u}d_{\\mathrm{KL}}\\big|_{u=0} \u0026amp;= -\\mathbb{E}_{x,a\\sim\\rho_{\\theta}}\\big[\\nabla^{2}\\log\\pi_{\\theta}(a\\mid x)\\big] \\\\[6pt] \u0026amp;= -\\mathbb{E}_{x}\\Big[\\sum_{a}\\nabla^{2}\\pi_{\\theta}(a\\mid x)\\Big]\u0026#43;\\mathbb{E}_{x,a\\sim\\rho_{\\theta}}\\big[\\nabla\\log\\pi_{\\theta}\\nabla\\log\\pi_{\\theta}^{\\top}\\big] \u0026amp;\u0026amp; \\pi\\cdot\\tfrac{\\nabla^{2}\\pi}{\\pi}=\\nabla^{2}\\pi \\\\[6pt] \u0026amp;= -\\mathbb{E}_{x}\\big[\\nabla^{2}1\\big]\u0026#43;F(\\theta) = F(\\theta) \u0026amp;\u0026amp; \\text{probabilities sum to one} \\end{aligned} \\]So the expansion is \\(0\u0026#43;0\u0026#43;u^{\\top}F(\\theta)u\u0026#43;\\dots\\) (the notes absorb the \\(\\tfrac{1}{2}\\) into \\(\\varepsilon\\)).\nWith the quadratic approximation the KL-constrained step is \\[ \\max_{u}\\ \\nabla J(\\theta)^{\\top}u\\quad\\text{s.t.}\\quad u^{\\top}F(\\theta)u\\le\\varepsilon^{2} \\] Natural policy gradient The solution is \\(u^{\\star}=\\alpha F(\\theta)^{-1}\\nabla J(\\theta)\\) with \\(\\alpha=\\varepsilon\\big(\\nabla J(\\theta)^{\\top}F(\\theta)^{-1}\\nabla J(\\theta)\\big)^{-1/2}\\), giving \\[\\theta\\leftarrow\\theta\u0026#43;\\alpha F(\\theta)^{-1}\\nabla J(\\theta)\\] \\(F^{-1}\\) preconditions the gradient so equal steps in parameter space are equal steps in policy space. The update is (approximately) invariant to how \\(\\pi_{\\theta}\\) is parameterized.\nSolving the KL-constrained step The trick is a change of variables that turns the ellipsoid \\(u^{\\top}Fu\\le\\varepsilon^{2}\\) into a ball, after which it is the same problem as gradient ascent. Let \\(v=F^{1/2}u\\) and \\(c=F^{-1/2}\\nabla J(\\theta)\\). \\[ \\begin{aligned} \\nabla J(\\theta)^{\\top}u \u0026amp;= \\big(F^{1/2}c\\big)^{\\top}F^{-1/2}v = c^{\\top}v \u0026amp;\u0026amp; F^{1/2} \\text{ symmetric, cancels} \\\\[6pt] u^{\\top}Fu \u0026amp;= v^{\\top}v=\\Vert v\\Vert^{2} \u0026amp;\u0026amp; \\text{substitute } u=F^{-1/2}v \\\\[6pt] \\Rightarrow\\quad v^{\\star} \u0026amp;= \\varepsilon\\frac{c}{\\Vert c\\Vert} \u0026amp;\u0026amp; \\text{Euclidean problem from before} \\\\[6pt] \\Rightarrow\\quad u^{\\star} \u0026amp;= F^{-1/2}v^{\\star}=\\varepsilon\\frac{F^{-1}\\nabla J(\\theta)}{\\Vert F^{-1/2}\\nabla J(\\theta)\\Vert} \u0026amp;\u0026amp; \\text{undo change of variables} \\\\[6pt] \u0026amp;= \\varepsilon\\big(\\nabla J^{\\top}F^{-1}\\nabla J\\big)^{-1/2}F^{-1}\\nabla J(\\theta) \u0026amp;\u0026amp; \\Vert F^{-1/2}g\\Vert^{2}=g^{\\top}F^{-1}g \\end{aligned} \\] Why it fixes saturation (two-action sigmoid) For \\(\\pi(1)=\\sigma(\\theta)\\) with rewards \\(r_{1},r_{0}\\): \\(\\nabla J=\\sigma(1-\\sigma)(r_{1}-r_{0})\\) and \\(F=\\sigma(1-\\sigma)\\). Plain PG steps shrink like \\(\\sigma(1-\\sigma)\\to 0\\) as the policy commits, so reaching \\(\\pi(1)=1-\\epsilon\\) takes time exponential in \\(1/\\epsilon\\). NPG\u0026rsquo;s step is \\(F^{-1}\\nabla J=r_{1}-r_{0}\\), a constant, so it takes time logarithmic in \\(1/\\epsilon\\).\nIn practice everything is estimated from on-policy samples \\(x_{i},a_{i},r_{i}\\sim\\rho_{\\theta}\\): \\(\\nabla J\\) by REINFORCE, and \\[ \\hat{F}_{\\theta}=\\frac{1}{n}\\sum_{i}\\nabla\\log\\pi_{\\theta}(a_{i}\\mid x_{i})\\nabla\\log\\pi_{\\theta}(a_{i}\\mid x_{i})^{\\top} \\]Proximal Policy Optimization NPG gives a principled direction but requires inverting an estimate of \\(F(\\theta)\\), a (number of parameters)\\(^{2}\\) matrix. For deep networks with millions or billions of parameters this is out of the question. PPO keeps the idea that closeness should be measured in KL, but swaps the hard constraint for a soft penalty. Having collected data with \\(\\theta\u0026#39;\\), solve \\[ \\max_{\\theta}\\ J(\\theta)-\\lambda\\,d_{\\mathrm{KL}}(\\theta\u0026#39;,\\theta)=\\max_{\\theta}\\ J_{\\lambda}(\\theta;\\theta\u0026#39;) \\] Rather than a first or second order approximation, we take several ordinary gradient steps on this objective. Since the data came from \\(\\pi_{\\theta\u0026#39;}\\), the \\(\\nabla J\\) term uses the off-policy REINFORCE estimate.\nGradient of the KL penalty \\[ \\begin{aligned} \\nabla_{\\theta}d_{\\mathrm{KL}}(\\theta\u0026#39;,\\theta) \u0026amp;= \\nabla_{\\theta}\\mathbb{E}_{x,a\\sim\\rho_{\\theta\u0026#39;}}\\big[\\log\\pi_{\\theta\u0026#39;}(a\\mid x)-\\log\\pi_{\\theta}(a\\mid x)\\big] \u0026amp;\u0026amp; \\text{definition of } d_{\\mathrm{KL}} \\\\[6pt] \u0026amp;= -\\mathbb{E}_{x,a\\sim\\rho_{\\theta\u0026#39;}}\\big[\\nabla_{\\theta}\\log\\pi_{\\theta}(a\\mid x)\\big] \u0026amp;\u0026amp; \\rho_{\\theta\u0026#39;},\\ \\pi_{\\theta\u0026#39;} \\text{ do not depend on } \\theta \\end{aligned} \\]Note this is over \\(\\rho_{\\theta\u0026#39;}\\), not \\(\\rho_{\\theta}\\), so the expected-score identity does not make it vanish. It is zero only at \\(\\theta=\\theta\u0026#39;\\).\nPPO gradient Combining the off-policy policy gradient with \\(-\\lambda\\nabla d_{\\mathrm{KL}}\\) and factoring out \\(\\nabla\\log\\pi_{\\theta}\\), \\[\\nabla J_{\\lambda}(\\theta;\\theta\u0026#39;)=\\mathbb{E}_{x,a\\sim\\rho_{\\theta\u0026#39;}}\\left[\\nabla\\log\\pi_{\\theta}(a\\mid x)\\left(\\frac{\\pi_{\\theta}(a\\mid x)}{\\pi_{\\theta\u0026#39;}(a\\mid x)}\\mathbb{E}_{r}[r]\u0026#43;\\lambda\\right)\\right]\\] Estimated from data collected with \\(\\pi_{\\theta\u0026#39;}\\), with the usual baseline, \\[\\widehat{\\nabla J}\u0026#39;_{\\lambda,\\theta}=\\frac{1}{n}\\sum_{i}\\nabla\\log\\pi_{\\theta}(a_{i}\\mid x_{i})\\left(\\frac{\\pi_{\\theta}(a_{i}\\mid x_{i})}{\\pi_{\\theta\u0026#39;}(a_{i}\\mid x_{i})}\\big(r_{i}-b(x_{i})\\big)\u0026#43;\\lambda\\right)\\] Compare with off-policy REINFORCE, which has only the importance-weighted reward term. The penalty adds a uniform bonus \\(\\lambda\\) to every (context, action) pair in the dataset: every observed action gets its likelihood nudged up regardless of reward. This encourages \\(\\pi_{\\theta}\\) to imitate the data collection policy \\(\\pi_{\\theta\u0026#39;}\\), keeps the KL small, and mitigates the variance blowup when \\(\\theta\\) wanders into regions not covered by \\(\\theta\u0026#39;\\). Because this is an off-policy estimate with the penalty holding the policy in place, we can now take multiple steps on one batch.\nThe most widely used variant, PPO-clip, replaces the penalty with switching logic on the importance weights: discard a data point from the gradient estimate as soon as (i) \\(w_{i}\u0026gt;1\u0026#43;\\epsilon\\) for a positive instance \\(r_{i}-b(x_{i})\u0026gt;0\\), or (ii) \\(w_{i}\u0026lt;1-\\epsilon\\) for a negative instance. This caps the variance and stops the policy from continuing to update in low-coverage directions. Combining clipping with a group baseline (for each context, sample \\(G\\) actions and use their mean reward as \\(b(x_{i})\\), optionally dividing by the group standard deviation) gives GRPO, which is what is widely used for post-training language models.\n","permalink":"https://shauryasen.com/posts/policy-gradients-from-the-ground-up/","summary":"From Monte Carlo estimates to REINFORCE.","title":"Policy Gradients from the Ground Up"},{"content":"Click or drag anywhere in the plot. Each dot is one handwritten digit, compressed to a point in a 2D latent space.\nLatent space (click or drag) Decoder output z = (0.00, 0.00) Show decoded grid A VAE trained on MNIST. Each dot is a test digit's encoder mean $\\mu(x)$, colored by its true label. Clicking runs the decoder on that point in your browser, including the empty regions no training digit landed on. Show decoded grid draws the whole space at once. A normal autoencoder can\u0026rsquo;t do this: click in the gaps of its latent space and you get static. The rest of this post explains why a VAE\u0026rsquo;s latent space is smooth enough to wander around in.\nData manifolds The manifold hypothesis says that real data, the inputs that actually make sense to us, sit on a low-dimensional manifold inside the high-dimensional space they\u0026rsquo;re written in. A \\(28\\times 28\\) image lives in \\(\\mathbb{R}^{784}\\), but almost every point in \\(\\mathbb{R}^{784}\\) is static. The handwritten digits are a thin, curved sheet inside it.\nLocally, that sheet can be \u0026ldquo;unrolled\u0026rdquo;: near any point, a few directions are enough to move along it. So sampling uniformly from the ambient space almost never lands on a digit.\nThe way around this is to sample from a simple distribution we can sample from, like a Gaussian, and learn a map from it onto the data manifold. VAEs learn a smooth probabilistic version of that map. GANs learn a deterministic one through an adversarial game.\nDimensionality reduction The core of a VAE is compressing \\(x\\in\\mathbb{R}^D\\) into a latent \\(z\\in\\mathbb{R}^d\\) with \\(d\\ll D\\). A good compression should:\nkeep the structure that matters, so similar inputs get nearby codes, keep enough information to reconstruct \\(x\\) from \\(z\\), and for a VAE specifically, make the latent space easy to sample from. PCA PCA is the linear version. It assumes the data lies near a linear subspace and picks the projection \\(W\\) that keeps the most variance. That works for feature extraction, but real manifolds curve, and a linear projection flattens the curve.\nAutoencoders An autoencoder does the nonlinear version with two networks:\nan encoder \\(e_\\phi(x)\\) that maps \\(x\\in\\mathbb{R}^D\\) to \\(z\\in\\mathbb{R}^d\\), and a decoder \\(d_\\theta(z)\\) that maps back to a reconstruction \\(x\u0026#39;\\approx x\\), trained to minimize reconstruction error:\n\\[ \\min_{\\phi,\\theta}\\ \\sum_{x\\in\\mathcal{D}} \\Vert x-x\u0026#39;\\Vert^2, \\qquad x\u0026#39;=d_\\theta(e_\\phi(x)). \\]Why you can\u0026rsquo;t sample from an autoencoder To generate new data, we\u0026rsquo;d like to draw \\(z\\sim p(z)\\) and decode it. But nothing in the autoencoder\u0026rsquo;s loss says what \\(p(z)\\) is. Training only cares about the codes of the training points, so the latent space ends up:\nsparse: codes cluster wherever is convenient, with big gaps between them, and unstructured: there\u0026rsquo;s no known shape to sample from, and points in the gaps decode to garbage. A VAE fixes both by making the encoder output a distribution and pinning those distributions to a known prior.\nKL divergence The tool for \u0026ldquo;pinning one distribution to another\u0026rdquo; is the KL divergence:\n\\[ D_{\\mathrm{KL}}(p\\,\\Vert\\,q)=\\mathbb{E}_{x\\sim p}\\left[\\log\\frac{p(x)}{q(x)}\\right]. \\]One reading is the expected extra surprise from modeling data from \\(p\\) with \\(q\\). It\u0026rsquo;s always \\(\\ge 0\\), it\u0026rsquo;s \\(0\\) only when \\(p=q\\), and it isn\u0026rsquo;t symmetric. The direction matters:\nForward KL, \\(D_{\\mathrm{KL}}(p\\Vert q)\\), is mass-covering. The expectation is over \\(p\\), so wherever \\(p\\) has mass and \\(q\\approx 0\\), \\(\\log\\frac{p}{q}\\) explodes. \\(q\\) is forced to cover all of \\(p\\), even if that means spreading thin. Reverse KL, \\(D_{\\mathrm{KL}}(q\\Vert p)\\), is mode-seeking. Now the expectation is over \\(q\\), so the penalty is for \\(q\\) putting mass where \\(p\\) has none. \\(q\\) stays tucked under \\(p\\), often inside a single mode. The VAE objective below uses \\(D_{\\mathrm{KL}}(q_\\phi(z\\mid x)\\,\\Vert\\,p(z))\\).\nMaking the autoencoder probabilistic Probabilistic encoder Instead of one code per input, the encoder outputs a Gaussian over codes:\n\\[ q_\\phi(z\\mid x)=\\mathcal{N}\\big(\\mu_\\phi(x),\\ \\mathrm{diag}(\\sigma^2_\\phi(x))\\big). \\]In practice the network predicts \\(\\mu\\) and \\(\\log\\sigma^2\\), since the log keeps the variance positive without a constraint.\nProbabilistic decoder The decoder defines \\(p_\\theta(x\\mid z)\\). With a Gaussian decoder of fixed variance, maximizing \\(\\log p_\\theta(x\\mid z)\\) is the same as minimizing squared error:\n\\[ \\max_{\\phi,\\theta}\\ \\mathbb{E}_{z\\sim q_\\phi(z\\mid x)}\\big[\\log p_\\theta(x\\mid z)\\big] \\quad\\Longleftrightarrow\\quad \\min_{\\phi,\\theta}\\ \\mathbb{E}_{z\\sim q_\\phi(z\\mid x)}\\big[\\Vert x-d_\\theta(z)\\Vert^2\\big]. \\]For MNIST I used a Bernoulli decoder instead, which treats each pixel as a coin flip. That makes the reconstruction term binary cross-entropy.\nKeeping every \\(q_\\phi(z\\mid x)\\) near one prior Each input now gets its own little Gaussian. If those Gaussians drift anywhere they like, we\u0026rsquo;re back to a sparse, unstructured latent space. So the VAE also penalizes each one for straying from a shared prior, \\(p(z)=\\mathcal{N}(0,I)\\):\n\\[ D_{\\mathrm{KL}}\\big(q_\\phi(z\\mid x)\\,\\Vert\\,p(z)\\big). \\]That\u0026rsquo;s why the cloud in the demo sits roughly inside the unit-Gaussian region and has no gaping holes.\nThe evidence lower bound What we\u0026rsquo;d actually like to maximize is the likelihood of the data, \\(\\log p_\\theta(x)=\\log\\int p_\\theta(x\\mid z)\\,p(z)\\,dz\\). That integral is over every possible \\(z\\), so it\u0026rsquo;s intractable. Instead we maximize a lower bound on it, the ELBO. The idea is to push a floor up as high as it will go; the quantity above it has to rise with it.\nDeriving the ELBO Start from the log-evidence and bring in \\(q_\\phi(z\\mid x)\\), which integrates to one:\n\\[ \\begin{aligned} \\log p_\\theta(x) \u0026amp;= \\int q_\\phi(z\\mid x)\\,\\log p_\\theta(x)\\,dz \u0026amp;\u0026amp; \\textstyle\\int q_\\phi = 1 \\\\[6pt] \u0026amp;= \\mathbb{E}_{q_\\phi}\\left[\\log\\frac{p_\\theta(x,z)}{p_\\theta(z\\mid x)}\\right] \u0026amp;\u0026amp; \\text{Bayes: } p(x)=\\tfrac{p(x,z)}{p(z\\mid x)} \\\\[6pt] \u0026amp;= \\mathbb{E}_{q_\\phi}\\left[\\log\\frac{p_\\theta(x,z)}{q_\\phi(z\\mid x)}\\cdot\\frac{q_\\phi(z\\mid x)}{p_\\theta(z\\mid x)}\\right] \u0026amp;\u0026amp; \\text{multiply and divide by } q_\\phi \\\\[6pt] \u0026amp;= \\underbrace{\\mathbb{E}_{q_\\phi}\\left[\\log\\frac{p_\\theta(x,z)}{q_\\phi(z\\mid x)}\\right]}_{\\text{ELBO}} \u0026#43; \\underbrace{D_{\\mathrm{KL}}\\big(q_\\phi(z\\mid x)\\,\\Vert\\,p_\\theta(z\\mid x)\\big)}_{\\ge\\,0} \u0026amp;\u0026amp; \\text{split the log} \\end{aligned} \\]The KL term is non-negative, so \\(\\log p_\\theta(x)\\ge\\text{ELBO}\\). The gap is exactly how far the encoder is from the true posterior. Raising the ELBO either raises the likelihood or tightens the gap, and both are things we want.\nExpanding \\(p_\\theta(x,z)=p_\\theta(x\\mid z)\\,p(z)\\) splits the ELBO into the two terms you actually implement:\nThe VAE objective \\[ \\text{ELBO}(x)=\\underbrace{\\mathbb{E}_{z\\sim q_\\phi(z\\mid x)}\\big[\\log p_\\theta(x\\mid z)\\big]}_{\\text{reconstruction}}-\\underbrace{D_{\\mathrm{KL}}\\big(q_\\phi(z\\mid x)\\,\\Vert\\,p(z)\\big)}_{\\text{stay near the prior}} \\]The first term asks the decoder to rebuild \\(x\\) from a sampled code. The second keeps every code distribution close to \\(\\mathcal{N}(0,I)\\). They pull in opposite directions: perfect reconstructions want tight, well-separated \\(q\\)\u0026rsquo;s, and the prior wants them all to be the same blob.\nFor a diagonal Gaussian encoder and a standard normal prior, the KL term has a closed form, so no sampling is needed for it:\nClosed-form KL for diagonal Gaussians Per latent dimension, with \\(q=\\mathcal{N}(\\mu,\\sigma^2)\\) and \\(p=\\mathcal{N}(0,1)\\):\n\\[ \\begin{aligned} D_{\\mathrm{KL}}(q\\Vert p) \u0026amp;= \\mathbb{E}_q\\big[\\log q(z)-\\log p(z)\\big] \\\\[6pt] \u0026amp;= \\mathbb{E}_q\\left[-\\tfrac12\\log\\sigma^2-\\frac{(z-\\mu)^2}{2\\sigma^2}\u0026#43;\\frac{z^2}{2}\\right] \u0026amp;\u0026amp; \\text{the } \\log\\sqrt{2\\pi} \\text{ terms cancel} \\\\[6pt] \u0026amp;= -\\tfrac12\\log\\sigma^2-\\tfrac12\u0026#43;\\tfrac12(\\sigma^2\u0026#43;\\mu^2) \u0026amp;\u0026amp; \\mathbb{E}_q[(z-\\mu)^2]=\\sigma^2,\\ \\mathbb{E}_q[z^2]=\\sigma^2\u0026#43;\\mu^2 \\\\[6pt] \u0026amp;= -\\tfrac12\\big(1\u0026#43;\\log\\sigma^2-\\mu^2-\\sigma^2\\big) \\end{aligned} \\]Sum over the \\(d\\) dimensions. This is the kl line in the code below.\nThe reparameterization trick The reconstruction term is an expectation over \\(z\\sim q_\\phi(z\\mid x)\\), and we need its gradient with respect to \\(\\phi\\). You can\u0026rsquo;t backpropagate through \u0026ldquo;draw a random sample.\u0026rdquo; The trick is to move the randomness outside the network:\n\\[ z=\\mu_\\phi(x)\u0026#43;\\sigma_\\phi(x)\\odot\\epsilon,\\qquad \\epsilon\\sim\\mathcal{N}(0,I). \\]Now \\(z\\) is a deterministic, differentiable function of \\(\\mu\\) and \\(\\sigma\\), plus noise that doesn\u0026rsquo;t depend on \\(\\phi\\). Gradients flow straight through. This is also much lower variance than the score-function (REINFORCE-style) estimator you\u0026rsquo;d otherwise need.\nThe implementation This is the model and loss behind the demo at the top. The full script, including the export to the browser, is train_vae.py.\nclass VAE(nn.Module): def __init__(self): super().__init__() self.encoder = nn.Sequential( nn.Linear(784, 512), nn.ReLU(), nn.Linear(512, 256), nn.ReLU(), ) self.to_mu = nn.Linear(256, 2) self.to_logvar = nn.Linear(256, 2) self.decoder = nn.Sequential( nn.Linear(2, 64), nn.ReLU(), nn.Linear(64, 256), nn.ReLU(), nn.Linear(256, 784), # logits ) def forward(self, x): h = self.encoder(x) mu, logvar = self.to_mu(h), self.to_logvar(h) z = mu + torch.exp(0.5 * logvar) * torch.randn_like(mu) # reparameterization return self.decoder(z), mu, logvar def neg_elbo(logits, x, mu, logvar): recon = F.binary_cross_entropy_with_logits(logits, x, reduction=\u0026#34;sum\u0026#34;) # -log p(x|z) kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) # KL(q(z|x) || N(0,I)) return recon + kl Trained with Adam (learning rate \\(10^{-3}\\), batch size 128) for 40 epochs, which took a couple of minutes on a laptop. It ends at a −ELBO of about 138 nats per image, roughly 132 of reconstruction and 7 of KL. The decoder is about 218k parameters. Stored as float16 that\u0026rsquo;s 426 KB, which the page loads and runs in plain JavaScript.\nCredit: these notes follow Cornell\u0026rsquo;s CS 4782: Deep Learning, which I took in Spring 2026, especially the course\u0026rsquo;s U-Nets and VAEs notes. The interactive demo here is my own reimplementation of the idea behind the course\u0026rsquo;s VAE latent-space demo, which is worth a look too. Mistakes are mine.\n","permalink":"https://shauryasen.com/posts/variational-autoencoders/","summary":"From Autoencoder to VAE trained on MNIST","title":"Variational Autoencoders from the Ground Up"}]