admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-30 13:09:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(train): keep checkpoint processor stats on resume (#4022)

Browse Source 
Co-authored-by: Martino Russi <77496684+nepyope@users.noreply.github.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
This commit is contained in:
Jash Shah
2026-07-30 04:47:56 -07:00
committed by GitHub
parent 643ffb4785
commit 0cef9cd197
1 changed files with 7 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/scripts/lerobot_train.py
+7 -2
View File
@@ -348,7 +348,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
preprocessor_overrides = { preprocessor_overrides = {
"device_processor": {"device": device.type}, "device_processor": {"device": device.type},
"normalizer_processor": { "normalizer_processor": {
"stats": dataset.meta.stats,
"features": {**policy.config.input_features, **policy.config.output_features}, "features": {**policy.config.input_features, **policy.config.output_features},
"norm_map": policy.config.normalization_mapping, "norm_map": policy.config.normalization_mapping,
}, },
@@ -356,11 +355,17 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
} }
postprocessor_overrides = { postprocessor_overrides = {
"unnormalizer_processor": { "unnormalizer_processor": {
"stats": dataset.meta.stats,
"features": policy.config.output_features, "features": policy.config.output_features,
"norm_map": policy.config.normalization_mapping, "norm_map": policy.config.normalization_mapping,
}, },
} }
# On resume, the checkpoint's saved processor stats are authoritative: they may have
# been adapted by the policy (e.g. EVO1 pads state/action stats to max_state_dim),
# and force-feeding raw dataset stats over them crashes normalization (#4006).
# This mirrors the `dataset_stats` kwarg above, which is also skipped on resume.
if not cfg.resume:
preprocessor_overrides["normalizer_processor"]["stats"] = dataset.meta.stats
postprocessor_overrides["unnormalizer_processor"]["stats"] = dataset.meta.stats
if getattr(active_cfg, "use_relative_actions", False): if getattr(active_cfg, "use_relative_actions", False):
preprocessor_overrides["relative_actions_processor"] = { preprocessor_overrides["relative_actions_processor"] = {
"enabled": True, "enabled": True,