luckyrobots 0.1.72__py3-none-any.whl → 0.1.74__py3-none-any.whl

luckyrobots/client.py

@@ -9,8 +9,6 @@ from __future__ import annotations
 
 import logging
 import time
-import statistics
-from dataclasses import dataclass
 from types import SimpleNamespace
 from typing import Any, Optional
 
@@ -19,26 +17,20 @@ logger = logging.getLogger("luckyrobots.client")
 try:
     from .grpc.generated import agent_pb2  # type: ignore
     from .grpc.generated import agent_pb2_grpc  # type: ignore
-    from .grpc.generated import camera_pb2  # type: ignore
-    from .grpc.generated import camera_pb2_grpc  # type: ignore
     from .grpc.generated import common_pb2  # type: ignore
-    from .grpc.generated import media_pb2  # type: ignore
+    from .grpc.generated import debug_pb2  # type: ignore
+    from .grpc.generated import debug_pb2_grpc  # type: ignore
     from .grpc.generated import mujoco_pb2  # type: ignore
     from .grpc.generated import mujoco_pb2_grpc  # type: ignore
     from .grpc.generated import scene_pb2  # type: ignore
     from .grpc.generated import scene_pb2_grpc  # type: ignore
-    from .grpc.generated import telemetry_pb2  # type: ignore
-    from .grpc.generated import telemetry_pb2_grpc  # type: ignore
-    from .grpc.generated import viewport_pb2  # type: ignore
-    from .grpc.generated import viewport_pb2_grpc  # type: ignore
 except Exception as e:  # pragma: no cover
     raise ImportError(
         "Missing generated gRPC stubs. Regenerate them from the protos in "
         "src/luckyrobots/grpc/proto into src/luckyrobots/grpc/generated."
     ) from e
 
-# Import Pydantic models for type-checked responses
-from .models import ObservationResponse, StateSnapshot, DomainRandomizationConfig
+from .models import ObservationResponse
 
 
 class GrpcConnectionError(Exception):
@@ -53,24 +45,18 @@ class LuckyEngineClient:
     """
     Client for connecting to the LuckyEngine gRPC server.
 
-    Provides access to all gRPC services defined by the protos under
-    `src/luckyrobots/rpc/proto`:
-    - SceneService
-    - MujocoService
-    - TelemetryService
-    - AgentService
-    - ViewportService
-    - CameraService
+    Provides access to gRPC services for RL training:
+    - AgentService: observations, stepping, resets
+    - SceneService: simulation mode control
+    - MujocoService: health checks
 
     Usage:
         client = LuckyEngineClient(host="127.0.0.1", port=50051)
         client.connect()
+        client.wait_for_server()
 
-        # Access services
-        scene_info = client.scene.GetSceneInfo(client.pb.scene.GetSceneInfoRequest())
-        joint_state = client.mujoco.GetJointState(
-            client.pb.mujoco.GetJointStateRequest()
-        )
+        schema = client.get_agent_schema()
+        obs = client.step(actions=[0.0] * 12)
 
         client.close()
     """
@@ -102,25 +88,19 @@ class LuckyEngineClient:
         # Service stubs (populated after connect)
         self._scene = None
         self._mujoco = None
-        self._telemetry = None
         self._agent = None
-        self._viewport = None
-        self._camera = None
+        self._debug = None
 
         # Cached agent schemas: agent_name -> (observation_names, action_names)
-        # Populated lazily by get_agent_schema() or fetch_schema()
         self._schema_cache: dict[str, tuple[list[str], list[str]]] = {}
 
         # Protobuf modules (for discoverability + explicit imports).
         self._pb = SimpleNamespace(
             common=common_pb2,
-            media=media_pb2,
             scene=scene_pb2,
             mujoco=mujoco_pb2,
-            telemetry=telemetry_pb2,
             agent=agent_pb2,
-            viewport=viewport_pb2,
-            camera=camera_pb2,
+            debug=debug_pb2,
         )
 
     def connect(self) -> None:
@@ -145,10 +125,8 @@ class LuckyEngineClient:
         # Create service stubs
         self._scene = scene_pb2_grpc.SceneServiceStub(self._channel)
         self._mujoco = mujoco_pb2_grpc.MujocoServiceStub(self._channel)
-        self._telemetry = telemetry_pb2_grpc.TelemetryServiceStub(self._channel)
         self._agent = agent_pb2_grpc.AgentServiceStub(self._channel)
-        self._viewport = viewport_pb2_grpc.ViewportServiceStub(self._channel)
-        self._camera = camera_pb2_grpc.CameraServiceStub(self._channel)
+        self._debug = debug_pb2_grpc.DebugServiceStub(self._channel)
 
         logger.info(f"Channel opened to {target} (server not verified yet)")
 
@@ -162,10 +140,8 @@ class LuckyEngineClient:
             self._channel = None
             self._scene = None
             self._mujoco = None
-            self._telemetry = None
             self._agent = None
-            self._viewport = None
-            self._camera = None
+            self._debug = None
             logger.info("gRPC channel closed")
 
     def is_connected(self) -> bool:
@@ -209,8 +185,6 @@ class LuckyEngineClient:
         Returns:
             True if server became available, False if timeout.
         """
-        import time
-
         start = time.perf_counter()
 
         while time.perf_counter() - start < timeout:
@@ -228,8 +202,6 @@ class LuckyEngineClient:
 
         return False
 
-    # --- Protobuf modules (discoverable + explicit) ---
-
     @property
     def pb(self) -> Any:
         """Access protobuf modules grouped by domain (e.g., `client.pb.scene`)."""
@@ -244,21 +216,9 @@ class LuckyEngineClient:
         """Set the default robot name used by calls that accept an optional robot_name."""
         self._robot_name = robot_name
 
-    # --- Service stubs ---
-    #
-    # Confirmed working:
-    #   - mujoco: GetMujocoInfo, GetJointState, SendControl
-    #   - agent: GetObservation, ResetAgent, GetAgentSchema, StreamAgent
-    #
-    # Placeholders (not yet confirmed working - use at your own risk):
-    #   - scene: GetSceneInfo
-    #   - telemetry: StreamTelemetry
-    #   - viewport: (no methods implemented)
-    #   - camera: ListCameras, StreamCamera
-
     @property
     def scene(self) -> Any:
-        """SceneService stub. [PLACEHOLDER - not confirmed working]"""
+        """SceneService stub."""
         if self._scene is None:
             raise GrpcConnectionError("Not connected. Call connect() first.")
         return self._scene
@@ -270,13 +230,6 @@ class LuckyEngineClient:
             raise GrpcConnectionError("Not connected. Call connect() first.")
         return self._mujoco
 
-    @property
-    def telemetry(self) -> Any:
-        """TelemetryService stub. [PLACEHOLDER - not confirmed working]"""
-        if self._telemetry is None:
-            raise GrpcConnectionError("Not connected. Call connect() first.")
-        return self._telemetry
-
     @property
     def agent(self) -> Any:
         """AgentService stub."""
@@ -285,80 +238,23 @@ class LuckyEngineClient:
         return self._agent
 
     @property
-    def viewport(self) -> Any:
-        """ViewportService stub. [PLACEHOLDER - not confirmed working]"""
-        if self._viewport is None:
-            raise GrpcConnectionError("Not connected. Call connect() first.")
-        return self._viewport
-
-    @property
-    def camera(self) -> Any:
-        """CameraService stub. [PLACEHOLDER - not confirmed working]"""
-        if self._camera is None:
+    def debug(self) -> Any:
+        """DebugService stub."""
+        if self._debug is None:
             raise GrpcConnectionError("Not connected. Call connect() first.")
-        return self._camera
-
-    # --- Convenience methods ---
-
-    def get_scene_info(self, timeout: Optional[float] = None):
-        """Get scene information. [PLACEHOLDER - not confirmed working]"""
-        raise NotImplementedError(
-            "get_scene_info() is not yet confirmed working. "
-            "Remove this check if you want to test it."
-        )
-        timeout = timeout or self.timeout
-        return self.scene.GetSceneInfo(
-            self.pb.scene.GetSceneInfoRequest(),
-            timeout=timeout,
-        )
+        return self._debug
 
     def get_mujoco_info(self, robot_name: str = "", timeout: Optional[float] = None):
-        """Get MuJoCo model information."""
+        """Get MuJoCo model information (joint names, limits, etc.)."""
         timeout = timeout or self.timeout
         robot_name = robot_name or self._robot_name
         if not robot_name:
-            raise ValueError(
-                "robot_name is required (pass `robot_name=` or set it once via "
-                "LuckyEngineClient(robot_name=...) / client.set_robot_name(...))."
-            )
+            raise ValueError("robot_name is required")
         return self.mujoco.GetMujocoInfo(
             self.pb.mujoco.GetMujocoInfoRequest(robot_name=robot_name),
             timeout=timeout,
         )
 
-    def get_joint_state(self, robot_name: str = "", timeout: Optional[float] = None):
-        """Get current joint state."""
-        timeout = timeout or self.timeout
-        robot_name = robot_name or self._robot_name
-        if not robot_name:
-            raise ValueError(
-                "robot_name is required (pass `robot_name=` or set it once via "
-                "LuckyEngineClient(robot_name=...) / client.set_robot_name(...))."
-            )
-        return self.mujoco.GetJointState(
-            self.pb.mujoco.GetJointStateRequest(robot_name=robot_name),
-            timeout=timeout,
-        )
-
-    def send_control(
-        self,
-        controls: list[float],
-        robot_name: str = "",
-        timeout: Optional[float] = None,
-    ):
-        """Send control commands to the robot."""
-        timeout = timeout or self.timeout
-        robot_name = robot_name or self._robot_name
-        if not robot_name:
-            raise ValueError(
-                "robot_name is required (pass `robot_name=` or set it once via "
-                "LuckyEngineClient(robot_name=...) / client.set_robot_name(...))."
-            )
-        return self.mujoco.SendControl(
-            self.pb.mujoco.SendControlRequest(robot_name=robot_name, controls=controls),
-            timeout=timeout,
-        )
-
     def get_agent_schema(self, agent_name: str = "", timeout: Optional[float] = None):
         """Get agent schema (observation/action sizes and names).
 
@@ -395,18 +291,6 @@ class LuckyEngineClient:
 
         return resp
 
-    def fetch_schema(self, agent_name: str = "", timeout: Optional[float] = None) -> None:
-        """Fetch and cache agent schema for named observation access.
-
-        Call this once before get_observation() to enable accessing observations
-        by name (e.g., obs["proj_grav_x"]).
-
-        Args:
-            agent_name: Agent name (empty = default agent).
-            timeout: RPC timeout.
-        """
-        self.get_agent_schema(agent_name=agent_name, timeout=timeout)
-
     def get_observation(
         self,
         agent_name: str = "",
@@ -415,10 +299,6 @@ class LuckyEngineClient:
         """
         Get the RL observation vector for an agent.
 
-        This returns only the flat observation vector defined by the agent's
-        observation spec in LuckyEngine. For sensor data (joints, telemetry,
-        cameras), use the dedicated methods.
-
         Args:
             agent_name: Agent name (empty = default agent).
             timeout: RPC timeout.
@@ -435,30 +315,17 @@ class LuckyEngineClient:
                 "LuckyEngineClient(robot_name=...) / client.set_robot_name(...))."
             )
 
-        # Request only the agent's RL observation (agent_frame in gRPC terms).
-        # Joint state and telemetry have their own dedicated methods.
         resp = self.agent.GetObservation(
             self.pb.agent.GetObservationRequest(
                 robot_name=resolved_robot_name,
                 agent_name=agent_name,
                 include_joint_state=False,
-                include_agent_frame=True,  # The RL observation vector
+                include_agent_frame=True,
                 include_telemetry=False,
             ),
             timeout=timeout,
         )
 
-        # Extract observation from agent_frame.
-        #
-        # In gRPC, the "agent_frame" is the message containing the RL observation
-        # data from LuckyEngine's agent system. It includes:
-        #   - observations: flat float vector matching the agent's observation spec
-        #   - actions: the last action vector sent to the agent (echoed back)
-        #   - timestamp_ms: wall-clock time when the observation was captured
-        #   - frame_number: monotonic counter for ordering observations
-        #
-        # This is distinct from joint_state (raw MuJoCo qpos/qvel) and telemetry
-        # (debugging data like contact forces, energy, etc).
         agent_frame = getattr(resp, "agent_frame", None)
         observations = []
         actions = []
@@ -471,7 +338,6 @@ class LuckyEngineClient:
             timestamp_ms = getattr(agent_frame, "timestamp_ms", timestamp_ms)
             frame_number = getattr(agent_frame, "frame_number", frame_number)
 
-        # Look up cached schema for named access
         cache_key = agent_name or "agent_0"
         obs_names, action_names = self._schema_cache.get(cache_key, (None, None))
 
@@ -485,117 +351,6 @@ class LuckyEngineClient:
             action_names=action_names,
         )
 
-    def get_state(
-        self,
-        agent_name: str = "",
-        include_observation: bool = True,
-        include_joint_state: bool = True,
-        camera_names: Optional[list[str]] = None,
-        width: int = 0,
-        height: int = 0,
-        format: str = "raw",
-        timeout: Optional[float] = None,
-    ) -> StateSnapshot:
-        """
-        Get a bundled snapshot of multiple data sources.
-
-        Use this for efficiency when you need multiple data types in one call.
-        For single data types, prefer the dedicated methods:
-        - get_observation() for RL observation vector
-        - get_joint_state() for joint positions/velocities
-        - stream_telemetry() for telemetry (streaming only)
-        - stream_camera() for camera frames (streaming)
-
-        Args:
-            agent_name: Agent name (empty = default agent).
-            include_observation: Include RL observation vector.
-            include_joint_state: Include joint positions/velocities.
-            camera_names: List of camera names to include frames from.
-            width: Desired width for camera frames (0 = native).
-            height: Desired height for camera frames (0 = native).
-            format: Image format ("raw" or "jpeg").
-            timeout: RPC timeout.
-
-        Returns:
-            StateSnapshot with requested data.
-        """
-        timeout = timeout or self.timeout
-
-        resolved_robot_name = self._robot_name
-        if not resolved_robot_name:
-            raise ValueError(
-                "robot_name is required (set it once via "
-                "LuckyEngineClient(robot_name=...) / client.set_robot_name(...))."
-            )
-
-        cameras = []
-        if camera_names:
-            for name in camera_names:
-                cameras.append(
-                    self.pb.agent.GetCameraFrameRequest(
-                        name=name,
-                        width=width,
-                        height=height,
-                        format=format,
-                    )
-                )
-
-        resp = self.agent.GetObservation(
-            self.pb.agent.GetObservationRequest(
-                robot_name=resolved_robot_name,
-                agent_name=agent_name,
-                include_joint_state=include_joint_state,
-                include_agent_frame=include_observation,
-                include_telemetry=False,  # Telemetry is streaming-only
-                cameras=cameras,
-            ),
-            timeout=timeout,
-        )
-
-        # Build ObservationResponse if requested
-        obs_response = None
-        if include_observation:
-            agent_frame = getattr(resp, "agent_frame", None)
-            observations = []
-            actions = []
-            if agent_frame is not None:
-                observations = list(agent_frame.observations) if agent_frame.observations else []
-                actions = list(agent_frame.actions) if agent_frame.actions else []
-
-            # Look up cached schema for named access
-            cache_key = agent_name or "agent_0"
-            obs_names, action_names = self._schema_cache.get(cache_key, (None, None))
-
-            obs_response = ObservationResponse(
-                observation=observations,
-                actions=actions,
-                timestamp_ms=getattr(resp, "timestamp_ms", 0),
-                frame_number=getattr(resp, "frame_number", 0),
-                agent_name=cache_key,
-                observation_names=obs_names,
-                action_names=action_names,
-            )
-
-        return StateSnapshot(
-            observation=obs_response,
-            joint_state=getattr(resp, "joint_state", None) if include_joint_state else None,
-            camera_frames=list(getattr(resp, "camera_frames", [])) if camera_names else None,
-            timestamp_ms=getattr(resp, "timestamp_ms", 0),
-            frame_number=getattr(resp, "frame_number", 0),
-        )
-
-    def stream_agent(self, agent_name: str = "", target_fps: int = 30):
-        """
-        Start streaming agent observations.
-
-        Returns an iterator of AgentFrame messages.
-        """
-        return self.agent.StreamAgent(
-            self.pb.agent.StreamAgentRequest(
-                agent_name=agent_name, target_fps=target_fps
-            ),
-        )
-
     def reset_agent(
         self,
         agent_name: str = "",
@@ -603,13 +358,10 @@ class LuckyEngineClient:
         timeout: Optional[float] = None,
     ):
         """
-        Reset a specific agent (full reset: clear buffers, reset state, resample commands, apply MuJoCo state).
-
-        Useful for multi-env RL where individual agents need to be reset without resetting the entire scene.
+        Reset a specific agent.
 
         Args:
-            agent_name: Agent logical name. Convention is `agent_0`, `agent_1`, ...
-                Empty string means "default agent" (agent_0).
+            agent_name: Agent logical name. Empty string means default agent.
             randomization_cfg: Optional domain randomization config for this reset.
             timeout: Timeout in seconds (uses default if None).
 
@@ -618,7 +370,6 @@ class LuckyEngineClient:
         """
         timeout = timeout or self.timeout
 
-        # Build request with optional randomization config
         request_kwargs = {"agent_name": agent_name}
 
         if randomization_cfg is not None:
@@ -640,13 +391,10 @@ class LuckyEngineClient:
         """
         Synchronous RL step: apply action, wait for physics, return observation.
 
-        This is the recommended interface for RL training as it eliminates
-        one network round-trip compared to separate SendControl + GetObservation.
-
         Args:
             actions: Action vector to apply for this step.
             agent_name: Agent name (empty = default agent).
-            timeout_ms: Server-side timeout for the step in milliseconds (0 = use default).
+            timeout_ms: Server-side timeout for the step in milliseconds.
             timeout: RPC timeout in seconds.
 
         Returns:
@@ -666,14 +414,12 @@ class LuckyEngineClient:
         if not resp.success:
             raise RuntimeError(f"Step failed: {resp.message}")
 
-        # Extract observation from AgentFrame
         agent_frame = resp.observation
         observations = list(agent_frame.observations) if agent_frame.observations else []
         actions_out = list(agent_frame.actions) if agent_frame.actions else []
         timestamp_ms = getattr(agent_frame, "timestamp_ms", 0)
         frame_number = getattr(agent_frame, "frame_number", 0)
 
-        # Look up cached schema for named access
         cache_key = agent_name or "agent_0"
         obs_names, action_names = self._schema_cache.get(cache_key, (None, None))
 
@@ -697,8 +443,8 @@ class LuckyEngineClient:
 
         Args:
             mode: "realtime", "deterministic", or "fast"
-                - realtime: Physics runs at 1x wall-clock speed (for visualization)
-                - deterministic: Physics runs at fixed rate (for reproducibility)
+                - realtime: Physics runs at 1x wall-clock speed
+                - deterministic: Physics runs at fixed rate
                 - fast: Physics runs as fast as possible (for RL training)
             timeout: RPC timeout in seconds.
 
@@ -712,39 +458,19 @@ class LuckyEngineClient:
             "deterministic": 1,
             "fast": 2,
         }
-        mode_value = mode_map.get(mode.lower(), 2)  # Default to fast
+        mode_value = mode_map.get(mode.lower(), 2)
 
         return self.scene.SetSimulationMode(
             self.pb.scene.SetSimulationModeRequest(mode=mode_value),
             timeout=timeout,
         )
 
-    def get_simulation_mode(self, timeout: Optional[float] = None):
-        """
-        Get current simulation timing mode.
-
-        Returns:
-            GetSimulationModeResponse with mode field (0=realtime, 1=deterministic, 2=fast).
-        """
-        timeout = timeout or self.timeout
-
-        return self.scene.GetSimulationMode(
-            self.pb.scene.GetSimulationModeRequest(),
-            timeout=timeout,
-        )
-
     def _randomization_to_proto(self, randomization_cfg: Any):
-        """Convert domain randomization config to proto message.
-
-        Accepts any object with randomization config fields (DomainRandomizationConfig,
-        PhysicsDRCfg from luckylab, or similar).
-        """
+        """Convert domain randomization config to proto message."""
         proto_kwargs = {}
 
-        # Helper to get attribute value, checking for None
         def get_val(name: str, default=None):
             val = getattr(randomization_cfg, name, default)
-            # Handle both None and empty tuples/lists
             if val is None or (isinstance(val, (tuple, list)) and len(val) == 0):
                 return None
             return val
@@ -766,7 +492,7 @@ class LuckyEngineClient:
         if joint_vel is not None and joint_vel != 0.0:
             proto_kwargs["joint_velocity_noise"] = joint_vel
 
-        # Physics parameters (ranges)
+        # Physics parameters
         friction = get_val("friction_range")
         if friction is not None:
             proto_kwargs["friction_range"] = list(friction)
@@ -812,186 +538,140 @@ class LuckyEngineClient:
 
         return self.pb.agent.DomainRandomizationConfig(**proto_kwargs)
 
-    def stream_telemetry(self, target_fps: int = 30):
+    def draw_velocity_command(
+        self,
+        origin: tuple[float, float, float],
+        lin_vel_x: float,
+        lin_vel_y: float,
+        ang_vel_z: float,
+        scale: float = 1.0,
+        clear_previous: bool = True,
+        timeout: Optional[float] = None,
+    ) -> bool:
         """
-        Start streaming telemetry data. [PLACEHOLDER - not confirmed working]
+        Draw velocity command visualization in LuckyEngine.
 
-        Returns an iterator of TelemetryFrame messages.
+        Args:
+            origin: (x, y, z) position of the robot.
+            lin_vel_x: Forward velocity command.
+            lin_vel_y: Lateral velocity command.
+            ang_vel_z: Angular velocity command (yaw rate).
+            scale: Scale factor for visualization.
+            clear_previous: Clear previous debug draws before drawing.
+            timeout: RPC timeout in seconds.
+
+        Returns:
+            True if draw succeeded, False otherwise.
         """
-        raise NotImplementedError(
-            "stream_telemetry() is not yet confirmed working. "
-            "Remove this check if you want to test it."
-        )
-        return self.telemetry.StreamTelemetry(
-            self.pb.telemetry.StreamTelemetryRequest(target_fps=target_fps),
-        )
+        timeout = timeout or self.timeout
 
-    def list_cameras(self, timeout: Optional[float] = None):
-        """List available cameras. [PLACEHOLDER - not confirmed working]"""
-        raise NotImplementedError(
-            "list_cameras() is not yet confirmed working. "
-            "Remove this check if you want to test it."
+        velocity_cmd = self.pb.debug.DebugVelocityCommand(
+            origin=self.pb.debug.DebugVector3(x=origin[0], y=origin[1], z=origin[2]),
+            lin_vel_x=lin_vel_x,
+            lin_vel_y=lin_vel_y,
+            ang_vel_z=ang_vel_z,
+            scale=scale,
         )
-        timeout = timeout or self.timeout
-        return self.camera.ListCameras(
-            self.pb.camera.ListCamerasRequest(),
-            timeout=timeout,
+
+        request = self.pb.debug.DebugDrawRequest(
+            velocity_command=velocity_cmd,
+            clear_previous=clear_previous,
         )
 
-    def stream_camera(
+        try:
+            resp = self.debug.Draw(request, timeout=timeout)
+            return resp.success
+        except Exception as e:
+            logger.debug(f"Debug draw failed: {e}")
+            return False
+
+    def draw_arrow(
         self,
-        camera_id: Optional[int] = None,
-        camera_name: Optional[str] = None,
-        target_fps: int = 30,
-        width: int = 640,
-        height: int = 480,
-        format: str = "raw",
-    ):
+        origin: tuple[float, float, float],
+        direction: tuple[float, float, float],
+        color: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 1.0),
+        scale: float = 1.0,
+        clear_previous: bool = False,
+        timeout: Optional[float] = None,
+    ) -> bool:
         """
-        Start streaming camera frames. [PLACEHOLDER - not confirmed working]
+        Draw a debug arrow in LuckyEngine.
 
         Args:
-            camera_id: Camera entity ID (use either id or name).
-            camera_name: Camera name (use either id or name).
-            target_fps: Desired frames per second.
-            width: Desired width (0 = native).
-            height: Desired height (0 = native).
-            format: Image format ("raw" or "jpeg").
-
-        Returns an iterator of ImageFrame messages.
+            origin: (x, y, z) start position.
+            direction: (x, y, z) direction and magnitude.
+            color: (r, g, b, a) color values (0-1 range).
+            scale: Scale factor for visualization.
+            clear_previous: Clear previous debug draws before drawing.
+            timeout: RPC timeout in seconds.
+
+        Returns:
+            True if draw succeeded, False otherwise.
         """
-        raise NotImplementedError(
-            "stream_camera() is not yet confirmed working. "
-            "Remove this check if you want to test it."
+        timeout = timeout or self.timeout
+
+        arrow = self.pb.debug.DebugArrow(
+            origin=self.pb.debug.DebugVector3(x=origin[0], y=origin[1], z=origin[2]),
+            direction=self.pb.debug.DebugVector3(
+                x=direction[0], y=direction[1], z=direction[2]
+            ),
+            color=self.pb.debug.DebugColor(
+                r=color[0], g=color[1], b=color[2], a=color[3]
+            ),
+            scale=scale,
         )
-        if camera_id is not None:
-            request = self.pb.camera.StreamCameraRequest(
-                id=self.pb.common.EntityId(id=camera_id),
-                target_fps=target_fps,
-                width=width,
-                height=height,
-                format=format,
-            )
-        elif camera_name is not None:
-            request = self.pb.camera.StreamCameraRequest(
-                name=camera_name,
-                target_fps=target_fps,
-                width=width,
-                height=height,
-                format=format,
-            )
-        else:
-            raise ValueError("Either camera_id or camera_name must be provided")
 
-        return self.camera.StreamCamera(request)
+        request = self.pb.debug.DebugDrawRequest(
+            arrows=[arrow],
+            clear_previous=clear_previous,
+        )
 
-    def benchmark(
+        try:
+            resp = self.debug.Draw(request, timeout=timeout)
+            return resp.success
+        except Exception as e:
+            logger.debug(f"Debug draw failed: {e}")
+            return False
+
+    def draw_line(
         self,
-        duration_seconds: float = 5.0,
-        method: str = "get_observation",
-        print_results: bool = True,
-    ) -> "BenchmarkResult":
+        start: tuple[float, float, float],
+        end: tuple[float, float, float],
+        color: tuple[float, float, float, float] = (1.0, 1.0, 1.0, 1.0),
+        clear_previous: bool = False,
+        timeout: Optional[float] = None,
+    ) -> bool:
         """
-        Benchmark gRPC performance.
-
-        Measures actual FPS and latency for observation calls.
+        Draw a debug line in LuckyEngine.
 
         Args:
-            duration_seconds: How long to run the benchmark.
-            method: Which method to benchmark ("get_observation" or "stream_agent").
-            print_results: Whether to print results to console.
+            start: (x, y, z) start position.
+            end: (x, y, z) end position.
+            color: (r, g, b, a) color values (0-1 range).
+            clear_previous: Clear previous debug draws before drawing.
+            timeout: RPC timeout in seconds.
 
         Returns:
-            BenchmarkResult with FPS and latency statistics.
+            True if draw succeeded, False otherwise.
         """
-        if method not in ("get_observation", "stream_agent"):
-            raise ValueError(f"Unknown method: {method}. Use 'get_observation' or 'stream_agent'.")
-
-        latencies: list[float] = []
-        frame_times: list[float] = []
-        frame_count = 0
-        start_time = time.perf_counter()
-        last_frame_time = start_time
-
-        if method == "get_observation":
-            while time.perf_counter() - start_time < duration_seconds:
-                call_start = time.perf_counter()
-                self.get_observation()
-                call_end = time.perf_counter()
-
-                latencies.append((call_end - call_start) * 1000)  # ms
-                frame_times.append(call_end - last_frame_time)
-                last_frame_time = call_end
-                frame_count += 1
-
-        elif method == "stream_agent":
-            stream = self.stream_agent(target_fps=1000)  # Request max FPS
-            for frame in stream:
-                now = time.perf_counter()
-                if now - start_time >= duration_seconds:
-                    break
-
-                frame_times.append(now - last_frame_time)
-                last_frame_time = now
-                frame_count += 1
-
-        total_time = time.perf_counter() - start_time
-
-        # Calculate statistics
-        actual_fps = frame_count / total_time if total_time > 0 else 0
-
-        if len(frame_times) > 1:
-            # Remove first frame (warmup)
-            frame_times = frame_times[1:]
-            fps_from_frames = 1.0 / statistics.mean(frame_times) if frame_times else 0
-        else:
-            fps_from_frames = actual_fps
-
-        result = BenchmarkResult(
-            method=method,
-            duration_seconds=total_time,
-            frame_count=frame_count,
-            actual_fps=actual_fps,
-            avg_latency_ms=statistics.mean(latencies) if latencies else 0,
-            min_latency_ms=min(latencies) if latencies else 0,
-            max_latency_ms=max(latencies) if latencies else 0,
-            std_latency_ms=statistics.stdev(latencies) if len(latencies) > 1 else 0,
-            p50_latency_ms=statistics.median(latencies) if latencies else 0,
-            p99_latency_ms=sorted(latencies)[int(len(latencies) * 0.99)] if latencies else 0,
+        timeout = timeout or self.timeout
+
+        line = self.pb.debug.DebugLine(
+            start=self.pb.debug.DebugVector3(x=start[0], y=start[1], z=start[2]),
+            end=self.pb.debug.DebugVector3(x=end[0], y=end[1], z=end[2]),
+            color=self.pb.debug.DebugColor(
+                r=color[0], g=color[1], b=color[2], a=color[3]
+            ),
         )
 
-        if print_results:
-            print(f"\n{'=' * 50}")
-            print(f"Benchmark Results ({method})")
-            print(f"{'=' * 50}")
-            print(f"Duration:     {result.duration_seconds:.2f}s")
-            print(f"Frames:       {result.frame_count}")
-            print(f"Actual FPS:   {result.actual_fps:.1f}")
-            if latencies:
-                print(f"{'─' * 50}")
-                print(f"Latency (ms):")
-                print(f"  avg:        {result.avg_latency_ms:.2f}")
-                print(f"  min:        {result.min_latency_ms:.2f}")
-                print(f"  max:        {result.max_latency_ms:.2f}")
-                print(f"  std:        {result.std_latency_ms:.2f}")
-                print(f"  p50:        {result.p50_latency_ms:.2f}")
-                print(f"  p99:        {result.p99_latency_ms:.2f}")
-            print(f"{'=' * 50}\n")
-
-        return result
-
-
-@dataclass
-class BenchmarkResult:
-    """Results from a benchmark run."""
-
-    method: str
-    duration_seconds: float
-    frame_count: int
-    actual_fps: float
-    avg_latency_ms: float
-    min_latency_ms: float
-    max_latency_ms: float
-    std_latency_ms: float
-    p50_latency_ms: float
-    p99_latency_ms: float
+        request = self.pb.debug.DebugDrawRequest(
+            lines=[line],
+            clear_previous=clear_previous,
+        )
+
+        try:
+            resp = self.debug.Draw(request, timeout=timeout)
+            return resp.success
+        except Exception as e:
+            logger.debug(f"Debug draw failed: {e}")
+            return False