Skip to content

Apply rename_map to fresh-built processor pipelines for all policies#3428

Open
ihlej wants to merge 2 commits intohuggingface:mainfrom
ihlej:fix/rename-map-fresh-training
Open

Apply rename_map to fresh-built processor pipelines for all policies#3428
ihlej wants to merge 2 commits intohuggingface:mainfrom
ihlej:fix/rename-map-fresh-training

Conversation

@ihlej
Copy link
Copy Markdown

@ihlej ihlej commented Apr 21, 2026

fix: apply --rename_map on fresh-built processor pipelines

Summary / Motivation

Any PI05/PI0/PI0Fast training run that combines use_relative_actions=true with --rename_map crashes at step 0. The CLI rename map is silently dropped on the fresh-build path. --rename_map is only wired into the processor pipeline on the pretrained-load path, via preprocessor_overrides["rename_observations_processor"] inside the if processor_pretrained_path is not None: branch at scripts/lerobot_train.py:288. On the fresh-build path (when processor_pretrained_path is None), the per-policy make_*_pre_post_processors functions (for example make_pi05_pre_post_processors()) all construct the rename step from a hardcoded RenameObservationsProcessorStep(rename_map={}), so the CLI rename_map is never used.

Two ways to hit the fresh-build path:

  1. Training from scratch (--policy.path not set), which affects any policy.
  2. use_relative_actions=true on a non-resume run, which forces processor_pretrained_path = None at scripts/lerobot_train.py:275 so processors are rebuilt from the current config (otherwise old checkpoints can silently drop the relative-action transform steps).

In both cases, --rename_map is silently ignored. For PI05 the
symptom is ValueError: All image features are missing from the batch
at step 0; for other policies this typically surfaces as a KeyError
on the renamed key when the policy tries to read its image features.
This PR applies rename_map on the fresh-build path; the
pretrained-load path is unchanged.

Related issues

Fixes: #3425

What changed

  • scripts/lerobot_train.py: forward cfg.rename_map into processor_kwargs unconditionally, not only when processor_pretrained_path is not None.
  • policies/factory.py: after the fresh-build branch constructs the pipeline, walk preprocessor.steps and set rename_map on the RenameObservationsProcessorStep. This is the fresh-build counterpart to the overrides argument that PolicyProcessorPipeline.from_pretrained already applies on the pretrained-load path.
  • No breaking changes. The pretrained-load path is untouched and still applies rename_map via preprocessor_overrides["rename_observations_processor"].

How was this tested (or how to run locally)

  • PI05 (hardware): The --rename_map silently ignored for PI05 when use_relative_actions=true #3425 reproduction previously crashed at step 0. With this fix it trains end to end.
  • ACT (zero-setup): the command below downloads lerobot/pusht (~115 kB of meta + video), builds policy.input_features with a renamed image key so the dataset and policy keys intentionally disagree, and passes --rename_map to bridge that mismatch. It then runs 5 training steps. Without the fix, --rename_map is silently ignored and the run crashes at step 0 with a KeyError on the renamed key; with the fix, the rename is applied and the 5 steps complete.
OUT=/tmp/act_rename_map_repro && rm -rf "$OUT"
REPO=lerobot/pusht

readarray -t L < <(python3 - "$REPO" <<'PY'
import sys, json
from huggingface_hub import hf_hub_download
info = json.loads(open(hf_hub_download(sys.argv[1], "meta/info.json", repo_type="dataset")).read())
feats = info["features"]
img_key, img = next((k, v) for k, v in feats.items() if v["dtype"] in ("image", "video"))
h, w, c = img["shape"]
renamed = img_key + "_renamed"
print(json.dumps({
    "observation.state": {"type": "STATE", "shape": feats["observation.state"]["shape"]},
    renamed: {"type": "VISUAL", "shape": [c, h, w]},
}, separators=(",", ":")))
print(json.dumps({img_key: renamed}, separators=(",", ":")))
PY
)

lerobot-train \
  --dataset.repo_id=$REPO \
  --policy.type=act --policy.device=cuda \
  --policy.input_features="${L[0]}" \
  --policy.push_to_hub=false --policy.repo_id=local/act_rename_map_repro \
  --output_dir="$OUT" --job_name=act_rename_map_repro \
  --steps=5 --save_freq=9999 --wandb.enable=false \
  --rename_map="${L[1]}"

Checklist (required before merge)

Reviewer notes

  • Hardware-verified on PI05 only; the ACT reproducer above demonstrates the fix works on at least one additional policy via the shared code path. I have not tested by hand every policy.

Previously --rename_map only took effect on the pretrained-load path via
preprocessor_overrides. On the fresh-build path (pretrained_path is None),
per-policy make_*_pre_post_processors hardcode
RenameObservationsProcessorStep(rename_map={}), so --rename_map was silently
dropped.
@github-actions github-actions Bot added the policies Items related to robot policies label Apr 21, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

policies Items related to robot policies

Projects

None yet

Development

Successfully merging this pull request may close these issues.

--rename_map silently ignored for PI05 when use_relative_actions=true

1 participant