From fc7b7bc6946583e7c8c8fcb48bcbd83e8bc714d9 Mon Sep 17 00:00:00 2001
From: Ury Zhilinsky <ury@physicalintelligence.company>
Date: Tue, 24 Dec 2024 19:38:19 -0800
Subject: [PATCH] Fix simple_client docker flow and documentation

---
 examples/simple_client/Dockerfile  |  2 +-
 examples/simple_client/README.md   | 12 +++++++++---
 examples/simple_client/compose.yml |  2 ++
 examples/simple_client/main.py     | 24 ++++++++++++++++++------
 4 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/examples/simple_client/Dockerfile b/examples/simple_client/Dockerfile
index ed3fcf6..eebca39 100644
--- a/examples/simple_client/Dockerfile
+++ b/examples/simple_client/Dockerfile
@@ -29,4 +29,4 @@ 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"]
+CMD /bin/bash -c "source /.venv/bin/activate && python examples/simple_client/main.py $SERVER_ARGS"
diff --git a/examples/simple_client/README.md b/examples/simple_client/README.md
index 7bb4ad5..c246752 100644
--- a/examples/simple_client/README.md
+++ b/examples/simple_client/README.md
@@ -2,10 +2,16 @@
 
 A minimal client that sends observations to the server and prints the inference rate.
 
+You can specifiy which runtime environment to use using the `--env` flag. You can see the available options by running:
+
+```bash
+uv run examples/simple_client/main.py --help
+```
+
 ## With Docker
 
 ```bash
-export SERVER_ARGS="--example aloha"
+export SERVER_ARGS="--env ALOHA_SIM"
 docker compose -f examples/simple_client/compose.yml up --build
 ```
 
@@ -14,11 +20,11 @@ docker compose -f examples/simple_client/compose.yml up --build
 Terminal window 1:
 
 ```bash
-uv run examples/simple_client/main.py
+uv run examples/simple_client/main.py --env DROID
 ```
 
 Terminal window 2:
 
 ```bash
-uv run scripts/serve_policy.py
+uv run scripts/serve_policy.py --env DROID
 ```
\ No newline at end of file
diff --git a/examples/simple_client/compose.yml b/examples/simple_client/compose.yml
index 92a479b..3401b3f 100644
--- a/examples/simple_client/compose.yml
+++ b/examples/simple_client/compose.yml
@@ -13,6 +13,8 @@ services:
     network_mode: host
     volumes:
       - $PWD:/app
+    environment:
+      - SERVER_ARGS  
 
   openpi_server:
     image: openpi_server
diff --git a/examples/simple_client/main.py b/examples/simple_client/main.py
index 3f36073..0b4f265 100644
--- a/examples/simple_client/main.py
+++ b/examples/simple_client/main.py
@@ -1,4 +1,5 @@
 import dataclasses
+import enum
 import logging
 import time
 
@@ -7,21 +8,32 @@ from openpi_client import websocket_client_policy as _websocket_client_policy
 import tyro
 
 
+class EnvMode(enum.Enum):
+    """Supported environments."""
+
+    ALOHA = "aloha"
+    ALOHA_SIM = "aloha_sim"
+    DROID = "droid"
+    CALVIN = "calvin"
+    LIBERO = "libero"
+
+
 @dataclasses.dataclass
 class Args:
     host: str = "0.0.0.0"
     port: int = 8000
 
-    example: str = "droid"
+    env: EnvMode = EnvMode.ALOHA_SIM
 
 
 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]
+        EnvMode.ALOHA: _random_observation_aloha,
+        EnvMode.ALOHA_SIM: _random_observation_aloha,
+        EnvMode.DROID: _random_observation_droid,
+        EnvMode.CALVIN: _random_observation_calvin,
+        EnvMode.LIBERO: _random_observation_libero,
+    }[args.env]
 
     policy = _websocket_client_policy.WebsocketClientPolicy(
         host=args.host,