Fix: apply rename_map on fresh-built processor pipelines (complete)#3514
Open
wadeKeith wants to merge 3 commits intohuggingface:mainfrom
Open
Fix: apply rename_map on fresh-built processor pipelines (complete)#3514wadeKeith wants to merge 3 commits intohuggingface:mainfrom
wadeKeith wants to merge 3 commits intohuggingface:mainfrom
Conversation
When use_relative_actions=true with a pretrained model (pi05, pi0, pi0_fast, smolvla, xvla), the code sets processor_pretrained_path=None to force building processors from the current policy config rather than loading from the checkpoint. However, rename_map was only being passed via preprocessor_overrides which only applies when loading from pretrained. This meant rename_map was silently ignored when use_relative_actions=true. The fix: - Add rename_map to ProcessorConfigKwargs TypedDict - Always pass rename_map to make_pre_post_processors (when not reward model) - Update all affected processor factories to accept and use rename_map: - make_pi05_pre_post_processors - make_pi0_pre_post_processors - make_pi0_fast_pre_post_processors - make_smolvla_pre_post_processors - make_xvla_pre_post_processors This ensures the RenameObservationsProcessorStep uses the user-provided rename_map regardless of whether processors are loaded from pretrained or built from scratch. Fixes huggingface#3425
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes rename_map handling when building processor pipelines from scratch (not loading from a pretrained processor), ensuring observation-key renames are consistently applied alongside normalization statistics so processor steps (notably normalization) can resolve the correct stats keys. It also improves backward compatibility for third‑party policy processor factories by conditionally passing rename_map only when supported.
Changes:
- In training (
lerobot_train.py), passrename_mapinto processor creation and renamedataset.meta.statskeys viarename_stats()in the fresh-build path. - In policy processor factories (PI0/PI05/PI0Fast/SmolVLA/XVLA), accept
rename_mapand wire it intoRenameObservationsProcessorStep. - In
policies/factory.py, extend processor factory kwargs to includerename_map, document it, and add signature-based backward-compat when calling dynamically imported (plugin) processor factories.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| src/lerobot/scripts/lerobot_train.py | Passes rename_map into processor construction and renames dataset stats keys in the fresh-build path. |
| src/lerobot/policies/factory.py | Adds rename_map to processor kwargs, forwards it to built-in factories, and adds backward-compatible dynamic factory invocation. |
| src/lerobot/policies/pi0/processor_pi0.py | Adds rename_map param and applies it in RenameObservationsProcessorStep. |
| src/lerobot/policies/pi05/processor_pi05.py | Adds rename_map param and applies it in RenameObservationsProcessorStep; updates pipeline order comment. |
| src/lerobot/policies/pi0_fast/processor_pi0_fast.py | Adds rename_map param and applies it in RenameObservationsProcessorStep; updates pipeline order comment. |
| src/lerobot/policies/smolvla/processor_smolvla.py | Adds rename_map param and applies it in RenameObservationsProcessorStep. |
| src/lerobot/policies/xvla/processor_xvla.py | Adds rename_map param and applies it in RenameObservationsProcessorStep. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| """ | ||
|
|
||
| import dataclasses | ||
| import inspect |
Comment on lines
+318
to
+323
| processor_kwargs["rename_map"] = cfg.rename_map | ||
| # When rename_map is provided and stats are used (fresh-build path), rename stats keys | ||
| # so NormalizerProcessorStep can find them after observations are renamed. | ||
| if cfg.rename_map and ((processor_pretrained_path and not cfg.resume) or not processor_pretrained_path): | ||
| processor_kwargs["dataset_stats"] = rename_stats(dataset.meta.stats, cfg.rename_map) | ||
|
|
Comment on lines
+222
to
+223
| policy feature names. When provided, stats are renamed accordingly so that | ||
| NormalizerProcessorStep can locate them after observation keys are remapped. |
Comment on lines
+625
to
+627
| if rename_map is not None: | ||
| sig = inspect.signature(function) | ||
| if "rename_map" in sig.parameters: |
5 tasks
…Copilot review comments 1. lerobot_train.py: remove unused 'inspect' import (was used in original fix but refactored away) 2. lerobot_train.py: call rename_stats() on dataset.meta.stats when both rename_map and fresh-build path are active, so NormalizerProcessorStep can locate stats for renamed keys. 3. factory.py _make_processors_from_policy_config: inspect signature before passing rename_map; fall back to (config, dataset_stats) for third-party factories that don't declare the parameter (backward-compatible). 4. factory.py ProcessorConfigKwargs: document rename_map in docstring.
wadeKeith
force-pushed
the
fix/3512-complete
branch
from
a4402df to
7beeca6
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fix: apply rename_map on fresh-built processor pipelines (complete fix)
This supersedes #3512 and addresses all Copilot review comments:
Changes
1.
lerobot_train.py— Rename stats alongside observation keysWhen
rename_mapis provided in the fresh-build path, callrename_stats()ondataset.meta.statsbefore passing tomake_pre_post_processors. This ensuresNormalizerProcessorStepcan locate stats for renamed keys.2.
factory.py— Backward-compat for third-party plugin factories_make_processors_from_policy_config()now inspects the signature of the dynamically-imported factory function. Ifrename_mapis not in its parameters, it falls back to calling with only(config, dataset_stats)— no breaking change for plugins.3.
factory.py— DocumentProcessorConfigKwargs.rename_mapUpdated the TypedDict docstring to mention
rename_mapand explain its interaction withdataset_stats.Previous fixes (included from #3511 and #3512)
rename_mapforwarded tomake_pi0/pi05/smolvla/xvla_pre_post_processors()rename_mapforwarded to_make_processors_from_policy_config()rename_mapinRenameObservationsProcessorStepRelated