Update Wall-X to 1.1.0 (#104)
This commit is contained in:
Starrick Liu
@@ -0,0 +1,5 @@
|
||||
"""Vendored subset of x2robot_utils for Wall-X."""
|
||||
|
||||
__version__ = "0.2.0-vendored"
|
||||
|
||||
__all__ = ["__version__"]
|
||||
@@ -0,0 +1,309 @@
|
||||
"""Rotation / pose geometry utilities (pure numpy + numba).
|
||||
|
||||
Shared between wall-x and internal_dataset_backend:
|
||||
- ``euler_to_matrix_zyx_6d_nb``: ZYX Euler -> flattened top-2-rows of R (N, 6)
|
||||
- ``so3_to_euler_zyx_batch_nb``: 6D rotation -> ZYX Euler (canonicalized)
|
||||
- ``compose_state_and_delta_to_abs_{rpy,6d}``: state + delta -> absolute pose
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
from numba import jit, prange
|
||||
|
||||
|
||||
@jit(nopython=True)
|
||||
def euler_to_matrix_zyx_6d_nb(eulers):
|
||||
"""Euler angles (N, 3) -> flattened top two rows of rotation matrix (N, 6)."""
|
||||
N = eulers.shape[0]
|
||||
R6 = np.empty((N, 6), dtype=np.float64)
|
||||
for i in prange(N):
|
||||
roll = eulers[i, 0]
|
||||
pitch = eulers[i, 1]
|
||||
yaw = eulers[i, 2]
|
||||
|
||||
cy, sy = np.cos(yaw), np.sin(yaw)
|
||||
cp, sp = np.cos(pitch), np.sin(pitch)
|
||||
cr, sr = np.cos(roll), np.sin(roll)
|
||||
|
||||
r00 = cy * cp
|
||||
r01 = cy * sp * sr - sy * cr
|
||||
r02 = cy * sp * cr + sy * sr
|
||||
|
||||
r10 = sy * cp
|
||||
r11 = sy * sp * sr + cy * cr
|
||||
r12 = sy * sp * cr - cy * sr
|
||||
|
||||
R6[i, 0] = r00
|
||||
R6[i, 1] = r01
|
||||
R6[i, 2] = r02
|
||||
R6[i, 3] = r10
|
||||
R6[i, 4] = r11
|
||||
R6[i, 5] = r12
|
||||
return R6
|
||||
|
||||
|
||||
@jit(nopython=True)
|
||||
def euler_to_matrix_zyx_batch_nb(eulers):
|
||||
N = eulers.shape[0]
|
||||
R = np.empty((N, 3, 3), dtype=np.float64)
|
||||
for i in prange(N):
|
||||
roll = eulers[i, 0]
|
||||
pitch = eulers[i, 1]
|
||||
yaw = eulers[i, 2]
|
||||
|
||||
cy, sy = np.cos(yaw), np.sin(yaw)
|
||||
cp, sp = np.cos(pitch), np.sin(pitch)
|
||||
cr, sr = np.cos(roll), np.sin(roll)
|
||||
|
||||
R[i, 0, 0] = cy * cp
|
||||
R[i, 0, 1] = cy * sp * sr - sy * cr
|
||||
R[i, 0, 2] = cy * sp * cr + sy * sr
|
||||
|
||||
R[i, 1, 0] = sy * cp
|
||||
R[i, 1, 1] = sy * sp * sr + cy * cr
|
||||
R[i, 1, 2] = sy * sp * cr - cy * sr
|
||||
|
||||
R[i, 2, 0] = -sp
|
||||
R[i, 2, 1] = cp * sr
|
||||
R[i, 2, 2] = cp * cr
|
||||
return R
|
||||
|
||||
|
||||
@jit(nopython=True)
|
||||
def matrix_to_euler_zyx_batch_nb(Rs):
|
||||
"""R = Rz(yaw) * Ry(pitch) * Rx(roll) -> (roll, pitch, yaw)."""
|
||||
N = Rs.shape[0]
|
||||
eulers = np.empty((N, 3), dtype=np.float64)
|
||||
for i in prange(N):
|
||||
r00 = Rs[i, 0, 0]
|
||||
r10 = Rs[i, 1, 0]
|
||||
r20 = Rs[i, 2, 0]
|
||||
r21 = Rs[i, 2, 1]
|
||||
r22 = Rs[i, 2, 2]
|
||||
|
||||
x = -r20
|
||||
if x > 1.0:
|
||||
x = 1.0
|
||||
elif x < -1.0:
|
||||
x = -1.0
|
||||
|
||||
pitch = np.arcsin(x)
|
||||
roll = np.arctan2(r21, r22)
|
||||
yaw = np.arctan2(r10, r00)
|
||||
|
||||
eulers[i, 0] = roll
|
||||
eulers[i, 1] = pitch
|
||||
eulers[i, 2] = yaw
|
||||
return eulers
|
||||
|
||||
|
||||
@jit(nopython=True)
|
||||
def so3_to_matrix_batch_nb(batch_so3):
|
||||
N = batch_so3.shape[0]
|
||||
R_all = np.empty((N, 3, 3), dtype=np.float64)
|
||||
eps = 1e-12
|
||||
for i in prange(N):
|
||||
r1x, r1y, r1z = batch_so3[i, 0], batch_so3[i, 1], batch_so3[i, 2]
|
||||
r2x, r2y, r2z = batch_so3[i, 3], batch_so3[i, 4], batch_so3[i, 5]
|
||||
|
||||
n1 = np.sqrt(r1x * r1x + r1y * r1y + r1z * r1z) + eps
|
||||
r1x /= n1
|
||||
r1y /= n1
|
||||
r1z /= n1
|
||||
|
||||
dot12 = r1x * r2x + r1y * r2y + r1z * r2z
|
||||
r2x -= dot12 * r1x
|
||||
r2y -= dot12 * r1y
|
||||
r2z -= dot12 * r1z
|
||||
n2 = np.sqrt(r2x * r2x + r2y * r2y + r2z * r2z) + eps
|
||||
r2x /= n2
|
||||
r2y /= n2
|
||||
r2z /= n2
|
||||
|
||||
r3x = r1y * r2z - r1z * r2y
|
||||
r3y = r1z * r2x - r1x * r2z
|
||||
r3z = r1x * r2y - r1y * r2x
|
||||
|
||||
R_all[i, 0, 0] = r1x
|
||||
R_all[i, 0, 1] = r1y
|
||||
R_all[i, 0, 2] = r1z
|
||||
R_all[i, 1, 0] = r2x
|
||||
R_all[i, 1, 1] = r2y
|
||||
R_all[i, 1, 2] = r2z
|
||||
R_all[i, 2, 0] = r3x
|
||||
R_all[i, 2, 1] = r3y
|
||||
R_all[i, 2, 2] = r3z
|
||||
return R_all
|
||||
|
||||
|
||||
@jit(nopython=True)
|
||||
def canonicalize_euler_zyx_batch_nb(rpy_batch):
|
||||
"""Canonicalize ZYX Euler angles so each component falls in (-pi, pi]."""
|
||||
N = rpy_batch.shape[0]
|
||||
out = np.empty_like(rpy_batch)
|
||||
two_pi = 2.0 * np.pi
|
||||
|
||||
for i in prange(N):
|
||||
r = rpy_batch[i, 0]
|
||||
p = rpy_batch[i, 1]
|
||||
y = rpy_batch[i, 2]
|
||||
|
||||
r = (r + np.pi) % two_pi - np.pi
|
||||
p = (p + np.pi) % two_pi - np.pi
|
||||
y = (y + np.pi) % two_pi - np.pi
|
||||
|
||||
if p > np.pi / 2.0:
|
||||
p = np.pi - p
|
||||
r = r + np.pi
|
||||
y = y + np.pi
|
||||
elif p <= -np.pi / 2.0:
|
||||
p = -np.pi - p
|
||||
r = r + np.pi
|
||||
y = y + np.pi
|
||||
|
||||
r = (r + np.pi) % two_pi - np.pi
|
||||
p = (p + np.pi) % two_pi - np.pi
|
||||
y = (y + np.pi) % two_pi - np.pi
|
||||
|
||||
out[i, 0] = r
|
||||
out[i, 1] = p
|
||||
out[i, 2] = y
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def so3_to_euler_zyx_batch_nb(batch_so3):
|
||||
matrix = so3_to_matrix_batch_nb(batch_so3)
|
||||
eulers = matrix_to_euler_zyx_batch_nb(matrix)
|
||||
return canonicalize_euler_zyx_batch_nb(eulers)
|
||||
|
||||
|
||||
@jit(nopython=True)
|
||||
def compose_state_and_delta_to_abs_rpy(delta, state):
|
||||
"""Compose a delta (ZYX rpy or 6D) with an absolute state -> absolute rpy(ZYX).
|
||||
|
||||
delta: (N, 3) deltarpy or (N, 6) delta6D. state: (3,) rpy or (6,) 6D.
|
||||
Output: (N, 3) rpy canonicalized into (-pi, pi].
|
||||
"""
|
||||
if delta.shape[-1] == 3:
|
||||
R_delta = euler_to_matrix_zyx_batch_nb(delta)
|
||||
elif delta.shape[-1] == 6:
|
||||
R_delta = so3_to_matrix_batch_nb(delta)
|
||||
else:
|
||||
raise ValueError(f"delta last dim must be 3 or 6, got {delta.shape[-1]}")
|
||||
|
||||
if state.shape[-1] == 3:
|
||||
R_state = euler_to_matrix_zyx_batch_nb(state[np.newaxis, :])[0]
|
||||
elif state.shape[-1] == 6:
|
||||
R_state = so3_to_matrix_batch_nb(state[np.newaxis, :])[0]
|
||||
else:
|
||||
raise ValueError(f"state last dim must be 3 or 6, got {state.shape[-1]}")
|
||||
|
||||
N = R_delta.shape[0]
|
||||
R_abs = np.empty((N, 3, 3), dtype=np.float64)
|
||||
|
||||
S00 = R_state[0, 0]
|
||||
S01 = R_state[0, 1]
|
||||
S02 = R_state[0, 2]
|
||||
S10 = R_state[1, 0]
|
||||
S11 = R_state[1, 1]
|
||||
S12 = R_state[1, 2]
|
||||
S20 = R_state[2, 0]
|
||||
S21 = R_state[2, 1]
|
||||
S22 = R_state[2, 2]
|
||||
|
||||
for i in prange(N):
|
||||
A00 = R_delta[i, 0, 0]
|
||||
A01 = R_delta[i, 0, 1]
|
||||
A02 = R_delta[i, 0, 2]
|
||||
A10 = R_delta[i, 1, 0]
|
||||
A11 = R_delta[i, 1, 1]
|
||||
A12 = R_delta[i, 1, 2]
|
||||
A20 = R_delta[i, 2, 0]
|
||||
A21 = R_delta[i, 2, 1]
|
||||
A22 = R_delta[i, 2, 2]
|
||||
|
||||
R_abs[i, 0, 0] = A00 * S00 + A01 * S10 + A02 * S20
|
||||
R_abs[i, 0, 1] = A00 * S01 + A01 * S11 + A02 * S21
|
||||
R_abs[i, 0, 2] = A00 * S02 + A01 * S12 + A02 * S22
|
||||
|
||||
R_abs[i, 1, 0] = A10 * S00 + A11 * S10 + A12 * S20
|
||||
R_abs[i, 1, 1] = A10 * S01 + A11 * S11 + A12 * S21
|
||||
R_abs[i, 1, 2] = A10 * S02 + A11 * S12 + A12 * S22
|
||||
|
||||
R_abs[i, 2, 0] = A20 * S00 + A21 * S10 + A22 * S20
|
||||
R_abs[i, 2, 1] = A20 * S01 + A21 * S11 + A22 * S21
|
||||
R_abs[i, 2, 2] = A20 * S02 + A21 * S12 + A22 * S22
|
||||
|
||||
abs_rpy = matrix_to_euler_zyx_batch_nb(R_abs)
|
||||
abs_rpy = canonicalize_euler_zyx_batch_nb(abs_rpy)
|
||||
|
||||
return abs_rpy
|
||||
|
||||
|
||||
@jit(nopython=True)
|
||||
def compose_state_and_delta_to_abs_6d(delta, state):
|
||||
"""Compose a 6D delta with a 6D state -> absolute 6D rotation.
|
||||
|
||||
delta: (N, 6). state: (6,). Output: (N, 6).
|
||||
"""
|
||||
R_delta = so3_to_matrix_batch_nb(delta)
|
||||
R_state = so3_to_matrix_batch_nb(state[np.newaxis, :])[0]
|
||||
|
||||
N = R_delta.shape[0]
|
||||
R_abs = np.empty((N, 3, 3), dtype=np.float64)
|
||||
|
||||
S00 = R_state[0, 0]
|
||||
S01 = R_state[0, 1]
|
||||
S02 = R_state[0, 2]
|
||||
S10 = R_state[1, 0]
|
||||
S11 = R_state[1, 1]
|
||||
S12 = R_state[1, 2]
|
||||
S20 = R_state[2, 0]
|
||||
S21 = R_state[2, 1]
|
||||
S22 = R_state[2, 2]
|
||||
|
||||
for i in prange(N):
|
||||
A00 = R_delta[i, 0, 0]
|
||||
A01 = R_delta[i, 0, 1]
|
||||
A02 = R_delta[i, 0, 2]
|
||||
A10 = R_delta[i, 1, 0]
|
||||
A11 = R_delta[i, 1, 1]
|
||||
A12 = R_delta[i, 1, 2]
|
||||
A20 = R_delta[i, 2, 0]
|
||||
A21 = R_delta[i, 2, 1]
|
||||
A22 = R_delta[i, 2, 2]
|
||||
|
||||
R_abs[i, 0, 0] = A00 * S00 + A01 * S10 + A02 * S20
|
||||
R_abs[i, 0, 1] = A00 * S01 + A01 * S11 + A02 * S21
|
||||
R_abs[i, 0, 2] = A00 * S02 + A01 * S12 + A02 * S22
|
||||
|
||||
R_abs[i, 1, 0] = A10 * S00 + A11 * S10 + A12 * S20
|
||||
R_abs[i, 1, 1] = A10 * S01 + A11 * S11 + A12 * S21
|
||||
R_abs[i, 1, 2] = A10 * S02 + A11 * S12 + A12 * S22
|
||||
|
||||
R_abs[i, 2, 0] = A20 * S00 + A21 * S10 + A22 * S20
|
||||
R_abs[i, 2, 1] = A20 * S01 + A21 * S11 + A22 * S21
|
||||
R_abs[i, 2, 2] = A20 * S02 + A21 * S12 + A22 * S22
|
||||
|
||||
abs_6d = np.empty((N, 6), dtype=np.float64)
|
||||
for i in prange(N):
|
||||
abs_6d[i, 0] = R_abs[i, 0, 0]
|
||||
abs_6d[i, 1] = R_abs[i, 0, 1]
|
||||
abs_6d[i, 2] = R_abs[i, 0, 2]
|
||||
abs_6d[i, 3] = R_abs[i, 1, 0]
|
||||
abs_6d[i, 4] = R_abs[i, 1, 1]
|
||||
abs_6d[i, 5] = R_abs[i, 1, 2]
|
||||
|
||||
return abs_6d
|
||||
|
||||
|
||||
__all__ = [
|
||||
"euler_to_matrix_zyx_6d_nb",
|
||||
"euler_to_matrix_zyx_batch_nb",
|
||||
"matrix_to_euler_zyx_batch_nb",
|
||||
"so3_to_matrix_batch_nb",
|
||||
"canonicalize_euler_zyx_batch_nb",
|
||||
"so3_to_euler_zyx_batch_nb",
|
||||
"compose_state_and_delta_to_abs_rpy",
|
||||
"compose_state_and_delta_to_abs_6d",
|
||||
]
|
||||
@@ -0,0 +1,172 @@
|
||||
"""Grounding-point helpers (point / bbox coordinate remap, pure regex)."""
|
||||
|
||||
import re
|
||||
from typing import List, Optional
|
||||
|
||||
|
||||
def process_grounding_points(
|
||||
text: str, orig_height, orig_width, resized_height, resized_width, model_type
|
||||
) -> str:
|
||||
"""Remap <point>/<box>/<bbox> coordinates inside ``text`` from the original
|
||||
image size to the resized space used by the given model type.
|
||||
"""
|
||||
point_pattern = re.compile(r"<(point|box|bbox)>(.*?)</\1>")
|
||||
|
||||
def process_match(match):
|
||||
tag_name = match.group(1)
|
||||
coords_str = match.group(2)
|
||||
try:
|
||||
coords = list(map(int, re.findall(r"\d+", coords_str)))
|
||||
|
||||
scale_w = resized_width / orig_width
|
||||
scale_h = resized_height / orig_height
|
||||
|
||||
if len(coords) == 2:
|
||||
x, y = coords
|
||||
if model_type == "qwen2_5":
|
||||
new_x = max(0, min(round(x * scale_w), resized_width - 1))
|
||||
new_y = max(0, min(round(y * scale_h), resized_height - 1))
|
||||
elif model_type in ["qwen2"]:
|
||||
new_x = max(0, min(999.999, (x / orig_width) * 1000))
|
||||
new_y = max(0, min(999.999, (y / orig_height) * 1000))
|
||||
else:
|
||||
raise ValueError("Unsupported model type")
|
||||
coords = [new_x, new_y]
|
||||
|
||||
if len(coords) == 4:
|
||||
x1, y1, x2, y2 = coords
|
||||
if model_type == "qwen2_5":
|
||||
new_x1 = max(0, min(round(x1 * scale_w), resized_width - 1))
|
||||
new_y1 = max(0, min(round(y1 * scale_h), resized_height - 1))
|
||||
new_x2 = max(0, min(round(x2 * scale_w), resized_width - 1))
|
||||
new_y2 = max(0, min(round(y2 * scale_h), resized_height - 1))
|
||||
elif model_type in ["qwen2"]:
|
||||
new_x1 = max(0, min(999.999, (x1 / orig_width) * 1000))
|
||||
new_y1 = max(0, min(999.999, (y1 / orig_height) * 1000))
|
||||
new_x2 = max(0, min(999.999, (x2 / orig_width) * 1000))
|
||||
new_y2 = max(0, min(999.999, (y2 / orig_height) * 1000))
|
||||
else:
|
||||
raise ValueError("Unsupported model type")
|
||||
coords = [new_x1, new_y1, new_x2, new_y2]
|
||||
|
||||
return f'<{tag_name}>[{", ".join(map(str, coords))}]</{tag_name}>'
|
||||
|
||||
except (ValueError, TypeError):
|
||||
return match.group(0)
|
||||
|
||||
return point_pattern.sub(process_match, text)
|
||||
|
||||
|
||||
def extract_grounding_points(text: str) -> List[List[float]]:
|
||||
"""Extract all <point>/<box>/<bbox> coordinates from ``text`` as list-of-list."""
|
||||
point_pattern = re.compile(r"<(point|box|bbox)>\s*\[([^\]]+)\]\s*</\1>")
|
||||
|
||||
points: List[List[float]] = []
|
||||
for match in point_pattern.finditer(text):
|
||||
coords_str = match.group(2)
|
||||
raw_values = re.findall(r"-?\d+\.?\d*", coords_str)
|
||||
converted: List[float] = []
|
||||
for value in raw_values:
|
||||
number = float(value)
|
||||
converted.append(int(number) if number.is_integer() else number)
|
||||
if converted:
|
||||
points.append(converted)
|
||||
|
||||
return points
|
||||
|
||||
|
||||
def reverse_grounding_points(
|
||||
text: str, orig_height, orig_width, resized_height, resized_width, model_type
|
||||
) -> str:
|
||||
"""Inverse of ``process_grounding_points`` - map resized coords back to original."""
|
||||
point_pattern = re.compile(r"<(point|box|bbox)>(.*?)</\1>")
|
||||
|
||||
def reverse_match(match):
|
||||
tag_name = match.group(1)
|
||||
coords_str = match.group(2)
|
||||
try:
|
||||
coords = list(map(float, re.findall(r"-?\d+\.?\d*", coords_str)))
|
||||
|
||||
scale_w = resized_width / orig_width
|
||||
scale_h = resized_height / orig_height
|
||||
|
||||
if len(coords) == 2:
|
||||
x, y = coords
|
||||
if model_type == "qwen2_5":
|
||||
orig_x = max(0, min(orig_width - 1, round(x / scale_w)))
|
||||
orig_y = max(0, min(orig_height - 1, round(y / scale_h)))
|
||||
elif model_type in ["qwen2"]:
|
||||
orig_x = max(
|
||||
0, min(orig_width - 1, round((x / 1000) * orig_width))
|
||||
)
|
||||
orig_y = max(
|
||||
0, min(orig_height - 1, round((y / 1000) * orig_height))
|
||||
)
|
||||
else:
|
||||
raise ValueError("Unsupported model type")
|
||||
coords = [orig_x, orig_y]
|
||||
|
||||
if len(coords) == 4:
|
||||
x1, y1, x2, y2 = coords
|
||||
if model_type == "qwen2_5":
|
||||
orig_x1 = max(0, min(orig_width - 1, round(x1 / scale_w)))
|
||||
orig_y1 = max(0, min(orig_height - 1, round(y1 / scale_h)))
|
||||
orig_x2 = max(0, min(orig_width - 1, round(x2 / scale_w)))
|
||||
orig_y2 = max(0, min(orig_height - 1, round(y2 / scale_h)))
|
||||
elif model_type in ["qwen2"]:
|
||||
orig_x1 = max(
|
||||
0, min(orig_width - 1, round((x1 / 1000) * orig_width))
|
||||
)
|
||||
orig_y1 = max(
|
||||
0, min(orig_height - 1, round((y1 / 1000) * orig_height))
|
||||
)
|
||||
orig_x2 = max(
|
||||
0, min(orig_width - 1, round((x2 / 1000) * orig_width))
|
||||
)
|
||||
orig_y2 = max(
|
||||
0, min(orig_height - 1, round((y2 / 1000) * orig_height))
|
||||
)
|
||||
else:
|
||||
raise ValueError("Unsupported model type")
|
||||
coords = [orig_x1, orig_y1, orig_x2, orig_y2]
|
||||
|
||||
return f'<{tag_name}>[{", ".join(map(str, map(int, coords)))}]</{tag_name}>'
|
||||
|
||||
except (ValueError, TypeError):
|
||||
return match.group(0)
|
||||
|
||||
return point_pattern.sub(reverse_match, text)
|
||||
|
||||
|
||||
def calculate_point_l1_distance(gt_text: str, pred_text: str) -> Optional[float]:
|
||||
"""Average L1 distance between 2D points extracted from ``gt_text`` / ``pred_text``.
|
||||
|
||||
Returns None if either side has no points or counts differ.
|
||||
"""
|
||||
point_pattern = re.compile(r"<point>\[(\d+),\s*(\d+)\]</point>")
|
||||
|
||||
gt_matches = point_pattern.findall(gt_text)
|
||||
pred_matches = point_pattern.findall(pred_text)
|
||||
|
||||
if not gt_matches or not pred_matches or len(gt_matches) != len(pred_matches):
|
||||
return None
|
||||
|
||||
total_l1_distance = 0.0
|
||||
for (gt_x, gt_y), (pred_x, pred_y) in zip(gt_matches, pred_matches):
|
||||
try:
|
||||
gt_x, gt_y = int(gt_x), int(gt_y)
|
||||
pred_x, pred_y = int(pred_x), int(pred_y)
|
||||
l1_dist = abs(gt_x - pred_x) + abs(gt_y - pred_y)
|
||||
total_l1_distance += l1_dist
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
return total_l1_distance / len(gt_matches) if gt_matches else None
|
||||
|
||||
|
||||
__all__ = [
|
||||
"process_grounding_points",
|
||||
"extract_grounding_points",
|
||||
"reverse_grounding_points",
|
||||
"calculate_point_l1_distance",
|
||||
]
|
||||
@@ -0,0 +1,235 @@
|
||||
"""Multimodal text preprocessing helpers.
|
||||
|
||||
This file is generated by ``scripts/export_opensource.py``. It keeps only the
|
||||
processor wrapper needed by harrix inference and uses a generic system prompt.
|
||||
Internal robot-id, dataset, camera, and frequency maps are not bundled here.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from transformers import BatchFeature
|
||||
from transformers.tokenization_utils_base import BatchEncoding
|
||||
|
||||
|
||||
def pad_text_input_to_target_length(
|
||||
text_inputs, target_length, pad_token_id=151643, padding_side="right"
|
||||
):
|
||||
"""Pad or truncate tokenized text to ``target_length``."""
|
||||
batch_size, current_length = text_inputs.input_ids.shape
|
||||
if current_length < target_length:
|
||||
padding_size = target_length - current_length
|
||||
padding = torch.full(
|
||||
(batch_size, padding_size),
|
||||
pad_token_id,
|
||||
dtype=text_inputs.input_ids.dtype,
|
||||
device=text_inputs.input_ids.device,
|
||||
)
|
||||
attention_padding = torch.zeros(
|
||||
(batch_size, padding_size),
|
||||
dtype=text_inputs.attention_mask.dtype,
|
||||
device=text_inputs.attention_mask.device,
|
||||
)
|
||||
if padding_side == "right":
|
||||
text_inputs["input_ids"] = torch.cat([text_inputs.input_ids, padding], dim=1)
|
||||
text_inputs["attention_mask"] = torch.cat(
|
||||
[text_inputs.attention_mask, attention_padding], dim=1
|
||||
)
|
||||
else:
|
||||
text_inputs["input_ids"] = torch.cat([padding, text_inputs.input_ids], dim=1)
|
||||
text_inputs["attention_mask"] = torch.cat(
|
||||
[attention_padding, text_inputs.attention_mask], dim=1
|
||||
)
|
||||
elif current_length > target_length:
|
||||
if padding_side == "right":
|
||||
text_inputs["input_ids"] = text_inputs.input_ids[:, :target_length]
|
||||
text_inputs["attention_mask"] = text_inputs.attention_mask[:, :target_length]
|
||||
else:
|
||||
text_inputs["input_ids"] = text_inputs.input_ids[:, -target_length:]
|
||||
text_inputs["attention_mask"] = text_inputs.attention_mask[:, -target_length:]
|
||||
return text_inputs
|
||||
|
||||
|
||||
def _replace_media_placeholders(text, grid_thw, token, merge_length):
|
||||
if grid_thw is None:
|
||||
return text
|
||||
index = 0
|
||||
for i in range(len(text)):
|
||||
while token in text[i]:
|
||||
if index >= len(grid_thw):
|
||||
raise ValueError(
|
||||
f"More {token} placeholders than media tensors in sample {i}"
|
||||
)
|
||||
token_count = int(grid_thw[index].prod() // merge_length)
|
||||
text[i] = text[i].replace(token, "<|placeholder|>" * token_count, 1)
|
||||
index += 1
|
||||
text[i] = text[i].replace("<|placeholder|>", token)
|
||||
return text
|
||||
|
||||
|
||||
_PUBLIC_CAMERA_LABELS = {
|
||||
"face_view": "front view",
|
||||
"right_wrist_view": "right wrist view",
|
||||
"left_wrist_view": "left wrist view",
|
||||
}
|
||||
|
||||
|
||||
def _camera_label(cam_name):
|
||||
return _PUBLIC_CAMERA_LABELS.get(str(cam_name), str(cam_name).replace("_", " "))
|
||||
|
||||
|
||||
def preprocesser_call(
|
||||
processor,
|
||||
norm_state=None,
|
||||
agent_pos_mask=None,
|
||||
images=None,
|
||||
prefix_text=None,
|
||||
postfix_text=None,
|
||||
videos=None,
|
||||
padding=False,
|
||||
padding_side="left",
|
||||
truncation=None,
|
||||
max_length=None,
|
||||
return_tensors="pt",
|
||||
pad_prefix_to_same_length=False,
|
||||
pad_to_128_multiple=True,
|
||||
state_augmentation_prob=0.0,
|
||||
state_augmentation_ratio=0.0,
|
||||
state_bins=256,
|
||||
inference_mode=False,
|
||||
**_,
|
||||
):
|
||||
"""Build a ``BatchFeature`` for Wall-X VLA inference.
|
||||
|
||||
This is the inference subset of the internal preprocessing helper: text,
|
||||
image/video placeholder expansion, optional discretized proprioception
|
||||
strings, padding, and labels=None for inference.
|
||||
"""
|
||||
if prefix_text is None:
|
||||
raise ValueError("prefix_text is required")
|
||||
if postfix_text is None:
|
||||
postfix_text = [""] * len(prefix_text)
|
||||
if not isinstance(prefix_text, list):
|
||||
prefix_text = [prefix_text]
|
||||
if not isinstance(postfix_text, list):
|
||||
postfix_text = [postfix_text]
|
||||
batch_size = len(prefix_text)
|
||||
|
||||
if images is not None and len(images) > 0:
|
||||
image_inputs = processor.image_processor(images=images, return_tensors=return_tensors)
|
||||
image_grid_thw = image_inputs["image_grid_thw"]
|
||||
else:
|
||||
image_inputs = {}
|
||||
image_grid_thw = None
|
||||
|
||||
if videos is not None:
|
||||
if hasattr(processor, "video_processor"):
|
||||
videos_inputs = processor.video_processor(videos=videos, return_tensors=return_tensors)
|
||||
else:
|
||||
videos_inputs = processor.image_processor(
|
||||
images=None, videos=videos, return_tensors=return_tensors
|
||||
)
|
||||
video_grid_thw = videos_inputs["video_grid_thw"]
|
||||
else:
|
||||
videos_inputs = {}
|
||||
video_grid_thw = None
|
||||
|
||||
merge_length = processor.image_processor.merge_size**2
|
||||
prefix_text = _replace_media_placeholders(
|
||||
list(prefix_text), image_grid_thw, "<|image_pad|>", merge_length
|
||||
)
|
||||
prefix_text = _replace_media_placeholders(
|
||||
prefix_text, video_grid_thw, "<|video_pad|>", merge_length
|
||||
)
|
||||
|
||||
if norm_state is not None:
|
||||
norm_state = norm_state.cpu().numpy() if isinstance(norm_state, torch.Tensor) else norm_state
|
||||
agent_pos_mask = (
|
||||
agent_pos_mask[:, 0, :].cpu().numpy().astype(bool)
|
||||
if isinstance(agent_pos_mask, torch.Tensor)
|
||||
else agent_pos_mask[:, 0, :].astype(bool)
|
||||
)
|
||||
discretized = np.digitize(norm_state, bins=np.linspace(-1, 1, state_bins + 1)[:-1]) - 1
|
||||
discretized = discretized[:, 0, :]
|
||||
for i in range(batch_size):
|
||||
if "<|propri|>" not in prefix_text[i]:
|
||||
continue
|
||||
state_str = " ".join(map(str, discretized[i, agent_pos_mask[i]]))
|
||||
prefix_text[i] = prefix_text[i].replace("<|propri|>", state_str)
|
||||
|
||||
if not pad_prefix_to_same_length:
|
||||
text = [pre + post for pre, post in zip(prefix_text, postfix_text)]
|
||||
text_inputs = processor.tokenizer(
|
||||
text,
|
||||
return_tensors=return_tensors,
|
||||
padding=padding,
|
||||
padding_side=padding_side,
|
||||
truncation=truncation,
|
||||
max_length=max_length,
|
||||
)
|
||||
text_inputs["prefix_length"] = None
|
||||
else:
|
||||
prefix_inputs = processor.tokenizer(
|
||||
prefix_text,
|
||||
return_tensors=return_tensors,
|
||||
padding=padding,
|
||||
padding_side="left",
|
||||
truncation=truncation,
|
||||
max_length=max_length,
|
||||
)
|
||||
postfix_inputs = processor.tokenizer(
|
||||
postfix_text,
|
||||
return_tensors=return_tensors,
|
||||
padding=padding,
|
||||
padding_side="right",
|
||||
truncation=truncation,
|
||||
max_length=max_length,
|
||||
)
|
||||
text_inputs = BatchEncoding(
|
||||
data={
|
||||
"input_ids": torch.cat([prefix_inputs.input_ids, postfix_inputs.input_ids], dim=1),
|
||||
"attention_mask": torch.cat(
|
||||
[prefix_inputs.attention_mask, postfix_inputs.attention_mask], dim=1
|
||||
),
|
||||
"prefix_length": prefix_inputs.input_ids.shape[1],
|
||||
}
|
||||
)
|
||||
|
||||
pad_token_id = processor.tokenizer.pad_token_id
|
||||
if pad_token_id is None:
|
||||
pad_token_id = processor.tokenizer.eos_token_id
|
||||
if pad_to_128_multiple:
|
||||
target_length = 128 * ((max(len(t) for t in text_inputs.input_ids) + 127) // 128)
|
||||
text_inputs = pad_text_input_to_target_length(
|
||||
text_inputs, target_length, pad_token_id=pad_token_id, padding_side=padding_side
|
||||
)
|
||||
|
||||
text_inputs["labels"] = None if inference_mode else None
|
||||
return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs})
|
||||
|
||||
|
||||
def get_prologue_with_embodied_information(dataset_name, cam_mapping, robot_id, uid, config):
|
||||
"""Return a generic VLA system prompt without private robot maps."""
|
||||
role_start = "<|im_start|>"
|
||||
role_end = "<|im_end|>"
|
||||
prologue = (
|
||||
f"{role_start}system\n"
|
||||
"You are an embodied vision-language-action model controlling a robot "
|
||||
"with language instructions."
|
||||
)
|
||||
if cam_mapping:
|
||||
cameras = ", ".join(_camera_label(name) for name in cam_mapping.values())
|
||||
prologue += f"\nCamera Setup: {cameras}"
|
||||
if not getattr(config, "use_relative_action", False):
|
||||
prologue += "\nAction Space: Abs EEF"
|
||||
else:
|
||||
prologue += "\nAction Space: Rel EEF"
|
||||
return f"{prologue}\n{role_end}\n"
|
||||
|
||||
|
||||
__all__ = [
|
||||
"preprocesser_call",
|
||||
"get_prologue_with_embodied_information",
|
||||
"pad_text_input_to_target_length",
|
||||
]
|
||||
Reference in New Issue
Block a user