rewrite base policy/platform

Author: CrafterSteve

eval/base_platform.py

@@ -1,78 +1,86 @@
-"""Abstract robot client: policy predict + platform-specific hardware hooks."""
+"""Base control layer for robot platforms (hardware or sim adapters)."""
 
 from __future__ import annotations
 
-from abc import ABC, abstractmethod
+import importlib
+import os
+from types import ModuleType
 from typing import Any
 
-from model_client import ModelClient, ModelResponse, image_to_rgb_hwc_u8_bytes, state_to_list
+import numpy as np
 
 
-class BasePolicy(ABC):
-    """Base class for robot policy clients."""
+PLATFORM_MODULES = {
+    "so101": "eval.lerobot_so101.so101_client",
+    "lerobot_so101": "eval.lerobot_so101.so101_client",
+}
 
-    def __init__(self, policy: ModelClient):
-        self._policy = policy
 
-    @property
-    def policy(self) -> ModelClient:
-        return self._policy
-
-    def predict(self, image: Any, state: Any, prompt: str, *, image_name: str = "camera1") -> ModelResponse:
-        """Synchronous policy inference."""
-        rgb, width, height, stride = image_to_rgb_hwc_u8_bytes(image)
-        observation = {
-            "images": [
-                {
-                    "name": image_name,
-                    "rgb_hwc_u8": rgb,
-                    "width": width,
-                    "height": height,
-                    "stride_bytes": stride,
-                }
-            ],
-            "state": state_to_list(state),
-            "prompt": prompt,
-        }
-        return self._policy.predict(observation)
-
-    def health(self) -> str:
-        return self._policy.health()
-
-    def reset_policy(self) -> str:
-        return self._policy.reset()
-
-    @abstractmethod
+class BasePlatform:
+    """Platform-side observe / act hooks. Override only what the deployment uses."""
+
     def connect(self) -> None:
-        """Connect robot hardware and capture platform-specific state."""
+        """Open platform resources (serial, cameras, sim env, etc.)."""
 
-    @abstractmethod
     def disconnect(self) -> None:
-        """Disconnect robot hardware."""
+        """Release platform resources."""
 
-    @abstractmethod
-    def get_observation(self) -> dict[str, Any]:
-        """Return the latest robot observation dict."""
+    def __enter__(self) -> BasePlatform:
+        self.connect()
+        return self
 
-    @abstractmethod
-    def send_action(self, action: dict[str, float]) -> None:
-        """Send one control step to the robot."""
+    def __exit__(self, exc_type: object, exc: object, tb: object) -> None:
+        self.disconnect()
+
+    def get_observation(self) -> dict[str, Any]:
+        raise NotImplementedError(f"{type(self).__name__} must implement get_observation()")
+
+    def send_action(self, action: Any) -> None:
+        if isinstance(action, dict):
+            payload = action
+        else:
+            keys = self.action_keys
+            if not keys:
+                raise RuntimeError("platform.action_keys is empty; connect the robot first")
+            values = np.asarray(action, dtype=np.float32).reshape(-1)
+            payload = {key: float(values[index]) for index, key in enumerate(keys)}
+        self._send_action(payload)
+
+    def _send_action(self, action: dict[str, float]) -> None:
+        raise NotImplementedError(f"{type(self).__name__} must implement _send_action()")
 
-    @abstractmethod
     def reset_home(self) -> None:
-        """Move robot back to the platform home pose and reset policy state."""
+        """Optional platform homing motion."""
+        self.on_reset_home()
+
+    def on_reset_home(self) -> None:
+        """Override for platform-specific homing."""
 
     @property
-    @abstractmethod
-    def camera_key(self) -> str:
-        """Primary camera key inside ``get_observation()``."""
+    def camera_key(self) -> str | None:
+        return None
 
     @property
     def model_image_name(self) -> str:
-        """Image key sent to the model server (defaults to ``camera_key``)."""
-        return self.camera_key
+        key = self.camera_key
+        return key if key is not None else "image"
 
     @property
-    @abstractmethod
     def action_keys(self) -> list[str]:
-        """Ordered action feature keys for this robot."""
+        return []
+
+
+def load_platform_module(name: str | None = None) -> ModuleType:
+    platform_name = name or os.environ.get("ROBOT_PLATFORM") or os.environ.get("PLATFORM") or "so101"
+    module_name = PLATFORM_MODULES.get(platform_name)
+    if module_name is None:
+        supported = ", ".join(sorted(PLATFORM_MODULES))
+        raise SystemExit(f"Unknown platform {platform_name!r}. Supported: {supported}")
+    return importlib.import_module(module_name)
+
+
+def create_platform(cfg: Any = None) -> BasePlatform:
+    module = load_platform_module()
+    if cfg is None:
+        cfg = module.config_from_env()
+    return module.create_platform(cfg)

eval/lerobot_so101/run_sync.py

@@ -3,33 +3,24 @@
 
 from __future__ import annotations
 
-import importlib
 import logging
 import os
 import sys
 from pathlib import Path
-from types import ModuleType
 
-_BASE_POLICY_DIR = Path(__file__).resolve().parent
-_ROBOT_CLIENT_DIR = _BASE_POLICY_DIR.parent
-_VLA_CPP_ROOT = _ROBOT_CLIENT_DIR.parent
-_MODEL_CLIENT_PYTHON = _ROBOT_CLIENT_DIR / "python"
-_LEROBOT_SO101 = _VLA_CPP_ROOT / "eval" / "lerobot_so101"
+_ROBOT_CPP_ROOT = Path(__file__).resolve().parents[2]
+_LEROBOT_SO101 = Path(__file__).resolve().parent
+_MODEL_CLIENT_PYTHON = _ROBOT_CPP_ROOT / "robot_client" / "python"
 
-if str(_ROBOT_CLIENT_DIR) not in sys.path:
-    sys.path.insert(0, str(_ROBOT_CLIENT_DIR))
-if str(_MODEL_CLIENT_PYTHON) not in sys.path:
-    sys.path.insert(0, str(_MODEL_CLIENT_PYTHON))
-if str(_LEROBOT_SO101) not in sys.path:
-    sys.path.insert(0, str(_LEROBOT_SO101))
+for path in (_ROBOT_CPP_ROOT, _LEROBOT_SO101, _MODEL_CLIENT_PYTHON):
+    text = str(path)
+    if text not in sys.path:
+        sys.path.insert(0, text)
 
 from model_client import ModelClient
-from base_policy.sync_loop import SyncControlLoop, SyncLoopConfig
-
-DEFAULT_PLATFORM = "lerobot_so101"
-PLATFORM_MODULES = {
-    "lerobot_so101": "so101_client",
-}
+from eval.base_platform import create_platform
+from robot_client.policy.base_policy import RobotPolicy
+from robot_client.policy.sync_loop import SyncControlLoop, SyncLoopConfig
 
 
 def _parse_host_port(target: str, default_port: int = 5555) -> tuple[str, int]:
@@ -47,25 +38,29 @@ def server_from_env(default_host: str = "127.0.0.1", default_port: int = 5555) -
     return host, port, server_timeout
 
 
-def load_platform_module(name: str | None = None) -> ModuleType:
-    platform = name or os.environ.get("ROBOT_PLATFORM", DEFAULT_PLATFORM)
-    module_name = PLATFORM_MODULES.get(platform)
-    if module_name is None:
-        supported = ", ".join(sorted(PLATFORM_MODULES))
-        raise SystemExit(f"Unknown ROBOT_PLATFORM={platform!r}. Supported: {supported}")
-    return importlib.import_module(module_name)
-
-
 def main() -> int:
     logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
 
-    platform = load_platform_module()
     host, port, timeout = server_from_env()
-    policy = ModelClient(host=host, port=port, timeout=timeout)
-    cfg = platform.config_from_env()
-    robot = platform.create_robot_client(policy, cfg)
+    client = ModelClient(host=host, port=port, timeout=timeout)
+    policy = RobotPolicy(client)
+
+    from eval.lerobot_so101.so101_client import config_from_env
+
+    cfg = config_from_env()
+    platform = create_platform(cfg)
+
+    logging.info(
+        "platform=SO101 server=%s:%s camera_key=%s model_image_name=%s",
+        host,
+        port,
+        platform.camera_key,
+        platform.model_image_name,
+    )
+
     SyncControlLoop(
-        robot,
+        platform,
+        policy,
         SyncLoopConfig(task=cfg.task, fps=cfg.fps, loops=cfg.loops),
     ).run()
     return 0

eval/lerobot_so101/shell/run_robot_client.sh

@@ -7,4 +7,4 @@ echo "[robot_sync] platform=${ROBOT_PLATFORM} server=${SERVER}"
 echo "[robot_sync] robot_port=${ROBOT_PORT} camera_key=${CAMERA_KEY} fps=${FPS}"
 echo "[robot_sync] Start server first: bash robot_server/shell/launch_robot_server_mac_cpu.sh"
 
-run_python -m base_policy
+run_python "${ROBOT_CPP_ROOT}/eval/lerobot_so101/run_sync.py"

eval/lerobot_so101/shell/so101_env.sh

@@ -3,11 +3,11 @@
 # Edit the values below for your machine (serial ports, camera index, dataset hub id, etc.).
 
 ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
-VLA_CPP_ROOT="$(cd "${ROOT}/../.." && pwd)"
+ROBOT_CPP_ROOT="$(cd "${ROOT}/../.." && pwd)"
 
 export CONDA_ENV="${CONDA_ENV:-lerobot-demo}"
 
-export PYTHONPATH="${ROOT}:${ROOT}/lerobot_camera_opencv_crop:${VLA_CPP_ROOT}/robot_client/python:${VLA_CPP_ROOT}/robot_client"
+export PYTHONPATH="${ROOT}:${ROOT}/lerobot_camera_opencv_crop:${ROBOT_CPP_ROOT}/robot_client/python:${ROBOT_CPP_ROOT}/robot_client"
 
 # --- Robot serial ports ---
 export ROBOT_PORT="/dev/tty.usbmodem5B3E1195731"

eval/lerobot_so101/so101_client.py

@@ -1,4 +1,4 @@
-"""LeRobot SO101 implementation of ``BasePolicy``."""
+"""LeRobot SO101 platform adapter."""
 
 from __future__ import annotations
 
@@ -10,8 +10,7 @@
 from lerobot.robots.so_follower.config_so_follower import SOFollowerRobotConfig
 from lerobot.robots.so_follower.so_follower import SOFollower
 
-from model_client import ModelClient
-from base_policy.base import BasePolicy
+from eval.base_platform import BasePlatform
 from utils.robot import build_camera_config, extract_home_action
 
 DEFAULT_FPS = 25
@@ -55,24 +54,14 @@ def config_from_env() -> SO101ClientConfig:
     )
 
 
-class SO101RobotClient(BasePolicy):
-
-    def __init__(self, policy: ModelClient, cfg: SO101ClientConfig | None = None):
-        super().__init__(policy)
+class SO101Platform(BasePlatform):
+    def __init__(self, cfg: SO101ClientConfig | None = None):
         self.cfg = cfg or config_from_env()
         self._robot: SOFollower | None = None
         self._home_action: dict[str, float] = {}
         self._dt = 1.0 / max(1, self.cfg.fps)
 
     def connect(self) -> None:
-        health = self._policy.health()
-        logging.info(
-            "Connected to robot server at %s:%s (%s)",
-            self._policy.host,
-            self._policy.port,
-            health,
-        )
-
         cams = build_camera_config(self.cfg.robot_cameras)
         robot_cfg = SOFollowerRobotConfig(
             port=self.cfg.robot_port,
@@ -101,17 +90,16 @@ def get_observation(self) -> dict:
             raise RuntimeError("SO101 robot is not connected")
         return self._robot.get_observation()
 
-    def send_action(self, action: dict[str, float]) -> None:
+    def _send_action(self, action: dict[str, float]) -> None:
         if self._robot is None:
             raise RuntimeError("SO101 robot is not connected")
         self._robot.send_action(action)
 
-    def reset_home(self) -> None:
+    def on_reset_home(self) -> None:
         if self._robot is None:
             raise RuntimeError("SO101 robot is not connected")
-        self.reset_policy()
         for _ in range(max(6, int(self.cfg.fps * 0.8))):
-            self._robot.send_action(dict(self._home_action))
+            self._send_action(dict(self._home_action))
             time.sleep(max(0.01, self._dt))
 
     @property
@@ -128,5 +116,6 @@ def action_keys(self) -> list[str]:
             return []
         return list(self._robot.action_features.keys())
 
-def create_robot_client(policy: ModelClient, cfg: SO101ClientConfig | None = None) -> SO101RobotClient:
-    return SO101RobotClient(policy, cfg)
+
+def create_platform(cfg: SO101ClientConfig | None = None) -> SO101Platform:
+    return SO101Platform(cfg or config_from_env())