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

scripts/compute_norm_stats.py

@@ -0,0 +1,76 @@
+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
+import numpy as np
+import argparse
+
+
+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 load_lerobot_dataset(repo_id, action_horizon, args):
+    dataset_meta = LeRobotDatasetMetadata(repo_id)
+    dataset = LeRobotDataset(
+        repo_id,
+        delta_timestamps={
+            key: [t / dataset_meta.fps for t in range(action_horizon)]
+            for key in [KEY_MAPPINGS[repo_id]["action"]]
+        },
+    )
+    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
+
+
+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)
+    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, 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)

scripts/draw_openloop_plot.py

@@ -1,22 +1,13 @@
 import os
 import yaml
 import torch
+import argparse
+from tqdm import tqdm
 import matplotlib.pyplot as plt
 from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction
 from wall_x.data.load_lerobot_dataset import load_test_dataset, get_data_configs
 
 
-model_path = "path/to/model"
-action_tokenizer_path = "path/to/action_tokenizer"
-save_dir = "path/to/plot"
-model = Qwen2_5_VLMoEForAction.from_pretrained(
-    model_path, action_tokenizer_path=action_tokenizer_path
-)
-model.eval()
-model = model.to("cuda")
-model = model.bfloat16()
-
-
 def load_config(config_path):
     """Load configuration from YAML file."""
     with open(config_path, "r") as f:
@@ -27,60 +18,97 @@ def load_config(config_path):
     return config
 
 
-# get test dataloader
-path = "path/to/config"
-config = load_config(path)
-dataload_config = get_data_configs(config["data"])
-lerobot_config = dataload_config.get("lerobot_config", {})
-dataset = load_test_dataset(config, lerobot_config, seed=42)
-dataloader = dataset.get_dataloader()
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--pred_horizon", type=int, default=32)
+    parser.add_argument("--origin_action_dim", type=int, default=7)
+    args = parser.parse_args()
 
-total_frames = len(dataloader)
+    origin_action_dim = args.origin_action_dim
+    pred_horizon = args.pred_horizon
 
-pred_horizon = 32
-action_dim = 14
-gt_traj = torch.zeros((total_frames, action_dim))
-pred_traj = torch.zeros((total_frames, action_dim))
+    # get train config
+    model_path = "/path/to/model"
+    action_tokenizer_path = "/path/to/action/tokenizer"
+    save_dir = "/path/to/save/dir"
+    path = "/path/to/train/config"
+    config = load_config(path)
 
-for idx, batch in enumerate(dataloader):
-    if idx % pred_horizon == 0 and idx + pred_horizon < total_frames:
-        batch = batch.to("cuda")
-        with torch.no_grad():
-            outputs = model(
-                **batch,
-                action_dim=action_dim,
-                pred_horizon=pred_horizon,
-                mode="predict",
-                predict_mode="fast",
+    # load model with customized robot config
+    model = Qwen2_5_VLMoEForAction.from_pretrained(
+        model_path, train_config=config, action_tokenizer_path=action_tokenizer_path
+    )
+    model.eval()
+    model = model.to("cuda")
+    model = model.bfloat16()
+
+    # get test dataloader
+    dataload_config = get_data_configs(config["data"])
+    lerobot_config = dataload_config.get("lerobot_config", {})
+    dataset = load_test_dataset(config, lerobot_config, seed=42)
+    dataloader = dataset.get_dataloader()
+
+    total_frames = len(dataloader)
+
+    predict_mode = "fast" if config.get("use_fast_tokenizer", False) else "diffusion"
+    action_dim = 20 if predict_mode == "diffusion" else origin_action_dim
+    gt_traj = torch.zeros((total_frames, origin_action_dim))
+    pred_traj = torch.zeros((total_frames, origin_action_dim))
+
+    # use tqdm to show the progress
+    for idx, batch in tqdm(
+        enumerate(dataloader), total=total_frames, desc="predicting"
+    ):
+        if idx % pred_horizon == 0 and idx + pred_horizon < total_frames:
+            batch = batch.to("cuda")
+            with torch.no_grad():
+                outputs = model(
+                    **batch,
+                    action_dim=action_dim,
+                    pred_horizon=pred_horizon,
+                    mode="predict",
+                    predict_mode=predict_mode,
+                )
+                pred_traj[idx : idx + pred_horizon] = (
+                    outputs["predict_action"][:, :, :origin_action_dim]
+                    .detach()
+                    .cpu()
+                    .squeeze(0)
+                )
+
+            # Denormalize ground truth actions
+            gt_action_chunk = batch["action_chunk"][:, :, :origin_action_dim]
+            dof_mask = batch["dof_mask"].to(gt_action_chunk.dtype)
+            denormalized_gt = (
+                model.action_preprocessor.normalizer_action.unnormalize_data(
+                    gt_action_chunk,
+                    [lerobot_config.get("repo_id", "physical-intelligence/libero")],
+                    dof_mask,
+                ).squeeze(0)
             )
-        pred_traj[idx : idx + pred_horizon] = outputs["predict_action"].detach().cpu()
+            gt_traj[idx : idx + pred_horizon] = denormalized_gt.detach().cpu()
 
-        # Denormalize ground truth actions
-        gt_action_chunk = batch["action_chunk"][:, :, :action_dim]
-        dof_mask = batch["dof_mask"].to(gt_action_chunk.dtype)
-        denormalized_gt = model.action_preprocessor.normalizer_action.unnormalize_data(
-            gt_action_chunk, ["x2_normal"], dof_mask
-        )
-        gt_traj[idx : idx + pred_horizon] = denormalized_gt.detach().cpu()
+    gt_traj_np = gt_traj.numpy()
+    pred_traj_np = pred_traj.numpy()
 
+    timesteps = gt_traj.shape[0]
 
-gt_traj_np = gt_traj.numpy()
-pred_traj_np = pred_traj.numpy()
+    fig, axs = plt.subplots(
+        origin_action_dim, 1, figsize=(15, 5 * origin_action_dim), sharex=True
+    )
+    fig.suptitle("Action Comparison for lerobot", fontsize=16)
 
-timesteps = gt_traj.shape[0]
+    for i in range(origin_action_dim):
+        axs[i].plot(range(timesteps), gt_traj_np[:, i], label="Ground Truth")
+        axs[i].plot(range(timesteps), pred_traj_np[:, i], label="Prediction")
+        axs[i].set_ylabel(f"Action Dim {i+1}")
+        axs[i].legend()
+        axs[i].grid(True)
 
-fig, axs = plt.subplots(action_dim, 1, figsize=(15, 5 * action_dim), sharex=True)
-fig.suptitle("Action Comparison for lerobot", fontsize=16)
-
-for i in range(action_dim):
-    axs[i].plot(range(timesteps), gt_traj_np[:, i], label="Ground Truth")
-    axs[i].plot(range(timesteps), pred_traj_np[:, i], label="Prediction")
-    axs[i].set_ylabel(f"Action Dim {i+1}")
-    axs[i].legend()
-    axs[i].grid(True)
-
-axs[-1].set_xlabel("Timestep")
-plt.tight_layout(rect=[0, 0.03, 1, 0.95])
-os.makedirs(save_dir, exist_ok=True)
-plt.savefig(os.path.join(save_dir, "lerobot_comparison.png"))
-plt.close()
+    axs[-1].set_xlabel("Timestep")
+    plt.tight_layout(rect=[0, 0.03, 1, 0.95])
+    os.makedirs(save_dir, exist_ok=True)
+    save_path = os.path.join(save_dir, "lerobot_comparison.png")
+    plt.savefig(save_path)
+    print(f"Saved plot to {save_path}")
+    plt.close()

scripts/normalize.py

@@ -0,0 +1,161 @@
+# This file is copied from openpi
+import json
+import pathlib
+
+import numpy as np
+import numpydantic
+import pydantic
+
+
+@pydantic.dataclasses.dataclass
+class NormStats:
+    mean: numpydantic.NDArray
+    std: numpydantic.NDArray
+    q01: numpydantic.NDArray | None = None  # 1st quantile
+    q99: numpydantic.NDArray | None = None  # 99th quantile
+
+
+class RunningStats:
+    """Compute running statistics of a batch of vectors."""
+
+    def __init__(self):
+        self._count = 0
+        self._mean = None
+        self._mean_of_squares = None
+        self._min = None
+        self._max = None
+        self._histograms = None
+        self._bin_edges = None
+        self._num_quantile_bins = 5000  # for computing quantiles on the fly
+
+    def update(self, batch: np.ndarray) -> None:
+        """
+        Update the running statistics with a batch of vectors.
+
+        Args:
+            vectors (np.ndarray): An array where all dimensions except the last are batch dimensions.
+        """
+        batch = batch.reshape(-1, batch.shape[-1])
+        num_elements, vector_length = batch.shape
+        if self._count == 0:
+            self._mean = np.mean(batch, axis=0)
+            self._mean_of_squares = np.mean(batch**2, axis=0)
+            self._min = np.min(batch, axis=0)
+            self._max = np.max(batch, axis=0)
+            self._histograms = [
+                np.zeros(self._num_quantile_bins) for _ in range(vector_length)
+            ]
+            self._bin_edges = [
+                np.linspace(
+                    self._min[i] - 1e-10,
+                    self._max[i] + 1e-10,
+                    self._num_quantile_bins + 1,
+                )
+                for i in range(vector_length)
+            ]
+        else:
+            if vector_length != self._mean.size:
+                raise ValueError(
+                    "The length of new vectors does not match the initialized vector length."
+                )
+            new_max = np.max(batch, axis=0)
+            new_min = np.min(batch, axis=0)
+            max_changed = np.any(new_max > self._max)
+            min_changed = np.any(new_min < self._min)
+            self._max = np.maximum(self._max, new_max)
+            self._min = np.minimum(self._min, new_min)
+
+            if max_changed or min_changed:
+                self._adjust_histograms()
+
+        self._count += num_elements
+
+        batch_mean = np.mean(batch, axis=0)
+        batch_mean_of_squares = np.mean(batch**2, axis=0)
+
+        # Update running mean and mean of squares.
+        self._mean += (batch_mean - self._mean) * (num_elements / self._count)
+        self._mean_of_squares += (batch_mean_of_squares - self._mean_of_squares) * (
+            num_elements / self._count
+        )
+
+        self._update_histograms(batch)
+
+    def get_statistics(self) -> NormStats:
+        """
+        Compute and return the statistics of the vectors processed so far.
+
+        Returns:
+            dict: A dictionary containing the computed statistics.
+        """
+        if self._count < 2:
+            raise ValueError("Cannot compute statistics for less than 2 vectors.")
+
+        variance = self._mean_of_squares - self._mean**2
+        stddev = np.sqrt(np.maximum(0, variance))
+        q01, q99 = self._compute_quantiles([0.01, 0.99])
+        return NormStats(mean=self._mean, std=stddev, q01=q01, q99=q99)
+
+    def _adjust_histograms(self):
+        """Adjust histograms when min or max changes."""
+        for i in range(len(self._histograms)):
+            old_edges = self._bin_edges[i]
+            new_edges = np.linspace(
+                self._min[i], self._max[i], self._num_quantile_bins + 1
+            )
+
+            # Redistribute the existing histogram counts to the new bins
+            new_hist, _ = np.histogram(
+                old_edges[:-1], bins=new_edges, weights=self._histograms[i]
+            )
+
+            self._histograms[i] = new_hist
+            self._bin_edges[i] = new_edges
+
+    def _update_histograms(self, batch: np.ndarray) -> None:
+        """Update histograms with new vectors."""
+        for i in range(batch.shape[1]):
+            hist, _ = np.histogram(batch[:, i], bins=self._bin_edges[i])
+            self._histograms[i] += hist
+
+    def _compute_quantiles(self, quantiles):
+        """Compute quantiles based on histograms."""
+        results = []
+        for q in quantiles:
+            target_count = q * self._count
+            q_values = []
+            for hist, edges in zip(self._histograms, self._bin_edges, strict=True):
+                cumsum = np.cumsum(hist)
+                idx = np.searchsorted(cumsum, target_count)
+                q_values.append(edges[idx])
+            results.append(np.array(q_values))
+        return results
+
+
+class _NormStatsDict(pydantic.BaseModel):
+    norm_stats: dict[str, NormStats]
+
+
+def serialize_json(norm_stats: dict[str, NormStats]) -> str:
+    """Serialize the running statistics to a JSON string."""
+    return _NormStatsDict(norm_stats=norm_stats).model_dump_json(indent=2)
+
+
+def deserialize_json(data: str) -> dict[str, NormStats]:
+    """Deserialize the running statistics from a JSON string."""
+    return _NormStatsDict(**json.loads(data)).norm_stats
+
+
+def save(directory: pathlib.Path | str, norm_stats: dict[str, NormStats]) -> None:
+    """Save the normalization stats to a directory."""
+    path = pathlib.Path(directory) / "norm_stats.json"
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(serialize_json(norm_stats))
+
+
+def load(directory: pathlib.Path | str) -> dict[str, NormStats]:
+    """Load the normalization stats from a directory."""
+    path = pathlib.Path(directory) / "norm_stats.json"
+    if not path.exists():
+        raise FileNotFoundError(f"Norm stats file not found at: {path}")
+    return deserialize_json(path.read_text())