466 lines
17 KiB
Python
466 lines
17 KiB
Python
|
"""Shared adapter base for Qwen-VL action models.
|
||
|
|
|
||
|
|
Variant-specific subclasses live in ``harrix.adapters.variants`` and provide
|
||
|
|
the Wall-X training adapter class via ``_training_adapter``. The shared base
|
||
|
|
handles checkpoint loading, normalizers, LIBERO observation encoding, prompt
|
||
|
|
construction, and flow-action inference.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import copy
|
||
|
|
import logging
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
from PIL import Image
|
||
|
|
|
||
|
|
from qwen_vl_utils.vision_process import smart_resize
|
||
|
|
|
||
|
|
from wall_x._vendor.x2robot_utils.text_templates import (
|
||
|
|
get_prologue_with_embodied_information,
|
||
|
|
preprocesser_call,
|
||
|
|
)
|
||
|
|
|
||
|
|
from wall_x._vendor.harrix.adapters.base import BaseInferAdapter
|
||
|
|
from wall_x._vendor.harrix.envs.libero_common import decode_chunk, encode_proprio
|
||
|
|
from wall_x._vendor.harrix.eval_config import EvalConfig
|
||
|
|
from wall_x._vendor.harrix.utils.ckpt_load import (
|
||
|
|
load_state_dict,
|
||
|
|
resolve_checkpoint_dir,
|
||
|
|
)
|
||
|
|
from wall_x._vendor.harrix.utils.normalizer import build_normalizers
|
||
|
|
from wall_x._vendor.harrix.utils.train_config import (
|
||
|
|
build_data_config,
|
||
|
|
build_model_config,
|
||
|
|
load_train_config_with_ckpt_overlay,
|
||
|
|
normalize_train_config_for_inference,
|
||
|
|
register_data_backend,
|
||
|
|
resolve_state_bins,
|
||
|
|
resolve_use_state_string_representation,
|
||
|
|
)
|
||
|
|
from wall_x.trainer.trainer_utils import load_wallx_processors
|
||
|
|
|
||
|
|
|
||
|
|
_SUPPORTED_MODES = {"flow", "ar", "dllm", "vqa", "subtask"}
|
||
|
|
_IMPLEMENTED_MODES = {"flow"}
|
||
|
|
|
||
|
|
# Special tokens used by the Qwen-VL action model.
|
||
|
|
_ROLE_START = "<|im_start|>"
|
||
|
|
_ROLE_END = "<|im_end|>"
|
||
|
|
_VISION_START = "<|vision_start|>"
|
||
|
|
_VISION_END = "<|vision_end|>"
|
||
|
|
_IMAGE_PAD = "<|image_pad|>"
|
||
|
|
_PROPRI = "<|propri|>"
|
||
|
|
_ACTION = "<|action|>"
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
_PUBLIC_CAMERA_LABELS = {
|
||
|
|
"face_view": "front view",
|
||
|
|
"right_wrist_view": "right wrist view",
|
||
|
|
"left_wrist_view": "left wrist view",
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def _camera_label(cam_name: str) -> str:
|
||
|
|
return _PUBLIC_CAMERA_LABELS.get(cam_name, cam_name.replace("_", " "))
|
||
|
|
|
||
|
|
|
||
|
|
def _normalizer_width(normalizer, norm_key: str) -> int | None:
|
||
|
|
if normalizer is None or norm_key not in normalizer.delta:
|
||
|
|
return None
|
||
|
|
return int(normalizer.delta[norm_key].shape[0])
|
||
|
|
|
||
|
|
|
||
|
|
def _normalize_real_prefix(normalizer, tensor, dataset_names):
|
||
|
|
if normalizer is None:
|
||
|
|
return tensor
|
||
|
|
widths = [_normalizer_width(normalizer, name) for name in dataset_names]
|
||
|
|
if any(width is None for width in widths) or len(set(widths)) != 1:
|
||
|
|
return normalizer.normalize_data(tensor, dataset_names)
|
||
|
|
width = widths[0]
|
||
|
|
if tensor.shape[-1] == width:
|
||
|
|
return normalizer.normalize_data(tensor, dataset_names)
|
||
|
|
if tensor.shape[-1] < width:
|
||
|
|
raise ValueError(
|
||
|
|
f"normalizer width {width} exceeds tensor dim {tensor.shape[-1]}"
|
||
|
|
)
|
||
|
|
out = tensor.clone()
|
||
|
|
out[..., :width] = normalizer.normalize_data(
|
||
|
|
tensor[..., :width], dataset_names
|
||
|
|
)
|
||
|
|
out[..., width:] = 0
|
||
|
|
return out
|
||
|
|
|
||
|
|
class QwenVLActInferAdapter(BaseInferAdapter):
|
||
|
|
"""Shared constructor and batched flow inference implementation."""
|
||
|
|
|
||
|
|
# ---- subclass hook ----
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def _training_adapter(cls):
|
||
|
|
"""Return the training-side ModelAdapter subclass."""
|
||
|
|
raise NotImplementedError(f"{cls.__name__} must override _training_adapter()")
|
||
|
|
|
||
|
|
# ---- ctor ----
|
||
|
|
|
||
|
|
def __init__(self, cfg: EvalConfig) -> None:
|
||
|
|
mode = cfg.model.action_mode
|
||
|
|
if mode not in _SUPPORTED_MODES:
|
||
|
|
raise ValueError(
|
||
|
|
f"{type(self).__name__}: unknown action_mode={mode!r}, "
|
||
|
|
f"supported={sorted(_SUPPORTED_MODES)}"
|
||
|
|
)
|
||
|
|
if mode not in _IMPLEMENTED_MODES:
|
||
|
|
raise NotImplementedError(
|
||
|
|
f"{type(self).__name__}: action_mode={mode!r} is not implemented; "
|
||
|
|
"currently supported="
|
||
|
|
f"{sorted(_IMPLEMENTED_MODES)}"
|
||
|
|
)
|
||
|
|
self._action_mode = mode
|
||
|
|
|
||
|
|
ta = self._training_adapter()
|
||
|
|
device = "cuda"
|
||
|
|
self._device = device
|
||
|
|
self._checkpoint_path = resolve_checkpoint_dir(cfg.model.checkpoint_path)
|
||
|
|
|
||
|
|
# 1) train_config with checkpoint overlays, then data backend registration.
|
||
|
|
train_config = load_train_config_with_ckpt_overlay(
|
||
|
|
cfg.model.train_config_path, self._checkpoint_path
|
||
|
|
)
|
||
|
|
train_config = normalize_train_config_for_inference(
|
||
|
|
train_config, cfg.model.train_config_path
|
||
|
|
)
|
||
|
|
register_data_backend(train_config)
|
||
|
|
self._train_config = train_config
|
||
|
|
|
||
|
|
# 2) normalizers first - action_tokenizer setup needs them during processor load.
|
||
|
|
normalizer_action, normalizer_propri, resolved_norm_key = build_normalizers(
|
||
|
|
self._checkpoint_path, train_config, cfg.model.norm_key
|
||
|
|
)
|
||
|
|
self._normalizer_action = normalizer_action
|
||
|
|
self._normalizer_propri = normalizer_propri
|
||
|
|
self._norm_key = resolved_norm_key
|
||
|
|
if normalizer_propri is not None and resolved_norm_key in normalizer_propri.delta:
|
||
|
|
train_config["_libero_proprio_norm_dim"] = int(
|
||
|
|
normalizer_propri.delta[resolved_norm_key].shape[0]
|
||
|
|
)
|
||
|
|
|
||
|
|
# 3) data_config. Image resizing needs resolution/image_factor/min/max.
|
||
|
|
self._data_config = build_data_config(cfg.model.train_config_path, train_config)
|
||
|
|
|
||
|
|
# 4) HF model config
|
||
|
|
ConfigClass = ta.config_class()
|
||
|
|
self._model_config = build_model_config(
|
||
|
|
ConfigClass,
|
||
|
|
self._checkpoint_path,
|
||
|
|
train_config,
|
||
|
|
cfg.model.train_config_path,
|
||
|
|
)
|
||
|
|
|
||
|
|
# 5) processor + tokenizer_mixin (may extend vocab via action_tokenizer)
|
||
|
|
procs = load_wallx_processors(
|
||
|
|
train_config, normalizer=normalizer_action, device=device
|
||
|
|
)
|
||
|
|
self._processor = procs["processor"]
|
||
|
|
self._tokenizer_mixin = procs.get("tokenizer_mixin")
|
||
|
|
logger.info(
|
||
|
|
"processor vocab size: %s (action_tokenizer_type=%r)",
|
||
|
|
len(self._processor.tokenizer),
|
||
|
|
train_config.get("action_tokenizer_type"),
|
||
|
|
)
|
||
|
|
|
||
|
|
# 6) model
|
||
|
|
ModelClass = ta.inference_model_class()
|
||
|
|
self._model_class = ModelClass
|
||
|
|
model = ModelClass(self._model_config, self._processor, self._tokenizer_mixin)
|
||
|
|
model.resize_token_embeddings(len(self._processor.tokenizer))
|
||
|
|
model.to_bfloat16_for_selected_params()
|
||
|
|
|
||
|
|
# 7) checkpoint weights and model finalization.
|
||
|
|
state_dict = load_state_dict(self._checkpoint_path, ModelClass)
|
||
|
|
embed_key = "model.embed_tokens.weight"
|
||
|
|
if embed_key in state_dict:
|
||
|
|
ckpt_vocab = state_dict[embed_key].shape[0]
|
||
|
|
cur_vocab = model.model.embed_tokens.weight.shape[0]
|
||
|
|
if cur_vocab != ckpt_vocab:
|
||
|
|
logger.info(
|
||
|
|
"resize_token_embeddings from %d to %d to match checkpoint",
|
||
|
|
cur_vocab,
|
||
|
|
ckpt_vocab,
|
||
|
|
)
|
||
|
|
model.resize_token_embeddings(ckpt_vocab)
|
||
|
|
msg = model.load_state_dict(state_dict, strict=False)
|
||
|
|
logger.info(
|
||
|
|
"%s load_state_dict: missing=%s unexpected=%s",
|
||
|
|
type(self).__name__,
|
||
|
|
len(msg.missing_keys),
|
||
|
|
len(msg.unexpected_keys),
|
||
|
|
)
|
||
|
|
model.set_normalizer(
|
||
|
|
copy.deepcopy(normalizer_action),
|
||
|
|
copy.deepcopy(normalizer_propri),
|
||
|
|
)
|
||
|
|
model.eval()
|
||
|
|
model.to(device)
|
||
|
|
model.to_bfloat16_for_selected_params()
|
||
|
|
self._model = model
|
||
|
|
|
||
|
|
# 8) cached runtime fields
|
||
|
|
self._cam_names = list(cfg.model.cam_names)
|
||
|
|
self._action_horizon = int(cfg.model.action_horizon)
|
||
|
|
self._action_dim = sum(train_config["dof_config"].values())
|
||
|
|
self._num_inference_timesteps = 10
|
||
|
|
self._robot_id = "10000"
|
||
|
|
|
||
|
|
# ---- BaseInferAdapter API ----
|
||
|
|
|
||
|
|
@property
|
||
|
|
def chunk_horizon(self) -> int:
|
||
|
|
return self._action_horizon
|
||
|
|
|
||
|
|
@property
|
||
|
|
def action_mode(self) -> str:
|
||
|
|
return self._action_mode
|
||
|
|
|
||
|
|
def predict_batch(self, payloads: list[dict]) -> list[np.ndarray]:
|
||
|
|
if self._action_mode == "flow":
|
||
|
|
return self._flow_batch(payloads)
|
||
|
|
raise NotImplementedError(
|
||
|
|
f"{type(self).__name__}.predict_batch: action_mode={self._action_mode!r} "
|
||
|
|
"is not dispatched"
|
||
|
|
)
|
||
|
|
|
||
|
|
# ---- batched flow inference ----
|
||
|
|
|
||
|
|
def _flow_batch(self, payloads: list[dict]) -> list[np.ndarray]:
|
||
|
|
# 1) Encode each observation into proprioception, masks, and views.
|
||
|
|
observations: list[dict] = []
|
||
|
|
instructions: list[str] = []
|
||
|
|
noises: list[np.ndarray | None] = []
|
||
|
|
any_noise = False
|
||
|
|
for p in payloads:
|
||
|
|
observations.append(
|
||
|
|
encode_proprio(
|
||
|
|
p["observation"], self._train_config, self._action_horizon
|
||
|
|
)
|
||
|
|
)
|
||
|
|
instructions.append(p["instruction"])
|
||
|
|
n = p.get("noise")
|
||
|
|
if n is not None:
|
||
|
|
any_noise = True
|
||
|
|
noises.append(n)
|
||
|
|
|
||
|
|
# 2) Stack noise when provided; all-None delegates sampling to the model.
|
||
|
|
if any_noise:
|
||
|
|
if any(n is None for n in noises):
|
||
|
|
raise ValueError(
|
||
|
|
"payload noise must be provided for every payload or for none"
|
||
|
|
)
|
||
|
|
batch_noise = torch.stack(
|
||
|
|
[torch.from_numpy(n).to(dtype=torch.float32) for n in noises], dim=0
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
batch_noise = None
|
||
|
|
|
||
|
|
# 3) Prompts and model inputs.
|
||
|
|
prefix_list, postfix_list = [], []
|
||
|
|
for ins in instructions:
|
||
|
|
prefix, postfix = self._get_flow_prompt(ins)
|
||
|
|
prefix_list.append(prefix)
|
||
|
|
postfix_list.append(postfix)
|
||
|
|
batch_inputs = self._construct_model_input(
|
||
|
|
observations, prefix_list, postfix_list
|
||
|
|
)
|
||
|
|
|
||
|
|
# 4) Normalized zero action as the flow starting point.
|
||
|
|
padding = torch.zeros(
|
||
|
|
(
|
||
|
|
len(batch_inputs["dataset_names"]),
|
||
|
|
1,
|
||
|
|
self._action_dim,
|
||
|
|
),
|
||
|
|
dtype=torch.float32,
|
||
|
|
)
|
||
|
|
padding_action = _normalize_real_prefix(
|
||
|
|
self._normalizer_action, padding, batch_inputs["dataset_names"]
|
||
|
|
).to(batch_inputs["input_ids"].device)
|
||
|
|
|
||
|
|
# 5) Flow forward.
|
||
|
|
model_output = self._model.generate_flow_action(
|
||
|
|
action_horizon=self._action_horizon,
|
||
|
|
action_dim=self._action_dim,
|
||
|
|
num_inference_timesteps=self._num_inference_timesteps,
|
||
|
|
padding_action=padding_action,
|
||
|
|
noise=batch_noise,
|
||
|
|
**batch_inputs,
|
||
|
|
)
|
||
|
|
|
||
|
|
# 6) Decode one action chunk per payload.
|
||
|
|
predict_action = model_output["predict_action"] # (B, H, D_action)
|
||
|
|
if isinstance(predict_action, torch.Tensor):
|
||
|
|
predict_action = predict_action.detach().cpu().numpy()
|
||
|
|
return [
|
||
|
|
decode_chunk(predict_action[i : i + 1], self._train_config)
|
||
|
|
for i in range(len(payloads))
|
||
|
|
]
|
||
|
|
|
||
|
|
# ---- prompt template ----
|
||
|
|
|
||
|
|
def _get_flow_prompt(self, instruction: str) -> tuple[str, str]:
|
||
|
|
"""Build the flow-action prompt."""
|
||
|
|
if self._train_config["data"].get("use_embodied_system_prompt_ratio", 0) > 0:
|
||
|
|
robot_id = (
|
||
|
|
self._robot_id if self._norm_key in ("x2_normal", "ex_normal") else 0
|
||
|
|
)
|
||
|
|
cam_name_mapping = {cn: cn for cn in self._cam_names}
|
||
|
|
prologue = get_prologue_with_embodied_information(
|
||
|
|
dataset_name=self._norm_key,
|
||
|
|
cam_mapping=cam_name_mapping,
|
||
|
|
robot_id=robot_id,
|
||
|
|
uid="",
|
||
|
|
config=self._data_config,
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
prologue = f"{_ROLE_START}system\nYou are a helpful assistant.{_ROLE_END}\n"
|
||
|
|
user_request = f"{_ROLE_START}user\nObservation:"
|
||
|
|
for cn in self._cam_names:
|
||
|
|
user_request += (
|
||
|
|
f" {_camera_label(cn)}: " f"{_VISION_START}{_IMAGE_PAD}{_VISION_END}"
|
||
|
|
)
|
||
|
|
user_request += "\nInstruction:"
|
||
|
|
text_prompt = (
|
||
|
|
f"\nPredict the next action in robot action.\nProprioception: {_PROPRI}\n"
|
||
|
|
)
|
||
|
|
user_message = f"{user_request} {instruction}{text_prompt}{_ROLE_END}\n"
|
||
|
|
assistant_message = f"{_ROLE_START}assistant\n"
|
||
|
|
flow_action = _ACTION * self._action_horizon
|
||
|
|
|
||
|
|
prefix_text = prologue + user_message + assistant_message
|
||
|
|
postfix_text = flow_action
|
||
|
|
return prefix_text, postfix_text
|
||
|
|
|
||
|
|
# ---- batched model input construction ----
|
||
|
|
|
||
|
|
def _construct_model_input(
|
||
|
|
self,
|
||
|
|
observations: list[dict],
|
||
|
|
prefix_list: list[str],
|
||
|
|
postfix_list: list[str],
|
||
|
|
) -> dict:
|
||
|
|
"""Build model inputs for the flow inference path."""
|
||
|
|
batch_size = len(observations)
|
||
|
|
dataset_names = [self._norm_key] * batch_size
|
||
|
|
|
||
|
|
# Proprioception and masks are prepared as ndarrays in encode_proprio.
|
||
|
|
agent_pos = torch.cat(
|
||
|
|
[torch.from_numpy(o["proprioception"]) for o in observations], dim=0
|
||
|
|
)
|
||
|
|
agent_pos_mask = torch.cat(
|
||
|
|
[torch.from_numpy(o["agent_pos_mask"]) for o in observations], dim=0
|
||
|
|
)
|
||
|
|
dof_mask = torch.cat(
|
||
|
|
[torch.from_numpy(o["dof_mask"]) for o in observations], dim=0
|
||
|
|
)
|
||
|
|
agent_pos = _normalize_real_prefix(
|
||
|
|
self._normalizer_propri, agent_pos, dataset_names
|
||
|
|
)
|
||
|
|
|
||
|
|
# Resize images per sample.
|
||
|
|
image_inputs: list[torch.Tensor] = []
|
||
|
|
all_image_sizes: list[tuple[int, int]] = []
|
||
|
|
for o in observations:
|
||
|
|
for cn in self._cam_names:
|
||
|
|
if cn not in o:
|
||
|
|
continue
|
||
|
|
tensor = self._resize_image(o[cn], cn)
|
||
|
|
image_inputs.append(tensor)
|
||
|
|
# Tensor (H, W, C) to PIL-compatible (W, H).
|
||
|
|
all_image_sizes.append((tensor.shape[1], tensor.shape[0]))
|
||
|
|
|
||
|
|
inputs = preprocesser_call(
|
||
|
|
processor=self._processor,
|
||
|
|
prefix_text=prefix_list,
|
||
|
|
postfix_text=postfix_list,
|
||
|
|
images=image_inputs,
|
||
|
|
videos=None,
|
||
|
|
padding=True,
|
||
|
|
truncation=True,
|
||
|
|
return_tensors="pt",
|
||
|
|
max_length=1000,
|
||
|
|
pad_to_128_multiple=False,
|
||
|
|
pad_prefix_to_same_length=False,
|
||
|
|
norm_state=(
|
||
|
|
agent_pos
|
||
|
|
if resolve_use_state_string_representation(self._train_config)
|
||
|
|
else None
|
||
|
|
),
|
||
|
|
agent_pos_mask=agent_pos_mask,
|
||
|
|
state_augmentation_prob=0.0,
|
||
|
|
state_drop_prob=0.0,
|
||
|
|
state_augmentation_ratio=0.0,
|
||
|
|
state_bins=resolve_state_bins(self._train_config),
|
||
|
|
inference_mode=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
action_token_id = self._processor.tokenizer.convert_tokens_to_ids(_ACTION)
|
||
|
|
moe_token_types = inputs["input_ids"] == action_token_id
|
||
|
|
|
||
|
|
extra = {
|
||
|
|
"proprioception": agent_pos.detach(),
|
||
|
|
"agent_pos_mask": agent_pos_mask,
|
||
|
|
"dof_mask": dof_mask,
|
||
|
|
"image_size": all_image_sizes,
|
||
|
|
"moe_token_types": moe_token_types,
|
||
|
|
"dataset_names": dataset_names,
|
||
|
|
}
|
||
|
|
inputs.update(extra)
|
||
|
|
return _move_to_device(inputs, self._device)
|
||
|
|
|
||
|
|
# ---- image resizing ----
|
||
|
|
|
||
|
|
def _resize_image(self, img: np.ndarray, cam_name: str) -> torch.Tensor:
|
||
|
|
"""Resize one image with the train-time image config."""
|
||
|
|
if isinstance(img, np.ndarray):
|
||
|
|
pil = Image.fromarray(img)
|
||
|
|
elif isinstance(img, Image.Image):
|
||
|
|
pil = img
|
||
|
|
else:
|
||
|
|
raise ValueError(f"unsupported image type: {type(img)}")
|
||
|
|
orig_w, orig_h = pil.size
|
||
|
|
|
||
|
|
target = self._data_config.resolution.get(cam_name, -1)
|
||
|
|
if target != -1:
|
||
|
|
if orig_w > orig_h:
|
||
|
|
new_w, new_h = target, int(target * orig_h / orig_w)
|
||
|
|
else:
|
||
|
|
new_h, new_w = target, int(target * orig_w / orig_h)
|
||
|
|
pil = pil.resize((new_w, new_h))
|
||
|
|
|
||
|
|
cur_w, cur_h = pil.size
|
||
|
|
resized_h, resized_w = smart_resize(
|
||
|
|
cur_h,
|
||
|
|
cur_w,
|
||
|
|
factor=self._data_config.image_factor,
|
||
|
|
min_pixels=self._data_config.min_pixels,
|
||
|
|
max_pixels=self._data_config.max_pixels,
|
||
|
|
)
|
||
|
|
resized = pil.resize((resized_w, resized_h))
|
||
|
|
return torch.from_numpy(np.array(resized)).to(self._device)
|
||
|
|
|
||
|
|
|
||
|
|
def _move_to_device(obj, device):
|
||
|
|
"""Recursively move tensors inside common containers to ``device``."""
|
||
|
|
from transformers import BatchFeature
|
||
|
|
|
||
|
|
if isinstance(obj, torch.Tensor):
|
||
|
|
return obj.to(device)
|
||
|
|
if isinstance(obj, (dict, BatchFeature)):
|
||
|
|
return {k: _move_to_device(v, device) for k, v in obj.items()}
|
||
|
|
if isinstance(obj, list):
|
||
|
|
return [_move_to_device(v, device) for v in obj]
|
||
|
|
if isinstance(obj, tuple):
|
||
|
|
return tuple(_move_to_device(v, device) for v in obj)
|
||
|
|
return obj
|