diff --git a/mjx/training_apg.ipynb b/mjx/training_apg.ipynb
index 512e6329..bb7d35ed 100644
--- a/mjx/training_apg.ipynb
+++ b/mjx/training_apg.ipynb
@@ -559,7 +559,7 @@
     "    obs_list.append(state_info['last_action'])\n",
     "    # kinematic reference\n",
     "    kin_ref = self.kinematic_ref_qpos[jp.array(state_info['steps']%self.l_cycle, int)]\n",
-    "    obs_list.append(kin_ref[7:]) # Gait schedule, per actuator.\n",
+    "    obs_list.append(kin_ref[7:]) # First 7 indicies are fixed\n",
     "\n",
     "    obs = jp.clip(jp.concatenate(obs_list), -100.0, 100.0)\n",
     "\n",
@@ -1213,9 +1213,9 @@
     "    hidden_layer_sizes=(256, 128)\n",
     ")\n",
     "\n",
-    "nets = make_networks_factory(observation_size=47,\n",
-    "                                action_size=12,\n",
-    "                                preprocess_observations_fn=running_statistics.normalize)\n",
+    "nets = make_networks_factory(observation_size=1, # Observation_size argument doesn't matter since it's only used for param init.\n",
+    "                             action_size=12,\n",
+    "                             preprocess_observations_fn=running_statistics.normalize)\n",
     "\n",
     "make_inference_fn = apg_networks.make_inference_fn(nets)\n",
     "\n",