Merge pull request #189 from youwyu:main

PiperOrigin-RevId: 798364011
Change-Id: I5962ac3de20148cbd952c6e3676ef0f5b8a2b8f3

Author: copybara-github

mujoco_playground/experimental/brax_network_to_onnx.ipynb

@@ -37,7 +37,8 @@
     "import tensorflow as tf\n",
     "from tensorflow.keras import layers\n",
     "import onnxruntime as rt\n",
-    "from brax.training.acme import running_statistics"
+    "from brax.training.acme import running_statistics\n",
+    "from brax.training.checkpoint import load"
    ]
   },
   {
@@ -108,42 +109,20 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# ckpt_path = \"/home/kevin/mujoco_playground/mujoco_playground/experimental/learning/checkpoints/G1JoystickFlatTerrain-20250109-015111\"  # g1_fulljntrange\n",
-    "# ckpt_path = \"/home/kevin/mujoco_playground/mujoco_playground/experimental/learning/checkpoints/G1JoystickFlatTerrain-20250109-021034\"  # g1_fulljntrange_energy\n",
-    "# ckpt_path = \"/home/kevin/mujoco_playground/mujoco_playground/experimental/learning/checkpoints/G1JoystickFlatTerrain-20250109-022739\"  # g1_fulljntrange_energy_biggerangvel (bad)\n",
-    "# ckpt_path = \"/home/kevin/mujoco_playground/mujoco_playground/experimental/learning/checkpoints/G1JoystickFlatTerrain-20250109-031036\"  # g1_fulljntrange_energy_biggerangvel_again (redo)\n",
-    "ckpt_path = \"/home/kevin/mujoco_playground/mujoco_playground/experimental/learning/checkpoints/G1JoystickFlatTerrain-20250109-040028\"\n",
-    "ckpt_path = \"/home/kevin/mujoco_playground/mujoco_playground/experimental/learning/checkpoints/G1JoystickFlatTerrain-20250109-041537\"\n",
-    "ckpt_path = \"/home/kevin/mujoco_playground/mujoco_playground/experimental/learning/checkpoints/G1JoystickFlatTerrain-20250109-043559\"\n",
-    "ckpt_path = \"/home/kevin/mujoco_playground/mujoco_playground/experimental/learning/checkpoints/G1JoystickFlatTerrain-20250109-045516\"\n",
-    "ckpt_path = \"/home/kevin/mujoco_playground/mujoco_playground/experimental/learning/checkpoints/G1JoystickFlatTerrain-20250109-052430\"\n",
-    "ckpt_path = \"/home/kevin/mujoco_playground/mujoco_playground/experimental/learning/checkpoints/G1JoystickFlatTerrain-20250109-135615\"\n",
-    "ckpt_path = \"/home/kevin/mujoco_playground/mujoco_playground/experimental/learning/checkpoints/G1JoystickFlatTerrain-20250109-152932\"\n",
-    "ckpt_path = \"/home/kevin/mujoco_playground/mujoco_playground/experimental/learning/checkpoints/G1JoystickFlatTerrain-20250109-185338\"\n",
-    "ckpt_path = \"/home/kevin/mujoco_playground/mujoco_playground/experimental/learning/checkpoints/G1JoystickRoughTerrain-20250109-192738\"\n",
-    "ckpt_path = \"/home/kevin/mujoco_playground/mujoco_playground/experimental/learning/checkpoints/G1JoystickFlatTerrain-20250109-232834\"\n",
-    "ckpt_path = \"/home/kevin/mujoco_playground/mujoco_playground/experimental/learning/checkpoints/G1JoystickFlatTerrain-20250109-170057\"\n",
-    "ckpt_path = \"/home/kevin/mujoco_playground/mujoco_playground/experimental/learning/checkpoints/G1JoystickFlatTerrain-20250109-235213\"\n",
-    "\n",
-    "ckpt_path = \"/home/kevin/mujoco_playground/mujoco_playground/experimental/learning/checkpoints/BerkeleyHumanoidJoystickFlatTerrain-20250110-231813\"\n",
-    "ckpt_path = \"/home/kevin/mujoco_playground/mujoco_playground/experimental/learning/checkpoints/BerkeleyHumanoidJoystickFlatTerrain-20250110-233839\"\n",
     "ckpt_path = \"/home/kevin/mujoco_playground/mujoco_playground/experimental/learning/checkpoints/BerkeleyHumanoidJoystickFlatTerrain-20250111-001442\"\n",
-    "ckpt_path = ckpt_path + \"/params.pkl\"\n",
     "\n",
-    "with open(ckpt_path, 'rb') as f:\n",
-    "    params = pickle.load(f)\n",
-    "print(params.keys())\n",
+    "params = load(ckpt_path)\n",
     "\n",
     "output_path = f\"bh_policy.onnx\""
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "params = (params[\"normalizer_params\"], params[\"policy_params\"])"
+    "params = (params[0], params[1])"
    ]
   },
   {