diff --git a/scripts/compute_norm_stats.py b/scripts/compute_norm_stats.py index 5bee282..b191d99 100644 --- a/scripts/compute_norm_stats.py +++ b/scripts/compute_norm_stats.py @@ -1,79 +1,183 @@ -import yaml -import torch -import tqdm -from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata -from wall_x.data.load_lerobot_dataset import KEY_MAPPINGS -import normalize +#!/usr/bin/env python3 + +import json +import logging +from collections import defaultdict +from pathlib import Path +from typing import Dict, List +from tqdm import tqdm + import numpy as np -import argparse + +from lerobot.datasets.lerobot_dataset import LeRobotDataset -def load_config(config_path): - """Load configuration from YAML file.""" - with open(config_path, "r") as f: - config = yaml.load(f, Loader=yaml.FullLoader) - - config["data"]["model_type"] = config.get("model_type") - - return config +def write_json(path: Path, data: Dict) -> None: + path.write_text( + json.dumps(data, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) -def load_lerobot_dataset(repo_id, root, action_horizon, args): - dataset_meta = LeRobotDatasetMetadata(repo_id) - dataset = LeRobotDataset( - repo_id, - root=root, - delta_timestamps={ - key: [t / dataset_meta.fps for t in range(action_horizon)] - for key in [KEY_MAPPINGS[repo_id]["action"]] +def compute_action_statistics( + action_data_by_robot: Dict[str, Dict[str, List]] +) -> Dict[str, Dict[str, Dict]]: + """ + Compute statistics (min, q01, q99, max) for each action type and dimension. + + Args: + action_data_by_robot: Dict[robot_id][action_type] -> list of arrays/lists + + Returns: + Dict[robot_id][action_type] -> { + "min": [min for each dim], + "q01": [quantile 1% for each dim], + "q99": [quantile 99% for each dim], + "max": [max for each dim], + "delta": [max - min for each dim] + "delta_q99_q01": [q99 - q01 for each dim] + } + """ + stats = {} + + for robot_id, action_data in action_data_by_robot.items(): + stats[robot_id] = {} + + for action_type, values_list in action_data.items(): + if not values_list: + continue + + # Convert to numpy array: shape (num_samples, num_dims) + try: + values_array = np.array(values_list) + if values_array.size == 0: + continue + + # Handle both 1D and 2D cases + if values_array.ndim == 1: + values_array = values_array.reshape(-1, 1) + elif values_array.ndim == 2: + pass + else: + logging.warning( + f"Unexpected shape for {robot_id}/{action_type}: {values_array.shape}" + ) + continue + + # Compute statistics for each dimension + min_vals = np.min(values_array, axis=0).tolist() + max_vals = np.max(values_array, axis=0).tolist() + q01_vals = np.quantile(values_array, 0.01, axis=0).tolist() + q99_vals = np.quantile(values_array, 0.99, axis=0).tolist() + delta_vals = (np.array(max_vals) - np.array(min_vals)).tolist() + delta_q99_q01_vals = (np.array(q99_vals) - np.array(q01_vals)).tolist() + + stats[robot_id][action_type] = { + "min": min_vals, + "q01": q01_vals, + "q99": q99_vals, + "max": max_vals, + "delta": delta_vals, + "delta_q99_q01": delta_q99_q01_vals, + } + + except Exception as e: + logging.warning( + f"Error computing statistics for {robot_id}/{action_type}: {e}" + ) + continue + + return stats + + +def load_lerobot_dataset( + repo_id: str, + trajectory_keys: Dict, + base_dir: Path, +) -> None: + + # Load local or remote dataset + dataset = LeRobotDataset(base_dir) + + # Iterate through all data + frames: Dict[str, Dict[str, List]] = defaultdict(lambda: defaultdict(list)) + + all_features = dataset.features + non_image_columns = [col for col in all_features if "image" not in col] + + print(f"Reading the following fields:{non_image_columns}") + fast_dataset = dataset.hf_dataset.select_columns(non_image_columns) + + for i in tqdm(range(len(fast_dataset))): + sample = fast_dataset[i] + action = sample["action"] # torch.Tensor + propri = sample["observation.state"] + + for key, action_keys in trajectory_keys.items(): + for action_key, action_range in action_keys.items(): + if key == "action": + frames[repo_id][action_key].append( + action[action_range[0] : action_range[1]].numpy().tolist() + ) + else: + frames[repo_id][action_key].append( + propri[action_range[0] : action_range[1]].numpy().tolist() + ) + + return frames + + +def compute_action_normalizer( + repo_id: str, trajectory_keys: Dict, base_dir: Path, output_dir: Path +) -> None: + """ + Compute action normalizer statistics for all robot_ids. + """ + logging.info("Starting action normalizer computation...") + + frames = load_lerobot_dataset(repo_id, trajectory_keys, base_dir) + + # Compute statistics + stats = compute_action_statistics(frames) + + # Save statistics for each robot_id + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + # for robot_id, robot_stats in stats.items(): + # output_file = output_dir / f"{robot_id}_action_stats.json" + # write_json(output_file, robot_stats) + # logging.info(f"Saved action statistics for {robot_id} to {output_file}") + + # Also save a combined file + combined_output = output_dir / "all_robots_action_stats.json" + write_json(combined_output, stats) + logging.info(f"Saved combined action statistics to {combined_output}") + + +def main() -> None: + + repo_id = "xxx" # your dataset name + data_root_path = "/path/to/lerobot/dataset" + output_stats_dir = "/path/to/save/action_stats" + trajectory_keys = { # your dataset keys + "action": { + "follow_right_ee_cartesian_pos": [0, 3], + "follow_right_ee_rotation": [3, 6], + "follow_right_gripper": [6, 7], }, - video_backend="pyav", + "propri": { + "master_right_ee_cartesian_pos": [0, 3], + "master_right_ee_rotation": [3, 6], + "master_right_gripper": [6, 7], + }, + } + + compute_action_normalizer( + repo_id, trajectory_keys, data_root_path, output_stats_dir ) - num_batches = len(dataset) // args.batch_size - generator = torch.Generator() - generator.manual_seed(args.seed) - data_loader = torch.utils.data.DataLoader( - dataset, - batch_size=args.batch_size, - shuffle=False, - drop_last=True, - generator=generator, - num_workers=args.num_workers, - persistent_workers=True if args.num_workers > 0 else False, - ) - return data_loader, num_batches + logging.info("Action normalizer computation completed.") if __name__ == "__main__": - # set args - parser = argparse.ArgumentParser() - parser.add_argument("--batch_size", type=int, default=256) - parser.add_argument("--num_workers", type=int, default=2) - parser.add_argument("--seed", type=int, default=0) - args = parser.parse_args() - - # Configs - path = "/path/to/config.yml" - output_path = "/path/to/output" - config = load_config(path) - lerobot_config = config["data"]["lerobot_config"] - repo_id = lerobot_config.get("repo_id", None) - root = lerobot_config.get("root", None) - assert repo_id is not None, "repo id is required" - action_horizon = config["data"].get("action_horizon", 32) - - data_loader, num_batches = load_lerobot_dataset(repo_id, root, action_horizon, args) - - keys = ["state", "action"] - stats = {key: normalize.RunningStats() for key in keys} - for batch in tqdm.tqdm(data_loader, total=num_batches, desc="Computing stats"): - for key in keys: - stats[key].update(np.asarray(batch[KEY_MAPPINGS[repo_id][key]])) - norm_stats = { - KEY_MAPPINGS[repo_id][key]: stats.get_statistics() - for key, stats in stats.items() - } - - output_path = output_path + "/" + repo_id - print(f"Writing stats to: {output_path}") - normalize.save(output_path, norm_stats) + main() diff --git a/scripts/draw_openloop_plot.py b/scripts/draw_openloop_plot.py index 950bf8d..07b3d4f 100644 --- a/scripts/draw_openloop_plot.py +++ b/scripts/draw_openloop_plot.py @@ -6,6 +6,8 @@ from tqdm import tqdm import matplotlib.pyplot as plt from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction from wall_x.data.load_lerobot_dataset import load_test_dataset, get_data_configs +from wall_x.model.model_utils import register_normalizers +import copy def load_config(config_path): @@ -21,37 +23,46 @@ def load_config(config_path): if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--pred_horizon", type=int, default=32) - parser.add_argument("--origin_action_dim", type=int, default=7) + parser.add_argument("--origin_action_dim", type=int, default=14) args = parser.parse_args() origin_action_dim = args.origin_action_dim pred_horizon = args.pred_horizon # get train config - model_path = "/path/to/model" - action_tokenizer_path = "/path/to/action/tokenizer" + model_path = "/path/to/your/checkpoint" + action_tokenizer_path = "/path/to/Models/fast" save_dir = "/path/to/save/dir" - path = "/path/to/train/config" + path = f"{model_path}/config.yml" config = load_config(path) + normalizer_action, normalizer_propri = register_normalizers(config, model_path) + # load model with customized robot config model = Qwen2_5_VLMoEForAction.from_pretrained( model_path, train_config=config, action_tokenizer_path=action_tokenizer_path ) + + model.set_normalizer( + copy.deepcopy(normalizer_action), copy.deepcopy(normalizer_propri) + ) model.eval() model = model.to("cuda") - model = model.bfloat16() + model.to_bfloat16_for_selected_params() # get test dataloader dataload_config = get_data_configs(config["data"]) lerobot_config = dataload_config.get("lerobot_config", {}) - dataset = load_test_dataset(config, lerobot_config, seed=42) + dataset = load_test_dataset( + config, lerobot_config, normalizer_action, normalizer_propri, seed=42 + ) dataloader = dataset.get_dataloader() + # dataloader = dataset.get_train_dataloader() total_frames = len(dataloader) predict_mode = "fast" if config.get("use_fast_tokenizer", False) else "diffusion" - action_dim = 20 if predict_mode == "diffusion" else origin_action_dim + action_dim = 14 if predict_mode == "diffusion" else origin_action_dim gt_traj = torch.zeros((total_frames, origin_action_dim)) pred_traj = torch.zeros((total_frames, origin_action_dim)) @@ -65,7 +76,7 @@ if __name__ == "__main__": outputs = model( **batch, action_dim=action_dim, - pred_horizon=pred_horizon, + action_horizon=pred_horizon, mode="predict", predict_mode=predict_mode, ) diff --git a/scripts/infer_libero.py b/scripts/infer_libero.py new file mode 100644 index 0000000..4b10725 --- /dev/null +++ b/scripts/infer_libero.py @@ -0,0 +1,222 @@ +import argparse +import time +import os + +# from wall_x.utils.baseline_utils import check_baseline_dump, update_baseline +from wall_x.infer.utils_libero import set_seed_everywhere, TaskSuite, TASK_MAX_STEPS +from wall_x.infer.infer_config import InferConfig +from wall_x.infer.env_libero import LiberoRobotEnv + + +if __name__ == "__main__": + args = argparse.ArgumentParser(description="Wall-X Libero evaluation script") + args.add_argument("--seed", type=int, default=42, help="Random seed") + args.add_argument("--id", type=int, default=None, help="Unique index id") + args.add_argument("--name", type=int, default=None, help="Launch command name") + args.add_argument( + "--baseline_path", type=str, default=None, help="Path to baseline record table" + ) + args.add_argument( + "--update_baseline", + type=bool, + default=False, + help="Whether to update baseline table", + ) + args.add_argument( + "--mode", type=str, default="flow", choices=["flow", "ar"], help="Running mode" + ) + args.add_argument( + "--checkpoint_path", type=str, required=True, help="Model checkpoint path" + ) + args.add_argument( + "--train_config_path", + type=str, + required=False, + default=None, + help="Path to training config .yml file", + ) + args.add_argument( + "--norm_key", + type=str, + default="physical-intelligence/libero", + help="Key for normalization statistics", + ) + args.add_argument( + "--cam_names", + nargs="+", + default=["face_view", "right_wrist_view"], + help="List of camera names (e.g., --cam_names face_view right_wrist_view)", + ) + args.add_argument( + "--task_suite_name", + type=str, + default=TaskSuite.LIBERO_SPATIAL, + choices=[e.value for e in TaskSuite], + help="Libero task suite to load", + ) + args.add_argument( + "--initial_states_path", + type=str, + default="DEFAULT", + help="Path to initial states .json file, or 'DEFAULT' to use default states.", + ) + args.add_argument( + "--num_trials_per_task", + type=int, + default=50, + help="Number of evaluation episodes to run per task", + ) + args.add_argument( + "--rollout_dir", + type=str, + default="./rollouts", + help="Directory to save rollout videos", + ) + args = args.parse_args() + + print(f"Using random seed: {args.seed}") + set_seed_everywhere(args.seed) + + print("Initializing InferConfig...") + if args.train_config_path is None: + args.train_config_path = os.path.join(args.checkpoint_path, "config.yml") + + config = InferConfig( + checkpoint_path=args.checkpoint_path, + train_config_path=args.train_config_path, + norm_key=args.norm_key, + cam_names=args.cam_names, + ) + if args.mode == "flow": + config.action_horizon = config.train_config.get("data", {}).get( + "action_horizon_flow", 10 + ) + elif args.mode == "ar": + config.action_horizon = config.train_config.get("data", {}).get( + "action_horizon_ar", 10 + ) + else: + raise ValueError(f"Invalid mode: {args.mode}") + config.model_device = "cuda" + + print("Initializing LiberoRobotEnv (Evaluator)...") + + config.action_dim = 7 + config.pred_horizon = 10 + + evaluator = LiberoRobotEnv( + config=config, + task_suite_name=args.task_suite_name, + initial_states_path=args.initial_states_path, + rollout_dir=args.rollout_dir, + seed=args.seed, + ) + + print(f"\n{'='*20} Starting Evaluation {'='*20}") + print(f"Task suite: {args.task_suite_name}") + print(f"Number of tasks: {evaluator.num_tasks}") + print(f"Trials per task: {args.num_trials_per_task}") + print(f"Initial states: {args.initial_states_path}") + print(f"Videos will be saved to: {evaluator.rollout_dir}") + print(f"{'='*50}\n") + + total_successes = 0 + total_episodes_run = 0 + start_time = time.time() + + for task_id in range(evaluator.num_tasks): + task_successes = 0 + task_episodes_attempted = 0 + + libero_env_instance = None + task_desc = "" + initial_states = None + + max_infer_times = TASK_MAX_STEPS[args.task_suite_name] + print( + f"{args.task_suite_name} TASK_MAX_STEPS: {TASK_MAX_STEPS[args.task_suite_name]}" + ) + for ep_idx in range(args.num_trials_per_task): + print(f" > Running trial {ep_idx + 1} / {args.num_trials_per_task}...") + + try: + print(f"\nCreating environment for Task {task_id}...") + libero_env_instance, task_desc, initial_states = ( + evaluator.create_env_for_task(task_id) + ) + print( + f"--- Starting task {task_id + 1} / {evaluator.num_tasks}: {task_desc} ---" + ) + except Exception as e: + print( + f"\n[CRITICAL ERROR] Failed to create environment for task {task_id}: {e}. Skipping entire task." + ) + continue + + task_episodes_attempted += 1 + total_episodes_run += 1 + + success = False + try: + if args.mode == "flow": + success = evaluator.run_infer_flow_action( + env=libero_env_instance, + task_id=task_id, + task_desc=task_desc, + default_initial_states=initial_states, + episode_idx=ep_idx, + max_infer_times=max_infer_times, + ) + elif args.mode == "ar": + success = evaluator.run_infer_ar_action( + env=libero_env_instance, + task_id=task_id, + task_desc=task_desc, + default_initial_states=initial_states, + episode_idx=ep_idx, + max_infer_times=max_infer_times, + ) + except Exception as e: + print(f" [EXCEPTION] Episode run error: {e}") + + if success: + task_successes += 1 + total_successes += 1 + print(" > Trial result: SUCCESS") + else: + print(" > Trial result: FAILURE") + + if task_episodes_attempted > 0: + print( + f" > Task {task_id} current success rate: {task_successes / task_episodes_attempted * 100:.1f}% ({task_successes}/{task_episodes_attempted})" + ) + if total_episodes_run > 0: + print( + f" > Overall current success rate: {total_successes / total_episodes_run * 100:.1f}% ({total_successes}/{total_episodes_run})" + ) + + task_success_rate = ( + task_successes / task_episodes_attempted + if task_episodes_attempted > 0 + else 0 + ) + print(f"\n--- Task {task_id} ({task_desc}) Summary ---") + print( + f"Success rate: {task_success_rate * 100:.1f}% ({task_successes}/{task_episodes_attempted})" + ) + print(f"{'-'*40}\n") + + end_time = time.time() + total_time = end_time - start_time + final_success_rate = ( + total_successes / total_episodes_run if total_episodes_run > 0 else 0 + ) + + print(f"\n{'='*20} Final Evaluation Summary {'='*20}") + print(f"Total runtime: {total_time:.2f} seconds ({total_time / 60:.1f} minutes)") + print(f"Total trials run: {total_episodes_run}") + print(f"Total successes: {total_successes}") + print(f"Overall success rate: {final_success_rate * 100:.2f}%") + print(f"{'='*56}") + + print("Evaluation completed.") diff --git a/scripts/infer_robochallenge.py b/scripts/infer_robochallenge.py new file mode 100644 index 0000000..f25b3d3 --- /dev/null +++ b/scripts/infer_robochallenge.py @@ -0,0 +1,1164 @@ +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() diff --git a/scripts/vqa_inference.py b/scripts/vqa_inference.py index 77748e5..adbef89 100644 --- a/scripts/vqa_inference.py +++ b/scripts/vqa_inference.py @@ -2,14 +2,22 @@ import torch from PIL import Image from transformers import AutoProcessor import yaml +import os from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction class VQAWrapper(object): - def __init__(self, model_path: str, train_config: dict): + def __init__(self, model_path: str, train_config: dict = None): + self.device = self._setup_device() - self.processor = self._load_processor(model_path) + if train_config is None: + try: + with open(os.path.join(model_path, "config.yml"), "r") as f: + train_config = yaml.load(f, Loader=yaml.FullLoader) + except Exception as e: + print(f"load train_config.yml fail: {e}") + self.processor = self._load_processor(train_config["processor_path"]) self.model = self._load_model(model_path, train_config) def _setup_device(self) -> str: @@ -69,8 +77,8 @@ class VQAWrapper(object): if __name__ == "__main__": - MODEL_PATH_FOR_MODULE_TEST = "/path/to/model" - train_config_path = "/path/to/config.yaml" + MODEL_PATH_FOR_MODULE_TEST = "/path/to/model_path" + train_config_path = "/path/to/model_path/config.yml" with open(train_config_path, "r") as f: train_config = yaml.load(f, Loader=yaml.FullLoader) wrapper = VQAWrapper( @@ -81,7 +89,9 @@ if __name__ == "__main__": test_question = "To move the red block in the plate with same color, what should you do next? Think step by step." # Local Image - img = Image.open("/path/to/wall-x/assets/cot_example_frame.png").convert("RGB") + img = Image.open( + "/x2robot_v2/yangping/github/wall-x/assets/cot_example_frame.png" + ).convert("RGB") # Internet Image # import requests # test_image_url = "https://www.ilankelman.org/stopsigns/australia.jpg" diff --git a/wall_x/data/load_lerobot_dataset.py b/wall_x/data/load_lerobot_dataset.py index cd510af..6bef686 100644 --- a/wall_x/data/load_lerobot_dataset.py +++ b/wall_x/data/load_lerobot_dataset.py @@ -17,7 +17,7 @@ from wall_x.data.utils import ( ) from transformers import AutoProcessor -from .utils import load_norm_stats, KEY_MAPPINGS +from .utils import KEY_MAPPINGS T_co = TypeVar("T_co", covariant=True) @@ -39,7 +39,8 @@ class PreprocessedDataset(Dataset[T_co]): dataset, config, dataload_config, - norm_stats, + normalizer_action, + normalizer_propri, lerobot_config, seed=42, rank=0, @@ -67,7 +68,9 @@ class PreprocessedDataset(Dataset[T_co]): self.config = config self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False) self.dataload_config = dataload_config - self.norm_stats = norm_stats + self.normalizer_action = (normalizer_action,) + self.normalizer_propri = normalizer_propri + # self.norm_stats = norm_stats self.lerobot_config = lerobot_config self.data_config = X2RDataProcessingConfig().update( @@ -128,6 +131,7 @@ class PreprocessedDataset(Dataset[T_co]): frame_index = data["frame_index"] instruction_info = {"instruction": data["task"]} generate_subtask_ratio = self.data_config.generate_subtask_ratio + complete_text, generate_subtask = get_wallx_normal_text( instruction_info, self.dataload_config.get("action_horizon", 33) - 1, @@ -188,7 +192,11 @@ class PreprocessedDataset(Dataset[T_co]): sampler=sampler, # Use distributed sampler instead of shuffle=True num_workers=num_workers, collate_fn=DataCollator( - self.config, self.dataload_config, self.norm_stats, self.lerobot_config + self.config, + self.dataload_config, + self.normalizer_action, + self.normalizer_propri, + self.lerobot_config, ), pin_memory=True, # Enable for GPU training persistent_workers=num_workers > 0, # Only if num_workers > 0 @@ -240,22 +248,29 @@ class DataCollator: _processor_cache = {} _action_tokenizer_cache = {} - def __init__(self, config, dataload_config, stats, lerobot_config): + def __init__( + self, + config, + dataload_config, + normalizer_action, + normalizer_propri, + lerobot_config, + ): self.config = config self.dataload_config = dataload_config - self.stats = stats - self.action_min_stat = stats["action"].min - self.action_delta = stats["action"].delta - self.state_min_stat = stats["state"].min - self.state_delta = stats["state"].delta + + self.normalizer_action = normalizer_action[0] + self.normalizer_propri = normalizer_propri self.lerobot_config = lerobot_config self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False) + self.dataset_name = self.config["data"]["lerobot_config"].get("repo_id", "") + self.dataset_name = [self.dataset_name] * self.config["batch_size_per_gpu"] self.load_processor() def load_processor(self): processor_path = self.config["pretrained_wallx_path"] - action_tokenizer_path = self.config["action_tokenizer_path"] + action_tokenizer_path = self.config.get("action_tokenizer_path", None) if ( self.use_fast_tokenizer @@ -319,33 +334,35 @@ class DataCollator: if agent_pos.dim() == 2: agent_pos = agent_pos.unsqueeze(1) agent_pos_mask = (~torch.isnan(agent_pos)).float() + # print("agent_pos_mask",agent_pos_mask.shape) agent_pos.nan_to_num_(nan=0.0) - agent_pos = self._normalize( - agent_pos, self.state_min_stat, self.state_delta + + # if agent_pos.shape[-1] != 20: + # agent_pos = torch.cat( + # [ + # agent_pos, + # torch.zeros( + # agent_pos.shape[0], + # agent_pos.shape[1], + # 20 - agent_pos.shape[-1], + # ), + # ], + # dim=-1, + # ) + # agent_pos_mask = torch.cat( + # [ + # agent_pos_mask, + # torch.zeros( + # agent_pos_mask.shape[0], + # agent_pos_mask.shape[1], + # 20 - agent_pos_mask.shape[-1], + # ), + # ], + # dim=-1, + # ) + agent_pos = self.normalizer_propri.normalize_data( + agent_pos, self.dataset_name ) - if agent_pos.shape[-1] != 20: - agent_pos = torch.cat( - [ - agent_pos, - torch.zeros( - agent_pos.shape[0], - agent_pos.shape[1], - 20 - agent_pos.shape[-1], - ), - ], - dim=-1, - ) - agent_pos_mask = torch.cat( - [ - agent_pos_mask, - torch.zeros( - agent_pos_mask.shape[0], - agent_pos_mask.shape[1], - 20 - agent_pos_mask.shape[-1], - ), - ], - dim=-1, - ) additional_inputs["proprioception"] = agent_pos additional_inputs["agent_pos_mask"] = agent_pos_mask elif key == "action": @@ -354,30 +371,31 @@ class DataCollator: action = action.unsqueeze(1) dof_mask = (~torch.isnan(action)).float() action.nan_to_num_(nan=0.0) - action = self._normalize( - action, self.action_min_stat, self.action_delta + + # if action.shape[-1] != 20: + # action = torch.cat( + # [ + # action, + # torch.zeros( + # action.shape[0], action.shape[1], 20 - action.shape[-1] + # ), + # ], + # dim=-1, + # ) + # dof_mask = torch.cat( + # [ + # dof_mask, + # torch.zeros( + # dof_mask.shape[0], + # dof_mask.shape[1], + # 20 - dof_mask.shape[-1], + # ), + # ], + # dim=-1, + # ) + action = self.normalizer_action.normalize_data( + action, self.dataset_name ) - if action.shape[-1] != 20: - action = torch.cat( - [ - action, - torch.zeros( - action.shape[0], action.shape[1], 20 - action.shape[-1] - ), - ], - dim=-1, - ) - dof_mask = torch.cat( - [ - dof_mask, - torch.zeros( - dof_mask.shape[0], - dof_mask.shape[1], - 20 - dof_mask.shape[-1], - ), - ], - dim=-1, - ) additional_inputs["action_chunk"] = action additional_inputs["dof_mask"] = dof_mask elif key == "image_inputs": @@ -431,6 +449,8 @@ class DataCollator: def load_lerobot_data( config, lerobot_config, + normalizer_action, + normalizer_propri, rank=0, world_size=1, seed=42, @@ -462,11 +482,11 @@ def load_lerobot_data( dataset_fps = meta_info.fps episodes_num = meta_info.total_episodes - norm_stats_path = config.get("norm_stats_path", None) - assert ( - norm_stats_path is not None - ), "norm stats is required, please refer to 'wall-x/scripts/compute_norm_stats.py' to compute stats" - norm_stats = load_norm_stats(norm_stats_path, repo_id) + # norm_stats_path = config.get("norm_stats_path", None) + # assert ( + # norm_stats_path is not None + # ), "norm stats is required, please refer to 'wall-x/scripts/compute_norm_stats.py' to compute stats" + # norm_stats = load_norm_stats(norm_stats_path, repo_id) delta_timestamps = { # action chunk @@ -500,7 +520,8 @@ def load_lerobot_data( train_dataset, config, dataload_config, - norm_stats, + normalizer_action, + normalizer_propri, lerobot_config, seed=seed, rank=rank, @@ -584,13 +605,21 @@ def get_data_configs(config): class TestDataset(PreprocessedDataset): def __init__( - self, dataset, config, dataload_config, norm_stats, lerobot_config, seed=42 + self, + dataset, + config, + dataload_config, + normalizer_action, + normalizer_propri, + lerobot_config, + seed=42, ): super().__init__( dataset, config, dataload_config, - norm_stats, + normalizer_action, + normalizer_propri, lerobot_config, seed=seed, rank=0, @@ -607,7 +636,11 @@ class TestDataset(PreprocessedDataset): self, batch_size=1, collate_fn=DataCollator( - self.config, self.dataload_config, self.norm_stats, self.lerobot_config + self.config, + self.dataload_config, + self.normalizer_action, + self.normalizer_propri, + self.lerobot_config, ), ) @@ -617,6 +650,8 @@ class TestDataset(PreprocessedDataset): def load_test_dataset( config, lerobot_config, + normalizer_action, + normalizer_propri, seed=42, episode=0, ): @@ -645,7 +680,7 @@ def load_test_dataset( assert ( norm_stats_path is not None ), "norm stats is required, please refer to 'wall-x/scripts/compute_norm_stats.py' to compute stats" - norm_stats = load_norm_stats(norm_stats_path, repo_id) + # norm_stats = load_norm_stats(norm_stats_path, repo_id) delta_timestamps = { # action chunk @@ -668,7 +703,13 @@ def load_test_dataset( print(f"Number of frames selected: {dataset.num_frames}") dataset = TestDataset( - dataset, config, dataload_config, norm_stats, lerobot_config, seed=seed + dataset, + config, + dataload_config, + normalizer_action, + normalizer_propri, + lerobot_config, + seed=seed, ) return dataset diff --git a/wall_x/data/utils.py b/wall_x/data/utils.py index 2ffec80..9cb0efc 100644 --- a/wall_x/data/utils.py +++ b/wall_x/data/utils.py @@ -34,7 +34,7 @@ KEY_MAPPINGS = { "state": "state", "action": "actions", }, - "x2": { + "x2_normal": { "camera": { "observation.images.faceImg": "face_view", "observation.images.leftImg": "left_wrist_view", @@ -43,6 +43,23 @@ KEY_MAPPINGS = { "state": "observation.state", "action": "action", }, + "libero": { + "camera": { + "observation.images.faceImg": "face_view", + "observation.images.rightImg": "right_wrist_view", + }, + "state": "observation.state", + "action": "action", + }, + "robochallenge_aloha": { + "camera": { + "observation.images.cam_high_rgb": "face_view", + "observation.images.cam_wrist_left_rgb": "left_wrist_view", + "observation.images.cam_wrist_right_rgb": "right_wrist_view", + }, + "state": "observation.state", + "action": "action", + }, } CAMERA_NAME_MAPPING = { diff --git a/wall_x/infer/base_dataclass.py b/wall_x/infer/base_dataclass.py new file mode 100644 index 0000000..b550c3b --- /dev/null +++ b/wall_x/infer/base_dataclass.py @@ -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 diff --git a/wall_x/infer/env.py b/wall_x/infer/env.py new file mode 100644 index 0000000..35106f7 --- /dev/null +++ b/wall_x/infer/env.py @@ -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) diff --git a/wall_x/infer/env_libero.py b/wall_x/infer/env_libero.py new file mode 100644 index 0000000..07fac47 --- /dev/null +++ b/wall_x/infer/env_libero.py @@ -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 diff --git a/wall_x/infer/infer_config.py b/wall_x/infer/infer_config.py new file mode 100644 index 0000000..bf7dc82 --- /dev/null +++ b/wall_x/infer/infer_config.py @@ -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) diff --git a/wall_x/infer/logger.py b/wall_x/infer/logger.py new file mode 100644 index 0000000..ec8bf3f --- /dev/null +++ b/wall_x/infer/logger.py @@ -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) diff --git a/wall_x/infer/utils.py b/wall_x/infer/utils.py new file mode 100644 index 0000000..e4fba55 --- /dev/null +++ b/wall_x/infer/utils.py @@ -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]) diff --git a/wall_x/infer/utils_libero.py b/wall_x/infer/utils_libero.py new file mode 100644 index 0000000..b6fd5cb --- /dev/null +++ b/wall_x/infer/utils_libero.py @@ -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 diff --git a/wall_x/model/action_head.py b/wall_x/model/action_head.py index cd5bb89..2f3ab73 100644 --- a/wall_x/model/action_head.py +++ b/wall_x/model/action_head.py @@ -1,143 +1,131 @@ -import math import torch + import torch.nn as nn + +from typing import Union +import math + +from diffusers.schedulers.scheduling_ddpm import DDPMScheduler from torch.distributions import Beta -from wall_x.utils.constant import action_statistic_dof -import logging + + +def print_rank_last(message): + """If distributed is initialized, print only on last rank.""" + if torch.distributed.is_initialized(): + if torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1): + print(message, flush=True) + else: + print(message, flush=True) class Normalizer(nn.Module): - """ - Action data normalizer for multi-robot systems. + @classmethod + def from_ckpt(cls, ckpt_path): + instance = cls.__new__(cls) + nn.Module.__init__(instance) - This module handles normalization and denormalization of action data for different robot - configurations. It maintains per-robot statistics (min values and deltas) and applies - normalization to map actions to the [-1, 1] range. - """ + instance.min = nn.ParameterDict() + instance.delta = nn.ParameterDict() + instance.min_key = "min" + instance.delta_key = "delta" - def _pad_to_action_dim(self, xs, action_dim): - """ - Pad the action data to the action dimension. - """ - if xs.shape[-1] < action_dim: - padding_shape = list(xs.shape) - padding_shape[-1] = action_dim - padding_shape[-1] - xs = torch.cat([xs, torch.zeros(padding_shape).to(xs.device)], dim=-1) - return xs + ckpt = torch.load(ckpt_path, map_location="cpu") - def __init__(self, action_statistic_dof, dof_config): - """ - Initialize the normalizer with robot-specific action statistics. + for key, value in ckpt.items(): + # Parse key: "min.robot_name" -> prefix="min", name="robot_name" + try: + prefix, name = key.split(".", 1) + if hasattr(instance, prefix): + getattr(instance, prefix)[name] = nn.Parameter( + value, requires_grad=False + ) + print("prefix", prefix) + print("name", name) + except ValueError: + continue - Args: - action_statistic_dof (dict): Statistical data for each robot's degrees of freedom - dof_config (dict): Configuration mapping for degrees of freedom per robot - """ + return instance + + def __init__( + self, action_statistic_dof, dof_config, min_key="min", delta_key="delta" + ): super(Normalizer, self).__init__() - action_statistic = {} - # hard code the action dimension to 20 - action_dim = 20 + self.min_key = min_key + self.delta_key = delta_key - # Process statistics for each robot + action_statistic = {} for robot_name in action_statistic_dof.keys(): action_statistic[robot_name] = {} all_dof_min = [] all_dof_delta = [] - - # Collect min and delta values for all DOFs for k in dof_config: if k in action_statistic_dof[robot_name]: - all_dof_min.extend(action_statistic_dof[robot_name][k]["min"]) - all_dof_delta.extend(action_statistic_dof[robot_name][k]["delta"]) + if ( + min_key in action_statistic_dof[robot_name][k] + and delta_key in action_statistic_dof[robot_name][k] + ): + all_dof_min.extend(action_statistic_dof[robot_name][k][min_key]) + all_dof_delta.extend( + action_statistic_dof[robot_name][k][delta_key] + ) + else: + if robot_name == "x2_normal" or "libero" in robot_name: + print_rank_last( + f"Normalizer (Warning): min_key {min_key} or delta_key {delta_key} " + ) + print_rank_last( + f"not in action_statistic_dof[{robot_name}][{k}], use default min 0.0 and delta 1.0" + ) + all_dof_min.extend([0.0] * dof_config[k]) + all_dof_delta.extend([1.0] * dof_config[k]) else: - # Use default values if statistics not available - # raise ValueError(f"Statistics not available for {k} of {robot_name}") - logging.warning( - f"Statistics not available for {k} of {robot_name}, using default values" - ) + if robot_name == "x2_normal" or "libero" in robot_name: + print_rank_last( + f"Normalizer (Warning): Action {k} not in action_statistic_dof for {robot_name}, use default min 0.0 and delta 1.0" + ) all_dof_min.extend([0.0] * dof_config[k]) all_dof_delta.extend([1.0] * dof_config[k]) + all_dof_min = torch.tensor(all_dof_min) + all_dof_delta = torch.tensor(all_dof_delta) + action_statistic[robot_name][min_key] = all_dof_min + action_statistic[robot_name][delta_key] = all_dof_delta - all_dof_min = self._pad_to_action_dim(torch.tensor(all_dof_min), action_dim) - all_dof_delta = self._pad_to_action_dim( - torch.tensor(all_dof_delta), action_dim - ) - action_statistic[robot_name]["min"] = all_dof_min - action_statistic[robot_name]["delta"] = all_dof_delta - - # Register statistics as non-trainable parameters self.min = nn.ParameterDict( { - k: nn.Parameter(action_statistic[k]["min"], requires_grad=False) + k: nn.Parameter(action_statistic[k][min_key], requires_grad=False) for k in action_statistic.keys() } ) self.delta = nn.ParameterDict( { - k: nn.Parameter(action_statistic[k]["delta"], requires_grad=False) + k: nn.Parameter(action_statistic[k][delta_key], requires_grad=False) for k in action_statistic.keys() } ) - def normalize_data(self, xs, dataset_names, dof_mask=None): - """ - Normalize action data to [-1, 1] range using robot-specific statistics. + for k, v in action_statistic.items(): + print_rank_last( + f"Normalizer: {k} min {action_statistic[k][min_key]} delta {action_statistic[k][delta_key]}" + ) - Args: - xs: Input action data tensors - dataset_names: List of dataset/robot names corresponding to each tensor - - Returns: - torch.Tensor: Normalized action data in [-1, 1] range - """ + def normalize_data(self, xs, dataset_names): new_xs = [] - # Filter out multimodal dataset entries dataset_names = [name for name in dataset_names if name != "x2_multimodal"] - dof_mask = dof_mask if dof_mask is not None else [None] * len(xs) - - for x, dataset_name, mask in zip(xs, dataset_names, dof_mask): - # Apply DOF mask if provided - if mask is not None: - mask = mask[0].bool() - action_space_delta = self.delta[dataset_name][mask] - action_space_min = self.min[dataset_name][mask] - else: - action_space_delta = self.delta[dataset_name] - action_space_min = self.min[dataset_name] - # Apply min-max normalization - x = (x - action_space_min) / (action_space_delta) - # Scale to [-1, 1] range + for x, dataset_name in zip(xs, dataset_names): + x = (x - self.min[dataset_name]) / (self.delta[dataset_name]) x = x * 2 - 1 - # Clamp to ensure bounds x = torch.clamp(x, -1, 1) new_xs.append(x) - new_xs = torch.stack(new_xs) return new_xs def unnormalize_data(self, xs, dataset_names, dof_mask=None): - """ - Convert normalized data back to original action space. - - Args: - xs: Normalized action data in [-1, 1] range - dataset_names: List of dataset/robot names - dof_mask: Optional mask to select specific degrees of freedom - - Returns: - torch.Tensor: Denormalized action data in original scale - """ new_xs = [] - # Filter out multimodal dataset entries dataset_names = [name for name in dataset_names if name != "x2_multimodal"] dof_mask = dof_mask if dof_mask is not None else [None] * len(xs) - for x, dataset_name, mask in zip(xs, dataset_names, dof_mask): - # Convert from [-1, 1] to [0, 1] range x = (x + 1) / 2 - - # Apply DOF mask if provided if mask is not None: mask = mask[0].bool() action_space_delta = self.delta[dataset_name][mask] @@ -145,166 +133,498 @@ class Normalizer(nn.Module): else: action_space_delta = self.delta[dataset_name] action_space_min = self.min[dataset_name] - - # Scale back to original range x = x * action_space_delta + action_space_min new_xs.append(x) - new_xs = torch.stack(new_xs) return new_xs class SinusoidalPosEmb(nn.Module): - """ - Sinusoidal positional embedding for diffusion timesteps. - - Generates sinusoidal embeddings commonly used in diffusion models to encode - timestep information with different frequencies. - """ - - def __init__(self, dim): - """ - Initialize sinusoidal positional embedding. - - Args: - dim (int): Embedding dimension (must be even) - """ + def __init__(self, dim: int, min_period: float = 4e-3, max_period: float = 4.0): super().__init__() self.dim = dim + if dim % 2 != 0: + raise ValueError(f"embedding_dim ({dim}) must be divisible by 2") + self.min_period = min_period + self.max_period = max_period def forward(self, x): - """ - Generate sinusoidal embeddings for input timesteps. - - Args: - x (torch.Tensor): Input timesteps - - Returns: - torch.Tensor: Sinusoidal embeddings of shape (..., dim) - """ device = x.device half_dim = self.dim // 2 emb = math.log(10000) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, device=device) * -emb) + emb = torch.exp( + torch.arange(half_dim, device=device, dtype=torch.float32) * -emb + ) emb = x[:, None] * emb[None, :] emb = torch.cat((emb.sin(), emb.cos()), dim=-1) return emb -class ActionProcessor(nn.Module): +class Downsample1d(nn.Module): + def __init__(self, dim): + super().__init__() + self.conv = nn.Conv1d(dim, dim, 3, 2, 1) + + def forward(self, x): + return self.conv(x) + + +class Upsample1d(nn.Module): + def __init__(self, dim): + super().__init__() + self.conv = nn.ConvTranspose1d(dim, dim, 4, 2, 1) + + def forward(self, x): + return self.conv(x) + + +class Conv1dBlock(nn.Module): """ - Action sequence processor for robotic control with flow matching. - - This module handles action sequence processing for robotic systems with the following capabilities: - 1. Adds controlled noise to action sequences using Beta distribution scheduling - 2. Generates temporal embeddings for timestep conditioning - 3. Projects actions to model hidden space for transformer processing - 4. Supports proprioceptive data integration and multi-robot configurations - - The Beta distribution provides more flexible noise injection strategies compared to - traditional linear schedules, allowing better control over the noise scheduling process. + Conv1d --> GroupNorm --> Mish """ - def __init__(self, config): - """ - Initialize the action processor with multi-robot support. - - Args: - config: Configuration object containing: - - dof_config (dict): Degrees of freedom configuration per robot type - - agent_pos_config (dict): Agent position/proprioception configuration - - hidden_size (int): Model hidden layer dimension - - noise_scheduler (dict): Noise scheduler configuration with Beta parameters - """ + def __init__(self, inp_channels, out_channels, kernel_size, n_groups=8): super().__init__() - # Calculate action and proprioception dimensions from configuration + self.block = nn.Sequential( + nn.Conv1d( + inp_channels, out_channels, kernel_size, padding=kernel_size // 2 + ), + nn.GroupNorm(n_groups, out_channels), + nn.Mish(), + ) + + def forward(self, x): + return self.block(x) + + +class ConditionalResidualBlock1D(nn.Module): + def __init__(self, in_channels, out_channels, cond_dim, kernel_size=3, n_groups=8): + super().__init__() + + self.blocks = nn.ModuleList( + [ + Conv1dBlock(in_channels, out_channels, kernel_size, n_groups=n_groups), + Conv1dBlock(out_channels, out_channels, kernel_size, n_groups=n_groups), + ] + ) + + # FiLM modulation https://arxiv.org/abs/1709.07871 + # predicts per-channel scale and bias + cond_channels = out_channels * 2 + self.out_channels = out_channels + self.cond_encoder = nn.Sequential( + nn.Mish(), + nn.Linear(cond_dim, cond_channels), + nn.Dropout(0.1), + nn.Unflatten(-1, (-1, 1)), + ) + + # make sure dimensions compatible + self.residual_conv = ( + nn.Conv1d(in_channels, out_channels, 1) + if in_channels != out_channels + else nn.Identity() + ) + + def forward(self, x, cond): + """ + x : [ batch_size x in_channels x horizon ] + cond : [ batch_size x cond_dim] + + returns: + out : [ batch_size x out_channels x horizon ] + """ + out = self.blocks[0](x) + embed = self.cond_encoder(cond) + + embed = embed.reshape(embed.shape[0], 2, self.out_channels, 1) + scale = embed[:, 0, ...] + bias = embed[:, 1, ...] + out = scale * out + bias + + out = self.blocks[1](out) + out = out + self.residual_conv(x) + return out + + +class ConditionalUnet1D(nn.Module): + def __init__( + self, + input_dim, + global_cond_dim, + diffusion_step_embed_dim=256, + down_dims=[256, 512, 1024], + # down_dims=[512, 1024, 2048], + kernel_size=5, + n_groups=8, + ): + """ + input_dim: Dim of actions. + global_cond_dim: Dim of global conditioning applied with FiLM + in addition to diffusion step embedding. This is usually obs_horizon * obs_dim + diffusion_step_embed_dim: Size of positional encoding for diffusion iteration k + down_dims: Channel size for each UNet level. + The length of this array determines numebr of levels. + kernel_size: Conv kernel size + n_groups: Number of groups for GroupNorm + """ + + super().__init__() + all_dims = [input_dim] + list(down_dims) + start_dim = down_dims[0] + + dsed = diffusion_step_embed_dim + diffusion_step_encoder = nn.Sequential( + SinusoidalPosEmb(dsed), + nn.Linear(dsed, dsed * 4), + nn.Mish(), + nn.Linear(dsed * 4, dsed), + ) + cond_dim = dsed + global_cond_dim + + in_out = list(zip(all_dims[:-1], all_dims[1:])) + mid_dim = all_dims[-1] + self.mid_modules = nn.ModuleList( + [ + ConditionalResidualBlock1D( + mid_dim, + mid_dim, + cond_dim=cond_dim, + kernel_size=kernel_size, + n_groups=n_groups, + ), + ConditionalResidualBlock1D( + mid_dim, + mid_dim, + cond_dim=cond_dim, + kernel_size=kernel_size, + n_groups=n_groups, + ), + ] + ) + + down_modules = nn.ModuleList([]) + for ind, (dim_in, dim_out) in enumerate(in_out): + is_last = ind >= (len(in_out) - 1) + down_modules.append( + nn.ModuleList( + [ + ConditionalResidualBlock1D( + dim_in, + dim_out, + cond_dim=cond_dim, + kernel_size=kernel_size, + n_groups=n_groups, + ), + ConditionalResidualBlock1D( + dim_out, + dim_out, + cond_dim=cond_dim, + kernel_size=kernel_size, + n_groups=n_groups, + ), + Downsample1d(dim_out) if not is_last else nn.Identity(), + ] + ) + ) + + up_modules = nn.ModuleList([]) + for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])): + is_last = ind >= (len(in_out) - 1) + up_modules.append( + nn.ModuleList( + [ + ConditionalResidualBlock1D( + dim_out * 2, + dim_in, + cond_dim=cond_dim, + kernel_size=kernel_size, + n_groups=n_groups, + ), + ConditionalResidualBlock1D( + dim_in, + dim_in, + cond_dim=cond_dim, + kernel_size=kernel_size, + n_groups=n_groups, + ), + Upsample1d(dim_in) if not is_last else nn.Identity(), + ] + ) + ) + + final_conv = nn.Sequential( + Conv1dBlock(start_dim, start_dim, kernel_size=kernel_size), + nn.Conv1d(start_dim, input_dim, 1), + ) + + self.diffusion_step_encoder = diffusion_step_encoder + self.up_modules = up_modules + self.down_modules = down_modules + self.final_conv = final_conv + + # print("number of parameters: {:e}".format( + # sum(p.numel() for p in self.parameters())) + # ) + + def forward( + self, + sample: torch.Tensor, + timestep: Union[torch.Tensor, float, int], + global_cond=None, + ): + """ + x: (B,T,input_dim) + timestep: (B,) or int, diffusion step + global_cond: (B,global_cond_dim) + output: (B,T,input_dim) + """ + # (B,T,C) + sample = sample.moveaxis(-1, -2) + # (B,C,T) + + # 1. time + timesteps = timestep + if not torch.is_tensor(timesteps): + # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can + timesteps = torch.tensor( + [timesteps], dtype=torch.long, device=sample.device + ) + elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0: + timesteps = timesteps[None].to(sample.device) + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timesteps = timesteps.expand(sample.shape[0]) + + global_feature = self.diffusion_step_encoder(timesteps) + + if global_cond is not None: + global_feature = torch.cat([global_feature, global_cond], axis=-1) + x = sample + h = [] + for idx, (resnet, resnet2, downsample) in enumerate(self.down_modules): + x = resnet(x, global_feature) + x = resnet2(x, global_feature) + h.append(x) + x = downsample(x) # bs, 2048, 5 + + for mid_module in self.mid_modules: + x = mid_module(x, global_feature) # bs, 2048, 5 + + for idx, (resnet, resnet2, upsample) in enumerate(self.up_modules): + x = torch.cat((x, h.pop()), dim=1) + x = resnet(x, global_feature) + x = resnet2(x, global_feature) + x = upsample(x) # bs, 512, 20 + + x = self.final_conv(x) # bs, 14, 20 + + # (B,C,T) + x = x.moveaxis(-1, -2) + # (B,T,C) + return x + + +class DP_Action_head(nn.Module): + def __init__( + self, + action_dim=14, + transformer_dim=896, + global_cond_dim=1806, + load_pretrained=True, + ): + super().__init__() + self.action_dim = action_dim + self.transformer_dim = transformer_dim + self.load_pretrained = load_pretrained + self.noise_scheduler = DDPMScheduler( + num_train_timesteps=132, + beta_schedule="squaredcos_cap_v2", + clip_sample=True, + prediction_type="epsilon", + ) + if self.load_pretrained: + self.global_cond_dim = global_cond_dim + self.condition_proj = nn.Sequential( + nn.Linear(self.transformer_dim, 2 * self.transformer_dim), + nn.ReLU(), + nn.Linear(2 * self.transformer_dim, 2 * self.global_cond_dim), + nn.ReLU(), + nn.Linear(2 * self.global_cond_dim, self.global_cond_dim), + ) + + self.noise_pred_net = ConditionalUnet1D( + input_dim=self.action_dim, + # down_dims=[256,512,1024], + down_dims=[512, 1024, 2048], + global_cond_dim=self.global_cond_dim, + ) + + # load pretrained model + action_pretrained_path = "/x2robot/liangyuxin/workspace/DiffusionPolicy/big_mix_0718_mn/30_noise_pred_net.pth" + print("load noise_pred_net from:", action_pretrained_path, flush=True) + self.noise_pred_net.load_state_dict(torch.load(action_pretrained_path)) + else: + self.noise_pred_net = ConditionalUnet1D( + input_dim=self.action_dim, + down_dims=[256, 512, 1024], + global_cond_dim=self.transformer_dim, + ) + + def forward(self, naction, condition, sample_times): + bs = naction.shape[0] + noise_shape = ( + naction.shape[0] * sample_times, + naction.shape[1], + naction.shape[2], + ) + noise = torch.randn(noise_shape, device=naction.device) + naction = ( + naction.unsqueeze(1) + .repeat(1, sample_times, 1, 1) + .reshape(bs * sample_times, naction.shape[1], naction.shape[2]) + ) + + timesteps = torch.randint( + 0, + self.noise_scheduler.config.num_train_timesteps, + (bs * sample_times,), + device=naction.device, + ).long() + condition = condition.to(self.condition_proj[0].weight.data.dtype) + if self.load_pretrained: + condition = self.condition_proj(condition) + + noisy_actions = self.noise_scheduler.add_noise(naction, noise, timesteps) + noise_pred = self.noise_pred_net( + noisy_actions, timesteps, global_cond=condition + ) + return noise, noise_pred + + @torch.no_grad() + def predict(self, condition, naction=None): + bs = condition.shape[0] + condition = condition.to(self.condition_proj[0].weight.data.dtype) + if self.load_pretrained: + condition = self.condition_proj(condition) + + if naction is not None: + noise_shape = (naction.shape[0], naction.shape[1], naction.shape[2]) + else: + noise_shape = (bs, 16, self.action_dim) # tobe parameterized + noise = torch.randn(noise_shape, device=condition.device) + naction_pred = noise + # init scheduler + self.noise_scheduler.set_timesteps( + self.noise_scheduler.config.num_train_timesteps + ) + + for k in self.noise_scheduler.timesteps: + # predict noise + noise_pred = self.noise_pred_net( + sample=naction_pred, timestep=k, global_cond=condition + ) + + # inverse diffusion step (remove noise) + naction_pred = self.noise_scheduler.step( + model_output=noise_pred, timestep=k, sample=naction_pred + ).prev_sample + + return naction, naction_pred + + +class ActionProcessor(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config self.dof_config = config.dof_config self.agent_pos_config = config.agent_pos_config self.action_dim = sum([v for k, v in self.dof_config.items()]) self.propri_dim = sum([v for k, v in self.agent_pos_config.items()]) - # Log configuration details for debugging - print("ActionProcessor Configuration:", flush=True) - print(f" Action dimension: {self.action_dim}", flush=True) - print(f" Proprioception dimension: {self.propri_dim}", flush=True) - print(" DOF configuration:", flush=True) - for key, value in self.dof_config.items(): - print(f" {key}: {value}", flush=True) - print(" Agent position configuration:", flush=True) - for key, value in self.agent_pos_config.items(): - print(f" {key}: {value}", flush=True) + print_rank_last( + f"self.dof_config: {self.dof_config}; action_dim: {self.action_dim}; self.agent_pos_config: {self.agent_pos_config}; propri_dim: {self.propri_dim}" + ) + self.action_hidden_size = config.action_hidden_size + self.state_hidden_size = config.state_hidden_size self.hidden_size = config.hidden_size - # Initialize data normalizers for actions and proprioception - self.normalizer_action = Normalizer( - action_statistic_dof, - ( - config.customized_dof_config - if hasattr(config, "customized_dof_config") - else config.dof_config - ), - ) - self.normalizer_propri = Normalizer( - action_statistic_dof, - ( - config.customized_agent_pos_config - if hasattr(config, "customized_agent_pos_config") - else config.agent_pos_config - ), - ) + if not self.config.use_state_string_representation: + if self.config.proj_with_mask: + self.propri_proj = nn.Linear( + self.propri_dim * 2, self.state_hidden_size, bias=False + ) + else: + self.propri_proj = nn.Linear( + self.propri_dim, self.state_hidden_size, bias=False + ) - # Proprioception projection layer (includes history/current state) - self.propri_proj = nn.Linear(self.propri_dim * 2, self.hidden_size, bias=False) + # noise scheduler configing + if getattr(self.config, "use_flow_action_expert", True): + noise_scheduler_config = config.noise_scheduler + self.beta_alpha = noise_scheduler_config.get("beta_alpha", 1.5) + self.beta_beta = noise_scheduler_config.get("beta_beta", 1.0) + self.s = noise_scheduler_config.get("s", 0.999) + alpha_tensor = torch.tensor(self.beta_alpha, dtype=torch.float32).to("cuda") + beta_tensor = torch.tensor(self.beta_beta, dtype=torch.float32).to("cuda") + self.beta_dist = Beta(alpha_tensor, beta_tensor) + self.time_embed = SinusoidalPosEmb(self.action_hidden_size) - # Beta distribution noise scheduler configuration - noise_scheduler_config = config.noise_scheduler - self.beta_alpha = noise_scheduler_config.get( - "beta_alpha", 1.5 - ) # Beta distribution α parameter - self.beta_beta = noise_scheduler_config.get( - "beta_beta", 1.0 - ) # Beta distribution β parameter - self.s = noise_scheduler_config.get("s", 0.999) # Scaling factor + # project to hidden space + if self.config.proj_with_mask: + self.w1 = nn.Linear( + self.action_dim * 2, self.action_hidden_size, bias=False + ) + else: + self.w1 = nn.Linear( + self.action_dim, self.action_hidden_size, bias=False + ) + if not self.config.use_adarms: + self.w2 = nn.Linear( + self.action_hidden_size * 2, self.action_hidden_size, bias=False + ) + self.w3 = nn.Linear( + self.action_hidden_size, self.action_hidden_size, bias=False + ) + self.act_fn = nn.SiLU() + else: + self.time_mlp_in = nn.Linear( + self.action_hidden_size, self.action_hidden_size + ) + self.time_mlp_out = nn.Linear( + self.action_hidden_size, self.action_hidden_size + ) + self.act_fn = nn.SiLU() - # Initialize Beta distribution for noise scheduling - alpha_tensor = torch.tensor(self.beta_alpha, dtype=torch.float32).to("cuda") - beta_tensor = torch.tensor(self.beta_beta, dtype=torch.float32).to("cuda") - self.beta_dist = Beta(alpha_tensor, beta_tensor) + # project back to action space + self.action_proj_back = nn.Linear( + self.action_hidden_size, self.action_dim, bias=False + ) + self.mse_loss = nn.MSELoss(reduction="none") - # Sinusoidal positional embedding for timesteps - self.time_embed = SinusoidalPosEmb(config.hidden_size) + def set_normalizer(self, normalizer_action, normalizer_propri): + self.normalizer_action = normalizer_action + self.normalizer_propri = normalizer_propri - # Action embedding network: project to hidden space - self.w1 = nn.Linear( - self.action_dim * 2, self.hidden_size, bias=False - ) # *2 for action + DOF mask - self.w2 = nn.Linear( - self.hidden_size * 2, self.hidden_size, bias=False - ) # *2 for action + time embeddings - self.w3 = nn.Linear(self.hidden_size, self.hidden_size, bias=False) - self.act_fn = nn.SiLU() - - # Project back to action space for flow matching loss - self.action_proj_back = nn.Linear(self.hidden_size, self.action_dim, bias=False) - self.mse_loss = nn.MSELoss(reduction="none") + # dataset_name = self.config["data"]["lerobot_config"]["repo_id"] + # print("normalizer_propri min", self.normalizer_propri.min.__getattr__(dataset_name), flush=True) + # print("normalizer_propri delta", self.normalizer_propri.delta.__getattr__(dataset_name), flush=True) + # print("normalizer_action min", self.normalizer_action.min.__getattr__(dataset_name), flush=True) + # print("normalizer_action delta", self.normalizer_action.delta.__getattr__(dataset_name), flush=True) def sample_time(self, batch_size, device, dtype): """ - Sample timesteps using Beta distribution for noise scheduling. + Sampling Time Step + Generates random numbers in the range [0, 1] using a Beta distribution, and then scales them. - Generates random timesteps in [0,1] range using Beta distribution, then scales them. - This provides more flexible control over the noise injection schedule compared to - uniform sampling. - - Args: - batch_size (int): Number of timesteps to sample - device: Target device for tensors - dtype: Target data type for tensors + Parameters: + batch_size (int): Batch size + device: Device type + dtype: Data type Returns: - torch.Tensor: Sampled timesteps of shape [batch_size] + torch.Tensor: Sampled time steps, with shape [batch_size] """ sample = self.beta_dist.sample([batch_size]).to(device=device, dtype=dtype) time = (1 - sample) * self.s @@ -314,152 +634,175 @@ class ActionProcessor(nn.Module): self, proprioception, dataset_names=None, dof_mask=None, use_history=False ): """ - Project proprioceptive data (joint positions, orientations) to hidden space. - - Args: - proprioception (torch.Tensor): Proprioceptive data of shape [batch_size, seq_len, propri_dim] - dataset_names (list, optional): Dataset names for normalization. Defaults to None. - dof_mask (torch.Tensor, optional): DOF mask of shape [batch_size, propri_dim]. Defaults to None. - use_history (bool, optional): Whether to use historical proprioceptive data. Defaults to False. - - Returns: - torch.Tensor: Projected proprioceptive features of shape [batch_size, seq_len, hidden_size] + proprioception: [batch_size, 1, action_dim] + dataset_names: [batch_size] + dof_mask: [batch_size, action_dim] """ - # Ensure proper device and dtype alignment proprioception = proprioception.to(device=self.propri_proj.weight.device).to( dtype=self.propri_proj.weight.dtype ) - if dof_mask is not None: - # Concatenate proprioception with DOF mask - # TODO: Use variable-based dimension checking for better flexibility - if use_history: - proprioception = torch.cat([proprioception, dof_mask], dim=-1) - else: - proprioception = torch.cat([proprioception, dof_mask], dim=-1) - + if self.config.proj_with_mask: + proprioception = torch.cat( + [proprioception, dof_mask], dim=-1 + ) # .unsqueeze(1) proprioception = proprioception.to(device=self.propri_proj.weight.device).to( dtype=self.propri_proj.weight.dtype ) - return self.propri_proj(proprioception) + proprio_embed = self.propri_proj( + proprioception + ) # [batch_size, 1, state_hidden_size] + + if self.state_hidden_size < self.hidden_size: + # padding to hidden size + padding_size = self.hidden_size - self.state_hidden_size + padding = torch.zeros( + (proprio_embed.shape[0], 1, padding_size), + device=proprio_embed.device, + dtype=proprio_embed.dtype, + ) + proprio_embed = torch.cat([proprio_embed, padding], dim=-1) + + return proprio_embed # [batch_size, 1, hidden_size] def forward(self, action_chunk, dataset_names, dof_mask=None): """ - Process action sequences with noise injection and temporal embedding. - - This method implements the forward pass for flow matching training: - 1. Adds Beta-distributed noise to action sequences - 2. Generates sinusoidal timestep embeddings - 3. Projects noisy actions to hidden space - 4. Combines action and temporal features - - Args: - action_chunk (torch.Tensor): Action sequences of shape [batch_size, seq_len, action_dim] - dataset_names (list): Dataset names for normalization - dof_mask (torch.Tensor, optional): DOF mask of shape [batch_size, seq_len, action_dim]. - Defaults to None. + Parameters: + action_chunk (torch.Tensor): Action sequence, shape [batch_size, action_chunk_len, action_dim] + dataset_names: [batch_size] + dof_mask: [batch_size, action_dim] Returns: - tuple: (action_embeddings, flow_target) where: - - action_embeddings: Processed action features of shape [batch_size, seq_len, hidden_size] - - flow_target: Flow matching target (action_chunk - noise) for loss computation + torch.Tensor: Processed action representation, shape [batch_size, seq_len, hidden_size] """ - batch_size = action_chunk.shape[0] - device = action_chunk.device - dtype = action_chunk.dtype + with torch.autocast("cuda", dtype=torch.float32): + action_chunk = action_chunk.to(dtype=torch.float32) + batch_size = action_chunk.shape[0] + device = action_chunk.device + dtype = action_chunk.dtype - # 1. Add noise to action sequences using flow matching - noise = torch.randn_like(action_chunk) - time = self.sample_time(batch_size, device, dtype) - t = time.unsqueeze(-1).unsqueeze(-1) # Broadcast to match action dimensions + # 1. add noise to action_chunk + noise = torch.randn_like(action_chunk) + time = self.sample_time(batch_size, device, dtype) + time_expanded = time.unsqueeze(-1).unsqueeze(-1) + noisy_action = (1 - time_expanded) * noise + time_expanded * action_chunk + flow = action_chunk - noise - # Linear interpolation between noise and action (flow matching) - noisy_action = (1 - t) * noise + t * action_chunk - flow = action_chunk - noise # Flow target for loss computation + # 2. sinusoidal positional encoding for timesteps + time_embed = self.time_embed(time).to(torch.float32) - # 2. Generate sinusoidal positional encoding for timesteps - time_embed = self.time_embed(time) + self.noise = noise + self.noisy_action = noisy_action # for new x-pred - # 3. Project noisy actions with DOF mask to hidden space - if dof_mask is not None: - noisy_action = torch.cat([noisy_action, dof_mask], dim=-1) + # 3.action_chunk_nosiy + t_pos_emb -> MLP_act_chunk -> action_chunk_nosiy_emb_with_t (dim=trans * chunk) + if dof_mask is not None: + noisy_action = torch.cat([noisy_action, dof_mask], dim=-1) - noisy_action = noisy_action.to(dtype=self.w1.weight.dtype) - action_embed = self.w1(noisy_action) + noisy_action = noisy_action.to(dtype=self.w1.weight.dtype) + action_embed = self.w1(noisy_action) - # Repeat time embedding for each sequence position - time_embed = ( - time_embed.unsqueeze(1) - .repeat(1, action_embed.shape[1], 1) - .to(dtype=self.w2.weight.dtype) - ) + self.time_expanded = time_expanded # for new x-pred - # Combine action and temporal embeddings - concat_embed = torch.cat([action_embed, time_embed], dim=-1) - concat_embed = self.w2(concat_embed) - embed = self.w3(self.act_fn(concat_embed)) + if not self.config.use_adarms: + time_embed = ( + time_embed.unsqueeze(1) + .repeat(1, action_embed.shape[1], 1) + .to(dtype=self.w2.weight.dtype) + ) + concat_embed = torch.cat([action_embed, time_embed], dim=-1) + concat_embed = self.w2(concat_embed) + action_time_embed = self.w3(self.act_fn(concat_embed)) + adarms_cond = None + else: + time_embed = self.time_mlp_in(time_embed) + time_embed = self.act_fn(time_embed) + time_embed = self.time_mlp_out(time_embed) + time_embed = self.act_fn(time_embed) + action_time_embed = action_embed + adarms_cond = time_embed - return embed, flow + if self.action_hidden_size < self.hidden_size: + # padding to hidden size + padding_size = self.hidden_size - self.action_hidden_size + padding = torch.zeros( + ( + action_time_embed.shape[0], + action_time_embed.shape[1], + padding_size, + ), + device=action_time_embed.device, + dtype=action_time_embed.dtype, + ) + action_time_embed = torch.cat([action_time_embed, padding], dim=-1) + + return action_time_embed, flow, adarms_cond def step(self, timestep, noisy_action, dof_mask=None): - """ - Single denoising step for diffusion inference. + # noisy_action: bs, pred_horizon, action_dim + # timestep: bs + with torch.autocast("cuda", dtype=torch.float32): + if dof_mask is not None and self.config.proj_with_mask: + if dof_mask.shape[1] == 1: + dof_mask = dof_mask.unsqueeze(1).repeat(1, noisy_action.shape[1], 1) + noisy_action = torch.cat([noisy_action, dof_mask], dim=-1) - Processes noisy actions at a specific timestep for iterative denoising during inference. + noisy_action = noisy_action.to(dtype=self.w1.weight.dtype) + time_embed = self.time_embed(timestep).to(torch.float32) # bs,hidden_size + action_embed = self.w1(noisy_action) - Args: - timestep (torch.Tensor): Current timesteps of shape [batch_size] - noisy_action (torch.Tensor): Noisy actions of shape [batch_size, seq_len, action_dim] - dof_mask (torch.Tensor, optional): DOF mask for action space. Defaults to None. + if not self.config.use_adarms: + time_embed = time_embed.unsqueeze(1).repeat(1, action_embed.shape[1], 1) + time_embed = time_embed.to(device=noisy_action.device).to( + dtype=noisy_action.dtype + ) + concat_embed = torch.cat([action_embed, time_embed], dim=-1) + concat_embed = self.w2(concat_embed) + embed = self.w3(self.act_fn(concat_embed)) # is this right? + adarms_cond = None + else: + time_embed = time_embed.to(dtype=self.time_mlp_in.weight.dtype) + time_embed = self.time_mlp_in(time_embed) + time_embed = self.act_fn(time_embed) + time_embed = self.time_mlp_out(time_embed) + time_embed = self.act_fn(time_embed) + embed = action_embed + adarms_cond = time_embed - Returns: - torch.Tensor: Processed action embeddings of shape [batch_size, seq_len, hidden_size] - """ - # Concatenate noisy action with DOF mask if provided - if dof_mask is not None: - noisy_action = torch.cat([noisy_action, dof_mask], dim=-1) + if self.action_hidden_size < self.hidden_size: + # padding to hidden size + padding_size = self.hidden_size - self.action_hidden_size + padding = torch.zeros( + (embed.shape[0], embed.shape[1], padding_size), + device=embed.device, + dtype=embed.dtype, + ) + embed = torch.cat([embed, padding], dim=-1) - # Generate timestep embeddings - time_embed = self.time_embed(timestep) # [batch_size, hidden_size] + return embed, adarms_cond - # Project noisy actions - action_embed = self.w1(noisy_action) + def flow_loss( + self, + action_hidden_states, + flow, + action_chunk, + dof_mask=None, + flow_loss_mask=None, + ): + with torch.autocast("cuda", dtype=torch.float32): + action_pred = self.action_proj_back( + action_hidden_states[:, : self.action_hidden_size] + ) + v_pred = action_pred + loss = self.mse_loss(v_pred, flow) + if dof_mask is not None: + dof_mask = dof_mask.reshape(-1, dof_mask.shape[-1]) + loss = loss * dof_mask - # Broadcast time embeddings to sequence length - time_embed = time_embed.unsqueeze(1).repeat(1, action_embed.shape[1], 1) - time_embed = time_embed.to(device=noisy_action.device).to( - dtype=noisy_action.dtype - ) - - # Combine embeddings and process through MLP - concat_embed = torch.cat([action_embed, time_embed], dim=-1) - concat_embed = self.w2(concat_embed) - embed = self.w3(self.act_fn(concat_embed)) - - return embed - - def flow_loss(self, action_hidden_states, flow, dof_mask=None): - """ - Compute flow matching loss between predicted and target actions. - - Args: - action_hidden_states (torch.Tensor): Hidden states from transformer - flow (torch.Tensor): Target flow (action - noise) for matching - dof_mask (torch.Tensor, optional): DOF mask to weight loss per dimension. Defaults to None. - - Returns: - torch.Tensor: Flow matching loss (no reduction for channel loss computation) - """ - # Project hidden states back to action space - action_pred = self.action_proj_back(action_hidden_states) - - # Compute MSE loss between predicted and target flow - loss = self.mse_loss(action_pred, flow) - - # Apply DOF mask if provided - if dof_mask is not None: - dof_mask = dof_mask.reshape(-1, dof_mask.shape[-1]) - loss = loss * dof_mask - - # Return loss without reduction for channel-wise loss computation + if flow_loss_mask is not None: + flow_loss_mask = ( + flow_loss_mask.unsqueeze(-1) + .reshape(-1, 1) + .expand(-1, loss.shape[-1]) + ) + loss = loss * flow_loss_mask return loss diff --git a/wall_x/model/joint_attention.py b/wall_x/model/joint_attention.py new file mode 100644 index 0000000..5ba9546 --- /dev/null +++ b/wall_x/model/joint_attention.py @@ -0,0 +1,651 @@ +import torch +import torch.nn as nn +from typing import Optional, Tuple + +from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig +from transformers.cache_utils import Cache +from transformers.utils import logging +from wall_x.fusions import ops +from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl import ( + apply_multimodal_rotary_pos_emb, +) +from flash_attn import flash_attn_func +from transformers.modeling_flash_attention_utils import ( + is_flash_attn_greater_or_equal_2_10, +) +from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( + Qwen2_5_VLRotaryEmbedding, + repeat_kv, +) + +logger = logging.get_logger(__name__) + + +# def rotate_half(x): +# x1 = x[..., : x.shape[-1] // 2] +# x2 = x[..., x.shape[-1] // 2 :] +# return torch.cat((-x2, x1), dim=-1) + + +# def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=2): +# mrope_section = mrope_section * 2 +# cos_split = torch.cat( +# [m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1 +# ).unsqueeze(unsqueeze_dim) +# sin_split = torch.cat( +# [m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1 +# ).unsqueeze(unsqueeze_dim) +# q_embed = (q * cos_split) + (rotate_half(q) * sin_split) +# k_embed = (k * cos_split) + (rotate_half(k) * sin_split) +# return q_embed, k_embed + + +class JointQwen2VLAttention(nn.Module): + def __init__(self, config: Qwen2_5_VLConfig, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " + "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + if not hasattr(config, "dim_inputs") or not config.dim_inputs: + raise ValueError("Configuration must contain a valid dim_inputs") + + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = config.rope_theta + self.is_causal = True + self.attention_dropout = config.attention_dropout + self.rope_scaling = config.rope_scaling + + self.dim_inputs = config.dim_inputs # Tuple[int, ...] + + if config.model_type == "qwen2_5_vl": + bias_qkv = True + else: + bias_qkv = False + + self.q_proj_experts = nn.ModuleList( + [ + nn.Linear(dim_input, self.num_heads * self.head_dim, bias=bias_qkv) + for dim_input in self.dim_inputs + ] + ) + self.k_proj_experts = nn.ModuleList( + [ + nn.Linear( + dim_input, self.num_key_value_heads * self.head_dim, bias=bias_qkv + ) + for dim_input in self.dim_inputs + ] + ) + self.v_proj_experts = nn.ModuleList( + [ + nn.Linear( + dim_input, self.num_key_value_heads * self.head_dim, bias=bias_qkv + ) + for dim_input in self.dim_inputs + ] + ) + self.o_proj_experts = nn.ModuleList( + [ + nn.Linear(self.num_heads * self.head_dim, dim_input, bias=False) + for dim_input in self.dim_inputs + ] + ) + + # Rotary embedding init + if config.model_type == "qwen2_5_vl": + self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) + else: + raise NotImplementedError(f"Unsupported model type: {config.model_type}") + + def repeat_kv(self, hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + Repeat key/value heads along the num_key_value_heads dimension (which is dim=2). + Input shape: (batch, seqlen, num_key_value_heads, head_dim) + Output shape: (batch, seqlen, num_key_value_heads * n_rep, head_dim) + Equivalent to torch.repeat_interleave(x, dim=2, repeats=n_rep) + """ + if n_rep == 1: + return hidden_states + + batch, slen, num_key_value_heads, head_dim = hidden_states.shape + + hidden_states = hidden_states.unsqueeze(3) + + hidden_states = hidden_states.expand( + batch, slen, num_key_value_heads, n_rep, head_dim + ) + + return hidden_states.reshape(batch, slen, num_key_value_heads * n_rep, head_dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + token_types: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + probs: Optional[torch.Tensor] = None, + row_id_map: Optional[torch.Tensor] = None, + orig_shape: Optional[Tuple[int]] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if token_types is None: + raise ValueError("token_types can not be None") + if token_types.max() >= len(self.dim_inputs): + raise ValueError( + f"token_types contains invalid expert indices: {token_types.max()}" + ) + + if self.config.mot_opt: + bsz, q_len, _ = orig_shape + query_states, key_states, value_states = self._generate_qkv_mot_opt( + hidden_states, + token_types, + start_indices, + end_indices, + probs, + row_id_map, + bsz, + q_len, + ) + else: + bsz, q_len, _ = hidden_states.size() + masks = [ + (token_types == expert_idx) + for expert_idx in range(len(self.dim_inputs)) + ] + query_states, key_states, value_states = self._generate_qkv( + hidden_states, masks + ) + + # Because the input can be padded, the absolute sequence length depends on the max position id. + cos, sin = position_embeddings + query_states, key_states = self._apply_rotary_pos_embed( + query_states, key_states, cos, sin, unsqueeze_dim=2 + ) + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + if past_key_value is not None: + cache_kwargs = { + "sin": sin, + "cos": cos, + "cache_position": cache_position, + } # Specific to RoPE models + if use_cache: + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + else: + past_key_states, past_value_states = past_key_value[self.layer_idx] + key_states = torch.cat([past_key_states, key_states], dim=-2) + value_states = torch.cat([past_value_states, value_states], dim=-2) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + causal_mask = attention_mask + if attention_mask is not None: + # Ensure that the attention_mask correctly matches across the head dimension. + if len(attention_mask.shape) == 2: # [batch_size, seq_len] + # Expanded to a causal mask format of [batch_size, 1, seq_len, seq_len] + bsz, seq_len = attention_mask.shape + causal_mask = attention_mask.view(bsz, 1, 1, seq_len).expand( + bsz, 1, seq_len, seq_len + ) + elif len(attention_mask.shape) == 3: # [batch_size, seq_len, seq_len] + # add head dimension: [batch_size, 1, seq_len, seq_len] + causal_mask = attention_mask.unsqueeze(1) + elif ( + len(attention_mask.shape) == 4 + ): # [batch_size, num_heads, seq_len, seq_len] + causal_mask = attention_mask + else: + raise ValueError( + f"Unsupported attention_mask dim: {attention_mask.shape}" + ) + + # convert the attention mask to bool type + causal_mask = causal_mask.to(torch.bool) + + # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, + # Reference: https://github.com/pytorch/pytorch/issues/112577. + if query_states.device.type == "cuda" and attention_mask is not None: + query_states = query_states.contiguous() + key_states = key_states.contiguous() + value_states = value_states.contiguous() + + # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment + # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. + # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1. + is_causal = True if causal_mask is None and q_len > 1 else False + + if q_len == 1: + is_causal = False + causal_mask = torch.ones( + bsz, + 1, + 1, + key_states.shape[2], + device=hidden_states.device, + dtype=hidden_states.dtype, + ).contiguous() + causal_mask = causal_mask.to(torch.bool) + + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=causal_mask, + dropout_p=self.attention_dropout if self.training else 0.0, + is_causal=is_causal, + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.view(bsz, q_len, -1) + + if self.config.mot_opt: + output = self._generate_output_mot_opt( + attn_output, token_types, start_indices, end_indices + ) + else: + output = self._generate_output(attn_output, masks) + + return output, None, past_key_value + + def _generate_qkv(self, hidden_states, masks): + bsz, q_len, _ = hidden_states.size() + + query_states = torch.zeros( + bsz, + q_len, + self.num_heads, + self.head_dim, + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + key_states = torch.zeros( + bsz, + q_len, + self.num_key_value_heads, + self.head_dim, + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + value_states = torch.zeros( + bsz, + q_len, + self.num_key_value_heads, + self.head_dim, + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + # for expert_idx in range(len(self.dim_inputs)): + for expert_idx, (q_proj, k_proj, v_proj, mask) in enumerate( + zip(self.q_proj_experts, self.k_proj_experts, self.v_proj_experts, masks) + ): + if not mask.any(): + continue + dim_input = self.dim_inputs[expert_idx] + + selected_hidden = hidden_states[mask].clone() + + q_out = q_proj(selected_hidden[:, :dim_input]).view( + -1, self.num_heads, self.head_dim + ) + k_out = k_proj(selected_hidden[:, :dim_input]).view( + -1, self.num_key_value_heads, self.head_dim + ) + v_out = v_proj(selected_hidden[:, :dim_input]).view( + -1, self.num_key_value_heads, self.head_dim + ) + + if self.config.model_type == "qwen3_vl_text": + q_out = self.q_norms[expert_idx](q_out)[0] + k_out = self.k_norms[expert_idx](k_out)[0] + + query_states[mask] = q_out + key_states[mask] = k_out + value_states[mask] = v_out + + return query_states, key_states, value_states + + def _generate_qkv_mot_opt( + self, + hidden_states: torch.Tensor, + experts_indices: torch.Tensor, + start_indices: torch.Tensor, + end_indices: torch.Tensor, + probs: torch.Tensor, + row_id_map: torch.Tensor, + batch_size: int, + seq_length: int, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Generate Q, K, V based on expert-sharded segments (start_indices / end_indices), + then restore them to the original sequence order. + + Args: + hidden_states: [total_tokens, hidden_dim], tokens already permuted and grouped by experts + experts_indices: [B, S], expert index for each token + start_indices: start token index for each expert (in the permuted token space) + end_indices: end token index for each expert (in the permuted token space) + probs: probability vector for each token (used for unpermute) + batch_size, seq_length: original batch size and sequence length + + Returns: + query_states: [B, num_heads, S, head_dim] + key_states: [B, num_key_value_heads, S, head_dim] + value_states: [B, num_key_value_heads, S, head_dim] + """ + + total_tokens, hidden_dim = hidden_states.shape + device, dtype = hidden_states.device, hidden_states.dtype + + # Initialize Q/K/V buffers in the permuted token space + q_buffer = torch.zeros(total_tokens, hidden_dim, device=device, dtype=dtype) + k_buffer = torch.zeros( + total_tokens, + self.num_key_value_heads * self.head_dim, + device=device, + dtype=dtype, + ) + v_buffer = torch.zeros( + total_tokens, + self.num_key_value_heads * self.head_dim, + device=device, + dtype=dtype, + ) + + # === Each expert processes its own token slice === + for expert_idx, (q_proj, k_proj, v_proj) in enumerate( + zip(self.q_proj_experts, self.k_proj_experts, self.v_proj_experts) + ): + start, end = start_indices[expert_idx], end_indices[expert_idx] + if start == end: + continue + + dim_input = self.dim_inputs[expert_idx] + expert_input = hidden_states[start:end, :dim_input] + + # Compute Q/K/V + q_out = q_proj(expert_input) + k_out = k_proj(expert_input) + v_out = v_proj(expert_input) + + if getattr(self.config, "model_type", None) == "qwen3_vl_text": + q_out = self.q_norms[expert_idx](q_out) + q_out = q_out[0] if isinstance(q_out, (tuple, list)) else q_out + k_out = self.k_norms[expert_idx](k_out) + k_out = k_out[0] if isinstance(k_out, (tuple, list)) else k_out + + q_buffer[start:end] = q_out + k_buffer[start:end] = k_out + v_buffer[start:end] = v_out + + # === Restore tokens to the original order === + # unpermute (using the same unpermute operation) + q_unpermuted = ops.unpermute(q_buffer, row_id_map, probs) + k_unpermuted = ops.unpermute(k_buffer, row_id_map, probs) + v_unpermuted = ops.unpermute(v_buffer, row_id_map, probs) + + # === Reshape to final form === + query_states = q_unpermuted.view( + batch_size, seq_length, self.num_heads, self.head_dim + ) + key_states = k_unpermuted.view( + batch_size, seq_length, self.num_key_value_heads, self.head_dim + ) + value_states = v_unpermuted.view( + batch_size, seq_length, self.num_key_value_heads, self.head_dim + ) + + return query_states, key_states, value_states + + def _apply_rotary_pos_embed( + self, query_states, key_states, cos, sin, unsqueeze_dim=1 + ): + if self.config.model_type == "qwen2_5_vl": + query_states, key_states = apply_multimodal_rotary_pos_emb( + query_states.contiguous(), + key_states.contiguous(), + cos.contiguous(), + sin.contiguous(), + self.rope_scaling["mrope_section"], + unsqueeze_dim, + ) + else: + raise NotImplementedError( + f"Unsupported model type: {self.config.model_type}" + ) + return query_states, key_states + + def _generate_output(self, attn_output, masks): + output = torch.zeros( + *attn_output.shape[:2], + self.hidden_size, + device=attn_output.device, + dtype=attn_output.dtype, + ) + for expert_idx, (o_proj, mask) in enumerate(zip(self.o_proj_experts, masks)): + if not mask.any(): + continue + dim_input = self.dim_inputs[expert_idx] + + # Obtain all necessary indexes in a single operation. + mask_indices = mask.nonzero(as_tuple=False) + if mask_indices.numel() == 0: + continue + + batch_indices = mask_indices[:, 0] + seq_indices = mask_indices[:, 1] + + # Use advanced indexing directly to avoid intermediate tensors. + selected_attn_output = attn_output[batch_indices, seq_indices] + projected_output = o_proj(selected_attn_output) + + output[batch_indices, seq_indices, :dim_input] = projected_output + + return output + + def _generate_output_mot_opt( + self, + attn_output: torch.Tensor, + experts_indices: torch.Tensor, + start_indices: torch.Tensor, + end_indices: torch.Tensor, + ) -> torch.Tensor: + """ + Expert-sharded version of attn_output processing based on start_indices / end_indices. + Rearranges the [B, S, H] attn_output according to expert order (permute), + applies the o_proj projection for each expert individually, + and keeps the final output in expert order ([Tokens, Hidden]) + instead of restoring it back to [B, S, H]. + + Args: + attn_output: [B, S, hidden_dim] + experts_indices: [B, S], expert index for each token + start_indices, end_indices: start and end token indices for each expert + (in the permuted token space) + + Returns: + output_buffer: [TotalTokens, hidden_dim], arranged in expert order + """ + + _, _, hidden_dim = attn_output.shape + device, dtype = attn_output.device, attn_output.dtype + + # === 1. Flatten and reorder by expert assignment === + flat_attn_output = attn_output.view(-1, hidden_dim) # [B*S, H] + flat_expert_indices = experts_indices.reshape(-1) # [B*S] + permuted_inputs, _ = ops.permute(flat_attn_output, flat_expert_indices) + total_tokens = permuted_inputs.shape[0] + + # === 2. Initialize output buffer (still in permuted token space) === + output_buffer = torch.zeros( + total_tokens, hidden_dim, device=device, dtype=dtype + ) + + # === 3. Each expert processes its own token segment independently === + for expert_idx, o_proj in enumerate(self.o_proj_experts): + start, end = start_indices[expert_idx], end_indices[expert_idx] + if start == end: + continue + + dim_input = self.dim_inputs[expert_idx] + expert_input = permuted_inputs[start:end] # [N_e, dim_input] + expert_output = o_proj(expert_input) # [N_e, hidden_dim] + + # Write results into the buffer (overwrite only valid dimension region) + output_buffer[start:end, :dim_input] = expert_output[:, :dim_input] + + # === 4. Return the output ordered by expert sequence === + return output_buffer + + +class JointQwen2VLFlashAttention(JointQwen2VLAttention): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + token_types: Optional[torch.LongTensor] = None, + position_embeddings: Optional[ + Tuple[torch.Tensor, torch.Tensor] + ] = None, # necessary, but kept here for BC + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + probs: Optional[torch.Tensor] = None, + row_id_map: Optional[torch.Tensor] = None, + orig_shape: Optional[Tuple[int]] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + + if token_types is None: + raise ValueError("token_types cannot be empty") + + if self.config.mot_opt: + bsz, q_len, _ = orig_shape + query_states, key_states, value_states = self._generate_qkv_mot_opt( + hidden_states, + token_types, + start_indices, + end_indices, + probs, + row_id_map, + bsz, + q_len, + ) + else: + bsz, q_len, _ = hidden_states.size() + masks = [ + (token_types == expert_idx) + for expert_idx in range(len(self.dim_inputs)) + ] + query_states, key_states, value_states = self._generate_qkv( + hidden_states, masks + ) + + # Because the input can be padded, the absolute sequence length depends on the max position id. + cos, sin = position_embeddings + query_states, key_states = self._apply_rotary_pos_embed( + query_states, key_states, cos, sin, unsqueeze_dim=2 + ) + + if past_key_value is not None: + cache_kwargs = { + "sin": sin, + "cos": cos, + "cache_position": cache_position, + } # Specific to RoPE models + key_states, value_states = past_key_value.update( + key_states.transpose(1, 2), + value_states.transpose(1, 2), + self.layer_idx, + cache_kwargs, + ) + key_states, value_states = key_states.transpose( + 1, 2 + ), value_states.transpose(1, 2) + + dropout_rate = 0.0 if not self.training else self.attention_dropout + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in float16 just to be sure everything works as expected. + input_dtype = query_states.dtype + if input_dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self.q_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + attn_output = flash_attn_func( + query_states, + key_states, + value_states, + dropout_rate, + softmax_scale=None, + causal=self.is_causal, + ) + + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() + + if self.config.mot_opt: + output = self._generate_output_mot_opt( + attn_output, token_types, start_indices, end_indices + ) + else: + output = self._generate_output(attn_output, masks) + + return output, None, past_key_value + + +JOINT_QWEN_ATTENTION_CLASSES = { + "eager": JointQwen2VLAttention, + "flash_attention_2": JointQwen2VLFlashAttention, + "sdpa": JointQwen2VLAttention, +} diff --git a/wall_x/model/model_utils.py b/wall_x/model/model_utils.py new file mode 100644 index 0000000..5aaeecc --- /dev/null +++ b/wall_x/model/model_utils.py @@ -0,0 +1,319 @@ +import torch +import os +import numpy as np +from transformers import AutoProcessor +from wall_x.model.action_head import Normalizer + + +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 load_wallx_processors(config): + processor = AutoProcessor.from_pretrained(config["processor_path"], use_fast=True) + # pad side = left + processor.tokenizer.padding_side = "left" + + new_tokens = ["<|propri|>", "<|action|>"] + # special_tokens = [] + action_tokenizer_type = config.get("action_tokenizer_type", None) + if action_tokenizer_type == "fast": + train_action_tokenizer = AutoProcessor.from_pretrained( + config["action_tokenizer_path"], trust_remote_code=True + ) + val_action_tokenizer = AutoProcessor.from_pretrained( + config["action_tokenizer_path"], trust_remote_code=True + ) + new_tokens += [ + f"<|action_token_{i}|>" for i in range(train_action_tokenizer.vocab_size) + ] + elif action_tokenizer_type == "spatialvla": + raise NotImplementedError("SpatialActionTokenizer is not implemented") + else: + train_action_tokenizer = None + val_action_tokenizer = None + + num_added_tokens = processor.tokenizer.add_tokens(new_tokens) + + if action_tokenizer_type and train_action_tokenizer.vocab_size > 0: + action_mapper = {} + for i in range(train_action_tokenizer.vocab_size): + token = f"<|action_token_{i}|>" + token_id = processor.tokenizer.convert_tokens_to_ids(token) + action_mapper[token_id] = i + else: + action_mapper = None + + return { + "processor": processor, + "train_action_tokenizer": train_action_tokenizer, + "val_action_tokenizer": val_action_tokenizer, + "action_mapper": action_mapper, + "num_added_tokens": num_added_tokens, + } + + +def register_normalizers(config, model_path): + # if config.get("customized_action_statistic_dof", None): + # action_statistic_dof = json.load(open(config["customized_action_statistic_dof"], "r")) + # else: + # action_statistic_dof = default_action_statistic_dof + + action_statistic_dof = None + + if os.path.exists(model_path + "/normalizer_action.pth"): + print( + "Loading normalizer_action from checkpoint", + model_path + "/normalizer_action.pth", + flush=True, + ) + normalizer_action = Normalizer.from_ckpt(model_path + "/normalizer_action.pth") + else: + normalizer_action = Normalizer( + action_statistic_dof, + config["dof_config"], + min_key=config.get("min_key", "min"), + delta_key=config.get("delta_key", "delta"), + ) + + # print("action_statistic_dof",action_statistic_dof) + + if os.path.exists(model_path + "/normalizer_propri.pth"): + print( + "Loading normalizer_propri from checkpoint", + model_path + "/normalizer_propri.pth", + flush=True, + ) + normalizer_propri = Normalizer.from_ckpt(model_path + "/normalizer_propri.pth") + else: + normalizer_propri = Normalizer( + action_statistic_dof, + config["agent_pos_config"], + min_key=config.get("min_key", "min"), + delta_key=config.get("delta_key", "delta"), + ) + + return normalizer_action, normalizer_propri + + +def find_first_last_ones(tensor): + """ + Input: a tensor of shape (bs, seq_len) containing 0s and 1s + Output: (first_indices, last_indices), each of shape (bs,) + where first_indices[i] is the index of the first 1 in the i-th batch, or -1 if none exists. + last_indices[i] is the index of the last 1 in the i-th batch, or -1 if none exists. + """ + bs, seq_len = tensor.shape + masks = tensor == 1 + has_ones = masks.any(dim=1) + + first = torch.full((bs,), -1, dtype=torch.long, device=tensor.device) + last = first.clone() + + first[has_ones] = torch.argmax(masks[has_ones].float(), dim=1) + + flipped_masks = masks.flip(dims=[1]) + last_argmax = torch.argmax(flipped_masks[has_ones].float(), dim=1) + last[has_ones] = seq_len - 1 - last_argmax + + return first, last + + +def flashmask_to_densemask(startend_row_indices, dtype, causal=True): + if startend_row_indices is None: + return None + bz, num_head, seq_len, bound_num = startend_row_indices.shape + m = np.ones((bz, num_head, seq_len, seq_len), dtype=dtype) + has_end = (causal and bound_num == 2) or ((not causal) and bound_num == 4) + for bi in range(bz): + for hi in range(num_head): + for j in range(seq_len): + downstart = startend_row_indices[bi, hi, j, 0] + if has_end: + downend = startend_row_indices[bi, hi, j, 1] + m[bi, hi, downstart:downend, j] = 0 + else: + m[bi, hi, downstart:, j] = 0 + if causal: + m[bi, hi, :j, j] = 0 + else: + if has_end: + upstart = startend_row_indices[bi, hi, j, 2] + upend = startend_row_indices[bi, hi, j, 3] + m[bi, hi, upstart:upend, j] = 0 + else: + upend = startend_row_indices[bi, hi, j, 1] + m[bi, hi, :upend, j] = 0 + return m + + +def num_floating_point_operations( + args, + batch_size: int, + num_lang_tokens: int, + num_action_tokens: int, + vision_seq_length: int = 756, +): + """ + Accurately estimate the training FLOPs of Transformer + MoE + MoT + Vision. + + Supported: + - expert0 = language tokens + - expert1 = action tokens + - MoE MLP (2 experts) + - MoT Attention (2 experts) + - GQA + - Vision Transformer (full+window attention) + """ + assert args.num_experts == 2, "The current model only supports 2 experts." + + dim_lang, dim_act = args.dim_inputs + + # Number of tokens per layer (flattened across batch) + N_lang = batch_size * num_lang_tokens + N_action = batch_size * num_action_tokens + N_total = N_lang + N_action # Used for non-MoT attention + + # ================================================================ + # Text MLP FLOPs + # ================================================================ + hidden_size = args.hidden_size + ffn_hidden_size = args.intermediate_size + num_layers = args.num_hidden_layers + + use_moe_mlp = getattr(args, "mlp_moe", False) + + # ---------- Forward-only MLP FLOPs ---------- + def forward_mlp_flops(N, d_in, d_ff): + """ + SwiGLU forward: + gate = x @ W1 (2*N*d_in*d_ff) + up = x @ W2 (2*N*d_in*d_ff) + act = silu + mul (~2*N*d_ff) + down = h @ W3 (2*N*d_ff*d_in) + + Forward ≈ 4*N*d_in*d_ff + 2*N*d_ff*d_in = 6*N*d_in*d_ff + 2*N*d_ff + """ + return 6 * N * d_in * d_ff + 2 * N * d_ff + + if not use_moe_mlp: + # Dense MLP + F_fwd = forward_mlp_flops(N_total, hidden_size, ffn_hidden_size) + total_mlp_flops_text = 3 * num_layers * F_fwd # <-- training FLOPs + else: + # MoE: expert0(language) + expert1(action) + hid_lang = args.experts[0]["intermediate_size"] + hid_act = args.experts[1]["intermediate_size"] + + F_lang_fwd = forward_mlp_flops(N_lang, dim_lang, hid_lang) + F_act_fwd = forward_mlp_flops(N_action, dim_act, hid_act) + + total_mlp_flops_text = 3 * num_layers * (F_lang_fwd + F_act_fwd) + + # ================================================================ + # Text Attention FLOPs + # ================================================================ + num_heads = args.num_attention_heads + num_kv = args.num_key_value_heads + H = hidden_size + B = batch_size + S = num_lang_tokens + num_action_tokens + N = B * S + + use_mot = getattr(args, "attention_moe", False) + + # ---------- attention matmul ---------- + F_matmul_fwd = 4 * B * (S**2) * H + + if not use_mot: + # -------- GQA + QKV / O -------- + F_q_fwd = 2 * N * H * H + F_kv_fwd = 4 * N * H * H * (num_kv / num_heads) + F_o_fwd = 2 * N * H * H + + F_attn_fwd = F_q_fwd + F_kv_fwd + F_o_fwd + F_matmul_fwd + + else: + # -------- MoT: expert0 + expert1 QKV -------- + F_lang_qkv = N_lang * dim_lang * H * (2 + 4 * num_kv / num_heads) + F_act_qkv = N_action * dim_act * H * (2 + 4 * num_kv / num_heads) + F_attn_fwd = F_lang_qkv + F_act_qkv + F_matmul_fwd + + # Training FLOPs + total_attn_flops_text = 3 * num_layers * F_attn_fwd + + # ================================================================ + # Logits projection FLOPs + # ================================================================ + vocab_size = getattr(args, "padded_vocab_size", args.vocab_size) + + F_logits_fwd = 2 * N * H * vocab_size + total_logits_flops = 3 * F_logits_fwd + + total_text_flops = total_mlp_flops_text + total_attn_flops_text + total_logits_flops + + # ================================================================ + # Vision Transformer FLOPs + # ================================================================ + total_vision_flops = 0 + + if hasattr(args, "vision_config") and vision_seq_length is not None: + vcfg = args.vision_config + + Bv = batch_size + Sv = vision_seq_length + Nv = Bv * Sv + + Hv = vcfg.hidden_size + Iv = vcfg.intermediate_size + num_heads_v = vcfg.num_heads + window_size = vcfg.window_size + out_hidden = vcfg.out_hidden_size + + depth_v = vcfg.depth + fullatt = set(vcfg.fullatt_block_indexes) + num_full = len(fullatt) + num_local = depth_v - num_full + + # ---------- forward FLOPs ---------- + def forward_vit_mlp(N, H, Inner): + return 6 * N * H * Inner + 2 * N * Inner + + F_mlp_v = forward_vit_mlp(Nv, Hv, Iv) + F_qkv_v = 6 * Nv * Hv * Hv + F_o_v = 2 * Nv * Hv * Hv + + F_full = 4 * Bv * (Sv**2) * Hv + num_windows = Sv / window_size + F_win = 4 * Bv * num_windows * (window_size**2) * (Hv / num_heads_v) + + F_block_full_fwd = F_mlp_v + F_qkv_v + F_o_v + F_full + F_block_local_fwd = F_mlp_v + F_qkv_v + F_o_v + F_win + + # ---------- train FLOPs ---------- + total_vision_flops = 3 * ( + num_full * F_block_full_fwd + num_local * F_block_local_fwd + ) + + # merger + total_vision_flops += 3 * (2 * Nv * Hv * out_hidden) + + # ================================================================ + # TOTAL TRAIN FLOPs + # ================================================================ + return total_text_flops + total_vision_flops diff --git a/wall_x/model/qwen2_5_based/configuration_qwen2_5_vl.py b/wall_x/model/qwen2_5_based/configuration_qwen2_5_vl.py index 439a369..a2d51a1 100644 --- a/wall_x/model/qwen2_5_based/configuration_qwen2_5_vl.py +++ b/wall_x/model/qwen2_5_based/configuration_qwen2_5_vl.py @@ -21,6 +21,9 @@ class Qwen2_5_VLVisionConfig(PretrainedConfig): window_size=112, out_hidden_size=3584, fullatt_block_indexes=[7, 15, 23, 31], + initializer_range=0.02, + _attn_implementation="flash_attention_2", + attn_deterministic=False, **kwargs, ): super().__init__(**kwargs) @@ -38,6 +41,9 @@ class Qwen2_5_VLVisionConfig(PretrainedConfig): self.window_size = window_size self.fullatt_block_indexes = fullatt_block_indexes self.out_hidden_size = out_hidden_size + self.initializer_range = initializer_range + self._attn_implementation = _attn_implementation + self.attn_deterministic = attn_deterministic class Qwen2_5_VLConfig(PretrainedConfig): @@ -169,6 +175,8 @@ class Qwen2_5_VLConfig(PretrainedConfig): self, vocab_size=152064, hidden_size=8192, + action_hidden_size=2048, + state_hidden_size=2048, intermediate_size=29568, num_hidden_layers=80, num_attention_heads=64, @@ -193,16 +201,25 @@ class Qwen2_5_VLConfig(PretrainedConfig): dim_inputs=(1536, 1536), attention_moe=False, mlp_moe=False, + norm_moe=False, + mot_opt=False, + flow_loss_weight=10, + use_state_string_representation=False, + use_adarms=False, + proj_with_mask=True, + use_flow_action_expert=True, + adarms_cond_dim=None, + action_horizon_flow=32, + causal_action_attention_mask=False, + use_x_pred=False, + attn_deterministic=False, **kwargs, ): - if isinstance(vision_config, dict): - self.vision_config = self.sub_configs["vision_config"](**vision_config) - elif vision_config is None: - self.vision_config = self.sub_configs["vision_config"]() - self.vocab_size = vocab_size self.max_position_embeddings = max_position_embeddings self.hidden_size = hidden_size + self.action_hidden_size = action_hidden_size + self.state_hidden_size = state_hidden_size self.intermediate_size = intermediate_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads @@ -230,6 +247,19 @@ class Qwen2_5_VLConfig(PretrainedConfig): self.dim_inputs = tuple(dim_inputs) self.attention_moe = attention_moe self.mlp_moe = mlp_moe + self.norm_moe = norm_moe + self.mot_opt = mot_opt + self.flow_loss_weight = flow_loss_weight + + self.use_state_string_representation = use_state_string_representation + self.use_adarms = use_adarms + self.adarms_cond_dim = adarms_cond_dim + self.proj_with_mask = proj_with_mask + self.use_flow_action_expert = use_flow_action_expert + self.action_horizon_flow = action_horizon_flow + self.causal_action_attention_mask = causal_action_attention_mask + self.use_x_pred = use_x_pred + self.attn_deterministic = attn_deterministic # Validate the correctness of rotary position embeddings parameters # BC: if there is a 'type' field, move it to 'rope_type'. @@ -244,5 +274,12 @@ class Qwen2_5_VLConfig(PretrainedConfig): super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + # move vision config initialization after super init to avoid recursively set in latest transformers version + # TODO: make it better + if isinstance(vision_config, dict): + self.vision_config = self.sub_configs["vision_config"](**vision_config) + elif vision_config is None: + self.vision_config = self.sub_configs["vision_config"]() + __all__ = ["Qwen2_5_VLConfig"] diff --git a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py b/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py index d09ad1b..951bed9 100644 --- a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py +++ b/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py @@ -120,23 +120,68 @@ class Qwen2_5_VisionRotaryEmbedding(nn.Module): class Qwen2RMSNorm(nn.Module): - def __init__(self, hidden_size, eps=1e-6): + def __init__( + self, hidden_size: int, eps: float = 1e-6, cond_dim: Optional[int] = None + ): """ - Qwen2RMSNorm is equivalent to T5LayerNorm + Qwen2RMSNorm with optional conditional input support, equivalent to T5LayerNorm """ super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps + self.hidden_size = hidden_size + self.cond_dim = cond_dim - def forward(self, hidden_states): + # Dense layer for adaptive normalization (if cond_dim is provided) + if cond_dim is not None: + self.dense = nn.Linear(cond_dim, hidden_size * 3, bias=True) + nn.init.zeros_(self.dense.weight) + else: + self.dense = None + self.weight = nn.Parameter(torch.ones(hidden_size)) + + def _norm(self, x): + # Compute variance in float32 for numerical stability + variance = x.pow(2).mean(-1, keepdim=True) + # Compute normalization + normed_inputs = x * torch.rsqrt(variance + self.variance_epsilon) + return normed_inputs + + def forward(self, hidden_states, cond=None): input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) - variance = hidden_states.pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) - return self.weight * hidden_states.to(input_dtype) + normed_inputs = self._norm(hidden_states) + + if cond is None or self.dense is None: + # Regular RMSNorm + normed_inputs = self.weight * normed_inputs + return normed_inputs.to(input_dtype), None + + # Adaptive RMSNorm + if cond.shape[-1] != self.cond_dim: + raise ValueError( + f"Expected cond dimension {self.cond_dim}, got {cond.shape[-1]}" + ) + + # Compute modulation parameters + cond = cond.to(dtype=self.dense.weight.dtype) + modulation = self.dense(cond) + if len(hidden_states.shape) == 3: # [batch, seq, features] + modulation = modulation.unsqueeze(1) + + scale, shift, gate = torch.chunk(modulation, 3, dim=-1) + + # Apply adaptive normalization + normed_inputs = normed_inputs * (1 + scale.to(torch.float32)) + shift.to( + torch.float32 + ) + + return normed_inputs.to(input_dtype), gate.to(input_dtype) def extra_repr(self): - return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + repr_str = f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + if self.dense is not None: + repr_str += f", adaptive=True, cond_dim={self.cond_dim}" + return repr_str class Qwen2_5_VLPatchMerger(nn.Module): @@ -151,7 +196,7 @@ class Qwen2_5_VLPatchMerger(nn.Module): ) def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.mlp(self.ln_q(x).view(-1, self.hidden_size)) + x = self.mlp(self.ln_q(x)[0].view(-1, self.hidden_size)) return x @@ -381,13 +426,13 @@ class Qwen2_5_VLVisionBlock(nn.Module): position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor: hidden_states = hidden_states + self.attn( - self.norm1(hidden_states), + self.norm1(hidden_states)[0], cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, rotary_pos_emb=rotary_pos_emb, position_embeddings=position_embeddings, ) - hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)[0]) return hidden_states @@ -1082,16 +1127,41 @@ class Qwen2_5_VLSdpaAttention(Qwen2_5_VLAttention): "cos": cos, "cache_position": cache_position, } # Specific to RoPE models - key_states, value_states = past_key_value.update( - key_states, value_states, self.layer_idx, cache_kwargs - ) + if use_cache: + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + else: + past_key_states, past_value_states = past_key_value[self.layer_idx] + key_states = torch.cat([past_key_states, key_states], dim=-2) + value_states = torch.cat([past_value_states, value_states], dim=-2) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) causal_mask = attention_mask if attention_mask is not None: # no matter the length, we just slice it - causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + # Ensure the attention_mask correctly matches the head dimension + if len(attention_mask.shape) == 2: # [batch_size, seq_len] + # Expand to [batch_size, 1, seq_len, seq_len] causal mask format + bsz, seq_len = attention_mask.shape + causal_mask = attention_mask.view(bsz, 1, 1, seq_len).expand( + bsz, 1, seq_len, seq_len + ) + elif len(attention_mask.shape) == 3: # [batch_size, seq_len, seq_len] + # Add head dimension: [batch_size, 1, seq_len, seq_len] + causal_mask = attention_mask.unsqueeze(1) + elif ( + len(attention_mask.shape) == 4 + ): # [batch_size, num_heads, seq_len, seq_len] + causal_mask = attention_mask + else: + raise ValueError( + f"Unsupported attention_mask dim: {attention_mask.shape}" + ) + + # Convert the attention mask to boolean type + causal_mask = causal_mask.to(torch.bool) # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, # Reference: https://github.com/pytorch/pytorch/issues/112577. @@ -1105,6 +1175,18 @@ class Qwen2_5_VLSdpaAttention(Qwen2_5_VLAttention): # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1. is_causal = True if causal_mask is None and q_len > 1 else False + if q_len == 1: + is_causal = False + causal_mask = torch.ones( + bsz, + 1, + 1, + key_states.shape[2], + device=hidden_states.device, + dtype=hidden_states.dtype, + ).contiguous() + causal_mask = causal_mask.to(torch.bool) + attn_output = torch.nn.functional.scaled_dot_product_attention( query_states, key_states, diff --git a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py b/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py index a01c198..996f879 100644 --- a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py +++ b/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py @@ -1,5 +1,6 @@ import os import torch +import yaml import numpy as np import glob import torch.nn as nn @@ -9,9 +10,8 @@ from torch.nn import CrossEntropyLoss from safetensors.torch import load_file from peft import LoraConfig, get_peft_model from typing import Optional, List, Tuple, Any, Dict, Union - +import time from transformers import AutoConfig, AutoProcessor -from transformers.activations import ACT2FN from transformers.utils import logging, is_torchdynamo_compiling from transformers.cache_utils import ( Cache, @@ -20,13 +20,12 @@ from transformers.cache_utils import ( StaticCache, ) from transformers.modeling_attn_mask_utils import AttentionMaskConverter -from transformers.models.qwen2_vl.modeling_qwen2_vl import ( - Qwen2RMSNorm, -) -from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( + +from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl import ( Qwen2_5_VLMLP, Qwen2_5_VLRotaryEmbedding, Qwen2_5_VLPreTrainedModel, + Qwen2RMSNorm, Qwen2_5_VLForConditionalGeneration, ) from transformers.modeling_outputs import ( @@ -37,13 +36,20 @@ from transformers.modeling_outputs import ( from wall_x.fusions import ops from wall_x.model.action_head import ActionProcessor from wall_x.model.qwen2_5_based.configuration_qwen2_5_vl import Qwen2_5_VLConfig +from wall_x.model.model_utils import load_wallx_processors, update_model_config from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl import ( Qwen2_5_VisionTransformerPretrainedModel, Qwen2_5_VLAttention, Qwen2_5_VLFlashAttention2, Qwen2_5_VLSdpaAttention, ) -from wall_x.data.config import ACTION_DATASET_NAMES, MULTIMODAL_DATASET_NAMES +from wall_x.model.vla_mixin import ActionGenerationMixin, ActionModelMixMin +from wall_x.model.vla_mixin import TokenTypeRouter, SparseMoeBlock +from wall_x.model.vla_mixin import ( + ATTENTION_TYPES_WITH_2D_MASK, +) +from wall_x.model.joint_attention import JOINT_QWEN_ATTENTION_CLASSES + from wall_x.data.utils import update_action_statistics from wall_x.utils.constant import action_statistic_dof from pprint import pprint @@ -66,104 +72,6 @@ class Qwen2_5_VLACausalLMOutputWithPast(ModelOutput): channel_loss_count_dict: Optional[dict[torch.FloatTensor]] = None -class BlockSparseMLP(nn.Module): - def __init__(self, config): - super().__init__() - - self.hidden_size = config["hidden_size"] - self.intermediate_size = config["intermediate_size"] - self.hidden_act = config["hidden_act"] - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) - self.act_fn = ACT2FN[self.hidden_act] - - def forward(self, hidden_state): - return self.down_proj( - self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state) - ) - - -class SparseMoeBlock(nn.Module): - """Sparse Mixture of Experts (MoE) Block optimized with Grouped GEMM. - - This module implements a sparse MoE layer where tokens are dynamically routed - to different expert networks. Uses grouped GEMM operations for efficient - computation across multiple experts. - - Args: - config: Configuration object containing expert specifications - num_experts: Number of expert networks in the MoE block - """ - - def __init__(self, config, num_experts: int): - super().__init__() - self.num_experts = num_experts - # Initialize expert networks based on configuration - self.experts = nn.ModuleList( - [BlockSparseMLP(config.experts[i]) for i in range(num_experts)] - ) - - def forward( - self, - hidden_states: torch.Tensor, - experts_indices: torch.Tensor, - start_indices: torch.Tensor, - end_indices: torch.Tensor, - ) -> torch.Tensor: - """Forward pass through the Sparse MoE block. - - Routes different hidden states to their corresponding expert networks - for processing. Uses efficient grouped operations to minimize overhead. - - Args: - hidden_states: Input tensor of shape (batch_size, seq_length, hidden_dim) - experts_indices: Expert assignment indices of shape (batch_size, seq_length) - indicating which expert each token should be routed to - start_indices: Starting indices for each expert's assigned tokens - end_indices: Ending indices for each expert's assigned tokens - - Returns: - output: Processed tensor of shape (batch_size, seq_length, hidden_dim) - after expert processing and token reordering - """ - batch_size, seq_length, hidden_dim = hidden_states.size() - - # Flatten inputs for efficient grouped processing - hidden_states = hidden_states.view(-1, hidden_dim) # [total_tokens, hidden_dim] - experts_indices = experts_indices.view(-1) # [total_tokens] - - # Create uniform probabilities for all tokens (can be modified for weighted routing) - probs = torch.ones_like(experts_indices, dtype=torch.float32).view(-1, 1) - - # Permute inputs to group tokens by expert assignment - permuted_inputs, row_id_map = ops.permute(hidden_states, experts_indices) - final_output = torch.zeros_like(permuted_inputs) - - # Process tokens through their assigned experts - for expert_idx, expert in enumerate(self.experts): - # Skip experts with no assigned tokens - if start_indices[expert_idx] == end_indices[expert_idx]: - continue - - # Process tokens assigned to this expert - expert_input = permuted_inputs[ - start_indices[expert_idx] : end_indices[expert_idx] - ] - expert_output = expert(expert_input) - final_output[start_indices[expert_idx] : end_indices[expert_idx]] = ( - expert_output - ) - - # Restore original token ordering - unpermuted_outputs = ops.unpermute(final_output, row_id_map, probs) - - # Reshape back to original dimensions - output = unpermuted_outputs.view(batch_size, seq_length, hidden_dim) - - return output - - QWEN2_5_VL_ATTENTION_CLASSES = { "eager": Qwen2_5_VLAttention, "flash_attention_2": Qwen2_5_VLFlashAttention2, @@ -171,10 +79,17 @@ QWEN2_5_VL_ATTENTION_CLASSES = { } -class Qwen2_5_VLDecoderLayer_with_MoE(nn.Module): - def __init__(self, config: Qwen2_5_VLConfig, layer_idx: int, num_experts: int): +class Qwen2_5_VLDecoderLayer_with_MoE(nn.Module, ActionModelMixMin): + def __init__( + self, + config: Qwen2_5_VLConfig, + layer_idx: int, + num_experts: int, + use_selective_recompute: bool = True, + ): super().__init__() self.hidden_size = config.hidden_size + self.use_selective_recompute = use_selective_recompute if ( config.use_sliding_window @@ -184,21 +99,64 @@ class Qwen2_5_VLDecoderLayer_with_MoE(nn.Module): f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; " "unexpected results may be encountered." ) + if config.attention_moe: + self.self_attn = JOINT_QWEN_ATTENTION_CLASSES[config._attn_implementation]( + config, layer_idx + ) + else: + self.self_attn = QWEN2_5_VL_ATTENTION_CLASSES[config._attn_implementation]( + config, layer_idx + ) - self.self_attn = QWEN2_5_VL_ATTENTION_CLASSES[config._attn_implementation]( - config, layer_idx - ) + if config.use_adarms: + adarms_cond_dims = [None, config.adarms_cond_dim] + else: + adarms_cond_dims = [None, None] - self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = Qwen2RMSNorm( - config.hidden_size, eps=config.rms_norm_eps - ) + if config.norm_moe: + self.input_layernorms = nn.ModuleList( + [ + Qwen2RMSNorm( + config.dim_inputs[i], + eps=config.rms_norm_eps, + cond_dim=adarms_cond_dims[i], + ) + for i in range(num_experts) + ] + ) + self.post_attention_layernorms = nn.ModuleList( + [ + Qwen2RMSNorm( + config.dim_inputs[i], + eps=config.rms_norm_eps, + cond_dim=adarms_cond_dims[i], + ) + for i in range(num_experts) + ] + ) + self.input_layernorm, self.post_attention_layernorm = None, None + else: + self.input_layernorm = Qwen2RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_attention_layernorm = Qwen2RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.input_layernorms, self.post_attention_layernorms = None, None if config.mlp_moe: - self.moe = SparseMoeBlock(config, num_experts=num_experts) + self.router = TokenTypeRouter(num_experts=num_experts) + self.moe = SparseMoeBlock( + config, + num_experts=num_experts, + use_selective_recompute=use_selective_recompute, + ) self.mlp = None else: self.mlp = Qwen2_5_VLMLP(config) + self.moe, self.router = None, None + + self.config = config def forward( self, @@ -206,13 +164,20 @@ class Qwen2_5_VLDecoderLayer_with_MoE(nn.Module): attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Tuple[torch.Tensor]] = None, - token_types=None, - start_indices=None, - end_indices=None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + position_embeddings: Optional[ + Tuple[torch.Tensor, torch.Tensor] + ] = None, # necessary, but kept here for BC + # for vla + token_types: Optional[torch.LongTensor] = None, + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + probs: Optional[torch.Tensor] = None, + row_id_map: Optional[torch.Tensor] = None, + orig_shape: Optional[Tuple[int, int, int]] = None, + adarms_conds: Optional[List[torch.Tensor]] = [None, None], **kwargs, ) -> Tuple[ torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] @@ -239,32 +204,73 @@ class Qwen2_5_VLDecoderLayer_with_MoE(nn.Module): into the model """ residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) + + hidden_states, gate, _ = self._apply_norm_moe( + hidden_states, + token_types, + adarms_conds, + self.input_layernorms, + self.input_layernorm, + start_indices, + end_indices, + self.use_selective_recompute, + ) # Self Attention - hidden_states, self_attn_weights, present_key_value = self.self_attn( - hidden_states=hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_value=past_key_value, - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - position_embeddings=position_embeddings, + if self.config.attention_moe: + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + token_types=token_types, + start_indices=start_indices, + end_indices=end_indices, + probs=probs, + row_id_map=row_id_map, + orig_shape=orig_shape, + position_embeddings=position_embeddings, + ) + else: + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + + hidden_states = self._gated_residual( + residual, hidden_states, gate, start_indices, end_indices ) - hidden_states = residual + hidden_states # Fully Connected residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) - if self.mlp is None: # using moe mlp - hidden_states = self.moe( - hidden_states, token_types, start_indices, end_indices - ) - else: - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states + hidden_states, gate, gate_mask = self._apply_norm_moe( + hidden_states, + token_types, + adarms_conds, + self.post_attention_layernorms, + self.post_attention_layernorm, + start_indices, + end_indices, + self.use_selective_recompute, + ) + + hidden_states = self._apply_mlp_moe( + hidden_states, token_types, start_indices, end_indices + ) + + hidden_states = self._gated_residual( + residual, hidden_states, gate, start_indices, end_indices + ) outputs = (hidden_states,) @@ -275,96 +281,81 @@ class Qwen2_5_VLDecoderLayer_with_MoE(nn.Module): return outputs -class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): - """Qwen2.5-VL model with Mixture of Experts (MoE) architecture. - - This model extends the base Qwen2.5-VL model by incorporating MoE layers - for improved scalability and specialization across different token types. - """ - +class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel, ActionModelMixMin): @classmethod def from_pretrained( - cls, - pretrained_model_name_or_path: str, - num_experts: Optional[int] = None, - *args, - **kwargs, + cls, pretrained_model_name_or_path, num_experts=None, *args, **kwargs ): - """Load a pretrained model with optional MoE configuration. - - Args: - pretrained_model_name_or_path: Path or name of the pretrained model - num_experts: Number of experts for MoE layers (if not in config) - *args: Additional arguments passed to parent class - **kwargs: Additional keyword arguments passed to parent class - - Returns: - Initialized model instance with MoE configuration - """ + # If `num_experts` is provided, ensure it is added to the config. config = kwargs.get("config", None) if config is None: config = AutoConfig.from_pretrained(pretrained_model_name_or_path) - # Override number of experts if specified if num_experts is not None: config.num_experts = num_experts kwargs["config"] = config return super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs) - def __init__(self, config: Qwen2_5_VLConfig): - """Initialize the Qwen2.5-VL MoE model. - - Args: - config: Model configuration containing architecture parameters - """ + def __init__(self, config: Qwen2_5_VLConfig, use_selective_recompute=False): super().__init__(config) - - # Basic model parameters + self.config = config + self.use_selective_recompute = use_selective_recompute self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size - # Model components self.embed_tokens = nn.Embedding( config.vocab_size, config.hidden_size, self.padding_idx ) - - # Decoder layers with MoE support self.layers = nn.ModuleList( [ - Qwen2_5_VLDecoderLayer_with_MoE(config, layer_idx, config.num_experts) + Qwen2_5_VLDecoderLayer_with_MoE( + config, + layer_idx, + config.num_experts, + use_selective_recompute=use_selective_recompute, + ) for layer_idx in range(config.num_hidden_layers) ] ) - - # Model configuration self._attn_implementation = config._attn_implementation - self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) - self.gradient_checkpointing = False + if config.use_adarms: + adarms_cond_dims = [None, config.adarms_cond_dim] + else: + adarms_cond_dims = [None, None] + + if config.norm_moe: + self.norms = nn.ModuleList( + [ + Qwen2RMSNorm( + config.dim_inputs[i], + eps=config.rms_norm_eps, + cond_dim=adarms_cond_dims[i], + ) + for i in range(config.num_experts) + ] + ) + self.norm = None + else: + self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.norms = None + + self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) + + self.gradient_checkpointing = False # Initialize weights and apply final processing self.post_init() - def get_input_embeddings(self) -> nn.Embedding: - """Get the input embedding layer. - - Returns: - The token embedding layer - """ + def get_input_embeddings(self): return self.embed_tokens - def set_input_embeddings(self, value: nn.Embedding) -> None: - """Set the input embedding layer. - - Args: - value: New embedding layer to use - """ + def set_input_embeddings(self, value): self.embed_tokens = value def forward( self, - input_ids: Optional[torch.LongTensor] = None, + input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -372,14 +363,15 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): moe_token_types: Optional[torch.LongTensor] = None, start_indices: Optional[torch.Tensor] = None, end_indices: Optional[torch.Tensor] = None, + positional_masks: Optional[dict] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + adarms_conds: Optional[List[torch.Tensor]] = [None, None], **kwargs, ) -> Union[Tuple, BaseModelOutputWithPast]: - # Set default output options output_attentions = ( output_attentions if output_attentions is not None @@ -391,20 +383,22 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): else self.config.output_hidden_states ) use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) - # Validate inputs if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError( "You must specify exactly one of input_ids or inputs_embeds" ) - if moe_token_types is None: - raise ValueError("moe_token_types must be provided for MoE routing") + raise ValueError("moe_token_types must be provided for MoE routing.") + if start_indices is None or end_indices is None: + raise ValueError( + "start_indices and end_indices must be provided for MoE routing" + ) - # Handle gradient checkpointing compatibility if self.gradient_checkpointing and self.training: if use_cache: logger.warning_once( @@ -412,15 +406,13 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): ) use_cache = False - # Initialize cache if needed + # torch.jit.trace() doesn't support cache objects in the output if use_cache and past_key_values is None and not torch.jit.is_tracing(): past_key_values = DynamicCache() - # Get input embeddings if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) - # Set up cache position if cache_position is None: past_seen_tokens = ( past_key_values.get_seq_length() if past_key_values is not None else 0 @@ -431,7 +423,7 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): device=inputs_embeds.device, ) - # Set up position IDs (hardcoded 3 dimensions for temporal, height, width) + # the hard coded `3` is for temporal, height and width. if position_ids is None: position_ids = cache_position.view(1, 1, -1).expand( 3, inputs_embeds.shape[0], -1 @@ -439,81 +431,142 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): elif position_ids.dim() == 2: position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) - # Create causal attention mask - causal_mask = self._update_causal_mask( - attention_mask, - inputs_embeds, - cache_position, - past_key_values, - output_attentions, - moe_token_types, - ) + if not self.config.attention_moe: + causal_mask = self._update_causal_mask( + attention_mask, + inputs_embeds, + cache_position, + past_key_values, + output_attentions, + moe_token_types, + ) + else: + causal_mask = attention_mask hidden_states = inputs_embeds - # Create position embeddings to be shared across decoder layers + if ( + self.config._attn_implementation != "flash_attention_2" + and self.config.attention_moe is True + ): + position_ids = self._update_position_ids( + position_ids, moe_token_types, positional_masks + ) + + # create position embeddings to be shared across the decoder layers position_embeddings = self.rotary_emb(hidden_states, position_ids) - # Initialize output collections + # If `mot_opt` is enabled, the tokens from different experts will be permuted first, resulting in a dimension of [Tokens, HiddenSize]. + orig_shape = hidden_states.shape + if self.config.mot_opt: + hidden_states = hidden_states.view(-1, hidden_states.size(-1)) + hidden_states, row_id_map = ops.permute( + hidden_states, moe_token_types.view(-1) + ) + else: + row_id_map = None + probs = torch.ones_like(moe_token_types.view(-1), dtype=torch.float32).view( + -1, 1 + ) + + # decoder layers all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None next_decoder_cache = None - # Process through decoder layers + # generate 2d attention mask if needed + if ( + self.config._attn_implementation in ATTENTION_TYPES_WITH_2D_MASK + and self.config.attention_moe is True + ): + if causal_mask is not None and inputs_embeds.shape[1] > 1: + causal_mask = self._update_joint_attention_mask_2d( + attention_mask=causal_mask, + moe_token_types=moe_token_types, + positional_masks=positional_masks, + ) + for decoder_layer in self.layers: if output_hidden_states: + assert ( + self.config.mot_opt is False + ), "When using mot_opt, output_hidden_states is not supported yet." all_hidden_states += (hidden_states,) if self.gradient_checkpointing and self.training: - # Use gradient checkpointing during training layer_outputs = self._gradient_checkpointing_func( decoder_layer.__call__, hidden_states, causal_mask, position_ids, past_key_values, - moe_token_types, output_attentions, use_cache, cache_position, position_embeddings, + # for vla + moe_token_types, + start_indices, + end_indices, + probs, + row_id_map, + orig_shape, + adarms_conds, ) else: - # Regular forward pass layer_outputs = decoder_layer( hidden_states, attention_mask=causal_mask, position_ids=position_ids, past_key_value=past_key_values, - token_types=moe_token_types, - start_indices=start_indices, - end_indices=end_indices, output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, + # for vla + token_types=moe_token_types, + start_indices=start_indices, + end_indices=end_indices, + probs=probs, + row_id_map=row_id_map, + orig_shape=orig_shape, + adarms_conds=adarms_conds, ) - hidden_states = layer_outputs[0] - # Update cache if using it if use_cache: next_decoder_cache = layer_outputs[2 if output_attentions else 1] - # Collect attention weights if requested if output_attentions: + assert ( + self.config.mot_opt is False + ), "When using mot_opt, output_hidden_states is not supported yet." all_self_attns += (layer_outputs[1],) - # Apply final layer normalization - hidden_states = self.norm(hidden_states) + hidden_states, _, _ = self._apply_norm_moe( + hidden_states, + moe_token_types, + adarms_conds, + self.norms, + self.norm, + start_indices, + end_indices, + self.use_selective_recompute, + ) - # Add final hidden states if collecting all states + # add hidden states from the last decoder layer if output_hidden_states: + assert ( + self.config.mot_opt is False + ), "When using mot_opt, output_hidden_states is not supported yet." all_hidden_states += (hidden_states,) next_cache = next_decoder_cache if use_cache else None - # Return outputs in requested format + if self.config.mot_opt: + hidden_states = ops.unpermute(hidden_states, row_id_map, probs) + hidden_states = hidden_states.view(orig_shape) + if not return_dict: return tuple( v @@ -537,71 +590,63 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): output_attentions: bool, moe_token_types: Optional[torch.LongTensor] = None, ): - """Update causal attention mask with support for bidirectional attention for specific token types. - - This method creates and modifies attention masks to support different attention patterns: - - Standard causal (unidirectional) attention for most tokens - - Bidirectional attention for specific token types (e.g., MoE routing tokens) - - Args: - attention_mask: Input attention mask to avoid attending to padding tokens - input_tensor: Input embeddings tensor for shape and device information - cache_position: Position indices for caching mechanisms - past_key_values: Cached key-value pairs from previous forward passes - output_attentions: Whether attention weights will be returned - moe_token_types: Optional tensor indicating token types for MoE routing - (type 1 tokens will use bidirectional attention) - - Returns: - Updated causal attention mask, or None if using Flash Attention 2 - """ - # Flash Attention 2 handles masking internally if self.config._attn_implementation == "flash_attention_2": + if attention_mask is not None and past_key_values is not None: + is_padding_right = ( + attention_mask[:, -1].sum().item() != input_tensor.size()[0] + ) + if is_padding_right: + raise ValueError( + "You are attempting to perform batched generation with padding_side='right'" + " this may lead to unexpected behaviour for Flash Attention version of Qwen2_5_VL. Make sure to " + " call `tokenizer.padding_side = 'left'` before tokenizing the input. " + ) + if attention_mask is not None and 0.0 in attention_mask: + return attention_mask return None - # Calculate sequence lengths for cache management + # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in + # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail + # to infer the attention mask. past_seen_tokens = ( past_key_values.get_seq_length() if past_key_values is not None else 0 ) using_static_cache = isinstance(past_key_values, StaticCache) using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache) - # For SDPA (Scaled Dot Product Attention), use `is_causal` argument when possible - # instead of explicit attention mask to enable Flash Attention 2 dispatch - # Note: This optimization is not compatible with static cache + # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward if ( self.config._attn_implementation == "sdpa" and not (using_static_cache or using_sliding_window_cache) and not output_attentions ): - # Check if we can ignore the causal mask and rely on SDPA's internal handling - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - sliding_window=self.config.sliding_window, - is_training=self.training, - ): - return None + if attention_mask.ndim == 2: + if AttentionMaskConverter._ignore_causal_mask_sdpa( + attention_mask, + inputs_embeds=input_tensor, + past_key_values_length=past_seen_tokens, + sliding_window=self.config.sliding_window, + is_training=self.training, + ): + return None + elif attention_mask.ndim == 3: + return attention_mask - # Extract tensor properties for mask creation dtype, device = input_tensor.dtype, input_tensor.device min_dtype = torch.finfo(dtype).min sequence_length = input_tensor.shape[1] - - # Determine target length based on cache type + # SlidingWindowCache or StaticCache if using_sliding_window_cache or using_static_cache: - # Use maximum cache shape for sliding window or static caches target_length = past_key_values.get_max_cache_shape() + # DynamicCache or no cache else: - # For dynamic cache or no cache, calculate based on attention mask or sequence length target_length = ( attention_mask.shape[-1] if isinstance(attention_mask, torch.Tensor) else past_seen_tokens + sequence_length + 1 ) - # Generate 4D causal attention mask from 2D input mask if provided + # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( attention_mask, sequence_length=sequence_length, @@ -613,41 +658,32 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): config=self.config, past_key_values=past_key_values, ) - - # Modify mask to support bidirectional attention for specific token types + # Modify the mask to support bidirectional attention. if moe_token_types is not None: - # Identify positions of type 1 tokens (MoE routing tokens) + # Find the positions of all tokens of type 1. type1_tokens = ( (moe_token_types == 1).unsqueeze(1).unsqueeze(2) - ) # Shape: [B, 1, 1, S] + ) # [B, 1, 1, S] - # Create bidirectional attention region for type 1 tokens - # This allows type 1 tokens to attend to each other bidirectionally - type1_mask = torch.zeros_like(causal_mask) # Shape: [B, num_heads, S, S] - type1_region = type1_tokens & type1_tokens.transpose( - -1, -2 - ) # Shape: [B, 1, S, S] + # Create a square mask for the type1 region. + type1_mask = torch.zeros_like(causal_mask) # [B, num_heads, S, S] + type1_region = type1_tokens & type1_tokens.transpose(-1, -2) # [B, 1, S, S] type1_mask = type1_mask.masked_fill(type1_region, 1.0).to(torch.bool) - - # Apply bidirectional attention: zero out causal constraints in type 1 regions + # Set the original causal_mask to zero in the type1 region, and then add the type1_mask. causal_mask = torch.where( - type1_mask, # Where type 1 tokens interact with each other - torch.zeros_like( - causal_mask - ), # Remove causal masking (allow bidirectional) - causal_mask, # Keep original causal masking for other regions + type1_mask, + torch.zeros_like(causal_mask), + causal_mask, ) - - # Handle special case for SDPA with CUDA/XPU devices if ( self.config._attn_implementation == "sdpa" and attention_mask is not None and attention_mask.device.type in ["cuda", "xpu"] and not output_attentions ): - # Ensure attention to all tokens in fully masked rows for memory-efficient attention - # This is required for F.scaled_dot_product_attention's memory-efficient path - # when using left padding. See: https://github.com/pytorch/pytorch/issues/110213 + # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when + # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + # Details: https://github.com/pytorch/pytorch/issues/110213 causal_mask = AttentionMaskConverter._unmask_unattended( causal_mask, min_dtype ) @@ -734,7 +770,9 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel): return causal_mask -class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): +class Qwen2_5_VLMoEForAction( + Qwen2_5_VLForConditionalGeneration, ActionGenerationMixin, ActionModelMixMin +): """ Qwen2.5 Vision-Language Mixture of Experts model for action processing. @@ -792,10 +830,11 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): def from_pretrained( cls, pretrained_model_path, - train_config, + train_config=None, config_path=None, processor_path=None, action_tokenizer_path=None, + is_train=False, **kwargs, ): """ @@ -812,9 +851,30 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): Qwen2_5_VLMoEForAction: Loaded model instance """ # Load model components from pretrained path - config_path = os.path.join(pretrained_model_path, "config.json") - config = cls.config_class.from_pretrained(config_path) - processor = AutoProcessor.from_pretrained(pretrained_model_path, use_fast=True) + + if train_config is None: + try: + with open(os.path.join(pretrained_model_path, "config.yml"), "r") as f: + train_config = yaml.load(f, Loader=yaml.FullLoader) + except Exception as e: + print(f"load train_config.yml fail: {e}") + train_config = None + + model_config_path = os.path.join(pretrained_model_path, "config.json") + model_config = cls.config_class.from_pretrained(model_config_path) + + if train_config is not None: + model_config = update_model_config(train_config, model_config) + processors_dict = load_wallx_processors(train_config) + processor = processors_dict["processor"] + else: + processor = AutoProcessor.from_pretrained( + pretrained_model_path, use_fast=True + ) + + if not is_train: + model_config._attn_implementation = "sdpa" + if action_tokenizer_path is not None: processor.action_processor = AutoProcessor.from_pretrained( action_tokenizer_path, trust_remote_code=True @@ -822,18 +882,19 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): # Set the customized robot configuration to ensure consistency between cross-embodiment # representations and the Wall-X action dimensionality. - cls._set_customized_config(train_config) - customized_dof_config = train_config["customized_robot_config"][ - "customized_dof_config" - ] - customized_agent_pos_config = train_config["customized_robot_config"][ - "customized_agent_pos_config" - ] - setattr(config, "customized_dof_config", customized_dof_config) - setattr(config, "customized_agent_pos_config", customized_agent_pos_config) + # if not train_config: + # cls._set_customized_config(train_config) + # customized_dof_config = train_config["customized_robot_config"][ + # "customized_dof_config" + # ] + # customized_agent_pos_config = train_config["customized_robot_config"][ + # "customized_agent_pos_config" + # ] + # setattr(model_config, "customized_dof_config", customized_dof_config) + # setattr(model_config, "customized_agent_pos_config", customized_agent_pos_config) # Initialize model with configuration and processor - model = cls(config, processor=processor, **kwargs) + model = cls(model_config, processor=processor, **kwargs) # Resize token embeddings to match processor tokenizer vocabulary size model.resize_token_embeddings(len(processor.tokenizer)) @@ -843,6 +904,7 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): os.path.join(pretrained_model_path, "*.safetensors") ) state_dict = {} + embed_tokens_size = len(processor.tokenizer) for file in safetensor_files: sd = load_file(file, device="cpu") # filter normalizer statistic params @@ -851,10 +913,14 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): if "action_preprocessor.normalizer" in key: print(f"filter load model weight {key}") del_keys.append(key) + if "embed_tokens.weight" in key: + embed_tokens_size = sd[key].shape[0] + # if train_config is not None: for key in del_keys: del sd[key] state_dict.update(sd) - + if embed_tokens_size != len(processor.tokenizer): + model.resize_token_embeddings(embed_tokens_size) model.load_state_dict(state_dict, strict=False) return model @@ -867,6 +933,7 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): action_tokenizer=None, action_mapper=None, flow_loss_weight=1.0, + use_selective_recompute=False, ): """ Initialize the Qwen2.5 VLMoE model for action processing. @@ -885,7 +952,9 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config( config.vision_config ) - self.model = Qwen2_5_VLMoEModel(config) + self.model = Qwen2_5_VLMoEModel( + config, use_selective_recompute=use_selective_recompute + ) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) @@ -897,6 +966,7 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): # Define action token IDs self.define_action_token_id() + self.times_cache = {} # cache times linspace for each num_inference_timesteps # Cache for rope deltas self.rope_deltas = None @@ -1002,46 +1072,69 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): attention_mask: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """ - Calculate 3D RoPE (Rotary Position Embedding) indices for vision and text tokens. + Calculate the 3D rope index based on image and video's temporal, height and width in LLM. - This method computes position embeddings that account for the temporal, height, and width - dimensions of vision tokens (images/videos) while maintaining standard 1D position embeddings - for text tokens. + Explanation: + Each embedding sequence contains vision embedding and text embedding or just contains text embedding. - For vision tokens, 3D position embeddings are calculated based on: - - Temporal dimension: Time patches in videos - - Height dimension: Vertical patches in images/video frames - - Width dimension: Horizontal patches in images/video frames + For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. + Examples: + input_ids: [T T T T T], here T is for text. + temporal position_ids: [0, 1, 2, 3, 4] + height position_ids: [0, 1, 2, 3, 4] + width position_ids: [0, 1, 2, 3, 4] - For text tokens, standard 1D position embeddings are used, continuing from the maximum - vision position ID plus 1. + For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part + and 1D rotary position embeddin for text part. + Examples: + Temporal (Time): 3 patches, representing different segments of the video in time. + Height: 2 patches, dividing each frame vertically. + Width: 2 patches, dividing each frame horizontally. + We also have some important parameters: + fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second. + tokens_per_second: This is a crucial parameter. It dictates how many "time-steps" or "temporal tokens" are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity. + temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames. + interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs. + input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. + vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100] + vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] + vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] + text temporal position_ids: [101, 102, 103, 104, 105] + text height position_ids: [101, 102, 103, 104, 105] + text width position_ids: [101, 102, 103, 104, 105] + Here we calculate the text start position_ids as the max vision position_ids plus 1. Args: - input_ids (torch.LongTensor, optional): Input token IDs of shape (batch_size, sequence_length) - image_grid_thw (torch.LongTensor, optional): Image grid dimensions (num_images, 3) for [temporal, height, width] - video_grid_thw (torch.LongTensor, optional): Video grid dimensions (num_videos, 3) for [temporal, height, width] - second_per_grid_ts (torch.Tensor, optional): Time interval per temporal grid (num_videos,) - attention_mask (torch.Tensor, optional): Attention mask (batch_size, sequence_length) + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*): + The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. Returns: - tuple: - - position_ids (torch.LongTensor): 3D position IDs of shape (3, batch_size, sequence_length) - - mrope_position_deltas (torch.Tensor): Position deltas for mRoPE of shape (batch_size, 1) + position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) + mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) """ spatial_merge_size = self.config.vision_config.spatial_merge_size image_token_id = self.config.image_token_id video_token_id = self.config.video_token_id vision_start_token_id = self.config.vision_start_token_id mrope_position_deltas = [] - if input_ids is not None and ( image_grid_thw is not None or video_grid_thw is not None ): total_input_ids = input_ids if attention_mask is None: attention_mask = torch.ones_like(total_input_ids) - - # Initialize 3D position IDs tensor position_ids = torch.ones( 3, input_ids.shape[0], @@ -1049,44 +1142,31 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): dtype=input_ids.dtype, device=input_ids.device, ) - image_index, video_index = 0, 0 attention_mask = attention_mask.to(total_input_ids.device) - - # Process each sequence in the batch for i, input_ids in enumerate(total_input_ids): input_ids = input_ids[attention_mask[i] == 1] image_nums, video_nums = 0, 0 - - # Find vision tokens and count images/videos vision_start_indices = torch.argwhere( input_ids == vision_start_token_id ).squeeze(1) vision_tokens = input_ids[vision_start_indices + 1] image_nums = (vision_tokens == image_token_id).sum() video_nums = (vision_tokens == video_token_id).sum() - input_tokens = input_ids.tolist() llm_pos_ids_list: list = [] st = 0 remain_images, remain_videos = image_nums, video_nums - - # Process each vision token (image or video) for _ in range(image_nums + video_nums): - # Find next image or video token if image_token_id in input_tokens and remain_images > 0: ed_image = input_tokens.index(image_token_id, st) else: ed_image = len(input_tokens) + 1 - if video_token_id in input_tokens and remain_videos > 0: ed_video = input_tokens.index(video_token_id, st) else: ed_video = len(input_tokens) + 1 - - # Determine if processing image or video token if ed_image < ed_video: - # Process image token t, h, w = ( image_grid_thw[image_index][0], image_grid_thw[image_index][1], @@ -1096,8 +1176,8 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): image_index += 1 remain_images -= 1 ed = ed_image + else: - # Process video token t, h, w = ( video_grid_thw[video_index][0], video_grid_thw[video_index][1], @@ -1110,8 +1190,6 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): video_index += 1 remain_videos -= 1 ed = ed_video - - # Calculate grid dimensions after spatial merging llm_grid_t, llm_grid_h, llm_grid_w = ( t.item(), h.item() // spatial_merge_size, @@ -1119,7 +1197,6 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): ) text_len = ed - st - # Add position IDs for text tokens before vision token st_idx = ( llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 @@ -1129,20 +1206,18 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx ) - # Calculate 3D position embeddings for vision tokens range_tensor = torch.arange(llm_grid_t).view(-1, 1) expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w) - # Calculate temporal position IDs with time scaling time_tensor = ( expanded_range * second_per_grid_t * self.config.vision_config.tokens_per_second ) + time_tensor_long = time_tensor.long() t_index = time_tensor_long.flatten() - # Calculate spatial position IDs h_index = ( torch.arange(llm_grid_h) .view(1, -1, 1) @@ -1155,14 +1230,11 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): .expand(llm_grid_t, llm_grid_h, -1) .flatten() ) - - # Add 3D position IDs for vision tokens llm_pos_ids_list.append( torch.stack([t_index, h_index, w_index]) + text_len + st_idx ) st = ed + llm_grid_t * llm_grid_h * llm_grid_w - # Add position IDs for remaining text tokens if st < len(input_tokens): st_idx = ( llm_pos_ids_list[-1].max() + 1 @@ -1174,7 +1246,6 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx ) - # Concatenate all position IDs for this sequence llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) position_ids[..., i, attention_mask[i] == 1] = llm_positions.to( position_ids.device @@ -1182,13 +1253,11 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): mrope_position_deltas.append( llm_positions.max() + 1 - len(total_input_ids[i]) ) - mrope_position_deltas = torch.tensor( mrope_position_deltas, device=input_ids.device ).unsqueeze(1) return position_ids, mrope_position_deltas else: - # Handle case without vision tokens - use standard 1D position embeddings if attention_mask is not None: position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) @@ -1222,9 +1291,6 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, - moe_token_types: Optional[ - torch.LongTensor - ] = None, # MoE token type assignments labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, @@ -1234,58 +1300,25 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): pixel_values_videos: Optional[torch.FloatTensor] = None, image_grid_thw: Optional[torch.LongTensor] = None, video_grid_thw: Optional[torch.LongTensor] = None, - action_chunk: Optional[torch.FloatTensor] = None, # Action trajectory chunks - proprioception: Optional[ - torch.FloatTensor - ] = None, # Joint position/orientation data rope_deltas: Optional[torch.LongTensor] = None, cache_position: Optional[torch.LongTensor] = None, second_per_grid_ts: Optional[torch.Tensor] = None, + # for vla + moe_token_types: Optional[torch.LongTensor] = None, + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + positional_masks: Optional[dict] = None, + action_chunk: Optional[torch.FloatTensor] = None, + proprioception: Optional[torch.FloatTensor] = None, dataset_names: Optional[str] = None, dof_mask: Optional[torch.FloatTensor] = None, agent_pos_mask: Optional[torch.FloatTensor] = None, + flow_loss_mask: Optional[torch.FloatTensor] = None, **kwargs, ) -> Union[Tuple, Qwen2_5_VLACausalLMOutputWithPast]: - """ - Forward pass for training with multi-modal inputs including vision, text, and action data. + if input_ids is not None: + batch_size, seq_length = input_ids.shape - This method handles the complete forward pass during training, processing various input modalities - including images, videos, text, proprioceptive data, and action sequences. It computes losses - for both language modeling and action prediction using flow matching. - - Args: - input_ids (torch.LongTensor, optional): Input token IDs - attention_mask (torch.Tensor, optional): Attention mask for input tokens - position_ids (torch.LongTensor, optional): Position IDs for tokens - past_key_values (List[torch.FloatTensor], optional): Cached key-value pairs for generation - inputs_embeds (torch.FloatTensor, optional): Pre-computed input embeddings - moe_token_types (torch.LongTensor, optional): Token type assignments for MoE routing - labels (torch.LongTensor, optional): Target labels for loss computation - use_cache (bool, optional): Whether to use key-value caching - output_attentions (bool, optional): Whether to return attention weights - output_hidden_states (bool, optional): Whether to return hidden states - return_dict (bool, optional): Whether to return structured output - pixel_values (torch.Tensor, optional): Image pixel values - pixel_values_videos (torch.FloatTensor, optional): Video pixel values - image_grid_thw (torch.LongTensor, optional): Image grid dimensions (temporal, height, width) - video_grid_thw (torch.LongTensor, optional): Video grid dimensions (temporal, height, width) - action_chunk (torch.FloatTensor, optional): Action trajectory data chunks - proprioception (torch.FloatTensor, optional): Proprioceptive sensor data (joint positions, etc.) - rope_deltas (torch.LongTensor, optional): RoPE position deltas - cache_position (torch.LongTensor, optional): Cache position indices - second_per_grid_ts (torch.Tensor, optional): Time interval per temporal grid - dataset_names (str, optional): Names of datasets in the current batch - dof_mask (torch.FloatTensor, optional): Degrees of freedom mask for action tokens - agent_pos_mask (torch.FloatTensor, optional): Agent position mask for proprioceptive data - **kwargs: Additional keyword arguments - - Returns: - Union[Tuple, Qwen2_5_VLACausalLMOutputWithPast]: Model outputs including losses, logits, - and auxiliary information, or tuple if return_dict=False - """ - batch_size, seq_length = input_ids.shape - - # Set output configuration from model config if not specified output_attentions = ( output_attentions if output_attentions is not None @@ -1300,60 +1333,53 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): return_dict if return_dict is not None else self.config.use_return_dict ) - # Calculate RoPE position IDs if not provided - # Note: Cannot calculate rope deltas with 4D attention mask. TODO: Fix this limitation + if start_indices is None or end_indices is None: + # Calculate the start and end positions of each expert group's tokens after permutation + group_size = torch.zeros( + self.config.num_experts, dtype=torch.long, device="cpu" + ) + for i in range(self.config.num_experts): + group_size[i] = (moe_token_types == i).sum() + + # Calculate start and end indices for each expert group + start_indices = torch.cumsum(group_size, dim=0) - group_size + end_indices = torch.cumsum(group_size, dim=0) + + # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme if position_ids is None and ( attention_mask is None or attention_mask.ndim == 2 ): - # Calculate RoPE index once per generation in the pre-fill stage only + # calculate RoPE index once per generation in the pre-fill stage only if ( (cache_position is not None and cache_position[0] == 0) or self.rope_deltas is None or (past_key_values is None or past_key_values.get_seq_length() == 0) ): - position_ids, rope_deltas = ops.get_rope_index( + position_ids, rope_deltas = self.get_rope_index( input_ids=input_ids, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, second_per_grid_ts=second_per_grid_ts, attention_mask=attention_mask, - spatial_merge_size=self.config.vision_config.spatial_merge_size, - image_token_id=self.config.image_token_id, - video_token_id=self.config.video_token_id, - vision_start_token_id=self.config.vision_start_token_id, - tokens_per_second=self.config.vision_config.tokens_per_second, ) self.rope_deltas = rope_deltas - # Use previously calculated rope deltas to get correct position IDs + # then use the prev pre-calculated rope-deltas to get the correct position ids else: + # batch_size, seq_length, _ = inputs_embeds.shape delta = ( - (cache_position[0] + self.rope_deltas).to(self.device) + (cache_position[0] + self.rope_deltas).to(cache_position.device) if cache_position is not None else 0 ) - position_ids = torch.arange(seq_length, device=self.device) + position_ids = torch.arange(seq_length, device=cache_position.device) position_ids = position_ids.view(1, -1).expand(batch_size, -1) if cache_position is not None: # otherwise `deltas` is an int `0` delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) position_ids = position_ids.add(delta) position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) - # Calculate token distribution across MoE expert groups - group_size = torch.zeros( - self.config.num_experts, dtype=torch.long, device="cpu" - ) - for i in range(self.config.num_experts): - group_size[i] = (moe_token_types == i).sum() - - # Calculate start and end indices for each expert group - start_indices = torch.cumsum(group_size, dim=0) - group_size - end_indices = torch.cumsum(group_size, dim=0) - - # Process input embeddings with multi-modal data if inputs_embeds is None: inputs_embeds = self.model.embed_tokens(input_ids) - - # Process image embeddings if pixel_values is not None: pixel_values = pixel_values.type(self.visual.dtype) image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) @@ -1367,18 +1393,16 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): ) inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) - # Process video embeddings if pixel_values_videos is not None: pixel_values_videos = pixel_values_videos.type(self.visual.dtype) video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) n_video_tokens = (input_ids == self.config.video_token_id).sum().item() n_video_features = video_embeds.shape[0] - - # Validate video token and feature count match if n_video_tokens != n_video_features: raise ValueError( f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}" ) + mask = input_ids == self.config.video_token_id mask_unsqueezed = mask.unsqueeze(-1) mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) @@ -1389,183 +1413,56 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): ) inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) - # Process proprioceptive data (joint positions, orientations, etc.) - if proprioception is not None: - proprioception = proprioception.to(inputs_embeds.device).to( - inputs_embeds.dtype - ) - agent_pos_mask = agent_pos_mask.to(inputs_embeds.device).to( - inputs_embeds.dtype - ) - proprioception = self.action_preprocessor.proprioception_proj( - proprioception, - dataset_names, - agent_pos_mask, - use_history=proprioception.shape[1] > 1, - ) - mask = input_ids == self.action_token_id_set["propri_token_id"] - mask_unsqueezed = mask.unsqueeze(-1) - mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) - proprioception_mask = mask_expanded.to(inputs_embeds.device) + inputs_embeds = self.scatter_proprioception_embeddings( + input_ids, inputs_embeds, proprioception, dataset_names, agent_pos_mask + ) - proprioception = proprioception.to( - inputs_embeds.device, inputs_embeds.dtype - ) - inputs_embeds = inputs_embeds.masked_scatter( - proprioception_mask, proprioception - ) - elif self.training: - # Dummy forward pass to ensure gradient registration in DDP - # This handles cases where one process has proprioception data while another doesn't - # Without this, DDP would hang waiting for a gradient that will never be computed - dummy_input = torch.randn( - 2, - self.action_preprocessor.propri_dim * 2, - device=inputs_embeds.device, - ) - dummy_forward = self.action_preprocessor.proprioception_proj( - dummy_input - ) - dummy_loss = sum(p.sum() for p in dummy_forward) - inputs_embeds = inputs_embeds + 0 * dummy_loss - - # Process action chunk data - if action_chunk is not None: - action_chunk = action_chunk.to(inputs_embeds.device).to( - inputs_embeds.dtype - ) - dof_mask = dof_mask.to(inputs_embeds.device).to(inputs_embeds.dtype) - noisy_action_emb, flow = self.action_preprocessor( - action_chunk, dataset_names, dof_mask - ) - mask = input_ids == self.action_token_id_set["action_token_id"] - mask_unsqueezed = mask.unsqueeze(-1) - mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) - action_mask = mask_expanded.to(inputs_embeds.device) - - noisy_action_emb = noisy_action_emb.to( - inputs_embeds.device, inputs_embeds.dtype - ) - inputs_embeds = inputs_embeds.masked_scatter( - action_mask, noisy_action_emb - ) + inputs_embeds, flow, adarms_cond = self.scatter_flow_action_embeddings( + input_ids, inputs_embeds, action_chunk, dataset_names, dof_mask + ) if attention_mask is not None: attention_mask = attention_mask.to(inputs_embeds.device) - # Forward pass through the main model outputs = self.model( input_ids=None, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, - moe_token_types=moe_token_types, # Pass token types for MoE routing + moe_token_types=moe_token_types, start_indices=start_indices, end_indices=end_indices, + positional_masks=positional_masks, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + adarms_conds=[None, adarms_cond], + # cache_position=cache_position, ) hidden_states = outputs[0] logits = self.lm_head(hidden_states) - # Initialize loss computation variables - loss = None - cross_entropy_loss, flow_loss = None, None - channel_loss_dict = None - channel_loss_count_dict = None + ( + loss, + cross_entropy_loss, + flow_loss, + channel_loss_dict, + channel_loss_count_dict, + ) = self.compute_loss( + hidden_states=hidden_states, + logits=logits, + input_ids=input_ids, + dataset_names=dataset_names, + labels=labels, + action_chunk=action_chunk, + dof_mask=dof_mask, + flow=flow, + flow_loss_mask=flow_loss_mask, + ) - # Compute losses if labels are provided - if labels is not None: - loss = 0 - action_accuracy = 0 - unique_datasets_name = list(set(dataset_names)) - - # Initialize per-dataset loss tracking dictionaries - channel_loss_dict = { - dataset_name: torch.tensor(0.0, device=logits.device) - for dataset_name in ACTION_DATASET_NAMES + MULTIMODAL_DATASET_NAMES - } - channel_loss_count_dict = { - dataset_name: torch.tensor(0, device=logits.device) - for dataset_name in ACTION_DATASET_NAMES + MULTIMODAL_DATASET_NAMES - } - - # Compute standard cross-entropy loss for language modeling - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1) - - # Enable model parallelism by moving labels to correct device - shift_labels = shift_labels.to(shift_logits.device) - non_ignored_mask = shift_labels != -100 - _cross_entropy_loss = self.loss_fct(shift_logits, shift_labels) - cross_entropy_loss = ( - _cross_entropy_loss[non_ignored_mask].mean() - if non_ignored_mask.any() - else torch.tensor(0.0, device=shift_logits.device) - ) - - # Compute per-dataset channel losses - _cross_entropy_loss = _cross_entropy_loss.view(batch_size, seq_length - 1) - non_ignored_mask = non_ignored_mask.view(batch_size, seq_length - 1) - for dataset_name_i in unique_datasets_name: - dataset_mask = torch.tensor( - [name == dataset_name_i for name in dataset_names], - device=logits.device, - ) - combined_mask = dataset_mask.unsqueeze(1) & non_ignored_mask - channel_loss_dict[dataset_name_i] = ( - _cross_entropy_loss[combined_mask].sum() - if combined_mask.any() - else torch.tensor(0.0, device=shift_logits.device) - ) - channel_loss_count_dict[dataset_name_i] += combined_mask.sum() - - # Add cross-entropy loss to total loss if valid - if not torch.isnan(cross_entropy_loss): - loss += cross_entropy_loss - else: - with torch.no_grad(): - cross_entropy_loss.detach() - - # Compute action token prediction accuracy - shift_logits = logits[..., :-1, :].contiguous() - action_preds = shift_logits.argmax(dim=-1) - shift_labels = labels[..., 1:].contiguous() - if self.use_fast_tokenizer: - action_mask = ( - shift_labels > self.action_token_id_set["fast_action_token_list"][0] - ) - correct_preds = (action_preds == shift_labels) & action_mask - action_accuracy = ( - correct_preds.sum().float() / action_mask.sum().float() - ) - channel_loss_dict["action_accuracy"] = action_accuracy - - if action_chunk is not None: - action_mask = input_ids == self.action_token_id_set["action_token_id"] - if action_mask.any(): - action_hidden_states = hidden_states[action_mask] - flow = flow.reshape(-1, flow.shape[-1]) - _flow_loss = self.action_preprocessor.flow_loss( - action_hidden_states, flow, dof_mask - ) - if isinstance(_flow_loss, torch.Tensor): - flow_loss = _flow_loss.mean() - if loss is not None: - loss += self.flow_loss_weight * flow_loss - else: - loss = self.flow_loss_weight * flow_loss - _flow_loss = _flow_loss.view( - dof_mask.shape[0], dof_mask.shape[1], dof_mask.shape[2] - ) - - # Return outputs based on return_dict setting if not return_dict: output = (logits,) + outputs[1:] return (loss,) + output if loss is not None else output @@ -1776,7 +1673,7 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): ) inputs_embeds[proprioception_mask] = proprio_embed.reshape( -1, inputs_embeds.shape[-1] - ) + ).to(inputs_embeds.dtype) if attention_mask is not None: attention_mask = attention_mask.to(inputs_embeds.device) @@ -1986,14 +1883,16 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): # Prepare timestep for batch processing timestep = timestep.unsqueeze(0).repeat(noisy_action.shape[0]) - action_embed = self.action_preprocessor.step( + action_embed, _ = self.action_preprocessor.step( timestep=timestep, noisy_action=noisy_action, dof_mask=dof_mask ) action_embed = action_embed.reshape(-1, inputs_embeds.shape[-1]) # Create temporary copy of embeddings for thread safety temp_inputs_embeds = inputs_embeds.clone() - temp_inputs_embeds[action_mask] = action_embed + temp_inputs_embeds[action_mask] = action_embed.to( + temp_inputs_embeds.dtype + ) # Forward pass through transformer transformer_outputs = self.model( @@ -2015,7 +1914,11 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): hidden_states = transformer_outputs.last_hidden_state action_mask = input_ids == self.action_token_id_set["action_token_id"] action_hidden_states = hidden_states[action_mask] - pred = self.action_preprocessor.action_proj_back(action_hidden_states) + pred = self.action_preprocessor.action_proj_back( + action_hidden_states[ + :, : self.action_preprocessor.action_hidden_size + ] + ) return pred.reshape(batch_size, pred_horizon, action_dim) # Perform ODE integration for diffusion sampling @@ -2026,7 +1929,12 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): device=inputs_embeds.device, dtype=inputs_embeds.dtype, ) - action_trajectory = odeint(step, noisy_action, times, method="euler") + action_trajectory = odeint( + step, + noisy_action.to(torch.float32), + times.to(torch.float32), + method="euler", + ) # Extract final predicted action and unnormalize predict_action = action_trajectory[-1] @@ -2047,6 +1955,411 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): return output + @torch.no_grad() + def generate_flow_action( + self, + input_ids, + action_horizon, + action_dim, + num_inference_timesteps: int = 10, + padding_action: Optional[torch.Tensor] = None, + prefix_length: Optional[int] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + moe_token_types: Optional[torch.LongTensor] = None, + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + positional_masks: Optional[torch.LongTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + action_chunk: Optional[torch.FloatTensor] = None, + proprioception: Optional[torch.FloatTensor] = None, + unnorm_proprioception: Optional[torch.FloatTensor] = None, + rope_deltas: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + dataset_names: Optional[str] = None, + dof_mask: Optional[torch.FloatTensor] = None, + agent_pos_mask: Optional[torch.FloatTensor] = None, + unnorm: Optional[bool] = True, + **kwargs, + ): + + total_start_time = time.time() + timing_results = {} + + batch_size = ( + input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0] + ) + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + embed_start_time = time.time() + if inputs_embeds is None: + inputs_embeds = self.model.embed_tokens(input_ids) + if pixel_values is not None: + pixel_values = pixel_values.type(self.visual.dtype) + image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + n_image_tokens = (input_ids == self.config.image_token_id).sum().item() + n_image_features = image_embeds.shape[0] + if n_image_tokens != n_image_features: + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" + ) + + mask = input_ids == self.config.image_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + image_mask = mask_expanded.to(inputs_embeds.device) + + image_embeds = image_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if pixel_values_videos is not None: + pixel_values_videos = pixel_values_videos.type(self.visual.dtype) + video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) + n_video_tokens = (input_ids == self.config.video_token_id).sum().item() + n_video_features = video_embeds.shape[0] + if n_video_tokens != n_video_features: + raise ValueError( + f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}" + ) + + mask = input_ids == self.config.video_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + video_mask = mask_expanded.to(inputs_embeds.device) + + video_embeds = video_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + if ( + proprioception is not None + and not self.config.use_state_string_representation + ): + proprioception = proprioception.to(inputs_embeds.device) + agent_pos_mask = agent_pos_mask.to(inputs_embeds.device) + proprio_embed = self.action_preprocessor.proprioception_proj( + proprioception, + dataset_names, + agent_pos_mask, + use_history=proprioception.shape[1] > 1, + ) + proprioception_mask = ( + input_ids == self.action_token_id_set["propri_token_id"] + ) + inputs_embeds[proprioception_mask] = proprio_embed.reshape( + -1, inputs_embeds.shape[-1] + ).to(inputs_embeds.dtype) + + if attention_mask is not None: + attention_mask = attention_mask.to(inputs_embeds.device) + + timing_results["embed_processing"] = time.time() - embed_start_time + + position_start_time = time.time() + # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme + if position_ids is None and ( + attention_mask is None or attention_mask.ndim == 2 + ): + # calculate RoPE index once per generation in the pre-fill stage only + if ( + (cache_position is not None and cache_position[0] == 0) + or self.rope_deltas is None + or (past_key_values is None or past_key_values.get_seq_length() == 0) + ): + position_ids, rope_deltas = self.get_rope_index( + input_ids, + image_grid_thw, + video_grid_thw, + second_per_grid_ts, + attention_mask, + ) + self.rope_deltas = rope_deltas + # then use the prev pre-calculated rope-deltas to get the correct position ids + else: + batch_size, seq_length, _ = inputs_embeds.shape + delta = ( + (cache_position[0] + self.rope_deltas).to(inputs_embeds.device) + if cache_position is not None + else 0 + ) + position_ids = torch.arange(seq_length, device=inputs_embeds.device) + position_ids = position_ids.view(1, -1).expand(batch_size, -1) + if cache_position is not None: # otherwise `deltas` is an int `0` + delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) + position_ids = position_ids.add(delta) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + + if start_indices is None or end_indices is None: + # Calculate the start and end positions of each expert group's tokens after permutation (the dataset does not contain `num_expert` information, so this calculation must be done here). + group_size = torch.zeros( + self.config.num_experts, dtype=torch.long, device="cpu" + ) + for i in range(self.config.num_experts): + group_size[i] = (moe_token_types == i).sum() + + # Calculate start and end indices for each expert group + start_indices = torch.cumsum(group_size, dim=0) - group_size + end_indices = torch.cumsum(group_size, dim=0) + + timing_results["position_encoding"] = time.time() - position_start_time + + action_init_start_time = time.time() + if action_chunk is not None: + action_chunk = action_chunk.to(inputs_embeds.device).to(torch.float32) + + output = {} + # Reproduce + # torch.manual_seed(0) + noise = torch.randn( + size=(batch_size, action_horizon, action_dim), + dtype=torch.float32, + device=inputs_embeds.device, + ) + noisy_action = noise.clone() + dof_mask = dof_mask.to(inputs_embeds.device).to(torch.float32) + + if num_inference_timesteps not in self.times_cache: + self.times_cache[num_inference_timesteps] = torch.linspace( + 0.0, + 1.0, + num_inference_timesteps + 1, + device=inputs_embeds.device, + dtype=torch.float32, + ) + times = self.times_cache[num_inference_timesteps] + dt = times[1] - times[0] + time_0 = times[0].unsqueeze(0).repeat(noisy_action.shape[0]) + action_embed, adarms_cond = self.action_preprocessor.step( + timestep=time_0, noisy_action=noisy_action, dof_mask=dof_mask + ) + action_embed = action_embed.reshape(-1, inputs_embeds.shape[-1]).to( + inputs_embeds.dtype + ) + flow_action_mask = input_ids == self.action_token_id_set["action_token_id"] + + inputs_embeds[flow_action_mask] = action_embed + + timing_results["action_initialization"] = time.time() - action_init_start_time + + prefetch_start_time = time.time() + prefetch_output = self.model( + input_ids=None, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=None, + inputs_embeds=inputs_embeds, + moe_token_types=moe_token_types, + start_indices=start_indices, + end_indices=end_indices, + positional_masks=positional_masks, + use_cache=True, + output_attentions=False, + output_hidden_states=False, + return_dict=True, + adarms_conds=[None, adarms_cond], + ) + hidden_states = prefetch_output.last_hidden_state + prefix_kv_cache = prefetch_output.past_key_values + + action_hidden_states = hidden_states[flow_action_mask].to(torch.float32) + action_pred = self.action_preprocessor.action_proj_back( + action_hidden_states[:, : self.action_preprocessor.action_hidden_size] + ) + if getattr(self.config, "use_x_pred", False): + v_0 = action_pred - noise.reshape(-1, noise.shape[-1]) + else: + v_0 = action_pred + + if (not dof_mask.all()) and (padding_action is not None): + print("use padding action", flush=True) + v_padding = padding_action - noisy_action + v_0 = (v_padding) * (1 - dof_mask) + v_0 * dof_mask + + noisy_action = noisy_action + dt * v_0.reshape( + batch_size, action_horizon, action_dim + ) + + timing_results["prefetch_forward"] = time.time() - prefetch_start_time + + cache_prep_start_time = time.time() + + if prefix_length is None: + has_true = flow_action_mask.any(dim=1) + prefix_length = torch.argmax(flow_action_mask.float(), dim=1, keepdim=True) + prefix_length[~has_true] = flow_action_mask.shape[1] + prefix_length = prefix_length[0] + + # support different transformers version + if hasattr(prefix_kv_cache, "key_cache"): + for layer_i in range(len(prefix_kv_cache.key_cache)): + prefix_kv_cache.key_cache[layer_i] = prefix_kv_cache.key_cache[layer_i][ + :, :, :prefix_length, : + ] + prefix_kv_cache.value_cache[layer_i] = prefix_kv_cache.value_cache[ + layer_i + ][:, :, :prefix_length, :] + else: + for layer_i in range(len(prefix_kv_cache.layers)): + prefix_kv_cache.layers[layer_i].keys = prefix_kv_cache.layers[ + layer_i + ].keys[:, :, :prefix_length, :] + prefix_kv_cache.layers[layer_i].values = prefix_kv_cache.layers[ + layer_i + ].values[:, :, :prefix_length, :] + + postfix_position_ids = position_ids[:, :, prefix_length:] + postfix_inputs_embeds = inputs_embeds[:, prefix_length:, :] + postfix_attention_mask = attention_mask[:, prefix_length:] + postfix_moe_token_types = moe_token_types[:, prefix_length:] + postfix_input_ids = input_ids[:, prefix_length:] + + group_size = torch.zeros( + self.config.num_experts, dtype=torch.long, device="cpu" + ) + for i in range(self.config.num_experts): + group_size[i] = (postfix_moe_token_types == i).sum() + + # Calculate start and end indices for each expert group + postfix_start_indices = torch.cumsum(group_size, dim=0) - group_size + postfix_end_indices = torch.cumsum(group_size, dim=0) + + pad_token_id = self.processor.tokenizer.pad_token_id + padding_mask = input_ids == pad_token_id + + # prefix_length, postfix_length = prefix_indices.shape[0], postfix_indices.shape[0] + + postfix_length = input_ids.shape[-1] - prefix_length + _postfix_attention_mask = torch.ones( + (batch_size, postfix_length, prefix_length + postfix_length), + dtype=torch.bool, + device=postfix_attention_mask.device, + ) + + # Use a padding mask to set the corresponding rows and columns to false. + # Get the padding mask for the postfix portion. + postfix_padding_mask = padding_mask[ + :, prefix_length: + ] # [batch_size, postfix_length] + full_padding_mask = padding_mask # [batch_size, prefix_length + postfix_length] + + # causal mask for postfix attention + if self.config.causal_action_attention_mask: + _postfix_attention_mask[:, :, prefix_length:] = torch.tril( + torch.ones( + (postfix_length, postfix_length), + dtype=torch.bool, + device=postfix_attention_mask.device, + ) + ) + + for batch_idx in range(padding_mask.shape[0]): + # Set the rows corresponding to the padding positions to False (where the query position is padding). + _postfix_attention_mask[batch_idx, postfix_padding_mask[batch_idx], :] = ( + False + ) + # Set the columns corresponding to the padding positions to False (the key position is the padding). + _postfix_attention_mask[batch_idx, :, full_padding_mask[batch_idx]] = False + + timing_results["cache_preprocessing"] = time.time() - cache_prep_start_time + + ode_start_time = time.time() + + def step_with_kvcache(timestep, noisy_action): + action_mask = ( + postfix_input_ids == self.action_token_id_set["action_token_id"] + ) + assert action_mask.any(), "No action token found in input_ids" + timestep = timestep.unsqueeze(0).repeat(noisy_action.shape[0]) + action_embed, adarms_cond = self.action_preprocessor.step( + timestep=timestep, noisy_action=noisy_action, dof_mask=dof_mask + ) + action_embed = action_embed.reshape(-1, postfix_inputs_embeds.shape[-1]) + + temp_inputs_embeds = postfix_inputs_embeds.clone() + temp_inputs_embeds[action_mask] = action_embed.to(temp_inputs_embeds.dtype) + transformer_outputs = self.model( + input_ids=None, + attention_mask=_postfix_attention_mask, + position_ids=postfix_position_ids, + past_key_values=prefix_kv_cache, + inputs_embeds=temp_inputs_embeds, + moe_token_types=postfix_moe_token_types, + start_indices=postfix_start_indices, + end_indices=postfix_end_indices, + use_cache=False, + output_attentions=False, + output_hidden_states=False, + return_dict=True, + adarms_conds=[None, adarms_cond], + ) + + hidden_states = transformer_outputs.last_hidden_state + action_hidden_states = hidden_states[action_mask].to(torch.float32) + action_pred = self.action_preprocessor.action_proj_back( + action_hidden_states[:, : self.action_preprocessor.action_hidden_size] + ) + if getattr(self.config, "use_x_pred", False): + v_t = action_pred - noise.reshape(-1, noise.shape[-1]) + else: + v_t = action_pred + return v_t.reshape(batch_size, action_horizon, action_dim) + + action_trajectory = odeint( + step_with_kvcache, noisy_action, times[1:], method="euler" + ) + + timing_results["ode_integration"] = time.time() - ode_start_time + + postprocess_start_time = time.time() + predict_action = action_trajectory[-1] + if unnorm: + predict_action = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + predict_action, dataset_names + ) + ) + output["predict_action"] = predict_action + # normalize action chunk to get gt_action + if action_chunk is not None: + output["gt_action"] = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + action_chunk, dataset_names + ) + ) + + timing_results["postprocessing"] = time.time() - postprocess_start_time + timing_results["total_time"] = time.time() - total_start_time + + output["timing_results"] = timing_results + + return output + def forward( self, mode: Optional[str] = None, predict_mode: Optional[str] = "text", **kwargs ): @@ -2074,7 +2387,7 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): with torch.no_grad(): return self.train_step_forward(**kwargs) elif mode == "predict": - return self.predict(predict_mode=predict_mode, **kwargs) + return self.generate_flow_action(predict_mode=predict_mode, **kwargs) elif mode == "train": return self.train_step_forward(use_cache=False, **kwargs) elif mode == "validate": diff --git a/wall_x/model/vla_mixin.py b/wall_x/model/vla_mixin.py new file mode 100644 index 0000000..1c98aa2 --- /dev/null +++ b/wall_x/model/vla_mixin.py @@ -0,0 +1,987 @@ +import torch +import torch.nn as nn +import torch.utils.checkpoint as cp + +from torch.distributed.fsdp import MixedPrecision as MP +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + +from wall_x.fusions import ops + +from peft import LoraConfig, get_peft_model +from typing import Optional, Union, Dict +from packaging import version + +from transformers import GenerationMixin +from transformers.activations import ACT2FN +from transformers.modeling_utils import AttentionInterface + +from transformers.utils import logging, is_torch_xla_available + +from wall_x.model.action_head import ActionProcessor +from wall_x.model.model_utils import find_first_last_ones + +ALL_ATTENTION_FUNCTIONS: AttentionInterface = AttentionInterface() +logger = logging.get_logger(__name__) + + +X2ROBOT_ATTENTION_FUNCTIONS = [] +ATTENTION_TYPES_WITH_2D_MASK = [ + "sdpa", +] +ATTENTION_TYPES_WITH_FLASH_MASK = [] + + +class TokenTypeRouter(nn.Module): + def __init__(self, num_experts: int): + super().__init__() + self.num_experts = num_experts + + def forward(self, token_types: torch.Tensor) -> torch.Tensor: + """ + Assigns tokens to different experts based on `token_type`. + Args: + token_types (torch.Tensor): A tensor of shape (batch_size, seq_length) representing the type of each token. + + Returns: + experts_indices (torch.Tensor): A tensor of shape (batch_size, seq_length) representing the expert index assigned to each token. + """ + experts_indices = token_types % self.num_experts + return experts_indices + + +class BlockSparseMLP(nn.Module): + def __init__(self, config, use_selective_recompute: bool = False): + super().__init__() + self.hidden_size = config["hidden_size"] + self.intermediate_size = config["intermediate_size"] + self.hidden_act = config["hidden_act"] + + self.use_selective_recompute = use_selective_recompute + + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + + self.act_fn = ACT2FN[self.hidden_act] + + def _full_mlp(self, hidden_state): + gate_out = self.gate_proj(hidden_state) + up_out = self.up_proj(hidden_state) + act_out = self.act_fn(gate_out) * up_out + return self.down_proj(act_out) + + def forward(self, hidden_state): + if self.use_selective_recompute: + # Perform checkpoint recalculation for the entire expert MLP. + return cp.checkpoint( + self._full_mlp, + hidden_state, + use_reentrant=False, + ) + else: + return self._full_mlp(hidden_state) + + +class SparseMoeBlock(nn.Module): + def __init__(self, config, num_experts: int, use_selective_recompute: bool = False): + super().__init__() + self.num_experts = num_experts + self.use_selective_recompute = use_selective_recompute + + # Pass the `use_selective_recompute` parameter to each expert. + self.experts = nn.ModuleList( + [ + BlockSparseMLP( + config.experts[i], use_selective_recompute=use_selective_recompute + ) + for i in range(num_experts) + ] + ) + + if not hasattr(config, "dim_inputs") or not config.dim_inputs: + raise ValueError("Configuration must contain a valid dim_inputs") + + self.dim_inputs = config.dim_inputs + self.permuted = config.mot_opt + + def forward( + self, + hidden_states: torch.Tensor, + experts_indices: torch.Tensor, + start_indices: torch.Tensor, + end_indices: torch.Tensor, + ) -> torch.Tensor: + + if self.permuted: + permuted_inputs = hidden_states + else: + batch_size, seq_length, hidden_dim = hidden_states.shape + + flat_hidden = hidden_states.reshape(-1, hidden_dim) + experts_indices = experts_indices.reshape(-1) + probs = torch.ones_like(experts_indices, dtype=torch.float32).reshape(-1, 1) + permuted_inputs, row_id_map = ops.permute(flat_hidden, experts_indices) + + # buffer + final_output = torch.zeros_like(permuted_inputs) + + # Expert forward contain selective recompute + for expert_idx, expert in enumerate(self.experts): + start, end = start_indices[expert_idx], end_indices[expert_idx] + if start == end: + continue + + dim_input = self.dim_inputs[expert_idx] + expert_input = permuted_inputs[start:end, :dim_input] + + partial_output = expert(expert_input) + final_output[start:end, :dim_input] = partial_output[:, :dim_input] + + if self.permuted: + return final_output + else: + final_output = ops.unpermute(final_output, row_id_map, probs) + return final_output.reshape(batch_size, seq_length, hidden_dim) + + +class ActionModelMixMin: + # config: Qwen2_5_VLConfig + action_preprocessor: ActionProcessor + router: TokenTypeRouter + moe: SparseMoeBlock + + def __init__(self, config, action_preprocessor, router, moe): + self.config = config + self.action_preprocessor = action_preprocessor + self.router = router + self.moe = moe + self._mot_opt_warned = False + + def set_normalizer(self, normalizer_action, normalizer_propri): + if hasattr(self, "action_preprocessor"): + self.action_preprocessor.set_normalizer( + normalizer_action, normalizer_propri + ) + else: + logger.warning( + "ActionModelMixMin.set_normalizer is called but action_preprocessor is not set" + ) + + def _apply_mlp_moe(self, hidden_states, token_types, start_indices, end_indices): + if self.config.mlp_moe: + hidden_states = self.moe( + hidden_states, token_types, start_indices, end_indices + ) + else: + hidden_states = self.mlp(hidden_states) + return hidden_states + + def _apply_norm_moe( + self, + hidden_states, + token_types, + adarms_conds, + norms, # list of norm layers (expert-wise) + norm, # shared norm if not norm_moe + start_indices=None, + end_indices=None, + use_selective_recompute=False, + ): + """ + MoE-aware LayerNorm with optional selective activation recomputation. + + Only activation math is recomputed. No GEMM is recomputed. + Safe for FSDP (use_reentrant=False). + """ + + gate = None + gate_mask = None + + # ------------------------- + # Case 1: norm_moe=True (expert-wise norm) + # ------------------------- + if self.config.norm_moe: + + # --------------------------------------------------------- + # Case 1A: mot_opt=True (segments assigned by start/end) + # --------------------------------------------------------- + if self.config.mot_opt: + new_hidden_states = torch.zeros_like(hidden_states) + + for expert_idx, expert_norm in enumerate(norms): + start = start_indices[expert_idx] + end = end_indices[expert_idx] + if start == end: + continue + + dim_input = self.config.dim_inputs[expert_idx] + selected = hidden_states[start:end] # [K, D] + + # ====== reshape if adarms on flow expert ====== + if self.config.use_adarms and expert_idx == 1: + selected = selected.view( + -1, + self.config.action_horizon_flow, + selected.shape[-1], + ) + input_slice = selected[:, :, :dim_input] + cond = adarms_conds[expert_idx] + else: + input_slice = selected[:, :dim_input] + cond = adarms_conds[expert_idx] + + if use_selective_recompute: + + def norm_chunk(t_x, t_cond, expert_norm=expert_norm): + if t_cond is None or ( + isinstance(t_cond, torch.Tensor) and t_cond.numel() == 0 + ): + out, _ = expert_norm(t_x) + else: + out, _ = expert_norm(t_x, t_cond) + return out + + cond_for_cp = ( + cond + if cond is not None + else torch.empty(0, device=input_slice.device) + ) + processed = cp.checkpoint( + norm_chunk, + input_slice, + cond_for_cp, + use_reentrant=False, + ) + else: + processed, gate = expert_norm(input_slice, cond) + + # reshape back if needed + if self.config.use_adarms and expert_idx == 1: + processed = processed.view(-1, dim_input) + + new_hidden_states[start:end, :dim_input] = processed.to( + hidden_states.dtype + ) + + hidden_states = new_hidden_states + + # --------------------------------------------------------- + # Case 1B: mot_opt=False (token-level mask) + # --------------------------------------------------------- + else: + + new_hidden_states = torch.zeros_like(hidden_states) + B, S, D = hidden_states.shape + + for expert_idx, expert_norm in enumerate(norms): + mask = token_types == expert_idx + if mask.sum() == 0: + continue + + dim_input = self.config.dim_inputs[expert_idx] + selected = hidden_states[mask] # [K, D] + + if self.config.use_adarms and expert_idx == 1: + gate_mask = mask + selected = selected.view( + -1, + self.config.action_horizon_flow, + selected.shape[-1], + ) + input_slice = selected[:, :, :dim_input] + cond = adarms_conds[expert_idx] + else: + input_slice = selected[:, :dim_input] + cond = adarms_conds[expert_idx] + + if use_selective_recompute: + + def norm_chunk(t_x, t_cond, expert_norm=expert_norm): + if t_cond is None or ( + isinstance(t_cond, torch.Tensor) and t_cond.numel() == 0 + ): + out, _ = expert_norm(t_x) + else: + out, _ = expert_norm(t_x, t_cond) + return out + + cond_for_cp = ( + cond + if cond is not None + else torch.empty(0, device=input_slice.device) + ) + + processed = cp.checkpoint( + norm_chunk, + input_slice, + cond_for_cp, + use_reentrant=False, + ) + else: + processed, gate = expert_norm(input_slice, cond) + + if self.config.use_adarms and expert_idx == 1: + processed = processed.view(-1, dim_input) + + # scatter back + b_id, s_id = torch.where(mask) + new_hidden_states[b_id, s_id, :dim_input] = processed.to( + hidden_states.dtype + ) + + hidden_states = new_hidden_states + + # ------------------------- + # Case 2: norm_moe=False (single LN) + # ------------------------- + else: + + def norm_chunk_shared(t_x, dummy, norm_module=norm): + out, _ = norm_module(t_x) + return out + + if use_selective_recompute: + dummy = torch.empty(0, device=hidden_states.device) + hidden_states = cp.checkpoint( + norm_chunk_shared, + hidden_states, + dummy, + use_reentrant=False, + ) + else: + hidden_states, gate = norm(hidden_states) + + return hidden_states, gate, gate_mask + + def _gated_residual(self, x, y, gate, start_indices=None, end_indices=None): + """ + Applies gated residual connection with optional gate parameter. + + Args: + x: Input tensor (residual) + y: Output tensor to be added + gate: Optional gate tensor to modulate the addition + + Returns: + x + y if gate is None, otherwise x + y * gate + """ + if x is None and y is None: + return None + if x is None or y is None: + return x if x is not None else y + if gate is None: + return x + y + + new_y = y.clone() + selected_y = y[start_indices[1] : end_indices[1]] + selected_y = selected_y.view( + -1, self.config.action_horizon_flow, selected_y.shape[-1] + )[:, :, : self.config.dim_inputs[1]] + selected_y = selected_y.to(torch.float32) * gate + new_y[start_indices[1] : end_indices[1], : self.config.dim_inputs[1]] = ( + selected_y.view(-1, self.config.dim_inputs[1]).to(new_y.dtype) + ) + + return x + new_y + + def scatter_proprioception_embeddings( + self, input_ids, inputs_embeds, proprioception, dataset_names, agent_pos_mask + ): + if ( + proprioception is not None + and not self.config.use_state_string_representation + ): + proprioception = proprioception.to(inputs_embeds.device).to( + inputs_embeds.dtype + ) + agent_pos_mask = agent_pos_mask.to(inputs_embeds.device).to( + inputs_embeds.dtype + ) + proprioception = self.action_preprocessor.proprioception_proj( + proprioception, + dataset_names, + agent_pos_mask, + use_history=proprioception.shape[1] > 1, + ) + mask = input_ids == self.action_token_id_set["propri_token_id"] + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + proprioception_mask = mask_expanded.to(inputs_embeds.device) + + proprioception = proprioception.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter( + proprioception_mask, proprioception + ) + + return inputs_embeds + + def scatter_flow_action_embeddings( + self, input_ids, inputs_embeds, action_chunk, dataset_names, dof_mask + ): + if not self.config.use_flow_action_expert: + return inputs_embeds, None, None + adarms_cond, flow = None, None + if action_chunk is not None: + action_chunk = action_chunk.to(inputs_embeds.device) + dof_mask = dof_mask.to(inputs_embeds.device) + noisy_action_emb, flow, adarms_cond = self.action_preprocessor( + action_chunk, dataset_names, dof_mask + ) + mask = input_ids == self.action_token_id_set["action_token_id"] + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + action_mask = mask_expanded.to(inputs_embeds.device) + + noisy_action_emb = noisy_action_emb.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(action_mask, noisy_action_emb) + + return inputs_embeds, flow, adarms_cond + + @staticmethod + def _update_position_ids( + position_ids, + moe_token_types, + positional_masks, + ): + if ( + positional_masks is None + or "ar_predict_token_positions" not in positional_masks + ): + return position_ids + + new_position_ids = position_ids.clone() + ar_predict_token_positions = positional_masks["ar_predict_token_positions"] + flow_mask = moe_token_types == 1 + + start_ar_pos, end_ar_pos = find_first_last_ones(ar_predict_token_positions) + start_flow_pos, end_flow_pos = find_first_last_ones(flow_mask) + + for bs_i in range(position_ids.shape[1]): + if start_ar_pos[bs_i] != -1 and end_ar_pos[bs_i] != -1: + start_ar_ids = new_position_ids[:, bs_i, start_ar_pos[bs_i]] + start_flow_ids = new_position_ids[:, bs_i, start_flow_pos[bs_i]] + diff = start_flow_ids - start_ar_ids + new_position_ids[:, bs_i, start_flow_pos[bs_i] :] = position_ids[ + :, bs_i, start_flow_pos[bs_i] : + ] - diff.unsqueeze(-1) + + return new_position_ids + + def _update_joint_attention_mask_2d( + self, + attention_mask, + moe_token_types, + positional_masks, + ): + if attention_mask.dim() == 3: # bs, seq_len, seq_len + return attention_mask + + bs, seq_len = moe_token_types.shape[0], moe_token_types.shape[1] + # Create a lower triangular matrix as a causal mask. + causal_mask = torch.tril( + torch.ones( + (seq_len, seq_len), dtype=torch.bfloat16, device=moe_token_types.device + ) + ) + # Extended to the batch dimension. + attention_mask = causal_mask.unsqueeze(0).expand(bs, -1, -1) + + if positional_masks is not None and "padding_positions" in positional_masks: + padding_positions = positional_masks["padding_positions"] + # The padding is set to zero. + attention_mask = torch.where( + padding_positions[:, None, :], + torch.zeros_like(attention_mask), + attention_mask, + ) + # The padding is set to zero. + attention_mask = torch.where( + padding_positions[:, :, None], + torch.zeros_like(attention_mask), + attention_mask, + ) + + # Set all values ​​in the moe1 section to 1, and disable the fast section. + moe1_mask = (moe_token_types[:, :, None]) & (moe_token_types[:, None, :]) + + if ( + not self.config.causal_action_attention_mask + ): # If a causal action attention mask is not used, then all elements in the moe1 section are set to 1. + attention_mask = torch.where( + moe1_mask, torch.ones_like(attention_mask), attention_mask + ) + + if ( + positional_masks is not None + and "ar_predict_token_positions" in positional_masks + ): + ar_predict_token_positions = positional_masks["ar_predict_token_positions"] + moe1_mask = (moe_token_types[:, :, None]) & ( + ar_predict_token_positions[:, None, :] + ) + attention_mask = torch.where( + moe1_mask, torch.zeros_like(attention_mask), attention_mask + ) + + if ( + positional_masks is not None + and "valid_flow_action_positions" in positional_masks + ): + # true in moe_token_types but false in valid_flow_action_positions + nonvalid_flow_action_positions = ( + moe_token_types & ~positional_masks["valid_flow_action_positions"] + ) + attention_mask = torch.where( + nonvalid_flow_action_positions[:, None, :], + torch.zeros_like(attention_mask), + attention_mask, + ) + attention_mask = torch.where( + nonvalid_flow_action_positions[:, :, None], + torch.zeros_like(attention_mask), + attention_mask, + ) + + return attention_mask + + def _update_joint_attention_flash_mask( + self, + attention_mask, + moe_token_types, + positional_masks, + debug=False, + ): + device = moe_token_types.device + B, S = moe_token_types.shape + i32 = torch.int32 + + # ---- Return vector initialization ---- + LTS = torch.ones((B, S), device=device, dtype=i32) * S + UTE = ( + torch.arange(S, device=device, dtype=i32).unsqueeze(0).expand(B, S).clone() + ) + + # Handling padding positions + if positional_masks is not None and "padding_positions" in positional_masks: + padding_positions = positional_masks["padding_positions"] + LTS[padding_positions] = 0 + UTE[padding_positions] = S + + # Handling ar predict tokens + if ( + positional_masks is not None + and "ar_predict_token_positions" in positional_masks + ): + start_ar_pos, end_ar_pos = find_first_last_ones( + positional_masks["ar_predict_token_positions"] + ) + for bs_i in range(B): + if end_ar_pos[bs_i] != -1: + LTS[bs_i, positional_masks["ar_predict_token_positions"][bs_i]] = ( + end_ar_pos[bs_i].to(i32) + 1 + ) + + # Handling flow action bidirectional mask + flow_mask = moe_token_types == 1 + if not self.config.causal_action_attention_mask: + start_flow_pos, end_flow_pos = find_first_last_ones(flow_mask) + for bs_i in range(B): + if start_flow_pos[bs_i] != -1: + UTE[bs_i, flow_mask[bs_i]] = start_flow_pos[bs_i].to(i32) + + # Handling validate flow + if ( + positional_masks is not None + and "valid_flow_action_positions" in positional_masks + ): + flow_mask = moe_token_types == 1 + nonvalid_flow_action_positions = ( + flow_mask & ~positional_masks["valid_flow_action_positions"] + ) + if nonvalid_flow_action_positions.any(): + LTS[nonvalid_flow_action_positions] = 0 + UTE[nonvalid_flow_action_positions] = S + + LTS = LTS.unsqueeze(-1) + UTE = UTE.unsqueeze(-1) + + startend_row_indices = torch.cat([LTS, UTE], dim=-1) + # startend_row_indices = LTS + + # add num_heads dimension + startend_row_indices = startend_row_indices.unsqueeze(1) + + return startend_row_indices + + +class ActionGenerationMixin(GenerationMixin): + action_preprocessor: ActionProcessor + + def to_bfloat16_for_selected_params(self, fsdp_plugin=None, accelerator=None): + """ + Keep some model parameters as float32, and convert others to bfloat16. + - If `fsdp_plugin` exists, use FSDP v1's `mixed_precision` wrapper. + - Otherwise, directly modify the parameter dtype. + """ + + def _assign_child(root_module, dotted_name: str, new_child): + parts = dotted_name.split(".") + parent = root_module + for p in parts[:-1]: + parent = getattr(parent, p) + setattr(parent, parts[-1], new_child) + + if fsdp_plugin: + fsdp_version = getattr(fsdp_plugin, "fsdp_version", None) + if fsdp_version != 1: + raise RuntimeError("Only FSDP v1 is supported (fsdp_version=1).") + + device = getattr( + accelerator, "device", torch.device("cuda", torch.cuda.current_device()) + ) + if isinstance(device, torch.device) and device.type == "cuda": + if device.index is not None: + torch.cuda.set_device(device.index) + device_id = device.index + + # move model to device + self = self.to(device) + + # Define the mixed-precision strategy. + bf16_policy = MP( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + buffer_dtype=torch.bfloat16, + cast_forward_inputs=False, + cast_root_forward_inputs=False, + ) + + fp32_policy = MP( + param_dtype=torch.float32, + reduce_dtype=torch.float32, + buffer_dtype=torch.float32, + cast_forward_inputs=False, + cast_root_forward_inputs=False, + ) + + # Step 1️⃣: Identify the top-level ActionProcessor module and wrap it separately with FSDP (FP32). + for name, module in list(self.named_modules()): + if isinstance(module, nn.Module) and any( + k in name.lower() for k in ["action_preprocessor"] + ): + if any(True for _ in module.children()): + continue + if getattr(module, "_fsdp_wrapped", False): + continue + + print(f"[FSDP v1] wrapping module in FP32: {name}") + wrapped = FSDP( + module, + mixed_precision=fp32_policy, + sharding_strategy=torch.distributed.fsdp.ShardingStrategy.SHARD_GRAD_OP, + backward_prefetch="BACKWARD_PRE", + device_id=device_id, + use_orig_params=True, + ) + _assign_child(self, name, wrapped) + setattr(wrapped, "_fsdp_wrapped", True) + + # Step 2️⃣: The outermost layer uses unified FSDP (BF16 strategy). + print("[FSDP v1] wrapping root model with bf16 mixed precision...") + self = FSDP( + self, + mixed_precision=bf16_policy, + sharding_strategy=torch.distributed.fsdp.ShardingStrategy.SHARD_GRAD_OP, + backward_prefetch="BACKWARD_PRE", + device_id=device_id, + use_orig_params=True, + ) + + return self + + # ----------------- Non-FSDP scenarios ----------------- + else: + print("[INFO] Running manual dtype conversion (no FSDP).") + self.to(dtype=torch.float32) + + params_to_keep_float32 = [] + for name, _ in self.named_parameters(): + if any( + k in name + for k in [ + "input_layernorm", + "post_attention_layernorm", + "model.norm", + "action_preprocessor", + ] + ): + params_to_keep_float32.append(name) + + for name, param in self.named_parameters(): + if name not in params_to_keep_float32: + param.data = param.data.to(torch.bfloat16) + + return self + + def define_action_token_id(self): + action_token_list = [] + if self.action_tokenizer_type: + for i in range(self.action_tokenizer.vocab_size): + action_token_id = self.processor.tokenizer.convert_tokens_to_ids( + f"<|action_token_{i}|>" + ) + action_token_list.append(action_token_id) + + action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>") + propri_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|propri|>") + self.action_token_id_set = { + "action_token_list": action_token_list, + "propri_token_id": propri_token_id, + "action_token_id": action_token_id, + } + + def add_lora( + self, r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.1 + ): + """Add LoRA adapter""" + config = LoraConfig( + r=r, + lora_alpha=lora_alpha, + target_modules=target_modules, + lora_dropout=lora_dropout, + bias="none", + task_type="CAUSAL_LM", + ) + self.model = get_peft_model(self.model, config) + # Print trainable parameter information. + self.model.print_trainable_parameters() + + def compute_loss( + self, + hidden_states, + logits, + input_ids=None, + dataset_names=None, + labels=None, + action_chunk=None, + dof_mask=None, + flow=None, + flow_loss_mask=None, + **kwargs, + ): + if input_ids is not None: + batch_size, seq_length = input_ids.shape + + loss = 0 + cross_entropy_loss, flow_loss = None, None + + # if dataset_names is not None: + # unique_datasets_name = list(set(dataset_names)) + # channel_loss_dict = { + # dataset_name: torch.tensor(0.0, device=logits.device) + # for dataset_name in _ACTION_DATASET_NAMES + _MULTIMODAL_DATASET_NAMES + # } + # channel_loss_count_dict = { + # dataset_name: torch.tensor(0, device=logits.device) + # for dataset_name in _ACTION_DATASET_NAMES + _MULTIMODAL_DATASET_NAMES + # } + # else: + unique_datasets_name, channel_loss_dict, channel_loss_count_dict = ( + None, + None, + None, + ) + + if labels is not None: + action_accuracy = 0 + + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + non_ignored_mask = shift_labels != -100 + _cross_entropy_loss = self.loss_fct(shift_logits, shift_labels) + cross_entropy_loss = ( + _cross_entropy_loss[non_ignored_mask].mean() + if non_ignored_mask.any() + else torch.tensor(0.0, device=shift_logits.device) + ) + + # compute channel loss + _cross_entropy_loss = _cross_entropy_loss.view(batch_size, seq_length - 1) + non_ignored_mask = non_ignored_mask.view(batch_size, seq_length - 1) + for dataset_name_i in unique_datasets_name: + dataset_mask = torch.tensor( + [name == dataset_name_i for name in dataset_names], + device=logits.device, + ) + combined_mask = dataset_mask.unsqueeze(1) & non_ignored_mask + channel_loss_dict[dataset_name_i] = ( + _cross_entropy_loss[combined_mask].sum() + if combined_mask.any() + else torch.tensor(0.0, device=shift_logits.device) + ) + channel_loss_count_dict[dataset_name_i] += combined_mask.sum() + + if not torch.isnan(cross_entropy_loss): + loss += cross_entropy_loss + else: + with torch.no_grad(): + cross_entropy_loss.detach() + + # compute action token accuracy + if len(self.action_token_id_set["action_token_list"]) > 0: + shift_logits = logits[..., :-1, :].contiguous() + action_preds = shift_logits.argmax(dim=-1) + shift_labels = labels[..., 1:].contiguous() + action_mask = ( + shift_labels > self.action_token_id_set["action_token_list"][0] + ) + correct_preds = (action_preds == shift_labels) & action_mask + action_accuracy = ( + correct_preds.sum().float() / action_mask.sum().float() + ) + channel_loss_dict["action_accuracy"] = action_accuracy + + if action_chunk is not None: + action_mask = input_ids == self.action_token_id_set["action_token_id"] + if action_mask.any(): + action_hidden_states = hidden_states[action_mask].to(torch.float32) + flow = flow.reshape(-1, flow.shape[-1]) + _flow_loss = self.action_preprocessor.flow_loss( + action_hidden_states, flow, action_chunk, dof_mask, flow_loss_mask + ) + if isinstance(_flow_loss, torch.Tensor): + flow_loss = _flow_loss.mean() + loss += flow_loss * self.config.flow_loss_weight + _flow_loss = _flow_loss.view( + dof_mask.shape[0], dof_mask.shape[1], dof_mask.shape[2] + ) + + return ( + loss, + cross_entropy_loss, + flow_loss, + channel_loss_dict, + channel_loss_count_dict, + ) + + +class AttentionsSelectorMixin: + + @classmethod + def _autoset_attn_implementation( + cls, + config, + use_flash_attention_2: bool = False, + torch_dtype: Optional[torch.dtype] = None, + device_map: Optional[Union[str, Dict[str, int]]] = None, + check_device_map: bool = True, + ): + """ + Automatically checks and dispatches to a default attention implementation. In order of priority: + 1. An implementation specified in `config._attn_implementation` (due for example to the argument attn_implementation="sdpa" in from_pretrained). + 2. DEPRECATED: if use_flash_attention_2 is set to `True` and `flash_attn` is available, flash attention. (`LlamaFlashAttention` for example) + 3. SDPA implementation, if available and supported by the model type. (`LlamaSdpaAttention` for example) + 4. The default model's implementation otherwise (`LlamaAttention` for example) . + """ + # Here we use config._attn_implementation_internal to check whether the attention implementation was explicitly set by the user. + # The property `PretrainedConfig._attn_implementation` is never `None`, for backward compatibility (always fall back on "eager"). + # The `hasattr` here is used as some Transformers tests for some reason do not call PretrainedConfig __init__ (e.g. test_no_super_init_config_and_model) + requested_attn_implementation = None + if ( + hasattr(config, "_attn_implementation_internal") + and config._attn_implementation_internal is not None + ): + if ( + config._attn_implementation != "flash_attention_2" + and use_flash_attention_2 + ): + raise ValueError( + f'Both attn_implementation="{config._attn_implementation}" and `use_flash_attention_2=True` were used when loading the model, which are not compatible.' + ' We recommend to just use `attn_implementation="flash_attention_2"` when loading the model.' + ) + + if ( + not isinstance(config._attn_implementation, dict) + and config._attn_implementation + not in ["eager"] + + ALL_ATTENTION_FUNCTIONS.valid_keys() + + X2ROBOT_ATTENTION_FUNCTIONS + ): + message = f'Specified `attn_implementation="{config._attn_implementation}"` is not supported. The only possible arguments are `attn_implementation="eager"` (manual attention implementation)' + if cls._supports_flash_attn_2: + message += ', `"attn_implementation=flash_attention_2"` (implementation using flash attention 2)' + if cls._supports_sdpa: + message += ', `"attn_implementation=sdpa"` (implementation using torch.nn.functional.scaled_dot_product_attention)' + if cls._supports_flex_attn: + message += ', `"attn_implementation=flex_attention"` (implementation using torch\'s flex_attention)' + raise ValueError(message + ".") + + # If a config is passed with a preset attn_implementation, we skip the automatic dispatch and use the user-provided config, with hard checks that the requested attention implementation is available. + requested_attn_implementation = config._attn_implementation_internal + + if use_flash_attention_2: + logger.warning_once( + 'The model was loaded with use_flash_attention_2=True, which is deprecated and may be removed in a future release. Please use `attn_implementation="flash_attention_2"` instead.' + ) + config._attn_implementation = "flash_attention_2" + + if config._attn_implementation == "flash_attention_2": + cls._check_and_enable_flash_attn_2( + config, + torch_dtype=torch_dtype, + device_map=device_map, + hard_check_only=False, + check_device_map=check_device_map, + ) + elif requested_attn_implementation == "flex_attention": + config = cls._check_and_enable_flex_attn(config, hard_check_only=True) + elif ( + requested_attn_implementation in [None, "sdpa"] + and not is_torch_xla_available() + ): + # use_flash_attention_2 takes priority over SDPA, hence SDPA treated in this elif. + config = cls._check_and_enable_sdpa( + config, + hard_check_only=( + False if requested_attn_implementation is None else True + ), + ) + + if ( + torch.version.hip is not None + and config._attn_implementation == "sdpa" + and torch.cuda.device_count() > 1 + and version.parse(torch.__version__) < version.parse("2.4.1") + ): + logger.warning_once( + "Using the `SDPA` attention implementation on multi-gpu setup with ROCM may lead to performance issues due to the FA backend. Disabling it to use alternative backends." + ) + torch.backends.cuda.enable_flash_sdp(False) + elif requested_attn_implementation in ALL_ATTENTION_FUNCTIONS.valid_keys(): + config._attn_implementation = requested_attn_implementation + elif isinstance(requested_attn_implementation, dict): + config._attn_implementation = None + elif config._attn_implementation in X2ROBOT_ATTENTION_FUNCTIONS: + pass + else: + config._attn_implementation = "eager" + + config._attn_implementation_autoset = True + return config + + def _check_and_adjust_attn_implementation( + self, attn_implementation: Optional[str], is_init_check: bool = False + ) -> str: + assert ( + attn_implementation + in ["eager", "flash_attention_2", "sdpa"] + X2ROBOT_ATTENTION_FUNCTIONS + ) + return attn_implementation diff --git a/wall_x/serving/policy/utils.py b/wall_x/serving/policy/utils.py index 6973885..d13d768 100644 --- a/wall_x/serving/policy/utils.py +++ b/wall_x/serving/policy/utils.py @@ -13,6 +13,7 @@ logger = logging.getLogger(__name__) def prepare_batch( obs: Dict, processor, + normalizer_propri, camera_key: List[str], agent_pos_dim, action_dim, @@ -84,6 +85,7 @@ def prepare_batch( img = Image.fromarray((img * 255).astype(np.uint8)) processed_images.append(img) + # print("processed_images:",processed_images) # Apply smart resize to images resized_images = process_images( processed_images, image_factor, min_pixels, max_pixels @@ -109,7 +111,9 @@ def prepare_batch( action_token_id = processor.tokenizer.convert_tokens_to_ids("<|action|>") moe_token_types = inputs.input_ids == action_token_id - inputs["moe_token_types"] = moe_token_types + inputs["moe_token_types"] = torch.tensor(moe_token_types) + + # obs["dataset_names"]="libero_all" # Handle robot state/proprioception if available if "state" in obs: @@ -126,20 +130,22 @@ def prepare_batch( state = state.unsqueeze(1) # [batch, 1, state_dim] # Pad to 20 dimensions if needed (same as training) - if state.shape[-1] < 20: - padding = torch.zeros(state.shape[0], state.shape[1], 20 - state.shape[-1]) - state = torch.cat([state, padding], dim=-1) + # if state.shape[-1] < 20: + # padding = torch.zeros(state.shape[0], state.shape[1], 20 - state.shape[-1]) + # state = torch.cat([state, padding], dim=-1) # Create mask for valid dimensions agent_pos_mask = torch.ones_like(state) if state.shape[-1] > agent_pos_dim: agent_pos_mask[:, :, agent_pos_dim:] = 0 + normalizer_propri.normalize_data(state, [obs["dataset_names"]] * state.shape[0]) + inputs["proprioception"] = state inputs["agent_pos_mask"] = agent_pos_mask # Add dataset name (required by model) - inputs["dataset_names"] = obs["dataset_names"] + inputs["dataset_names"] = [obs["dataset_names"]] * state.shape[0] # Move all tensors to device for key in inputs: @@ -168,9 +174,21 @@ def process_images( """ resized_images = [] for img_pil in images: - current_width, current_height = img_pil.size + + orig_width, orig_height = img_pil.size + target_size = 256 + if target_size != -1: + # Maintain aspect ratio logic + if orig_width > orig_height: # Landscape image + new_width = target_size + new_height = int(target_size * orig_height / orig_width) + else: # Portrait image + new_height = target_size + new_width = int(target_size * orig_width / orig_height) + img_pil = img_pil.resize((new_width, new_height)) # Apply smart scaling (Qwen logic) + current_width, current_height = img_pil.size resized_height, resized_width = smart_resize( current_height, current_width, @@ -188,7 +206,7 @@ def process_images( def format_text_with_vision_tokens( instruction: str, camera_key: List[str], - predict_mode: str = "fast", + predict_mode: str = "diffusion", pred_horizon: int = 32, ) -> str: """Format text prompt with vision tokens for the model. @@ -208,7 +226,7 @@ def format_text_with_vision_tokens( image_pad_symbol = "<|image_pad|>" propri_symbol = "<|propri|>" action_symbol = "<|action|>" - # action_fast_symbol = "<|action_fast|>" + action_fast_symbol = "<|action_fast|>" # Camera name mapping camera_name_mapping = { @@ -237,9 +255,11 @@ def format_text_with_vision_tokens( f"\nPredict the next action in robot action.\nProprioception: {propri_symbol}\n" ) user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n" - assistant_output = f"{role_start_symbol}assistant\n" + assistant_output = ( + f"{role_start_symbol}assistant\n{action_fast_symbol}{role_end_symbol}\n" + ) if predict_mode == "diffusion": - assistant_output += f"{action_symbol * pred_horizon}" + assistant_output = f"{role_start_symbol}assistant\n{action_symbol * pred_horizon}{role_end_symbol}\n" complete_text = prologue + user_message + assistant_output return complete_text diff --git a/wall_x/serving/policy/wall_x_policy.py b/wall_x/serving/policy/wall_x_policy.py index ed670d4..e2d6fa9 100644 --- a/wall_x/serving/policy/wall_x_policy.py +++ b/wall_x/serving/policy/wall_x_policy.py @@ -1,12 +1,12 @@ import logging from typing import Dict, Any, List import torch +import copy import numpy as np -from transformers import AutoProcessor - from wall_x.serving.websocket_policy_server import BasePolicy from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction from wall_x.serving.policy.utils import prepare_batch +from wall_x.model.model_utils import load_wallx_processors, register_normalizers logger = logging.getLogger(__name__) @@ -25,7 +25,7 @@ class WallXPolicy(BasePolicy): camera_key: List[str], device: str = "cuda", dtype: str = "bfloat16", - predict_mode: str = "fast", + predict_mode: str = "diffusion", default_prompt: str | None = None, min_pixels: int = 4 * 28 * 28, max_pixels: int = 16384 * 28 * 28, @@ -50,21 +50,27 @@ class WallXPolicy(BasePolicy): """ logger.info(f"Loading Wall-X model from {model_path}") + self.normalizer_action, self.normalizer_propri = register_normalizers( + train_config, model_path + ) + self.model = Qwen2_5_VLMoEForAction.from_pretrained( model_path, train_config=train_config, action_tokenizer_path=action_tokenizer_path, ) + self.model.set_normalizer( + copy.deepcopy(self.normalizer_action), copy.deepcopy(self.normalizer_propri) + ) self.model.eval() self.model = self.model.to(device) - - self.model = self.model.bfloat16() + self.model.to_bfloat16_for_selected_params() # hard code the action dim to 20 for align to wall-x configuration - self.fixed_action_dim = 20 + self.fixed_action_dim = action_dim self.action_dim = action_dim - self.agent_pos_dim = agent_pos_dim + self.agent_pos_dim = action_dim self.pred_horizon = pred_horizon self.device = device self.predict_mode = predict_mode @@ -77,10 +83,14 @@ class WallXPolicy(BasePolicy): self.image_factor = image_factor self.max_length = max_length + print("predict_mode", predict_mode) + print("camera_key", camera_key) + # Load processor logger.info("Loading processor and tokenizer...") - self.processor = AutoProcessor.from_pretrained(model_path, use_fast=True) - self.processor.tokenizer.padding_side = "left" + + processors_dict = load_wallx_processors(train_config) + self.processor = processors_dict["processor"] # Action buffer for multi-step predictions self.action_buffer = [] @@ -126,6 +136,7 @@ class WallXPolicy(BasePolicy): input_batch = prepare_batch( obs, self.processor, + self.normalizer_propri, self.camera_key, self.agent_pos_dim, self.action_dim, @@ -147,7 +158,7 @@ class WallXPolicy(BasePolicy): if self.predict_mode == "fast" else self.fixed_action_dim ), - pred_horizon=self.pred_horizon, + action_horizon=self.pred_horizon, mode="predict", predict_mode=self.predict_mode, ) @@ -164,9 +175,7 @@ class WallXPolicy(BasePolicy): .to(torch.float32) .numpy() ) - - print(predicted_actions.shape) - return {"action": predicted_actions} + return {"predict_action": predicted_actions} except Exception as e: logger.error(f"Error during inference: {e}") diff --git a/wall_x/trainer/qwen_vl_act_trainer.py b/wall_x/trainer/qwen_vl_act_trainer.py index 044d781..c981e24 100644 --- a/wall_x/trainer/qwen_vl_act_trainer.py +++ b/wall_x/trainer/qwen_vl_act_trainer.py @@ -1,12 +1,14 @@ import os import gc import time +import yaml +import shutil import torch import random import numpy as np import torch.nn as nn import torch.distributed as dist - +import json from tqdm import tqdm from functools import wraps from datetime import datetime @@ -16,14 +18,17 @@ from accelerate import Accelerator from safetensors.torch import load_file from transformers.optimization import get_cosine_with_min_lr_schedule_with_warmup from transformers import AutoProcessor +from wall_x.model.action_head import Normalizer from wall_x.utils.timers import Timers from wall_x.model.qwen2_5_based import Qwen2_5_VLMoEForAction, Qwen2_5_VLConfig +from wall_x.utils.constant import action_statistic_dof as default_action_statistic_dof from wall_x.data.config import ACTION_DATASET_NAMES, MULTIMODAL_DATASET_NAMES from wall_x.data.load_lerobot_dataset import ( PreprocessedDataset, get_data_configs, load_lerobot_data, ) +import copy def timer(func): @@ -76,6 +81,35 @@ def seed_all(seed): torch.manual_seed(seed) +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"] + + if train_config.get("attn_deterministic", None) is not None: + model_config.attn_deterministic = train_config["attn_deterministic"] + model_config.vision_config.attn_deterministic = train_config[ + "attn_deterministic" + ] + print("[DEBUG] Attention is using deterministic kernel for this run!") + else: + model_config.attn_deterministic = False + model_config.vision_config.attn_deterministic = False + + return model_config + + class QwenVlAct_Trainer: """ Vision-Language-Action trainer for Qwen-VL models with robotic action prediction. @@ -145,8 +179,10 @@ class QwenVlAct_Trainer: self.dataload_config = get_data_configs(self.config["data"]) self.data_config_path = data_config_path self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False) + self.use_selective_recompute = self.config.get("use_selective_recompute", False) # Load model and initialize training components + self.load_normalizer() self.load_model() self.action_dim = sum(self.config["dof_config"].values()) @@ -190,6 +226,33 @@ class QwenVlAct_Trainer: "gradient_accumulation_steps", 1 ) + def load_normalizer(self): + if self.config.get("norm_stats_path", None): + self.print_rank0( + f"loading customized action statistic dof from {self.config['norm_stats_path']}" + ) + action_statistic_dof = json.load(open(self.config["norm_stats_path"], "r")) + else: + self.print_rank0( + "loading default action statistic dof from default_action_statistic_dof" + ) + action_statistic_dof = default_action_statistic_dof + + 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("self.normalizer_action.min: ", self.normalizer_action) + 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"), + ) + def print_rank0(self, msg, flush=True): """ Print message only on rank 0 (main process). @@ -228,7 +291,7 @@ class QwenVlAct_Trainer: self.save_checkpoint(epoch) # Validation after each epoch - self.val_loop() + # self.val_loop() self.accelerator.wait_for_everyone() # Memory cleanup @@ -532,7 +595,7 @@ class QwenVlAct_Trainer: model = model.to(torch.bfloat16) elif model_type == "qwen2_5": - config = Qwen2_5_VLConfig.from_pretrained( + model_config = Qwen2_5_VLConfig.from_pretrained( self.config["qwen_vl_act_config_path"] ) flow_loss_weight = self.config.get("flow_loss_weight", 1.0) @@ -565,21 +628,25 @@ class QwenVlAct_Trainer: # Set the customized robot configuration to ensure consistency between cross-embodiment # representations and the Wall-X action dimensionality. - Qwen2_5_VLMoEForAction._set_customized_config(self.config) + # Qwen2_5_VLMoEForAction._set_customized_config(self.config) customized_dof_config = self.config["customized_robot_config"][ "customized_dof_config" ] customized_agent_pos_config = self.config["customized_robot_config"][ "customized_agent_pos_config" ] - setattr(config, "customized_dof_config", customized_dof_config) - setattr(config, "customized_agent_pos_config", customized_agent_pos_config) + setattr(model_config, "customized_dof_config", customized_dof_config) + setattr( + model_config, "customized_agent_pos_config", customized_agent_pos_config + ) + model_config = update_model_config(self.config, model_config) model = Qwen2_5_VLMoEForAction( - config, + model_config, self.use_fast_tokenizer, self.processor, flow_loss_weight=flow_loss_weight, + use_selective_recompute=self.use_selective_recompute, ) model = model.to(torch.bfloat16) @@ -688,8 +755,10 @@ class QwenVlAct_Trainer: # Load LeRobot dataset self.dataset, self.train_num = load_lerobot_data( - self.config, - self.dataload_config.get("lerobot_config", {}), + config=self.config, + lerobot_config=self.dataload_config.get("lerobot_config", {}), + normalizer_action=copy.deepcopy(self.normalizer_action), + normalizer_propri=copy.deepcopy(self.normalizer_propri), rank=self.rank, world_size=self.world_size, ) @@ -754,8 +823,8 @@ class QwenVlAct_Trainer: renamed_weights[key] = value # Load weights into model - err = model.load_state_dict(renamed_weights, strict=False) - self.print_rank0(f"Weight loading report: {err}", flush=True) + # err = model.load_state_dict(renamed_weights, strict=False) + # self.print_rank0(f"Weight loading report: {err}", flush=True) if self.accelerator.is_main_process: self.print_rank0(f"Loaded pretrained weights from: {pretrain_weight_path}") @@ -802,36 +871,110 @@ class QwenVlAct_Trainer: self.timers.log(timers_to_log, normalizer=1) def save_checkpoint(self, epoch, step=0): - """ - Save training checkpoint. - - Args: - epoch (int): Current epoch number - step (int, optional): Current step number. Defaults to 0. - - Saves model state, optimizer state, and training progress information. - """ save_path = self.config["save_path"] if step == 0: ckpt_path = f"{save_path}/{epoch}" else: ckpt_path = f"{save_path}/{epoch}_{step}" - - # If FSDP SHARDED_STATE_DICT is used, please refer to the wall-x/workspace/README.md - # merge checkpoint section to merge the weights into a single safetensors if needed. self.accelerator.save_state(ckpt_path) + # Save random seed if self.accelerator.is_main_process: - self.processor.save_pretrained(os.path.join(ckpt_path, "processor")) + # FIXME the dataset does not have random seed now. Should the dataset set the random seed? + torch.save( + {"seed": self.seed}, os.path.join(ckpt_path, "seed.pth") + ) # seed is shared by all ranks; seed follows dataset + torch.save( + {"global_step": self.global_step}, + os.path.join(ckpt_path, "global_step.pth"), + ) + torch.save( + {"current_epoch": epoch}, os.path.join(ckpt_path, "current_epoch.pth") + ) - # Save current iteration steps for dataset resuming - if step != 0: + # Save configuration in YAML format + config_path = os.path.join(ckpt_path, "config.yml") + with open(config_path, "w", encoding="utf-8") as f: + yaml.dump( + self.config, + f, + default_flow_style=False, + allow_unicode=True, + indent=2, + sort_keys=False, + ) + + pretrained_dir = self.config.get("pretrained_qwen_vl_path", None) + if pretrained_dir is not None: + files_to_copy = [ + "preprocessor_config.json", + "tokenizer_config.json", + "tokenizer.json", + "vocab.json", + ] + + for filename in files_to_copy: + src = os.path.join(pretrained_dir, filename) + dst = os.path.join(ckpt_path, filename) + + if os.path.exists(src): + shutil.copy(src, dst) + print(f"[Checkpoint] Copied {filename} to {ckpt_path}") + else: + print(f"[Checkpoint] WARNING: {src} not found, skip copying.") + + act_config_path = self.config.get("qwen_vl_act_config_path", None) + if act_config_path is not None: + dst = os.path.join(ckpt_path, "config.json") + + if os.path.exists(act_config_path): + shutil.copy(act_config_path, dst) + print(f"[Checkpoint] Copied act config to {dst}") + else: + print( + f"[Checkpoint] WARNING: {act_config_path} not found, skipping." + ) + # Save normalizer + torch.save( + self.normalizer_action.state_dict(), + os.path.join(ckpt_path, "normalizer_action.pth"), + ) + torch.save( + self.normalizer_propri.state_dict(), + os.path.join(ckpt_path, "normalizer_propri.pth"), + ) + + # Save current iter steps + if step != 0: # step==0, no need for dataset resume _rank = self.accelerator.process_index - if isinstance(self.dataset, PreprocessedDataset): + if self.data_config["multimodal_data_ratio"] != 1: torch.save( - {"epoch": epoch, "step": step}, + { + "episode_start_index": self.dataset.primary_pool_start_index.value + }, + os.path.join(ckpt_path, f"episode_start_index_rank_{_rank}.pth"), + ) + torch.save( + { + "multimodal_episode_start_index": self.dataset.secondary_pool_start_index.value + }, os.path.join( - ckpt_path, f"epoch_{epoch}_step_{step}_rank_{_rank}.pth" + ckpt_path, f"multimodal_episode_start_index_rank_{_rank}.pth" + ), + ) + else: + torch.save( + { + "episode_start_index": self.dataset.secondary_pool_start_index.value + }, + os.path.join(ckpt_path, f"episode_start_index_rank_{_rank}.pth"), + ) + torch.save( + { + "multimodal_episode_start_index": self.dataset.primary_pool_start_index.value + }, + os.path.join( + ckpt_path, f"multimodal_episode_start_index_rank_{_rank}.pth" ), ) @@ -841,7 +984,6 @@ class QwenVlAct_Trainer: Handles both full checkpoint loading and model-only loading based on configuration. """ - checkpoint_path = self.config["resume"]["ckpt"] if self.config.get("resume", {}).get("load_ckpt_only", False): if self.config.get("FSDP2", False): @@ -862,9 +1004,59 @@ class QwenVlAct_Trainer: self.model.load_state_dict(new_state_dict, strict=False) else: # Load full checkpoint including optimizer and scheduler states - self.accelerator.load_state(checkpoint_path) - self.print_rank0(f"\033[32mResumed from checkpoint: {checkpoint_path}\033[0m") + # self.accelerator.load_state(checkpoint_path) + state_dict = load_file( + self.config["resume"]["ckpt"] + "/model.safetensors", device="cpu" + ) + + filtered_state_dict = { + k: v + for k, v in state_dict.items() + if not k.startswith("action_preprocessor.normalizer") + } + + if self.config["resume"].get("try_harder", False): + new_state_dict = {} + for name, param in filtered_state_dict.items(): + if name in self.model.state_dict(): + if param.size() == self.model.state_dict()[name].size(): + new_state_dict[name] = param + else: + size_0 = param.size() + size_1 = self.model.state_dict()[name].size() + new_state_dict[name] = self.model.state_dict()[name] + slices = [ + slice(0, min(old_dim, new_dim)) + for old_dim, new_dim in zip(size_0, size_1) + ] + new_state_dict[name][slices] = param[slices] + self.print_rank0( + f"Not match key: {name}, required shape: {size_1}, loaded shape: {size_0}, new shape: {new_state_dict[name].size()}" + ) + elif "module." + name in self.model.state_dict(): + name = "module." + name + if param.size() == self.model.state_dict()[name].size(): + new_state_dict[name] = param + else: + size_0 = param.size() + size_1 = self.model.state_dict()[name].size() + new_state_dict[name] = self.model.state_dict()[name] + slices = [ + slice(0, min(old_dim, new_dim)) + for old_dim, new_dim in zip(size_0, size_1) + ] + new_state_dict[name][slices] = param[slices] + self.print_rank0( + f"Not match key: {name}, required shape: {size_1}, loaded shape: {size_0}, new shape: {new_state_dict[name].size()}" + ) + else: + self.print_rank0(f"Not used parameter: {name}") + err = self.model.load_state_dict(new_state_dict, strict=False) + else: + err = self.model.load_state_dict(filtered_state_dict, strict=False) + + self.print_rank0(f"err in load model: {err}", err) def _load_fsdp_state_dict_with_distribute_tensor(self): diff --git a/wall_x/utils/timers.py b/wall_x/utils/timers.py index 32b03b6..4233d19 100644 --- a/wall_x/utils/timers.py +++ b/wall_x/utils/timers.py @@ -1,35 +1,101 @@ import time -import torch from torch.cuda import nvtx from abc import ABC, abstractmethod from typing import List +import torch +from functools import wraps +from contextlib import nullcontext +import os +ENABLE_PERFORMANCE_TIMING = ( + os.environ.get("ENABLE_PERFORMANCE_TIMING", "True").lower() == "true" +) + +ENABLE_CUDA_SYNC_IN_TIMER = ( + os.environ.get("ENABLE_CUDA_SYNC_IN_TIMER", "False").lower() == "true" +) + + +class ScopeTimerContext: + def __init__(self, msg): + self.msg = msg + + def __enter__(self): + if ENABLE_CUDA_SYNC_IN_TIMER and torch.cuda.is_available(): + torch.cuda.synchronize() + self.start_time = time.perf_counter() + return self + + def __exit__(self, exc_type, exc_value, traceback): + if ENABLE_CUDA_SYNC_IN_TIMER and torch.cuda.is_available(): + torch.cuda.synchronize() + end_time = time.perf_counter() + cost_ms = (end_time - self.start_time) * 1e3 + print(f"\033[92m{self.msg} took {cost_ms:.3f} ms to execute\033[0m") + + +ScopeTimer = ScopeTimerContext if ENABLE_PERFORMANCE_TIMING else nullcontext + + +def timer(func, msg=None): + """ + Decorator to measure function execution time. + + Args: + func: Function to be timed + + Returns: + Wrapped function with timing functionality + """ + + if msg is None: + msg = func.__name__ + else: + msg = f"{func.__name__:} {msg}" + + @wraps(func) + def wrapper(*args, **kwargs): + with ScopeTimer(msg): + result = func(*args, **kwargs) + + return result + + return wrapper + + +# Helper functions to check for distributed environment def _is_distributed(): + """Checks if the current environment is set up for distributed training.""" return torch.distributed.is_available() and torch.distributed.is_initialized() def _get_world_size(): + """Safely retrieves the world size (number of processes).""" if _is_distributed(): return torch.distributed.get_world_size() return 1 def _get_rank(): + """Safely retrieves the rank of the current process.""" if _is_distributed(): return torch.distributed.get_rank() return 0 def _barrier(group=None): + """Safely executes a distributed barrier to synchronize processes.""" if _is_distributed(): torch.distributed.barrier(group=group) +# Dynamically set the all_gather function if torch.distributed.is_available(): try: dist_all_gather_func = torch.distributed.all_gather_into_tensor except AttributeError: + # Fallback to standard all_gather if all_gather_into_tensor is missing dist_all_gather_func = torch.distributed.all_gather else: dist_all_gather_func = None @@ -144,7 +210,7 @@ class Timer(TimerBase): """ self._barrier_group = barrier_group - def start(self, barrier=False, nvtx_push=False): + def start(self, barrier=False, nvtx_push=False, sync=False): """Start the timer. Args: @@ -153,7 +219,7 @@ class Timer(TimerBase): assert not self._started, "timer has already been started" if barrier: _barrier(group=self._barrier_group) - if torch.cuda.is_available(): + if torch.cuda.is_available() and sync: torch.cuda.synchronize() self._start_time = time.time() self._started = True @@ -272,10 +338,13 @@ class Timers: def _get_elapsed_time_all_ranks(self, names, reset, barrier): """Returns elapsed times of timers in names. + For single-node/single-GPU cases, directly returns the time for the current rank. + For distributed cases, maintains the existing all_gather logic. + Args: names (List[str]): list of timer names reset (bool): reset the timer after recording the elapsed time - barrier (bool): if set, do a global barrier before time measurments + barrier (bool): if set, do a global barrier before time measurements Returns: torch.tensor: Tensor of size [world_size, len(names)] with times in float. @@ -288,6 +357,7 @@ class Timers: world_size = _get_world_size() rank = _get_rank() + # Create device tensor if torch.cuda.is_available(): device = torch.cuda.current_device() else: @@ -297,16 +367,19 @@ class Timers: (world_size, len(names)), dtype=torch.float, device=device ) + # Fill timing data for the current rank for i, name in enumerate(names): if name in self._timers: rank_name_to_time[rank, i] = self._timers[name].elapsed(reset=reset) + # Return directly for single-node; perform all_gather for distributed setup if world_size > 1 and _is_distributed() and dist_all_gather_func is not None: try: dist_all_gather_func( rank_name_to_time.view(-1), rank_name_to_time[rank, :].view(-1) ) except Exception as e: + # If all_gather fails, print a warning and proceed with single rank timing print(f"Warning: all_gather failed: {e}. Using single rank timing.") return rank_name_to_time @@ -340,13 +413,17 @@ class Timers: world_size = _get_world_size() if world_size == 1: + # Simplified output for single-node setup output_string = "time (ms):" for name in name_to_min_max_time: - _, max_time = name_to_min_max_time[name] + _, max_time = name_to_min_max_time[ + name + ] # min and max are identical for a single rank output_string += "\n {}: {:.2f}".format( (name + " ").ljust(48, "."), max_time ) else: + # Maintain original output format for multi-node setup if max_only: output_string = "max time across ranks (ms):" else: diff --git a/workspace/lerobot_example/libero/config_qact_libero_from_vlm.yml b/workspace/lerobot_example/libero/config_qact_libero_from_vlm.yml index 3ab0a64..21198e4 100644 --- a/workspace/lerobot_example/libero/config_qact_libero_from_vlm.yml +++ b/workspace/lerobot_example/libero/config_qact_libero_from_vlm.yml @@ -31,38 +31,26 @@ epoch_save_interval: 1 # Robot configuration - Define degrees of freedom for each component dof_config: - follow_left_ee_cartesian_pos: 3 # Left end-effector Cartesian position - follow_left_ee_rotation: 3 # Left end-effector rotation - follow_left_gripper: 1 # Left gripper control - follow_right_ee_cartesian_pos: 3 # Right end-effector Cartesian position - follow_right_ee_rotation: 3 # Right end-effector rotation - follow_right_gripper: 1 # Right gripper control - head_actions: 2 # Head/camera movement - height: 1 # Mobile base height control - car_pose: 3 # Mobile base pose (x, y, theta) + master_right_ee_cartesian_pos: 3 # Right end-effector Cartesian position + master_right_ee_rotation: 3 # Right end-effector rotation + master_right_gripper: 1 # Right gripper control # Agent proprioception configuration (typically matches DOF config) agent_pos_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 -norm_stats_path: "wall-x/workspace/lerobot_example/libero/libero_norm_stats.json" +norm_stats_path: "/path/to/libero_norm_stats.json" enable_customized_robot_config: true customized_robot_config: - name: "physical-intelligence/libero" + name: "libero_all" customized_dof_config: "panda_action_eef_with_gripper": 7 customized_agent_pos_config: - "panda_state_eef_with_gripper": 8 + "panda_state_eef_with_gripper": 7 # Checkpoint resuming configuration # resume: @@ -75,7 +63,7 @@ data: # LeRobot dataset configuration lerobot_config: - repo_id: "physical-intelligence/libero" + repo_id: "libero_all" root: null episodes: null image_transforms: null @@ -86,31 +74,19 @@ data: download_videos: true video_backend: null - action_horizon: 32 + action_horizon: 10 train_test_split: 0.95 # Action keys for observation and prediction 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 - - head_actions - - height - - car_pose predict_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 - - head_actions - - height - - car_pose + - master_right_ee_cartesian_pos + - master_right_ee_rotation + - master_right_gripper # Image resolution configuration for different camera views resolution: