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
@@ -3,6 +3,7 @@ import torch
 import torch.nn as nn
 from torch.distributions import Beta
 from wall_x.utils.constant import action_statistic_dof
+import logging


 class Normalizer(nn.Module):
@@ -14,6 +15,16 @@ class Normalizer(nn.Module):
    normalization to map actions to the [-1, 1] range.
    """

+    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
+
    def __init__(self, action_statistic_dof, dof_config):
        """
        Initialize the normalizer with robot-specific action statistics.
@@ -25,6 +36,8 @@ class Normalizer(nn.Module):
        super(Normalizer, self).__init__()

        action_statistic = {}
+        # hard code the action dimension to 20
+        action_dim = 20

        # Process statistics for each robot
        for robot_name in action_statistic_dof.keys():
@@ -39,11 +52,17 @@ class Normalizer(nn.Module):
                    all_dof_delta.extend(action_statistic_dof[robot_name][k]["delta"])
                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"
+                    )
                    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)
+            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

@@ -61,7 +80,7 @@ class Normalizer(nn.Module):
            }
        )

-    def normalize_data(self, xs, dataset_names):
+    def normalize_data(self, xs, dataset_names, dof_mask=None):
        """
        Normalize action data to [-1, 1] range using robot-specific statistics.

@@ -75,10 +94,19 @@ class Normalizer(nn.Module):
        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 in zip(xs, dataset_names):
+        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 - self.min[dataset_name]) / (self.delta[dataset_name])
+            x = (x - action_space_min) / (action_space_delta)
            # Scale to [-1, 1] range
            x = x * 2 - 1
            # Clamp to ensure bounds
@@ -210,9 +238,21 @@ class ActionProcessor(nn.Module):
        self.hidden_size = config.hidden_size

        # Initialize data normalizers for actions and proprioception
-        self.normalizer_action = Normalizer(action_statistic_dof, config.dof_config)
+        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.agent_pos_config
+            action_statistic_dof,
+            (
+                config.customized_agent_pos_config
+                if hasattr(config, "customized_agent_pos_config")
+                else config.agent_pos_config
+            ),
        )

        # Proprioception projection layer (includes history/current state)
@@ -44,7 +44,9 @@ from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl import (
    Qwen2_5_VLSdpaAttention,
 )
 from wall_x.data.config import ACTION_DATASET_NAMES, MULTIMODAL_DATASET_NAMES
-
+from wall_x.utils.constant import action_statistic_dof
+from wall_x.data.utils import load_norm_stats
+from pprint import pprint

 logger = logging.get_logger(__name__)

@@ -744,10 +746,77 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
    config_class = Qwen2_5_VLConfig
    _no_split_modules = ["Qwen2_5_VLDecoderLayer_with_MoE", "Qwen2_5_VLVisionBlock"]

+    @classmethod
+    def _set_customized_config(cls, config):
+        """
+        Processing norm_stats.json and reconstruct the DoF mapping
+        """
+        dataload_config = config["data"]
+        if not dataload_config.get("use_lerobot", False):
+            raise NotImplementedError(
+                "Not implemented for non-lerobot dataset currently"
+            )
+
+        enable_customized_robot_config = config.get(
+            "enable_customized_robot_config", False
+        )
+        assert (
+            enable_customized_robot_config
+        ), "enable_customized_robot_config must be true when use lerobot dataset"
+
+        customized_dof_config = config["customized_robot_config"][
+            "customized_dof_config"
+        ]
+        customized_agent_pos_config = config["customized_robot_config"][
+            "customized_agent_pos_config"
+        ]
+        norm_stats_path = config["norm_stats_path"]
+        norm_stats = load_norm_stats(
+            norm_stats_path, config["data"]["lerobot_config"]["repo_id"]
+        )
+        action_min = norm_stats["action"].min.numpy().tolist()
+        action_delta = norm_stats["action"].delta.numpy().tolist()
+        state_min = norm_stats["state"].min.numpy().tolist()
+        state_delta = norm_stats["state"].delta.numpy().tolist()
+
+        name = config["customized_robot_config"]["name"]
+
+        dof_key = []
+        agent_pos_key = []
+        dof_value = []
+        agent_pos_value = []
+        stats_dict = {}
+        for k, v in customized_dof_config.items():
+            dof_key.append(k)
+            dof_value.append(v)
+        for k, v in customized_agent_pos_config.items():
+            agent_pos_key.append(k)
+            agent_pos_value.append(v)
+
+        dof_idx = np.array([0] + dof_value).cumsum()
+        for i in range(len(dof_idx) - 1):
+            stats_dict[dof_key[i]] = {
+                "min": action_min[dof_idx[i] : dof_idx[i + 1]],
+                "delta": action_delta[dof_idx[i] : dof_idx[i + 1]],
+            }
+
+        agent_pos_idx = np.array([0] + agent_pos_value).cumsum()
+        for i in range(len(agent_pos_idx) - 1):
+            stats_dict[agent_pos_key[i]] = {
+                "min": state_min[agent_pos_idx[i] : agent_pos_idx[i + 1]],
+                "delta": state_delta[agent_pos_idx[i] : agent_pos_idx[i + 1]],
+            }
+
+        action_statistic_dof[name] = stats_dict
+
+        print("Customized robot config added")
+        pprint(action_statistic_dof)
+
    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_path,
+        train_config,
        config_path=None,
        processor_path=None,
        action_tokenizer_path=None,
@@ -766,7 +835,6 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
        Returns:
            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)
@@ -776,6 +844,18 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
                action_tokenizer_path, trust_remote_code=True
            )

+        # 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)
+
        # Initialize model with configuration and processor
        model = cls(config, processor=processor, **kwargs)

@@ -789,6 +869,14 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
        state_dict = {}
        for file in safetensor_files:
            sd = load_file(file, device="cpu")
+            # filter normalizer statistic params
+            del_keys = []
+            for key in sd.keys():
+                if "action_preprocessor.normalizer" in key:
+                    print(f"filter load model weight {key}")
+                    del_keys.append(key)
+            for key in del_keys:
+                del sd[key]
            state_dict.update(sd)

        model.load_state_dict(state_dict, strict=False)
@@ -860,11 +948,14 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
        """
        # Create list of fast action token IDs
        fast_action_token_list = []
-        for i in range(self.processor.tokenizer.init_kwargs["action_token_vocab_size"]):
-            action_token_id = self.processor.tokenizer.convert_tokens_to_ids(
-                f"<|action_token_{i}|>"
-            )
-            fast_action_token_list.append(action_token_id)
+        if self.use_fast_tokenizer:
+            for i in range(
+                self.processor.tokenizer.init_kwargs["action_token_vocab_size"]
+            ):
+                action_token_id = self.processor.tokenizer.convert_tokens_to_ids(
+                    f"<|action_token_{i}|>"
+                )
+                fast_action_token_list.append(action_token_id)

        # Get special action token IDs
        action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>")