opentau 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

opentau/scripts/grpc/robot_inference_pb2_grpc.py

@@ -0,0 +1,210 @@
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+
+import grpc
+from opentau.scripts.grpc import robot_inference_pb2 as robot__inference__pb2
+
+GRPC_GENERATED_VERSION = "1.76.0"
+GRPC_VERSION = grpc.__version__
+_version_not_supported = False
+
+try:
+    from grpc._utilities import first_version_is_lower
+
+    _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION)
+except ImportError:
+    _version_not_supported = True
+
+if _version_not_supported:
+    raise RuntimeError(
+        f"The grpc package installed is at version {GRPC_VERSION},"
+        + " but the generated code in robot_inference_pb2_grpc.py depends on"
+        + f" grpcio>={GRPC_GENERATED_VERSION}."
+        + f" Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}"
+        + f" or downgrade your generated code using grpcio-tools<={GRPC_VERSION}."
+    )
+
+
+class RobotPolicyServiceStub:
+    """Service for robot policy inference"""
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.GetActionChunk = channel.unary_unary(
+            "/robot_inference.RobotPolicyService/GetActionChunk",
+            request_serializer=robot__inference__pb2.ObservationRequest.SerializeToString,
+            response_deserializer=robot__inference__pb2.ActionChunkResponse.FromString,
+            _registered_method=True,
+        )
+        self.StreamActionChunks = channel.stream_stream(
+            "/robot_inference.RobotPolicyService/StreamActionChunks",
+            request_serializer=robot__inference__pb2.ObservationRequest.SerializeToString,
+            response_deserializer=robot__inference__pb2.ActionChunkResponse.FromString,
+            _registered_method=True,
+        )
+        self.HealthCheck = channel.unary_unary(
+            "/robot_inference.RobotPolicyService/HealthCheck",
+            request_serializer=robot__inference__pb2.HealthCheckRequest.SerializeToString,
+            response_deserializer=robot__inference__pb2.HealthCheckResponse.FromString,
+            _registered_method=True,
+        )
+
+
+class RobotPolicyServiceServicer:
+    """Service for robot policy inference"""
+
+    def GetActionChunk(self, request, context):
+        """Get action chunk from observations (single request-response)"""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details("Method not implemented!")
+        raise NotImplementedError("Method not implemented!")
+
+    def StreamActionChunks(self, request_iterator, context):
+        """Streaming version for continuous inference (robot sends observations, server sends action chunks)"""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details("Method not implemented!")
+        raise NotImplementedError("Method not implemented!")
+
+    def HealthCheck(self, request, context):
+        """Health check"""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details("Method not implemented!")
+        raise NotImplementedError("Method not implemented!")
+
+
+def add_RobotPolicyServiceServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+        "GetActionChunk": grpc.unary_unary_rpc_method_handler(
+            servicer.GetActionChunk,
+            request_deserializer=robot__inference__pb2.ObservationRequest.FromString,
+            response_serializer=robot__inference__pb2.ActionChunkResponse.SerializeToString,
+        ),
+        "StreamActionChunks": grpc.stream_stream_rpc_method_handler(
+            servicer.StreamActionChunks,
+            request_deserializer=robot__inference__pb2.ObservationRequest.FromString,
+            response_serializer=robot__inference__pb2.ActionChunkResponse.SerializeToString,
+        ),
+        "HealthCheck": grpc.unary_unary_rpc_method_handler(
+            servicer.HealthCheck,
+            request_deserializer=robot__inference__pb2.HealthCheckRequest.FromString,
+            response_serializer=robot__inference__pb2.HealthCheckResponse.SerializeToString,
+        ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+        "robot_inference.RobotPolicyService", rpc_method_handlers
+    )
+    server.add_generic_rpc_handlers((generic_handler,))
+    server.add_registered_method_handlers("robot_inference.RobotPolicyService", rpc_method_handlers)
+
+
+# This class is part of an EXPERIMENTAL API.
+class RobotPolicyService:
+    """Service for robot policy inference"""
+
+    @staticmethod
+    def GetActionChunk(
+        request,
+        target,
+        options=(),
+        channel_credentials=None,
+        call_credentials=None,
+        insecure=False,
+        compression=None,
+        wait_for_ready=None,
+        timeout=None,
+        metadata=None,
+    ):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            "/robot_inference.RobotPolicyService/GetActionChunk",
+            robot__inference__pb2.ObservationRequest.SerializeToString,
+            robot__inference__pb2.ActionChunkResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True,
+        )
+
+    @staticmethod
+    def StreamActionChunks(
+        request_iterator,
+        target,
+        options=(),
+        channel_credentials=None,
+        call_credentials=None,
+        insecure=False,
+        compression=None,
+        wait_for_ready=None,
+        timeout=None,
+        metadata=None,
+    ):
+        return grpc.experimental.stream_stream(
+            request_iterator,
+            target,
+            "/robot_inference.RobotPolicyService/StreamActionChunks",
+            robot__inference__pb2.ObservationRequest.SerializeToString,
+            robot__inference__pb2.ActionChunkResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True,
+        )
+
+    @staticmethod
+    def HealthCheck(
+        request,
+        target,
+        options=(),
+        channel_credentials=None,
+        call_credentials=None,
+        insecure=False,
+        compression=None,
+        wait_for_ready=None,
+        timeout=None,
+        metadata=None,
+    ):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            "/robot_inference.RobotPolicyService/HealthCheck",
+            robot__inference__pb2.HealthCheckRequest.SerializeToString,
+            robot__inference__pb2.HealthCheckResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True,
+        )

opentau/scripts/grpc/server.py

@@ -0,0 +1,313 @@
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""gRPC server for robot policy inference on GPU.
+
+This server loads an ML policy model and serves inference requests from
+robots running ROS 2. Designed to run on a server with a ML GPU.
+
+Usage:
+    python src/opentau/scripts/grpc/server.py --config_path=/path/to/config.json \\
+        --server.port=50051 --server.max_workers=4
+"""
+
+import io
+import logging
+from concurrent import futures
+from dataclasses import asdict
+from pprint import pformat
+from typing import Iterator
+
+import numpy as np
+import torch
+from PIL import Image
+
+import grpc
+from opentau.configs import parser
+from opentau.configs.train import TrainPipelineConfig
+from opentau.policies.factory import get_policy_class
+from opentau.scripts.grpc import robot_inference_pb2, robot_inference_pb2_grpc
+from opentau.utils.random_utils import set_seed
+from opentau.utils.utils import (
+    attempt_torch_compile,
+    auto_torch_device,
+    init_logging,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class RobotPolicyServicer(robot_inference_pb2_grpc.RobotPolicyServiceServicer):
+    """gRPC servicer implementing the RobotPolicyService."""
+
+    def __init__(self, cfg: TrainPipelineConfig):
+        """Initialize the servicer with model and configuration.
+
+        Args:
+            cfg: Training pipeline configuration including policy settings.
+        """
+        self.cfg = cfg
+        self.device = auto_torch_device()
+        self.dtype = torch.bfloat16
+
+        logger.info(f"Initializing policy on device: {self.device}")
+
+        # Load the policy model
+        self._load_policy()
+
+    def _load_policy(self):
+        """Load the policy model from pretrained weights."""
+        logger.info(f"Loading policy from: {self.cfg.policy.pretrained_path}")
+
+        policy_class = get_policy_class(self.cfg.policy.type)
+        self.policy = policy_class.from_pretrained(self.cfg.policy.pretrained_path, config=self.cfg.policy)
+        self.policy.to(device=self.device, dtype=self.dtype)
+        self.policy.eval()
+        self.policy = attempt_torch_compile(self.policy, device_hint=self.device)
+        self.policy.reset()
+        logger.info("Policy loaded successfully")
+
+    def _decode_image(self, camera_image: robot_inference_pb2.CameraImage) -> torch.Tensor:
+        """Decode an image from the protobuf message.
+
+        Args:
+            camera_image: CameraImage protobuf message.
+
+        Returns:
+            Tensor of shape (1, C, H, W) normalized to [0, 1].
+        """
+        if camera_image.encoding in ["jpeg", "png"]:
+            # Decode compressed image
+            image = Image.open(io.BytesIO(camera_image.image_data))
+            image = image.convert("RGB")
+            image = image.resize(self.cfg.resolution[::-1])  # PIL uses (W, H)
+            image_array = np.array(image, dtype=np.float32) / 255.0
+        elif camera_image.encoding == "raw":
+            # Raw image data - assume it's already in the right shape
+            image_array = np.frombuffer(camera_image.image_data, dtype=np.float32)
+            # Reshape assuming square image with 3 channels
+            side = int(np.sqrt(len(image_array) / 3))
+            image_array = image_array.reshape(side, side, 3)
+            # Resize if needed
+            if (side, side) != self.cfg.resolution:
+                image = Image.fromarray((image_array * 255).astype(np.uint8))
+                image = image.resize(self.cfg.resolution[::-1])
+                image_array = np.array(image, dtype=np.float32) / 255.0
+        else:
+            raise ValueError(f"Unknown image encoding: {camera_image.encoding}")
+
+        # Convert to (C, H, W) tensor
+        image_tensor = torch.from_numpy(image_array).permute(2, 0, 1).unsqueeze(0)
+        return image_tensor.to(device=self.device, dtype=self.dtype)
+
+    def _prepare_observation(
+        self, request: robot_inference_pb2.ObservationRequest
+    ) -> dict[str, torch.Tensor]:
+        """Convert a protobuf observation request to the policy input format.
+
+        Args:
+            request: ObservationRequest protobuf message.
+
+        Returns:
+            Dictionary of tensors matching the policy's expected input format.
+        """
+        batch = {}
+
+        # Process camera images
+        img_is_pad = []
+        for i, camera_image in enumerate(request.images):
+            camera_name = f"camera{i}"
+            batch[camera_name] = self._decode_image(camera_image)
+            img_is_pad.append(False)
+
+        # Fill in missing cameras with zeros
+        for i in range(len(request.images), self.cfg.num_cams):
+            batch[f"camera{i}"] = torch.zeros(
+                (1, 3, *self.cfg.resolution),
+                dtype=self.dtype,
+                device=self.device,
+            )
+            img_is_pad.append(True)
+
+        batch["img_is_pad"] = torch.tensor([img_is_pad], dtype=torch.bool, device=self.device)
+
+        # Process robot state
+        if request.robot_state.state:
+            state = list(request.robot_state.state)
+            # Pad to max_state_dim if needed
+            if len(state) < self.cfg.max_state_dim:
+                state.extend([0.0] * (self.cfg.max_state_dim - len(state)))
+            batch["state"] = torch.tensor(
+                [state[: self.cfg.max_state_dim]],
+                dtype=self.dtype,
+                device=self.device,
+            )
+        else:
+            raise ValueError("Robot state is required but was not provided in the request")
+
+        # Process prompt
+        batch["prompt"] = [request.prompt] if request.prompt else [""]
+
+        return batch
+
+    def GetActionChunk(
+        self,
+        request: robot_inference_pb2.ObservationRequest,
+        context: grpc.ServicerContext,
+    ) -> robot_inference_pb2.ActionChunkResponse:
+        """Handle a single action chunk inference request.
+
+        Args:
+            request: ObservationRequest containing observations.
+            context: gRPC context.
+
+        Returns:
+            ActionChunkResponse containing the predicted action chunk.
+        """
+        import time
+
+        start_time = time.perf_counter()
+        response = robot_inference_pb2.ActionChunkResponse()
+        response.request_id = request.request_id
+        response.timestamp_ns = time.time_ns()
+
+        try:
+            # Prepare observation batch
+            batch = self._prepare_observation(request)
+
+            # Run inference
+            with torch.inference_mode():
+                action_chunk = self.policy.sample_actions(batch)
+                # action_chunk shape: (n_action_steps, batch_size=1, action_dim)
+                # Remove batch dimension and convert to numpy
+                action_chunk = action_chunk.squeeze(1).to("cpu", torch.float32).numpy()
+
+            # Populate 2D action chunk structure
+            for action_vector in action_chunk:
+                action_vec_msg = robot_inference_pb2.ActionVector()
+                action_vec_msg.values.extend(action_vector.tolist())
+                response.action_chunk.append(action_vec_msg)
+
+        except ValueError as e:
+            # Invalid request (e.g., missing required fields)
+            logger.error(f"Invalid request: {e}")
+            context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(e))
+
+        except Exception as e:
+            # Unexpected error during inference
+            logger.exception("Error during inference")
+            context.abort(grpc.StatusCode.INTERNAL, f"Inference error: {e}")
+
+        response.inference_time_ms = (time.perf_counter() - start_time) * 1000
+        return response
+
+    def StreamActionChunks(
+        self,
+        request_iterator: Iterator[robot_inference_pb2.ObservationRequest],
+        context: grpc.ServicerContext,
+    ) -> Iterator[robot_inference_pb2.ActionChunkResponse]:
+        """Handle streaming action chunk inference requests.
+
+        Args:
+            request_iterator: Iterator of ObservationRequest messages.
+            context: gRPC context.
+
+        Yields:
+            ActionChunkResponse messages for each observation.
+        """
+        for request in request_iterator:
+            if context.is_active():
+                yield self.GetActionChunk(request, context)
+            else:
+                break
+
+    def HealthCheck(
+        self,
+        request: robot_inference_pb2.HealthCheckRequest,
+        context: grpc.ServicerContext,
+    ) -> robot_inference_pb2.HealthCheckResponse:
+        """Check server health and GPU status.
+
+        Args:
+            request: HealthCheckRequest message.
+            context: gRPC context.
+
+        Returns:
+            HealthCheckResponse with server status.
+        """
+        response = robot_inference_pb2.HealthCheckResponse()
+        response.healthy = True
+        response.status = "Server is running"
+        response.model_name = self.cfg.policy.type
+        response.device = str(self.device)
+
+        if torch.cuda.is_available():
+            response.gpu_memory_used_gb = torch.cuda.memory_allocated() / 1e9
+            response.gpu_memory_total_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
+
+        return response
+
+
+def serve(cfg: TrainPipelineConfig):
+    """Start the gRPC server.
+
+    Args:
+        cfg: Training pipeline configuration including server settings.
+    """
+    server_cfg = cfg.server
+    server = grpc.server(
+        futures.ThreadPoolExecutor(max_workers=server_cfg.max_workers),
+        options=[
+            ("grpc.max_send_message_length", server_cfg.max_send_message_length),
+            ("grpc.max_receive_message_length", server_cfg.max_receive_message_length),
+        ],
+    )
+
+    servicer = RobotPolicyServicer(cfg)
+    robot_inference_pb2_grpc.add_RobotPolicyServiceServicer_to_server(servicer, server)
+
+    server.add_insecure_port(f"[::]:{server_cfg.port}")
+    server.start()
+
+    logger.info(f"Server started on port {server_cfg.port}")
+    logger.info(f"Policy: {cfg.policy.type}")
+    logger.info(f"Device: {servicer.device}")
+    logger.info(f"Max workers: {server_cfg.max_workers}")
+    logger.info("Waiting for requests...")
+
+    try:
+        server.wait_for_termination()
+    except KeyboardInterrupt:
+        logger.info("Shutting down server...")
+        server.stop(grace=5)
+
+
+@parser.wrap()
+def server_main(cfg: TrainPipelineConfig):
+    """Main entry point for the gRPC server.
+
+    Args:
+        cfg: Training pipeline configuration parsed from CLI/config file.
+    """
+    logging.info(pformat(asdict(cfg)))
+
+    if cfg.seed is not None:
+        set_seed(cfg.seed)
+
+    serve(cfg)
+
+
+if __name__ == "__main__":
+    init_logging()
+    server_main()

opentau/scripts/launch.py

@@ -18,10 +18,6 @@ import sys
 from pathlib import Path
 from types import ModuleType
 
-import opentau.scripts.eval as eval_script
-import opentau.scripts.export_to_onnx as export_script
-import opentau.scripts.train as train_script
-
 
 def launch(script_module: ModuleType, description: str, use_accelerate: bool = True):
     """Generic launcher for OpenTau scripts using Accelerate or Python."""
@@ -68,12 +64,24 @@ def launch(script_module: ModuleType, description: str, use_accelerate: bool = T
 
 
 def train():
+    import opentau.scripts.train as train_script
+
     launch(train_script, "Launch OpenTau training with Accelerate")
 
 
 def eval():
+    import opentau.scripts.eval as eval_script
+
     launch(eval_script, "Launch OpenTau evaluation with Accelerate")
 
 
 def export():
+    import opentau.scripts.export_to_onnx as export_script
+
     launch(export_script, "Launch OpenTau ONNX export", use_accelerate=False)
+
+
+def visualize():
+    import opentau.scripts.visualize_dataset as visualize_script
+
+    launch(visualize_script, "Launch OpenTau visualization", use_accelerate=False)