diff --git a/mjx/training_apg.ipynb b/mjx/training_apg.ipynb index eb7f1b6a..7d4cb990 100644 --- a/mjx/training_apg.ipynb +++ b/mjx/training_apg.ipynb @@ -16,7 +16,7 @@ "\n", "#### Zeroth-Order Policy Gradients (ZoPG)\n", "\n", - "Referring to mjx.step as the simulation function f, we borrowing some [terminology](https://arxiv.org/abs/2202.00817) to differentiate between zeroth-order gradients, which only depend on values of f, and first-order gradients, which depend on its jacobian.\n", + "Referring to `mjx.step` as the simulation function f, we borrow some [terminology](https://arxiv.org/abs/2202.00817) to differentiate between zeroth-order gradients, which only depend on values of f, and first-order gradients, which depend on its jacobian.\n", "\n", "Reinforcement learning (RL) algorithms such as the standard [PPO](https://github.com/google/brax/blob/main/brax/training/agents/ppo/train.py) assume the stochastic state transition model $x_{t+1} \\sim P(\\cdot | x_t, a_t)$. This leads to a ZoPG of the form:\n", "\n", @@ -27,22 +27,22 @@ "Where $R(\\tau)$ is some function depending on the rollout $\\tau = \\{x_t, a_t\\}_{t=0}^{T}$. Despite this method's popularity and extensive research into its refinement, a fundamental property is that the gradient has high variance. This allows the optimizer to thoroughly explore the space of policies, leading to the robust and often surprisingly good policies that have been achieved. However, the variance comes at the cost of requiring many samples $(x_t, a_t)$ to converge.\n", "\n", "#### First-Order Policy Gradients (FoPG)\n", - "On the other hand, if you assume a deterministic state transition model $x_{t+1} = f(x_t, a_t)$, you end up with the first-order policy gradient. Other common names include Analytical Policy Gradients (APG) and Backpropogation through Time (BPTT). Unlike ZoPG methods, which model the state evolution as a probabilistic black box, the FoPG explicitly contains the jacobians of the simulation function f. For example, let's look at the gradient of $r_t$, in the case that it only depends on state.\n", + "On the other hand, if you assume a deterministic state transition model $x_{t+1} = f(x_t, a_t)$, you end up with the first-order policy gradient. Other common names include Analytical Policy Gradients (APG) and Backpropogation through Time (BPTT). Unlike ZoPG methods, which model the state evolution as a probabilistic black box, the FoPG explicitly contains the jacobians of the simulation function f. For example, let's look at the gradient of the reward $r_t$, in the case that it only depends on state.\n", "$$\n", "\\frac{\\partial r_t}{\\partial \\theta} = \\frac{\\partial r_t}{\\partial x_t}\\frac{\\partial x_t}{\\partial \\theta} \n", "$$\n", "\n", "$$\n", - "\\frac{\\partial x_t}{\\partial \\theta} = \\textcolor{Navy}{\\frac{\\partial f(x_t, a_t)}{\\partial x_{t-1}}}\\frac{\\partial x_{t-1}}{\\partial \\theta} + \\textcolor{Navy}{\\frac{\\partial f(x_t, a_t)}{\\partial a_{t-1}}} \\frac{\\partial a_{t-1}}{\\partial \\theta}\n", + "\\frac{\\partial x_t}{\\partial \\theta} = \\textcolor{Navy}{\\frac{\\partial f(x_{t-1}, a_{t-1})}{\\partial x_{t-1}}}\\frac{\\partial x_{t-1}}{\\partial \\theta} + \\textcolor{Navy}{\\frac{\\partial f(x_{t-1}, a_{t-1})}{\\partial a_{t-1}}} \\frac{\\partial a_{t-1}}{\\partial \\theta}\n", "$$\n", "\n", "The navy-colored terms in the above expression are enabled by MJX's differentiability and are the key difference between FoPG's and ZoPG's. An important consideration is what these jacobians look like near contact points. To see why certain gradients within the jacobian can be pathological, imagine a hard sphere falling toward a block of marble. How does its velocity change with respect to distance ($\\frac{\\partial \\dot{z}_t}{\\partial z_t}$), the instant before it touches the ground? This is the case of an **uninformative gradient**, due to [hard contact](https://arxiv.org/html/2404.02887v1). Fortunately, the default contact settings in Mujoco are sufficiently [soft](https://mujoco.readthedocs.io/en/stable/computation/index.html#soft-contact-model) for learning via FoPG's. With soft contacts, the ground applies an increasing force on the ball as it penetrates it, unlike rigid contacts, which instantly provide enough force for deflection.\n", "\n", - "A helpful way to think about FoPG's is via the chain rule, as illustrated below for how $r_2$ influences the policy gradient, again for the case that the reward does not depend on action:\n", + "A helpful way to think about FoPG's is via the chain rule and computation graphs, as illustrated below for how $r_2$ influences the policy gradient, again for the case that the reward does not depend on action:\n", "\n", "\"drawing\"\n", "\n", - "Note that there three distinct gradient chains in this example. The red pathway considers how the immediately prior action affected the state. The blue path explains the name *Backpropogation through Time*, capturing how actions affect downstream rewards. The least intuitive may be the green chain, which shows how the reward depends on how actions depend on previous actions - experience shows that blocking this pathway via jax.lax.stop_grad can badly hinder policy learning. As the length of $x_t$ backbone increases, [gradient explosion](https://arxiv.org/abs/2111.05803) becomes a crucial consideration. In practice, this can be resolved via decaying downstream gradients or periodically truncating the gradient.\n", + "Note that there three distinct gradient chains in this example. The red pathway considers how the immediately prior action affected the state. The blue path explains the name *Backpropogation through Time*, capturing how actions affect downstream rewards. The least intuitive may be the green chain, which shows how the reward depends on how actions depend on previous actions.Experience shows that blocking *any* of these three pathways via jax.lax.stop_grad can badly hinder policy learning. As the length of $x_t$ backbone increases, [gradient explosion](https://arxiv.org/abs/2111.05803) becomes a crucial consideration. In practice, this can be resolved via decaying downstream gradients or periodically truncating the gradient.\n", "\n", "**The Sharp Bits of FoPG's**\n", "\n", @@ -50,7 +50,7 @@ "\n", "Additionally, discontinuous reward formulations are ubiquitious in RL, for instance, a large penalty when the robot falls. It can be significantly more [challenging](https://arxiv.org/abs/2403.14864) to design robust policies with FoPG's, since they cannot backprop through such penalties.\n", "\n", - "Last, despite the sample efficiency, FoPG methods can still struggle with wall-clock time. Because the gradients have low variance, they do not benefit significantly from massive parallelization of data collection - unlike [RL](https://arxiv.org/abs/2109.11978). Additionally, the policy gradient is typically calculated via autodifferentiation. This can be 3-5x slower than unrolling the simulation forward, and memory intensive, with memory requirements scaling with $O((m+n) \\cdot m \\cdot T)$, where m and n are the state and control dimensions, $(m+n) \\cdot m$ is the jacobian dimension, and T is the number of steps propogated through.\n", + "Last, despite the sample efficiency, FoPG methods can still struggle with wall-clock time. Because the gradients have low variance, they do not benefit significantly from massive parallelization of data collection - unlike [RL](https://arxiv.org/abs/2109.11978). Additionally, the policy gradient is typically calculated via autodifferentiation. This can be 3-5x slower than unrolling the simulation forward, and memory intensive, with memory requirements scaling with $O(m \\cdot (m+n) \\cdot T)$, where m and n are the state and control dimensions, $m \\cdot (m+n)$ is the jacobian dimension, and T is the number of steps propogated through.\n", "\n", "Note that with certain models, using autodifferentiation through mjx.step currently causes [nan gradients](https://github.com/google-deepmind/mujoco/issues/1517). For now, we address this issue by using double-precision floats, at the cost of doubling the memory requirements and training time." ] @@ -68,9 +68,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "## Setup: Imports and installations\n", - "\n", - "The examples below require a modified version of Brax's Analytical Policy Gradient (APG) algorithm." + "## Setup: Imports and installations" ] }, {