Wrong normalization of the images

[Cadene]

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]

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

[Cadene]

Addressed by #64