Files
lerobot/tests/training
alejodosr 266be2bd17 feat(train): add opt-in EMA of the policy weights (--ema.enable=true) (#4323)
* feat(train): add opt-in EMA of the policy weights (--ema.enable=true)

Maintain an EMA shadow via diffusers' EMAModel (lazy import, no new
dependency) with the reference Diffusion Policy schedule. Saves the
shadow for exact resume plus a loadable pretrained_model_ema/ per
checkpoint, evaluates the EMA weights during env eval, and pushes them
to a sibling <repo_id>-ema repo. Fixes huggingface/lerobot#4259.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* docs(diffusion): document the --ema.enable training flag

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* fix(tests): skip EMA training tests when accelerate/diffusers are missing

* feat(train): support constant EMA decay (--ema.decay) for openpi-style policies

* fix(train): gate EMA step on sync_gradients; use parallel_dims.is_sharded guard

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
2026-08-07 13:52:23 +02:00
..