Migration mangles feature names containing an underscore (observation.environment_state)

[saime428]

What happens

extract_normalization_stats in src/lerobot/processor/migrate_policy_normalization.py derives a feature name from a checkpoint buffer name with an unconditional replace("_", ".") (:112). A feature whose own name contains an underscore is mangled, so its migrated statistics land under a key the runtime never looks up.

observation.environment_state (OBS_ENV_STATE, used by ACT, TDMPC and gaussian_actor) is the concrete case:

import torch
from lerobot.processor.migrate_policy_normalization import extract_normalization_stats

state_dict = {
    "normalize_inputs.buffer_observation_environment_state.mean": torch.tensor([1.0]),
    "normalize_inputs.buffer_observation_environment_state.std": torch.tensor([2.0]),
}
print(sorted(extract_normalization_stats(state_dict)))

['observation.environment.state']        # expected: ['observation.environment_state']

Why it matters

This hits plain single-dataset checkpoints — no multi-dataset prefix involved. The statistics are extracted and written to the migrated processor config, but under observation.environment.state, so at runtime the normalizer never finds them for observation.environment_state and the feature silently goes unnormalized.

Why it is not a one-line fix

The mangling is not reversible from the string alone. buffer_observation_environment_state could be observation.environment_state or observation.environment.state, and nothing in the key distinguishes them. Resolving it needs the policy's declared feature keys as the authority, rather than more string splitting.

Related

Raised by @takashi-jp while reviewing #4441, which fixes a different defect in the same function (dataset-prefixed keys being dropped entirely). Filed separately because the two need different fixes, and #4441 deliberately leaves this one alone.

[ShivamThareja]

Hi saime428, I'd like to work on this issue. I plan to trace how the policy's declared feature keys are exposed during normalization-stat migration and use those keys as the source of truth rather than relying on the lossy _ → . conversion. I'll also add a regression test covering observation.environment_state and verify that existing feature-name migrations remain unchanged.

[saime428]

Thanks for the offer — before you sink time into it, two things that aren't visible from the issue alone.

I have #4441 open on the same function. It changes the prefix matching inside extract_normalization_stats so dataset-scoped keys (normalize_inputs.so100_buffer_observation_state.mean) survive migration. That's the same loop this issue is about, so whoever lands second rebases on top of the other. I filed this issue while working on that PR, and I was planning to take it once #4441 settles.

If you'd still rather drive it, I'm happy to hand it over and review — just say so and it's yours. Either way, here's the part that cost me the most time:

Using the policy's declared feature keys as the source of truth is the right call, but they aren't reachable from where the mangling happens. extract_normalization_stats receives only the state_dict — there is no policy object and no feature list in scope — so the key set has to be threaded in from the caller rather than recovered inside the function. That plumbing, not the string handling, is the actual body of this fix.

[takashi-jp]

Thanks for filing this, @saime428 — I flagged it as an aside in a review comment on #4441 and it could easily have been lost there. Turning it into a standalone issue with a runnable reproduction, and being explicit about why it is not a one-line fix, is a much better home for it than a paragraph buried in a review thread.

Two small things to add on top of what you already laid out.

Disclosure, since it affects the rebase picture: I also have #4446 open, which touches extract_normalization_stats / migrate_policy_normalization.py — dataset-prefix handling, not this bug. It leaves this one alone on purpose and now links here from its reviewer notes. So whoever picks this up will end up rebasing on top of both #4441 and #4446, not just #4441. Worth knowing before starting rather than discovering it at conflict time.

To be clear, I'm not looking to take this one — @ShivamThareja / @saime428, it's yours to sort out between you. Just adding what I know.

On the plumbing you flagged (feature keys not being in scope inside extract_normalization_stats) — agreed that this is the body of the fix, and one detail that might make it less daunting than it reads: main() already loads config before calling it, at line 524, a few lines above the extract_normalization_stats(state_dict) call at line 528. config["input_features"] / config["output_features"] — the modern, unambiguous feature-name schema, distinct from the free-form config["features"] that detect_features_and_norm_modes falls back to — sit in that dict from the moment it's loaded. So for any checkpoint with a modern config.json, the exact declared feature names are already available at the call site; they just aren't threaded into extract_normalization_stats as a parameter yet. The plumbing is real, but it looks like one extra argument rather than having to reach for a policy object.

That doesn't fully close it, though: checkpoints old enough to lack input_features/output_features in their config have no config-level source of truth, and the ambiguity in the issue body is real for those — the same gap detect_features_and_norm_modes's stats-based fallback already has today.