huangyanpei/VLA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add mot (#83)

Browse Source 
* add mot

* update libero example

* translate zh to en

* fix load model from hf

* lint

* lint

---------

Co-authored-by: yangping <yangping@x2robot.com>
This commit is contained in:
suolyer
2026-02-03 11:35:25 +08:00
committed by GitHub
co-authored by yangping
parent 05b6d8dcf7
commit d18fa65fa1
26 changed files with 8509 additions and 1179 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/compute_norm_stats.py
+173 -69
View File
@@ -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()
scripts/draw_openloop_plot.py
+19 -8
View File
@@ -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,
)
scripts/infer_libero.py
+222
View File
@@ -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.")
scripts/infer_robochallenge.py
+1164
View File
File diff suppressed because it is too large Load Diff
scripts/vqa_inference.py
+15 -5
View File
@@ -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"