Update app.py

app.py

@@ -43,6 +43,20 @@ def tail_file(path: str, n=200):
     return "".join(lines[-n:])
 
 # ---------- RUN DIR HELPERS ----------
+def newest_run(prefer_checkpoint: bool = True) -> str:
+    """Return the newest pusht_* folder. If prefer_checkpoint=True, pick the newest that has checkpoints/last/."""
+    root = pathlib.Path(RUN_ROOT)
+    if not root.exists():
+        return ""
+    runs = sorted(root.glob("pusht_*"), key=lambda r: r.stat().st_mtime, reverse=True)
+    if not runs:
+        return ""
+    if prefer_checkpoint:
+        for r in runs:
+            if has_checkpoint(str(r)):
+                return str(r)
+    return str(runs[0])
+    
 def new_run_dir():
     """Return a unique run dir path WITHOUT creating it (so LeRobot can create it)."""
     base = pathlib.Path(RUN_ROOT) / f"pusht_{int(time.time())}"
@@ -55,12 +69,43 @@ def new_run_dir():
     return str(d)
 
 def current_run_dir(user_override: str | None):
-    """Prefer user text if given, else use the LAST pointer if present."""
+    """
+    Resolve which run to use:
+    - If user typed something, accept folder name or full path.
+    - Else try LAST pointer.
+    - Else pick newest run with a checkpoint.
+    - Else pick newest run (even without checkpoint).
+    - Else return "" (none).
+    """
+    # A) explicit user input
     if user_override and user_override.strip():
-        return user_override.strip()
+        p = user_override.strip()
+        # allow just "pusht_123..." as well as absolute path
+        if not p.startswith("/"):
+            p = str(pathlib.Path(RUN_ROOT) / p)
+        return p
+
+    # B) LAST pointer if present and valid
     if LAST_PTR.exists():
-        return LAST_PTR.read_text().strip()
+        p = LAST_PTR.read_text().strip()
+        if p and os.path.isdir(p):
+            return p
+
+    # C) newest run WITH checkpoint
+    p = newest_run(prefer_checkpoint=True)
+    if p:
+        LAST_PTR.write_text(p)
+        return p
+
+    # D) newest run (no checkpoint yet)
+    p = newest_run(prefer_checkpoint=False)
+    if p:
+        LAST_PTR.write_text(p)
+        return p
+
+    # E) nothing found
     return ""
+    
 
 def has_checkpoint(run_dir: str):
     """We consider a checkpoint present once checkpoints/last/ exists (first save is at step 500)."""
@@ -105,12 +150,12 @@ def start_training(steps, batch_size, push_to_hub, repo_id):
 def resume_training(extra_steps, push_to_hub, repo_id, run_dir_text):
     run_dir = current_run_dir(run_dir_text)
     if not run_dir:
-        return "No run found yet. Start a fresh training first.", "", "(no log)"
+        return "No run found on disk. Start a fresh training first (let it pass step 500 to create a checkpoint).", "", "(no log)"
     log = train_log_path(run_dir)
 
     if not has_checkpoint(run_dir):
-        return f"No checkpoint in {run_dir}/checkpoints/last/ yet — run at least 500 steps once.", run_dir, tail_file(log)
-
+        return f"Selected run: {run_dir}\nNo checkpoint in {run_dir}/checkpoints/last/ yet — run at least 500 steps once.", run_dir, tail_file(log)
+        
     push_flags = (f"--policy.push_to_hub=true --policy.repo_id='{repo_id.strip()}'"
                   if push_to_hub and repo_id.strip() else
                   "--policy.push_to_hub=false")