migrate_policy_normalization crashes on lerobot/vqbet_pusht: config schema drift (unknown field mlp_hidden_dim)
Description
migrate_policy_normalization cannot migrate lerobot/vqbet_pusht — a first-party checkpoint — because its serialized config.json carries mlp_hidden_dim, a field the current VQBeTConfig no longer declares. draccus's strict decoding rejects unknown fields:
draccus.utils.DecodingError: The fields `mlp_hidden_dim` are not valid for VQBeTConfigThis is independent of the save-path fix in #4654 (and of the 0.6.1 release): the pre-#4654 code path fails on the same checkpoint too, just earlier and less clearly (TypeError: __init__() got an unexpected keyword argument 'mlp_hidden_dim' from config_cls(**cleaned_config)).
Schema drift is expected for old checkpoints — handling it is arguably part of the migration tool's job.
Reproduction
python -m lerobot.processor.migrate_policy_normalization --pretrained-path lerobot/vqbet_pusht --output-dir /tmp/vqbet
# draccus.utils.DecodingError: The fields `mlp_hidden_dim` are not valid for VQBeTConfigWorkaround (verified)
Intersect the old config's keys with the target config class's dataclass fields, warn on drops, write a local dir, and migrate that:
import dataclasses, json
from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
cfg = json.load(open(config_path))
known = {f.name for f in dataclasses.fields(VQBeTConfig)} | {"type"}
dropped = [k for k in cfg if k not in known] # -> ['mlp_hidden_dim']After dropping the one unknown field, migration completes and the cleaned state dict is bit-identical to the original snapshot (207 tensors, exactly the 6 normalization buffers removed).
Suggested fix
Have the script filter config.json keys against the resolved config class's dataclass fields before parsing, logging a WARNING per dropped key (mlp_hidden_dim is not used by the current architecture, so dropping is semantics-preserving here; a warning keeps renamed-field cases visible). Happy to PR this on top of the #4654 approach.
Aside (separate from the crash)
For reference, after the lossless migration this checkpoint scores near zero in the current inference pipeline (max coverage 0.016 on PushT seed 1000, identical via direct in-process inference and via a ROS transport bridge — so it is not a transport/loading artifact). If that's simply the demo checkpoint's quality level, please ignore this remark; flagging it in case it's unexpectedly low.
Source: huggingface/lerobot