Enable inference serving && fix train stability (#59)
* enable serving * lint * update * update
This commit is contained in:
Lufang Chen
@@ -0,0 +1,263 @@
|
||||
# Wall-X Model Serving
|
||||
|
||||
This directory contains scripts for serving Wall-X models via a websocket server, allowing remote clients to connect and get action predictions from observations.
|
||||
|
||||
## Overview
|
||||
|
||||
The serving infrastructure consists of three main components:
|
||||
|
||||
1. **WebsocketPolicyServer** (`wall_x/serving/websocket_policy_server.py`): Generic websocket server that can serve any policy implementing the `BasePolicy` interface
|
||||
2. **WallXPolicy** (`wall_x/serving/policy/wall_x_policy.py`): Policy wrapper that adapts the Wall-X model to the `BasePolicy` interface
|
||||
3. **launch_serving.py**: Main script for starting the server with various configurations
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Basic Usage
|
||||
|
||||
Serve a model with default LIBERO configuration:
|
||||
|
||||
```bash
|
||||
cd /x2robot_v2/vincent/workspace/opensource
|
||||
python -m wall_x.serving.launch_serving \
|
||||
--env libero \
|
||||
--model-config.model-path /path/to/libero_model_stuff \
|
||||
--model-config.action-tokenizer-path /path/to/fast/ \
|
||||
--model-config.train-config-path /path/to/config.yml
|
||||
```
|
||||
|
||||
### Specify Environment
|
||||
|
||||
Serve with a specific environment preset:
|
||||
|
||||
```bash
|
||||
# LIBERO (single arm, 7 DOF)
|
||||
python -m wall_x.serving.launch_serving --env libero
|
||||
|
||||
# ALOHA (dual arm, 14 DOF)
|
||||
python -m wall_x.serving.launch_serving --env aloha
|
||||
```
|
||||
|
||||
### Custom Configuration
|
||||
|
||||
Serve with custom model paths and settings:
|
||||
|
||||
```bash
|
||||
python -m wall_x.serving.launch_serving \
|
||||
--model-config.model-path /path/to/model \
|
||||
--model-config.action-tokenizer-path /path/to/tokenizer \
|
||||
--model-config.train-config-path /path/to/train_config.yml \
|
||||
--model-config.action-dim 7 \
|
||||
--model-config.state-dim 8 \
|
||||
--model-config.pred-horizon 32 \
|
||||
--model-config.camera-key front_view left_wrist_view \
|
||||
--port 8000
|
||||
```
|
||||
|
||||
## Command Line Arguments
|
||||
|
||||
### Basic Arguments
|
||||
|
||||
- `--env {libero,aloha}`: Environment mode (default: libero)
|
||||
- `--port PORT`: Port to serve on (default: 8000)
|
||||
- `--host HOST`: Host to bind to (default: 0.0.0.0)
|
||||
- `--default-prompt TEXT`: Default text prompt if not provided in observation
|
||||
- `--debug`: Enable debug logging
|
||||
|
||||
### Model Configuration
|
||||
|
||||
All model configuration arguments use the `--model-config.` prefix:
|
||||
|
||||
- `--model-config.model-path PATH`: Path to pretrained model checkpoint (required)
|
||||
- `--model-config.action-tokenizer-path PATH`: Path to action tokenizer (required)
|
||||
- `--model-config.train-config-path PATH`: Path to train config YAML file (required)
|
||||
- `--model-config.action-dim INT`: Action space dimension (default: 7)
|
||||
- `--model-config.state-dim INT`: Robot state dimension (default: 8)
|
||||
- `--model-config.pred-horizon INT`: Prediction horizon (default: 32)
|
||||
- `--model-config.device {cuda,cpu}`: Device to run on (default: cuda)
|
||||
- `--model-config.dtype {bfloat16,float16,float32}`: Model dtype (default: bfloat16)
|
||||
- `--model-config.predict-mode {fast,diffusion}`: Prediction mode (default: fast)
|
||||
- `--model-config.camera-key KEY1 KEY2 ...`: Camera keys for observation images
|
||||
|
||||
### Camera Keys
|
||||
|
||||
The `camera-key` parameter specifies which camera views are expected in the observation dictionary. This is **critical** for proper operation:
|
||||
|
||||
- Keys must match between server configuration and client observations
|
||||
- Order matters: keys are processed in the order specified
|
||||
- Common keys: `front_view`, `left_wrist_view`, `right_wrist_view`, `face_view`
|
||||
|
||||
Example:
|
||||
```bash
|
||||
--model-config.camera-key front_view left_wrist_view
|
||||
```
|
||||
|
||||
Client must send observations with matching keys:
|
||||
```python
|
||||
obs = {
|
||||
"front_view": image1, # Must match camera-key[0]
|
||||
"left_wrist_view": image2, # Must match camera-key[1]
|
||||
"prompt": "task description",
|
||||
"state": robot_state,
|
||||
}
|
||||
```
|
||||
|
||||
## Default Configurations
|
||||
|
||||
### LIBERO (Single Arm)
|
||||
|
||||
```python
|
||||
ModelConfig(
|
||||
model_path="/path/to/model",
|
||||
action_tokenizer_path="/path/to/action_tokenizer",
|
||||
train_config_path="/path/to/train_config",
|
||||
state_dim=8,
|
||||
action_dim=7,
|
||||
pred_horizon=32,
|
||||
device="cuda",
|
||||
dtype="bfloat16",
|
||||
predict_mode="fast",
|
||||
camera_key=["front_view", "left_wrist_view"],
|
||||
)
|
||||
```
|
||||
|
||||
### ALOHA (Dual Arm)
|
||||
|
||||
```python
|
||||
ModelConfig(
|
||||
model_path="/path/to/model",
|
||||
action_tokenizer_path="/path/to/action_tokenizer",
|
||||
train_config_path="/path/to/train_config",
|
||||
state_dim=14,
|
||||
action_dim=14,
|
||||
pred_horizon=32,
|
||||
device="cuda",
|
||||
dtype="bfloat16",
|
||||
predict_mode="fast",
|
||||
camera_key=["face_view", "left_wrist_view", "right_wrist_view"],
|
||||
)
|
||||
```
|
||||
|
||||
## Server Protocol
|
||||
|
||||
### Connection Flow
|
||||
|
||||
1. Client connects to `ws://host:port`
|
||||
2. Server sends metadata JSON with policy information
|
||||
3. Client sends observation (msgpack-encoded)
|
||||
4. Server responds with action prediction (msgpack-encoded)
|
||||
5. Repeat steps 3-4 for each inference
|
||||
|
||||
### Observation Format
|
||||
|
||||
Observations must be a dictionary with camera keys matching server configuration:
|
||||
|
||||
```python
|
||||
obs = {
|
||||
# Image observations - keys must match server's camera_key configuration
|
||||
"front_view": np.ndarray, # (H, W, 3) uint8 or float
|
||||
"left_wrist_view": np.ndarray, # (H, W, 3) uint8 or float
|
||||
|
||||
# Required fields
|
||||
"prompt": str, # Task description
|
||||
"dataset_names": List[str], # Dataset/robot name, e.g., ["physical-intelligence/libero"]
|
||||
"state": np.ndarray, # Robot proprioception state (state_dim,)
|
||||
}
|
||||
```
|
||||
|
||||
**Important**: The image keys (`front_view`, `left_wrist_view`, etc.) must exactly match the `camera_key` parameter configured on the server.
|
||||
|
||||
### Action Response Format
|
||||
|
||||
Actions are returned as a dictionary:
|
||||
|
||||
```python
|
||||
{
|
||||
"action": np.ndarray, # Predicted action [pred_horizon, action_dim]
|
||||
"server_timing": {
|
||||
"infer_ms": float, # Inference time in milliseconds
|
||||
"prev_total_ms": float, # Total time for previous request
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Server Metadata
|
||||
|
||||
When connecting, the server sends metadata:
|
||||
|
||||
```python
|
||||
{
|
||||
"action_dim": int, # Action space dimension
|
||||
"pred_horizon": int, # Number of future actions predicted
|
||||
"device": str, # Device model runs on
|
||||
"predict_mode": str, # Prediction mode (fast/diffusion)
|
||||
"env": str, # Environment name
|
||||
}
|
||||
```
|
||||
|
||||
### Health Check
|
||||
|
||||
HTTP health check endpoint available at:
|
||||
```
|
||||
http://host:port/healthz
|
||||
```
|
||||
|
||||
Returns `200 OK` if the server is running.
|
||||
|
||||
## Client Example
|
||||
|
||||
### Synchronous Python Client
|
||||
|
||||
For synchronous usage, see `wall_x/serving/client.py`:
|
||||
|
||||
```python
|
||||
from wall_x.serving.client import WallXClient
|
||||
|
||||
# Create and connect
|
||||
client = WallXClient(uri="ws://localhost:8000")
|
||||
client.connect_sync()
|
||||
|
||||
# Prepare observation
|
||||
obs = {
|
||||
"front_view": image1,
|
||||
"left_wrist_view": image2,
|
||||
"prompt": "task description",
|
||||
"state": robot_state,
|
||||
"dataset_names": ["physical-intelligence/libero"],
|
||||
}
|
||||
|
||||
# Get prediction
|
||||
response = client.predict_sync(obs)
|
||||
action = response["action"]
|
||||
|
||||
# Close connection
|
||||
client.close_sync()
|
||||
```
|
||||
|
||||
## Architecture
|
||||
|
||||
### WebsocketPolicyServer
|
||||
|
||||
Generic websocket server that:
|
||||
- Handles websocket connections with msgpack serialization
|
||||
- Tracks inference timing and performance metrics
|
||||
- Provides health check endpoint
|
||||
- Handles errors gracefully with proper logging
|
||||
- Supports concurrent client connections
|
||||
|
||||
### WallXPolicy
|
||||
|
||||
Policy wrapper that:
|
||||
- Loads and manages the Wall-X model from pretrained checkpoint
|
||||
- Processes multi-camera observations
|
||||
- Handles image preprocessing (smart resize, normalization)
|
||||
- Manages device placement and dtype conversion
|
||||
- Provides policy metadata to clients
|
||||
- Supports both fast tokenizer and diffusion prediction modes
|
||||
|
||||
### Image Processing Pipeline
|
||||
|
||||
1. **Camera Key Matching**: Extracts images from observation dict using configured camera keys
|
||||
2. **Format Conversion**: Converts numpy arrays to PIL Images
|
||||
3. **Smart Resize**: Applies Qwen's smart resize algorithm based on min/max pixels
|
||||
4. **Vision Token Formatting**: Inserts vision tokens in text prompt
|
||||
5. **Batch Preparation**: Creates model-ready BatchFeature input
|
||||
Reference in New Issue
Block a user