added local checkpoints support

Author: omarrayyann

src/mjlab/scripts/play.py

@@ -196,6 +196,46 @@ def __call__(self, obs) -> torch.Tensor:
     policy = runner.get_inference_policy(device=device)
 
   ckpt_manager = None
+  if TRAINED_MODE and cfg.wandb_run_path is None and resume_path is not None:
+    import time as _time
+
+    ckpt_dir = resume_path.parent
+    _runner_local = runner
+
+    def fetch_available_local() -> list[tuple[str, str]]:
+      now = _time.time()
+      entries: list[tuple[str, str, int]] = []
+      for f in sorted(ckpt_dir.glob("*.pt")):
+        try:
+          step = int(f.stem.split("_")[1])
+        except (IndexError, ValueError):
+          step = 0
+        s = int(now - f.stat().st_mtime)
+        for div, unit in ((86400, "d"), (3600, "h"), (60, "m")):
+          if s >= div:
+            t = f"{s // div}{unit} ago"
+            break
+        else:
+          t = f"{s}s ago"
+        entries.append((f.name, t, step))
+      entries.sort(key=lambda x: x[2])
+      return [(name, t) for name, t, _ in entries]
+
+    def load_checkpoint_local(name: str):
+      _runner_local.load(
+        str(ckpt_dir / name),
+        load_cfg={"actor": True},
+        strict=True,
+        map_location=device,
+      )
+      return _runner_local.get_inference_policy(device=device)
+
+    ckpt_manager = CheckpointManager(
+      current_name=resume_path.name,
+      fetch_available=fetch_available_local,
+      load_checkpoint=load_checkpoint_local,
+    )
+
   if TRAINED_MODE and cfg.wandb_run_path is not None:
     from datetime import datetime, timezone
 

src/mjlab/viewer/viser/viewer.py

@@ -36,12 +36,12 @@
 
 @dataclass
 class CheckpointManager:
-  run_name: str
-  run_url: str
-  run_status: str
   current_name: str
   fetch_available: Callable[[], list[tuple[str, str]]]
   load_checkpoint: Callable[[str], PolicyProtocol]
+  run_name: str | None = None
+  run_url: str | None = None
+  run_status: str | None = None
 
 
 class UpdateReason(Enum):
@@ -186,24 +186,34 @@ def _debug_viz_extra() -> None:
     self._scene.create_groups_gui(tabs)
 
     if self._ckpt_mgr is not None:
-      with tabs.add_tab("W&B Run", icon=viser.Icon.CLOUD):
-        self._server.gui.add_html(
-          f'<div style="font-size: 0.85em; line-height: 1.25;'
-          f' padding: 0 1em 0.5em 1em;">'
-          f"<strong>Name:</strong> {self._ckpt_mgr.run_name}<br/>"
-          f"<strong>Status:</strong> {self._ckpt_mgr.run_status}"
-          f"</div>"
-        )
-
-        open_button = self._server.gui.add_button(
-          "Open Run",
-          icon=viser.Icon.EXTERNAL_LINK,
-        )
-
-        @open_button.on_click
-        def _(_) -> None:
-          assert self._ckpt_mgr is not None
-          webbrowser.open(self._ckpt_mgr.run_url)
+      is_wandb = self._ckpt_mgr.run_url is not None
+      with tabs.add_tab("Checkpoints", icon=viser.Icon.DATABASE):
+        if is_wandb:
+          self._server.gui.add_html(
+            f'<div style="font-size: 0.85em; line-height: 1.25;'
+            f' padding: 0 1em 0.5em 1em;">'
+            f"<strong>Source:</strong> W&B<br/>"
+            f"<strong>Run:</strong> {self._ckpt_mgr.run_name}<br/>"
+            f"<strong>Status:</strong> {self._ckpt_mgr.run_status}"
+            f"</div>"
+          )
+
+          open_button = self._server.gui.add_button(
+            "Open Run",
+            icon=viser.Icon.EXTERNAL_LINK,
+          )
+
+          @open_button.on_click
+          def _(_) -> None:
+            assert self._ckpt_mgr is not None
+            webbrowser.open(self._ckpt_mgr.run_url)
+        else:
+          self._server.gui.add_html(
+            '<div style="font-size: 0.85em; line-height: 1.25;'
+            ' padding: 0 1em 0.5em 1em;">'
+            "<strong>Source:</strong> Local"
+            "</div>"
+          )
 
         self._ckpt_dropdown = self._server.gui.add_dropdown(
           "Checkpoint",
@@ -259,6 +269,13 @@ def _handle_custom_action(self, action: ViewerAction, payload: Optional[Any]) ->
       print(f"[INFO]: Loading {name}...")
       self.policy = self._ckpt_mgr.load_checkpoint(name)
       self._ckpt_mgr.current_name = name
+      self._ckpt_updating = True
+      cur = next(
+        (lbl for lbl in self._ckpt_dropdown.options if lbl.startswith(name)),
+        name,
+      )
+      self._ckpt_dropdown.value = cur
+      self._ckpt_updating = False
       self.reset_environment()
       print(f"[INFO]: Loaded {name}")
     return True