remote_inference.md

Running openpi models remotely

We provide utilities for running openpi models remotely. This is useful for running inference on more powerful GPUs off-robot, and also helps keep the robot and policy environments separate (and e.g. avoid dependency hell with robot software).

Starting a remote policy server

To start a remote policy server, you can simply run the following command:

uv run scripts/serve_policy.py --env=[DROID | ALOHA | LIBERO]

The env argument specifies which $\pi_0$ checkpoint should be loaded. Under the hood, this script will execute a command like the following, which you can use to start a policy server, e.g. for checkpoints you trained yourself (here an example for the DROID environment):

uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi0_fast_droid --policy.dir=s3://openpi-assets/checkpoints/pi0_fast_droid

This will start a policy server that will serve the policy specified by the config and dir arguments. The policy will be served on the specified port (default: 8000).

Querying the remote policy server from your robot code

We provide a client utility with minimal dependencies that you can easily embed into any robot codebase.

First, install the openpi-client package in your robot environment:

cd $OPENPI_ROOT/packages/openpi-client
pip install -e .

Then, you can use the client to query the remote policy server from your robot code. Here's an example of how to do this:

from openpi_client import websocket_client_policy

policy_client = websocket_client_policy.WebsocketClientPolicy(host="10.32.255.0", port=8000)
action_chunk = policy_client.infer(example)["actions"]

Here, the host and port arguments specify the IP address and port of the remote policy server. You can also specify these as command-line arguments to your robot code, or hard-code them in your robot codebase. The example is a dictionary of observations and the prompt, following the specification of the policy inputs for the policy you are serving. We have concrete examples of how to construct this dictionary for different environments in the simple client example.