Fix simple_client docker flow and documentation

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"

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
 ```

examples/simple_client/compose.yml

@@ -13,6 +13,8 @@ services:
     network_mode: host
     volumes:
       - $PWD:/app
+    environment:
+      - SERVER_ARGS  
 
   openpi_server:
     image: openpi_server

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,