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/infer/base_dataclass.py
+376
View File
@@ -0,0 +1,376 @@
from wall_x.infer.infer_config import InferConfig
from typing import Optional, List
from dataclasses import dataclass, field
import numpy as np
import torch
import wall_x.infer.data_utils as data_utils
from wall_x.infer.logger import InferLogger
dof_dims = {
"left_ee_cartesian_pos": 3,
"left_ee_cartesian_pos_relative": 3,
"left_ee_rotation": 3,
"left_ee_rotation_relative": 3,
"left_ee_rotation_6D": 6,
"left_ee_rotation_6D_relative": 6,
"left_arm_joint_pos": 7,
"left_gripper": 1,
"left_gripper_cur": 1,
"left_arm_joint_cur": 1,
"right_ee_cartesian_pos": 3,
"right_ee_cartesian_pos_relative": 3,
"right_ee_rotation": 3,
"right_ee_rotation_relative": 3,
"right_ee_rotation_6D": 6,
"right_ee_rotation_6D_relative": 6,
"right_arm_joint_pos": 7,
"right_gripper": 1,
"right_gripper_cur": 1,
"right_arm_joint_cur": 1,
"head_actions": 2,
"height": 1,
"car_pose": 3,
"velocity_decomposed": 3,
}
class ComputedDict(dict):
"""Smart dictionary that supports registering computation rules and auto-computes None values on get"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._compute_rules = {} # key -> compute_function
def register_compute_rule(self, key, compute_func):
"""
Register a computation rule.
Args:
key: The key that needs computation
compute_func: Computation function that takes self as argument and returns the computed result
"""
self._compute_rules[key] = compute_func
def get(self, key, default=None):
"""Override get method to support auto-computation"""
value = super().get(key, default)
# If value is None and there's a compute rule, try to compute
if value is None and key in self._compute_rules:
try:
computed_value = self._compute_rules[key](self)
if computed_value is not None:
# Cache the computed result
self[key] = computed_value
return computed_value
except Exception:
pass # If computation fails, return None or default
return value if value is not None else default
def __getitem__(self, key):
"""Override [] operator to support auto-computation"""
value = super().get(key, None)
# If value is None and there's a compute rule, try to compute
if value is None and key in self._compute_rules:
try:
computed_value = self._compute_rules[key](self)
if computed_value is not None:
# Cache the computed result
self[key] = computed_value
return computed_value
except Exception:
pass # If computation fails, raise original KeyError or return None
if key in self:
return super().__getitem__(key)
raise KeyError(key)
@dataclass
class RobotStateActionData:
config: InferConfig = None
data: ComputedDict = field(
default_factory=lambda: ComputedDict(
{
# State (formerly pose) - using state_ prefix
"state_left_ee_cartesian_pos": None, # (1, 3)
"state_left_ee_rotation": None, # (1, 3)
"state_left_ee_rotation_6D": None,
"state_left_arm_joint_pos": None,
"state_left_gripper": None, # (1, 1)
"state_left_arm_joint_cur": None,
"state_left_gripper_cur": None,
"state_right_ee_cartesian_pos": None, # (1, 3)
"state_right_ee_rotation": None,
"state_right_ee_rotation_6D": None, # (1, 6)
"state_right_arm_joint_pos": None,
"state_right_gripper": None,
"state_right_gripper_cur": None,
"state_right_arm_joint_cur": None, # (1, 1)
"state_head_actions": None,
"state_height": None,
"state_car_pose": None,
"state_velocity_decomposed": None,
# Action - using action_ prefix
"action_left_ee_cartesian_pos": None,
"action_left_ee_cartesian_pos_relative": None,
"action_left_ee_rotation": None,
"action_left_ee_rotation_relative": None,
"action_left_ee_rotation_6D": None,
"action_left_ee_rotation_6D_relative": None,
"action_left_gripper": None,
"action_left_arm_joint_pos": None,
"action_right_ee_cartesian_pos": None,
"action_right_ee_cartesian_pos_relative": None,
"action_right_ee_rotation": None,
"action_right_ee_rotation_relative": None,
"action_right_ee_rotation_6D": None,
"action_right_ee_rotation_6D_relative": None,
"action_right_gripper": None,
"action_right_arm_joint_pos": None,
"action_head_actions": None,
"action_height": None,
"action_car_pose": None,
"action_velocity_decomposed": None,
}
)
)
dof_mask: np.ndarray = None
logger = InferLogger.get_robot_logger("RobotStateActionData")
def __post_init__(self):
"""Register computation rules"""
# State computation rules - euler angles -> 6D rotation
self.data.register_compute_rule(
"state_left_ee_rotation_6D",
lambda d: (
data_utils.euler_to_matrix_zyx_6d_nb(d["state_left_ee_rotation"])
if d.get("state_left_ee_rotation") is not None
else None
),
)
self.data.register_compute_rule(
"state_right_ee_rotation_6D",
lambda d: (
data_utils.euler_to_matrix_zyx_6d_nb(d["state_right_ee_rotation"])
if d.get("state_right_ee_rotation") is not None
else None
),
)
# Action computation rules - absolute position computed from relative + state
self.data.register_compute_rule(
"action_left_ee_cartesian_pos",
lambda d: (
d.get("state_left_ee_cartesian_pos")
+ d.get("action_left_ee_cartesian_pos_relative")
if d.get("state_left_ee_cartesian_pos") is not None
and d.get("action_left_ee_cartesian_pos_relative") is not None
else None
),
)
self.data.register_compute_rule(
"action_right_ee_cartesian_pos",
lambda d: (
d.get("state_right_ee_cartesian_pos")
+ d.get("action_right_ee_cartesian_pos_relative")
if d.get("state_right_ee_cartesian_pos") is not None
and d.get("action_right_ee_cartesian_pos_relative") is not None
else None
),
)
# Action computation rules - get absolute rpy
self.data.register_compute_rule( # delta rpy -> abs rpy
"action_left_ee_rotation",
lambda d: (
data_utils.compose_state_and_delta_to_abs_rpy(
d["action_left_ee_rotation_relative"],
d["state_left_ee_rotation"][0],
)
if d.get("action_left_ee_rotation_relative") is not None
and d.get("state_left_ee_rotation") is not None
else None
),
)
self.data.register_compute_rule( # abs 6D -> abs rpy
"action_left_ee_rotation",
lambda d: (
data_utils.so3_to_euler_zyx_batch_nb(d["action_left_ee_rotation_6D"])
if d.get("action_left_ee_rotation_6D") is not None
else None
),
)
self.data.register_compute_rule( # delta 6D -> abs 6D -> abs rpy
"action_left_ee_rotation_6D",
lambda d: (
data_utils.compose_state_and_delta_to_abs_rpy(
d["action_left_ee_rotation_6D_relative"],
d["state_left_ee_rotation_6D"][0],
)
if d.get("action_left_ee_rotation_6D_relative") is not None
and d.get("state_left_ee_rotation_6D") is not None
else None
),
)
self.data.register_compute_rule( # delta rpy -> abs rpy
"action_right_ee_rotation",
lambda d: (
data_utils.compose_state_and_delta_to_abs_rpy(
d["action_right_ee_rotation_relative"],
d["state_right_ee_rotation"][0],
)
if d.get("action_right_ee_rotation_relative") is not None
and d.get("state_right_ee_rotation") is not None
else None
),
)
self.data.register_compute_rule( # abs 6D -> abs rpy
"action_right_ee_rotation",
lambda d: (
data_utils.so3_to_euler_zyx_batch_nb(d["action_right_ee_rotation_6D"])
if d.get("action_right_ee_rotation_6D") is not None
else None
),
)
self.data.register_compute_rule( # delta 6D -> abs 6D -> abs rpy
"action_right_ee_rotation_6D",
lambda d: (
data_utils.compose_state_and_delta_to_abs_rpy(
d["action_right_ee_rotation_6D_relative"],
d["state_right_ee_rotation_6D"][0],
)
if d.get("action_right_ee_rotation_6D_relative") is not None
and d.get("state_right_ee_rotation_6D") is not None
else None
),
)
def get_agent_pos(self, obs_action_keys=None):
if obs_action_keys is None:
obs_action_keys = self.config.train_config["data"]["obs_action_keys"]
agent_pose_data = []
for key in obs_action_keys:
# Remove follow_ or master_ prefix
if key.startswith("follow_"):
key = key.replace("follow_", "")
elif key.startswith("master_"):
key = key.replace("master_", "")
# Add state_ prefix to access state data
state_key = f"state_{key}"
if state_key in self.data:
# Use get method, which will auto-handle None value computation
value = self.data.get(state_key)
if value is None:
# If still None after computation, use zero vector
agent_pose_data.append(np.zeros((1, dof_dims[key])))
else:
agent_pose_data.append(value)
else:
raise ValueError(f"Key {state_key} not found in data")
agent_pose_data = np.concatenate(agent_pose_data, axis=1)[None] # (1, 1, D)
return agent_pose_data
def get_agent_pos_mask(self, obs_action_keys=None):
if obs_action_keys is None:
obs_action_keys = self.config.train_config["data"]["obs_action_keys"]
agent_pos_mask_data = []
for key in obs_action_keys:
# Remove follow_ or master_ prefix
if key.startswith("follow_"):
key = key.replace("follow_", "")
elif key.startswith("master_"):
key = key.replace("master_", "")
# Add state_ prefix to access state data
state_key = f"state_{key}"
if state_key in self.data:
# Use get method, which will auto-handle None value computation
value = self.data.get(state_key)
if value is None:
agent_pos_mask_data.append(np.zeros((1, dof_dims[key])))
else:
agent_pos_mask_data.append(np.ones((1, dof_dims[key])))
else:
raise ValueError(f"Key {state_key} not found in data")
return np.concatenate(agent_pos_mask_data, axis=1)[None] # (1, 1, D)
def save_state_data_with_key(self, value, key):
# Remove follow_ or master_ prefix
key = key.replace("follow_", "")
key = key.replace("master_", "")
# if torch, convert to numpy
if isinstance(value, torch.Tensor):
value = value.detach().cpu().numpy()
if f"state_{key}" not in self.data: # TODO: joint angle control
self.logger.warning(f"{key} is not a valid state key, not recorded")
return
# Shape validation for value, expected shape is (1, D)
if value.shape == (1, dof_dims[key]):
self.data[f"state_{key}"] = value
elif value.shape == (1, 1, dof_dims[key]):
self.data[f"state_{key}"] = value[0]
elif value.shape == (dof_dims[key],):
self.data[f"state_{key}"] = value[None]
else:
raise ValueError(f"Value shape {value.shape} is not legal")
def save_action_data_with_key(self, value, key):
key = key.replace("follow_", "")
key = key.replace("master_", "")
if isinstance(value, torch.Tensor):
value = value.detach().cpu().numpy()
if value.shape == (dof_dims[key],):
self.data[f"action_{key}"] = value[None]
else:
self.data[f"action_{key}"] = value
def save_action_data(
self, predict_action, predict_action_keys: Optional[List[str]] = None
):
if predict_action_keys is None:
predict_action_keys = self.config.data_config["predict_action_keys"]
if isinstance(predict_action, torch.Tensor):
predict_action = predict_action.detach().cpu().numpy()
if predict_action.ndim == 3:
predict_action = predict_action[0]
dof_start = 0
for action_key in predict_action_keys:
action_key = action_key.replace("follow_", "")
action_key = action_key.replace("master_", "")
dof_dim = dof_dims[action_key]
action_key = f"action_{action_key}"
self.data[action_key] = predict_action[:, dof_start : dof_start + dof_dim]
dof_start += dof_dim
# For compatibility, provide convenient property access
@property
def agent_pos(self):
return self.get_agent_pos()
@property
def agent_pos_mask(self):
return self.get_agent_pos_mask()
@property
def action(self):
pass # TODO: support action access
wall_x/infer/env.py
+143
View File
@@ -0,0 +1,143 @@
"""
Base Environment Class for Robot Control and Inference
"""
from typing import Dict, Any, List
from abc import ABC, abstractmethod
import time
from wall_x.infer.infer_config import InferConfig
from wall_x.infer.utils import KeyboardThread
from wall_x.infer.logger import InferLogger
class BaseEnv(ABC):
def __init__(self, config: InferConfig):
self.config = config
self.logger = InferLogger.get_env_logger("Env")
@abstractmethod
def get_observation(self) -> Dict[str, Any]:
raise NotImplementedError
@abstractmethod
def apply_action(self, input: dict) -> None:
raise NotImplementedError
@abstractmethod
def get_instruction(self) -> str:
raise NotImplementedError
def reset(self) -> Dict[str, Any]:
raise NotImplementedError
def stop(self) -> None:
raise NotImplementedError
class RealRobotEnv(BaseEnv):
def __init__(
self, config: InferConfig, instructions: List[str], enable_keyboard: bool = True
):
"""
Args:
config: Inference configuration
instruction: Task instruction
"""
super().__init__(config)
self.instruction = "test"
self.model = self._register_model()
self.robot = self._register_robot()
# Keyboard control
self.keyboard_thread = None
if enable_keyboard:
self.keyboard_thread = KeyboardThread()
# Instruction list
self.instructions = instructions
self.instruction_index = 0
# def _register_model(self) -> WallxModelWrapper:
# return WallxModelWrapper(self.config)
def _register_robot(self):
from wall_x.infer.robot import DesktopRobot, TurtleRobot
if self.config.robot_type == "desktop":
return DesktopRobot(self.config)
elif self.config.robot_type == "turtle":
return TurtleRobot(self.config)
else:
raise ValueError(f"Invalid robot type: {self.config.robot_type}")
def get_observation(self):
return self.robot.get_observation()
def apply_action(self, input: dict):
self.robot.apply_action(input)
def get_instruction(self) -> str:
"""Return task instruction"""
return self.instructions[self.instruction_index]
def reset(self):
self.robot.go_home()
def listen_to_keyboard(self):
if self.keyboard_thread is not None:
if self.keyboard_thread.should_stop:
time.sleep(1)
return True
if self.keyboard_thread.should_reset:
self.reset()
self.keyboard_thread.should_reset = False
time.sleep(1)
return True
if self.keyboard_thread.new_instruction_index is not None:
new_index = self.keyboard_thread.new_instruction_index
# Check if index is valid
if 0 <= new_index < len(self.instructions):
self.instruction_index = new_index
self.logger.info(
f"[Keyboard] Instruction index switched to {new_index}: {self.instructions[new_index]}"
)
else:
self.logger.info(
f"[Keyboard] Invalid instruction index {new_index}, valid range: 0-{len(self.instructions)-1}"
)
# Reset flag
self.keyboard_thread.new_instruction_index = None
time.sleep(1)
return True
return False
def run_infer_flow_action(self):
while True:
if self.listen_to_keyboard():
continue
observation = self.get_observation()
instruction = self.get_instruction()
model_output = self.model.infer_flow_action(observation, instruction)
self.apply_action(model_output)
def run_infer_flow_action_with_subtask(self, subtask_interval: int = 2):
step = 0
subtask = ""
while True:
if self.listen_to_keyboard():
continue
observation = self.get_observation()
instruction = self.get_instruction()
if step == 0 or step % subtask_interval == 0:
subtask = self.model.infer_subtask(observation, instruction)
model_output = self.model.infer_flow_action(observation, subtask)
self.apply_action(model_output)
def run_infer_ar_action(self):
while True:
if self.listen_to_keyboard():
continue
observation = self.get_observation()
instruction = self.get_instruction()
model_output = self.model.infer_ar_action(observation, instruction)
self.apply_action(model_output)
wall_x/infer/env_libero.py
+744
View File
@@ -0,0 +1,744 @@
import os
import json
import numpy as np
from typing import Dict, Any, Tuple, List
from libero.libero import benchmark
from wall_x.infer.env import BaseEnv, InferConfig
from wall_x.serving.policy.wall_x_policy import WallXPolicy
from wall_x.infer.base_dataclass import RobotStateActionData
from wall_x.infer.utils_libero import (
get_libero_env,
get_libero_dummy_action,
get_libero_image,
get_libero_wrist_image,
quat2axisangle,
TaskSuite,
save_rollout_video,
)
from robosuite.wrappers import VisualizationWrapper
def _create_libero_env_standalone(
task_id: int,
task_suite_name: str,
model_family: str = "wallx",
resolution: int = 256,
seed: int = 7,
) -> Any:
"""
Standalone function to create a Libero environment, independent of LiberoRobotEnv instance.
Used for creating environments in subprocess during multi-batch inference,
avoiding serialization of large objects containing the model.
Args:
task_id: Task ID
task_suite_name: Task suite name
model_family: Model family
resolution: Resolution
seed: Random seed
Returns:
Environment instance
"""
from libero.libero import benchmark
# Get task suite and task
benchmark_dict = benchmark.get_benchmark_dict()
task_suite = benchmark_dict[task_suite_name]()
task = task_suite.get_task(task_id)
# Create environment
env, _ = get_libero_env(
task,
model_family=model_family,
resolution=resolution,
seed=seed,
)
# Wrap environment
env.env = VisualizationWrapper(env.env)
env.env.set_visualization_setting(setting="grippers", visible=False)
return env
class LiberoRobotEnv(BaseEnv):
def __init__(
self,
config: InferConfig,
task_suite_name: str = TaskSuite.LIBERO_SPATIAL,
initial_states_path: str = "DEFAULT",
rollout_dir: str = "./rollouts",
model_family: str = "wallx",
resolution: int = 256,
seed: int = 7,
):
super().__init__(config)
self.logger.info(
f"Initializing LiberoRobotEnv (Stateless), task suite: {task_suite_name}"
)
self.model = self._register_model()
self.model_family = model_family
self.resolution = resolution
self.seed = seed
self.logger.info("Importing Libero and related utils...")
self.RobotStateActionData = RobotStateActionData
self.rollout_dir = os.path.join(rollout_dir, task_suite_name)
os.makedirs(self.rollout_dir, exist_ok=True)
if save_rollout_video is not None:
self.save_rollout_video = save_rollout_video
else:
self.save_rollout_video = None
self.logger.warning("save_rollout_video not found, video saving disabled.")
self.task_suite_name = task_suite_name
benchmark_dict = benchmark.get_benchmark_dict()
self.task_suite = benchmark_dict[self.task_suite_name]()
self.num_tasks = self.task_suite.n_tasks
self.initial_states_path = initial_states_path
self.all_initial_states = None
if self.initial_states_path != "DEFAULT":
try:
with open(self.initial_states_path, "r") as f:
self.all_initial_states = json.load(f)
self.logger.info(
f"Loaded custom initial states from {self.initial_states_path}"
)
except Exception as e:
self.logger.error(f"Failed to load initial states file: {e}")
raise
def _register_model(self) -> WallXPolicy:
return WallXPolicy(
model_path=self.config.model_path,
train_config=self.config.train_config,
action_tokenizer_path=self.config.action_tokenizer_path,
action_dim=self.config.action_dim,
agent_pos_dim=self.config.action_dim,
pred_horizon=self.config.pred_horizon,
camera_key=self.config.cam_names,
predict_mode=self.config.predict_mode,
)
def get_instruction(self, task_desc: str) -> str:
return task_desc
def get_observation(self, raw_obs: Dict[str, Any]) -> Dict[str, Any]:
if raw_obs is None:
raise ValueError("Raw observation is None")
data_obj = self.RobotStateActionData(config=self.config)
pos = raw_obs["robot0_eef_pos"]
rot = quat2axisangle(raw_obs["robot0_eef_quat"])
grip = raw_obs["robot0_gripper_qpos"][0:1]
data_obj.save_state_data_with_key(pos[None], "follow_right_ee_cartesian_pos")
data_obj.save_state_data_with_key(rot[None], "follow_right_ee_rotation")
data_obj.save_state_data_with_key(grip[None], "follow_right_gripper")
data_obj.dof_mask = self._get_dof_mask()
face_view = get_libero_image(raw_obs)
right_wrist_view = get_libero_wrist_image(raw_obs)
return {
"robot_state_action_data": data_obj,
"face_view": face_view,
"right_wrist_view": right_wrist_view,
}
def apply_action(
self, input_data: Dict[str, Any], env: Any = None, replay_images: list = None
) -> bool:
if env is None:
raise ValueError(
"In Stateless mode, apply_action must be called with explicit 'env' parameter"
)
action_data = input_data["robot_state_action_data"]
right_arm_traj = self._get_right_arm_action(action_data)
while (
right_arm_traj is not None
and right_arm_traj.ndim > 2
and right_arm_traj.shape[0] == 1
):
right_arm_traj = right_arm_traj.squeeze(0)
done = False
t = 0
try:
for i in range(len(right_arm_traj)):
if done:
break
action_7d = right_arm_traj[i]
obs, reward, done, info = env.step(action_7d)
t += 1
if obs is not None and replay_images is not None:
replay_images.append(get_libero_image(obs))
input_data["_last_obs"] = obs
except Exception as e:
self.logger.error(f"Env step error: {e}")
return False # Error treated as failure
return done, t
def apply_action_batch(
self,
vec_env: Any,
trajectories: List[np.ndarray],
active_indices: List[int],
status_list: List[Dict[str, Any]],
model_outputs: List[Dict[str, Any]],
) -> None:
"""
Execute action trajectories in parallel batch.
Uses SubprocVectorEnv to execute actions in parallel for all active environments.
Integrates vec_env.step(batch_actions, id=still_active) in this function.
"""
if not trajectories:
return
max_traj_len = max(len(traj) for traj in trajectories)
if max_traj_len == 0:
return
for step_idx in range(max_traj_len):
# Check if there are still active environments
still_active = [
idx
for idx in active_indices
if (not status_list[idx]["done"])
and status_list[idx]["count"] > 0
and step_idx < len(trajectories[active_indices.index(idx)])
]
if not still_active:
break
# Build batch actions (only includes actions for still_active environments)
batch_actions = []
for idx in still_active:
traj_idx = active_indices.index(idx)
action_7d = trajectories[traj_idx][step_idx]
# Ensure action_7d is numpy array or list
if isinstance(action_7d, np.ndarray):
batch_actions.append(action_7d)
else:
batch_actions.append(np.array(action_7d))
# Convert to numpy array with shape (batch_size, action_dim)
batch_actions = np.array(batch_actions)
# Execute step in parallel (only for still_active environments)
obs_list, reward_list, done_list, info_list = vec_env.step(
batch_actions, id=still_active
)
# Process returned results
if obs_list.dtype == object:
obs_list = list(obs_list)
else:
obs_list = [obs_list[i] for i in range(len(obs_list))]
done_list = [bool(done_list[i]) for i in range(len(done_list))]
# Update each environment's status
for i, idx in enumerate(still_active):
st = status_list[idx]
obs = obs_list[i]
done = done_list[i]
if obs is not None:
st["current_obs"] = obs
if st["replay_images"] is not None:
st["replay_images"].append(get_libero_image(obs))
st["count"] -= 1
st["success"] = done
st["done"] = done or st["count"] <= 0
# Update model_output's _last_obs
model_outputs[active_indices.index(idx)]["_last_obs"] = obs
def get_task_info(self, task_id: int) -> Tuple[str, Any]:
"""
Get task information (task description and initial states) without creating environment.
Returns:
Tuple[str, Any]: (task_desc, default_initial_states)
"""
if task_id < 0 or task_id >= self.num_tasks:
raise ValueError(f"Invalid task ID: {task_id}")
task = self.task_suite.get_task(task_id)
task_desc = task.language
default_initial_states = self.task_suite.get_task_init_states(task_id)
return task_desc, default_initial_states
def create_env_for_task(self, task_id: int) -> Tuple[Any, str, Any]:
if task_id < 0 or task_id >= self.num_tasks:
raise ValueError(f"Invalid task ID: {task_id}")
task = self.task_suite.get_task(task_id)
default_initial_states = self.task_suite.get_task_init_states(task_id)
env, task_desc = get_libero_env(
task,
model_family=self.model_family,
resolution=self.resolution,
seed=self.seed,
)
env.env = VisualizationWrapper(env.env)
env.env.set_visualization_setting(setting="grippers", visible=False)
return env, task_desc, default_initial_states
def _get_initial_state_for_episode(
self, task_desc: str, default_states: Any, episode_idx: int
) -> np.ndarray:
if self.initial_states_path == "DEFAULT":
if default_states is None:
raise ValueError("Default states missing")
return default_states[episode_idx]
else:
if self.all_initial_states is None:
raise ValueError("Custom states not loaded")
initial_states_task_key = task_desc.replace(" ", "_")
episode_key = f"demo_{episode_idx}"
if not self.all_initial_states[initial_states_task_key][episode_key][
"success"
]:
raise ValueError(f"Expert demo failed for {episode_key}")
return np.array(
self.all_initial_states[initial_states_task_key][episode_key][
"initial_state"
]
)
def reset_env(
self, env: Any, task_desc: str, default_states: Any, episode_idx: int
) -> Any:
try:
if episode_idx >= 0:
state = self._get_initial_state_for_episode(
task_desc, default_states, episode_idx
)
obs = env.set_init_state(state)
return obs
else:
return env.reset()
except Exception as e:
self.logger.error(f"Reset failed: {e}, falling back to default reset")
return env.reset()
def _get_dof_mask(self):
dof_config = self.config.train_config["dof_config"]
total_dof = sum(dof_config.values())
dof_mask = np.ones((1, self.config.action_horizon, total_dof))
mask_keys = [
"follow_left_ee_cartesian_pos",
"follow_left_ee_rotation",
"follow_left_gripper",
"head_actions",
"height",
"velocity_decomposed",
]
start_idx = 0
for key, dof_size in dof_config.items():
if key in mask_keys:
dof_mask[:, :, start_idx : start_idx + dof_size] = 0
start_idx += dof_size
return dof_mask
def _get_right_arm_action(
self, robot_state_action_data: RobotStateActionData
) -> np.ndarray:
right_ee_cartesian_pos = robot_state_action_data.data[
"action_right_ee_cartesian_pos"
]
right_ee_rotation = robot_state_action_data.data["action_right_ee_rotation"]
right_gripper = robot_state_action_data.data["action_right_gripper"]
return np.concatenate(
[right_ee_cartesian_pos, right_ee_rotation, right_gripper], axis=1
)
def _get_left_arm_action(
self, robot_state_action_data: RobotStateActionData
) -> np.ndarray:
left_ee_cartesian_pos = robot_state_action_data.data[
"action_left_ee_cartesian_pos"
]
left_ee_rotation = robot_state_action_data.data["action_left_ee_rotation"]
left_gripper = robot_state_action_data.data["action_left_gripper"]
return np.concatenate(
[left_ee_cartesian_pos, left_ee_rotation, left_gripper], axis=1
)
def _save_rollout(
self,
replay_images: List[np.ndarray],
success: bool,
task_id: int,
task_desc: str,
episode_idx: int,
):
if not self.save_rollout_video or not replay_images:
return
try:
task_name_safe = task_desc.replace(" ", "_").replace(".", "")
filename = f"{episode_idx}{'_SUCCESS' if success else '_FAILURE'}--_{task_name_safe}.mp4"
self.save_rollout_video(
self.rollout_dir,
replay_images,
filename,
success=success,
task_description=task_desc,
log_file=None,
model_family=self.model_family,
)
self.logger.info(f"Saved video: {filename}")
except Exception as e:
self.logger.error(f"Save video failed: {e}")
def run_infer_flow_action(
self,
env: Any,
task_id: int,
task_desc: str,
default_initial_states: Any,
episode_idx: int,
max_infer_times: int = 5,
num_steps_wait: int = 10,
) -> bool:
replay_images = []
num_steps = 0
done = False
count = max_infer_times
current_obs = self.reset_env(
env, task_desc, default_initial_states, episode_idx
)
if current_obs is None:
return False
while num_steps < num_steps_wait:
obs, reward, done, info = env.step(
get_libero_dummy_action(self.model_family)
)
num_steps += 1
if obs is not None:
current_obs = obs
while not done and count > 0:
try:
model_input = self.get_observation(current_obs)
instruction = self.get_instruction(task_desc)
model_input["prompt"] = instruction
model_input["dataset_names"] = "libero_all"
state = np.concatenate(
[
model_input["robot_state_action_data"].data[
"state_right_ee_cartesian_pos"
],
model_input["robot_state_action_data"].data[
"state_right_ee_rotation"
],
model_input["robot_state_action_data"].data[
"state_right_gripper"
],
],
axis=-1,
)
model_input["state"] = state
model_output = self.model.infer(model_input)
model_output["robot_state_action_data"] = model_input[
"robot_state_action_data"
]
model_output["robot_state_action_data"].save_action_data(
model_output["predict_action"]
)
model_output["_last_obs"] = None
done, delta_t = self.apply_action(
model_output, env=env, replay_images=replay_images
)
if model_output.get("_last_obs") is not None:
current_obs = model_output["_last_obs"]
count -= delta_t
except Exception as e:
self.logger.error(f"Episode Error: {e}")
import traceback
traceback.print_exc()
break
success = done
if count <= 0 and not done:
self.logger.warning(
f"Timeout: reached {max_infer_times} steps without success."
)
success = False
self._save_rollout(replay_images, success, task_id, task_desc, episode_idx)
return success
def run_infer_flow_action_batch(
self,
vec_env: Any,
task_ids: List[int] = None,
task_descs: List[str] = None,
default_initial_states_list: List[Any] = None,
episode_indices: List[int] = None,
max_infer_times: int = 5,
num_steps_wait: int = 10,
) -> List[bool]:
"""
Support batch inference: model inference in parallel (batch), environment execution in parallel (SubprocVectorEnv).
Uses SubprocVectorEnv to run environments in subprocess during multi-batch inference, all environments execute actions in parallel.
Returns a list of success flags for each sample.
"""
if vec_env is None:
raise ValueError("vec_env must be specified")
batch_size = len(vec_env)
if task_ids is not None:
assert len(task_ids) == batch_size, "task_ids length must match envs"
if episode_indices is not None:
assert (
len(episode_indices) == batch_size
), "episode_indices length must match envs"
status_list = []
for i in range(batch_size):
status_list.append(
{
"vec_env": vec_env,
"env_id": i, # Index in vec_env
"task_desc": task_descs[i],
"replay_images": [],
"num_steps": 0,
"done": False, # Whether episode has ended
"success": False, # Whether successfully completed
"count": max_infer_times,
"current_obs": None,
"default_states": None,
}
)
# Initialize/reset: Use SubprocVectorEnv to batch set initial states
init_states_to_set = []
for i in range(batch_size):
task_desc = task_descs[i]
default_states = default_initial_states_list[i]
status_list[i]["default_states"] = default_states
ep_i = episode_indices[i]
init_state = self._get_initial_state_for_episode(
task_desc, default_states, ep_i
)
init_states_to_set.append(init_state)
# Batch set initial states
try:
obs_list = vec_env.set_init_state(init_states_to_set)
if obs_list.dtype == object:
obs_list = list(obs_list)
else:
obs_list = [obs_list[i] for i in range(len(obs_list))]
for i, obs in enumerate(obs_list):
if obs is None:
raise ValueError(
f"Reset environment returned None, task_id: {task_ids[i]}, episode_idx: {episode_indices[i]}"
)
status_list[i]["current_obs"] = obs
status_list[i]["done"] = False
status_list[i]["success"] = False
status_list[i]["count"] = max_infer_times
except Exception as e:
self.logger.error(f"Failed to batch set initial states: {e}")
raise
# Warmup steps (batch execution)
dummy_action = get_libero_dummy_action(self.model_family)
dummy_actions = np.array([dummy_action] * batch_size)
for _ in range(num_steps_wait):
obs_list, _, done_list, _ = vec_env.step(dummy_actions)
# Update current_obs
if obs_list.dtype == object:
obs_list = list(obs_list)
else:
obs_list = [obs_list[i] for i in range(len(obs_list))]
for i, obs in enumerate(obs_list):
if obs is not None:
status_list[i]["current_obs"] = obs
# Main loop: model parallel inference, environment parallel execution (SubprocVectorEnv)
while any((not st["done"]) and st["count"] > 0 for st in status_list):
active_indices = [
idx
for idx, st in enumerate(status_list)
if (not st["done"]) and st["count"] > 0
]
if not active_indices:
break
print(f"Batch infer loop, active indices: {active_indices}")
observations = []
instructions = []
for idx in active_indices:
st = status_list[idx]
observations.append(self.get_observation(st["current_obs"]))
instructions.append(self.get_instruction(st["task_desc"]))
# Model batch inference
model_outputs = self.model.infer_flow_action_batch(
observations, instructions
)
# Extract action trajectories for all active environments
trajectories = []
for out in model_outputs:
action_data = out["robot_state_action_data"]
right_arm_traj = self._get_right_arm_action(action_data)
while (
right_arm_traj is not None
and right_arm_traj.ndim > 2
and right_arm_traj.shape[0] == 1
):
right_arm_traj = right_arm_traj.squeeze(0)
if right_arm_traj is None or len(right_arm_traj) == 0:
# If trajectory is empty, create an empty trajectory
right_arm_traj = np.array([]).reshape(0, 7)
trajectories.append(right_arm_traj)
# Use apply_action_batch to execute action trajectories in parallel
try:
self.apply_action_batch(
vec_env=vec_env,
trajectories=trajectories,
active_indices=active_indices,
status_list=status_list,
model_outputs=model_outputs,
)
except Exception as e:
self.logger.error(f"Batch parallel action error: {e}")
# Mark all active environments as failed
for idx in active_indices:
status_list[idx]["done"] = True
status_list[idx]["success"] = False
# Save replay and results
success_list = []
for i, st in enumerate(status_list):
success = st.get("success", False)
if st["count"] <= 0 and not st["success"]:
self.logger.warning(
f"Batch timeout: reached {max_infer_times} steps without success (idx {i})."
)
st["replay_images"] = st.get("replay_images", [])
tid_i = task_ids[i]
epi_i = episode_indices[i]
self._save_rollout(
st["replay_images"],
success,
tid_i,
st["task_desc"],
epi_i,
)
success_list.append(success)
return success_list
def run_infer_ar_action(
self,
env: Any,
task_id: int,
task_desc: str,
default_initial_states: Any,
episode_idx: int,
max_infer_times: int = 10,
num_steps_wait: int = 10,
) -> bool:
replay_images = []
num_steps = 0
done = False
count = max_infer_times
current_obs = self.reset_env(
env, task_desc, default_initial_states, episode_idx
)
if current_obs is None:
self.logger.error("Environment reset returned None.")
return False
while num_steps < num_steps_wait:
obs, reward, done, info = env.step(
get_libero_dummy_action(self.model_family)
)
num_steps += 1
if obs is not None:
current_obs = obs
while not done and count > 0:
try:
model_input = self.get_observation(current_obs)
instruction = self.get_instruction(task_desc)
model_output = self.model.infer_ar_action(model_input, instruction)
model_output["_last_obs"] = None
done, delta_t = self.apply_action(
model_output, env=env, replay_images=replay_images
)
if model_output.get("_last_obs") is not None:
current_obs = model_output["_last_obs"]
else:
if not done:
self.logger.warning(
"Did not receive new observation after apply_action, but episode is not done."
)
count -= delta_t
except Exception as e:
self.logger.error(
f"AR Episode Run Error (Task {task_id}, Ep {episode_idx}): {e}"
)
import traceback
traceback.print_exc()
break
success = done
if count <= 0 and not done:
self.logger.warning(
f"Timeout: AR policy reached {max_infer_times} steps without success."
)
success = False
self._save_rollout(replay_images, success, task_id, task_desc, episode_idx)
return success
wall_x/infer/infer_config.py
+587
View File
@@ -0,0 +1,587 @@
import yaml
import os
from wall_x.model.model_utils import update_model_config
# from x2robot_dataset.configs.config import X2RDataConfig
import json
from typing import List, Dict, Optional, Any
from dataclasses import dataclass, field
@dataclass
class X2RDataConfig:
"""
Unified X2Robot data configuration class (reorganized by README's 5 modules):
1) Data I/O and caching
2) Visual input and sampling (image/camera)
3) Action and time series
4) Instruction and multimodal
5) Data cleaning and alignment (validation/augmentation/framework constraints)
"""
# ----------------------------------------------------------------------
# 1) Data I/O and caching
# ----------------------------------------------------------------------
cache_dir: str = "~/.cache/dataset_cache"
dataset_config_path: Optional[str] = None
use_cache: bool = True
check_mode: bool = True
preload_size: int = 128
buffer_size: int = 20000
batch_size: int = 32
train_test_split: float = 0.9
seed: int = 42
episode_chunk_size: int = (
500 # Commonly used on VG side (number of frames for episode chunking)
)
# ----------------------------------------------------------------------
# 2) Visual input and sampling (image/camera)
# ----------------------------------------------------------------------
# Camera mapping
cam_mapping: Dict[str, str] = field(
default_factory=lambda: {
"faceImg": "face_view",
"leftImg": "left_wrist_view",
"rightImg": "right_wrist_view",
}
)
# Image and augmentation
resolution: Dict[str, int] = field(
default_factory=lambda: {
"face_view": -1,
"left_wrist_view": 128,
"right_wrist_view": 128,
}
)
cam_augmentation_list: List[str] = field(default_factory=list)
# Image time series (history/future)
image_horizon: int = 1
image_history_length: int = 0
image_history_interval: int = 1
future_image_length: int = 0
future_image_interval: int = 1
future_image_indices: Optional[List[int]] = (
None # If provided, length must equal image_horizon
)
# Smart scaling
max_pixels: int = field(
default_factory=lambda: 1280 * 28 * 28
) # Will be replaced with MAX_PIXELS in __post_init__
min_pixels: int = field(
default_factory=lambda: 4 * 28 * 28
) # Will be replaced with MIN_PIXELS in __post_init__
image_factor: int = 28 # Will be replaced with IMAGE_FACTOR in __post_init__
# ----------------------------------------------------------------------
# 3) Action and time series
# ----------------------------------------------------------------------
predict_action_keys: List[str] = field(default_factory=list)
obs_action_keys: List[str] = field(default_factory=list)
# Action window
action_horizon: int = 21
action_history_length: int = 0
action_horizon_flow: int = 32
action_horizon_ar: int = 0
# Padding strategy
left_padding: bool = True
right_padding: bool = True
# Dimension configuration
dof_config: Dict[str, int] = field(default_factory=dict) # Input degrees of freedom
agent_pos_config: Dict[str, int] = field(
default_factory=dict
) # Output degrees of freedom
# State augmentation
state_augmentation_ratio: float = 1.0 # Ratio of augmented states
state_augmentation_prob: float = (
0.1 # Random dimension masking probability for state string
)
state_drop_prob: float = 0.0 # Probability of dropping entire state
# ----------------------------------------------------------------------
# 4) Instruction and multimodal
# ----------------------------------------------------------------------
default_instruction: str = ""
instruction_path: Optional[str] = None
instruction_key: Optional[List[Dict]] = None
multimodal_chunk_size: int = 500
generate_subtask_ratio: float = 0.0
cot_ratio: float = 0.0
multimodal_data_ratio: float = (
0.25 # Multimodal data ratio per batch in VLA dataset
)
instruction_key_prob: Optional[Dict[str, float]] = None
trunc_action_with_instruction: bool = True
use_embodied_system_prompt_ratio: float = 0.0
# ----------------------------------------------------------------------
# 5) Data cleaning and alignment (validation/augmentation/framework constraints)
# ----------------------------------------------------------------------
filter_angle_outliers: bool = False
trim_stationary: bool = False
use_state_string_representation: bool = False
pad_prefix_to_same_length: bool = False
put_ar_predict_in_postfix: bool = (
False # Whether to put ar prediction in postfix, set to True in prediction mode, False in training
)
pad_to_128_multiple: bool = (
False # Triton Attention requirement (deprecated, always set to False)
)
max_seqlen: int = 768
model_type: Optional[str] = None # qwen2_5, qwen2
model_config_path: Optional[str] = (
None # Model config path (used to derive PaddingSide)
)
low_dim_obs_horizon: int = 1 # To be deprecated
# ----------------------------------------------------------------------
# Validation and post-processing
# ----------------------------------------------------------------------
def __post_init__(self):
# TODO: Determine VGA model type validation here
# assert self.model_type in ["qwen2_5", "qwen3"], f"Unsupported model type: {self.model_type}"
if self.model_type == "qwen2_5":
self.max_pixels = 16384 * 28 * 28
self.min_pixels = 4 * 28 * 28
self.image_factor = 28
elif self.model_type == "qwen3":
self.max_pixels = 16384 * 32 * 32
self.min_pixels = 4 * 32 * 32
self.image_factor = 32
# Future image indices validation
if (
self.future_image_indices
and len(self.future_image_indices) != self.image_horizon
):
raise ValueError(
f"future_image_indices length must equal image_horizon: "
f"{len(self.future_image_indices)} != {self.image_horizon}"
)
# Auto-derive action window
if self.action_horizon == 0:
self.action_horizon = max(self.action_horizon_flow, self.action_horizon_ar)
# Auto-derive action keys
if not self.obs_action_keys:
self.obs_action_keys = list(self.agent_pos_config.keys())
if not self.predict_action_keys:
self.predict_action_keys = list(self.dof_config.keys())
# Derive PaddingSide
# @Ryan: Only FlashAttention can use RightPadding, other AttnImpl use LeftPadding
if self.model_config_path is not None:
with open(self.model_config_path, "r", encoding="utf-8") as f:
cfg = json.load(f)
attn_impl = cfg["_attn_implementation"]
if attn_impl == "flash_attention_2":
self.padding_side = "right"
else:
self.padding_side = "left"
# Convenience properties
@property
def use_6D_rotation(self) -> bool:
"""Whether to use 6D rotation (auto-determined from predict_action_keys)"""
if hasattr(self, "_use_6D_rotation"):
return self._use_6D_rotation
self._use_6D_rotation = any("6D" in key for key in self.predict_action_keys)
return self._use_6D_rotation
@property
def use_relative_action(self) -> bool:
"""Whether to use relative action (auto-determined from predict_action_keys)"""
if hasattr(self, "_use_relative_action"):
return self._use_relative_action
self._use_relative_action = any(
"relative" in key for key in self.predict_action_keys
)
return self._use_relative_action
# ----------------------------------------------------------------------
# YAML initialization
# ----------------------------------------------------------------------
@classmethod
def from_yaml_dict(cls, yaml_dict: Dict[str, Any]) -> "X2RDataConfig":
"""
Create config object from YAML config dict. Prioritizes data sub-config, then top-level fields.
"""
data_config = yaml_dict.get("data", {})
params: Dict[str, Any] = {}
# 1) Data I/O and caching
params.update(
{
"cache_dir": data_config.get(
"cache_dir", yaml_dict.get("cache_dir", "~/.cache/dataset_cache")
),
"dataset_config_path": data_config.get(
"dataset_config_path", yaml_dict.get("dataset_config_path", None)
),
"use_cache": data_config.get(
"use_cache", yaml_dict.get("use_cache", True)
),
"check_mode": data_config.get(
"check_mode", yaml_dict.get("check_mode", True)
),
"preload_size": data_config.get(
"preload_size", yaml_dict.get("preload_size", 128)
),
"buffer_size": data_config.get(
"buffer_size", yaml_dict.get("buffer_size", 20000)
),
"batch_size": data_config.get(
"batch_size",
yaml_dict.get(
"batch_size_per_gpu", yaml_dict.get("batch_size", 32)
),
),
"train_test_split": data_config.get("train_test_split", 0.9),
"seed": yaml_dict.get("seed", 42),
"episode_chunk_size": data_config.get("episode_chunk_size", 500),
}
)
# 2) Visual input and sampling (image/camera)
params.update(
{
"cam_mapping": data_config.get(
"cam_mapping",
{
"faceImg": "face_view",
"leftImg": "left_wrist_view",
"rightImg": "right_wrist_view",
},
),
"resolution": data_config.get(
"resolution",
{"face_view": -1, "left_wrist_view": 128, "right_wrist_view": 128},
),
"cam_augmentation_list": data_config.get("cam_augmentation_list", []),
"image_horizon": data_config.get("image_horizon", 1),
"image_history_length": data_config.get("image_history_length", 0),
"image_history_interval": data_config.get("image_history_interval", 1),
"future_image_length": data_config.get("future_image_length", 0),
"future_image_interval": data_config.get("future_image_interval", 1),
"future_image_indices": data_config.get("future_image_indices", None),
"max_pixels": data_config.get("max_pixels", 1280 * 28 * 28),
"min_pixels": data_config.get("min_pixels", 4 * 28 * 28),
"image_factor": data_config.get("image_factor", 28),
}
)
# 3) Action and time series
params.update(
{
"predict_action_keys": data_config.get("predict_action_keys", []),
"obs_action_keys": data_config.get("obs_action_keys", []),
"action_horizon": data_config.get("action_horizon", 0),
"action_history_length": data_config.get("action_history_length", 0),
"action_horizon_flow": data_config.get(
"action_horizon_flow", yaml_dict.get("action_horizon_flow", 32)
),
"action_horizon_ar": data_config.get("action_horizon_ar", 0),
"left_padding": data_config.get("left_padding", True),
"right_padding": data_config.get("right_padding", True),
"dof_config": yaml_dict.get(
"dof_config", data_config.get("dof_config", {})
),
"agent_pos_config": yaml_dict.get(
"agent_pos_config", data_config.get("agent_pos_config", {})
),
"state_augmentation_prob": data_config.get(
"state_augmentation_prob", 0.05
),
"state_drop_prob": data_config.get("state_drop_prob", 0.0),
}
)
# 4) Instruction and multimodal
params.update(
{
"default_instruction": data_config.get("default_instruction", ""),
"instruction_path": data_config.get("instruction_path", None),
"instruction_key": data_config.get("instruction_key", None),
"multimodal_chunk_size": data_config.get("multimodal_chunk_size", 500),
"generate_subtask_ratio": data_config.get(
"generate_subtask_ratio", 0.0
),
"cot_ratio": data_config.get("cot_ratio", 0.0),
"multimodal_data_ratio": data_config.get("multimodal_data_ratio", 0.25),
"instruction_key_prob": data_config.get("instruction_key_prob", None),
"trunc_action_with_instruction": data_config.get(
"trunc_action_with_instruction", True
),
"use_embodied_system_prompt_ratio": data_config.get(
"use_embodied_system_prompt_ratio",
yaml_dict.get("use_embodied_system_prompt_ratio", 0.0),
),
}
)
# 5) Data cleaning and alignment (validation/augmentation/framework constraints)
params.update(
{
"filter_angle_outliers": data_config.get(
"filter_angle_outliers", False
),
"trim_stationary": data_config.get("trim_stationary", False),
"use_state_string_representation": data_config.get(
"use_state_string_representation",
yaml_dict.get("use_state_string_representation", False),
),
"pad_prefix_to_same_length": data_config.get(
"pad_prefix_to_same_length", False
),
"put_ar_predict_in_postfix": data_config.get(
"put_ar_predict_in_postfix", False
),
# "pad_to_128_multiple": data_config.get("pad_to_128_multiple", True),
"padding_side": data_config.get("padding_side", "left"),
"max_seqlen": yaml_dict.get("max_seqlen", 768),
"model_type": yaml_dict.get("model_type", "qwen2_5"),
"model_config_path": yaml_dict.get("qwen_vl_act_config_path", None),
"low_dim_obs_horizon": data_config.get("low_dim_obs_horizon", 1),
}
)
# Only keep valid fields defined in dataclass
valid_fields = {f.name for f in cls.__dataclass_fields__.values()}
filtered = {k: v for k, v in params.items() if k in valid_fields}
return cls(**filtered)
# ----------------------------------------------------------------------
# Dict-style access (for compatibility with existing calls)
# ----------------------------------------------------------------------
def __getitem__(self, key: str):
try:
return getattr(self, key)
except AttributeError:
raise KeyError(f"'{key}' not found in {self.__class__.__name__}")
def __setitem__(self, key: str, value):
setattr(self, key, value)
def __contains__(self, key: str) -> bool:
return hasattr(self, key)
def keys(self):
return self.__dict__.keys()
def values(self):
return self.__dict__.values()
def items(self):
return self.__dict__.items()
class InferConfig:
def __init__(
self,
checkpoint_path: str | None = None,
train_config_path: str | None = None,
robot_host: str = "0.0.0.0",
robot_port: int = 33723,
robot_id: str = "10053",
robot_type: str = "desktop", # ["desktop", "turtle"]
robot_action_start_ratio: float = 0, # Action execution start ratio
robot_action_end_ratio: float = 0.8, # Action execution end ratio
robot_action_interpolate_multiplier: int = 70, # Action interpolation
robot_use_joint_angle_control: bool = False, # Use joint control (model must be joint prediction model)
turtle_as_desktop: bool = False, # Use turtle body for desktop operation, fixed chassis head movement, head camera, and chassis height
action_horizon: int = 10, # Please correctly fill in the model's horizon
action_dim: int | None = None,
model_device: str = "cuda:0",
num_inference_timesteps: int = 10,
norm_key: str = "x2_normal",
cam_names: list[str] = ["face_view", "right_wrist_view"],
):
# Private attribute for storing path
assert checkpoint_path is not None
self._checkpoint_path = checkpoint_path
if os.path.exists(os.path.join(checkpoint_path, "normalizer_action.pth")):
self.normalizer_action_path = os.path.join(
checkpoint_path, "normalizer_action.pth"
)
if os.path.exists(os.path.join(checkpoint_path, "normalizer_propri.pth")):
self.normalizer_propri_path = os.path.join(
checkpoint_path, "normalizer_propri.pth"
)
self.model_path = checkpoint_path
self.action_tokenizer_path = "/x2robot_v2/Models/fast/"
# Other configuration attributes
self.robot_host = robot_host
self.robot_port = robot_port
self.robot_type = robot_type # ["desktop", "turtle"]
self.robot_id = robot_id
self.robot_action_start_ratio = robot_action_start_ratio
self.robot_action_end_ratio = robot_action_end_ratio
self.robot_action_interpolate_multiplier = robot_action_interpolate_multiplier
self.robot_use_joint_angle_control = (
robot_use_joint_angle_control # Use joint angle control
)
self.turtle_as_desktop = turtle_as_desktop
self._action_horizon = (
action_horizon # Default controlled by train config's flow action horizon
)
self._action_dim = action_dim # Default determined by train config's dof config
self.action_dim = action_dim
self.pred_horizon = action_horizon
self.predict_mode = "diffusion"
self.camera_key = cam_names
self.model_device = model_device
self.num_inference_timesteps = (
num_inference_timesteps # flow matching related config
)
# Initialize config objects
self.train_config: dict = {}
self.model_config = None
self.data_config = None
self.norm_key = norm_key
self.cam_names = cam_names
# Load all configs
self._load_all_configs(train_config_path)
@property
def checkpoint_path(self) -> str | None:
return self._checkpoint_path
@checkpoint_path.setter
def checkpoint_path(self, value: str | None):
"""When checkpoint_path is updated, reload all configs"""
if self._checkpoint_path != value:
self._checkpoint_path = value
self._load_all_configs()
@property
def action_horizon(self) -> int:
return self._action_horizon
@action_horizon.setter
def action_horizon(self, value: int):
self._action_horizon = value
@property
def action_dim(self) -> int | None:
return self._action_dim
@action_dim.setter
def action_dim(self, value: int | None):
self._action_dim = value
def _load_all_configs(self, train_config_path=None):
"""Unified entry point for loading all configs"""
self._load_train_config(train_config_path)
self._load_model_config()
self._load_data_config()
# Update action_horizon and action_dim (if needed)
if self._action_horizon is None:
self._action_horizon = self.train_config.get("data", {}).get(
"action_horizon_flow", 32
)
assert self._action_horizon is not None and self._action_horizon > 0
if self._action_dim is None:
self._action_dim = sum(self.train_config.get("dof_config", {}).values())
def _load_train_config(self, train_config_path):
if train_config_path is None:
train_config_path = os.path.join(self._checkpoint_path, "config.yml")
with open(train_config_path, "r") as f:
self.train_config = yaml.load(f, Loader=yaml.FullLoader)
ckpt_dir = self._checkpoint_path
preprocessor_file = os.path.join(ckpt_dir, "preprocessor_config.json")
if os.path.exists(preprocessor_file):
print(f"[LoadConfig] Found {preprocessor_file}, override processor_path.")
self.train_config["processor_path"] = ckpt_dir
tokenizer_file = os.path.join(ckpt_dir, "tokenizer.json")
tokenizer_config_file = os.path.join(ckpt_dir, "tokenizer_config.json")
if "action_tokenizer_path" in self.train_config and not os.path.exists(
self.train_config["action_tokenizer_path"]
):
if os.path.exists(tokenizer_file) and os.path.exists(tokenizer_config_file):
print(
f"[LoadConfig] Found tokenizer files in {ckpt_dir}, override action_tokenizer_path."
)
self.train_config["action_tokenizer_path"] = ckpt_dir
else:
print("[LoadConfig] Cannot load action tokenizer! ")
def _load_model_config(self):
ckpt_config_path = os.path.join(self._checkpoint_path, "config.json")
resolved_cfg_path = None
if os.path.exists(ckpt_config_path):
# Prefer checkpoint config
resolved_cfg_path = ckpt_config_path
print(f"[LoadModelConfig] Using checkpoint config.json: {ckpt_config_path}")
else:
# Fallback to original config path
fallback_cfg = self.train_config.get("qwen_vl_act_config_path", None)
if fallback_cfg is not None:
resolved_cfg_path = fallback_cfg
print(f"[LoadModelConfig] Using fallback act config: {fallback_cfg}")
if resolved_cfg_path is None or (not os.path.exists(resolved_cfg_path)):
raise ValueError(
f"[LoadModelConfig] Cannot load model config! "
f"Checked:\n"
f" - Checkpoint config.json: {ckpt_config_path}\n"
f" - Fallback path: {self.train_config.get('qwen_vl_act_config_path', None)}"
)
# Save back to config for consistency
self.train_config["qwen_vl_act_config_path"] = resolved_cfg_path
model_type = self.train_config["model_type"]
if model_type == "qwen2_5":
from wall_x.model.qwen2_5_based import Qwen2_5_VLConfig
ConfigClass = Qwen2_5_VLConfig
# elif model_type == "qwen3":
# from wall_x.model.qwen3_based import Qwen3VLConfig
# ConfigClass = Qwen3VLConfig
else:
raise ValueError(f"[LoadModelConfig] Unsupported model type: {model_type}")
print(f"[LoadModelConfig] Loading model config from: {resolved_cfg_path}")
self.model_config = ConfigClass.from_pretrained(resolved_cfg_path)
self.model_config = update_model_config(self.train_config, self.model_config)
self.model_config._attn_implementation = "sdpa"
self.model_config.vision_config._attn_implementation = "flash_attention_2"
print("[LoadModelConfig] Model config loaded and updated successfully.")
def _load_data_config(self):
self.data_config = X2RDataConfig.from_yaml_dict(self.train_config)
if __name__ == "__main__":
config = InferConfig()
print(config.train_config)
wall_x/infer/logger.py
+299
View File
@@ -0,0 +1,299 @@
"""
Hierarchical Inference Logging System
Level structure:
- ENV: Environment layer (RealRobotEnv)
- ROBOT: Robot layer (Robot subclasses)
- CONTROLLER: Controller layer (RobotController, RobotCommunication)
- MODEL: Model layer (WallxModelWrapper)
- UTILS: Utility layer (various utility classes)
Usage examples:
# Method 1: Auto-detect level
from wall_x.infer.logger import get_logger
logger = get_logger(__name__)
logger.info("This is an info message")
# Method 2: Manually specify level
logger = get_logger(__name__, "ROBOT")
logger.debug("Robot state updated")
# Method 3: Use shortcut methods
from wall_x.infer.logger import InferLogger
logger = InferLogger.get_robot_logger("DesktopRobot")
logger.warning("Action out of bounds")
"""
import logging
import sys
from typing import Optional
from pathlib import Path
from datetime import datetime
try:
import colorlog
HAS_COLORLOG = True
except ImportError:
HAS_COLORLOG = False
print("[WARNING] colorlog not installed. Install with: pip install colorlog")
class InferLogger:
"""
Hierarchical inference logging system
"""
_loggers = {}
_initialized = False
# Level definitions
LEVEL_ENV = "ENV"
LEVEL_ROBOT = "ROBOT"
LEVEL_CONTROLLER = "CONTROLLER"
LEVEL_MODEL = "MODEL"
LEVEL_UTILS = "UTILS"
# Level color mapping (for terminal output)
LEVEL_COLORS = {
LEVEL_ENV: "cyan",
LEVEL_ROBOT: "green",
LEVEL_CONTROLLER: "yellow",
LEVEL_MODEL: "purple", # colorlog uses 'purple' not 'magenta'
LEVEL_UTILS: "blue",
}
@classmethod
def setup(
cls,
log_level: str = "INFO",
log_dir: Optional[str] = None,
console_output: bool = True,
file_output: bool = True,
colorful: bool = True,
):
"""
Initialize the logging system
Args:
log_level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
log_dir: Log file directory
console_output: Whether to output to console
file_output: Whether to output to file
colorful: Whether to use colorful output (requires colorlog)
"""
if cls._initialized:
return
cls.log_level = getattr(logging, log_level.upper())
cls.console_output = console_output
cls.file_output = file_output
cls.colorful = colorful and HAS_COLORLOG
# Create log directory
if file_output and log_dir:
cls.log_dir = Path(log_dir)
cls.log_dir.mkdir(parents=True, exist_ok=True)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
cls.log_file = cls.log_dir / f"infer_{timestamp}.log"
else:
cls.log_file = None
cls._initialized = True
@classmethod
def get_logger(cls, name: str, level: str = None) -> logging.Logger:
"""
Get logger for specified level
Args:
name: Logger name (usually module name or class name)
level: Level identifier (ENV, ROBOT, CONTROLLER, MODEL, UTILS)
Returns:
Configured logger instance
"""
if not cls._initialized:
cls.setup()
# Auto-detect level
if level is None:
level = cls._detect_level(name)
logger_key = f"{level}.{name}"
if logger_key in cls._loggers:
return cls._loggers[logger_key]
# Create new logger
logger = logging.getLogger(logger_key)
logger.setLevel(cls.log_level)
logger.propagate = False
# Clear existing handlers
logger.handlers.clear()
# Console output
if cls.console_output:
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setLevel(cls.log_level)
if cls.colorful:
# Colorful formatting
color = cls.LEVEL_COLORS.get(level, "white")
console_format = (
f"%(log_color)s[%(asctime)s]%(reset)s "
f"%(bold_{color})s[{level:^10}]%(reset)s "
f"%(bold_white)s[%(name)s]%(reset)s "
f"%(log_color)s%(levelname)-8s%(reset)s "
f"%(message)s"
)
console_formatter = colorlog.ColoredFormatter(
console_format,
datefmt="%H:%M:%S",
log_colors={
"DEBUG": "cyan",
"INFO": "green",
"WARNING": "yellow",
"ERROR": "red",
"CRITICAL": "bold_red",
},
)
else:
# Plain formatting
console_format = (
f"[%(asctime)s] [{level:^10}] [%(name)s] "
f"%(levelname)-8s %(message)s"
)
console_formatter = logging.Formatter(
console_format, datefmt="%H:%M:%S"
)
console_handler.setFormatter(console_formatter)
logger.addHandler(console_handler)
# File output
if cls.file_output and cls.log_file:
file_handler = logging.FileHandler(cls.log_file, encoding="utf-8")
file_handler.setLevel(cls.log_level)
file_format = (
f"[%(asctime)s] [{level:^10}] [%(name)s] "
f"%(levelname)-8s %(message)s"
)
file_formatter = logging.Formatter(file_format, datefmt="%Y-%m-%d %H:%M:%S")
file_handler.setFormatter(file_formatter)
logger.addHandler(file_handler)
cls._loggers[logger_key] = logger
return logger
@classmethod
def _detect_level(cls, name: str) -> str:
"""Auto-detect level based on name"""
name_lower = name.lower()
if "env" in name_lower:
return cls.LEVEL_ENV
elif "robot" in name_lower and "controller" not in name_lower:
return cls.LEVEL_ROBOT
elif (
"controller" in name_lower
or "communication" in name_lower
or "socket" in name_lower
):
return cls.LEVEL_CONTROLLER
elif "model" in name_lower or "wrapper" in name_lower:
return cls.LEVEL_MODEL
else:
return cls.LEVEL_UTILS
@classmethod
def get_env_logger(cls, name: str = "Environment") -> logging.Logger:
"""Get environment layer logger"""
return cls.get_logger(name, cls.LEVEL_ENV)
@classmethod
def get_robot_logger(cls, name: str = "Robot") -> logging.Logger:
"""Get robot layer logger"""
return cls.get_logger(name, cls.LEVEL_ROBOT)
@classmethod
def get_controller_logger(cls, name: str = "Controller") -> logging.Logger:
"""Get controller layer logger"""
return cls.get_logger(name, cls.LEVEL_CONTROLLER)
@classmethod
def get_model_logger(cls, name: str = "Model") -> logging.Logger:
"""Get model layer logger"""
return cls.get_logger(name, cls.LEVEL_MODEL)
@classmethod
def get_utils_logger(cls, name: str = "Utils") -> logging.Logger:
"""Get utility layer logger"""
return cls.get_logger(name, cls.LEVEL_UTILS)
@classmethod
def set_level(cls, level: str):
"""Dynamically modify log level for all loggers"""
new_level = getattr(logging, level.upper())
cls.log_level = new_level
for logger in cls._loggers.values():
logger.setLevel(new_level)
for handler in logger.handlers:
handler.setLevel(new_level)
@classmethod
def close_all(cls):
"""Close file handles for all loggers"""
for logger in cls._loggers.values():
for handler in logger.handlers[:]:
handler.close()
logger.removeHandler(handler)
cls._loggers.clear()
cls._initialized = False
# Convenience functions
def get_logger(name: str, level: str = None) -> logging.Logger:
"""
Convenience function to get logger
Args:
name: Logger name (usually use __name__)
level: Level identifier (optional, will auto-detect)
Returns:
Configured logger instance
Usage examples:
from wall_x.infer.logger import get_logger
logger = get_logger(__name__) # Auto-detect level
logger = get_logger(__name__, "ROBOT") # Manually specify level
"""
return InferLogger.get_logger(name, level)
def setup_logger(
log_level: str = "INFO",
log_dir: Optional[str] = None,
console_output: bool = True,
file_output: bool = True,
colorful: bool = True,
):
"""
Convenience function to setup logging system
Args:
log_level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
log_dir: Log file directory
console_output: Whether to output to console
file_output: Whether to output to file
colorful: Whether to use colorful output
Usage examples:
from wall_x.infer.logger import setup_logger
setup_logger(log_level="DEBUG", log_dir="./logs")
"""
InferLogger.setup(log_level, log_dir, console_output, file_output, colorful)
wall_x/infer/utils.py
+283
View File
@@ -0,0 +1,283 @@
import numpy as np
from scipy.signal import savgol_filter
from scipy.spatial.transform import Rotation as R # TODO: Convert to numba functions
from collections import deque
import threading
from wall_x.infer.logger import InferLogger
class KeyboardThread(threading.Thread):
"""
Simple keyboard listening thread that provides stop and reset functionality
"""
def __init__(self):
self.should_reset = False
self.should_stop = False
self.new_instruction_index = None # Used to store new instruction index
self.logger = InferLogger.get_utils_logger("KeyboardThread")
super(KeyboardThread, self).__init__(name="keyboard-thread", daemon=True)
self.show_help()
self.start()
def run(self):
"""Listen to keyboard input"""
while True:
try:
user_input = input().strip().lower()
if user_input in ["s", "stop"]:
self.should_stop = not self.should_stop
self.logger.info("[Keyboard] Stop signal sent")
elif user_input in ["r", "reset"]:
self.logger.info("[Keyboard] Executing reset...")
self.should_reset = True
self.logger.info("[Keyboard] Reset signal sent")
elif user_input.isdigit():
# Handle digit input, switch to corresponding instruction index
index = int(user_input)
self.new_instruction_index = index
self.logger.info(
f"[Keyboard] Switched to instruction index: {index}"
)
else:
self.logger.info(
f"[Keyboard] Received input: {user_input}. No action taken."
)
except EOFError:
break
except Exception as e:
self.logger.error(f"[Keyboard] Error: {e}")
def show_help(self):
self.logger.info(
"[Keyboard] Keyboard control: Enter 's' to stop, 'r' to reset, 'number' to switch instruction index"
)
# Robot arm trajectory parameters
ARM_MAX_VELOCITY = 0.02
ARM_EXECUTION_HZ = 20
ARM_MIN_EXECUTION_TIME = 5.0
ARM_MAX_EXECUTION_TIME = 15.0
class UnifiedTrajectoryProcessor:
"""Unified trajectory processor"""
@staticmethod
def interpolate_trajectory_batch(trajectories, target_length, smooth=True):
"""
Batch interpolate multiple trajectories to unified length
Args:
trajectories: list of np.array, each array with shape (N, D)
target_length: int, target length
smooth: bool, whether to smooth
Returns:
list of np.array, interpolated trajectories
"""
if not trajectories:
return []
results = []
for traj in trajectories:
if len(traj) == 0:
results.append(np.zeros((target_length, traj.shape[1])))
continue
if len(traj) == target_length:
results.append(traj)
continue
# Vectorized interpolation
original_indices = np.linspace(0, len(traj) - 1, len(traj))
target_indices = np.linspace(0, len(traj) - 1, target_length)
# Handle different types of data
if traj.shape[1] == 7: # Robot arm data [x,y,z,rx,ry,rz,gripper]
interpolated = UnifiedTrajectoryProcessor._interpolate_arm_trajectory(
traj, original_indices, target_indices, target_length
)
else: # Other data (height, current, etc.)
interpolated = np.zeros((target_length, traj.shape[1]))
for i in range(traj.shape[1]):
interpolated[:, i] = np.interp(
target_indices, original_indices, traj[:, i]
)
# Smooth processing
if smooth and len(interpolated) >= 5:
interpolated = UnifiedTrajectoryProcessor._smooth_trajectory(
interpolated
)
results.append(interpolated)
return results
@staticmethod
def _interpolate_arm_trajectory(
traj, original_indices, target_indices, target_length
):
"""Optimized robot arm trajectory interpolation"""
interpolated = np.zeros((target_length, 7))
# Vectorized interpolation for position and gripper
for i in [0, 1, 2, 6]: # x, y, z, gripper
interpolated[:, i] = np.interp(target_indices, original_indices, traj[:, i])
# Quaternion interpolation (vectorized)
quaternions = R.from_euler("xyz", traj[:, 3:6]).as_quat()
interpolated_quats = np.zeros((target_length, 4))
for i in range(4):
interpolated_quats[:, i] = np.interp(
target_indices, original_indices, quaternions[:, i]
)
# Batch normalization
norms = np.linalg.norm(interpolated_quats, axis=1, keepdims=True)
interpolated_quats = interpolated_quats / norms
# Batch convert back to Euler angles
interpolated[:, 3:6] = R.from_quat(interpolated_quats).as_euler("xyz")
return interpolated
@staticmethod
def _interpolate_position_trajectory(
traj, original_indices, target_indices, target_length
):
"""Optimized position trajectory interpolation"""
interpolated = np.zeros((target_length, 3))
for i in range(3):
interpolated[:, i] = np.interp(target_indices, original_indices, traj[:, i])
return interpolated
@staticmethod
def _smooth_trajectory(trajectory):
"""Vectorized smooth processing"""
if len(trajectory) < 5:
return trajectory
try:
# Batch smooth all dimensions
smoothed = np.zeros_like(trajectory)
for dim in range(trajectory.shape[1]):
smoothed[:, dim] = savgol_filter(
trajectory[:, dim],
min(
5,
(
len(trajectory)
if len(trajectory) % 2 == 1
else len(trajectory) - 1
),
),
3,
mode="nearest",
)
return smoothed
except Exception:
return trajectory
@staticmethod
def calculate_optimal_trajectory_length(left_traj, right_traj):
"""Calculate optimal trajectory length"""
# Vectorized distance calculation
def calc_distance(traj):
if len(traj) < 2:
return 0.0
pos_diff = traj[1:, :3] - traj[:-1, :3]
return np.sum(np.linalg.norm(pos_diff, axis=1))
distances = [calc_distance(left_traj), calc_distance(right_traj)]
max_distance = max(distances)
if max_distance > 1e-6:
execution_time = np.clip(
max_distance / ARM_MAX_VELOCITY,
ARM_MIN_EXECUTION_TIME,
ARM_MAX_EXECUTION_TIME,
)
else:
execution_time = ARM_MIN_EXECUTION_TIME
return max(int(execution_time * ARM_EXECUTION_HZ), len(left_traj))
class VehiclePoseHandler:
"""Vehicle pose and velocity calculation"""
def __init__(self):
self.current_pose = None
self.previous_pose = None
self.pose_history = deque(maxlen=10)
def update_pose(self, new_pose):
"""Update vehicle pose"""
if new_pose is not None:
self.previous_pose = self.current_pose
self.current_pose = np.array(new_pose)
self.pose_history.append(self.current_pose.copy())
print("current_pose", self.current_pose, flush=True)
return self.current_pose
def velocity_to_pose(self, vx_body, vy_body, vyaw, dt, start_pose=None):
"""Convert body frame velocity to global frame position"""
if start_pose is None:
if self.current_pose is not None:
start_pose = self.current_pose.copy()
else:
start_pose = np.array([0.0, 0.0, 0.0])
x, y, theta = start_pose
# Convert body frame velocity to global frame displacement
cos_theta = np.cos(theta)
sin_theta = np.sin(theta)
# Coordinate transformation: body frame -> global frame
dx_global = (vx_body * cos_theta - vy_body * sin_theta) * dt
dy_global = (vx_body * sin_theta + vy_body * cos_theta) * dt
dtheta = vyaw * dt
# Calculate new position
x_new = x + dx_global
y_new = y + dy_global
theta_new = theta + dtheta
# Constrain angle to [-pi, pi] range
theta_new = (theta_new + np.pi) % (2 * np.pi) - np.pi
return np.array([x_new, y_new, theta_new])
def compute_body_velocities_from_poses(
self, current_pose, previous_pose, dt=1 / 20
):
"""Compute body frame velocity from pose changes"""
if current_pose is None or previous_pose is None:
return np.array([0.0, 0.0, 0.0])
# Calculate displacement in global frame
dx_global = current_pose[0] - previous_pose[0]
dy_global = current_pose[1] - previous_pose[1]
dtheta = current_pose[2] - previous_pose[2]
# Use previous frame's angle for coordinate transformation
theta = previous_pose[2]
cos_theta = np.cos(theta)
sin_theta = np.sin(theta)
# Convert global frame displacement to body frame velocity
vx_body = (dx_global * cos_theta + dy_global * sin_theta) / dt
vy_body = (-dx_global * sin_theta + dy_global * cos_theta) / dt
vyaw = dtheta / dt
return np.array([vx_body, vy_body, vyaw])
wall_x/infer/utils_libero.py
+323
View File
@@ -0,0 +1,323 @@
"""Utils for evaluating policies in LIBERO simulation environments."""
import math
import os
from enum import Enum
import imageio
import numpy as np
import matplotlib.pyplot as plt
import torch
from transformers import BatchFeature
import random
import time
from libero.libero import get_libero_path
from libero.libero.envs import OffScreenRenderEnv
# Define task suite constants
class TaskSuite(str, Enum):
LIBERO_SPATIAL = "libero_spatial"
LIBERO_OBJECT = "libero_object"
LIBERO_GOAL = "libero_goal"
LIBERO_10 = "libero_10"
LIBERO_90 = "libero_90"
# Define max steps for each task suite
TASK_MAX_STEPS = {
TaskSuite.LIBERO_SPATIAL: 220, # longest training demo has 193 steps
TaskSuite.LIBERO_OBJECT: 280, # longest training demo has 254 steps
TaskSuite.LIBERO_GOAL: 300, # longest training demo has 270 steps
TaskSuite.LIBERO_10: 520, # longest training demo has 505 steps
TaskSuite.LIBERO_90: 400, # longest training demo has 373 steps
}
# Initialize important constants
ACTION_DIM = 7
DATE = time.strftime("%Y_%m_%d")
DATE_TIME = time.strftime("%Y_%m_%d-%H_%M_%S")
# Configure NumPy print settings
np.set_printoptions(formatter={"float": lambda x: "{0:0.3f}".format(x)})
def set_seed_everywhere(seed: int) -> None:
"""
Set random seed for all random number generators for reproducibility.
Args:
seed: The random seed to use
"""
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)
random.seed(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
os.environ["PYTHONHASHSEED"] = str(seed)
def normalize_gripper_action(action: np.ndarray, binarize: bool = True) -> np.ndarray:
"""
Normalize gripper action from [0,1] to [-1,+1] range.
This is necessary for some environments because the dataset wrapper
standardizes gripper actions to [0,1]. Note that unlike the other action
dimensions, the gripper action is not normalized to [-1,+1] by default.
Normalization formula: y = 2 * (x - orig_low) / (orig_high - orig_low) - 1
Args:
action: Action array with gripper action in the last dimension
binarize: Whether to binarize gripper action to -1 or +1
Returns:
np.ndarray: Action array with normalized gripper action
"""
# Create a copy to avoid modifying the original
normalized_action = action.copy()
# Normalize the last action dimension to [-1,+1]
orig_low, orig_high = 0.0, 1.0
normalized_action[..., -1] = (
2 * (normalized_action[..., -1] - orig_low) / (orig_high - orig_low) - 1
)
if binarize:
# Binarize to -1 or +1
normalized_action[..., -1] = np.sign(normalized_action[..., -1])
return normalized_action
def invert_gripper_action(action: np.ndarray) -> np.ndarray:
"""
Flip the sign of the gripper action (last dimension of action vector).
This is necessary for environments where -1 = open, +1 = close, since
the RLDS dataloader aligns gripper actions such that 0 = close, 1 = open.
Args:
action: Action array with gripper action in the last dimension
Returns:
np.ndarray: Action array with inverted gripper action
"""
# Create a copy to avoid modifying the original
inverted_action = action.copy()
# Invert the gripper action
inverted_action[..., -1] *= -1.0
return inverted_action
def move_to_cuda(obj, device="cuda"):
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 get_libero_env(task, model_family, resolution=256, seed=7):
"""Initializes and returns the LIBERO environment, along with the task description."""
task_description = task.language
task_bddl_file = os.path.join(
get_libero_path("bddl_files"), task.problem_folder, task.bddl_file
)
env_args = {
"bddl_file_name": task_bddl_file,
"camera_heights": resolution,
"camera_widths": resolution,
}
env = OffScreenRenderEnv(**env_args)
env.seed(
seed
) # IMPORTANT: seed seems to affect object positions even when using fixed initial state
return env, task_description
def get_libero_dummy_action(model_family: str):
"""Get dummy/no-op action, used to roll out the simulation while the robot does nothing."""
return [0, 0, 0, 0, 0, 0, -1]
def get_libero_image(obs):
"""Extracts third-person image from observations and preprocesses it."""
img = obs["agentview_image"]
img = img[::-1, ::-1] # IMPORTANT: rotate 180 degrees to match train preprocessing
return img
def get_libero_wrist_image(obs):
"""Extracts wrist camera image from observations and preprocesses it."""
img = obs["robot0_eye_in_hand_image"]
img = img[::-1, ::-1] # IMPORTANT: rotate 180 degrees to match train preprocessing
return img
def save_rollout_video(
rollout_dir,
rollout_images,
idx,
success,
task_description,
log_file=None,
model_family="openvla_oft",
):
"""Saves an MP4 replay of an episode."""
processed_task_description = (
task_description.lower()
.replace(" ", "_")
.replace("\n", "_")
.replace(".", "_")[:50]
)
mp4_path = f"{rollout_dir}/episode={idx}--success={success}--task={processed_task_description}.mp4"
video_writer = imageio.get_writer(mp4_path, fps=30)
for img in rollout_images:
video_writer.append_data(img)
video_writer.close()
print(f"Saved rollout MP4 at path {mp4_path}")
if log_file is not None:
log_file.write(f"Saved rollout MP4 at path {mp4_path}\n")
return mp4_path
def save_rollout_data(
rollout_dir,
rollout_data,
idx,
success,
task_description,
log_file=None,
model_family="openvla_oft",
):
"""
Saves an NPY file of the rollout data.
"""
# Process task description to make it suitable for filename
processed_task_description = (
task_description.lower()
.replace(" ", "_")
.replace("\n", "_")
.replace(".", "_")[:50]
)
# Build .npy file path
npy_path = f"{rollout_dir}/episode={idx}--success={success}--task={processed_task_description}--action.npy"
# Save rollout_data as .npy file
np.save(npy_path, rollout_data)
print(f"Saved rollout data at path {npy_path}")
fig, axes = plt.subplots(nrows=1, ncols=7, figsize=(20, 3))
titles = ["x", "y", "z", "roll", "pitch", "yaw", "grasp"]
for i in range(rollout_data.shape[1]):
ax = axes[i] # Select the i-th subplot
ax.plot(rollout_data[:, i], label=f"Feature {i+1}") # Plot line chart
ax.set_title(titles[i]) # Set subplot title
ax.set_xlabel("Time in one episode") # Set x-axis label
axes[-1].legend(["predicted action"], loc="upper right")
plt.tight_layout()
png_path = f"{rollout_dir}/episode={idx}--success={success}--task={processed_task_description}--action.png"
plt.savefig(png_path, dpi=300) # Save image as PNG file
# If log file is provided, record the save path
if log_file is not None:
log_file.write(f"Saved rollout data at path {npy_path}\n")
return npy_path
def save_rollout_observation(
rollout_dir,
rollout_data,
idx,
success,
task_description,
log_file=None,
model_family="openvla_oft",
):
"""
Saves an NPY file of the rollout data.
"""
# Process task description to make it suitable for filename
processed_task_description = (
task_description.lower()
.replace(" ", "_")
.replace("\n", "_")
.replace(".", "_")[:50]
)
# Build .npy file path
npy_path = f"{rollout_dir}/episode={idx}--success={success}--task={processed_task_description}--observation.npy"
# Save rollout_data as .npy file
np.save(npy_path, rollout_data)
print(f"Saved rollout data at path {npy_path}")
if rollout_data.shape[1] == 7:
fig, axes = plt.subplots(nrows=1, ncols=7, figsize=(20, 3))
titles = ["x", "y", "z", "roll", "pitch", "yaw", "grasp"]
else:
fig, axes = plt.subplots(nrows=1, ncols=8, figsize=(20, 3))
titles = ["x", "y", "z", "roll", "pitch", "yaw", "-", "grasp"]
for i in range(rollout_data.shape[1]):
ax = axes[i] # Select the i-th subplot
ax.plot(rollout_data[:, i], label=f"Feature {i+1}") # Plot line chart
ax.set_title(titles[i]) # Set subplot title
ax.set_xlabel("Time in one episode") # Set x-axis label
axes[-1].legend(["predicted action"], loc="upper right")
plt.tight_layout()
png_path = f"{rollout_dir}/episode={idx}--success={success}--task={processed_task_description}--observation.png"
plt.savefig(png_path, dpi=300) # Save image as PNG file
# If log file is provided, record the save path
if log_file is not None:
log_file.write(f"Saved rollout data at path {npy_path}\n")
return npy_path
def quat2axisangle(quat):
"""
Copied from robosuite: https://github.com/ARISE-Initiative/robosuite/blob/eafb81f54ffc104f905ee48a16bb15f059176ad3/robosuite/utils/transform_utils.py#L490C1-L512C55
Converts quaternion to axis-angle format.
Returns a unit vector direction scaled by its angle in radians.
Args:
quat (np.array): (x,y,z,w) vec4 float angles
Returns:
np.array: (ax,ay,az) axis-angle exponential coordinates
"""
# clip quaternion
if quat[3] > 1.0:
quat[3] = 1.0
elif quat[3] < -1.0:
quat[3] = -1.0
den = np.sqrt(1.0 - quat[3] * quat[3])
if math.isclose(den, 0.0):
# This is (close to) a zero degree rotation, immediately return
return np.zeros(3)
return (quat[:3] * 2.0 * math.acos(quat[3])) / den