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
@@ -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}