Align Normalizer zero-delta handling with LeRobot training collator

[Papaercold]

Summary

Fix a zero-delta normalization mismatch between the LeRobot training collator and the shared model normalizer.

During fine-tuning + evaluation, I observed a case where training loss looked normal, but evaluation produced all-NaN action vectors. I traced this to Normalizer.normalize_data() dividing by a zero normalization delta.

Problem

The LeRobot training collator already handles zero deltas safely in:

wall_x/data/backends/lerobot/loader.py:473-481

def _normalize(cls, action, min_stat, delta):
    """
    Normalize action data using min-max normalization.
    """
    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

This training-side normalization is used for agent_pos here:

wall_x/data/backends/lerobot/loader.py:535-553

and for action here:

wall_x/data/backends/lerobot/loader.py:560-578

So during training, if a dimension has delta == 0, the collator replaces that delta with 1 before division.

However, the shared normalizer in:

wall_x/model/core/action/normalizer.py:284-295

previously divided directly by self.delta[dataset_name]. For constant action/state dimensions, the normalization stats can have:

q99[i] == q01[i]
delta[i] == 0

Then normalization can produce 0 / 0 -> NaN. torch.clamp() does not remove NaNs, so these values can propagate into evaluation/inference actions.

This explains the behavior I saw:

Training path: delta == 0 is guarded -> loss stays normal
Evaluation/inference path: delta == 0 can divide by zero -> action becomes NaN

Fix

This PR updates Normalizer.normalize_data() to use the same zero-delta guard as the LeRobot training collator:

def normalize_data(self, xs, dataset_names):
    new_xs = []
    dataset_names = [name for name in dataset_names if name != "x2_multimodal"]
    for x, dataset_name in zip(xs, dataset_names):
        delta = self.delta[dataset_name]
        delta = torch.where(delta == 0, torch.ones_like(delta), delta)
        x = (x - self.min[dataset_name]) / delta
        x = x * 2 - 1
        x = torch.clamp(x, -1, 1)
        new_xs.append(x)
    new_xs = torch.stack(new_xs)
    return new_xs

This does not change tensor shapes or nonzero-delta behavior. It only prevents division by zero for zero-range dimensions.

Validation

I reproduced the issue in a fine-tuning + evaluation workflow where predicted actions became all NaNs. After this change, the same workflow no longer produced all-NaN action vectors.

[StarrickLiu]

Thanks @Papaercold — this is a clean, correct fix and the diagnosis in #113 is spot on. The training-side LeRobot collator already guards delta == 0, but the shared Normalizer.normalize_data() divided by delta directly, so any constant dimension (q99 == q01 → delta == 0) produced 0 / 0 → NaN at eval/inference time while training stayed stable.

I confirmed the fix is also complete: normalize_data() is the only divide-by-delta path, and unnormalize_data() multiplies by delta (multiply-by-zero is fine), so this single guard closes the NaN source. It also matches the exact guard already used in the collator, keeping train and eval normalization consistent.

Merging. Thanks for the careful report and fix — much appreciated! 🙏