add mot (#83)
* 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
+173
-69
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user