chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update text

Browse Source
This commit is contained in:
Andrew
2024-04-17 11:59:09 +02:00
parent 7a7c33a23b
commit 9b2274593d
1 changed files with 17 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/training_apg.ipynb
+17 -11
View File
@@ -36,15 +36,15 @@
"\\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",
"$$\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**. 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",
"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",
"\n",
"<img src=\"../doc/images/mjx/apg_diagram.png\" alt=\"drawing\" width=\"300\"/>\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 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",
"\n",
"**The Sharp bits of FoPG's**\n",
"**The Sharp Bits of FoPG's**\n",
"\n",
"While FoPG's have been shown to be very sample efficient, especially as the [dimension of the state space increases](https://arxiv.org/abs/2204.07137), one fundamental shortcoming is that due to the lower gradient variance, FoPG's also have less exploration power than ZoPG's and benefit from the practioner being more explicit in the problem formulation.\n",
"\n",
@@ -59,7 +59,9 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"In this tutorial, we demonstrate two ways to use FoPG's, using a simple variation of Brax's APG [algorithm](https://github.com/Andrew-Luo1/brax_new_apg/tree/main/brax/training/agents/apg). This algorithm essentially performs live stochastic gradient descent on the policy, unrolling it for a short window, doing a gradient update, then continuing where it left off."
"**Coming Up**\n",
"\n",
"In this tutorial, we demonstrate two ways to use FoPG's, using Brax's simple APG [algorithm](https://github.com/google/brax/tree/main/brax/training/agents/apg). This algorithm uses FoPG's to essentially perform live stochastic gradient descent on the policy, unrolling it for a short window, doing a policy update, then continuing where it left off."
]
},
{
@@ -727,7 +729,7 @@
"source": [
"**A note on sample efficiency**\n",
"\n",
"Above, we train using epochs * horizon_length * num_envs = 1.024e6 total simulator steps. Let's compare with PPO, using ten times more samples:"
"Above, we train using epochs * horizon_length * num_envs = 1.024e6 total simulator steps. Let's see what we get from PPO, using ten times more samples:"
]
},
{
@@ -785,7 +787,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"We see that PPO takes around 9e6 simulator steps to catch up to APG!"
"We see that PPO takes around 9e6 simulator steps to catch up to APG."
]
},
{
@@ -802,17 +804,17 @@
"p^* = h_0 + \\frac{\\Delta T}{2} v_0\n",
"$$\n",
"\n",
"Where $p^*$ is x, y component of the foot position target, $h_0$ is the x, y component of the corresponding hip at lift-off, $\\Delta T$ is the scheduled step duration, and $v_0$ is the base velocity at lift-off. \n",
"Where $p^*$ is x, y component of the foot's target position, $h_0$ is the x, y component of the corresponding hip at lift-off, $\\Delta T$ is the scheduled step duration, and $v_0$ is the base velocity at lift-off. \n",
"\n",
"Due to their limited exploration power, FoPG methods benefit greatly from having an good \"initial guess\" of the policy - familiar terminology in Model Predictive Control and Trajectory Optimization. We formulate the problem as [residual learning](https://arxiv.org/abs/1512.03385). Let $\\phi$ be the parameters of a baseline policy that we assume we already have, and let $f$ and $g$ be the neural networks for the learned and baseline policy. We freeze $\\phi$ and learn parameters $\\theta$, for the policy:\n",
"Due to their limited exploration power, FoPG methods benefit greatly from having an good \"initial guess\" of the policy - familiar terminology in Model Predictive Control and Trajectory Optimization. We formulate the problem as [residual learning](https://arxiv.org/abs/1512.03385). Let $\\phi$ be the parameters of a baseline policy that we already have, and let $f$ and $g$ be the neural networks for the learned and baseline policy. We freeze $\\phi$ and learn parameters $\\theta$, for the policy:\n",
"\n",
"$$\n",
"a_t = f(g(x_t; \\phi), x_t; \\theta) + g(x_t; \\phi)\n",
"$$\n",
"\n",
"We use the in-place trotting policy from last section as $\\phi$ and track a 0.75 m/s velocity target. Since locomotion is stabler at faster trots, you can experiment with $\\phi$ for faster velocity targets!\n",
"We use the in-place trotting policy from last section as $\\phi$ and $x_t$ denotes the state at time t. In this example we track a 0.75 m/s velocity target, but since locomotion is stabler at faster trots, you can experiment with $\\phi$ for faster velocity targets!\n",
"\n",
"While a more natural way to \"hotstart\" the learning would be to simply initialize the parameters $\\theta$ as $\\phi$ and use the policy $a_t = f(x_t; \\theta)$, the residual method has been found train more stably in practice."
"While it might seem more natural to \"hotstart\" the learning by simply initializing the parameters $\\theta$ as $\\phi$, with the policy $a_t = f(x_t; \\theta)$, we find that the residual method trains more stably in practice."
]
},
{
@@ -1380,7 +1382,11 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"We see that PPO struggles to learn locomotion in this setup, even with 10x the number of simulator steps."
"We see that PPO struggles to learn locomotion in this setup, even with over 10x the number of simulator steps. \n",
"\n",
"This doesn't mean that PPO is a bad algorithm! Rather, this study demonstrates one way we can *design policy-learning setups that leverage FoPG methods*. FoPG's perform well in this setup because it involves learning small, accurate perturbations from a good baseline - optimizing to the local minima in a deep valley. FoPG's are precise enough to guide these perturbations using nuanced reward signals, such as our foot placement spline.\n",
"\n",
"In contrast, RL algorithms such as PPO benefit from [policy-learning setups](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb) that have less structured rewards. Unlike FoPG methods, they [benefit greatly](https://www.science.org/doi/abs/10.1126/scirobotics.adg1462) from sparse, non-differentiable rewards such as a large penalty for falling."
]
}
],