281 lines
8.8 KiB
Python
281 lines
8.8 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Server script for Wall-X model.
|
|
|
|
This script serves a Wall-X model using a websocket server, allowing
|
|
clients to connect and get action predictions from observations.
|
|
|
|
"""
|
|
|
|
import dataclasses
|
|
import enum
|
|
import inspect
|
|
import logging
|
|
import os
|
|
import socket
|
|
import sys
|
|
import yaml
|
|
import traceback
|
|
import tyro
|
|
|
|
from wall_x._vendor.harrix.serving.websocket_policy_server import WebsocketPolicyServer
|
|
|
|
|
|
def _server_model_config_to_infer_kwargs(model_config) -> dict:
|
|
from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig
|
|
|
|
infer_params = inspect.signature(InferConfig.__init__).parameters
|
|
return {
|
|
k: v
|
|
for k, v in vars(model_config).items()
|
|
if v is not None and k in infer_params
|
|
}
|
|
|
|
|
|
def get_wallx_policy(model_config, image_passing_mode, serialize_actions=True):
|
|
from wall_x._vendor.harrix.serving.policy.wall_x_policy import WallXPolicy
|
|
from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig
|
|
|
|
config = InferConfig(**_server_model_config_to_infer_kwargs(model_config))
|
|
return WallXPolicy(config=config, image_passing_mode=image_passing_mode, serialize_actions=serialize_actions)
|
|
|
|
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class EnvMode(enum.Enum):
|
|
"""Supported environments/datasets."""
|
|
|
|
X2ROBOT = "x2robot"
|
|
LIBERO = "libero"
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class ServerModelConfig:
|
|
"""Configuration for loading a Wall-X model."""
|
|
|
|
checkpoint_path: str | None = None
|
|
train_config_path: str | None = None
|
|
# robot_host: str = '0.0.0.0'
|
|
# robot_port: int = 33723
|
|
robot_type: str = "desktop" # ["desktop", "turtle"]
|
|
robot_action_start_ratio: float = (
|
|
0.0 # proportion of action execution to start from
|
|
)
|
|
robot_action_end_ratio: float = 1.0 # proportion of action execution to end at
|
|
robot_action_interpolate_multiplier: int = 10 # action interpolation multiplier
|
|
robot_use_joint_angle_control: bool = (
|
|
False # use joint angle control (model must predict joints)
|
|
)
|
|
turtle_as_desktop: bool = (
|
|
False # use turtle platform as desktop with fixed base/head/camera/height
|
|
)
|
|
action_horizon: int = 32 # specify the correct horizon for the model
|
|
action_dim: int | None = None
|
|
model_device: str = "cuda"
|
|
num_inference_timesteps: int = 10
|
|
num_inference_steps: int | None = None
|
|
cfg_scale: float | None = None
|
|
seed: int | None = None
|
|
save_video_dir: str = "./videos"
|
|
|
|
# Please specify explicitly if the checkpoint was not trained on the x2robot dataset.
|
|
norm_key: str | None = None
|
|
# Model cameras; None = infer from train config ``data.key_mappings.camera``.
|
|
cam_names: list[str] | None = None
|
|
# Robot camera keys in incoming websocket observations.
|
|
camera_front_key: str = "camera_front"
|
|
camera_left_key: str = "camera_left"
|
|
camera_right_key: str = "camera_right"
|
|
# Serving prompt controls. If the client request has no instruction,
|
|
# default_instruction is used. prompt_template follows train-config semantics.
|
|
default_instruction: str | None = None
|
|
prompt_template: str | None = None
|
|
qwen25_prompt_template: str | None = None
|
|
prompt_priority_order: str | None = None
|
|
|
|
@dataclasses.dataclass
|
|
class Args:
|
|
"""Arguments for the serve_wall_x script."""
|
|
|
|
# Environment mode (used for default configurations)
|
|
env: EnvMode = EnvMode.X2ROBOT
|
|
|
|
# Model configuration. If not provided, uses default config for the environment
|
|
model_config: ServerModelConfig | None = None
|
|
|
|
# Default text prompt to use if not provided in observation
|
|
default_prompt: str | None = None
|
|
|
|
# Port to serve the policy on
|
|
port: int = 43007
|
|
|
|
# Host to bind the server to
|
|
host: str = "0.0.0.0"
|
|
|
|
# Enable debug logging
|
|
debug: bool = False
|
|
|
|
# Image passing mode
|
|
image_passing_mode: str = "base64" # ["numpy", "base64"]
|
|
|
|
# Model type
|
|
model_type: str = "wallx" # OSS supports qwen2.5 Wall-X only
|
|
|
|
# Serialize actions via robot_preprocessor (True for robot control, False for raw output)
|
|
serialize_actions: bool = True
|
|
|
|
# -- Dynamic batching -----------------------------------------
|
|
# Set max_batch_size to enable dynamic batching. None = single mode.
|
|
max_batch_size: int | None = None
|
|
max_wait_time_ms: float = 0
|
|
max_queue_size: int = 100
|
|
timeout_ms: float = 30000
|
|
|
|
# -- Engine flags ---------------------------------------------
|
|
enable_experimental_engine: bool = False
|
|
enable_cuda_graph: bool = False
|
|
|
|
|
|
# Default model configurations for each environment
|
|
DEFAULT_CONFIGS: dict[EnvMode, ServerModelConfig] = {
|
|
EnvMode.X2ROBOT: ServerModelConfig(
|
|
checkpoint_path=None,
|
|
train_config_path=None,
|
|
robot_action_start_ratio=0.0,
|
|
robot_action_end_ratio=1.0,
|
|
robot_action_interpolate_multiplier=10,
|
|
robot_use_joint_angle_control=False,
|
|
turtle_as_desktop=False,
|
|
action_horizon=32,
|
|
action_dim=None,
|
|
model_device="cuda",
|
|
num_inference_timesteps=10,
|
|
),
|
|
EnvMode.LIBERO: ServerModelConfig(
|
|
checkpoint_path=None,
|
|
train_config_path=None,
|
|
robot_type="desktop",
|
|
robot_action_start_ratio=0.0,
|
|
robot_action_end_ratio=1.0,
|
|
robot_action_interpolate_multiplier=1,
|
|
action_horizon=10,
|
|
action_dim=None,
|
|
model_device="cuda",
|
|
num_inference_timesteps=10,
|
|
cam_names=None,
|
|
),
|
|
}
|
|
|
|
|
|
def get_model_config(args: Args) -> ServerModelConfig:
|
|
"""Get model configuration from args or defaults."""
|
|
if args.model_config is not None:
|
|
return args.model_config
|
|
|
|
if config := DEFAULT_CONFIGS.get(args.env):
|
|
logger.info(f"Using default configuration for {args.env.value}")
|
|
return config
|
|
|
|
raise ValueError(
|
|
f"No default configuration for {args.env.value}. "
|
|
f"Please provide --model-config with model_path and action_tokenizer_path."
|
|
)
|
|
|
|
|
|
def create_policy(args: Args):
|
|
"""Create a policy from the given arguments."""
|
|
model_config = get_model_config(args)
|
|
if args.model_type != "wallx":
|
|
raise ValueError(
|
|
f"Unsupported model type: {args.model_type!r}. "
|
|
"The public package only supports model_type='wallx'."
|
|
)
|
|
policy = get_wallx_policy(
|
|
model_config, args.image_passing_mode, args.serialize_actions
|
|
)
|
|
return policy
|
|
|
|
|
|
def main(args: Args) -> None:
|
|
"""Main function to start the Wall-X model server."""
|
|
log_level = logging.DEBUG if args.debug else logging.INFO
|
|
logging.basicConfig(
|
|
level=log_level,
|
|
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
|
)
|
|
|
|
# Set engine environment variables
|
|
if args.enable_experimental_engine:
|
|
os.environ["ENABLE_EXPERIMENTAL_INFERENCE_ENGINE"] = "true"
|
|
logger.info("ENABLE_EXPERIMENTAL_INFERENCE_ENGINE=true")
|
|
|
|
if args.enable_cuda_graph:
|
|
os.environ["ENABLE_CUDA_GRAPH"] = "true"
|
|
logger.info("ENABLE_CUDA_GRAPH=true")
|
|
|
|
logger.info("Starting model server")
|
|
logger.info(f"Model type: {args.model_type}")
|
|
logger.info(f"Environment: {args.env.value}")
|
|
logger.info(f"Port: {args.port}")
|
|
logger.info(f"Host: {args.host}")
|
|
logger.info(f"Serialize actions: {args.serialize_actions}")
|
|
|
|
# Create policy
|
|
try:
|
|
policy = create_policy(args)
|
|
except Exception as e:
|
|
logger.error(f"Failed to create policy: {e}")
|
|
logger.error(traceback.format_exc())
|
|
sys.exit(1)
|
|
|
|
# Get policy metadata
|
|
policy_metadata = policy.metadata
|
|
policy_metadata["env"] = args.env.value
|
|
|
|
# Get network info
|
|
hostname = socket.gethostname()
|
|
try:
|
|
local_ip = socket.gethostbyname(hostname)
|
|
except Exception:
|
|
local_ip = "unknown"
|
|
|
|
logger.info(f"Server hostname: {hostname}")
|
|
logger.info(f"Server IP: {local_ip}")
|
|
logger.info(f"Server will be available at: ws://{args.host}:{args.port}")
|
|
logger.info(f"Health check endpoint: http://{args.host}:{args.port}/healthz")
|
|
|
|
batching_str = (
|
|
f"batch_size={args.max_batch_size}, wait={args.max_wait_time_ms}ms"
|
|
if args.max_batch_size
|
|
else "disabled"
|
|
)
|
|
logger.info(f"Batching: {batching_str}")
|
|
|
|
# Create and start server
|
|
server = WebsocketPolicyServer(
|
|
policy=policy,
|
|
host=args.host,
|
|
port=args.port,
|
|
metadata=policy_metadata,
|
|
max_batch_size=args.max_batch_size,
|
|
max_wait_time_ms=args.max_wait_time_ms,
|
|
max_queue_size=args.max_queue_size,
|
|
timeout_ms=args.timeout_ms,
|
|
)
|
|
|
|
logger.info("Starting server...")
|
|
try:
|
|
server.serve_forever()
|
|
except KeyboardInterrupt:
|
|
logger.info("Server stopped by user")
|
|
except Exception as e:
|
|
logger.error(f"Server error: {e}")
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main(tyro.cli(Args))
|