Add Wall-X serving and Turtle2 TCP WebSocket bridge
Pre-commit / pre-commit (push) Canceled after 0s
Pre-commit / pre-commit (push) Canceled after 0s
This commit is contained in:
Executable
+342
@@ -0,0 +1,342 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
Usage:
|
||||
bash scripts/run_serving_rtc.sh --checkpoint-path /path/to/checkpoint [options]
|
||||
|
||||
Required:
|
||||
--checkpoint-path PATH Checkpoint directory or checkpoint file.
|
||||
|
||||
Common options:
|
||||
--train-config-path PATH Training config used by the checkpoint.
|
||||
--port PORT WebSocket port. Default: 32195.
|
||||
--host HOST Bind host. Default: 0.0.0.0.
|
||||
--env X2ROBOT|LIBERO Serving environment. Default: X2ROBOT.
|
||||
--cuda-id ID Sets CUDA_VISIBLE_DEVICES. Default: 0.
|
||||
--image-passing-mode MODE base64 or numpy. Default: base64.
|
||||
--action-horizon N Model action horizon. Default: 32.
|
||||
--robot-type TYPE desktop, turtle, or ex001. Default: desktop.
|
||||
--raw-actions Alias for --no-serialize-actions.
|
||||
--serialize-actions Return robot-serialized actions.
|
||||
--no-serialize-actions Return raw model action chunks. Default.
|
||||
--max-batch-size N Enable dynamic batching.
|
||||
--enable-cuda-graph Enable CUDA graph in the serving runtime.
|
||||
--enable-experimental-engine Enable the experimental inference engine.
|
||||
--rtc-execution-horizon N Model steps between RTC replans. Default: 6.
|
||||
--rtc-max-guidance-weight F RTC guidance clamp. Default: 10.0.
|
||||
--rtc-prefix-schedule NAME zeros, ones, linear, or exp. Default: linear.
|
||||
--debug Enable debug logging.
|
||||
--dry-run Print the command without running it.
|
||||
|
||||
Additional arguments after "--" are forwarded to launch_serving_rtc.py, for example:
|
||||
bash scripts/run_serving_rtc.sh --checkpoint-path /ckpt -- \
|
||||
--model-config.norm-key libero_all
|
||||
|
||||
Environment variables can also be used, e.g. CHECKPOINT_PATH,
|
||||
TRAIN_CONFIG_PATH, PORT, CUDA_ID, ACTION_HORIZON, WALLX_ENV.
|
||||
EOF
|
||||
}
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
SOURCE_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
|
||||
|
||||
if [[ -d "${SOURCE_ROOT}/wall_x" ]]; then
|
||||
export PYTHONPATH="${SOURCE_ROOT}:${PYTHONPATH:-}"
|
||||
fi
|
||||
|
||||
PYTHON_BIN="${PYTHON_BIN:-python}"
|
||||
CHECKPOINT_PATH="${CHECKPOINT_PATH:-}"
|
||||
TRAIN_CONFIG_PATH="${TRAIN_CONFIG_PATH:-}"
|
||||
PORT="${PORT:-32195}"
|
||||
HOST="${HOST:-0.0.0.0}"
|
||||
WALLX_ENV="${WALLX_ENV:-X2ROBOT}"
|
||||
CUDA_ID="${CUDA_ID:-0}"
|
||||
IMAGE_PASSING_MODE="${IMAGE_PASSING_MODE:-base64}"
|
||||
ACTION_HORIZON="${ACTION_HORIZON:-32}"
|
||||
ROBOT_TYPE="${ROBOT_TYPE:-desktop}"
|
||||
ROBOT_ACTION_INTERPOLATE_MULTIPLIER="${ROBOT_ACTION_INTERPOLATE_MULTIPLIER:-1}"
|
||||
ROBOT_ACTION_END_RATIO="${ROBOT_ACTION_END_RATIO:-1.0}"
|
||||
MODEL_DEVICE="${MODEL_DEVICE:-cuda}"
|
||||
MAX_BATCH_SIZE="${MAX_BATCH_SIZE:-}"
|
||||
DEFAULT_PROMPT="${DEFAULT_PROMPT:-}"
|
||||
SERIALIZE_ACTIONS="${SERIALIZE_ACTIONS:-0}"
|
||||
ENABLE_CUDA_GRAPH="${ENABLE_CUDA_GRAPH:-0}"
|
||||
ENABLE_EXPERIMENTAL_ENGINE="${ENABLE_EXPERIMENTAL_ENGINE:-0}"
|
||||
RTC_EXECUTION_HORIZON="${RTC_EXECUTION_HORIZON:-6}"
|
||||
RTC_MAX_GUIDANCE_WEIGHT="${RTC_MAX_GUIDANCE_WEIGHT:-10.0}"
|
||||
RTC_PREFIX_SCHEDULE="${RTC_PREFIX_SCHEDULE:-linear}"
|
||||
DEBUG="${DEBUG:-0}"
|
||||
DRY_RUN=0
|
||||
EXTRA_ARGS=()
|
||||
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case "$1" in
|
||||
-h|--help)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
--checkpoint-path)
|
||||
CHECKPOINT_PATH="${2:?missing value for --checkpoint-path}"
|
||||
shift 2
|
||||
;;
|
||||
--checkpoint-path=*)
|
||||
CHECKPOINT_PATH="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--train-config-path)
|
||||
TRAIN_CONFIG_PATH="${2:?missing value for --train-config-path}"
|
||||
shift 2
|
||||
;;
|
||||
--train-config-path=*)
|
||||
TRAIN_CONFIG_PATH="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--port)
|
||||
PORT="${2:?missing value for --port}"
|
||||
shift 2
|
||||
;;
|
||||
--port=*)
|
||||
PORT="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--host)
|
||||
HOST="${2:?missing value for --host}"
|
||||
shift 2
|
||||
;;
|
||||
--host=*)
|
||||
HOST="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--env)
|
||||
WALLX_ENV="${2:?missing value for --env}"
|
||||
shift 2
|
||||
;;
|
||||
--env=*)
|
||||
WALLX_ENV="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--rtc-execution-horizon)
|
||||
RTC_EXECUTION_HORIZON="${2:?missing value for --rtc-execution-horizon}"
|
||||
shift 2
|
||||
;;
|
||||
--rtc-max-guidance-weight)
|
||||
RTC_MAX_GUIDANCE_WEIGHT="${2:?missing value for --rtc-max-guidance-weight}"
|
||||
shift 2
|
||||
;;
|
||||
--rtc-prefix-schedule)
|
||||
RTC_PREFIX_SCHEDULE="${2:?missing value for --rtc-prefix-schedule}"
|
||||
shift 2
|
||||
;;
|
||||
--cuda-id)
|
||||
CUDA_ID="${2:?missing value for --cuda-id}"
|
||||
shift 2
|
||||
;;
|
||||
--cuda-id=*)
|
||||
CUDA_ID="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--image-passing-mode)
|
||||
IMAGE_PASSING_MODE="${2:?missing value for --image-passing-mode}"
|
||||
shift 2
|
||||
;;
|
||||
--image-passing-mode=*)
|
||||
IMAGE_PASSING_MODE="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--action-horizon)
|
||||
ACTION_HORIZON="${2:?missing value for --action-horizon}"
|
||||
shift 2
|
||||
;;
|
||||
--action-horizon=*)
|
||||
ACTION_HORIZON="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--robot-type)
|
||||
ROBOT_TYPE="${2:?missing value for --robot-type}"
|
||||
shift 2
|
||||
;;
|
||||
--robot-type=*)
|
||||
ROBOT_TYPE="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--robot-action-interpolate-multiplier)
|
||||
ROBOT_ACTION_INTERPOLATE_MULTIPLIER="${2:?missing value for --robot-action-interpolate-multiplier}"
|
||||
shift 2
|
||||
;;
|
||||
--robot-action-interpolate-multiplier=*)
|
||||
ROBOT_ACTION_INTERPOLATE_MULTIPLIER="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--robot-action-end-ratio)
|
||||
ROBOT_ACTION_END_RATIO="${2:?missing value for --robot-action-end-ratio}"
|
||||
shift 2
|
||||
;;
|
||||
--robot-action-end-ratio=*)
|
||||
ROBOT_ACTION_END_RATIO="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--model-device)
|
||||
MODEL_DEVICE="${2:?missing value for --model-device}"
|
||||
shift 2
|
||||
;;
|
||||
--model-device=*)
|
||||
MODEL_DEVICE="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--max-batch-size)
|
||||
MAX_BATCH_SIZE="${2:?missing value for --max-batch-size}"
|
||||
shift 2
|
||||
;;
|
||||
--max-batch-size=*)
|
||||
MAX_BATCH_SIZE="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--default-prompt)
|
||||
DEFAULT_PROMPT="${2:?missing value for --default-prompt}"
|
||||
shift 2
|
||||
;;
|
||||
--default-prompt=*)
|
||||
DEFAULT_PROMPT="${1#*=}"
|
||||
shift
|
||||
;;
|
||||
--serialize-actions)
|
||||
SERIALIZE_ACTIONS=1
|
||||
shift
|
||||
;;
|
||||
--no-serialize-actions|--raw-actions)
|
||||
SERIALIZE_ACTIONS=0
|
||||
shift
|
||||
;;
|
||||
--enable-cuda-graph)
|
||||
ENABLE_CUDA_GRAPH=1
|
||||
shift
|
||||
;;
|
||||
--enable-experimental-engine)
|
||||
ENABLE_EXPERIMENTAL_ENGINE=1
|
||||
shift
|
||||
;;
|
||||
--debug)
|
||||
DEBUG=1
|
||||
shift
|
||||
;;
|
||||
--dry-run)
|
||||
DRY_RUN=1
|
||||
shift
|
||||
;;
|
||||
--)
|
||||
shift
|
||||
EXTRA_ARGS+=("$@")
|
||||
break
|
||||
;;
|
||||
*)
|
||||
EXTRA_ARGS+=("$1")
|
||||
shift
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
if [[ -z "${CHECKPOINT_PATH}" ]]; then
|
||||
echo "error: --checkpoint-path is required." >&2
|
||||
echo >&2
|
||||
usage >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
# The CLI action horizon overrides the checkpoint config. Reject mismatches
|
||||
# before loading a large model or handing an incompatible packet to Turtle2.
|
||||
if [[ -z "${TRAIN_CONFIG_PATH}" && -f "${CHECKPOINT_PATH%/}/config.yml" ]]; then
|
||||
TRAIN_CONFIG_PATH="${CHECKPOINT_PATH%/}/config.yml"
|
||||
fi
|
||||
if [[ -n "${TRAIN_CONFIG_PATH}" ]]; then
|
||||
if ! "${PYTHON_BIN}" - "${TRAIN_CONFIG_PATH}" "${ACTION_HORIZON}" <<'PY'
|
||||
import sys
|
||||
import yaml
|
||||
|
||||
path, requested = sys.argv[1], int(sys.argv[2])
|
||||
try:
|
||||
with open(path, encoding="utf-8") as stream:
|
||||
config = yaml.load(stream, Loader=yaml.FullLoader) or {}
|
||||
except (OSError, yaml.YAMLError) as exc:
|
||||
raise SystemExit(f"error: cannot read training config {path}: {exc}")
|
||||
|
||||
task = config.get("task") or {}
|
||||
data = config.get("data") or {}
|
||||
configured = (
|
||||
task.get("action_horizon_flow")
|
||||
or task.get("action_horizon")
|
||||
or data.get("action_horizon_flow")
|
||||
or data.get("action_horizon")
|
||||
)
|
||||
if configured is not None and requested != int(configured):
|
||||
raise SystemExit(
|
||||
f"error: action horizon {requested} disagrees with training config "
|
||||
f"{path}: {configured}; use --action-horizon {configured}"
|
||||
)
|
||||
PY
|
||||
then
|
||||
exit 2
|
||||
fi
|
||||
fi
|
||||
|
||||
export CUDA_VISIBLE_DEVICES="${CUDA_ID}"
|
||||
export ENABLE_FAST_PREPROCESS="${ENABLE_FAST_PREPROCESS:-true}"
|
||||
|
||||
CMD=(
|
||||
"${PYTHON_BIN}" -m wall_x._vendor.harrix.serving.launch_serving_rtc
|
||||
--env "${WALLX_ENV}"
|
||||
--host "${HOST}"
|
||||
--port "${PORT}"
|
||||
--image-passing-mode "${IMAGE_PASSING_MODE}"
|
||||
--rtc-execution-horizon "${RTC_EXECUTION_HORIZON}"
|
||||
--rtc-max-guidance-weight "${RTC_MAX_GUIDANCE_WEIGHT}"
|
||||
--rtc-prefix-attention-schedule "${RTC_PREFIX_SCHEDULE}"
|
||||
)
|
||||
|
||||
if [[ "${SERIALIZE_ACTIONS}" == "1" ]]; then
|
||||
CMD+=(--serialize-actions)
|
||||
else
|
||||
CMD+=(--no-serialize-actions)
|
||||
fi
|
||||
if [[ -n "${MAX_BATCH_SIZE}" ]]; then
|
||||
CMD+=(--max-batch-size "${MAX_BATCH_SIZE}")
|
||||
fi
|
||||
if [[ -n "${DEFAULT_PROMPT}" ]]; then
|
||||
CMD+=(--default-prompt "${DEFAULT_PROMPT}")
|
||||
fi
|
||||
if [[ "${ENABLE_CUDA_GRAPH}" == "1" ]]; then
|
||||
CMD+=(--enable-cuda-graph)
|
||||
fi
|
||||
if [[ "${ENABLE_EXPERIMENTAL_ENGINE}" == "1" ]]; then
|
||||
CMD+=(--enable-experimental-engine)
|
||||
fi
|
||||
if [[ "${DEBUG}" == "1" ]]; then
|
||||
CMD+=(--debug)
|
||||
fi
|
||||
|
||||
CMD+=(
|
||||
model-config:server-model-config
|
||||
--model-config.checkpoint-path "${CHECKPOINT_PATH}"
|
||||
--model-config.action-horizon "${ACTION_HORIZON}"
|
||||
--model-config.robot-type "${ROBOT_TYPE}"
|
||||
--model-config.robot-action-interpolate-multiplier "${ROBOT_ACTION_INTERPOLATE_MULTIPLIER}"
|
||||
--model-config.robot-action-end-ratio "${ROBOT_ACTION_END_RATIO}"
|
||||
--model-config.model-device "${MODEL_DEVICE}"
|
||||
)
|
||||
|
||||
if [[ -n "${TRAIN_CONFIG_PATH}" ]]; then
|
||||
CMD+=(--model-config.train-config-path "${TRAIN_CONFIG_PATH}")
|
||||
fi
|
||||
|
||||
CMD+=("${EXTRA_ARGS[@]}")
|
||||
|
||||
printf 'Launching Wall-X RTC serving:\n'
|
||||
printf ' %q' "${CMD[@]}"
|
||||
printf '\n'
|
||||
|
||||
if [[ "${DRY_RUN}" == "1" ]]; then
|
||||
exit 0
|
||||
fi
|
||||
|
||||
exec "${CMD[@]}"
|
||||
Reference in New Issue
Block a user