feat(policies): add EVO1 policy (#3908)

* feat(policies): add EVO1 policy

* fix(evo1): infer batch size after normalizing image dims

`_collect_image_batches` read `batch_size = batch[camera_keys[0]].shape[0]`
before normalizing per-camera tensors to `(B, C, H, W)`. For an unbatched
`(C, H, W)` input (which the function tries to support via the `image.dim() == 3`
branch), this picked up the channel count `C` instead of the real batch size,
making the subsequent per-sample loop iterate `C` times and indexing go
out of bounds.

Normalize each camera tensor up-front, then read `batch_size` from the
normalized batch dim. Adds `test_collect_image_batches_handles_unbatched_chw`
covering the regression.

Reported by Copilot review on huggingface/lerobot#3545.

* chore(lock): regenerate uv.lock for evo1 extra

Adds the `evo1` entry to `[package.metadata.requires-dist]` and the
`provides-extras` list so that `uv sync --locked --extra test` (used by
fast_tests.yml) no longer reports the lockfile as stale.

Generated with `uv 0.8.0` (matching `UV_VERSION` in fast_tests.yml).
The non-evo1 marker tweaks are produced by `uv lock` re-resolving the
existing dep graph and are not introduced by this PR.

* chore(evo1): align with policy contribution guide conventions

- Add `src/lerobot/policies/evo1/README.md` symlink into `docs/source/evo1.mdx`
  to match the in-tree README convention (mirroring the EO-1 layout).
- Convert `transformers` import in `internvl3_embedder.py` to the standard
  `TYPE_CHECKING + _transformers_available` two-step gating used by other
  optional-backbone policies (e.g. diffusion). The previous lazy-in-`__init__`
  import was functionally equivalent for runtime gating but didn't expose the
  real symbols to type checkers.
- Add `lerobot[evo1]` to the `all` extra in `pyproject.toml` so
  `pip install 'lerobot[all]'` keeps installing every optional policy.

Per the guidance in https://moon-ci-docs.huggingface.co/docs/lerobot/pr_3534/en/contributing_a_policy.

* fix(evo1): finalize policy guide alignment

* docs(evo1): format results table

* Fix EVO1 LIBERO rollout processors

* Fix EVO1 LIBERO eval action postprocessing

* Fix eval action conversion for bf16 policies

* fix(evo1): move LIBERO padding into policy processors

* refactor(evo1): use native HF InternVL3-1B-hf, drop trust_remote_code

- Switch from OpenGVLab/InternVL3-1B (requires trust_remote_code=True)
  to OpenGVLab/InternVL3-1B-hf (native transformers implementation).
- Replace manual _extract_feature + _prepare_and_fuse_embeddings with
  a single model.forward() call — verified bit-for-bit identical output.
- Remove ~170 lines of manual ViT/pixel-shuffle/projection logic.
- Symlink README.md to docs/source/ following repo convention.

Weights are byte-identical between both model variants; only the module
naming differs. All 12 existing unit tests pass. Local training (10 steps)
on maximellerbach/omx_pickandplace confirmed working.

* refactor(policy): evo1 GPU-batched preprocessing +  vectorized attention masking + remove dead code

* fix(style): pre-commit

oops

* chore(evo1): delete added test + reduce diff

* refactor(policies): use config for evo1 + local imports

* refactor(policies): multiple improvements

* chore: update docs + remove legacy codepaths

* feat(policies): implement RTC to EVO1

---------

Co-authored-by: javadcc_mac <javadcc1@sjtu.edu.cn>
Co-authored-by: Yiming Wang <145452074+JAVAdcc@users.noreply.github.com>
Co-authored-by: Martino Russi <nopyeps@gmail.com>

src/lerobot/optim/schedulers.py

@@ -105,6 +105,28 @@ class ConstantWithWarmupSchedulerConfig(LRSchedulerConfig):
         return LambdaLR(optimizer, lr_lambda, -1)
 
 
+@LRSchedulerConfig.register_subclass("cosine_annealing_with_warmup")
+@dataclass
+class CosineAnnealingWithWarmupSchedulerConfig(LRSchedulerConfig):
+    """Linear warmup followed by cosine annealing from the peak LR to zero.
+
+    Used by EVO1; the annealing phase always spans the remaining training steps.
+    """
+
+    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)
+
+
 @LRSchedulerConfig.register_subclass("cosine_decay_with_warmup")
 @dataclass
 class CosineDecayWithWarmupSchedulerConfig(LRSchedulerConfig):

src/lerobot/policies/__init__.py

@@ -17,6 +17,7 @@ from lerobot.utils.action_interpolator import ActionInterpolator as ActionInterp
 from .act.configuration_act import ACTConfig as ACTConfig
 from .diffusion.configuration_diffusion import DiffusionConfig as DiffusionConfig
 from .eo1.configuration_eo1 import EO1Config as EO1Config
+from .evo1.configuration_evo1 import Evo1Config as Evo1Config
 from .factory import get_policy_class, make_policy, make_policy_config, make_pre_post_processors
 from .fastwam.configuration_fastwam import FastWAMConfig as FastWAMConfig
 from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig as GaussianActorConfig
@@ -46,6 +47,7 @@ __all__ = [
     "EO1Config",
     "FastWAMConfig",
     "GaussianActorConfig",
+    "Evo1Config",
     "GrootConfig",
     "LingBotVAConfig",
     "MolmoAct2Config",

src/lerobot/policies/evo1/README.md

@@ -0,0 +1 @@
+../../../../docs/source/policy_evo1_README.md

src/lerobot/policies/evo1/__init__.py

@@ -0,0 +1,19 @@
+# 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 .configuration_evo1 import Evo1Config
+from .modeling_evo1 import Evo1Policy
+from .processor_evo1 import make_evo1_pre_post_processors
+
+__all__ = ["Evo1Config", "Evo1Policy", "make_evo1_pre_post_processors"]

src/lerobot/policies/evo1/configuration_evo1.py

@@ -0,0 +1,252 @@
+# 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 logging
+from dataclasses import dataclass, field
+
+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 CosineAnnealingWithWarmupSchedulerConfig
+from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
+
+from ..rtc.configuration_rtc import RTCConfig
+
+logger = logging.getLogger(__name__)
+
+
+@PreTrainedConfig.register_subclass("evo1")
+@dataclass
+class Evo1Config(PreTrainedConfig):
+    training_stage: str = "stage1"
+    # When True and the policy runs on CUDA, EVO1 wraps its own forward passes (training and
+    # inference) in a bfloat16 autocast block, so its numerics do not depend on the dtype of any
+    # outer autocast context opened by lerobot-train/lerobot-eval.
+    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
+    postprocess_action_dim: int | None = None
+    binarize_gripper: bool = False
+    gripper_index: int = 6
+    gripper_threshold: float = 0.5
+    gripper_below_threshold_value: float = 1.0
+    gripper_above_threshold_value: float = -1.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-hf"
+    vlm_num_layers: int | None = 14
+    vlm_dtype: str = "bfloat16"
+    # Max token length for tokenizing the (image placeholders + instruction) prompt. Prompts longer
+    # than this are right-truncated, so raise it for tasks with long language instructions or many views.
+    max_text_length: int = 1024
+    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
+    # When True, the action head is conditioned on a single pooled VL token (the last non-padding
+    # token of the causal decoder) instead of the full fused token sequence.
+    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
+
+    # Real-Time Chunking guidance for asynchronous inference (lerobot-rollout --inference.type=rtc
+    # sets this and calls init_rtc_processor()); None disables RTC.
+    rtc_config: RTCConfig | None = None
+
+    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
+
+    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:
+            stage_defaults = {
+                "stage1": {
+                    "finetune_vlm": False,
+                    "finetune_language_model": False,
+                    "finetune_vision_model": False,
+                    "finetune_action_head": True,
+                },
+                "stage2": {
+                    "finetune_vlm": True,
+                    "finetune_language_model": True,
+                    "finetune_vision_model": True,
+                    "finetune_action_head": True,
+                },
+            }[self.training_stage]
+            for flag_name, default_value in stage_defaults.items():
+                current_value = getattr(self, flag_name)
+                if current_value is not None and current_value != default_value:
+                    logger.warning(
+                        "EVO1 %s=%s is overridden by training_stage=%s default %s. "
+                        "Set apply_training_stage_defaults=false to keep explicit finetuning flags.",
+                        flag_name,
+                        current_value,
+                        self.training_stage,
+                        default_value,
+                    )
+                setattr(self, flag_name, default_value)
+        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:
+                # An explicit finetune_vlm decides both branches; otherwise stage2 defaults to a
+                # full-VLM finetune.
+                vlm_finetune = self.finetune_vlm if self.finetune_vlm is not None else True
+                self.finetune_vlm = vlm_finetune
+                self.finetune_language_model = vlm_finetune
+                self.finetune_vision_model = vlm_finetune
+            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})"
+            )
+        if len(self.image_resolution) != 2 or self.image_resolution[0] != self.image_resolution[1]:
+            raise ValueError(
+                "EVO1 currently expects a square image_resolution because InternVL3 preprocessing "
+                f"uses a scalar image_size, got {self.image_resolution}."
+            )
+        if not 0 <= self.default_embodiment_id < self.num_categories:
+            raise ValueError(
+                f"default_embodiment_id ({self.default_embodiment_id}) must be in "
+                f"[0, num_categories={self.num_categories})"
+            )
+
+    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 CosineAnnealingWithWarmupSchedulerConfig(
+            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

src/lerobot/policies/evo1/evo1_model.py

@@ -0,0 +1,210 @@
+# 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 torch
+import torch.nn as nn
+
+from .configuration_evo1 import Evo1Config
+from .flow_matching import FlowmatchingActionHead
+from .internvl3_embedder import InternVL3Embedder
+
+
+class Evo1Model(nn.Module):
+    def __init__(self, config: Evo1Config, vlm_hub_kwargs: dict | None = None):
+        super().__init__()
+        self.config = config
+        self._device = config.device
+        self.return_cls_only = config.return_cls_only
+        # Set by Evo1Policy.init_rtc_processor() when config.rtc_config is provided.
+        self.rtc_processor = None
+
+        # Gradient checkpointing only pays off when the VLM is actually being trained; keep it off
+        # whenever every VLM branch is frozen so the frozen forward stays cheap.
+        tracks_vlm_gradients = bool(
+            config.finetune_vlm or config.finetune_language_model or config.finetune_vision_model
+        )
+        enable_gradient_checkpointing = config.enable_gradient_checkpointing and tracks_vlm_gradients
+
+        self.embedder = InternVL3Embedder(
+            model_name=config.vlm_model_name,
+            image_size=int(config.image_resolution[0]),
+            device=self._device,
+            num_language_layers=config.vlm_num_layers,
+            model_dtype=config.vlm_dtype,
+            use_flash_attn=config.use_flash_attn,
+            max_text_length=config.max_text_length,
+            enable_gradient_checkpointing=enable_gradient_checkpointing,
+            gradient_checkpointing_use_reentrant=config.gradient_checkpointing_use_reentrant,
+            hub_kwargs=vlm_hub_kwargs,
+        )
+
+        action_head_type = config.action_head.lower()
+        if action_head_type != "flowmatching":
+            raise NotImplementedError(f"Unknown action_head: {action_head_type}")
+
+        horizon = config.chunk_size
+        per_action_dim = config.max_action_dim
+        action_dim = horizon * per_action_dim
+
+        self.horizon = horizon
+        self.per_action_dim = per_action_dim
+        self.action_head = FlowmatchingActionHead(
+            embed_dim=config.embed_dim,
+            hidden_dim=config.hidden_dim,
+            action_dim=action_dim,
+            horizon=horizon,
+            per_action_dim=per_action_dim,
+            num_heads=config.num_heads,
+            num_layers=config.num_layers,
+            dropout=config.dropout,
+            num_inference_timesteps=config.num_inference_timesteps,
+            num_categories=config.num_categories,
+            state_dim=config.max_state_dim,
+            state_hidden_dim=config.state_hidden_dim,
+        ).to(self._device)
+
+    def get_vl_embeddings(
+        self,
+        images: list[torch.Tensor],
+        image_mask: torch.Tensor,
+        prompt: str | list[str] | None = None,
+        return_cls_only: bool | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Fused VL embeddings from per-camera image batches.
+
+        Args:
+            images: list of per-camera tensors, each shaped ``(B, C, H, W)`` with values in ``[0, 1]``.
+            image_mask: bool tensor ``(B, max_views)`` marking present views.
+
+        Returns:
+            ``(embeddings, valid_mask)``: the fused tokens and the bool mask of attendable context
+            positions (None when a single pooled token is returned).
+        """
+        if return_cls_only is None:
+            return_cls_only = self.return_cls_only
+        if not images:
+            raise ValueError("EVO1 expects at least one image per sample.")
+
+        batch_size = images[0].shape[0]
+        if prompt is None:
+            prompts = [""] * batch_size
+        elif isinstance(prompt, str):
+            prompts = [prompt] * batch_size
+        else:
+            prompts = [str(p) for p in prompt]
+            if len(prompts) != batch_size:
+                raise ValueError(
+                    f"Prompt batch size {len(prompts)} does not match image batch size {batch_size}"
+                )
+
+        if image_mask.dim() == 1:
+            image_mask = image_mask.unsqueeze(0)
+        if image_mask.shape[0] != batch_size:
+            raise ValueError(
+                f"image_mask batch size {image_mask.shape[0]} does not match image batch size {batch_size}"
+            )
+
+        return self.embedder.get_fused_image_text_embedding_batched(
+            camera_images=images,
+            image_masks=image_mask,
+            text_prompts=prompts,
+            return_cls_only=return_cls_only,
+        )
+
+    def predict_action(
+        self,
+        fused_tokens: torch.Tensor,
+        state: torch.Tensor,
+        actions_gt: torch.Tensor | None = None,
+        action_mask: torch.Tensor | None = None,
+        embodiment_ids: torch.Tensor | None = None,
+        context_mask: torch.Tensor | None = None,
+        inference_delay: int | None = None,
+        prev_chunk_left_over: torch.Tensor | None = None,
+        execution_horizon: int | None = None,
+    ):
+        if actions_gt is None:
+            return self.action_head.get_action(
+                fused_tokens,
+                state=state,
+                action_mask=action_mask,
+                embodiment_id=embodiment_ids,
+                context_mask=context_mask,
+                inference_delay=inference_delay,
+                prev_chunk_left_over=prev_chunk_left_over,
+                execution_horizon=execution_horizon,
+                rtc_processor=self.rtc_processor,
+            )
+        return self.action_head(
+            fused_tokens,
+            state=state,
+            actions_gt=actions_gt,
+            action_mask=action_mask,
+            embodiment_id=embodiment_ids,
+            context_mask=context_mask,
+        )
+
+    def forward(
+        self,
+        fused_tokens: torch.Tensor,
+        state: torch.Tensor | None = None,
+        actions_gt: torch.Tensor | None = None,
+        action_mask: torch.Tensor | None = None,
+        embodiment_ids: torch.Tensor | None = None,
+        context_mask: torch.Tensor | None = None,
+        inference_delay: int | None = None,
+        prev_chunk_left_over: torch.Tensor | None = None,
+        execution_horizon: int | None = None,
+    ):
+        return self.predict_action(
+            fused_tokens,
+            state,
+            actions_gt,
+            action_mask,
+            embodiment_ids,
+            context_mask,
+            inference_delay,
+            prev_chunk_left_over,
+            execution_horizon,
+        )
+
+    def _set_module_trainable(self, module: nn.Module, trainable: bool):
+        for param in module.parameters():
+            param.requires_grad = trainable
+
+    def _vlm_submodule(self, name: str) -> nn.Module:
+        module = getattr(self.embedder.model, name, None)
+        if not isinstance(module, nn.Module):
+            raise AttributeError(
+                f"InternVL model {type(self.embedder.model).__name__} has no '{name}' submodule; "
+                "the native HF InternVL layout (language_model / vision_tower / "
+                "multi_modal_projector) is required to apply the EVO1 finetune flags."
+            )
+        return module
+
+    def set_finetune_flags(self):
+        # __post_init__ resolves every finetune flag to a concrete boolean, so branch-level flags
+        # are authoritative here. Freeze everything first, then re-enable the requested branches.
+        self._set_module_trainable(self.embedder, False)
+        self._set_module_trainable(
+            self._vlm_submodule("language_model"), bool(self.config.finetune_language_model)
+        )
+        finetune_vision = bool(self.config.finetune_vision_model)
+        self._set_module_trainable(self._vlm_submodule("vision_tower"), finetune_vision)
+        self._set_module_trainable(self._vlm_submodule("multi_modal_projector"), finetune_vision)
+
+        if not self.config.finetune_action_head:
+            self._set_module_trainable(self.action_head, False)

src/lerobot/policies/evo1/flow_matching.py

@@ -0,0 +1,483 @@
+# 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 logging
+import math
+
+import torch
+import torch.nn as nn
+
+logger = logging.getLogger(__name__)
+
+
+class SinusoidalPositionalEncoding(nn.Module):
+    def __init__(self, dim: int, max_len: int = 1000):
+        super().__init__()
+        pe = torch.zeros(max_len, dim)
+        position = torch.arange(0, max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, dim, 2) * -(math.log(10000.0) / dim))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer("pe", pe)
+
+    def forward(self, seq_len: int):
+        if seq_len > self.pe.size(1):
+            self._extend_pe(seq_len)
+        return self.pe[:, :seq_len, :]
+
+    def _extend_pe(self, new_max_len):
+        old_max_len, dim = self.pe.size(1), self.pe.size(2)
+        if new_max_len <= old_max_len:
+            return
+        extra_positions = torch.arange(old_max_len, new_max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, dim, 2, dtype=torch.float) * -(math.log(10000.0) / dim))
+        extra_pe = torch.zeros(new_max_len - old_max_len, dim)
+        extra_pe[:, 0::2] = torch.sin(extra_positions * div_term)
+        extra_pe[:, 1::2] = torch.cos(extra_positions * div_term)
+        extra_pe = extra_pe.unsqueeze(0)
+        new_pe = torch.cat([self.pe, extra_pe.to(self.pe.device)], dim=1)
+        self.pe = new_pe
+
+
+class CategorySpecificLinear(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, num_categories: int = 1):
+        super().__init__()
+        self.num_categories = num_categories
+        if num_categories <= 1:
+            self.linear = nn.Linear(in_dim, out_dim)
+        else:
+            self.weight = nn.Parameter(torch.empty(num_categories, in_dim, out_dim))
+            self.bias = nn.Parameter(torch.zeros(num_categories, out_dim))
+            # Initialize each per-category (in_dim, out_dim) matrix separately: xavier on the full
+            # 3D tensor would compute fan_in = in_dim * out_dim and badly under-scale the weights.
+            for category in range(num_categories):
+                nn.init.xavier_uniform_(self.weight[category])
+
+    def forward(self, x: torch.Tensor, category_id: torch.LongTensor):
+        if self.num_categories <= 1:
+            if x.dtype != self.linear.weight.dtype:
+                x = x.to(dtype=self.linear.weight.dtype)
+            return self.linear(x)
+
+        if x.dtype != self.weight.dtype:
+            x = x.to(dtype=self.weight.dtype)
+
+        orig_shape = x.shape
+        x_flat = x.reshape(-1, orig_shape[-1])
+        if category_id.dim() == 0:
+            cid = category_id.item()
+            out = x_flat @ self.weight[cid] + self.bias[cid]
+        else:
+            category_id = category_id.reshape(-1)
+            if category_id.numel() != x_flat.size(0):
+                raise ValueError(
+                    f"category_id length {category_id.numel()} does not match flattened batch {x_flat.size(0)}"
+                )
+            weight_selected = self.weight[category_id]
+            bias_selected = self.bias[category_id]
+            out = torch.bmm(x_flat.unsqueeze(1), weight_selected).squeeze(1) + bias_selected
+        out_shape = orig_shape[:-1] + (out.shape[-1],)
+        return out.view(out_shape)
+
+
+class CategorySpecificMLP(nn.Module):
+    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, num_categories: int = 1):
+        super().__init__()
+        self.fc1 = CategorySpecificLinear(input_dim, hidden_dim, num_categories)
+        self.fc2 = CategorySpecificLinear(hidden_dim, output_dim, num_categories)
+        self.activation = nn.ReLU(inplace=True)
+
+    def forward(self, x: torch.Tensor, category_id: torch.LongTensor):
+        out = self.activation(self.fc1(x, category_id))
+        out = self.fc2(out, category_id)
+        return out
+
+
+class MultiEmbodimentActionEncoder(nn.Module):
+    def __init__(
+        self, action_dim: int, embed_dim: int, hidden_dim: int, horizon: int, num_categories: int = 1
+    ):
+        super().__init__()
+        self.horizon = horizon
+        self.embed_dim = embed_dim
+        self.num_categories = num_categories
+
+        self.W1 = CategorySpecificLinear(action_dim, hidden_dim, num_categories)
+        self.W2 = CategorySpecificLinear(hidden_dim, hidden_dim, num_categories)
+        self.W3 = CategorySpecificLinear(hidden_dim, embed_dim, num_categories)
+
+        self.pos_encoding = SinusoidalPositionalEncoding(hidden_dim, max_len=horizon)
+        self.activation = nn.ReLU(inplace=True)
+
+    def forward(self, action_seq: torch.Tensor, category_id: torch.LongTensor):
+        batch_size, horizon, action_dim = action_seq.shape
+        if self.horizon != horizon:
+            raise ValueError(
+                f"Action sequence length must match horizon: got {horizon}, expected {self.horizon}."
+            )
+
+        x = action_seq.reshape(batch_size * horizon, action_dim)
+        if category_id.dim() == 0:
+            cat_ids = category_id.expand(horizon * batch_size)
+        else:
+            cat_ids = category_id.unsqueeze(1).expand(batch_size, horizon).reshape(batch_size * horizon)
+
+        out = self.activation(self.W1(x, cat_ids))
+        pos_enc = self.pos_encoding(horizon).to(device=out.device, dtype=out.dtype)
+        out = out.view(batch_size, horizon, -1) + pos_enc
+        out = out.view(batch_size * horizon, -1)
+        out = self.activation(self.W2(out, cat_ids))
+        out = self.W3(out, cat_ids)
+        return out.view(batch_size, horizon, self.embed_dim)
+
+
+class BasicTransformerBlock(nn.Module):
+    def __init__(self, embed_dim: int, num_heads: int, hidden_dim: int, dropout: float = 0.0):
+        super().__init__()
+        self.attn = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout, batch_first=True)
+        self.norm1 = nn.LayerNorm(embed_dim)
+        self.norm2 = nn.LayerNorm(embed_dim)
+        self.ff = nn.Sequential(nn.Linear(embed_dim, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, embed_dim))
+
+    def forward(
+        self,
+        action_tokens: torch.Tensor,
+        context_tokens: torch.Tensor,
+        time_emb: torch.Tensor,
+        context_key_padding_mask: torch.Tensor | None = None,
+    ):
+        x = self.norm1(action_tokens)
+        attn_out, _ = self.attn(x, context_tokens, context_tokens, key_padding_mask=context_key_padding_mask)
+        x = action_tokens + attn_out
+        x2 = self.norm2(x)
+        if time_emb is not None:
+            x2 = x2 + time_emb.unsqueeze(1)
+        ff_out = self.ff(x2)
+        return x + ff_out
+
+
+class FlowmatchingActionHead(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int = 896,
+        hidden_dim: int = 1024,
+        action_dim: int = 16 * 7,
+        horizon: int = 16,
+        per_action_dim: int = 7,
+        num_heads: int = 8,
+        num_layers: int = 8,
+        dropout: float = 0.0,
+        num_inference_timesteps: int = 20,
+        num_categories: int = 1,
+        state_dim: int | None = None,
+        state_hidden_dim: int | None = None,
+    ):
+        super().__init__()
+
+        logger.info("FlowmatchingActionHead num_inference_timesteps=%s", num_inference_timesteps)
+        self.embed_dim = embed_dim
+        self.horizon = horizon
+        self.per_action_dim = per_action_dim
+        self.action_dim = action_dim
+        self.num_inference_timesteps = num_inference_timesteps
+        self.num_categories = num_categories
+
+        self.time_pos_enc = SinusoidalPositionalEncoding(embed_dim, max_len=1000)
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    embed_dim=embed_dim,
+                    num_heads=num_heads,
+                    hidden_dim=embed_dim * 4,
+                    dropout=dropout,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.norm_out = nn.LayerNorm(embed_dim)
+        self.seq_pool_proj = nn.Linear(self.horizon * self.embed_dim, self.embed_dim)
+        self.mlp_head = CategorySpecificMLP(
+            input_dim=embed_dim,
+            hidden_dim=hidden_dim,
+            output_dim=action_dim,
+            num_categories=num_categories,
+        )
+
+        self.state_encoder = None
+        if state_dim is not None:
+            state_hidden = state_hidden_dim if state_hidden_dim is not None else embed_dim
+            self.state_encoder = CategorySpecificMLP(
+                input_dim=state_dim,
+                hidden_dim=state_hidden,
+                output_dim=embed_dim,
+                num_categories=num_categories,
+            )
+
+        if horizon > 1:
+            self.action_encoder = MultiEmbodimentActionEncoder(
+                action_dim=self.per_action_dim,
+                embed_dim=embed_dim,
+                hidden_dim=embed_dim,
+                horizon=horizon,
+                num_categories=num_categories,
+            )
+            self.single_action_proj = None
+        else:
+            self.action_encoder = None
+            self.single_action_proj = nn.Linear(self.per_action_dim, self.embed_dim)
+
+    def _project_actions(self, action_seq: torch.Tensor, embodiment_id: torch.LongTensor) -> torch.Tensor:
+        if self.horizon > 1 and self.action_encoder is not None:
+            return self.action_encoder(action_seq, embodiment_id)
+        if self.single_action_proj is None:
+            raise RuntimeError("single_action_proj is not initialized for horizon <= 1.")
+        return self.single_action_proj(action_seq)
+
+    def _expand_action_mask(
+        self,
+        action_mask: torch.Tensor,
+        batch_size: int,
+        per_action_dim: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        if action_mask is None:
+            raise ValueError("action_mask must be provided for flow matching inference.")
+
+        if action_mask.dim() == 2:
+            expected_last_dim = self.horizon * per_action_dim
+            if action_mask.shape == (batch_size, expected_last_dim):
+                expanded_mask = action_mask.reshape(batch_size, self.horizon, per_action_dim)
+            elif action_mask.shape == (batch_size, per_action_dim):
+                expanded_mask = action_mask.unsqueeze(1).expand(batch_size, self.horizon, per_action_dim)
+            else:
+                raise ValueError(
+                    f"Expected action_mask shape {(batch_size, expected_last_dim)} or "
+                    f"{(batch_size, per_action_dim)}, got {tuple(action_mask.shape)}"
+                )
+        elif action_mask.dim() == 3:
+            expected_shape = (batch_size, self.horizon, per_action_dim)
+            if tuple(action_mask.shape) != expected_shape:
+                raise ValueError(
+                    f"Expected action_mask shape {expected_shape}, got {tuple(action_mask.shape)}"
+                )
+            expanded_mask = action_mask
+        else:
+            raise ValueError(f"Unsupported action_mask rank: {action_mask.dim()}")
+
+        return expanded_mask.to(device=device, dtype=dtype)
+
+    def _prepare_context(
+        self,
+        fused_tokens: torch.Tensor,
+        state: torch.Tensor | None,
+        embodiment_id: torch.LongTensor | None,
+        context_mask: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, torch.LongTensor]:
+        """Normalize the VL context and embodiment ids shared by training and inference.
+
+        Returns the context tokens ``(B, S, E)``, a key_padding_mask for
+        ``nn.MultiheadAttention`` (True = ignore) or None, and the resolved embodiment ids.
+        """
+        batch_size = fused_tokens.size(0)
+        device = fused_tokens.device
+        if embodiment_id is None:
+            embodiment_id = torch.zeros(batch_size, dtype=torch.long, device=device)
+        elif self.num_categories > 1 and (
+            int(embodiment_id.min()) < 0 or int(embodiment_id.max()) >= self.num_categories
+        ):
+            raise ValueError(
+                f"embodiment ids must be in [0, num_categories={self.num_categories}), "
+                f"got range [{int(embodiment_id.min())}, {int(embodiment_id.max())}]"
+            )
+
+        context_tokens = fused_tokens
+        if context_tokens.dim() == 2:
+            # A single pooled VL token (return_cls_only): give it a sequence dim of 1.
+            context_tokens = context_tokens.unsqueeze(1)
+            context_mask = None
+        if state is not None and self.state_encoder is not None:
+            state_emb = self.state_encoder(state, embodiment_id).unsqueeze(1)
+            context_tokens = torch.cat([context_tokens, state_emb], dim=1)
+            if context_mask is not None:
+                state_valid = torch.ones(batch_size, 1, dtype=torch.bool, device=context_mask.device)
+                context_mask = torch.cat([context_mask.to(torch.bool), state_valid], dim=1)
+
+        key_padding_mask = None if context_mask is None else ~context_mask.to(torch.bool)
+        return context_tokens, key_padding_mask, embodiment_id
+
+    def forward(
+        self,
+        fused_tokens: torch.Tensor,
+        state: torch.Tensor = None,
+        actions_gt: torch.Tensor = None,
+        embodiment_id: torch.LongTensor = None,
+        action_mask: torch.Tensor = None,
+        context_mask: torch.Tensor = None,
+    ):
+        if actions_gt is None:
+            return self.get_action(
+                fused_tokens,
+                state=state,
+                embodiment_id=embodiment_id,
+                action_mask=action_mask,
+                context_mask=context_mask,
+            )
+
+        batch_size = fused_tokens.size(0)
+        device = fused_tokens.device
+        context_tokens, key_padding_mask, embodiment_id = self._prepare_context(
+            fused_tokens, state, embodiment_id, context_mask
+        )
+
+        t = (
+            torch.distributions.Beta(2, 2)
+            .sample((batch_size,))
+            .clamp(0.02, 0.98)
+            .to(device)
+            .to(dtype=self.dtype)
+        )
+        time_index = (t * 999).long().clamp_(0, 999)
+        time_emb = self.time_pos_enc(1000)[:, time_index, :].squeeze(0).to(dtype=context_tokens.dtype)
+
+        actions_gt_seq = actions_gt
+        noise = torch.rand_like(actions_gt) * 2 - 1
+        if action_mask is not None:
+            action_mask = action_mask.to(dtype=noise.dtype, device=noise.device)
+            if action_mask.shape != noise.shape:
+                raise ValueError(f"action_mask shape {action_mask.shape} != noise shape {noise.shape}")
+            actions_gt_seq = actions_gt_seq * action_mask
+            noise = noise * action_mask
+
+        if self.horizon > 1:
+            noise_seq = noise.view(batch_size, self.horizon, self.per_action_dim)
+        else:
+            noise_seq = noise if noise.dim() == 3 else noise.unsqueeze(1)
+        t_broadcast = t.view(batch_size, 1, 1)
+        action_intermediate_seq = (1 - t_broadcast) * noise_seq + t_broadcast * actions_gt_seq
+
+        action_tokens = self._project_actions(action_intermediate_seq, embodiment_id)
+        target_dtype = self.dtype
+        action_tokens = action_tokens.to(dtype=target_dtype)
+        context_tokens = context_tokens.to(dtype=target_dtype)
+        time_emb = time_emb.to(dtype=target_dtype)
+
+        x = action_tokens
+        for block in self.transformer_blocks:
+            x = block(x, context_tokens, time_emb, key_padding_mask)
+        x = self.norm_out(x)
+
+        if self.horizon > 1:
+            x_flat = x.reshape(batch_size, -1)
+            x_pooled = self.seq_pool_proj(x_flat)
+        else:
+            x_pooled = x.squeeze(1)
+
+        pred_velocity = self.mlp_head(x_pooled, embodiment_id)
+        return pred_velocity, noise
+
+    def get_action(
+        self,
+        fused_tokens: torch.Tensor,
+        state: torch.Tensor = None,
+        embodiment_id: torch.LongTensor = None,
+        action_mask: torch.Tensor = None,
+        context_mask: torch.Tensor = None,
+        inference_delay: int | None = None,
+        prev_chunk_left_over: torch.Tensor | None = None,
+        execution_horizon: int | None = None,
+        rtc_processor=None,
+    ):
+        batch_size = fused_tokens.size(0)
+        device = fused_tokens.device
+        context_tokens, key_padding_mask, embodiment_id = self._prepare_context(
+            fused_tokens, state, embodiment_id, context_mask
+        )
+
+        action_dim_total = self.action_dim
+        per_action_dim = self.per_action_dim
+
+        action = torch.rand(batch_size, action_dim_total, device=device, dtype=context_tokens.dtype) * 2 - 1
+        action_seq = action.view(batch_size, self.horizon, per_action_dim)
+        action_mask = self._expand_action_mask(
+            action_mask,
+            batch_size=batch_size,
+            per_action_dim=per_action_dim,
+            device=action_seq.device,
+            dtype=action_seq.dtype,
+        )
+        action_seq = action_seq * action_mask
+
+        target_dtype = self.dtype
+        context_tokens = context_tokens.to(dtype=target_dtype)
+
+        num_steps = int(self.num_inference_timesteps)
+        if num_steps <= 0:
+            raise ValueError(f"num_inference_timesteps must be positive, got {num_steps}")
+        dt = 1.0 / num_steps
+
+        use_rtc = rtc_processor is not None and (
+            inference_delay is not None or prev_chunk_left_over is not None
+        )
+
+        def predict_velocity(seq: torch.Tensor, step_time_emb: torch.Tensor) -> torch.Tensor:
+            """Predict the masked flow velocity (x1 - x0 convention) for one integration step."""
+            seq = seq * action_mask
+            action_tokens = self._project_actions(seq, embodiment_id).to(dtype=target_dtype)
+            x = action_tokens
+            for block in self.transformer_blocks:
+                x = block(x, context_tokens, step_time_emb, key_padding_mask)
+            x = self.norm_out(x)
+            x_pooled = self.seq_pool_proj(x.reshape(batch_size, -1)) if self.horizon > 1 else x.squeeze(1)
+            pred = self.mlp_head(x_pooled, embodiment_id)
+            return pred.view(batch_size, self.horizon, per_action_dim) * action_mask
+
+        for i in range(num_steps):
+            t = i / num_steps
+            time_index = min(int(t * 999), 999)
+            time_emb = self.time_pos_enc(1000)[:, time_index, :].to(device).squeeze(0).to(dtype=target_dtype)
+            time_emb = time_emb.unsqueeze(0).repeat(batch_size, 1)
+
+            if use_rtc:
+                # RTCProcessor assumes the pi0 flow convention: its `time` runs 1 -> 0 and the
+                # clean-action estimate is x1 = x_t - time * v. EVO1 integrates t: 0 -> 1 with
+                # velocity v = x1 - x0 (so x1 = x_t + (1 - t) * v); passing time = 1 - t and
+                # flipping the velocity sign in both directions maps one convention onto the other.
+                guided = rtc_processor.denoise_step(
+                    x_t=action_seq,
+                    prev_chunk_left_over=prev_chunk_left_over,
+                    inference_delay=inference_delay,
+                    time=1.0 - t,
+                    original_denoise_step_partial=lambda seq, emb=time_emb: -predict_velocity(seq, emb),
+                    execution_horizon=execution_horizon,
+                )
+                velocity = -guided
+            else:
+                velocity = predict_velocity(action_seq, time_emb)
+
+            action_seq = action_seq + dt * velocity
+
+        action_seq = action_seq * action_mask
+        return action_seq.reshape(batch_size, -1)
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype

src/lerobot/policies/evo1/internvl3_embedder.py

@@ -0,0 +1,369 @@
+# 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 logging
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+import torchvision.transforms.functional as tvf
+from torchvision.transforms.functional import InterpolationMode
+
+from lerobot.utils.import_utils import _transformers_available, require_package
+
+if TYPE_CHECKING or _transformers_available:
+    from transformers import AutoModel, AutoTokenizer
+else:
+    AutoModel = None
+    AutoTokenizer = None
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+IMG_CONTEXT_TOKEN = "<IMG_CONTEXT>"  # nosec B105
+IMG_START_TOKEN = "<img>"  # nosec B105
+IMG_END_TOKEN = "</img>"  # nosec B105
+
+logger = logging.getLogger(__name__)
+
+
+def _batched_resize_01(images: torch.Tensor, image_size: int) -> torch.Tensor:
+    """Resize a batch of ``[0, 1]`` images to ``(image_size, image_size)`` on-device.
+
+    Numerically mirrors InternVL3's reference PIL preprocessing
+    (``to_pil_image`` -> ``Image.resize`` -> ``to_tensor``): the float input is quantized to uint8
+    exactly as ``to_pil_image`` does, then resized with bicubic interpolation and antialiasing,
+    which matches PIL's default resampler. Matching the reference pixel-for-pixel keeps the policy
+    interchangeable with checkpoints produced by the upstream EVO1 preprocessing.
+
+    Args:
+        images: float tensor of shape ``(N, C, H, W)`` with values in ``[0, 1]``.
+
+    Returns:
+        float32 tensor of shape ``(N, C, image_size, image_size)`` with values in ``[0, 1]``.
+    """
+    # to_pil_image() quantizes float [0, 1] to uint8 (x * 255, truncated); replicate that so the
+    # bicubic resample sees the same integer pixels PIL would.
+    pixels_u8 = (images * 255.0).clamp(0, 255).to(torch.uint8)
+    resized = tvf.resize(
+        pixels_u8, [image_size, image_size], interpolation=InterpolationMode.BICUBIC, antialias=True
+    )
+    return resized.to(torch.float32) / 255.0
+
+
+def _batched_pixel_values(
+    camera_images: Sequence[torch.Tensor],
+    max_views: int,
+    image_size: int,
+    mean: torch.Tensor,
+    std: torch.Tensor,
+    dtype: torch.dtype,
+    device: torch.device | str,
+) -> torch.Tensor:
+    """Build InternVL3 ``pixel_values`` from per-camera ``[0, 1]`` image batches without leaving the device.
+
+    Each image is resized, converted to ``dtype``, and ImageNet-normalized (a single tile per
+    image), batched across the whole minibatch. Absent views (fewer cameras than ``max_views``)
+    are filled with zero images; their placeholder tokens are masked out of attention downstream
+    via ``_mask_absent_image_tokens``.
+
+    Returns:
+        ``pixel_values`` of shape ``(B * max_views, C, image_size, image_size)``, ordered row-major
+        over ``(sample, view)`` to line up with the per-view image placeholders in the prompt.
+    """
+    resized: list[torch.Tensor] = []
+    for image in camera_images:
+        resized.append(_batched_resize_01(image.to(device=device), image_size).to(dtype))
+
+    batch_size = resized[0].shape[0]
+    channels = resized[0].shape[1]
+    while len(resized) < max_views:
+        resized.append(torch.zeros(batch_size, channels, image_size, image_size, dtype=dtype, device=device))
+
+    stacked = torch.stack(resized[:max_views], dim=1)  # (B, V, C, H, W)
+    mean = mean.to(device=device, dtype=dtype).view(1, 1, -1, 1, 1)
+    std = std.to(device=device, dtype=dtype).view(1, 1, -1, 1, 1)
+    normalized = (stacked - mean) / std
+    return normalized.reshape(batch_size * max_views, channels, image_size, image_size)
+
+
+class InternVL3Embedder(nn.Module):
+    """Vision-language embedder using the native HF InternVL3 model (no trust_remote_code)."""
+
+    def __init__(
+        self,
+        model_name="OpenGVLab/InternVL3-1B-hf",
+        image_size=448,
+        device="cuda",
+        num_language_layers: int | None = 14,
+        model_dtype: str | torch.dtype = "bfloat16",
+        use_flash_attn: bool = True,
+        max_text_length: int = 1024,
+        enable_gradient_checkpointing: bool = True,
+        gradient_checkpointing_use_reentrant: bool = False,
+        hub_kwargs: dict | None = None,
+    ):
+        super().__init__()
+        self._requested_device = device
+        self.image_size = image_size
+        self.num_language_layers = num_language_layers
+        self.max_text_length = max_text_length
+        self.enable_gradient_checkpointing = bool(enable_gradient_checkpointing)
+        self.gradient_checkpointing_use_reentrant = bool(gradient_checkpointing_use_reentrant)
+        hub_kwargs = hub_kwargs or {}
+
+        require_package("transformers", extra="evo1")
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, **hub_kwargs)
+        if isinstance(model_dtype, str):
+            try:
+                model_dtype = getattr(torch, model_dtype)
+            except AttributeError as exc:
+                raise ValueError(f"Unsupported EVO1 vlm_dtype '{model_dtype}'") from exc
+        self.model_dtype = model_dtype
+
+        attn_implementation = "flash_attention_2" if (use_flash_attn and _flash_attn_available()) else "eager"
+        if use_flash_attn and attn_implementation == "eager":
+            logger.warning("flash_attn is not installed. Falling back to eager attention.")
+
+        self.model = AutoModel.from_pretrained(
+            model_name,
+            torch_dtype=model_dtype,
+            attn_implementation=attn_implementation,
+            low_cpu_mem_usage=True,
+            **hub_kwargs,
+        ).to(self._requested_device)
+
+        checkpoint_image_size = getattr(self.model.config.vision_config, "image_size", None)
+        if isinstance(checkpoint_image_size, (list, tuple)):
+            checkpoint_image_size = checkpoint_image_size[0]
+        if checkpoint_image_size is not None and int(checkpoint_image_size) != int(image_size):
+            raise ValueError(
+                f"EVO1 image_resolution ({image_size}) must match the InternVL checkpoint's native "
+                f"image size ({checkpoint_image_size}): the checkpoint's image_seq_length assumes "
+                "its native resolution, so other sizes would desync the image placeholder tokens "
+                "from the vision features."
+            )
+
+        self.num_image_token = self.model.config.image_seq_length
+
+        # Truncate language model to the requested number of layers
+        layers = self.model.language_model.layers
+        if self.num_language_layers is not None:
+            layers = layers[: self.num_language_layers]
+        self.model.language_model.layers = torch.nn.ModuleList(layers)
+
+        self._configure_memory_features()
+        self.img_context_token_id = self.tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
+
+    def _configure_memory_features(self) -> None:
+        checkpoint_kwargs = {"use_reentrant": self.gradient_checkpointing_use_reentrant}
+
+        if not self.enable_gradient_checkpointing:
+            language_model = self.model.language_model
+            if hasattr(language_model, "gradient_checkpointing_disable"):
+                language_model.gradient_checkpointing_disable()
+            vision_tower = getattr(self.model, "vision_tower", None)
+            if vision_tower is not None and hasattr(vision_tower, "encoder"):
+                vision_tower.encoder.gradient_checkpointing = False
+            return
+
+        def _enable_ckpt(module: nn.Module | None) -> bool:
+            if module is None:
+                return False
+            if hasattr(module, "gradient_checkpointing_enable"):
+                try:
+                    module.gradient_checkpointing_enable(gradient_checkpointing_kwargs=checkpoint_kwargs)
+                except TypeError:
+                    module.gradient_checkpointing_enable()
+                return True
+            if hasattr(module, "gradient_checkpointing"):
+                module.gradient_checkpointing = True
+                return True
+            return False
+
+        enabled_any = _enable_ckpt(self.model)
+
+        vision_tower = getattr(self.model, "vision_tower", None)
+        if vision_tower is not None:
+            enabled_any = _enable_ckpt(vision_tower) or enabled_any
+
+        language_model = self.model.language_model
+        enabled_any = _enable_ckpt(language_model) or enabled_any
+        if hasattr(language_model, "config"):
+            language_model.config.use_cache = False
+
+        if hasattr(self.model, "config"):
+            self.model.config.use_cache = False
+        if hasattr(self.model, "enable_input_require_grads"):
+            self.model.enable_input_require_grads()
+
+        if enabled_any:
+            logger.info("Gradient checkpointing enabled for InternVL3 embedder.")
+        else:
+            logger.warning(
+                "Requested gradient checkpointing, but model does not expose checkpointing controls."
+            )
+
+    def _build_multimodal_prompts(
+        self,
+        batch_num_tiles_list: list[list[int]],
+        text_prompts: Sequence[str],
+    ) -> list[str]:
+        prompts = []
+        for num_tiles_list, text_prompt in zip(batch_num_tiles_list, text_prompts, strict=True):
+            prompt_segments = []
+            for i, tile_count in enumerate(num_tiles_list):
+                token_count = self.num_image_token * tile_count
+                image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * token_count + IMG_END_TOKEN
+                prompt_segments.append(f"Image-{i + 1}: {image_tokens}\n")
+            prompts.append("".join(prompt_segments) + text_prompt.strip())
+        return prompts
+
+    def get_fused_image_text_embedding_batched(
+        self,
+        camera_images: Sequence[torch.Tensor],
+        image_masks: torch.Tensor,
+        text_prompts: Sequence[str],
+        return_cls_only: bool = True,
+    ):
+        """Fused VL embedding from per-camera ``[0, 1]`` image batches (no PIL, no host round-trip).
+
+        Args:
+            camera_images: list of per-camera tensors, each shaped ``(B, C, H, W)`` in ``[0, 1]``.
+            image_masks: bool tensor ``(B, max_views)`` marking present views.
+
+        Returns:
+            A ``(embeddings, valid_mask)`` tuple. With ``return_cls_only=False``, ``embeddings`` is
+            ``(B, L, H)`` and ``valid_mask`` is a ``(B, L)`` bool tensor marking tokens downstream
+            attention may attend to (padding and absent-view tokens are False). With
+            ``return_cls_only=True``, ``embeddings`` is the pooled ``(B, H)`` last-valid-token state
+            and ``valid_mask`` is None.
+        """
+        max_views = int(image_masks.shape[1])
+        batch_size = int(image_masks.shape[0])
+        mean = torch.tensor(IMAGENET_MEAN, device=self.device, dtype=self.model_dtype)
+        std = torch.tensor(IMAGENET_STD, device=self.device, dtype=self.model_dtype)
+        pixel_values = _batched_pixel_values(
+            camera_images, max_views, self.image_size, mean, std, self.model_dtype, self.device
+        )
+        # InternVL3 preprocessing uses a single tile per image (max_num=1).
+        batch_num_tiles_list = [[1] * max_views for _ in range(batch_size)]
+        return self._forward_vlm(
+            pixel_values, batch_num_tiles_list, image_masks, text_prompts, return_cls_only
+        )
+
+    def _mask_absent_image_tokens(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        image_masks: torch.Tensor,
+        batch_num_tiles_list: list[list[int]],
+    ) -> torch.Tensor:
+        """Zero attention over the image-context tokens of absent (zero-padded) views.
+
+        Fully vectorized: runs without any host<->device synchronization.
+        """
+        # A single tile per image (max_num=1), so every image occupies the same number of
+        # context tokens.
+        tiles_per_image = (
+            batch_num_tiles_list[0][0] if batch_num_tiles_list and batch_num_tiles_list[0] else 1
+        )
+        tokens_per_image = self.num_image_token * tiles_per_image
+
+        image_masks = image_masks.to(device=input_ids.device).bool()
+        img_token_mask = input_ids == self.img_context_token_id  # (B, L)
+        # keep[b, k] tells whether the k-th image-context token (ordered view0, view1, ...) survives.
+        per_token_keep = image_masks.repeat_interleave(tokens_per_image, dim=1)  # (B, V * tokens_per_image)
+        # Rank each context token by its running position among the row's context tokens.
+        ctx_index = img_token_mask.to(torch.long).cumsum(dim=1) - 1
+        ctx_index = ctx_index.clamp(min=0, max=per_token_keep.shape[1] - 1)
+        keep_here = torch.gather(per_token_keep, 1, ctx_index)  # (B, L)
+        drop = img_token_mask & ~keep_here
+        return attention_mask.masked_fill(drop, 0)
+
+    def _forward_vlm(
+        self,
+        pixel_values: torch.Tensor,
+        batch_num_tiles_list: list[list[int]],
+        image_masks: torch.Tensor,
+        text_prompts: Sequence[str],
+        return_cls_only: bool,
+    ):
+        if pixel_values.shape[0] == 0:
+            logger.warning("InternVL3 received an empty image batch after preprocessing.")
+            hidden_size = getattr(self.model.config, "hidden_size", None)
+            if hidden_size is None:
+                hidden_size = getattr(self.model.config.text_config, "hidden_size", None)
+            if hidden_size is None:
+                raise RuntimeError("Unable to infer hidden size for empty InternVL3 batch.")
+            return torch.empty(0, hidden_size, device=self.device, dtype=torch.float32), None
+
+        prompts = self._build_multimodal_prompts(batch_num_tiles_list, text_prompts)
+
+        model_inputs = self.tokenizer(
+            list(prompts),
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=self.max_text_length,
+        ).to(self.device)
+        input_ids = model_inputs["input_ids"]
+        if input_ids.shape[1] >= self.max_text_length:
+            # Truncation cuts from the right, so text is dropped before image placeholders — but a
+            # large max_views * image_seq_length budget can still eat into them. Fail loudly instead
+            # of letting the VLM crash on a placeholder/vision-feature count mismatch.
+            expected_image_tokens = self.num_image_token * sum(batch_num_tiles_list[0])
+            image_token_counts = (input_ids == self.img_context_token_id).sum(dim=1)
+            if not bool((image_token_counts == expected_image_tokens).all()):
+                raise ValueError(
+                    f"Prompt truncation at max_text_length={self.max_text_length} cut into the "
+                    f"image placeholder tokens ({expected_image_tokens} expected per sample). "
+                    "Increase max_text_length or reduce max_views."
+                )
+        attention_mask = self._mask_absent_image_tokens(
+            input_ids, model_inputs["attention_mask"], image_masks, batch_num_tiles_list
+        )
+
+        outputs = self.model(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+            return_dict=True,
+        )
+        fused_hidden = outputs.hidden_states[-1].to(torch.float32)
+        valid_mask = attention_mask.to(torch.bool)
+        if return_cls_only:
+            # Right-padded causal decoder: the last valid token is the only one that has attended
+            # to the full image + text prompt.
+            positions = torch.arange(valid_mask.shape[1], device=valid_mask.device)
+            last_valid = (valid_mask.long() * positions).argmax(dim=1)
+            batch_index = torch.arange(fused_hidden.shape[0], device=fused_hidden.device)
+            return fused_hidden[batch_index, last_valid], None
+        return fused_hidden, valid_mask
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.model.parameters()).device
+
+
+def _flash_attn_available() -> bool:
+    try:
+        import flash_attn  # noqa: F401
+    except ModuleNotFoundError:
+        return False
+    return True

src/lerobot/policies/evo1/modeling_evo1.py

@@ -0,0 +1,532 @@
+# 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 builtins
+from collections import deque
+from contextlib import nullcontext
+from pathlib import Path
+from typing import TypedDict, Unpack
+
+import torch
+from torch import Tensor
+
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.policies.pretrained import PreTrainedPolicy, T
+from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
+
+from ..rtc.modeling_rtc import RTCProcessor
+from .configuration_evo1 import Evo1Config
+from .evo1_model import Evo1Model
+
+
+class ActionSelectKwargs(TypedDict, total=False):
+    inference_delay: int | None
+    prev_chunk_left_over: Tensor | None
+    execution_horizon: int | None
+
+
+class Evo1Policy(PreTrainedPolicy):
+    config_class = Evo1Config
+    name = "evo1"
+
+    def __init__(self, config: Evo1Config, *, vlm_hub_kwargs: dict | None = None, **kwargs):
+        super().__init__(config)
+        config.validate_features()
+
+        if len(config.image_features) > config.max_views:
+            raise ValueError(
+                f"EVO1 supports at most {config.max_views} camera streams, got {len(config.image_features)}"
+            )
+
+        self.config = config
+        self.model = Evo1Model(config, vlm_hub_kwargs=vlm_hub_kwargs)
+        self.model.set_finetune_flags()
+        self._keep_frozen_embedder_eval()
+        self.init_rtc_processor()
+        self.reset()
+
+    def init_rtc_processor(self):
+        """Create the RTC processor when config.rtc_config is set.
+
+        The RTC rollout backend assigns config.rtc_config after loading the policy and re-invokes
+        this method.
+        """
+        self.rtc_processor = None
+        if self.config.rtc_config is not None:
+            self.rtc_processor = RTCProcessor(self.config.rtc_config)
+        model = getattr(self, "model", None)
+        if model is not None:
+            model.rtc_processor = self.rtc_processor
+
+    def _rtc_enabled(self) -> bool:
+        return self.config.rtc_config is not None and self.config.rtc_config.enabled
+
+    @classmethod
+    def from_pretrained(
+        cls: builtins.type[T],
+        pretrained_name_or_path: str | Path,
+        *,
+        config: PreTrainedConfig | None = None,
+        force_download: bool = False,
+        resume_download: bool | None = None,
+        proxies: dict | None = None,
+        token: str | bool | None = None,
+        cache_dir: str | Path | None = None,
+        local_files_only: bool = False,
+        revision: str | None = None,
+        strict: bool | None = None,
+        **kwargs,
+    ) -> T:
+        if strict is None:
+            strict = True
+        vlm_hub_kwargs = kwargs.pop("vlm_hub_kwargs", None)
+        if config is None:
+            config = PreTrainedConfig.from_pretrained(
+                pretrained_name_or_path=pretrained_name_or_path,
+                force_download=force_download,
+                resume_download=resume_download,
+                proxies=proxies,
+                token=token,
+                cache_dir=cache_dir,
+                local_files_only=local_files_only,
+                revision=revision,
+                **kwargs,
+            )
+        if vlm_hub_kwargs is None:
+            # Forward the hub download options to the base-VLM download as well; `revision` is not
+            # forwarded because it identifies the policy repo, not the VLM repo.
+            vlm_hub_kwargs = {
+                key: value
+                for key, value in (
+                    ("token", token),
+                    ("cache_dir", cache_dir),
+                    ("local_files_only", local_files_only),
+                    ("proxies", proxies),
+                )
+                if value not in (None, False)
+            }
+        kwargs["vlm_hub_kwargs"] = vlm_hub_kwargs
+        return super().from_pretrained(
+            pretrained_name_or_path=pretrained_name_or_path,
+            config=config,
+            force_download=force_download,
+            resume_download=resume_download,
+            proxies=proxies,
+            token=token,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+            revision=revision,
+            strict=strict,
+            **kwargs,
+        )
+
+    @property
+    def _camera_keys(self) -> list[str]:
+        return list(self.config.image_features)
+
+    @property
+    def _env_action_dim(self) -> int:
+        action_feature = self.config.action_feature
+        if action_feature is None:
+            return self.config.max_action_dim
+        return int(action_feature.shape[0])
+
+    @property
+    def _compute_dtype(self) -> torch.dtype:
+        return next(self.model.action_head.parameters()).dtype
+
+    @property
+    def _device(self) -> torch.device:
+        # The device the policy actually lives on. Derived from the parameters rather than
+        # config.device so the policy keeps working after accelerate (or a plain .to()) moves it.
+        return next(self.model.action_head.parameters()).device
+
+    @property
+    def _amp_enabled(self) -> bool:
+        return bool(self.config.use_amp) and self._device.type == "cuda"
+
+    def _maybe_autocast(self):
+        # EVO1 manages its own mixed precision: an explicit bf16 autocast that also overrides any
+        # outer autocast context (e.g. lerobot-eval's fp16 default), keeping train and eval
+        # numerics identical.
+        if self._amp_enabled:
+            return torch.autocast(device_type="cuda", dtype=torch.bfloat16)
+        return nullcontext()
+
+    def get_optim_params(self) -> list[dict]:
+        decay, no_decay = [], []
+        for name, param in self.named_parameters():
+            if not param.requires_grad:
+                continue
+            is_bias = name.endswith("bias") or ".bias" in name
+            is_norm = param.dim() == 1 or "norm" in name.lower()
+            if is_bias or is_norm:
+                no_decay.append(param)
+            else:
+                decay.append(param)
+        return [
+            {"params": decay, "weight_decay": self.config.optimizer_weight_decay},
+            {"params": no_decay, "weight_decay": 0.0},
+        ]
+
+    def reset(self):
+        self._action_queue = deque([], maxlen=self.config.n_action_steps)
+
+    def _normalize_task_batch(self, batch: dict[str, Tensor | list[str] | str]) -> list[str]:
+        prompts = batch.get(self.config.task_field)
+        if prompts is None and self.config.task_field != "task":
+            prompts = batch.get("task")
+        if prompts is None:
+            raise ValueError(f"EVO1 expects a '{self.config.task_field}' text field in the batch.")
+        if isinstance(prompts, str):
+            return [prompts]
+        if isinstance(prompts, (list, tuple)):
+            return [str(prompt) for prompt in prompts]
+        raise TypeError(f"Unsupported prompt batch type: {type(prompts)}")
+
+    def _prepare_state(self, batch: dict[str, Tensor]) -> tuple[Tensor, Tensor]:
+        if OBS_STATE not in batch:
+            raise ValueError(f"EVO1 requires '{OBS_STATE}' in the batch.")
+        state = batch[OBS_STATE]
+        if state.dim() == 1:
+            state = state.unsqueeze(0)
+        elif state.dim() == 3:
+            state = state[:, -1]
+        elif state.dim() != 2:
+            raise ValueError(f"Unsupported state tensor shape for EVO1: {tuple(state.shape)}")
+        batch_size, state_dim = state.shape
+        if state_dim > self.config.max_state_dim:
+            raise ValueError(
+                f"State dim {state_dim} exceeds configured max_state_dim {self.config.max_state_dim}"
+            )
+        explicit_mask = batch.get("state_mask")
+        if explicit_mask is not None:
+            if explicit_mask.dim() == 1:
+                explicit_mask = explicit_mask.unsqueeze(0)
+            elif explicit_mask.dim() == 3:
+                explicit_mask = explicit_mask[:, -1]
+            elif explicit_mask.dim() != 2:
+                raise ValueError(
+                    f"Unsupported state_mask tensor shape for EVO1: {tuple(explicit_mask.shape)}"
+                )
+            if explicit_mask.shape != (batch_size, state_dim):
+                raise ValueError(
+                    f"state_mask shape {tuple(explicit_mask.shape)} does not match state shape {(batch_size, state_dim)}"
+                )
+        device = self._device
+        padded = torch.zeros(
+            batch_size,
+            self.config.max_state_dim,
+            dtype=state.dtype,
+            device=device,
+        )
+        padded[:, :state_dim] = state.to(device=device)
+        mask = torch.zeros(
+            batch_size,
+            self.config.max_state_dim,
+            dtype=torch.bool,
+            device=device,
+        )
+        if explicit_mask is None:
+            mask[:, :state_dim] = True
+        else:
+            mask[:, :state_dim] = explicit_mask.to(device=device, dtype=torch.bool)
+        # Zero out masked state dims so an explicit state_mask actually affects the model input
+        # (the state encoder has no mask argument of its own).
+        padded = padded * mask.to(dtype=padded.dtype)
+        return padded.to(dtype=self._compute_dtype), mask
+
+    def _prepare_actions(self, batch: dict[str, Tensor]) -> tuple[Tensor, Tensor]:
+        if ACTION not in batch:
+            raise ValueError(f"EVO1 requires '{ACTION}' in the batch for training.")
+        action = batch[ACTION]
+        if action.dim() == 2:
+            action = action.unsqueeze(1)
+        batch_size, horizon, action_dim = action.shape
+        if horizon != self.config.chunk_size:
+            raise ValueError(
+                f"EVO1 expects chunk_size={self.config.chunk_size}, got action horizon {horizon}"
+            )
+        if action_dim > self.config.max_action_dim:
+            raise ValueError(
+                f"Action dim {action_dim} exceeds configured max_action_dim {self.config.max_action_dim}"
+            )
+        explicit_mask = batch.get("action_mask")
+        if explicit_mask is not None:
+            if explicit_mask.dim() == 2:
+                if horizon == 1:
+                    explicit_mask = explicit_mask.unsqueeze(1)
+                else:
+                    raise ValueError(
+                        f"2D action_mask is only supported when chunk_size=1, got action horizon {horizon}"
+                    )
+            elif explicit_mask.dim() != 3:
+                raise ValueError(
+                    f"Unsupported action_mask tensor shape for EVO1: {tuple(explicit_mask.shape)}"
+                )
+            if explicit_mask.shape != (batch_size, horizon, action_dim):
+                raise ValueError(
+                    "action_mask shape "
+                    f"{tuple(explicit_mask.shape)} does not match action shape {(batch_size, horizon, action_dim)}"
+                )
+        device = self._device
+        padded = torch.zeros(
+            batch_size,
+            horizon,
+            self.config.max_action_dim,
+            dtype=action.dtype,
+            device=device,
+        )
+        padded[:, :, :action_dim] = action.to(device=device)
+        mask = torch.zeros(
+            batch_size,
+            horizon,
+            self.config.max_action_dim,
+            dtype=torch.bool,
+            device=device,
+        )
+        if explicit_mask is None:
+            mask[:, :, :action_dim] = True
+        else:
+            mask[:, :, :action_dim] = explicit_mask.to(device=device, dtype=torch.bool)
+
+        # Timesteps beyond the episode end hold fabricated (repeated) actions; exclude them from
+        # the loss like the other chunked policies do.
+        action_is_pad = batch.get("action_is_pad")
+        if action_is_pad is not None:
+            if action_is_pad.shape != (batch_size, horizon):
+                raise ValueError(
+                    f"action_is_pad shape {tuple(action_is_pad.shape)} does not match "
+                    f"(batch_size, chunk_size)={(batch_size, horizon)}"
+                )
+            in_episode = ~action_is_pad.to(device=device, dtype=torch.bool)
+            mask = mask & in_episode.unsqueeze(-1)
+        return padded.to(dtype=self._compute_dtype), mask
+
+    def _prepare_inference_action_mask(self, batch_size: int) -> Tensor:
+        mask = torch.zeros(
+            batch_size,
+            self.config.max_action_dim,
+            dtype=torch.bool,
+            device=self._device,
+        )
+        mask[:, : self._env_action_dim] = True
+        return mask
+
+    def _get_embodiment_ids(self, batch: dict[str, Tensor], batch_size: int) -> Tensor:
+        embodiment_ids = batch.get("embodiment_id")
+        if embodiment_ids is None and self.config.embodiment_id_field:
+            embodiment_ids = batch.get(self.config.embodiment_id_field)
+        if embodiment_ids is None:
+            return torch.full(
+                (batch_size,),
+                self.config.default_embodiment_id,
+                dtype=torch.long,
+                device=self._device,
+            )
+        if embodiment_ids.dim() == 0:
+            embodiment_ids = embodiment_ids.unsqueeze(0)
+        elif embodiment_ids.dim() > 1:
+            embodiment_ids = embodiment_ids[:, -1]
+        return embodiment_ids.to(device=self._device, dtype=torch.long)
+
+    @property
+    def _tracks_vlm_gradients(self) -> bool:
+        return bool(
+            self.config.finetune_vlm
+            or self.config.finetune_language_model
+            or self.config.finetune_vision_model
+        )
+
+    def _keep_frozen_embedder_eval(self) -> None:
+        if self._tracks_vlm_gradients:
+            return
+        embedder = getattr(self.model, "embedder", None)
+        if embedder is not None:
+            embedder.eval()
+
+    def train(self, mode: bool = True):
+        super().train(mode)
+        self._keep_frozen_embedder_eval()
+        return self
+
+    def _collect_image_batches(self, batch: dict[str, Tensor]) -> tuple[list[Tensor], Tensor]:
+        camera_keys = self._camera_keys or sorted(key for key in batch if key.startswith(f"{OBS_IMAGES}."))
+        if not camera_keys:
+            raise ValueError("EVO1 requires at least one visual observation feature.")
+        camera_keys = list(camera_keys)[: self.config.max_views]
+
+        # Configured cameras may be absent from the batch up to the empty_cameras budget (e.g. the
+        # placeholder features added by validate_features); they become masked-out views that the
+        # embedder zero-pads. Any other absent camera is an error.
+        present_keys = [key for key in camera_keys if key in batch]
+        missing_keys = [key for key in camera_keys if key not in batch]
+        if len(missing_keys) > self.config.empty_cameras:
+            raise ValueError(
+                f"Missing camera features {missing_keys} in batch; at most "
+                f"empty_cameras={self.config.empty_cameras} may be absent."
+            )
+        if not present_keys:
+            raise ValueError("EVO1 requires at least one visual observation in the batch.")
+
+        # Keep each present camera as a batched (B, C, H, W) tensor on its current (GPU) device.
+        # Resizing/normalization and zero-padding of absent views happen batched inside the
+        # embedder, so images never leave the device here.
+        camera_images: list[Tensor] = []
+        for camera_key in present_keys:
+            image = batch[camera_key]
+            if image.dim() == 3:
+                # Promote an unbatched (C, H, W) frame so batch_size is read from a real batch dim.
+                image = image.unsqueeze(0)
+            elif image.dim() == 5:
+                image = image[:, -1]
+            elif image.dim() != 4:
+                raise ValueError(
+                    f"Unsupported image tensor shape for EVO1: key={camera_key} shape={tuple(image.shape)}"
+                )
+            camera_images.append(image)
+
+        batch_size = camera_images[0].shape[0]
+        n_present = len(camera_images)
+        image_masks = torch.zeros(
+            batch_size, self.config.max_views, dtype=torch.bool, device=camera_images[0].device
+        )
+        image_masks[:, :n_present] = True
+
+        return camera_images, image_masks
+
+    def _compute_fused_tokens(
+        self,
+        prompts: list[str],
+        image_batches: list[Tensor],
+        image_masks: Tensor,
+    ) -> tuple[Tensor, Tensor | None]:
+        track_vlm_gradients = self._tracks_vlm_gradients
+        grad_context = nullcontext() if track_vlm_gradients else torch.no_grad()
+        with grad_context:
+            fused_tokens, context_mask = self.model.get_vl_embeddings(
+                images=image_batches,
+                image_mask=image_masks,
+                prompt=prompts,
+                return_cls_only=self.config.return_cls_only,
+            )
+
+        if not track_vlm_gradients:
+            fused_tokens = fused_tokens.detach()
+        fused_tokens = fused_tokens.to(device=self._device, dtype=self._compute_dtype)
+        if context_mask is not None:
+            context_mask = context_mask.to(device=self._device)
+        return fused_tokens, context_mask
+
+    def _compute_masked_loss(
+        self,
+        pred_velocity: Tensor,
+        target_velocity: Tensor,
+        action_mask: Tensor,
+        reduction: str,
+    ) -> Tensor:
+        flat_mask = action_mask.view(action_mask.shape[0], -1).to(dtype=pred_velocity.dtype)
+        sq_error = ((pred_velocity - target_velocity) * flat_mask).pow(2)
+        active = flat_mask.sum(dim=1).clamp_min(1.0)
+        per_sample_loss = sq_error.sum(dim=1) / active
+        if reduction == "none":
+            return per_sample_loss
+        if reduction != "mean":
+            raise ValueError(f"Unsupported reduction '{reduction}'")
+        return sq_error.sum() / active.sum()
+
+    def forward(self, batch: dict[str, Tensor], reduction: str = "mean") -> tuple[Tensor, dict]:
+        prompts = self._normalize_task_batch(batch)
+        image_batches, image_masks = self._collect_image_batches(batch)
+        states, _state_mask = self._prepare_state(batch)
+        actions_gt, action_mask = self._prepare_actions(batch)
+        embodiment_ids = self._get_embodiment_ids(batch, states.shape[0])
+
+        with self._maybe_autocast():
+            fused_tokens, context_mask = self._compute_fused_tokens(prompts, image_batches, image_masks)
+            pred_velocity, noise = self.model(
+                fused_tokens,
+                state=states,
+                actions_gt=actions_gt,
+                action_mask=action_mask.to(device=self._device, dtype=self._compute_dtype),
+                embodiment_ids=embodiment_ids,
+                context_mask=context_mask,
+            )
+
+        # Compute the flow-matching regression loss in fp32, outside the autocast block.
+        pred_velocity = pred_velocity.float()
+        noise = noise.float()
+        flat_action_mask = action_mask.view(action_mask.shape[0], -1).to(dtype=torch.float32)
+        # Flow-matching velocity target. Padded (masked-out) action dims are already zero on both sides
+        # here (`actions_gt` is zero-padded in `_prepare_actions`, and `noise` is masked inside the head),
+        # and the whole difference is multiplied by `flat_action_mask`, so padded dims contribute nothing.
+        target_velocity = (actions_gt.float() - noise).view(actions_gt.shape[0], -1) * flat_action_mask
+        loss = self._compute_masked_loss(pred_velocity, target_velocity, action_mask, reduction)
+        loss_mean = loss.mean().item() if loss.ndim > 0 else loss.item()
+        return loss, {
+            "loss": loss_mean,
+            "active_action_dims": float(action_mask.sum(dim=(1, 2)).float().mean().item()),
+        }
+
+    @torch.no_grad()
+    def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]) -> Tensor:
+        inference_delay = kwargs.get("inference_delay")
+        prev_chunk_left_over = kwargs.get("prev_chunk_left_over")
+        execution_horizon = kwargs.get("execution_horizon")
+        if (inference_delay is not None or prev_chunk_left_over is not None) and not self._rtc_enabled():
+            raise RuntimeError(
+                "Received RTC arguments but RTC is not configured for this EVO1 policy: set "
+                "config.rtc_config and call init_rtc_processor() (lerobot-rollout does this for "
+                "--inference.type=rtc)."
+            )
+        self.eval()
+
+        prompts = self._normalize_task_batch(batch)
+        image_batches, image_masks = self._collect_image_batches(batch)
+        states, _state_mask = self._prepare_state(batch)
+        embodiment_ids = self._get_embodiment_ids(batch, states.shape[0])
+        action_mask = self._prepare_inference_action_mask(states.shape[0])
+        if prev_chunk_left_over is not None:
+            prev_chunk_left_over = prev_chunk_left_over.to(device=self._device)
+
+        with self._maybe_autocast():
+            fused_tokens, context_mask = self._compute_fused_tokens(prompts, image_batches, image_masks)
+            actions = self.model(
+                fused_tokens,
+                state=states,
+                action_mask=action_mask,
+                embodiment_ids=embodiment_ids,
+                context_mask=context_mask,
+                inference_delay=inference_delay,
+                prev_chunk_left_over=prev_chunk_left_over,
+                execution_horizon=execution_horizon,
+            )
+        actions = actions.view(states.shape[0], self.config.chunk_size, self.config.max_action_dim)
+        return actions.to(dtype=torch.float32)
+
+    @torch.no_grad()
+    def select_action(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
+        assert not self._rtc_enabled(), (
+            "RTC is not supported for select_action, use it with predict_action_chunk"
+        )
+        self.eval()
+        if len(self._action_queue) == 0:
+            action_chunk = self.predict_action_chunk(batch)[:, : self.config.n_action_steps]
+            self._action_queue.extend(action_chunk.transpose(0, 1))
+        # Returns one step of shape (B, max_action_dim): actions are emitted at the padded max_action_dim
+        # width and cropped to the real action dim downstream by the postprocessor (Evo1ActionProcessorStep).
+        # Callers that bypass the postprocessor receive the padded width.
+        return self._action_queue.popleft()

src/lerobot/policies/evo1/processor_evo1.py

@@ -0,0 +1,400 @@
+# 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
+
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
+from lerobot.processor import (
+    AddBatchDimensionProcessorStep,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    ObservationProcessorStep,
+    PolicyAction,
+    PolicyActionProcessorStep,
+    PolicyProcessorPipeline,
+    ProcessorStep,
+    ProcessorStepRegistry,
+    RenameObservationsProcessorStep,
+    UnnormalizerProcessorStep,
+)
+from lerobot.processor.converters import (
+    batch_to_transition,
+    create_transition,
+    policy_action_to_transition,
+    transition_to_policy_action,
+)
+from lerobot.types import EnvTransition, TransitionKey
+from lerobot.utils.constants import (
+    ACTION,
+    DONE,
+    INFO,
+    OBS_PREFIX,
+    OBS_STATE,
+    POLICY_POSTPROCESSOR_DEFAULT_NAME,
+    POLICY_PREPROCESSOR_DEFAULT_NAME,
+    REWARD,
+    TRUNCATED,
+)
+
+from .configuration_evo1 import Evo1Config
+
+
+def evo1_batch_to_transition(batch: dict[str, Any]):
+    transition = batch_to_transition(batch)
+    complementary_data = dict(transition.get("complementary_data") or {})
+    reserved = {ACTION, REWARD, DONE, TRUNCATED, INFO}
+    for key, value in batch.items():
+        if key in reserved or key.startswith(OBS_PREFIX):
+            continue
+        complementary_data.setdefault(key, value)
+    return create_transition(
+        observation=transition.get("observation"),
+        action=transition.get("action"),
+        reward=transition.get("reward", 0.0),
+        done=transition.get("done", False),
+        truncated=transition.get("truncated", False),
+        info=transition.get("info", {}),
+        complementary_data=complementary_data,
+    )
+
+
+@dataclass
+@ProcessorStepRegistry.register(name="evo1_pad_state_processor")
+class Evo1PadStateProcessorStep(ObservationProcessorStep):
+    """Pad policy observations to EVO1's fixed state width before normalization."""
+
+    max_state_dim: int = 24
+
+    def observation(self, observation: dict[str, Any]) -> dict[str, Any]:
+        if OBS_STATE not in observation:
+            return observation
+
+        state = observation[OBS_STATE]
+        state_dim = state.shape[-1]
+        if state_dim > self.max_state_dim:
+            raise ValueError(
+                f"EVO1 state has {state_dim} dims, which exceeds max_state_dim={self.max_state_dim}."
+            )
+        if state_dim < self.max_state_dim:
+            observation = observation.copy()
+            observation[OBS_STATE] = torch.nn.functional.pad(state, (0, self.max_state_dim - state_dim))
+        return observation
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        new_features = {ft: feats.copy() for ft, feats in features.items()}
+        obs_feats = new_features.setdefault(PipelineFeatureType.OBSERVATION, {})
+        if OBS_STATE in obs_feats:
+            obs_feats[OBS_STATE] = PolicyFeature(type=FeatureType.STATE, shape=(self.max_state_dim,))
+        return new_features
+
+    def get_config(self) -> dict[str, Any]:
+        return {"max_state_dim": self.max_state_dim}
+
+
+@dataclass
+@ProcessorStepRegistry.register(name="evo1_pad_action_processor")
+class Evo1PadActionProcessorStep(ProcessorStep):
+    """Pad training actions and preserve the active action dimensions with action_mask."""
+
+    max_action_dim: int = 24
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        action = transition.get(TransitionKey.ACTION)
+        if action is None:
+            return transition
+        if not isinstance(action, PolicyAction):
+            raise ValueError(f"EVO1 action should be a PolicyAction tensor, but got {type(action)}.")
+
+        action_dim = action.shape[-1]
+        if action_dim > self.max_action_dim:
+            raise ValueError(
+                f"EVO1 action has {action_dim} dims, which exceeds max_action_dim={self.max_action_dim}."
+            )
+
+        new_transition = transition.copy()
+        new_action = action
+        if action_dim < self.max_action_dim:
+            new_action = torch.nn.functional.pad(action, (0, self.max_action_dim - action_dim))
+
+        complementary_data = dict(new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {})
+        action_mask = complementary_data.get("action_mask")
+        if action_mask is None:
+            action_mask = torch.ones(action.shape, dtype=torch.bool, device=action.device)
+        else:
+            action_mask = torch.as_tensor(action_mask, dtype=torch.bool, device=action.device)
+            if action_mask.shape != action.shape:
+                raise ValueError(
+                    f"action_mask shape {tuple(action_mask.shape)} does not match action shape {tuple(action.shape)}."
+                )
+        if action_dim < self.max_action_dim:
+            action_mask = torch.nn.functional.pad(action_mask, (0, self.max_action_dim - action_dim))
+
+        complementary_data["action_mask"] = action_mask
+        new_transition[TransitionKey.ACTION] = new_action
+        new_transition[TransitionKey.COMPLEMENTARY_DATA] = complementary_data
+        return new_transition
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        new_features = {ft: feats.copy() for ft, feats in features.items()}
+        action_feats = new_features.setdefault(PipelineFeatureType.ACTION, {})
+        action_feats[ACTION] = PolicyFeature(type=FeatureType.ACTION, shape=(self.max_action_dim,))
+        return new_features
+
+    def get_config(self) -> dict[str, Any]:
+        return {"max_action_dim": self.max_action_dim}
+
+
+@dataclass
+@ProcessorStepRegistry.register(name="evo1_action_processor")
+class Evo1ActionProcessorStep(PolicyActionProcessorStep):
+    """Crop padded EVO1 actions and optionally binarize the LIBERO gripper channel."""
+
+    action_dim: int
+    binarize_gripper: bool = False
+    gripper_index: int = 6
+    gripper_threshold: float = 0.5
+    gripper_below_threshold_value: float = 1.0
+    gripper_above_threshold_value: float = -1.0
+
+    def action(self, action: PolicyAction) -> PolicyAction:
+        if action.shape[-1] < self.action_dim:
+            raise ValueError(
+                f"EVO1 action has {action.shape[-1]} dims, which is smaller than action_dim={self.action_dim}."
+            )
+
+        action = action[..., : self.action_dim]
+        if not self.binarize_gripper:
+            return action
+
+        if not 0 <= self.gripper_index < self.action_dim:
+            raise ValueError(
+                f"gripper_index={self.gripper_index} must be within action_dim={self.action_dim}."
+            )
+
+        action = action.clone()
+        below = torch.as_tensor(
+            self.gripper_below_threshold_value,
+            dtype=action.dtype,
+            device=action.device,
+        )
+        above = torch.as_tensor(
+            self.gripper_above_threshold_value,
+            dtype=action.dtype,
+            device=action.device,
+        )
+        action[..., self.gripper_index] = torch.where(
+            action[..., self.gripper_index] > self.gripper_threshold,
+            above,
+            below,
+        )
+        return action
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        new_features = {ft: feats.copy() for ft, feats in features.items()}
+        action_feats = new_features.setdefault(PipelineFeatureType.ACTION, {})
+        action_feats[ACTION] = PolicyFeature(type=FeatureType.ACTION, shape=(self.action_dim,))
+        return new_features
+
+    def get_config(self) -> dict[str, Any]:
+        return {
+            "action_dim": self.action_dim,
+            "binarize_gripper": self.binarize_gripper,
+            "gripper_index": self.gripper_index,
+            "gripper_threshold": self.gripper_threshold,
+            "gripper_below_threshold_value": self.gripper_below_threshold_value,
+            "gripper_above_threshold_value": self.gripper_above_threshold_value,
+        }
+
+
+def _evo1_action_dim(config: Evo1Config) -> int:
+    if config.postprocess_action_dim is not None:
+        return config.postprocess_action_dim
+    action_feature = config.action_feature
+    if action_feature is None:
+        return config.max_action_dim
+    return int(action_feature.shape[0])
+
+
+def _evo1_normalization_features(config: Evo1Config) -> dict[str, PolicyFeature]:
+    features = {**config.input_features, **config.output_features}
+    features[OBS_STATE] = PolicyFeature(type=FeatureType.STATE, shape=(config.max_state_dim,))
+    features[ACTION] = PolicyFeature(type=FeatureType.ACTION, shape=(config.max_action_dim,))
+    return features
+
+
+def _evo1_action_features(config: Evo1Config) -> dict[str, PolicyFeature]:
+    return {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(config.max_action_dim,))}
+
+
+_STAT_PAD_VALUES = {
+    "mean": 0.0,
+    "std": 1.0,
+    "min": -1.0,
+    "max": 1.0,
+    "q01": -1.0,
+    "q99": 1.0,
+    "q10": -1.0,
+    "q90": 1.0,
+}
+
+
+def _pad_stat_value(value: Any, target_dim: int, stat_name: str) -> torch.Tensor:
+    tensor = torch.as_tensor(value)
+    if not tensor.is_floating_point():
+        tensor = tensor.to(dtype=torch.float32)
+    if tensor.ndim == 0 or tensor.shape[-1] >= target_dim:
+        return tensor
+
+    pad_shape = (*tensor.shape[:-1], target_dim - tensor.shape[-1])
+    pad_value = _STAT_PAD_VALUES.get(stat_name, 0.0)
+    padding = torch.full(pad_shape, pad_value, dtype=tensor.dtype, device=tensor.device)
+    return torch.cat([tensor, padding], dim=-1)
+
+
+def _pad_feature_stats(
+    stats: dict[str, dict[str, Any]],
+    feature_key: str,
+    target_dim: int,
+) -> None:
+    if feature_key not in stats:
+        return
+    stats[feature_key] = {
+        stat_name: _pad_stat_value(stat_value, target_dim, stat_name)
+        for stat_name, stat_value in stats[feature_key].items()
+    }
+
+
+def _pad_evo1_stats(
+    config: Evo1Config,
+    stats: dict[str, dict[str, Any]] | None,
+) -> dict[str, dict[str, Any]] | None:
+    if stats is None:
+        return None
+
+    padded_stats = deepcopy(stats)
+    # Added dimensions represent zero-padding inside EVO1. These neutral stats keep
+    # padded observations at normalized zero and only provide shape compatibility.
+    _pad_feature_stats(padded_stats, OBS_STATE, config.max_state_dim)
+    _pad_feature_stats(padded_stats, ACTION, config.max_action_dim)
+    return padded_stats
+
+
+def reconcile_evo1_processors(
+    config: Evo1Config,
+    preprocessor: PolicyProcessorPipeline,
+    postprocessor: PolicyProcessorPipeline,
+) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
+    """Reconcile checkpoint-loaded pipelines with the current EVO1 config.
+
+    Two things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
+    (converters are plain functions and are never serialized), and eval-time CLI overrides of the
+    action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`). This
+    restores the converter and rebuilds the action step from the current config so those overrides
+    take effect.
+    """
+    # Pipelines reloaded from a checkpoint come back with the default batch converter, which drops
+    # non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1.
+    preprocessor.to_transition = evo1_batch_to_transition
+
+    action_step = Evo1ActionProcessorStep(
+        action_dim=_evo1_action_dim(config),
+        binarize_gripper=config.binarize_gripper,
+        gripper_index=config.gripper_index,
+        gripper_threshold=config.gripper_threshold,
+        gripper_below_threshold_value=config.gripper_below_threshold_value,
+        gripper_above_threshold_value=config.gripper_above_threshold_value,
+    )
+    steps = list(postprocessor.steps)
+    action_step_idx = next(
+        (idx for idx, step in enumerate(steps) if isinstance(step, Evo1ActionProcessorStep)), None
+    )
+    if action_step_idx is None:
+        insert_idx = next(
+            (idx + 1 for idx, step in enumerate(steps) if isinstance(step, UnnormalizerProcessorStep)),
+            0,
+        )
+        steps.insert(insert_idx, action_step)
+    else:
+        steps[action_step_idx] = action_step
+    postprocessor.steps = steps
+
+    return preprocessor, postprocessor
+
+
+def make_evo1_pre_post_processors(
+    config: Evo1Config,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    normalization_features = _evo1_normalization_features(config)
+    action_features = _evo1_action_features(config)
+    normalization_stats = _pad_evo1_stats(config, dataset_stats)
+
+    input_steps = [
+        RenameObservationsProcessorStep(rename_map={}),
+        AddBatchDimensionProcessorStep(),
+        Evo1PadStateProcessorStep(max_state_dim=config.max_state_dim),
+        Evo1PadActionProcessorStep(max_action_dim=config.max_action_dim),
+        NormalizerProcessorStep(
+            features=normalization_features,
+            norm_map=config.normalization_mapping,
+            stats=normalization_stats,
+        ),
+        DeviceProcessorStep(device=config.device),
+    ]
+    output_steps = [
+        UnnormalizerProcessorStep(
+            features=action_features,
+            norm_map=config.normalization_mapping,
+            stats=normalization_stats,
+        ),
+        Evo1ActionProcessorStep(
+            action_dim=_evo1_action_dim(config),
+            binarize_gripper=config.binarize_gripper,
+            gripper_index=config.gripper_index,
+            gripper_threshold=config.gripper_threshold,
+            gripper_below_threshold_value=config.gripper_below_threshold_value,
+            gripper_above_threshold_value=config.gripper_above_threshold_value,
+        ),
+        # float32 so downstream numpy conversion works even when the policy computes in bf16.
+        DeviceProcessorStep(device="cpu", float_dtype="float32"),
+    ]
+
+    return (
+        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
+            steps=input_steps,
+            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
+            to_transition=evo1_batch_to_transition,
+        ),
+        PolicyProcessorPipeline[PolicyAction, PolicyAction](
+            steps=output_steps,
+            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
+            to_transition=policy_action_to_transition,
+            to_output=transition_to_policy_action,
+        ),
+    )

src/lerobot/policies/factory.py

@@ -47,6 +47,7 @@ from lerobot.utils.feature_utils import dataset_to_policy_features
 from .act.configuration_act import ACTConfig
 from .diffusion.configuration_diffusion import DiffusionConfig
 from .eo1.configuration_eo1 import EO1Config
+from .evo1.configuration_evo1 import Evo1Config
 from .fastwam.configuration_fastwam import FastWAMConfig
 from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig
 from .groot.configuration_groot import GrootConfig
@@ -93,7 +94,7 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
     Args:
         name: The name of the policy. Supported names are "tdmpc", "diffusion", "act",
             "multi_task_dit", "vqbet", "pi0", "pi05", "gaussian_actor", "smolvla", "wall_x",
-            "molmoact2".
+            "molmoact2", "eo1", "evo1".
     Returns:
         The policy class corresponding to the given name.
 
@@ -172,6 +173,10 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
         from .fastwam.modeling_fastwam import FastWAMPolicy
 
         return FastWAMPolicy
+    elif name == "evo1":
+        from .evo1.modeling_evo1 import Evo1Policy
+
+        return Evo1Policy
     else:
         try:
             return _get_policy_cls_from_policy_name(name=name)
@@ -189,7 +194,7 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
     Args:
         policy_type: The type of the policy. Supported types include "tdmpc",
                      "multi_task_dit", "diffusion", "act", "vqbet", "pi0", "pi05", "gaussian_actor",
-                     "smolvla", "wall_x", "molmoact2".
+                     "smolvla", "wall_x", "molmoact2", "eo1", "evo1".
         **kwargs: Keyword arguments to be passed to the configuration class constructor.
 
     Returns:
@@ -232,6 +237,8 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
         return LingBotVAConfig(**kwargs)
     elif policy_type == "fastwam":
         return FastWAMConfig(**kwargs)
+    elif policy_type == "evo1":
+        return Evo1Config(**kwargs)
     else:
         try:
             config_cls = PreTrainedConfig.get_choice_class(policy_type)
@@ -334,6 +341,14 @@ def make_pre_post_processors(
             revision=pretrained_revision,
         )
         _reconnect_relative_absolute_steps(preprocessor, postprocessor)
+        if isinstance(policy_cfg, Evo1Config):
+            from .evo1.processor_evo1 import reconcile_evo1_processors
+
+            preprocessor, postprocessor = reconcile_evo1_processors(
+                policy_cfg,
+                preprocessor,
+                postprocessor,
+            )
         return preprocessor, postprocessor
 
     # Create a new processor based on policy type
@@ -445,6 +460,13 @@ def make_pre_post_processors(
             config=policy_cfg,
             dataset_stats=kwargs.get("dataset_stats"),
         )
+    elif isinstance(policy_cfg, Evo1Config):
+        from .evo1.processor_evo1 import make_evo1_pre_post_processors
+
+        processors = make_evo1_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
 
     elif isinstance(policy_cfg, MolmoAct2Config):
         from .molmoact2.processor_molmoact2 import make_molmoact2_pre_post_processors