huggingface / lerobot

🤗 LeRobot: Making AI for Robotics more accessible with end-to-end learning
Apache License 2.0
6.9k stars 631 forks source link

Wrong normalization of the images #61

Closed Cadene closed 6 months ago

Cadene commented 6 months ago

We compute stats on raw data. For instance:

stats["observation", "image", "top", "max"].max()
tensor(255.)

However, in lerobot/common/datasets/factory.py we apply a first transform to modify the pixel range of images from [0,255] to [0,1]:

transforms = [Prod(in_keys=img_keys, prod=1 / 255)]

then we apply our stats which expect images to be in the range [0,255].

This bug only affects tdmpc which uses image keys during its normalization:

        in_keys = [("observation", "state"), ("action")]

        if cfg.policy.name == "tdmpc":
            # TODO(rcadene): we add img_keys to the keys to normalize for tdmpc only, since diffusion and act policies normalize the image inside the model for now
            in_keys += img_keys
            # TODO(racdene): since we use next observations in tdmpc, we also add them to the normalization. We are wasting a bit of compute on this for now.
            in_keys += [("next", *key) for key in img_keys]
            in_keys.append(("next", "observation", "state"))

cc @alexander-soare @aliberts @qgallouedec

alexander-soare commented 6 months ago

Thanks @Cadene , I stumbled across this myself. I'll take care of it.

Cadene commented 6 months ago

Addressed by https://github.com/huggingface/lerobot/pull/64