VLA/wall_x/trainer/optimizer/utils.py

import inspect

from torch.optim import AdamW

from wall_x.config.hyperparams_config import LRGroupConfig

_OPTIMIZERS = {}


def register_optimizer(name, optimizer_fn):
    _OPTIMIZERS[name] = optimizer_fn


def get_optimizer(name, *args, **kwargs):
    if name not in _OPTIMIZERS:
        raise KeyError(f"Unknown optimizer '{name}'. Registered: {sorted(_OPTIMIZERS)}")
    return _OPTIMIZERS[name](*args, **kwargs)


def _group_weight_decay(opt_cfg):
    return getattr(opt_cfg, "weight_decay", None)


def resolve_lr_group_configs(opt_cfg, default_action_lr_keywords):
    """Return structured LR groups, with legacy action config fallback.

    ``optimizer.lr_groups`` is the preferred path. The legacy
    ``action_expert_learning_rate`` fields are still converted into a single
    action group so older configs keep working.
    """
    if opt_cfg.lr_groups:
        if (
            opt_cfg.action_expert_learning_rate is not None
            or opt_cfg.action_lr_keywords is not None
        ):
            raise ValueError(
                "Use either optimizer.lr_groups or legacy "
                "action_expert_learning_rate/action_lr_keywords, not both."
            )
        return opt_cfg.lr_groups

    if opt_cfg.action_expert_learning_rate is None:
        return []

    action_lr_keywords = (
        opt_cfg.action_lr_keywords
        if opt_cfg.action_lr_keywords is not None
        else default_action_lr_keywords
    )
    return [
        LRGroupConfig(
            name="action_lr_group",
            lr=opt_cfg.action_expert_learning_rate,
            include=action_lr_keywords,
            fail_on_empty=True,
        )
    ]


def uses_legacy_action_lr_groups(opt_cfg) -> bool:
    return not opt_cfg.lr_groups and opt_cfg.action_expert_learning_rate is not None


def _make_param_group(name, params, lr, opt_cfg):
    group = {
        "params": params,
        "lr": lr,
        "group_name": name,
    }
    weight_decay = _group_weight_decay(opt_cfg)
    if weight_decay is not None:
        group["weight_decay"] = weight_decay
    return group


def _validate_lr_group(group: LRGroupConfig, *, base_group_name: str):
    if not group.name:
        raise ValueError("optimizer.lr_groups entries must have a non-empty name")
    if group.name == base_group_name:
        raise ValueError(
            f"optimizer.lr_groups name {group.name!r} is reserved for the base group"
        )
    if "/" in group.name:
        raise ValueError(
            f"optimizer.lr_groups name {group.name!r} must not contain '/'. "
            "DMuon appends '/muon' and '/adamw' to group names."
        )
    if not group.include:
        raise ValueError(
            f"optimizer.lr_groups.{group.name} must define at least one include keyword"
        )


def build_lr_param_groups(model, opt_cfg, lr_groups, *, base_group_name="base"):
    """Split trainable params into named LR groups plus a base group.

    ``lr_groups`` is a list of :class:`LRGroupConfig`. Each group matches
    parameter names by substring. A parameter may match at most one explicit
    group; unmatched trainable parameters remain in the ``base`` group using
    ``opt_cfg.learning_rate``.

    Returns a list of torch.optim-compatible param_group dicts. The
    ``group_name`` key is non-standard but preserved by torch.optim via
    ``setdefault`` in ``add_param_group`` and is consumed downstream for
    per-group lr logging.

    For AdamW / native Muon, each returned group includes ``weight_decay``.
    For DMuon, weight decay is split between Muon and AdamW defaults, so the
    groups only carry lr and metadata; DMuon applies its own per-route defaults.
    The caller is expected to gate on ``opt_cfg.optimizer_type`` before calling
    this.
    """
    if not lr_groups:
        return None

    names = [group.name for group in lr_groups]
    duplicate_names = sorted({name for name in names if names.count(name) > 1})
    if duplicate_names:
        raise ValueError(
            f"optimizer.lr_groups contains duplicate names: {duplicate_names}"
        )

    for group in lr_groups:
        _validate_lr_group(group, base_group_name=base_group_name)

    base_params = []
    grouped_params = {group.name: [] for group in lr_groups}
    ambiguous = []

    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        matches = [
            group.name
            for group in lr_groups
            if any(keyword in name for keyword in group.include)
        ]
        if len(matches) > 1:
            ambiguous.append((name, matches))
            continue
        if matches:
            grouped_params[matches[0]].append(param)
        else:
            base_params.append(param)

    if ambiguous:
        examples = ", ".join(f"{name} -> {matches}" for name, matches in ambiguous[:10])
        raise ValueError(
            "Some parameters match multiple optimizer.lr_groups. Make group "
            f"include patterns disjoint. Examples: {examples}"
        )

    if opt_cfg.train_action_expert_only:
        assert len(base_params) == 0, (
            f"Expected 0 base_params after pre-wrap freeze, got {len(base_params)}. "
            "Ensure base params are frozen before building the optimizer."
        )

    param_groups = []
    if len(base_params) > 0:
        param_groups.append(
            _make_param_group(
                base_group_name, base_params, opt_cfg.learning_rate, opt_cfg
            )
        )

    for group in lr_groups:
        params = grouped_params[group.name]
        if len(params) == 0:
            if group.fail_on_empty:
                raise ValueError(
                    f"No params found for optimizer.lr_groups.{group.name}. "
                    f"Please check include={group.include!r}."
                )
            continue
        param_groups.append(_make_param_group(group.name, params, group.lr, opt_cfg))

    return param_groups


def build_action_expert_param_groups(model, opt_cfg, action_lr_keywords):
    """Compatibility wrapper for the legacy action-expert LR config."""
    return build_lr_param_groups(
        model,
        opt_cfg,
        [
            LRGroupConfig(
                name="action_lr_group",
                lr=opt_cfg.action_expert_learning_rate,
                include=action_lr_keywords,
                fail_on_empty=True,
            )
        ],
        base_group_name="base_lr_group",
    )


def get_adamw_optimizer(model, *, opt_cfg, param_groups=None):
    """Build AdamW from AdamWConfig."""
    from wall_x.config.hyperparams_config import AdamWConfig

    if not isinstance(opt_cfg, AdamWConfig):
        raise TypeError(
            f"get_adamw_optimizer expects AdamWConfig, got {type(opt_cfg).__name__}"
        )

    if param_groups is None:
        params = [p for p in model.parameters() if p.requires_grad]
    else:
        # Per-group lr / weight_decay are preserved by torch.optim (via
        # setdefault in add_param_group), so top-level values act only as
        # defaults. Extra keys like ``group_name`` are kept in-place and used
        # downstream for per-group lr logging.
        params = param_groups

    kw = {
        "lr": opt_cfg.learning_rate,
        "weight_decay": opt_cfg.weight_decay,
        "betas": tuple(opt_cfg.betas),
        "eps": opt_cfg.eps,
    }
    sig_params = inspect.signature(AdamW.__init__).parameters
    if "foreach" in sig_params and opt_cfg.foreach is not None:
        kw["foreach"] = opt_cfg.foreach
    if "fused" in sig_params:
        kw["fused"] = opt_cfg.fused
    return AdamW(params, **kw)


register_optimizer("adamw", get_adamw_optimizer)