From 43a7493d1739d5cb1618de6863f929d99c7a8822 Mon Sep 17 00:00:00 2001
From: Baruch Tabanpour <btaba@google.com>
Date: Wed, 16 Oct 2024 16:23:03 -0700
Subject: [PATCH] Fix embarrassing bug with ckpt.

PiperOrigin-RevId: 686679906
Change-Id: I61126ce94214b5f7b25b31e343a243d1a1dbb9f4
---
 mjx/tutorial.ipynb | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mjx/tutorial.ipynb b/mjx/tutorial.ipynb
index 6314245a..751a150a 100644
--- a/mjx/tutorial.ipynb
+++ b/mjx/tutorial.ipynb
@@ -1513,8 +1513,8 @@
    "source": [
     "# grab the latest checkpoint from the flat terrain joystick policy\n",
     "latest_ckpts = list(ckpt_path.glob('*'))\n",
-    "latest_ckpts.sort()\n",
-    "latest_ckpt = latest_ckpts[0]\n",
+    "latest_ckpts.sort(key=lambda x: int(x.as_posix().split('/')[-1]))\n",
+    "latest_ckpt = latest_ckpts[-1]\n",
     "\n",
     "train_fn = functools.partial(\n",
     "      ppo.train, num_timesteps=40_000_000, num_evals=5,\n",