Skip to content

Commit 1f62bd6

Browse files
kanghui0204hkang-nvxyao-nv
authored
[FEATURE] remote policy with server/client (#394)
## Summary old MR is [MR 310](#310) Hi Team, I added a remote server/client mechanism for IsaacLab-Arena that allows the IsaacLab environment and the policy model environment to run separately. Alex and I wrote a design document([design doc](https://docs.google.com/document/d/1uNcYswolxPD4PpHgvCapYuqwZ2QhzSnlJB9KRMFMQbQ/edit?tab=t.0)) for this feature, which includes our discussions about the design. Below is an overview of this PR. ### Functionality Previously, the IsaacLab-Arena pipeline ran entirely in a single process: **IsaacLab env** → **Env policy class** (e.g., Gr00tClosedloopPolicy) → **local policy model** (e.g., Gr00tPolicy), all in one environment. This PR adds remote policy/server–client support and decouples the env policy class from the local policy model. The new pipeline becomes: **Client**: **IsaacLab env** → **Env policy class** (e.g., ActionChunkingClientPolicy) **Server**: **VLA/VLN policy model** (e.g., Gr00tRemoteServerSidePolicy) The client and server exchange observations and actions via sockets. The server/client implementation lives in `isaaclab_arena/remote_policy` inside IsaacLab-Arena. Users can copy this directory and import it on the server side directly, without needing to install it as a separate package. --------- Co-authored-by: Hui Kang <hkang@nvidia.com> Co-authored-by: Xinjie Yao <xyao@nvidia.com>
1 parent a239d00 commit 1f62bd6

34 files changed

+3449
-329
lines changed

‎CONTRIBUTORS.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,3 +32,4 @@ Guidelines for modifications:
3232
* Rebecca Zhang <rebeccazhang0707, rebeccaz@nvidia.com> (NVIDIA)
3333
* Shiwei Sheng <shiweis, shiweis@nvidia.com> (NVIDIA)
3434
* Weihua Zhang <yami007007-weihuaz, weihuaz@nvidia.com> (NVIDIA)
35+
* Hui Kang <hkang@nvidia.com> (NVIDIA)

‎docker/Dockerfile.gr00t_server‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
FROM nvcr.io/nvidia/pytorch:24.07-py3
2+
3+
ARG WORKDIR="/workspace"
4+
ARG GROOT_DEPS_GROUP="base"
5+
ENV WORKDIR=${WORKDIR}
6+
ENV GROOT_DEPS_GROUP=${GROOT_DEPS_GROUP}
7+
WORKDIR "${WORKDIR}"
8+
9+
RUN apt-get update && apt-get install -y \
10+
git \
11+
git-lfs \
12+
cmake \
13+
&& rm -rf /var/lib/apt/lists/*
14+
15+
RUN pip install --upgrade pip
16+
17+
COPY ./submodules/Isaac-GR00T ${WORKDIR}/submodules/Isaac-GR00T
18+
19+
COPY docker/setup/install_gr00t_deps.sh /tmp/install_gr00t_deps.sh
20+
RUN chmod +x /tmp/install_gr00t_deps.sh && \
21+
/tmp/install_gr00t_deps.sh --server && \
22+
rm -f /tmp/install_gr00t_deps.sh
23+
24+
RUN pip install --no-cache-dir --upgrade "opencv-python-headless==4.8.0.74"
25+
26+
COPY isaaclab_arena/remote_policy ${WORKDIR}/isaaclab_arena/remote_policy
27+
COPY isaaclab_arena_gr00t ${WORKDIR}/isaaclab_arena_gr00t
28+
COPY isaaclab_arena_g1 ${WORKDIR}/isaaclab_arena_g1
29+
30+
RUN pip install --no-cache-dir pyzmq msgpack
31+
32+
ENV PYTHONPATH=${WORKDIR}
33+
34+
ENTRYPOINT ["python", "-u", "-m", "isaaclab_arena.remote_policy.remote_policy_server_runner"]

‎docker/run_gr00t_server.sh‎

Lines changed: 243 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,243 @@
1+
#!/usr/bin/env bash
2+
set -euo pipefail
3+
4+
# -------------------------
5+
# User-configurable defaults
6+
# -------------------------
7+
8+
# Default mount directories on the host machine
9+
DATASETS_DIR="${DATASETS_DIR:-$HOME/datasets}"
10+
MODELS_DIR="${MODELS_DIR:-$HOME/models}"
11+
EVAL_DIR="${EVAL_DIR:-$HOME/eval}"
12+
13+
# Docker image name and tag for the GR00T policy server
14+
DOCKER_IMAGE_NAME="${DOCKER_IMAGE_NAME:-gr00t_policy_server}"
15+
DOCKER_VERSION_TAG="${DOCKER_VERSION_TAG:-latest}"
16+
17+
# Rebuild controls
18+
FORCE_REBUILD="${FORCE_REBUILD:-false}"
19+
NO_CACHE=""
20+
21+
# Server parameters (can also be overridden via environment variables)
22+
HOST="${HOST:-0.0.0.0}"
23+
PORT="${PORT:-5555}"
24+
API_TOKEN="${API_TOKEN:-}"
25+
TIMEOUT_MS="${TIMEOUT_MS:-5000}"
26+
POLICY_TYPE="${POLICY_TYPE:-gr00t_closedloop}"
27+
POLICY_CONFIG_YAML_PATH="${POLICY_CONFIG_YAML_PATH:-/workspace/isaaclab_arena_gr00t/gr1_manip_gr00t_closedloop_config.yaml}"
28+
29+
# GPU selection for docker --gpus (can also be overridden via environment variables)
30+
# Examples:
31+
# all -> use all GPUs
32+
# 1 -> use 1 GPU (count)
33+
# "device=0" -> use GPU 0
34+
# "device=0,1" -> use GPU 0 and 1
35+
GPUS="${GPUS:-all}"
36+
37+
# -------------------------
38+
# Help message
39+
# -------------------------
40+
usage() {
41+
script_name=$(basename "$0")
42+
cat <<EOF
43+
Helper script to build and run the GR00T policy server Docker environment.
44+
45+
Usage:
46+
$script_name [options] [-- server-args...]
47+
48+
Options (Docker / paths; env vars with the same name take precedence):
49+
-v Verbose output (set -x).
50+
-d <datasets directory> Path to datasets on the host. Default: "$DATASETS_DIR".
51+
-m <models directory> Path to models on the host. Default: "$MODELS_DIR".
52+
-e <eval directory> Path to evaluation data on the host. Default: "$EVAL_DIR".
53+
-n <docker name> Docker image name. Default: "$DOCKER_IMAGE_NAME".
54+
-g <gpus> GPU selection for docker --gpus. Default: "all".
55+
Examples: "all", "1", "device=0", "device=0,1".
56+
-r Force rebuilding of the Docker image.
57+
-R Force rebuilding of the Docker image, without cache.
58+
59+
Server-specific options (passed through to the policy server entrypoint):
60+
--host HOST
61+
--port PORT
62+
--api_token TOKEN
63+
--timeout_ms MS
64+
--policy_type TYPE
65+
--policy_config_yaml_path PATH
66+
67+
Examples:
68+
# Minimal: use defaults, just build & run server
69+
bash $script_name
70+
71+
# Custom models directory, port and single GPU (GPU 0)
72+
bash $script_name -m /data/models -g "device=0" --port 6000 --api_token MY_TOKEN
73+
74+
# Custom image name, force rebuild, datasets/eval mounts, and multiple GPUs
75+
bash $script_name -n gr00t_server -r \\
76+
-d /data/datasets -m /data/models -e /data/eval \\
77+
-g "device=0,1" \\
78+
--policy_type isaaclab_arena_gr00t.policy.gr00t_remote_policy.Gr00tRemoteServerSidePolicy \\
79+
--policy_config_yaml_path isaaclab_arena_gr00t/policy/config/gr1_manip_gr00t_closedloop_config.yaml
80+
EOF
81+
}
82+
83+
# -------------------------
84+
# Parse docker/path options (short flags, like run_docker.sh)
85+
# -------------------------
86+
DOCKER_ARGS_DONE=false
87+
SERVER_ARGS=()
88+
89+
while [[ $# -gt 0 ]]; do
90+
if [ "$DOCKER_ARGS_DONE" = false ]; then
91+
case "$1" in
92+
-v)
93+
# Enable verbose mode for debugging
94+
set -x
95+
shift 1
96+
;;
97+
-d)
98+
# Set host datasets directory
99+
DATASETS_DIR="$2"
100+
shift 2
101+
;;
102+
-m)
103+
# Set host models directory
104+
MODELS_DIR="$2"
105+
shift 2
106+
;;
107+
-e)
108+
# Set host eval directory
109+
EVAL_DIR="$2"
110+
shift 2
111+
;;
112+
-n)
113+
# Set Docker image name
114+
DOCKER_IMAGE_NAME="$2"
115+
shift 2
116+
;;
117+
-g)
118+
# Set GPU selection for docker --gpus
119+
GPUS="$2"
120+
shift 2
121+
;;
122+
-r)
123+
# Force rebuild of Docker image
124+
FORCE_REBUILD="true"
125+
shift 1
126+
;;
127+
-R)
128+
# Force rebuild of Docker image, without cache
129+
FORCE_REBUILD="true"
130+
NO_CACHE="--no-cache"
131+
shift 1
132+
;;
133+
-h|--help)
134+
usage
135+
exit 0
136+
;;
137+
--host|--port|--api_token|--timeout_ms|--policy_type|--policy_config_yaml_path)
138+
# From here on, treat everything as server args and stop parsing docker flags
139+
DOCKER_ARGS_DONE=true
140+
SERVER_ARGS+=("$1")
141+
shift 1
142+
;;
143+
--*)
144+
# Unknown long option at docker level -> treat as server arg
145+
DOCKER_ARGS_DONE=true
146+
SERVER_ARGS+=("$1")
147+
shift 1
148+
;;
149+
*)
150+
# Anything else -> treat as server arg
151+
DOCKER_ARGS_DONE=true
152+
SERVER_ARGS+=("$1")
153+
shift 1
154+
;;
155+
esac
156+
else
157+
# Additional server arguments after docker/path args
158+
SERVER_ARGS+=("$1")
159+
shift 1
160+
fi
161+
done
162+
163+
# If no server args were passed, use defaults
164+
if [ ${#SERVER_ARGS[@]} -eq 0 ]; then
165+
SERVER_ARGS=(
166+
--host "${HOST}"
167+
--port "${PORT}"
168+
--api_token "${API_TOKEN}"
169+
--timeout_ms "${TIMEOUT_MS}"
170+
--policy_type "${POLICY_TYPE}"
171+
--policy_config_yaml_path "${POLICY_CONFIG_YAML_PATH}"
172+
)
173+
fi
174+
175+
echo "Host paths:"
176+
echo " DATASETS_DIR = ${DATASETS_DIR}"
177+
echo " MODELS_DIR = ${MODELS_DIR}"
178+
echo " EVAL_DIR = ${EVAL_DIR}"
179+
echo "Docker image:"
180+
echo " ${DOCKER_IMAGE_NAME}:${DOCKER_VERSION_TAG}"
181+
echo "GPU:"
182+
echo " --gpus ${GPUS}"
183+
echo "Rebuild:"
184+
echo " FORCE_REBUILD = ${FORCE_REBUILD}, NO_CACHE = '${NO_CACHE}'"
185+
echo "Server args:"
186+
printf ' %q ' "${SERVER_ARGS[@]}"; echo
187+
188+
# -------------------------
189+
# 1) Build the Docker image
190+
# -------------------------
191+
192+
IMAGE_TAG_FULL="${DOCKER_IMAGE_NAME}:${DOCKER_VERSION_TAG}"
193+
194+
# 1) Decide whether to build
195+
SHOULD_BUILD=false
196+
197+
if [ "${FORCE_REBUILD}" = "true" ]; then
198+
# -r or -R: force rebuild
199+
SHOULD_BUILD=true
200+
else
201+
# Without force flag: only build if the image does not exist locally
202+
if [ -z "$(docker images -q "${IMAGE_TAG_FULL}")" ]; then
203+
SHOULD_BUILD=true
204+
fi
205+
fi
206+
207+
# 2) Build or skip
208+
if [ "${SHOULD_BUILD}" = "true" ]; then
209+
echo "Building Docker image ${IMAGE_TAG_FULL}..."
210+
docker build \
211+
${NO_CACHE} \
212+
-f docker/Dockerfile.gr00t_server \
213+
-t "${IMAGE_TAG_FULL}" \
214+
.
215+
else
216+
echo "Docker image ${IMAGE_TAG_FULL} already exists. Skipping rebuild."
217+
echo "Use -r or -R to force rebuilding the image."
218+
fi
219+
220+
# -------------------------
221+
# 2) Run the container
222+
# -------------------------
223+
224+
DOCKER_RUN_ARGS=(
225+
--rm
226+
--gpus "${GPUS}"
227+
--net host
228+
--name gr00t_policy_server_container
229+
-v "${MODELS_DIR}":/models
230+
)
231+
232+
# Only mount datasets / eval if the directories exist on host
233+
if [ -d "${DATASETS_DIR}" ]; then
234+
DOCKER_RUN_ARGS+=(-v "${DATASETS_DIR}":/datasets)
235+
fi
236+
237+
if [ -d "${EVAL_DIR}" ]; then
238+
DOCKER_RUN_ARGS+=(-v "${EVAL_DIR}":/eval)
239+
fi
240+
241+
docker run "${DOCKER_RUN_ARGS[@]}" \
242+
"${IMAGE_TAG_FULL}" \
243+
"${SERVER_ARGS[@]}"

0 commit comments

Comments
 (0)