Initial commit

examples/simple_client/Dockerfile

@@ -0,0 +1,32 @@
+# Dockerfile for the simple client.
+
+# Build the container:
+# docker build . -t simple_client -f examples/simple_client/Dockerfile
+
+# Run the container:
+# docker run --rm -it --network=host -v .:/app simple_client /bin/bash
+
+FROM python:3.7-slim
+COPY --from=ghcr.io/astral-sh/uv:0.5.1 /uv /uvx /bin/
+
+WORKDIR /app
+
+# Copy from the cache instead of linking since it's a mounted volume
+ENV UV_LINK_MODE=copy
+
+# Write the virtual environment outside of the project directory so it doesn't
+# leak out of the container when we mount the application code.
+ENV UV_PROJECT_ENVIRONMENT=/.venv
+
+# Copy the requirements files so we can install dependencies.
+# The rest of the project is mounted as a volume, so we don't need to rebuild on changes.
+# This strategy is best for development-style usage.
+COPY ./examples/simple_client/requirements.txt /tmp/requirements.txt
+COPY ./packages/openpi-client/pyproject.toml /tmp/openpi-client/pyproject.toml
+
+# Install python dependencies.
+RUN uv venv --python 3.7 $UV_PROJECT_ENVIRONMENT
+RUN uv pip sync /tmp/requirements.txt /tmp/openpi-client/pyproject.toml
+ENV PYTHONPATH=/app:/app/src:/app/packages/openpi-client/src
+
+CMD ["/bin/bash", "-c", "source /.venv/bin/activate && python examples/simple_client/main.py"]

examples/simple_client/README.md

@@ -0,0 +1,24 @@
+# Simple Client
+
+A minimal client that sends observations to the server and prints the inference rate.
+
+## With Docker
+
+```bash
+export SERVER_ARGS="--example aloha"
+docker compose -f examples/simple_client/compose.yml up --build
+```
+
+## Without Docker
+
+Terminal window 1:
+
+```bash
+uv run examples/simple_client/main.py
+```
+
+Terminal window 2:
+
+```bash
+uv run scripts/serve_policy.py
+```

examples/simple_client/compose.yml

@@ -0,0 +1,37 @@
+# Run with:
+# docker compose -f examples/simple_client/compose.yml up --build
+services:
+  runtime:
+    image: simple_client
+    depends_on:
+      - openpi_server
+    build:
+      context: ../..
+      dockerfile: examples/simple_client/Dockerfile
+    init: true
+    tty: true
+    network_mode: host
+    volumes:
+      - $PWD:/app
+
+  openpi_server:
+    image: openpi_server
+    build:
+      context: ../..
+      dockerfile: scripts/serve_policy.Dockerfile
+    init: true
+    tty: true
+    network_mode: host
+    volumes:
+      - $PWD:/app
+    environment:
+      - SERVER_ARGS
+
+    # Comment out this block if not running on a machine with GPUs.
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [gpu]

examples/simple_client/main.py

@@ -0,0 +1,81 @@
+import dataclasses
+import logging
+import time
+
+import numpy as np
+from openpi_client import websocket_client_policy as _websocket_client_policy
+import tyro
+
+
+@dataclasses.dataclass
+class Args:
+    host: str = "0.0.0.0"
+    port: int = 8000
+
+    example: str = "droid"
+
+
+def main(args: Args) -> None:
+    obs_fn = {
+        "aloha": _random_observation_aloha,
+        "droid": _random_observation_droid,
+        "calvin": _random_observation_calvin,
+        "libero": _random_observation_libero,
+    }[args.example]
+
+    policy = _websocket_client_policy.WebsocketClientPolicy(
+        host=args.host,
+        port=args.port,
+    )
+
+    # Send 1 observation to make sure the model is loaded.
+    policy.infer(obs_fn())
+
+    start = time.time()
+    for _ in range(100):
+        policy.infer(obs_fn())
+    end = time.time()
+
+    print(f"Total time taken: {end - start}")
+    # Note that each inference returns many action chunks.
+    print(f"Inference rate: {100 / (end - start)} Hz")
+
+
+def _random_observation_aloha() -> dict:
+    return {
+        "qpos": np.ones((14,)),
+        "image": np.random.rand(4, 3, 480, 640).astype(np.float32),
+    }
+
+
+def _random_observation_droid() -> dict:
+    return {
+        "observation/exterior_image_1_left": np.random.randint(256, size=(224, 224, 3), dtype=np.uint8),
+        "observation/wrist_image_left": np.random.randint(256, size=(224, 224, 3), dtype=np.uint8),
+        "observation/joint_position": np.random.rand(7),
+        "observation/gripper_position": np.random.rand(1),
+        "prompt": "do something",
+    }
+
+
+def _random_observation_calvin() -> dict:
+    return {
+        "observation/state": np.random.rand(15),
+        "observation/rgb_static": np.random.rand(4, 3, 480, 640).astype(np.float32),
+        "observation/rgb_gripper": np.random.rand(4, 3, 480, 640).astype(np.float32),
+        "prompt": "do something",
+    }
+
+
+def _random_observation_libero() -> dict:
+    return {
+        "observation/state": np.random.rand(8),
+        "observation/image": np.random.rand(4, 3, 480, 640).astype(np.float32),
+        "observation/wrist_image": np.random.rand(4, 3, 480, 640).astype(np.float32),
+        "prompt": "do something",
+    }
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    tyro.cli(main)

examples/simple_client/requirements.in

@@ -0,0 +1,2 @@
+numpy
+tyro

examples/simple_client/requirements.txt

@@ -0,0 +1,27 @@
+# This file was autogenerated by uv via the following command:
+#    uv pip compile examples/simple_client/requirements.in -o examples/simple_client/requirements.txt --python-version 3.7
+backports-cached-property==1.0.2
+    # via tyro
+docstring-parser==0.16
+    # via tyro
+eval-type-backport==0.1.3
+    # via tyro
+markdown-it-py==2.2.0
+    # via rich
+mdurl==0.1.2
+    # via markdown-it-py
+numpy==1.21.6
+    # via -r examples/simple_client/requirements.in
+pygments==2.17.2
+    # via rich
+rich==13.8.1
+    # via tyro
+shtab==1.7.1
+    # via tyro
+typing-extensions==4.7.1
+    # via
+    #   markdown-it-py
+    #   rich
+    #   tyro
+tyro==0.9.1
+    # via -r examples/simple_client/requirements.in