Fix normalizer (#57)

* 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
2025-10-24 17:29:12 +08:00
parent 35399d187a
commit d821b0cb26
14 changed files with 837 additions and 119 deletions
@@ -36,6 +36,8 @@ ACTION_DATASET_NAMES = [
    "taco_play",
    "utaustin_mutex",
    "viola",
+    "physical-intelligence/libero",
+    "lerobot/aloha_mobile_cabinet",
 ]

 # Supported multimodal datasets
@@ -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
@@ -11,7 +11,28 @@ import random
 from collections import OrderedDict
 from typing import List, Dict, Any, Optional, Union, Tuple
 from transformers import BatchFeature
+from dataclasses import dataclass
+import json

+KEY_MAPPINGS = {
+    "lerobot/aloha_mobile_cabinet": {
+        "camera": {
+            "observation.images.cam_high": "face_view",
+            "observation.images.cam_left_wrist": "left_wrist_view",
+            "observation.images.cam_right_wrist": "right_wrist_view",
+        },
+        "state": "observation.state",
+        "action": "action",
+    },
+    "physical-intelligence/libero": {
+        "camera": {
+            "image": "face_view",
+            "wrist_image": "left_wrist_view",
+        },
+        "state": "state",
+        "action": "actions",
+    },
+}

 CAMERA_NAME_MAPPING = {
    "face_view": "front view",
@@ -609,3 +630,35 @@ def replace_action_token(
        text = [t.replace("<|action_fast|><|im_end|>\n", "") for t in text]

    return text
+
+
+@dataclass
+class NormStats:
+    min: torch.Tensor
+    max: torch.Tensor
+    delta: torch.Tensor
+
+
+def load_norm_stats(norm_stats_path, dataset_name):
+    with open(norm_stats_path, "r") as f:
+        norm_stats = json.load(f)
+    action_key = KEY_MAPPINGS[dataset_name]["action"]
+    state_key = KEY_MAPPINGS[dataset_name]["state"]
+    q01 = torch.tensor(norm_stats["norm_stats"][action_key]["q01"])
+    q99 = torch.tensor(norm_stats["norm_stats"][action_key]["q99"])
+    delta = q99 - q01
+    action_norm_stats = NormStats(
+        min=q01,
+        max=q99,
+        delta=delta,
+    )
+    q01 = torch.tensor(norm_stats["norm_stats"][state_key]["q01"])
+    q99 = torch.tensor(norm_stats["norm_stats"][state_key]["q99"])
+    delta = q99 - q01
+    state_norm_stats = NormStats(
+        min=q01,
+        max=q99,
+        delta=delta,
+    )
+
+    return {"action": action_norm_stats, "state": state_norm_stats}