From 147b8f248df00c1e2296828b5b815c57be08acea Mon Sep 17 00:00:00 2001 From: pepijn Date: Wed, 8 Jul 2026 11:15:33 +0000 Subject: [PATCH] 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 `-ema` hub push), and the ema-pytorch dependency. Co-authored-by: Cursor --- pyproject.toml | 5 -- src/lerobot/configs/default.py | 13 ---- src/lerobot/configs/train.py | 3 +- src/lerobot/scripts/lerobot_train.py | 107 +-------------------------- uv.lock | 15 ---- 5 files changed, 2 insertions(+), 141 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6a9b7087c..fffd11c6e 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/lerobot/configs/default.py b/src/lerobot/configs/default.py index bc0891437..b51d1000e 100644 --- a/src/lerobot/configs/default.py +++ b/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 diff --git a/src/lerobot/configs/train.py b/src/lerobot/configs/train.py index c9030f2d8..a977665a2 100644 --- a/src/lerobot/configs/train.py +++ b/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. diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 09427f338..34e7d2021 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/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 — - # 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) - ) + eval_target_policy = 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 ``-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() diff --git a/uv.lock b/uv.lock index 4a3304a2c..70d787bd8 100644 --- a/uv.lock +++ b/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" },