From a0d7627d81c57f9ca826029e581deca313e0d548 Mon Sep 17 00:00:00 2001 From: Adil Zouitine Date: Mon, 29 Sep 2025 17:37:26 +0200 Subject: [PATCH] feat(train): include input and output features in processor overrides for normalization (#2088) (#2090) Signed-off-by: AdilZouitine --- src/lerobot/scripts/lerobot_train.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 86b2bbae..12a1f53c 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/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(