feat(train): include input and output features in processor overrides for normalization (#2088) (#2090)

Signed-off-by: AdilZouitine <[email protected]>

Author: AdilZouitine

src/lerobot/scripts/lerobot_train.py

@@ -187,10 +187,16 @@ def train(cfg: TrainPipelineConfig):
     if cfg.policy.pretrained_path is not None:
         processor_kwargs["preprocessor_overrides"] = {
             "device_processor": {"device": device.type},
-            "normalizer_processor": {"stats": dataset.meta.stats},
+            "normalizer_processor": {
+                "stats": dataset.meta.stats,
+                "features": {**policy.config.input_features, **policy.config.output_features},
+            },
         }
         processor_kwargs["postprocessor_overrides"] = {
-            "unnormalizer_processor": {"stats": dataset.meta.stats},
+            "unnormalizer_processor": {
+                "stats": dataset.meta.stats,
+                "features": {**policy.config.input_features, **policy.config.output_features},
+            },
         }
 
     preprocessor, postprocessor = make_pre_post_processors(