huangyanpei/VLA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add mot (#83)

Browse Source 
* add mot

* update libero example

* translate zh to en

* fix load model from hf

* lint

* lint

---------

Co-authored-by: yangping <yangping@x2robot.com>
This commit is contained in:
suolyer
2026-02-03 11:35:25 +08:00
committed by GitHub
co-authored by yangping
parent 05b6d8dcf7
commit d18fa65fa1
26 changed files with 8509 additions and 1179 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wall_x/serving/policy/utils.py
+30 -10
View File
@@ -13,6 +13,7 @@ logger = logging.getLogger(__name__)
def prepare_batch(
obs: Dict,
processor,
normalizer_propri,
camera_key: List[str],
agent_pos_dim,
action_dim,
@@ -84,6 +85,7 @@ def prepare_batch(
img = Image.fromarray((img * 255).astype(np.uint8))
processed_images.append(img)
# print("processed_images:",processed_images)
# Apply smart resize to images
resized_images = process_images(
processed_images, image_factor, min_pixels, max_pixels
@@ -109,7 +111,9 @@ def prepare_batch(
action_token_id = processor.tokenizer.convert_tokens_to_ids("<|action|>")
moe_token_types = inputs.input_ids == action_token_id
inputs["moe_token_types"] = moe_token_types
inputs["moe_token_types"] = torch.tensor(moe_token_types)
# obs["dataset_names"]="libero_all"
# Handle robot state/proprioception if available
if "state" in obs:
@@ -126,20 +130,22 @@ def prepare_batch(
state = state.unsqueeze(1) # [batch, 1, state_dim]
# Pad to 20 dimensions if needed (same as training)
if state.shape[-1] < 20:
padding = torch.zeros(state.shape[0], state.shape[1], 20 - state.shape[-1])
state = torch.cat([state, padding], dim=-1)
# if state.shape[-1] < 20:
# padding = torch.zeros(state.shape[0], state.shape[1], 20 - state.shape[-1])
# state = torch.cat([state, padding], dim=-1)
# Create mask for valid dimensions
agent_pos_mask = torch.ones_like(state)
if state.shape[-1] > agent_pos_dim:
agent_pos_mask[:, :, agent_pos_dim:] = 0
normalizer_propri.normalize_data(state, [obs["dataset_names"]] * state.shape[0])
inputs["proprioception"] = state
inputs["agent_pos_mask"] = agent_pos_mask
# Add dataset name (required by model)
inputs["dataset_names"] = obs["dataset_names"]
inputs["dataset_names"] = [obs["dataset_names"]] * state.shape[0]
# Move all tensors to device
for key in inputs:
@@ -168,9 +174,21 @@ def process_images(
"""
resized_images = []
for img_pil in images:
current_width, current_height = img_pil.size
orig_width, orig_height = img_pil.size
target_size = 256
if target_size != -1:
# Maintain aspect ratio logic
if orig_width > orig_height: # Landscape image
new_width = target_size
new_height = int(target_size * orig_height / orig_width)
else: # Portrait image
new_height = target_size
new_width = int(target_size * orig_width / orig_height)
img_pil = img_pil.resize((new_width, new_height))
# Apply smart scaling (Qwen logic)
current_width, current_height = img_pil.size
resized_height, resized_width = smart_resize(
current_height,
current_width,
@@ -188,7 +206,7 @@ def process_images(
def format_text_with_vision_tokens(
instruction: str,
camera_key: List[str],
predict_mode: str = "fast",
predict_mode: str = "diffusion",
pred_horizon: int = 32,
) -> str:
"""Format text prompt with vision tokens for the model.
@@ -208,7 +226,7 @@ def format_text_with_vision_tokens(
image_pad_symbol = "<|image_pad|>"
propri_symbol = "<|propri|>"
action_symbol = "<|action|>"
# action_fast_symbol = "<|action_fast|>"
action_fast_symbol = "<|action_fast|>"
# Camera name mapping
camera_name_mapping = {
@@ -237,9 +255,11 @@ def format_text_with_vision_tokens(
f"\nPredict the next action in robot action.\nProprioception: {propri_symbol}\n"
)
user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n"
assistant_output = f"{role_start_symbol}assistant\n"
assistant_output = (
f"{role_start_symbol}assistant\n{action_fast_symbol}{role_end_symbol}\n"
)
if predict_mode == "diffusion":
assistant_output += f"{action_symbol * pred_horizon}"
assistant_output = f"{role_start_symbol}assistant\n{action_symbol * pred_horizon}{role_end_symbol}\n"
complete_text = prologue + user_message + assistant_output
return complete_text
wall_x/serving/policy/wall_x_policy.py
+22 -13
View File
@@ -1,12 +1,12 @@
import logging
from typing import Dict, Any, List
import torch
import copy
import numpy as np
from transformers import AutoProcessor
from wall_x.serving.websocket_policy_server import BasePolicy
from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction
from wall_x.serving.policy.utils import prepare_batch
from wall_x.model.model_utils import load_wallx_processors, register_normalizers
logger = logging.getLogger(__name__)
@@ -25,7 +25,7 @@ class WallXPolicy(BasePolicy):
camera_key: List[str],
device: str = "cuda",
dtype: str = "bfloat16",
predict_mode: str = "fast",
predict_mode: str = "diffusion",
default_prompt: str | None = None,
min_pixels: int = 4 * 28 * 28,
max_pixels: int = 16384 * 28 * 28,
@@ -50,21 +50,27 @@ class WallXPolicy(BasePolicy):
"""
logger.info(f"Loading Wall-X model from {model_path}")
self.normalizer_action, self.normalizer_propri = register_normalizers(
train_config, model_path
)
self.model = Qwen2_5_VLMoEForAction.from_pretrained(
model_path,
train_config=train_config,
action_tokenizer_path=action_tokenizer_path,
)
self.model.set_normalizer(
copy.deepcopy(self.normalizer_action), copy.deepcopy(self.normalizer_propri)
)
self.model.eval()
self.model = self.model.to(device)
self.model = self.model.bfloat16()
self.model.to_bfloat16_for_selected_params()
# hard code the action dim to 20 for align to wall-x configuration
self.fixed_action_dim = 20
self.fixed_action_dim = action_dim
self.action_dim = action_dim
self.agent_pos_dim = agent_pos_dim
self.agent_pos_dim = action_dim
self.pred_horizon = pred_horizon
self.device = device
self.predict_mode = predict_mode
@@ -77,10 +83,14 @@ class WallXPolicy(BasePolicy):
self.image_factor = image_factor
self.max_length = max_length
print("predict_mode", predict_mode)
print("camera_key", camera_key)
# Load processor
logger.info("Loading processor and tokenizer...")
self.processor = AutoProcessor.from_pretrained(model_path, use_fast=True)
self.processor.tokenizer.padding_side = "left"
processors_dict = load_wallx_processors(train_config)
self.processor = processors_dict["processor"]
# Action buffer for multi-step predictions
self.action_buffer = []
@@ -126,6 +136,7 @@ class WallXPolicy(BasePolicy):
input_batch = prepare_batch(
obs,
self.processor,
self.normalizer_propri,
self.camera_key,
self.agent_pos_dim,
self.action_dim,
@@ -147,7 +158,7 @@ class WallXPolicy(BasePolicy):
if self.predict_mode == "fast"
else self.fixed_action_dim
),
pred_horizon=self.pred_horizon,
action_horizon=self.pred_horizon,
mode="predict",
predict_mode=self.predict_mode,
)
@@ -164,9 +175,7 @@ class WallXPolicy(BasePolicy):
.to(torch.float32)
.numpy()
)
print(predicted_actions.shape)
return {"action": predicted_actions}
return {"predict_action": predicted_actions}
except Exception as e:
logger.error(f"Error during inference: {e}")