From 384feca91ae31cf13c43494eae752d5a74190301 Mon Sep 17 00:00:00 2001 From: pepijn223 Date: Fri, 5 Jun 2026 16:09:08 +0200 Subject: [PATCH] fix(ema): default EMAConfig.enable to False (opt-in) EMA was on by default, so every training run on the branch (incl. VLA-JEPA and other non-flow-matching policies) created a full fp32 shadow copy. EMA only benefits flow-matching/diffusion policies (pi0/pi05/pi052). Make it opt-in via --ema.enable=true; the pi05/pi052 recipes already pass that flag. Co-authored-by: Cursor --- src/lerobot/configs/default.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/lerobot/configs/default.py b/src/lerobot/configs/default.py index f741fc6ae..0d3231d69 100644 --- a/src/lerobot/configs/default.py +++ b/src/lerobot/configs/default.py @@ -98,14 +98,14 @@ class EMAConfig: Cost: 1× model params in fp32 shadow (~13 GB for pi052's 3.3B params) + one elementwise update per training step (~1% step time). - On by default — matches openpi (JAX) which ships EMA on for every - config, and closes the gap with the openpi PyTorch port which - explicitly lists EMA as unsupported. Set ``--ema.enable=false`` to - disable for short runs / memory-constrained training where the - extra fp32 shadow copy is the bottleneck. + Off by default (opt-in): EMA is only beneficial for flow-matching / + diffusion policies (pi0/pi05/pi052), and the fp32 shadow copy is pure + overhead for other policies (e.g. VLA-JEPA). Set ``--ema.enable=true`` + to turn it on (the pi05/pi052 training recipes do this). openpi (JAX) + ships EMA on for every config; enable it explicitly to match that. """ - enable: bool = True + enable: bool = False # Target EMA decay β in θ_ema ← β·θ_ema + (1-β)·θ_live (passed to # ema-pytorch as ``beta``). # 0.999 — last ~1000 steps; pi05_libero default in openpi