suolyerandyangping
d18fa65fa1
* add mot * update libero example * translate zh to en * fix load model from hf * lint * lint --------- Co-authored-by: yangping <yangping@x2robot.com>
1165 lines
42 KiB
Python
1165 lines
42 KiB
Python
import os
|
|
from scipy.fft import dct
|
|
from scipy.fft import idct
|
|
import yaml
|
|
import torch
|
|
import numpy as np
|
|
import dataclasses
|
|
import copy
|
|
import json
|
|
from PIL import Image
|
|
from safetensors.torch import load_file
|
|
from qwen_vl_utils.vision_process import smart_resize
|
|
from transformers import BatchFeature, AutoProcessor
|
|
|
|
from wall_x.model.action_head import Normalizer
|
|
from wall_x.utils.constant import action_statistic_dof as default_action_statistic_dof
|
|
from numba import jit, prange
|
|
|
|
try:
|
|
from spatial_tokenizer.spatial_tokenizer_kdisk import SpatialActionTokenizer
|
|
except ImportError:
|
|
SpatialActionTokenizer = None
|
|
|
|
device = "cuda"
|
|
|
|
dof_config = {
|
|
"follow_left_ee_cartesian_pos": 3,
|
|
"follow_left_ee_rotation": 3,
|
|
"follow_left_gripper": 1,
|
|
"follow_right_ee_cartesian_pos": 3,
|
|
"follow_right_ee_rotation": 3,
|
|
"follow_right_gripper": 1,
|
|
"head_actions": 2,
|
|
"height": 1,
|
|
"car_pose": 3,
|
|
"velocity_decomposed": 3,
|
|
}
|
|
_CAM_NAME_MAPPING = {
|
|
"face_view": "front view",
|
|
"left_wrist_view": "left wrist view",
|
|
"right_wrist_view": "right wrist view",
|
|
"move1_view": "move view",
|
|
"move2_view": "move view",
|
|
"wall_view": "wall view",
|
|
"top_view": "top view",
|
|
"side_view": "side view",
|
|
"global_view": "global view",
|
|
}
|
|
camera_to_view_mapping = {
|
|
"camera_front": "face_view",
|
|
"camera_left": "left_wrist_view",
|
|
"camera_right": "right_wrist_view",
|
|
"camera_side": "side_view",
|
|
"camera_global": "global_view",
|
|
}
|
|
|
|
action_key_mapping = {
|
|
"follow_left_ee_cartesian_pos": "follow1_pos[:3]",
|
|
"follow_left_ee_rotation": "follow1_pos[3:6]",
|
|
"follow_left_gripper": "follow1_pos[6:7]",
|
|
"follow_right_ee_cartesian_pos": "follow2_pos[:3]",
|
|
"follow_right_ee_rotation": "follow2_pos[3:6]",
|
|
"follow_right_gripper": "follow2_pos[6:7]",
|
|
"head_actions": "head_pos",
|
|
"height": "lift",
|
|
"velocity_decomposed": "velocity_decomposed",
|
|
"follow_left_arm_joint_cur": "follow1_joints_cur[-1:]",
|
|
"follow_right_arm_joint_cur": "follow2_joints_cur[-1:]",
|
|
"follow_left_arm_joint_pos": "follow1_pos",
|
|
"follow_right_arm_joint_pos": "follow2_pos",
|
|
}
|
|
|
|
dim_dof_config = {
|
|
"right_xyz": {"rpy": (0, 3), "so3": (0, 3)},
|
|
"right_rot": {"rpy": (3, 6), "so3": (3, 9)},
|
|
"right_gripper": {"rpy": (6, 7), "so3": (9, 10)},
|
|
}
|
|
SINGLE_ARM_DIM = 7
|
|
|
|
|
|
@jit(nopython=True, parallel=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, parallel=True)
|
|
def compose_state_and_delta_to_abs_rpy(delta, state):
|
|
"""
|
|
Input:
|
|
delta: (N,3) -> Δrpy(ZYX) or (N,6) -> Δ6D (first two rows flattened)
|
|
state: (3,) -> rpy(ZYX) or (6,) -> 6D (first two rows flattened)
|
|
Output:
|
|
abs_rpy: (N,3) Absolute pose rpy(ZYX, radians), normalized to (-π, π]
|
|
"""
|
|
if delta.shape[-1] == 3:
|
|
R_delta = euler_to_matrix_zyx_batch_nb(delta) # (N,3,3)
|
|
elif delta.shape[-1] == 6:
|
|
R_delta = so3_to_matrix_batch_nb(delta) # (N,3,3)
|
|
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] # (3,3)
|
|
elif state.shape[-1] == 6:
|
|
R_state = so3_to_matrix_batch_nb(state[np.newaxis, :])[0] # (3,3)
|
|
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
|
|
|
|
|
|
@jit(nopython=True, parallel=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]
|
|
|
|
# normalize r1
|
|
n1 = np.sqrt(r1x * r1x + r1y * r1y + r1z * r1z) + eps
|
|
r1x /= n1
|
|
r1y /= n1
|
|
r1z /= n1
|
|
|
|
# orthogonalize r2 to r1, then normalize
|
|
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
|
|
|
|
# r3 = r1 x r2
|
|
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, parallel=True)
|
|
def matrix_to_euler_zyx_batch_nb(Rs):
|
|
"""
|
|
R = Rz(yaw) * Ry(pitch) * Rx(roll)
|
|
extract:
|
|
pitch = asin(-R[2,0])
|
|
roll = atan2(R[2,1], R[2,2])
|
|
yaw = atan2(R[1,0], R[0,0])
|
|
"""
|
|
N = Rs.shape[0]
|
|
eulers = np.empty((N, 3), dtype=np.float64)
|
|
for i in prange(N):
|
|
r00 = Rs[i, 0, 0]
|
|
# r01 = Rs[i, 0, 1]
|
|
# r02 = Rs[i, 0, 2]
|
|
r10 = Rs[i, 1, 0]
|
|
# r11 = Rs[i, 1, 1]
|
|
# r12 = Rs[i, 1, 2]
|
|
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, parallel=True)
|
|
def canonicalize_euler_zyx_batch_nb(rpy_batch):
|
|
"""
|
|
Batch ZYX Euler Angle Normalization (Parallel Version)
|
|
Input: rpy_batch (N, 3) [roll, pitch, yaw] (radians)
|
|
Output: out (N, 3) Constrained to the same branch with each component in (-π, π]
|
|
Rules:
|
|
1) First, wrap each component to (-π, π]
|
|
2) If p > π/2: p = π - p; r += π; y += π
|
|
If p <= -π/2: p = -π - p; r += π; y += π
|
|
3) Finally, wrap each component to (-π, π] again.
|
|
"""
|
|
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)
|
|
|
|
|
|
def update_model_config(train_config, model_config):
|
|
model_config.use_state_string_representation = train_config["data"].get(
|
|
"use_state_string_representation", False
|
|
)
|
|
model_config.flow_loss_weight = train_config.get("flow_loss_weight", 1.0)
|
|
|
|
model_config.dof_config = train_config["dof_config"]
|
|
model_config.agent_pos_config = train_config["agent_pos_config"]
|
|
|
|
model_config.action_horizon_flow = train_config["data"].get(
|
|
"action_horizon_flow", 32
|
|
)
|
|
|
|
if train_config.get("_attn_implementation", None) is not None:
|
|
model_config._attn_implementation = train_config["_attn_implementation"]
|
|
|
|
return model_config
|
|
|
|
|
|
def move_to_cuda(obj, device=device):
|
|
if isinstance(obj, torch.Tensor):
|
|
return obj.to(device)
|
|
elif isinstance(obj, (dict, BatchFeature)):
|
|
return {k: move_to_cuda(v, device) for k, v in obj.items()}
|
|
elif isinstance(obj, list):
|
|
return [move_to_cuda(v, device) for v in obj]
|
|
elif isinstance(obj, tuple):
|
|
return tuple(move_to_cuda(v, device) for v in obj)
|
|
else:
|
|
return obj
|
|
|
|
|
|
def extract_components(data, config, is_rpy):
|
|
mode = "rpy" if is_rpy else "so3"
|
|
result = {}
|
|
for key, slice_config in config.items():
|
|
slice_range = slice_config[mode]
|
|
result[key] = data[:, slice_range[0] : slice_range[1]]
|
|
return result
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class WallxInferArgs:
|
|
config_path: str | None = None
|
|
checkpoint_path: str | None = None
|
|
action_mode: str = "diffusion" # "ar" or "diffusion"
|
|
|
|
max_time_step: int = 1000
|
|
action_start_ratio: float = 0
|
|
action_end_ratio: float = 0.6
|
|
|
|
model_action_dim: int = 14
|
|
action_horizon: int = 32
|
|
|
|
action_dim: int = 14
|
|
|
|
interpolate_action: bool = False
|
|
interpolate_multiplier: int = 1
|
|
turtle_as_desktop: bool = False
|
|
generate_subtask: bool = False
|
|
subtask_interval: int = 0
|
|
with_cur: bool = False
|
|
state_str: bool = True ### NOTE
|
|
wostate: bool = False
|
|
delta_action: bool = False
|
|
state_rpy: bool = True
|
|
action_rpy: bool = True
|
|
|
|
dataset_name: str = "robochallenge_aloha"
|
|
use_hard_prompt: bool = True
|
|
dct_scale: float = -1
|
|
|
|
|
|
class WallxModelWrapper:
|
|
def __init__(self, args: WallxInferArgs):
|
|
self.args = args
|
|
self.get_model_and_processor()
|
|
self.action_predict_mode = "ar" if args.action_mode == "ar" else "diffusion"
|
|
print("action_predict_mode", self.action_predict_mode, flush=True)
|
|
|
|
def get_model_and_processor(self):
|
|
if self.args.config_path is None:
|
|
self.args.config_path = os.path.join(
|
|
self.args.checkpoint_path, "config.yml"
|
|
)
|
|
with open(self.args.config_path, "r") as f:
|
|
config = yaml.load(f, Loader=yaml.FullLoader)
|
|
self.config = config
|
|
self.dof_config = config["dof_config"]
|
|
self.agent_pos_config = config["agent_pos_config"]
|
|
self.obs_action_keys = [
|
|
"follow_left_ee_cartesian_pos",
|
|
"follow_left_ee_rotation",
|
|
"follow_left_gripper",
|
|
"follow_right_ee_cartesian_pos",
|
|
"follow_right_ee_rotation",
|
|
"follow_right_gripper",
|
|
]
|
|
print("obs_action_keys", self.obs_action_keys, flush=True)
|
|
self.action_tokenizer_type = config.get("action_tokenizer_type", None)
|
|
config_path = config["qwen_vl_act_config_path"]
|
|
self.processor = AutoProcessor.from_pretrained(
|
|
config["processor_path"], use_fast=True
|
|
)
|
|
self.processor.tokenizer.padding_side = "left"
|
|
new_tokens = ["<|propri|>", "<|action|>"]
|
|
|
|
# load fast tokenizer
|
|
self.action_tokenizer_type = config.get("action_tokenizer_type", None)
|
|
print("self.action_tokenizer_type", self.action_tokenizer_type, flush=True)
|
|
self.action_tokenizer = None
|
|
self.action_mapper = None
|
|
if self.action_tokenizer_type:
|
|
# fast
|
|
if self.action_tokenizer_type == "fast":
|
|
print("Using fast tokenizer")
|
|
self.action_tokenizer = AutoProcessor.from_pretrained(
|
|
config["action_tokenizer_path"], trust_remote_code=True
|
|
)
|
|
elif self.action_tokenizer_type == "spatialvla":
|
|
print("Using spatialvla tokenizer")
|
|
assert (
|
|
SpatialActionTokenizer is not None
|
|
), "SpatialActionTokenizer is not installed"
|
|
self.action_tokenizer = SpatialActionTokenizer()
|
|
else:
|
|
raise ValueError(
|
|
f"Unsupported action tokenizer type: {self.action_tokenizer_type}"
|
|
)
|
|
new_tokens += [
|
|
f"<|action_token_{i}|>" for i in range(self.action_tokenizer.vocab_size)
|
|
]
|
|
|
|
# num_added_tokens = self.processor.tokenizer.add_tokens(new_tokens)
|
|
|
|
# define action mapper
|
|
if self.action_tokenizer_type:
|
|
self.action_mapper = {}
|
|
for i in range(self.action_tokenizer.vocab_size):
|
|
token = f"<|action_token_{i}|>"
|
|
token_id = self.processor.tokenizer.convert_tokens_to_ids(token)
|
|
self.action_mapper[token_id] = i
|
|
|
|
# action & propri normalizer
|
|
self._register_normalizers()
|
|
|
|
model_type = config["model_type"]
|
|
|
|
if model_type == "qwen2_5":
|
|
print("Using qwen2_5 model as base model")
|
|
from wall_x.model.qwen2_5_based import (
|
|
Qwen2_5_VLMoEForAction,
|
|
Qwen2_5_VLConfig,
|
|
)
|
|
|
|
ModelClass = Qwen2_5_VLMoEForAction
|
|
ConfigClass = Qwen2_5_VLConfig
|
|
|
|
model_config = ConfigClass.from_pretrained(config_path)
|
|
model_config = update_model_config(config, model_config)
|
|
|
|
print("model_config", model_config, flush=True)
|
|
|
|
# if self.args.action_mode == "ar":
|
|
# model_config._attn_implementation = "flash_attention_2"
|
|
# else
|
|
model_config._attn_implementation = "sdpa"
|
|
model_config.vision_config._attn_implementation = "flash_attention_2"
|
|
|
|
if model_config.model_type == "qwen3_vl":
|
|
self.MAX_PIXELS = 16384 * 32 * 32
|
|
self.MIN_PIXELS = 4 * 32 * 32
|
|
self.IMAGE_FACTOR = 32
|
|
elif model_config.model_type == "qwen2_5_vl":
|
|
self.MAX_PIXELS = 16384 * 28 * 28
|
|
self.MIN_PIXELS = 4 * 28 * 28
|
|
self.IMAGE_FACTOR = 28
|
|
|
|
model = ModelClass(
|
|
model_config,
|
|
self.action_tokenizer_type,
|
|
self.processor,
|
|
self.action_tokenizer,
|
|
self.action_mapper,
|
|
)
|
|
model.resize_token_embeddings(len(self.processor.tokenizer))
|
|
state_dict = load_file(
|
|
self.args.checkpoint_path + "/model.safetensors", device="cpu"
|
|
)
|
|
if os.path.exists(os.path.join(self.args.checkpoint_path, "global_step.pth")):
|
|
global_step = torch.load(
|
|
os.path.join(self.args.checkpoint_path, "global_step.pth")
|
|
)["global_step"]
|
|
print("Loaded global step:", global_step)
|
|
msg = model.load_state_dict(state_dict, strict=False)
|
|
print(msg)
|
|
model.eval()
|
|
model.set_normalizer(
|
|
copy.deepcopy(self.normalizer_action), copy.deepcopy(self.normalizer_propri)
|
|
)
|
|
model.to(device)
|
|
model.to_bfloat16_for_selected_params()
|
|
|
|
self.model = model
|
|
print("self.args.dataset_name", self.args.dataset_name, flush=True)
|
|
print(
|
|
"normalizer_action min",
|
|
self.normalizer_action.min.__getattr__(self.args.dataset_name),
|
|
flush=True,
|
|
)
|
|
print(
|
|
"normalizer_action delta",
|
|
self.normalizer_action.delta.__getattr__(self.args.dataset_name),
|
|
flush=True,
|
|
)
|
|
|
|
def _register_normalizers(self):
|
|
if self.config.get("customized_action_statistic_dof", None):
|
|
action_statistic_dof = json.load(
|
|
open(self.config["customized_action_statistic_dof"], "r")
|
|
)
|
|
else:
|
|
action_statistic_dof = default_action_statistic_dof
|
|
|
|
if os.path.exists(self.args.checkpoint_path + "/normalizer_action.pth"):
|
|
print(
|
|
"Loading normalizer_action from checkpoint",
|
|
self.args.checkpoint_path + "/normalizer_action.pth",
|
|
flush=True,
|
|
)
|
|
self.normalizer_action = Normalizer.from_ckpt(
|
|
self.args.checkpoint_path + "/normalizer_action.pth"
|
|
)
|
|
else:
|
|
self.normalizer_action = Normalizer(
|
|
action_statistic_dof,
|
|
self.config["dof_config"],
|
|
min_key=self.config.get("min_key", "min"),
|
|
delta_key=self.config.get("delta_key", "delta"),
|
|
)
|
|
|
|
# print("action_statistic_dof",action_statistic_dof)
|
|
|
|
if os.path.exists(self.args.checkpoint_path + "/normalizer_propri.pth"):
|
|
print(
|
|
"Loading normalizer_propri from checkpoint",
|
|
self.args.checkpoint_path + "/normalizer_propri.pth",
|
|
flush=True,
|
|
)
|
|
self.normalizer_propri = Normalizer.from_ckpt(
|
|
self.args.checkpoint_path + "/normalizer_propri.pth"
|
|
)
|
|
else:
|
|
self.normalizer_propri = Normalizer(
|
|
action_statistic_dof,
|
|
self.config["agent_pos_config"],
|
|
min_key=self.config.get("min_key", "min"),
|
|
delta_key=self.config.get("delta_key", "delta"),
|
|
)
|
|
|
|
print("self.args.dataset_name", self.args.dataset_name, flush=True)
|
|
print(
|
|
"normalizer_propri min",
|
|
self.normalizer_propri.min.__getattr__(self.args.dataset_name),
|
|
flush=True,
|
|
)
|
|
print(
|
|
"normalizer_propri delta",
|
|
self.normalizer_propri.delta.__getattr__(self.args.dataset_name),
|
|
flush=True,
|
|
)
|
|
print(
|
|
"normalizer_action min",
|
|
self.normalizer_action.min.__getattr__(self.args.dataset_name),
|
|
flush=True,
|
|
)
|
|
print(
|
|
"normalizer_action delta",
|
|
self.normalizer_action.delta.__getattr__(self.args.dataset_name),
|
|
flush=True,
|
|
)
|
|
|
|
def get_text_ar(self, instruction, camera_names, norm_state=None, state_mask=None):
|
|
role_start_symbol = "<|im_start|>"
|
|
role_end_symbol = "<|im_end|>"
|
|
vision_start_symbol = "<|vision_start|>"
|
|
vision_end_symbol = "<|vision_end|>"
|
|
image_pad_symbol = "<|image_pad|>"
|
|
propri_symbol = "<|propri|>"
|
|
|
|
prologue = f"{role_start_symbol}system\nYou are a helpful assistant.{role_end_symbol}\n"
|
|
user_request = f"{role_start_symbol}user\nObservation:"
|
|
print("camera_names", camera_names, flush=True)
|
|
for cam_name in camera_names:
|
|
user_request += f" {_CAM_NAME_MAPPING[cam_name]}: {vision_start_symbol}{image_pad_symbol}{vision_end_symbol}"
|
|
user_request += "\nInstruction:"
|
|
if self.args.state_str:
|
|
assert norm_state is not None
|
|
if isinstance(norm_state, torch.Tensor):
|
|
if state_mask is not None:
|
|
if isinstance(state_mask, torch.Tensor):
|
|
mask_1d = state_mask[0, 0].to(
|
|
dtype=torch.bool, device=norm_state.device
|
|
)
|
|
else:
|
|
mask_1d = torch.as_tensor(state_mask, device=norm_state.device)[
|
|
0, 0
|
|
].to(dtype=torch.bool)
|
|
|
|
norm_state = norm_state[..., mask_1d]
|
|
|
|
norm_state = norm_state.detach().cpu().numpy()
|
|
print("norm_state", norm_state, flush=True)
|
|
discretized_state = (
|
|
np.digitize(norm_state, bins=np.linspace(-1, 1, 256 + 1)[:-1]) - 1
|
|
)
|
|
propri = " ".join(map(str, discretized_state[0, 0]))
|
|
elif self.args.wostate:
|
|
propri = ""
|
|
else:
|
|
propri = propri_symbol
|
|
text_prompt = (
|
|
f"\nPredict the next action in robot action.\nProprioception: {propri}\n"
|
|
)
|
|
user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n"
|
|
assistant_message = f"{role_start_symbol}assistant\n"
|
|
text = prologue + user_message + assistant_message
|
|
|
|
return text
|
|
|
|
def get_text_flow(
|
|
self,
|
|
instruction,
|
|
camera_names,
|
|
action_chunk_size,
|
|
norm_state=None,
|
|
state_mask=None,
|
|
):
|
|
role_start_symbol = "<|im_start|>"
|
|
role_end_symbol = "<|im_end|>"
|
|
vision_start_symbol = "<|vision_start|>"
|
|
vision_end_symbol = "<|vision_end|>"
|
|
image_pad_symbol = "<|image_pad|>"
|
|
propri_symbol = "<|propri|>"
|
|
action_symbol = "<|action|>"
|
|
action_space = "Rel EEF" if self.args.delta_action else "Abs EEF"
|
|
_camera = ", ".join([_CAM_NAME_MAPPING[cam_name] for cam_name in camera_names])
|
|
prologue = f"<|im_start|>system\nYou are an embodied vision-language-action (VLA) model controlling the robot with language instructions.\n Embodiment: {self.args.dataset_name.split('_')[-1]}\n Camera Setup: {_camera},\n Frequency: 32HZ\n Action Space: {action_space}\n<|im_end|>\n"
|
|
|
|
user_request = f"{role_start_symbol}user\nObservation:"
|
|
print("camera_names", camera_names, flush=True)
|
|
for cam_name in camera_names:
|
|
user_request += f" {_CAM_NAME_MAPPING[cam_name]}: {vision_start_symbol}{image_pad_symbol}{vision_end_symbol}"
|
|
user_request += "\nInstruction:"
|
|
if self.args.state_str:
|
|
assert norm_state is not None
|
|
if isinstance(norm_state, torch.Tensor):
|
|
if state_mask is not None:
|
|
if isinstance(state_mask, torch.Tensor):
|
|
mask_1d = state_mask[0, 0].to(
|
|
dtype=torch.bool, device=norm_state.device
|
|
)
|
|
else:
|
|
mask_1d = torch.as_tensor(state_mask, device=norm_state.device)[
|
|
0, 0
|
|
].to(dtype=torch.bool)
|
|
norm_state = norm_state[..., mask_1d]
|
|
norm_state = norm_state.detach().cpu().numpy()
|
|
discretized_state = (
|
|
np.digitize(norm_state, bins=np.linspace(-1, 1, 256 + 1)[:-1]) - 1
|
|
)
|
|
propri = " ".join(map(str, discretized_state[0, 0]))
|
|
elif self.args.wostate:
|
|
propri = ""
|
|
else:
|
|
propri = propri_symbol
|
|
text_prompt = (
|
|
f"\nPredict the next action in robot action.\nProprioception: {propri}\n"
|
|
)
|
|
user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n"
|
|
assistant_message = f"{role_start_symbol}assistant\n"
|
|
action = f"{action_symbol * action_chunk_size}"
|
|
text = prologue + user_message + assistant_message + action
|
|
|
|
return text
|
|
|
|
def resize_images(self, observation):
|
|
image_inputs = []
|
|
view_candidates = [
|
|
"face_view",
|
|
"left_wrist_view",
|
|
"right_wrist_view",
|
|
"side_view",
|
|
"global_view",
|
|
]
|
|
for key in observation.keys():
|
|
if key not in view_candidates:
|
|
print("!!! key not in view_candidates", key, flush=True)
|
|
continue
|
|
# 1. Get the original image
|
|
current_obs = observation[key]
|
|
img_pil = Image.fromarray(current_obs)
|
|
orig_width, orig_height = img_pil.size
|
|
|
|
# 2. Apply resolution limits (if the configuration is not -1)
|
|
target_size = 256
|
|
if target_size != -1:
|
|
# Logic for maintaining aspect ratio constraints
|
|
if orig_width > orig_height:
|
|
new_width = target_size
|
|
new_height = int(target_size * orig_height / orig_width)
|
|
else:
|
|
new_height = target_size
|
|
new_width = int(target_size * orig_width / orig_height)
|
|
img_pil = img_pil.resize((new_width, new_height))
|
|
|
|
# 3. Apply intelligent scaling
|
|
current_width, current_height = img_pil.size
|
|
resized_height, resized_width = smart_resize(
|
|
current_height,
|
|
current_width,
|
|
factor=self.IMAGE_FACTOR,
|
|
min_pixels=self.MIN_PIXELS,
|
|
max_pixels=self.MAX_PIXELS,
|
|
)
|
|
resized_img = img_pil.resize((resized_width, resized_height))
|
|
print("resized_img", resized_img.size, flush=True)
|
|
|
|
image_inputs.append(resized_img)
|
|
|
|
return image_inputs
|
|
|
|
def _construct_input(
|
|
self,
|
|
observation,
|
|
instruction,
|
|
camera_names,
|
|
valid_action_dim=7,
|
|
mode="ar",
|
|
single_image=False,
|
|
):
|
|
additional_inputs = {}
|
|
|
|
agent_pos = torch.from_numpy(observation["agent_pos"])
|
|
agent_pos_mask = torch.from_numpy(observation["agent_pos_mask"])
|
|
dof_mask = torch.from_numpy(observation["dof_mask"])
|
|
additional_inputs["dof_mask"] = dof_mask
|
|
print("before normalizing agent_pos", agent_pos, flush=True)
|
|
|
|
if self.normalizer_propri is not None:
|
|
|
|
agent_pos = self.normalizer_propri.normalize_data(
|
|
agent_pos, [self.args.dataset_name]
|
|
)
|
|
additional_inputs["proprioception"] = agent_pos
|
|
additional_inputs["agent_pos_mask"] = agent_pos_mask
|
|
|
|
print(
|
|
f"normalizing agent_pos: {agent_pos}, {self.args.dataset_name}",
|
|
flush=True,
|
|
)
|
|
print("agent_pos_mask", agent_pos_mask, flush=True)
|
|
print("dof_mask", dof_mask[0, 0], flush=True)
|
|
if mode == "ar":
|
|
text = self.get_text_ar(
|
|
instruction, camera_names, agent_pos, agent_pos_mask
|
|
)
|
|
elif mode == "diffusion":
|
|
text = self.get_text_flow(
|
|
instruction,
|
|
camera_names,
|
|
self.args.action_horizon,
|
|
agent_pos,
|
|
agent_pos_mask,
|
|
)
|
|
elif mode == "subtask":
|
|
text = self.get_text_subtask(instruction, single_view=single_image)
|
|
else:
|
|
raise ValueError(f"Invalid mode: {mode}")
|
|
text = [text]
|
|
|
|
image_inputs = self.resize_images(observation)
|
|
if single_image:
|
|
image_inputs = [
|
|
image_inputs[0]
|
|
] # single view subtask/vqa use head view only
|
|
image_inputs = self.processor.image_processor(
|
|
images=image_inputs, videos=None, return_tensors="pt"
|
|
)
|
|
image_grid_thw = image_inputs["image_grid_thw"]
|
|
# Processing image placeholder tokens in the text
|
|
if image_grid_thw is not None:
|
|
merge_length = self.processor.image_processor.merge_size**2
|
|
index = 0
|
|
for i in range(len(text)):
|
|
while "<|image_pad|>" in text[i]:
|
|
# Replace image placeholders with actual quantities.
|
|
text[i] = text[i].replace(
|
|
"<|image_pad|>",
|
|
"<|placeholder|>"
|
|
* (image_grid_thw[index].prod() // merge_length),
|
|
1,
|
|
)
|
|
index += 1
|
|
text[i] = text[i].replace("<|placeholder|>", "<|image_pad|>")
|
|
|
|
text_inputs = self.processor.tokenizer(
|
|
text, return_tensors="pt", padding=True, truncation=True, max_length=1024
|
|
)
|
|
inputs = BatchFeature(data={**text_inputs, **image_inputs})
|
|
|
|
action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>")
|
|
additional_inputs["moe_token_types"] = inputs.input_ids == action_token_id
|
|
additional_inputs["dataset_names"] = [self.args.dataset_name]
|
|
|
|
inputs.update(additional_inputs)
|
|
inputs = move_to_cuda(inputs, device)
|
|
print("inputs", inputs.keys(), flush=True)
|
|
for k in inputs.keys():
|
|
if isinstance(inputs[k], torch.Tensor):
|
|
print(k, inputs[k].shape, flush=True)
|
|
return inputs
|
|
|
|
def preprocess(self, state, views, valid_action_dim):
|
|
# print("state",state, flush=True)
|
|
# state: dict, keys: follow1_pos, follow2_pos
|
|
model_action_dim = sum(self.dof_config.values())
|
|
|
|
if valid_action_dim not in (SINGLE_ARM_DIM, 2 * SINGLE_ARM_DIM):
|
|
raise ValueError(
|
|
f"Invalid valid_action_dim: {valid_action_dim}, expect 7 or 14"
|
|
)
|
|
|
|
# 1) First, prepare a container of size (1, 1, D), where D = model_action_dim.
|
|
agent_data = np.zeros((1, 1, model_action_dim), dtype=np.float32)
|
|
agent_pos_mask = np.zeros((1, 1, model_action_dim), dtype=np.float32)
|
|
dof_mask = np.zeros(
|
|
(1, self.args.action_horizon, model_action_dim), dtype=np.float32
|
|
)
|
|
|
|
# 2) Determine the interval to be filled [start:end)
|
|
if valid_action_dim == SINGLE_ARM_DIM:
|
|
start = 0 if model_action_dim == SINGLE_ARM_DIM else SINGLE_ARM_DIM
|
|
end = start + SINGLE_ARM_DIM
|
|
|
|
if end > model_action_dim:
|
|
raise ValueError(
|
|
f"model_action_dim={model_action_dim} too small for valid_action_dim=7 "
|
|
f"(need end={end})"
|
|
)
|
|
|
|
follow2 = np.asarray(state["follow2_pos"], dtype=np.float32).reshape(
|
|
1, 1, SINGLE_ARM_DIM
|
|
)
|
|
agent_data[:, :, start:end] = follow2
|
|
agent_pos_mask[:, :, start:end] = 1
|
|
dof_mask[:, :, start:end] = 1
|
|
|
|
else: # valid_action_dim == 14
|
|
end = 2 * SINGLE_ARM_DIM
|
|
if end > model_action_dim:
|
|
raise ValueError(
|
|
f"model_action_dim={model_action_dim} too small for valid_action_dim=14"
|
|
)
|
|
|
|
follow1 = np.asarray(state["follow1_pos"], dtype=np.float32).reshape(
|
|
1, 1, SINGLE_ARM_DIM
|
|
)
|
|
follow2 = np.asarray(state["follow2_pos"], dtype=np.float32).reshape(
|
|
1, 1, SINGLE_ARM_DIM
|
|
)
|
|
agent_data[:, :, :end] = np.concatenate([follow1, follow2], axis=-1)
|
|
agent_pos_mask[:, :, :end] = 1
|
|
dof_mask[:, :, :end] = 1
|
|
|
|
observation = {
|
|
camera_to_view_mapping[key]: views[key][0] for key in views.keys()
|
|
}
|
|
observation["agent_pos"] = agent_data
|
|
observation["agent_pos_mask"] = agent_pos_mask
|
|
observation["dof_mask"] = dof_mask
|
|
return observation
|
|
|
|
def model_output_process(self, action_pred, state):
|
|
|
|
if not self.args.delta_action and self.args.action_rpy:
|
|
return action_pred
|
|
|
|
pred_components = extract_components(
|
|
action_pred, dim_dof_config, self.args.action_rpy
|
|
)
|
|
|
|
pred_right_xyz = pred_components["right_xyz"]
|
|
pred_right_rot = pred_components["right_rot"]
|
|
pred_right_gripper = pred_components["right_gripper"]
|
|
if "left_xyz" in pred_components:
|
|
pred_left_xyz = pred_components["left_xyz"]
|
|
pred_left_rot = pred_components["left_rot"]
|
|
pred_left_gripper = pred_components["left_gripper"]
|
|
else:
|
|
pred_left_xyz = np.zeros((self.args.action_horizon, 3))
|
|
pred_left_rot = np.zeros((self.args.action_horizon, 3))
|
|
pred_left_gripper = np.zeros((self.args.action_horizon, 1))
|
|
|
|
post_action_pred = np.zeros((self.args.action_horizon, self.args.action_dim))
|
|
if self.args.delta_action:
|
|
assert (
|
|
self.args.action_dim == 14
|
|
), "Delta robot support and testing are not yet available."
|
|
|
|
state_components = extract_components(
|
|
state, dim_dof_config, self.args.state_rpy
|
|
)
|
|
left_xyz = state_components["left_xyz"]
|
|
left_rot = state_components["left_rot"]
|
|
right_xyz = state_components["right_xyz"]
|
|
right_rot = state_components["right_rot"]
|
|
|
|
post_action_pred[:, :3] = pred_left_xyz + left_xyz
|
|
post_action_pred[:, 3:6] = compose_state_and_delta_to_abs_rpy(
|
|
pred_left_rot, left_rot[0]
|
|
)
|
|
post_action_pred[:, 6:7] = pred_left_gripper
|
|
post_action_pred[:, 7:10] = pred_right_xyz + right_xyz
|
|
post_action_pred[:, 10:13] = compose_state_and_delta_to_abs_rpy(
|
|
pred_right_rot, right_rot[0]
|
|
)
|
|
post_action_pred[:, 13:14] = pred_right_gripper
|
|
|
|
elif not self.args.action_rpy:
|
|
post_action_pred[:, :3] = pred_left_xyz
|
|
post_action_pred[:, 3:6] = so3_to_euler_zyx_batch_nb(pred_left_rot)
|
|
post_action_pred[:, 6:7] = pred_left_gripper
|
|
post_action_pred[:, 7:10] = pred_right_xyz
|
|
post_action_pred[:, 10:13] = so3_to_euler_zyx_batch_nb(pred_right_rot)
|
|
post_action_pred[:, 13:14] = pred_right_gripper
|
|
else:
|
|
post_action_pred = action_pred
|
|
return post_action_pred
|
|
|
|
def postprocess(self, action_pred, interpolate_multiplier=None):
|
|
|
|
if interpolate_multiplier is None:
|
|
interpolate_multiplier = self.args.interpolate_multiplier
|
|
|
|
if isinstance(action_pred, torch.Tensor):
|
|
action_pred = action_pred.to(torch.float32).cpu().squeeze(0).numpy()
|
|
left_action_pred = action_pred[:, :7] # (32, 7)
|
|
right_action_pred = action_pred[:, 7:14] # (32, 7)
|
|
|
|
start_frame = int(self.args.action_start_ratio * len(left_action_pred))
|
|
end_frame = int(self.args.action_end_ratio * len(left_action_pred))
|
|
left_action_pred = left_action_pred[start_frame:end_frame]
|
|
right_action_pred = right_action_pred[start_frame:end_frame]
|
|
|
|
print("left_action_pred", left_action_pred[-1], flush=True)
|
|
print("right_action_pred", right_action_pred[-1], flush=True)
|
|
|
|
left_action_pred = left_action_pred.tolist()
|
|
right_action_pred = right_action_pred.tolist()
|
|
|
|
serialized_actions = {
|
|
"follow1_pos": left_action_pred,
|
|
"follow2_pos": right_action_pred,
|
|
## for joint-control
|
|
# "follow1_joints":left_action_pred,
|
|
# "follow2_joints":right_action_pred,
|
|
}
|
|
|
|
return serialized_actions
|
|
|
|
def predict_action_rtc(
|
|
self,
|
|
state,
|
|
views,
|
|
instruction=None,
|
|
valid_action_dim=7,
|
|
update_subtask=False,
|
|
action_predict_mode=None,
|
|
):
|
|
if action_predict_mode is not None:
|
|
self.action_predict_mode = action_predict_mode
|
|
|
|
observation = self.preprocess(state, views, valid_action_dim)
|
|
print("use instruction", instruction, flush=True)
|
|
camera_names = [camera_to_view_mapping[key] for key in views.keys()]
|
|
# camera_names = ["right_wrist_view", "global_view", "side_view"]
|
|
print("mode:", self.action_predict_mode, flush=True)
|
|
inputs = self._construct_input(
|
|
observation,
|
|
instruction,
|
|
camera_names=camera_names,
|
|
valid_action_dim=valid_action_dim,
|
|
mode=self.action_predict_mode,
|
|
)
|
|
model_action_dim = sum(self.dof_config.values())
|
|
padding = torch.zeros((1, model_action_dim))
|
|
norm_padding = self.normalizer_action.normalize_data(
|
|
padding, [self.args.dataset_name]
|
|
)
|
|
inputs["padding_action"] = norm_padding
|
|
inputs = move_to_cuda(inputs, device="cuda:0")
|
|
agent_data = observation["agent_pos"][..., : self.args.model_action_dim][0]
|
|
print("before generate_flow_action", flush=True)
|
|
print(inputs.keys(), flush=True)
|
|
print(inputs["dataset_names"], flush=True)
|
|
print(self.processor.tokenizer.decode(inputs["input_ids"][0]), flush=True)
|
|
print(inputs["agent_pos_mask"][0], flush=True)
|
|
print(inputs["dof_mask"][0, 0], flush=True)
|
|
if self.action_predict_mode == "ar":
|
|
action_pred = self.generate_ar_action(inputs)
|
|
else:
|
|
action_pred = self.generate_flow_action(inputs)
|
|
print("after generate_flow_action", flush=True)
|
|
if isinstance(action_pred, torch.Tensor):
|
|
action_pred = action_pred.float().cpu().squeeze(0).numpy()
|
|
|
|
if action_pred is None:
|
|
return None
|
|
|
|
if self.args.dct_scale > 0:
|
|
scale = self.args.dct_scale
|
|
dct_coeff = dct(action_pred, axis=0, norm="ortho")
|
|
dct_coeff = np.around(dct_coeff * scale)
|
|
action_pred = idct(dct_coeff / scale, axis=0, norm="ortho")
|
|
|
|
if action_pred.shape[-1] == 7:
|
|
print("Before concat action_pred", action_pred.shape, flush=True)
|
|
right_action_pred = action_pred
|
|
left_action_pred = np.zeros_like(right_action_pred)
|
|
action_pred = np.concatenate([left_action_pred, right_action_pred], axis=1)
|
|
print("After concat action_pred", action_pred.shape, flush=True)
|
|
# unnorm action_pred
|
|
# print("action_pred", action_pred[:, 3], flush=True)
|
|
action_pred = (
|
|
self.normalizer_action.unnormalize_data(
|
|
torch.tensor(action_pred).unsqueeze(0), [self.args.dataset_name]
|
|
)
|
|
.squeeze(0)
|
|
.numpy()
|
|
)
|
|
print("After unnormalize_data action_pred", action_pred.shape, flush=True)
|
|
|
|
action_pred = self.model_output_process(action_pred, agent_data)
|
|
action_pred = self.postprocess(action_pred)
|
|
|
|
return action_pred
|
|
|
|
def generate_flow_action(
|
|
self,
|
|
inputs,
|
|
last_action_chunk=None,
|
|
max_guidance_weight=20.0,
|
|
num_inference_timesteps=10,
|
|
sigma_action=0.2,
|
|
):
|
|
model_action_dim = sum(self.dof_config.values())
|
|
if last_action_chunk is None:
|
|
output = self.model.generate_flow_action(
|
|
action_horizon=self.args.action_horizon,
|
|
action_dim=model_action_dim,
|
|
num_inference_timesteps=num_inference_timesteps,
|
|
unnorm=False,
|
|
**inputs,
|
|
)
|
|
else:
|
|
output = self.model.generate_flow_action_rtc(
|
|
action_horizon=self.args.action_horizon,
|
|
action_dim=model_action_dim,
|
|
num_inference_timesteps=num_inference_timesteps,
|
|
inference_delay=self.args.rtc_inference_delay,
|
|
execution_horizon=self.args.rtc_execution_horizon - 1,
|
|
max_guidance_weight=max_guidance_weight,
|
|
last_action_chunk=last_action_chunk,
|
|
sigma_action=sigma_action,
|
|
unnorm=False,
|
|
**inputs,
|
|
)
|
|
action_pred = output["predict_action"] # (b, action_horizon, action_dim)
|
|
return action_pred
|
|
|
|
def generate_text(self, inputs):
|
|
return self.model.generate_text(**inputs)
|
|
|
|
def _preprocess_ar_batch(self, batch):
|
|
input_ids = batch["input_ids"]
|
|
attention_mask = batch["attention_mask"]
|
|
moe_token_types = batch["moe_token_types"]
|
|
labels = batch.get("labels", None)
|
|
prefix_length = batch.get("prefix_length", None)
|
|
|
|
generation_prompt_ids = torch.tensor(
|
|
[151644, 77091], device=input_ids.device, dtype=input_ids.dtype
|
|
) # <|im_start|>assistant
|
|
matches = (input_ids[0, :-1] == generation_prompt_ids[0]) & (
|
|
input_ids[0, 1:] == generation_prompt_ids[1]
|
|
)
|
|
if matches.any():
|
|
split_pos = torch.nonzero(matches, as_tuple=True)[0][0].item()
|
|
# construct output ids
|
|
gt_output_ids = input_ids[:, split_pos + 3 : prefix_length]
|
|
# remove output part from input
|
|
input_ids = input_ids[:, : split_pos + 3]
|
|
moe_token_types = moe_token_types[:, : split_pos + 3]
|
|
if attention_mask is not None:
|
|
attention_mask = attention_mask[:, : split_pos + 3]
|
|
if labels is not None:
|
|
labels = labels[:, split_pos + 3 : prefix_length]
|
|
|
|
batch.update(
|
|
{
|
|
"input_ids": input_ids,
|
|
"attention_mask": attention_mask,
|
|
"moe_token_types": moe_token_types,
|
|
"labels": labels,
|
|
"gt_output_ids": gt_output_ids,
|
|
"prefix_length": split_pos + 3,
|
|
}
|
|
)
|
|
|
|
return batch
|
|
|
|
def generate_ar_action(self, inputs):
|
|
# batch = self._preprocess_ar_batch(batch=inputs)
|
|
action_pred = None
|
|
count = 0
|
|
while action_pred is None:
|
|
if count > 5:
|
|
# raise ValueError("re-generate ar action failed")
|
|
return None
|
|
count += 1
|
|
output = self.model.generate_ar_action(
|
|
# action_dim=args.action_dim,
|
|
action_dim=14,
|
|
action_horizon=self.args.action_horizon,
|
|
unnorm=False,
|
|
**inputs,
|
|
)
|
|
action_pred = output["predict_action"]
|
|
|
|
action_pred = action_pred[0]
|
|
return action_pred
|
|
|
|
|
|
class WallxInfer:
|
|
def __init__(self, args: WallxInferArgs):
|
|
self.args = args
|
|
self.model_wrapper = WallxModelWrapper(args)
|
|
|
|
def run_infer_robochallenge(
|
|
self, state, views, instruction, valid_action_dim=7, action_predict_mode=None
|
|
):
|
|
action_pred = self.model_wrapper.predict_action_rtc(
|
|
state=state,
|
|
views=views,
|
|
instruction=instruction,
|
|
valid_action_dim=valid_action_dim,
|
|
action_predict_mode=action_predict_mode,
|
|
)
|
|
return action_pred
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = WallxInferArgs()
|
|
Infer = WallxInfer(args)
|
|
Infer.run_infer()
|