From 9f0e4dfb538783977ec45ec5ed61b4fc2d08856f Mon Sep 17 00:00:00 2001 From: pepijn223 Date: Mon, 13 Jul 2026 16:46:10 +0200 Subject: [PATCH] fix(pi052): accept use_flex_attention config field for checkpoint compat MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Newer PI052 training runs serialize use_flex_attention into config.json. This branch's attention path is SDPA/eager (mathematically equivalent), so the field is accepted as an inert no-op (mirrors the existing use_hf_kernels compat field) — otherwise loading those checkpoints raises DecodingError. Co-authored-by: Cursor --- src/lerobot/policies/pi052/configuration_pi052.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/lerobot/policies/pi052/configuration_pi052.py b/src/lerobot/policies/pi052/configuration_pi052.py index 1e8911604..b65f1cbae 100644 --- a/src/lerobot/policies/pi052/configuration_pi052.py +++ b/src/lerobot/policies/pi052/configuration_pi052.py @@ -246,6 +246,15 @@ class PI052Config(PI05Config): ``DecodingError: The fields use_hf_kernels are not valid for PI052Config`` (job 22164492). Remove in a future major bump.""" + use_flex_attention: bool = False + """Accepted for checkpoint-config compatibility only — no-op in this branch. + Newer training runs serialize ``use_flex_attention`` into ``config.json`` to + select the FlexAttention kernel at train time; this branch's attention path + is SDPA/eager, which is mathematically equivalent (same softmax attention, + different kernel), so inference/eval results are unchanged. Retained so those + checkpoints load instead of raising ``DecodingError: The fields + use_flex_attention are not valid for PI052Config``.""" + # Optimizer foreach/fused. pi052 carries these locally because the shared # PI05Config (kept identical to upstream main) does not define them; the # checkpoints we train serialize both keys into config.json, so they must