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

Fix normalizer (#57)

Browse Source 
* fix normalizer

* fix val

* update compute stats

* delete norm

* update readme

* minor fix

* fix action normalizer

* fix

* fix

* update

* update

* update

* update

* update

* lint

* lint

* lint

* lint
This commit is contained in:
Lufang Chen
2025-10-24 17:29:12 +08:00
committed by GitHub
parent 35399d187a
commit d821b0cb26
14 changed files with 837 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wall_x/data/load_lerobot_dataset.py
+69 -37
View File
@@ -17,17 +17,10 @@ from wall_x.data.utils import (
)
from transformers import AutoProcessor
from .utils import load_norm_stats, KEY_MAPPINGS
T_co = TypeVar("T_co", covariant=True)
CAMERA_KEY_MAPPINGS = {
"lerobot/aloha_mobile_cabinet": {
"observation.images.cam_high": "face_view",
"observation.images.cam_left_wrist": "left_wrist_view",
"observation.images.cam_right_wrist": "right_wrist_view",
},
}
# Abstract class for dataset
class Dataset(Protocol[T_co]):
@@ -46,6 +39,8 @@ class PreprocessedDataset(Dataset[T_co]):
dataset,
config,
dataload_config,
norm_stats,
lerobot_config,
seed=42,
rank=0,
world_size=1,
@@ -72,6 +67,8 @@ 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.lerobot_config = lerobot_config
self.data_config = X2RDataProcessingConfig().update(
train_test_split=self.dataload_config["train_test_split"],
@@ -82,7 +79,9 @@ class PreprocessedDataset(Dataset[T_co]):
priority_order=self.dataload_config.get("priority_order", None),
)
self._cam_key_mapping = CAMERA_KEY_MAPPINGS[self.hf_dataset.meta.repo_id]
self._cam_key_mapping = KEY_MAPPINGS[self.hf_dataset.meta.repo_id]["camera"]
self._state_key_mapping = KEY_MAPPINGS[self.hf_dataset.meta.repo_id]
self._action_key_mapping = KEY_MAPPINGS[self.hf_dataset.meta.repo_id]
def _vision_preprocess(self, frames):
processed_frames = []
@@ -124,14 +123,14 @@ class PreprocessedDataset(Dataset[T_co]):
def __getitem__(self, index):
data = self._dataset[index]
image_inputs, h, w, resize_h, resize_w = self._vision_preprocess(data)
agent_pos = data["observation.state"]
action = data["action"]
agent_pos = data[self._state_key_mapping["state"]]
action = data[self._action_key_mapping["action"]]
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,
33 - 1,
self.dataload_config.get("action_horizon", 33) - 1,
frame_index,
self.data_config.priority_order,
self._cam_key_mapping,
@@ -189,7 +188,7 @@ 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.hf_dataset.meta.stats
self.config, self.dataload_config, self.norm_stats, self.lerobot_config
),
pin_memory=True, # Enable for GPU training
persistent_workers=num_workers > 0, # Only if num_workers > 0
@@ -225,7 +224,7 @@ class PreprocessedDataset(Dataset[T_co]):
sampler=sampler,
num_workers=num_workers,
collate_fn=DataCollator(
self.config, self.dataload_config, self.hf_dataset.meta.stats
self.config, self.dataload_config, self.norm_stats, self.lerobot_config
),
pin_memory=True,
persistent_workers=num_workers > 0,
@@ -241,13 +240,16 @@ class DataCollator:
_processor_cache = {}
_action_tokenizer_cache = {}
def __init__(self, config, dataload_config, stats):
def __init__(self, config, dataload_config, stats, lerobot_config):
self.config = config
self.dataload_config = dataload_config
self.stats = stats
self.min_stat = stats["action"]["min"]
self.max_stat = stats["action"]["max"]
self.delta = self.max_stat - self.min_stat
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.lerobot_config = lerobot_config
self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False)
self.load_processor()
@@ -271,10 +273,11 @@ class DataCollator:
if self.config.get("padding_side", "left") == "left":
processor.tokenizer.padding_side = "left"
new_tokens = ["<|propri|>", "<|action|>"]
processor.tokenizer.add_tokens(new_tokens)
if self.use_fast_tokenizer and self.config.get("model_type") == "qwen2_5":
action_tokenizer = self._action_tokenizer_cache[action_tokenizer_path]
new_tokens = ["<|propri|>", "<|action|>"]
new_tokens += [
new_tokens = [
f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size)
]
processor.tokenizer.add_tokens(new_tokens)
@@ -301,7 +304,6 @@ class DataCollator:
"""
Normalize action data using min-max normalization.
"""
delta = torch.from_numpy(delta)
delta = torch.where(delta == 0, torch.ones_like(delta), delta)
x = (action - min_stat) / delta
x = x * 2 - 1
@@ -318,7 +320,9 @@ class DataCollator:
agent_pos = agent_pos.unsqueeze(1)
agent_pos_mask = (~torch.isnan(agent_pos)).float()
agent_pos.nan_to_num_(nan=0.0)
agent_pos = self._normalize(agent_pos, self.min_stat, self.delta)
agent_pos = self._normalize(
agent_pos, self.state_min_stat, self.state_delta
)
if agent_pos.shape[-1] != 20:
agent_pos = torch.cat(
[
@@ -350,7 +354,9 @@ class DataCollator:
action = action.unsqueeze(1)
dof_mask = (~torch.isnan(action)).float()
action.nan_to_num_(nan=0.0)
action = self._normalize(action, self.min_stat, self.delta)
action = self._normalize(
action, self.action_min_stat, self.action_delta
)
if action.shape[-1] != 20:
action = torch.cat(
[
@@ -393,7 +399,7 @@ class DataCollator:
additional_inputs["text"],
additional_inputs["action_chunk"],
self.train_action_tokenizer if self.use_fast_tokenizer else None,
["x2_normal"] * additional_inputs["text"].__len__(),
[self.lerobot_config["repo_id"]] * additional_inputs["text"].__len__(),
additional_inputs["dof_mask"],
)
@@ -415,7 +421,9 @@ class DataCollator:
inputs.update(additional_inputs)
inputs["dataset_names"] = ["x2_normal"] * inputs["action_chunk"].shape[0]
inputs["dataset_names"] = [self.lerobot_config["repo_id"]] * inputs[
"action_chunk"
].shape[0]
return inputs
@@ -447,18 +455,24 @@ def load_lerobot_data(
dataload_config = get_data_configs(config["data"])
# repo_id = "lerobot/aloha_mobile_cabinet"
repo_id = lerobot_config.get("repo_id", "lerobot/aloha_mobile_cabinet")
repo_id = lerobot_config.get("repo_id", None)
assert repo_id is not None, "repo id is required"
root = lerobot_config.get("root", None)
meta_info = LeRobotDatasetMetadata(repo_id)
meta_info = LeRobotDatasetMetadata(repo_id, root=root)
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)
delta_timestamps = {
# action chunk
"action": [
KEY_MAPPINGS[repo_id]["action"]: [
t / dataset_fps
for t in range(dataload_config.get("action_horizon", 32) - 1)
for t in range(dataload_config.get("action_horizon", 33) - 1)
],
}
batch_size = config.get("batch_size_per_gpu", 8)
@@ -486,6 +500,8 @@ def load_lerobot_data(
train_dataset,
config,
dataload_config,
norm_stats,
lerobot_config,
seed=seed,
rank=rank,
world_size=world_size,
@@ -567,11 +583,15 @@ def get_data_configs(config):
class TestDataset(PreprocessedDataset):
def __init__(self, dataset, config, dataload_config, seed=42):
def __init__(
self, dataset, config, dataload_config, norm_stats, lerobot_config, seed=42
):
super().__init__(
dataset,
config,
dataload_config,
norm_stats,
lerobot_config,
seed=seed,
rank=0,
world_size=1,
@@ -587,7 +607,7 @@ class TestDataset(PreprocessedDataset):
self,
batch_size=1,
collate_fn=DataCollator(
self.config, self.dataload_config, self.hf_dataset.meta.stats
self.config, self.dataload_config, self.norm_stats, self.lerobot_config
),
)
@@ -614,29 +634,41 @@ def load_test_dataset(
# Set seed for reproducibility
torch.manual_seed(seed)
dataset_fps = 50
repo_id = lerobot_config.get("repo_id", None)
assert repo_id is not None, "repo id is required"
root = lerobot_config.get("root", None)
meta_info = LeRobotDatasetMetadata(repo_id, root=root)
dataset_fps = meta_info.fps
dataload_config = get_data_configs(config["data"])
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
"action": [
KEY_MAPPINGS[repo_id]["action"]: [
t / dataset_fps
for t in range(dataload_config.get("action_horizon", 32) - 1)
for t in range(dataload_config.get("action_horizon", 33) - 1)
],
}
repo_id = lerobot_config.get("repo_id", "lerobot/aloha_mobile_cabinet")
dataset = LeRobotDataset(
repo_id,
episodes=[episode],
delta_timestamps=delta_timestamps,
video_backend="pyav",
root=root,
)
print(f"Selected episodes: {dataset.episodes}")
print(f"Number of episodes selected: {dataset.num_episodes}")
print(f"Number of frames selected: {dataset.num_frames}")
dataset = TestDataset(dataset, config, dataload_config, seed=seed)
dataset = TestDataset(
dataset, config, dataload_config, norm_stats, lerobot_config, seed=seed
)
return dataset