v1.1.7; Add go2 lab eval; Add server client connection checker

robogauge/scripts/client.py

@@ -6,12 +6,36 @@ from typing import Dict, Any, Optional
 from robogauge.scripts.server import ResponseStatus
 
 class RoboGaugeClient:
-    def __init__(self, base_url: str = "http://127.0.0.1:9973"):
+    def __init__(self, base_url: str = "http://127.0.0.1:9973", request_timeout: float = 5.0):
         self.base_url = base_url
+        self.request_timeout = request_timeout
         self.processing_ids = []
         self.task_id2info = {}
         self.response_data = {}
 
+    def is_server_available(self) -> bool:
+        """Return True when the RoboGauge API and evaluator main process are healthy."""
+        try:
+            response = requests.get(
+                f"{self.base_url}/health",
+                timeout=self.request_timeout,
+            )
+            if response.status_code != 200:
+                return False
+            return response.json().get("status") == "ok"
+        except (requests.exceptions.RequestException, ValueError):
+            return False
+
+    def wait_until_available(self, retry_interval: int = 2) -> None:
+        """Block until the RoboGauge API can be reached."""
+        while not self.is_server_available():
+            print(
+                f"[RoboGaugeClient]⏳ Waiting for RoboGauge server at {self.base_url} "
+                f"(retrying in {retry_interval}s)..."
+            )
+            time.sleep(retry_interval)
+        print(f"[RoboGaugeClient]✅ Connected to RoboGauge server at {self.base_url}.")
+
     def submit_task(self, 
         model_path: str, 
         step: int, 
@@ -40,9 +64,16 @@ class RoboGaugeClient:
 
         print(f"[RoboGaugeClient]📤 Preparing to submit task: {task_name}")
 
+        if wait_for_server:
+            self.wait_until_available(retry_interval=retry_interval)
+
         while True:
             try:
-                response = requests.post(f"{self.base_url}/submit_eval", json=payload)
+                response = requests.post(
+                    f"{self.base_url}/submit_eval",
+                    json=payload,
+                    timeout=self.request_timeout,
+                )
                 
                 if response.status_code == 200:
                     data = response.json()
@@ -55,9 +86,9 @@ class RoboGaugeClient:
                     print(f"[RoboGaugeClient]❌ Server returned error: {response.text}")
                     return None
 
-            except requests.exceptions.ConnectionError:
+            except requests.exceptions.RequestException:
                 if not wait_for_server:
-                    print("[RoboGaugeClient]❌ Unable to connect to server (Connection Refused).")
+                    print("[RoboGaugeClient]❌ Unable to connect to server.")
                     return None
                 
                 print(f"[RoboGaugeClient]⏳ Server not responding, retrying in {retry_interval} seconds... (Make sure server.py is running)")
@@ -66,11 +97,24 @@ class RoboGaugeClient:
                 print(f"[RoboGaugeClient]❌ Unknown error: {e}")
                 return None
 
-    def monitor_tasks(self):
+    def monitor_tasks(self, wait_for_server: bool = True, retry_interval: int = 2):
         print("[RoboGaugeClient]⏱️ Monitoring submitted tasks...")
         """ Monitor all submitted tasks until completion. """
+        if wait_for_server:
+            self.wait_until_available(retry_interval=retry_interval)
+
         for task_id in reversed(self.processing_ids):
-            respone = requests.get(f"{self.base_url}/get_result/{task_id}")
+            try:
+                respone = requests.get(f"{self.base_url}/get_result/{task_id}", timeout=self.request_timeout)
+            except requests.exceptions.RequestException:
+                if not wait_for_server:
+                    raise
+                print(
+                    f"[RoboGaugeClient]⏳ Lost connection to RoboGauge server at {self.base_url}; "
+                    f"waiting before checking task results..."
+                )
+                self.wait_until_available(retry_interval=retry_interval)
+                continue
             if respone.status_code != 200:
                 continue
             resp_data = respone.json()

robogauge/scripts/server.py

@@ -21,11 +21,12 @@ os.environ["OMP_NUM_THREADS"] = "1"
 os.environ["MKL_NUM_THREADS"] = "1"
 
 import multiprocessing
+import threading
 import uvicorn
 import queue
 import time
 import uuid
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from typing import Dict, Optional
 import argparse
@@ -47,6 +48,8 @@ default_args_list = [
     '--headless',
 ]
 
+HEARTBEAT_TIMEOUT_SECONDS = 10.0
+
 @dataclass
 class EvalTaskData:
     model_path: str
@@ -67,15 +70,42 @@ class ResponseStatus:
     ERROR = "error"
     NOT_FOUND = "not_found"
 
-def run_api_server(input_queue: multiprocessing.Queue, result_dict: dict, port=9973):
+def update_main_heartbeat(health_dict: dict, stop_event: threading.Event, interval: float = 1.0):
+    while not stop_event.is_set():
+        health_dict["main_pid"] = os.getpid()
+        health_dict["last_heartbeat"] = time.time()
+        stop_event.wait(interval)
+
+def is_main_process_healthy(health_dict: dict) -> bool:
+    last_heartbeat = float(health_dict.get("last_heartbeat", 0.0))
+    return time.time() - last_heartbeat <= HEARTBEAT_TIMEOUT_SECONDS
+
+def run_api_server(input_queue: multiprocessing.Queue, result_dict: dict, health_dict: dict, port=9973):
     """
     Running in a separate subprocess.
     I/O Process: submit requests -> put into queue -> return ID.
     """
     app = FastAPI()
 
+    @app.get("/health")
+    def health():
+        last_heartbeat = float(health_dict.get("last_heartbeat", 0.0))
+        heartbeat_age = time.time() - last_heartbeat
+        if heartbeat_age > HEARTBEAT_TIMEOUT_SECONDS:
+            raise HTTPException(
+                status_code=503,
+                detail=f"RoboGauge main process heartbeat is stale ({heartbeat_age:.1f}s).",
+            )
+        return {
+            "status": "ok",
+            "main_pid": health_dict.get("main_pid"),
+            "heartbeat_age": heartbeat_age,
+        }
+
     @app.post("/submit_eval")
     def submit_eval(req: EvalRequest):
+        if not is_main_process_healthy(health_dict):
+            raise HTTPException(status_code=503, detail="RoboGauge main process is not healthy.")
         task_id = str(uuid.uuid4())
         task_data = EvalTaskData(
             model_path=req.model_path,
@@ -110,10 +140,18 @@ def main():
     manager = ctx.Manager()
     task_queue = manager.Queue()
     results_store = manager.dict()
+    health_store = manager.dict()
+    heartbeat_stop = threading.Event()
+    heartbeat_thread = threading.Thread(
+        target=update_main_heartbeat,
+        args=(health_store, heartbeat_stop),
+        daemon=True,
+    )
+    heartbeat_thread.start()
 
     api_p = ctx.Process(
         target=run_api_server, 
-        args=(task_queue, results_store, args_cli.port),
+        args=(task_queue, results_store, health_store, args_cli.port),
         daemon=True
     )
     api_p.start()
@@ -165,6 +203,8 @@ def main():
         print("\n🛑 Shutting down...")
         api_p.terminate()
         api_p.join()
+        heartbeat_stop.set()
+        heartbeat_thread.join(timeout=2.0)
 
 if __name__ == "__main__":
     main()

robogauge/tasks/robots/go2/go2_lab_config.py

@@ -14,6 +14,7 @@ class Go2LabConfig(Go2Config):
     """Go2 Lab robot configuration aligned with RobotLab observation scaling."""
 
     class control(Go2Config.control):
+        p_gains = [25.0, 25.0, 25.0, 25.0, 25.0, 25.0, 25.0, 25.0, 25.0, 25.0, 25.0, 25.0]  # [N*m/rad]
         class scales(Go2Config.control.scales):
             # RobotLab policy command observation uses unit scale.
             cmd = [1.0, 1.0, 1.0]