From 852d5ff41780d481d74abe399f2fff02f86831b7 Mon Sep 17 00:00:00 2001
From: ByteRainx <88035662+ByteRainx@users.noreply.github.com>
Date: Thu, 11 Sep 2025 15:51:44 +0800
Subject: [PATCH] fix: improve numerical stability in action normalization
 (#18)

- Add zero delta handling in _normalize method
- Use torch.ones_like to replace zero delta values
---
 wall_x/data/load_lerobot_dataset.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/wall_x/data/load_lerobot_dataset.py b/wall_x/data/load_lerobot_dataset.py
index cd26b50..99ce239 100644
--- a/wall_x/data/load_lerobot_dataset.py
+++ b/wall_x/data/load_lerobot_dataset.py
@@ -263,7 +263,12 @@ class DataCollator:
 
     @classmethod
     def _normalize(cls, action, min_stat, delta):
-        x = (action - min_stat) / (delta)
+        """
+        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
         x = torch.clamp(x, -1, 1)
         return x