Make the timestep budget settable from the UI

app.py

@@ -28,7 +28,8 @@ AUTO_START = os.environ.get("AUTO_START", "1") == "1"
 OUT = Path("/data/ckpt") if Path("/data").is_dir() else Path("/tmp/ckpt")
 
 LOG: deque[str] = deque(maxlen=800)
-STATE = {"thread": None, "status": "idle", "stop": False, "step": 0}
+STATE = {"thread": None, "status": "idle", "stop": False, "step": 0,
+         "target": NUM_TIMESTEPS}
 
 
 def log(msg: str) -> None:
@@ -153,10 +154,10 @@ def train_worker() -> None:
         env = registry.load(ENV_NAME)
         env_cfg = registry.get_default_config(ENV_NAME)
         ppo_params = locomotion_params.brax_ppo_config(ENV_NAME)
-        ppo_params.num_timesteps = NUM_TIMESTEPS
+        ppo_params.num_timesteps = int(STATE["target"])
         log(f"num_envs={ppo_params.get('num_envs')} "
             f"batch_size={ppo_params.get('batch_size')} "
-            f"timesteps={NUM_TIMESTEPS:,}")
+            f"timesteps={int(STATE['target']):,}")
 
         OUT.mkdir(parents=True, exist_ok=True)
         init_upload()
@@ -218,10 +219,15 @@ def train_worker() -> None:
         log(traceback.format_exc()[-2000:])
 
 
-def start():
+def start(timesteps=None):
+    """Budget is settable from the UI so a rerun does not need a Space restart
+    (env vars are only injected at container start)."""
     t = STATE["thread"]
     if t is not None and t.is_alive():
         return status_md()
+    if timesteps:
+        STATE["target"] = int(timesteps)
+    STATE["step"] = 0
     STATE["stop"] = False
     STATE["thread"] = threading.Thread(target=train_worker, daemon=True)
     STATE["thread"].start()
@@ -236,10 +242,11 @@ def stop():
 
 def status_md() -> str:
     alive = STATE["thread"] is not None and STATE["thread"].is_alive()
-    pct = 100.0 * STATE["step"] / max(NUM_TIMESTEPS, 1)
+    target = int(STATE["target"])
+    pct = 100.0 * STATE["step"] / max(target, 1)
     return (f"**{ENV_NAME}** — status: `{STATE['status']}`"
             f"{' (running)' if alive else ''}  \n"
-            f"step {STATE['step']:,} / {NUM_TIMESTEPS:,}  ({pct:.1f}%)  \n"
+            f"step {STATE['step']:,} / {target:,}  ({pct:.1f}%)  \n"
             f"checkpoints: `{OUT}`"
             + (f" → pushing to `{HF_REPO}`" if HF_REPO else
                "  \n⚠️ `HF_REPO` unset — checkpoints are not pushed to the Hub"))
@@ -253,7 +260,9 @@ with gr.Blocks(title="G1 rough-terrain training") as demo:
     gr.Markdown("# Unitree G1 — rough-terrain locomotion training")
     st = gr.Markdown(status_md())
     with gr.Row():
-        gr.Button("Start", variant="primary").click(start, outputs=st)
+        steps_in = gr.Number(value=NUM_TIMESTEPS, precision=0, label="timesteps",
+                             minimum=100_000, maximum=2_000_000_000)
+        gr.Button("Start", variant="primary").click(start, inputs=steps_in, outputs=st)
         gr.Button("Stop", variant="stop").click(stop, outputs=st)
     out = gr.Textbox(label="training log", lines=26, max_lines=26,
                      autoscroll=True, value=logs())