343 lines
10 KiB
Bash
Executable File
343 lines
10 KiB
Bash
Executable File
#!/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[@]}"
|