diff --git a/mjx/training_apg.ipynb b/mjx/training_apg.ipynb
index bb7d35ed..baf44ccd 100644
--- a/mjx/training_apg.ipynb
+++ b/mjx/training_apg.ipynb
@@ -813,7 +813,7 @@
    "source": [
     "## Study 2: Quadruped Locomotion\n",
     "\n",
-    "As we saw in the imitation learning example, FoPG methods benefit from detailed reward signals. To teach locomotion, we reward the feet based on the Raibert Heuristic. We implement a gait schedule, that incentivizes opposite pairs of legs to move in sync at a fixed frequency. At the beginning of a new scheduled step, we calculate the target position for the feet at the end of the step. \n",
+    "As we saw in the imitation learning example, FoPG methods benefit from detailed reward signals. To teach locomotion, we reward the feet based on the Raibert Heuristic. Similarly to [prior work](https://arxiv.org/abs/2403.14864), we use a gait schedule to incentivize opposite pairs of legs to move in sync at a fixed frequency. At the beginning of a new scheduled step, we calculate the target position for the feet at the end of the step. \n",
     "\n",
     "For each foot, we calculate:\n",
     "\n",