lerobot/src/lerobot/policies/evo1/configuration_evo1.py

# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import math
from dataclasses import dataclass, field

from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR

from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.optim.optimizers import AdamWConfig
from lerobot.optim.schedulers import LRSchedulerConfig
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE


@LRSchedulerConfig.register_subclass("evo1_exact")
@dataclass
class Evo1SchedulerConfig(LRSchedulerConfig):
    num_warmup_steps: int

    def build(self, optimizer: Optimizer, num_training_steps: int) -> LambdaLR:
        def lr_lambda(current_step: int) -> float:
            if current_step < self.num_warmup_steps:
                return current_step / max(1, self.num_warmup_steps)
            progress = (current_step - self.num_warmup_steps) / max(
                1, num_training_steps - self.num_warmup_steps
            )
            return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))

        return LambdaLR(optimizer, lr_lambda, -1)


@PreTrainedConfig.register_subclass("evo1")
@dataclass
class Evo1Config(PreTrainedConfig):
    training_stage: str = "stage1"
    use_amp: bool = True

    n_obs_steps: int = 1
    chunk_size: int = 50
    n_action_steps: int = 50

    max_state_dim: int = 24
    max_action_dim: int = 24
    max_views: int = 3
    image_resolution: tuple[int, int] = (448, 448)
    empty_cameras: int = 0

    normalization_mapping: dict[str, NormalizationMode] = field(
        default_factory=lambda: {
            "VISUAL": NormalizationMode.IDENTITY,
            "STATE": NormalizationMode.MIN_MAX,
            "ACTION": NormalizationMode.MIN_MAX,
        }
    )

    vlm_model_name: str = "OpenGVLab/InternVL3-1B"
    vlm_num_layers: int | None = 14
    vlm_dtype: str = "bfloat16"
    use_flash_attn: bool = True
    action_head: str = "flowmatching"
    embed_dim: int = 896
    hidden_dim: int = 1024
    state_hidden_dim: int = 1024
    num_heads: int = 8
    num_layers: int = 8
    dropout: float = 0.0
    num_inference_timesteps: int = 32
    num_categories: int = 1
    return_cls_only: bool = False
    enable_gradient_checkpointing: bool = True
    gradient_checkpointing_use_reentrant: bool = False

    finetune_vlm: bool | None = None
    finetune_language_model: bool | None = None
    finetune_vision_model: bool | None = None
    finetune_action_head: bool | None = None
    # Reapply stage defaults after loading checkpoint configs so stage2 cannot
    # accidentally inherit the frozen VLM flags stored by a stage1 checkpoint.
    apply_training_stage_defaults: bool = True

    task_field: str = "task"
    embodiment_id_field: str | None = None
    default_embodiment_id: int = 0

    optimizer_lr: float = 1e-5
    optimizer_betas: tuple[float, float] = (0.9, 0.999)
    optimizer_eps: float = 1e-8
    optimizer_weight_decay: float = 1e-5
    optimizer_grad_clip_norm: float = 1.0

    scheduler_warmup_steps: int = 300
    drop_last: bool = True

    def __post_init__(self):
        super().__post_init__()
        if self.training_stage not in {"stage1", "stage2"}:
            raise ValueError(
                f"Unsupported EVO1 training_stage '{self.training_stage}', expected 'stage1' or 'stage2'"
            )

        if self.apply_training_stage_defaults:
            if self.training_stage == "stage1":
                self.finetune_vlm = False
                self.finetune_language_model = False
                self.finetune_vision_model = False
                self.finetune_action_head = True
            elif self.training_stage == "stage2":
                self.finetune_vlm = True
                self.finetune_language_model = True
                self.finetune_vision_model = True
                self.finetune_action_head = True
        elif self.training_stage == "stage1":
            if self.finetune_vlm is None:
                self.finetune_vlm = False
            if self.finetune_language_model is None:
                self.finetune_language_model = False
            if self.finetune_vision_model is None:
                self.finetune_vision_model = False
            if self.finetune_action_head is None:
                self.finetune_action_head = True
        elif self.training_stage == "stage2":
            has_explicit_branch_flags = any(
                flag is not None for flag in (self.finetune_language_model, self.finetune_vision_model)
            )
            if not has_explicit_branch_flags:
                if self.finetune_vlm is None:
                    self.finetune_vlm = True
                if self.finetune_language_model is None:
                    self.finetune_language_model = True
                if self.finetune_vision_model is None:
                    self.finetune_vision_model = True
            elif self.finetune_vlm is None:
                self.finetune_vlm = bool(self.finetune_language_model or self.finetune_vision_model)
            if self.finetune_action_head is None:
                self.finetune_action_head = True

        if self.finetune_vlm is None:
            self.finetune_vlm = False
        if self.finetune_language_model is None:
            self.finetune_language_model = False
        if self.finetune_vision_model is None:
            self.finetune_vision_model = False
        if self.finetune_action_head is None:
            self.finetune_action_head = False

        branch_vlm = self.finetune_language_model or self.finetune_vision_model
        if self.finetune_vlm != branch_vlm:
            raise ValueError(
                "Inconsistent EVO1 finetune config: "
                f"finetune_vlm={self.finetune_vlm} but "
                f"(finetune_language_model or finetune_vision_model)={branch_vlm}. "
                "When branch-level flags are used, finetune_vlm must match their effective union."
            )

        if self.n_action_steps > self.chunk_size:
            raise ValueError(
                f"n_action_steps ({self.n_action_steps}) must be <= chunk_size ({self.chunk_size})"
            )

    def validate_features(self) -> None:
        if self.input_features is None:
            self.input_features = {}
        if self.output_features is None:
            self.output_features = {}

        for i in range(self.empty_cameras):
            key = OBS_IMAGES + f".empty_camera_{i}"
            if key not in self.input_features:
                self.input_features[key] = PolicyFeature(
                    type=FeatureType.VISUAL,
                    shape=(3, *self.image_resolution),
                )

        if OBS_STATE not in self.input_features:
            self.input_features[OBS_STATE] = PolicyFeature(
                type=FeatureType.STATE,
                shape=(self.max_state_dim,),
            )

        if ACTION not in self.output_features:
            self.output_features[ACTION] = PolicyFeature(
                type=FeatureType.ACTION,
                shape=(self.max_action_dim,),
            )

    def get_optimizer_preset(self) -> AdamWConfig:
        return AdamWConfig(
            lr=self.optimizer_lr,
            betas=self.optimizer_betas,
            eps=self.optimizer_eps,
            weight_decay=self.optimizer_weight_decay,
            grad_clip_norm=self.optimizer_grad_clip_norm,
        )

    def get_scheduler_preset(self):
        return Evo1SchedulerConfig(
            num_warmup_steps=self.scheduler_warmup_steps,
        )

    @property
    def observation_delta_indices(self) -> list[int]:
        return [0]

    @property
    def action_delta_indices(self) -> list[int]:
        return list(range(self.chunk_size))

    @property
    def reward_delta_indices(self) -> None:
        return None