Initial commit

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)