mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-09 02:51:56 +00:00
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>
This commit is contained in:
pepijn
@@ -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",
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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 ``<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()
|
||||
|
||||
@@ -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" },
|
||||
|
||||
Reference in New Issue
Block a user