refactor(train): remove EMA support from training pipeline

Drop the opt-in EMA-shadow feature entirely: EMAConfig, the `ema` field on
TrainPipelineConfig, all EMA logic in lerobot_train.py (setup/resume, per-step
update, W&B observability, checkpoint save, EMA-model eval, and the sibling
`<repo_id>-ema` hub push), and the ema-pytorch dependency.

Co-authored-by: Cursor <cursoragent@cursor.com>

pyproject.toml

@@ -85,11 +85,6 @@ dependencies = [
     "termcolor>=2.4.0,<4.0.0",
     "tqdm>=4.66.0,<5.0.0",
 
-    # Training utilities
-    # EMA of policy parameters (Diffusion Policy / pi05 style). Tiny
-    # pure-python dependency — preferred over a hand-rolled implementation.
-    "ema-pytorch>=0.7.7,<1.0.0",
-
     # Build tools (required by opencv-python-headless on some platforms)
     "cmake>=3.29.0.1,<4.2.0",
     "setuptools>=71.0.0,<81.0.0",

src/lerobot/configs/default.py

@@ -80,19 +80,6 @@ class WandBConfig:
     log_examples_n: int = 4
 
 
-@dataclass
-class EMAConfig:
-    """EMA shadow for flow/diffusion policies. Off by default because it doubles model memory."""
-
-    enable: bool = False
-    # Target EMA decay beta in theta_ema <- beta * theta_ema + (1 - beta) * theta_live.
-    decay: float = 0.99
-    # Initial update calls that keep the shadow as a hard copy before averaging starts.
-    warmup_steps: int = 0
-    # Use the EMA model for periodic eval.
-    use_for_eval: bool = True
-
-
 @dataclass
 class EvalConfig:
     n_episodes: int = 50

src/lerobot/configs/train.py

@@ -31,7 +31,7 @@ from lerobot.utils.hub import HubMixin, find_latest_hub_checkpoint
 from lerobot.utils.sample_weighting import SampleWeightingConfig
 
 from . import parser
-from .default import DatasetConfig, EMAConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
+from .default import DatasetConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
 from .policies import PreTrainedConfig
 from .rewards import RewardModelConfig
 
@@ -118,7 +118,6 @@ class TrainPipelineConfig(HubMixin):
     scheduler: LRSchedulerConfig | None = None
     eval: EvalConfig = field(default_factory=EvalConfig)
     wandb: WandBConfig = field(default_factory=WandBConfig)
-    ema: EMAConfig = field(default_factory=EMAConfig)
     peft: PeftConfig | None = None
 
     # Where to run training (local default, or an HF Jobs flavor). See JobConfig.

src/lerobot/scripts/lerobot_train.py

@@ -600,53 +600,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
 
     policy.train()
 
-    # ------------------------------------------------------------------
-    # EMA setup
-    # ------------------------------------------------------------------
-    # Shadow copy of the trainable params for late-training averaging
-    # (Chi et al. 2023 Diffusion Policy §V.D; openpi JAX trainer ships
-    # this with decay=0.999 for pi05_libero; openpi PyTorch port and
-    # LeRobot main both skip it). Off by default; opt in with
-    # ``--ema.enable=true``. Implemented via ema-pytorch
-    # (https://github.com/lucidrains/ema-pytorch) — the standard PyTorch
-    # EMA library, also used by lucidrains' diffusion repos.
-    ema = None
-    if cfg.ema.enable and is_main_process:
-        from ema_pytorch import EMA  # noqa: PLC0415
-
-        ema = EMA(
-            accelerator.unwrap_model(policy),
-            beta=cfg.ema.decay,
-            update_after_step=cfg.ema.warmup_steps,
-            update_every=1,  # update on every ema.update() call
-            # Don't register the live model as an ema submodule — accelerator
-            # already owns its lifecycle, and double-registration would
-            # double-count its params in ``ema.state_dict()``.
-            include_online_model=False,
-        )
-        ema.to(accelerator.device)
-        logging.info(
-            "EMA enabled (ema-pytorch): beta=%g, update_after_step=%d, use_for_eval=%s",
-            cfg.ema.decay,
-            cfg.ema.warmup_steps,
-            cfg.ema.use_for_eval,
-        )
-
-        # Resume the EMA shadow if a previous run wrote one.
-        if cfg.checkpoint_path is not None:
-            ema_path = cfg.checkpoint_path / "training_state" / "ema_state.pt"
-            if ema_path.exists():
-                logging.info("Resuming EMA shadow from %s", ema_path)
-                try:
-                    ema.load_state_dict(
-                        torch.load(ema_path, map_location=accelerator.device, weights_only=True)
-                    )
-                except Exception as exc:  # noqa: BLE001
-                    logging.warning(
-                        "Failed to load EMA shadow (%s) — restarting EMA from current live weights",
-                        exc,
-                    )
-
     train_metrics = {
         # Per-rank loss reflects only one shard of the global batch; mean recovers the loss DDP
         # is actually optimizing. grad_norm and lr are already identical on every rank (post
@@ -714,14 +667,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             log_metrics=log_metrics,
         )
 
-        # EMA update: pull one step of the live weights into the shadow.
-        # Runs only on the main process (the shadow lives there); other
-        # ranks rely on the live model staying in sync via accelerator.
-        # ``ema-pytorch`` holds an internal reference to the online model
-        # (set at construction), so ``ema.update()`` takes no args.
-        if ema is not None:
-            ema.update()
-
         # Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
         # increment `step` here.
         step += 1
@@ -759,14 +704,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
                     if sample_weighter is not None:
                         weighter_stats = sample_weighter.get_stats()
                         wandb_log_dict.update({f"sample_weighting/{k}": v for k, v in weighter_stats.items()})
-                    # EMA observability: ``ema.step`` is the count of
-                    # ``ema.update()`` calls (= optimizer steps once EMA is
-                    # enabled); ``ema.initted`` flips to True once we've
-                    # crossed ``update_after_step``.
-                    if ema is not None:
-                        wandb_log_dict["ema/step"] = int(ema.step.item())
-                        wandb_log_dict["ema/initted"] = float(ema.initted.item())
-                        wandb_log_dict["ema/beta"] = float(cfg.ema.decay)
                     wandb_logger.log_dict(wandb_log_dict, step)
             train_tracker.reset_averages()
 
@@ -837,19 +774,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
                     optim_state_dict=optim_state_dict,
                 )
                 update_last_checkpoint(checkpoint_dir)
-                # Save the EMA shadow alongside the training state so a
-                # resumed run picks up exactly where the live EMA left off.
-                # ``ema-pytorch.state_dict()`` returns the full shadow
-                # nn.Module's state dict + step/initted buffers; saved as
-                # .pt (the rest of training_state mixes formats already).
-                if ema is not None:
-                    try:
-                        ema_path = checkpoint_dir / "training_state" / "ema_state.pt"
-                        ema_path.parent.mkdir(parents=True, exist_ok=True)
-                        torch.save(ema.state_dict(), ema_path)
-                    except Exception as exc:  # noqa: BLE001
-                        logging.warning("Failed to save EMA shadow: %s", exc)
-
                 if cfg.save_checkpoint_to_hub:
                     push_checkpoint_to_hub(
                         checkpoint_dir,
@@ -865,16 +789,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             if is_main_process:
                 step_id = get_step_identifier(step, cfg.steps)
                 logging.info(f"Eval policy at step {step}")
-                # Use the EMA shadow model for eval when enabled —
+                eval_target_policy = accelerator.unwrap_model(policy)
-                # standard practice for diffusion-style policies (~1–3%
-                # lift on closed-loop success). ``ema.ema_model`` is a
-                # full nn.Module clone, so we just pass it through; no
-                # swap/restore on the live policy needed.
-                eval_target_policy = (
-                    ema.ema_model
-                    if (ema is not None and cfg.ema.use_for_eval)
-                    else accelerator.unwrap_model(policy)
-                )
                 with torch.no_grad(), accelerator.autocast():
                     eval_info = eval_policy_all(
                         envs=eval_env,  # dict[suite][task_id] -> vec_env
@@ -941,26 +856,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             preprocessor.push_to_hub(active_cfg.repo_id)
             postprocessor.push_to_hub(active_cfg.repo_id)
 
-            # When EMA is on we *eval* the EMA weights but the push above
-            # ships the live weights — they're different models. Push the EMA
-            # weights too, to a sibling ``<repo_id>-ema`` repo, so both are
-            # fully loadable and you can benchmark/deploy whichever is better.
-            # Non-fatal: the live model is already up if this fails.
-            if ema is not None and not (not cfg.is_reward_model_training and cfg.policy.use_peft):
-                ema_model = ema.ema_model
-                ema_repo_id = f"{active_cfg.repo_id}-ema"
-                orig_repo_id = ema_model.config.repo_id
-                try:
-                    ema_model.config.repo_id = ema_repo_id
-                    ema_model.push_model_to_hub(cfg)
-                    preprocessor.push_to_hub(ema_repo_id)
-                    postprocessor.push_to_hub(ema_repo_id)
-                    logging.info("Pushed EMA weights to %s", ema_repo_id)
-                except Exception as exc:  # noqa: BLE001
-                    logging.warning("Failed to push EMA weights to %s: %s", ema_repo_id, exc)
-                finally:
-                    ema_model.config.repo_id = orig_repo_id
-
     # Properly clean up the distributed process group
     accelerator.wait_for_everyone()
     accelerator.end_training()

uv.lock

@@ -1430,19 +1430,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/64/cb/809f0c3e4e7bfe78c6dd468631896a8866c3ba853e3c855cc3fa58fae660/eiquadprog-1.2.9-0-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:416f4b584ea30072f166b2a6a3e0a63a2a260a378f9bcbd2dfc9cde13b810a50", size = 118538, upload-time = "2025-02-17T19:00:16.297Z" },
 ]
 
-[[package]]
-name = "ema-pytorch"
-version = "0.7.9"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" },
-    { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux'" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/c7/b8/de3bb1eacd77a81eb843fb098c2eaeac3bbe6777623a15c786e03fdc44be/ema_pytorch-0.7.9.tar.gz", hash = "sha256:a8ccdf2eeecce5489de02fc7c9776ef55400220afff92b8223f7516cb570b594", size = 10564, upload-time = "2025-12-19T21:30:04.154Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/9f/99/d2c69a90d2d666ff9ef45426992bd2c7d578068ea6324a1a46e6020c636d/ema_pytorch-0.7.9-py3-none-any.whl", hash = "sha256:fe6f236d16e879d7d3cf4946fa1eed395c01f42d3df6a61a9f0ebebbc9aae56d", size = 11538, upload-time = "2025-12-19T21:30:03.122Z" },
-]
-
 [[package]]
 name = "etils"
 version = "1.14.0"
@@ -2842,7 +2829,6 @@ dependencies = [
     { name = "cmake" },
     { name = "draccus" },
     { name = "einops" },
-    { name = "ema-pytorch" },
     { name = "gymnasium" },
     { name = "huggingface-hub" },
     { name = "numpy" },
@@ -3318,7 +3304,6 @@ requires-dist = [
     { name = "draccus", specifier = "==0.10.0" },
     { name = "dynamixel-sdk", marker = "extra == 'dynamixel'", specifier = ">=3.7.31,<3.9.0" },
     { name = "einops", specifier = ">=0.8.0,<0.9.0" },
-    { name = "ema-pytorch", specifier = ">=0.7.7,<1.0.0" },
     { name = "faker", marker = "extra == 'sarm'", specifier = ">=33.0.0,<35.0.0" },
     { name = "fastapi", marker = "extra == 'phone'", specifier = "<1.0" },
     { name = "feetech-servo-sdk", marker = "extra == 'feetech'", specifier = ">=1.0.0,<2.0.0" },