refactor(policies): multiple improvements

docs/source/evo1.mdx

@@ -34,7 +34,7 @@ The broader EVO1 project may include additional training scripts and dataset too
 
 3. Install a `flash-attn` wheel only if it is compatible with your Python, PyTorch, CUDA, and GPU stack. EVO1 falls back to standard attention when `flash_attn` is not available, but reproducing the official LIBERO checkpoint conversion result below requires the same FlashAttention path used by the original EVO1 checkpoint.
 
-EVO1 uses InternVL3 through the Hugging Face `transformers` remote-code path, so the first run may download the configured VLM checkpoint unless `policy.vlm_model_name` points to a local model directory.
+EVO1 uses the native Hugging Face `transformers` InternVL implementation (no `trust_remote_code`), so `policy.vlm_model_name` must point to a natively converted checkpoint such as `OpenGVLab/InternVL3-1B-hf` (note the `-hf` suffix; the original `OpenGVLab/InternVL3-1B` repo requires remote code and cannot be loaded). The first run may download the configured VLM checkpoint unless `policy.vlm_model_name` points to a local model directory.
 
 ## Data Requirements
 
@@ -58,7 +58,7 @@ policy.type=evo1
 By default, a new EVO1 policy initializes its VLM from:
 
 ```python
-policy.vlm_model_name=OpenGVLab/InternVL3-1B
+policy.vlm_model_name=OpenGVLab/InternVL3-1B-hf
 ```
 
 Once a LeRobot-format EVO1 checkpoint is available, load it with:
@@ -84,7 +84,7 @@ lerobot-train \
   --dataset.repo_id=your_org/your_dataset \
   --policy.type=evo1 \
   --policy.training_stage=stage1 \
-  --policy.vlm_model_name=OpenGVLab/InternVL3-1B \
+  --policy.vlm_model_name=OpenGVLab/InternVL3-1B-hf \
   --policy.device=cuda \
   --policy.chunk_size=50 \
   --policy.n_action_steps=50 \
@@ -105,7 +105,7 @@ lerobot-train \
   --dataset.repo_id=your_org/your_dataset \
   --policy.path=./outputs/evo1_stage1/checkpoints/005000/pretrained_model \
   --policy.training_stage=stage2 \
-  --policy.vlm_model_name=OpenGVLab/InternVL3-1B \
+  --policy.vlm_model_name=OpenGVLab/InternVL3-1B-hf \
   --policy.device=cuda \
   --policy.chunk_size=50 \
   --policy.n_action_steps=50 \
@@ -125,23 +125,23 @@ every finetuning flag.
 
 ### Key Training Parameters
 
-| Parameter                                     | Default                  | Description                                                       |
-| --------------------------------------------- | ------------------------ | ----------------------------------------------------------------- |
-| `policy.vlm_model_name`                       | `OpenGVLab/InternVL3-1B` | InternVL3 checkpoint or local model directory                     |
-| `policy.training_stage`                       | `stage1`                 | `stage1` trains the action head; `stage2` finetunes VLM branches  |
-| `policy.apply_training_stage_defaults`        | `true`                   | Reapplies stage finetuning defaults after loading a checkpoint    |
-| `policy.vlm_num_layers`                       | `14`                     | Number of InternVL3 language layers kept for the policy           |
-| `policy.vlm_dtype`                            | `bfloat16`               | Requested VLM dtype                                               |
-| `policy.use_flash_attn`                       | `true`                   | Requests FlashAttention when installed; otherwise falls back      |
-| `policy.enable_gradient_checkpointing`        | `true`                   | Enables checkpointing on supported InternVL3 modules              |
-| `policy.gradient_checkpointing_use_reentrant` | `false`                  | Reentrant setting passed to gradient checkpointing when supported |
-| `policy.chunk_size`                           | `50`                     | Number of future actions predicted per chunk                      |
-| `policy.n_action_steps`                       | `50`                     | Number of actions consumed from a sampled chunk                   |
-| `policy.max_state_dim`                        | `24`                     | State padding dimension                                           |
-| `policy.max_action_dim`                       | `24`                     | Action padding dimension                                          |
-| `policy.postprocess_action_dim`               | `null`                   | Optional action dimension returned after EVO1 postprocessing      |
-| `policy.binarize_gripper`                     | `false`                  | Binarizes the postprocessed gripper channel for LIBERO-style eval |
-| `policy.task_field`                           | `task`                   | Batch field used as the language prompt                           |
+| Parameter                                     | Default                     | Description                                                       |
+| --------------------------------------------- | --------------------------- | ----------------------------------------------------------------- |
+| `policy.vlm_model_name`                       | `OpenGVLab/InternVL3-1B-hf` | Natively converted InternVL3 checkpoint or local model directory  |
+| `policy.training_stage`                       | `stage1`                    | `stage1` trains the action head; `stage2` finetunes VLM branches  |
+| `policy.apply_training_stage_defaults`        | `true`                      | Reapplies stage finetuning defaults after loading a checkpoint    |
+| `policy.vlm_num_layers`                       | `14`                        | Number of InternVL3 language layers kept for the policy           |
+| `policy.vlm_dtype`                            | `bfloat16`                  | Requested VLM dtype                                               |
+| `policy.use_flash_attn`                       | `true`                      | Requests FlashAttention when installed; otherwise falls back      |
+| `policy.enable_gradient_checkpointing`        | `true`                      | Enables checkpointing on supported InternVL3 modules              |
+| `policy.gradient_checkpointing_use_reentrant` | `false`                     | Reentrant setting passed to gradient checkpointing when supported |
+| `policy.chunk_size`                           | `50`                        | Number of future actions predicted per chunk                      |
+| `policy.n_action_steps`                       | `50`                        | Number of actions consumed from a sampled chunk                   |
+| `policy.max_state_dim`                        | `24`                        | State padding dimension                                           |
+| `policy.max_action_dim`                       | `24`                        | Action padding dimension                                          |
+| `policy.postprocess_action_dim`               | `null`                      | Optional action dimension returned after EVO1 postprocessing      |
+| `policy.binarize_gripper`                     | `false`                     | Binarizes the postprocessed gripper channel for LIBERO-style eval |
+| `policy.task_field`                           | `task`                      | Batch field used as the language prompt                           |
 
 ## Results
 
@@ -151,6 +151,11 @@ The checkpoint [javadcc/evo1-libero-lerobot](https://huggingface.co/javadcc/evo1
 is the LeRobot-format conversion of the official EVO1 LIBERO checkpoint. The conversion was checked against
 the official EVO1 checkpoint with the same LIBERO Object initial states and action postprocessing.
 
+> [!NOTE]
+> This checkpoint is currently hosted in a community namespace and the upstream-to-LeRobot weight
+> conversion script is not part of this integration; a `lerobot`-hosted copy with a pinned revision
+> and the conversion tooling are planned follow-ups.
+
 | Checkpoint                   | Suite           | Episodes         | Success Rate |
 | ---------------------------- | --------------- | ---------------- | ------------ |
 | Official EVO1 checkpoint     | `libero_object` | 10, one per task | 100%         |
@@ -171,7 +176,7 @@ FlashAttention, and set the LIBERO action postprocessing flags:
 ```bash
 lerobot-eval \
   --policy.path=javadcc/evo1-libero-lerobot \
-  --policy.vlm_model_name=OpenGVLab/InternVL3-1B \
+  --policy.vlm_model_name=OpenGVLab/InternVL3-1B-hf \
   --policy.device=cuda \
   --policy.use_flash_attn=true \
   --policy.n_action_steps=14 \
@@ -189,7 +194,7 @@ lerobot-eval \
 ## References
 
 - [EVO1 repository](https://github.com/MINT-SJTU/Evo-1)
-- [InternVL3-1B](https://huggingface.co/OpenGVLab/InternVL3-1B)
+- [InternVL3-1B-hf](https://huggingface.co/OpenGVLab/InternVL3-1B-hf)
 
 ## License
 

docs/source/policy_evo1_README.md

@@ -12,7 +12,7 @@ The upstream EVO1 project is available at
 @misc{evo1,
   title = {EVO1},
   author = {{MINT-SJTU}},
-  year = {2026},
+  year = {2025},
   howpublished = {\url{https://github.com/MINT-SJTU/Evo-1}},
 }
 ```

src/lerobot/optim/schedulers.py

@@ -83,6 +83,28 @@ class VQBeTSchedulerConfig(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/evo1/__init__.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from .configuration_evo1 import Evo1Config
-from .modeling_evo1 import EVO1Policy
+from .modeling_evo1 import Evo1Policy
 from .processor_evo1 import make_evo1_pre_post_processors
 
-__all__ = ["Evo1Config", "EVO1Policy", "make_evo1_pre_post_processors"]
+__all__ = ["Evo1Config", "Evo1Policy", "make_evo1_pre_post_processors"]

src/lerobot/policies/evo1/configuration_evo1.py

@@ -15,42 +15,24 @@
 from __future__ import annotations
 
 import logging
-import math
 from dataclasses import dataclass, field
 
-from torch.optim import Optimizer
-from torch.optim.lr_scheduler import LambdaLR
-
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 from lerobot.optim.optimizers import AdamWConfig
-from lerobot.optim.schedulers import LRSchedulerConfig
+from lerobot.optim.schedulers import CosineAnnealingWithWarmupSchedulerConfig
 from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
 
 logger = logging.getLogger(__name__)
 
 
-@LRSchedulerConfig.register_subclass("evo1_exact")
-@dataclass
-class Evo1SchedulerConfig(LRSchedulerConfig):
-    num_warmup_steps: int
-
-    def build(self, optimizer: Optimizer, num_training_steps: int) -> LambdaLR:
-        def lr_lambda(current_step: int) -> float:
-            if current_step < self.num_warmup_steps:
-                return current_step / max(1, self.num_warmup_steps)
-            progress = (current_step - self.num_warmup_steps) / max(
-                1, num_training_steps - self.num_warmup_steps
-            )
-            return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))
-
-        return LambdaLR(optimizer, lr_lambda, -1)
-
-
 @PreTrainedConfig.register_subclass("evo1")
 @dataclass
 class Evo1Config(PreTrainedConfig):
     training_stage: str = "stage1"
+    # 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
@@ -93,6 +75,8 @@ class Evo1Config(PreTrainedConfig):
     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
@@ -116,6 +100,8 @@ class Evo1Config(PreTrainedConfig):
     optimizer_grad_clip_norm: float = 1.0
 
     scheduler_warmup_steps: int = 300
+    # Deprecated, has no effect. Kept only so configs serialized by earlier EVO1 checkpoints
+    # (which stored this field) can still be parsed; draccus rejects unknown fields.
     drop_last: bool = True
 
     def __post_init__(self):
@@ -166,12 +152,12 @@ class Evo1Config(PreTrainedConfig):
                 flag is not None for flag in (self.finetune_language_model, self.finetune_vision_model)
             )
             if not has_explicit_branch_flags:
-                if self.finetune_vlm is None:
-                    self.finetune_vlm = True
-                if self.finetune_language_model is None:
-                    self.finetune_language_model = True
-                if self.finetune_vision_model is None:
-                    self.finetune_vision_model = True
+                # 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:
@@ -204,6 +190,11 @@ class Evo1Config(PreTrainedConfig):
                 "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:
@@ -241,7 +232,7 @@ class Evo1Config(PreTrainedConfig):
         )
 
     def get_scheduler_preset(self):
-        return Evo1SchedulerConfig(
+        return CosineAnnealingWithWarmupSchedulerConfig(
             num_warmup_steps=self.scheduler_warmup_steps,
         )
 

src/lerobot/policies/evo1/evo1_model.py

@@ -22,8 +22,8 @@ from .flow_matching import FlowmatchingActionHead
 from .internvl3_embedder import InternVL3Embedder
 
 
-class EVO1(nn.Module):
-    def __init__(self, config: Evo1Config):
+class Evo1Model(nn.Module):
+    def __init__(self, config: Evo1Config, vlm_hub_kwargs: dict | None = None):
         super().__init__()
         self.config = config
         self._device = config.device
@@ -46,6 +46,7 @@ class EVO1(nn.Module):
             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()
@@ -79,12 +80,16 @@ class EVO1(nn.Module):
         image_mask: torch.Tensor,
         prompt: str | list[str] | None = None,
         return_cls_only: bool | None = None,
-    ) -> torch.Tensor:
+    ) -> 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
@@ -117,19 +122,6 @@ class EVO1(nn.Module):
             return_cls_only=return_cls_only,
         )
 
-    def prepare_state(self, state_input: list | torch.Tensor) -> torch.Tensor:
-        if isinstance(state_input, list):
-            state_tensor = torch.tensor(state_input)
-        elif isinstance(state_input, torch.Tensor):
-            state_tensor = state_input
-        else:
-            raise TypeError(f"Unsupported state input type: {type(state_input)}")
-
-        if state_tensor.ndim == 1:
-            state_tensor = state_tensor.unsqueeze(0)
-
-        return state_tensor.to(self._device)
-
     def predict_action(
         self,
         fused_tokens: torch.Tensor,
@@ -137,6 +129,7 @@ class EVO1(nn.Module):
         actions_gt: torch.Tensor | None = None,
         action_mask: torch.Tensor | None = None,
         embodiment_ids: torch.Tensor | None = None,
+        context_mask: torch.Tensor | None = None,
     ):
         if actions_gt is None:
             return self.action_head.get_action(
@@ -144,6 +137,7 @@ class EVO1(nn.Module):
                 state=state,
                 action_mask=action_mask,
                 embodiment_id=embodiment_ids,
+                context_mask=context_mask,
             )
         return self.action_head(
             fused_tokens,
@@ -151,6 +145,7 @@ class EVO1(nn.Module):
             actions_gt=actions_gt,
             action_mask=action_mask,
             embodiment_id=embodiment_ids,
+            context_mask=context_mask,
         )
 
     def forward(
@@ -160,32 +155,34 @@ class EVO1(nn.Module):
         actions_gt: torch.Tensor | None = None,
         action_mask: torch.Tensor | None = None,
         embodiment_ids: torch.Tensor | None = None,
+        context_mask: torch.Tensor | None = None,
     ):
-        return self.predict_action(fused_tokens, state, actions_gt, action_mask, embodiment_ids)
+        return self.predict_action(fused_tokens, state, actions_gt, action_mask, embodiment_ids, context_mask)
 
     def _set_module_trainable(self, module: nn.Module, trainable: bool):
         for param in module.parameters():
             param.requires_grad = trainable
 
-    def set_finetune_flags(self):
-        finetune_vlm = bool(self.config.finetune_vlm)
-        finetune_language_model = bool(self.config.finetune_language_model)
-        finetune_vision_model = bool(self.config.finetune_vision_model)
-        has_explicit_branch_flags = any(
-            flag is not None
-            for flag in (self.config.finetune_language_model, self.config.finetune_vision_model)
-        )
+    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
 
-        if has_explicit_branch_flags:
-            self._set_module_trainable(self.embedder, False)
-            if hasattr(self.embedder.model, "language_model"):
-                self._set_module_trainable(self.embedder.model.language_model, finetune_language_model)
-            if hasattr(self.embedder.model, "vision_model"):
-                self._set_module_trainable(self.embedder.model.vision_model, finetune_vision_model)
-            if hasattr(self.embedder.model, "mlp1"):
-                self._set_module_trainable(self.embedder.model.mlp1, finetune_vision_model)
-        elif not finetune_vlm:
-            self._set_module_trainable(self.embedder, False)
+    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

@@ -62,7 +62,10 @@ class CategorySpecificLinear(nn.Module):
         else:
             self.weight = nn.Parameter(torch.empty(num_categories, in_dim, out_dim))
             self.bias = nn.Parameter(torch.zeros(num_categories, out_dim))
-            nn.init.xavier_uniform_(self.weight)
+            # 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:
@@ -150,9 +153,15 @@ class BasicTransformerBlock(nn.Module):
         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):
+    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)
+        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:
@@ -185,6 +194,7 @@ class FlowmatchingActionHead(nn.Module):
         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(
@@ -271,29 +281,68 @@ class FlowmatchingActionHead(nn.Module):
 
         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,
-        state_mask: torch.Tensor = 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
+                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
-        if embodiment_id is None:
-            embodiment_id = torch.zeros(batch_size, dtype=torch.long, device=device)
-
-        context_tokens = fused_tokens
-        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)
+        context_tokens, key_padding_mask, embodiment_id = self._prepare_context(
+            fused_tokens, state, embodiment_id, context_mask
+        )
 
         t = (
             torch.distributions.Beta(2, 2)
@@ -329,7 +378,7 @@ class FlowmatchingActionHead(nn.Module):
 
         x = action_tokens
         for block in self.transformer_blocks:
-            x = block(x, context_tokens, time_emb)
+            x = block(x, context_tokens, time_emb, key_padding_mask)
         x = self.norm_out(x)
 
         if self.horizon > 1:
@@ -347,16 +396,13 @@ class FlowmatchingActionHead(nn.Module):
         state: torch.Tensor = None,
         embodiment_id: torch.LongTensor = None,
         action_mask: torch.Tensor = None,
+        context_mask: torch.Tensor = None,
     ):
         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)
-
-        context_tokens = fused_tokens
-        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)
+        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
@@ -398,7 +444,7 @@ class FlowmatchingActionHead(nn.Module):
 
             x = action_tokens
             for block in self.transformer_blocks:
-                x = block(x, context_tokens, time_emb)
+                x = block(x, context_tokens, time_emb, key_padding_mask)
             x = self.norm_out(x)
 
             if self.horizon > 1:

src/lerobot/policies/evo1/internvl3_embedder.py

@@ -114,6 +114,7 @@ class InternVL3Embedder(nn.Module):
         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
@@ -122,15 +123,17 @@ class InternVL3Embedder(nn.Module):
         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)
+        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":
@@ -141,8 +144,20 @@ class InternVL3Embedder(nn.Module):
             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
@@ -230,13 +245,20 @@ class InternVL3Embedder(nn.Module):
         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=torch.bfloat16)
-        std = torch.tensor(IMAGENET_STD, device=self.device, dtype=torch.bfloat16)
+        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, torch.bfloat16, self.device
+            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)]
@@ -289,18 +311,30 @@ class InternVL3Embedder(nn.Module):
                 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)
+            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="max_length",
+            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
         )
@@ -313,7 +347,15 @@ class InternVL3Embedder(nn.Module):
             return_dict=True,
         )
         fused_hidden = outputs.hidden_states[-1].to(torch.float32)
-        return fused_hidden[:, 0, :] if return_cls_only else fused_hidden
+        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:

src/lerobot/policies/evo1/modeling_evo1.py

@@ -27,14 +27,14 @@ from lerobot.policies.pretrained import PreTrainedPolicy, T
 from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
 
 from .configuration_evo1 import Evo1Config
-from .evo1_model import EVO1
+from .evo1_model import Evo1Model
 
 
-class EVO1Policy(PreTrainedPolicy):
+class Evo1Policy(PreTrainedPolicy):
     config_class = Evo1Config
     name = "evo1"
 
-    def __init__(self, config: Evo1Config, **kwargs):
+    def __init__(self, config: Evo1Config, *, vlm_hub_kwargs: dict | None = None, **kwargs):
         super().__init__(config)
         config.validate_features()
 
@@ -44,7 +44,7 @@ class EVO1Policy(PreTrainedPolicy):
             )
 
         self.config = config
-        self.model = EVO1(config)
+        self.model = Evo1Model(config, vlm_hub_kwargs=vlm_hub_kwargs)
         self.model.set_finetune_flags()
         self._keep_frozen_embedder_eval()
         self.reset()
@@ -67,6 +67,33 @@ class EVO1Policy(PreTrainedPolicy):
     ) -> 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,
@@ -97,16 +124,22 @@ class EVO1Policy(PreTrainedPolicy):
         return next(self.model.action_head.parameters()).dtype
 
     @property
-    def _training_compute_dtype(self) -> torch.dtype:
-        if str(self.config.device).startswith("cuda"):
-            return torch.bfloat16
-        return self._compute_dtype
+    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 _inference_compute_dtype(self) -> torch.dtype:
-        if str(self.config.device).startswith("cuda") and self.config.use_amp:
-            return torch.bfloat16
-        return self._compute_dtype
+    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 = [], []
@@ -168,23 +201,27 @@ class EVO1Policy(PreTrainedPolicy):
                 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=self.config.device,
+            device=device,
         )
-        padded[:, :state_dim] = state.to(device=self.config.device)
+        padded[:, :state_dim] = state.to(device=device)
         mask = torch.zeros(
             batch_size,
             self.config.max_state_dim,
             dtype=torch.bool,
-            device=self.config.device,
+            device=device,
         )
         if explicit_mask is None:
             mask[:, :state_dim] = True
         else:
-            mask[:, :state_dim] = explicit_mask.to(device=self.config.device, dtype=torch.bool)
+            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]:
@@ -220,25 +257,38 @@ class EVO1Policy(PreTrainedPolicy):
                     "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=self.config.device,
+            device=device,
         )
-        padded[:, :, :action_dim] = action.to(device=self.config.device)
+        padded[:, :, :action_dim] = action.to(device=device)
         mask = torch.zeros(
             batch_size,
             horizon,
             self.config.max_action_dim,
             dtype=torch.bool,
-            device=self.config.device,
+            device=device,
         )
         if explicit_mask is None:
             mask[:, :, :action_dim] = True
         else:
-            mask[:, :, :action_dim] = explicit_mask.to(device=self.config.device, dtype=torch.bool)
+            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:
@@ -246,7 +296,7 @@ class EVO1Policy(PreTrainedPolicy):
             batch_size,
             self.config.max_action_dim,
             dtype=torch.bool,
-            device=self.config.device,
+            device=self._device,
         )
         mask[:, : self._env_action_dim] = True
         return mask
@@ -260,13 +310,13 @@ class EVO1Policy(PreTrainedPolicy):
                 (batch_size,),
                 self.config.default_embodiment_id,
                 dtype=torch.long,
-                device=self.config.device,
+                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.config.device, dtype=torch.long)
+        return embodiment_ids.to(device=self._device, dtype=torch.long)
 
     @property
     def _tracks_vlm_gradients(self) -> bool:
@@ -294,11 +344,24 @@ class EVO1Policy(PreTrainedPolicy):
             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 (no per-sample .cpu() round-trip).
         camera_images: list[Tensor] = []
-        for camera_key in camera_keys:
+        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.
@@ -323,13 +386,13 @@ class EVO1Policy(PreTrainedPolicy):
     def _compute_fused_tokens(
         self,
         prompts: list[str],
-        image_batches: list[list[Tensor]],
+        image_batches: list[Tensor],
         image_masks: Tensor,
-    ) -> 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 = self.model.get_vl_embeddings(
+            fused_tokens, context_mask = self.model.get_vl_embeddings(
                 images=image_batches,
                 image_mask=image_masks,
                 prompt=prompts,
@@ -338,7 +401,10 @@ class EVO1Policy(PreTrainedPolicy):
 
         if not track_vlm_gradients:
             fused_tokens = fused_tokens.detach()
-        return fused_tokens.to(device=self.config.device, dtype=self._compute_dtype)
+        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,
@@ -362,24 +428,27 @@ class EVO1Policy(PreTrainedPolicy):
         image_batches, image_masks = self._collect_image_batches(batch)
         states, _state_mask = self._prepare_state(batch)
         actions_gt, action_mask = self._prepare_actions(batch)
-        fused_tokens = self._compute_fused_tokens(prompts, image_batches, image_masks)
-        states = states.to(dtype=self._training_compute_dtype)
-        actions_gt = actions_gt.to(dtype=self._training_compute_dtype)
-        fused_tokens = fused_tokens.to(dtype=self._training_compute_dtype)
         embodiment_ids = self._get_embodiment_ids(batch, states.shape[0])
 
-        pred_velocity, noise = self.model(
-            fused_tokens,
-            state=states,
-            actions_gt=actions_gt,
-            action_mask=action_mask.to(device=self.config.device, dtype=self._compute_dtype),
-            embodiment_ids=embodiment_ids,
-        )
-        flat_action_mask = action_mask.view(action_mask.shape[0], -1).to(dtype=actions_gt.dtype)
+        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 - noise).view(actions_gt.shape[0], -1) * flat_action_mask
+        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, {
@@ -389,30 +458,30 @@ class EVO1Policy(PreTrainedPolicy):
 
     @torch.no_grad()
     def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
+        if kwargs.get("inference_delay") is not None or kwargs.get("prev_chunk_left_over") is not None:
+            raise NotImplementedError(
+                "EVO1 does not implement real-time-chunking (RTC) inference; "
+                "use the synchronous inference backend."
+            )
         self.eval()
 
         prompts = self._normalize_task_batch(batch)
         image_batches, image_masks = self._collect_image_batches(batch)
         states, _state_mask = self._prepare_state(batch)
-        fused_tokens = self._compute_fused_tokens(prompts, image_batches, image_masks)
-        states = states.to(dtype=self._inference_compute_dtype)
-        fused_tokens = fused_tokens.to(dtype=self._inference_compute_dtype)
         embodiment_ids = self._get_embodiment_ids(batch, states.shape[0])
         action_mask = self._prepare_inference_action_mask(states.shape[0])
 
-        with (
-            torch.autocast(device_type="cuda", dtype=torch.bfloat16)
-            if self.config.use_amp and str(self.config.device).startswith("cuda")
-            else nullcontext()
-        ):
+        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,
             )
         actions = actions.view(states.shape[0], self.config.chunk_size, self.config.max_action_dim)
-        return actions
+        return actions.to(dtype=torch.float32)
 
     @torch.no_grad()
     def select_action(self, batch: dict[str, Tensor], **kwargs) -> Tensor:

src/lerobot/policies/evo1/processor_evo1.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 from copy import deepcopy
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import Any
 
 import torch
@@ -101,9 +101,9 @@ class Evo1PadStateProcessorStep(ObservationProcessorStep):
         self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
     ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
         new_features = {ft: feats.copy() for ft, feats in features.items()}
-        state_feats = new_features.setdefault(FeatureType.STATE, {})
-        if OBS_STATE in state_feats:
-            state_feats[OBS_STATE] = PolicyFeature(type=FeatureType.STATE, shape=(self.max_state_dim,))
+        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]:
@@ -157,7 +157,7 @@ class Evo1PadActionProcessorStep(ProcessorStep):
         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(FeatureType.ACTION, {})
+        action_feats = new_features.setdefault(PipelineFeatureType.ACTION, {})
         action_feats[ACTION] = PolicyFeature(type=FeatureType.ACTION, shape=(self.max_action_dim,))
         return new_features
 
@@ -214,7 +214,7 @@ class Evo1ActionProcessorStep(PolicyActionProcessorStep):
         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(FeatureType.ACTION, {})
+        action_feats = new_features.setdefault(PipelineFeatureType.ACTION, {})
         action_feats[ACTION] = PolicyFeature(type=FeatureType.ACTION, shape=(self.action_dim,))
         return new_features
 
@@ -328,7 +328,17 @@ def ensure_evo1_processor_steps(
     preprocessor: PolicyProcessorPipeline,
     postprocessor: PolicyProcessorPipeline,
 ) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
-    """Add EVO1 processor steps when loading older checkpoints that do not serialize them."""
+    """Reconcile checkpoint-loaded pipelines with the current EVO1 config.
+
+    Adds the EVO1 steps when loading older checkpoints that do not serialize them, restores the
+    EVO1 batch converter (converters are not serialized), and refreshes the config-driven step
+    parameters (padding widths, action cropping, gripper binarization) so CLI overrides at
+    load/eval time take effect on checkpoints that already serialize these steps.
+    """
+
+    # 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
 
     has_state_padding = any(isinstance(step, Evo1PadStateProcessorStep) for step in preprocessor.steps)
     if not has_state_padding:
@@ -350,25 +360,44 @@ def ensure_evo1_processor_steps(
         steps.insert(insert_idx, Evo1PadActionProcessorStep(max_action_dim=config.max_action_dim))
         preprocessor.steps = steps
 
-    has_action_processor = any(isinstance(step, Evo1ActionProcessorStep) for step in postprocessor.steps)
-    if not has_action_processor:
-        steps = list(postprocessor.steps)
+    preprocessor.steps = [
+        replace(step, max_state_dim=config.max_state_dim)
+        if isinstance(step, Evo1PadStateProcessorStep)
+        else replace(step, max_action_dim=config.max_action_dim)
+        if isinstance(step, Evo1PadActionProcessorStep)
+        else step
+        for step in preprocessor.steps
+    ]
+
+    current_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,
-            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,
-            ),
-        )
-        postprocessor.steps = steps
+        steps.insert(insert_idx, current_action_step)
+    else:
+        steps[action_step_idx] = current_action_step
+    # Actions must leave the postprocessor as float32 (numpy cannot represent bf16); older
+    # checkpoints serialized the device step without a float_dtype.
+    steps = [
+        replace(step, float_dtype="float32")
+        if isinstance(step, DeviceProcessorStep) and step.float_dtype is None
+        else step
+        for step in steps
+    ]
+    postprocessor.steps = steps
 
     _refresh_evo1_normalization_steps(config, preprocessor, postprocessor)
     return preprocessor, postprocessor
@@ -411,7 +440,8 @@ def make_evo1_pre_post_processors(
             gripper_below_threshold_value=config.gripper_below_threshold_value,
             gripper_above_threshold_value=config.gripper_above_threshold_value,
         ),
-        DeviceProcessorStep(device="cpu"),
+        # float32 so downstream numpy conversion works even when the policy computes in bf16.
+        DeviceProcessorStep(device="cpu", float_dtype="float32"),
     ]
 
     return (

src/lerobot/policies/factory.py

@@ -169,9 +169,9 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
 
         return FastWAMPolicy
     elif name == "evo1":
-        from .evo1.modeling_evo1 import EVO1Policy
+        from .evo1.modeling_evo1 import Evo1Policy
 
-        return EVO1Policy
+        return Evo1Policy
     else:
         try:
             return _get_policy_cls_from_policy_name(name=name)

src/lerobot/scripts/lerobot_eval.py

@@ -283,7 +283,7 @@ def rollout(
             action = action_transition[ACTION]
 
             # Convert to CPU / numpy.
-            action_numpy: np.ndarray = action.detach().to(device="cpu", dtype=torch.float32).numpy()
+            action_numpy: np.ndarray = action.to("cpu").numpy()
             assert action_numpy.ndim == 2, "Action dimensions should be (batch, action_dim)"
 
             # Apply the next action.

tests/envs/test_dispatch.py

@@ -64,28 +64,6 @@ def test_processors_delegation():
     assert len(pre.steps) == 0
 
 
-def test_processors_delegation_supports_legacy_override_signature():
-    """External EnvConfig subclasses with the old get_env_processors() signature keep working."""
-    from lerobot.processor.pipeline import DataProcessorPipeline
-
-    @EnvConfig.register_subclass("_dispatch_legacy_proc_test")
-    @dataclass
-    class _Env(EnvConfig):
-        task: str = "x"
-        features: dict[str, PolicyFeature] = field(default_factory=dict)
-
-        @property
-        def gym_kwargs(self):
-            return {}
-
-        def get_env_processors(self):
-            return DataProcessorPipeline(steps=[]), DataProcessorPipeline(steps=[])
-
-    pre, post = make_env_pre_post_processors(_Env(), policy_cfg=object())
-    assert isinstance(pre, DataProcessorPipeline)
-    assert isinstance(post, DataProcessorPipeline)
-
-
 def test_libero_processors_are_policy_agnostic():
     cfg = LiberoEnv()
     pre, post = make_env_pre_post_processors(cfg, policy_cfg=object())
@@ -186,7 +164,7 @@ def test_custom_get_env_processors_override():
         def gym_kwargs(self):
             return {}
 
-        def get_env_processors(self, policy_cfg=None):
+        def get_env_processors(self):
             return DataProcessorPipeline(steps=[]), DataProcessorPipeline(steps=[])
 
     pre, post = _Env().get_env_processors()

tests/policies/evo1/test_evo1.py

@@ -35,11 +35,29 @@ from lerobot.policies.evo1.processor_evo1 import (
     Evo1PadActionProcessorStep,
     Evo1PadStateProcessorStep,
     ensure_evo1_processor_steps,
+    evo1_batch_to_transition,
     make_evo1_pre_post_processors,
 )
 from lerobot.policies.factory import get_policy_class, make_policy_config
-from lerobot.processor import NormalizerProcessorStep, PolicyProcessorPipeline, UnnormalizerProcessorStep
-from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
+from lerobot.processor import (
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    PolicyProcessorPipeline,
+    UnnormalizerProcessorStep,
+)
+from lerobot.processor.converters import (
+    batch_to_transition,
+    policy_action_to_transition,
+    transition_to_batch,
+    transition_to_policy_action,
+)
+from lerobot.utils.constants import (
+    ACTION,
+    OBS_IMAGES,
+    OBS_STATE,
+    POLICY_POSTPROCESSOR_DEFAULT_NAME,
+    POLICY_PREPROCESSOR_DEFAULT_NAME,
+)
 
 STATE_DIM = 4
 ACTION_DIM = 3
@@ -49,8 +67,8 @@ CHUNK_SIZE = 2
 EMBED_DIM = 8
 
 
-class DummyEVO1(nn.Module):
-    def __init__(self, config):
+class DummyEvo1Model(nn.Module):
+    def __init__(self, config, vlm_hub_kwargs=None):
         super().__init__()
         self.config = config
         self.embedder = nn.Dropout(p=0.0)
@@ -68,7 +86,9 @@ class DummyEVO1(nn.Module):
         self.embedder_training_calls.append(self.embedder.training)
         # images is a list of per-camera (B, C, H, W) tensors, so the batch dim is images[0].shape[0].
         batch_size = images[0].shape[0]
-        return torch.ones(batch_size, 4, EMBED_DIM, requires_grad=torch.is_grad_enabled())
+        tokens = torch.ones(batch_size, 4, EMBED_DIM, requires_grad=torch.is_grad_enabled())
+        valid_mask = torch.ones(batch_size, 4, dtype=torch.bool)
+        return tokens, valid_mask
 
     def forward(
         self,
@@ -77,6 +97,7 @@ class DummyEVO1(nn.Module):
         actions_gt=None,
         action_mask=None,
         embodiment_ids=None,
+        context_mask=None,
     ):
         batch_size = fused_tokens.shape[0]
         if actions_gt is None:
@@ -86,6 +107,31 @@ class DummyEVO1(nn.Module):
         return pred_velocity, noise
 
 
+class ChunkCountingDummyModel(DummyEvo1Model):
+    """Emits per-step distinguishable actions so queue ordering and re-prediction are observable."""
+
+    def __init__(self, config, vlm_hub_kwargs=None):
+        super().__init__(config, vlm_hub_kwargs)
+        self.chunks_predicted = 0
+
+    def forward(
+        self,
+        fused_tokens,
+        state=None,
+        actions_gt=None,
+        action_mask=None,
+        embodiment_ids=None,
+        context_mask=None,
+    ):
+        if actions_gt is not None:
+            return super().forward(fused_tokens, state, actions_gt, action_mask, embodiment_ids, context_mask)
+        self.chunks_predicted += 1
+        batch_size = fused_tokens.shape[0]
+        step_values = torch.arange(CHUNK_SIZE, dtype=torch.float32) + 10.0 * self.chunks_predicted
+        chunk = step_values.repeat_interleave(MAX_ACTION_DIM).unsqueeze(0).repeat(batch_size, 1)
+        return chunk
+
+
 def make_config(training_stage="stage1", **kwargs):
     config_kwargs = {
         "device": "cpu",
@@ -138,6 +184,24 @@ def make_stats(state_dim=STATE_DIM, action_dim=ACTION_DIM):
     }
 
 
+def make_flowmatching_head(**overrides):
+    kwargs = {
+        "embed_dim": EMBED_DIM,
+        "hidden_dim": 16,
+        "action_dim": CHUNK_SIZE * ACTION_DIM,
+        "horizon": CHUNK_SIZE,
+        "per_action_dim": ACTION_DIM,
+        "num_heads": 2,
+        "num_layers": 1,
+        "num_inference_timesteps": 2,
+        "state_dim": STATE_DIM,
+        "state_hidden_dim": 16,
+        "num_categories": 1,
+    }
+    kwargs.update(overrides)
+    return FlowmatchingActionHead(**kwargs)
+
+
 def test_evo1_factory_registration():
     cfg = make_policy_config(
         "evo1",
@@ -151,7 +215,7 @@ def test_evo1_factory_registration():
     )
 
     assert isinstance(cfg, Evo1Config)
-    assert get_policy_class("evo1") is modeling_evo1.EVO1Policy
+    assert get_policy_class("evo1") is modeling_evo1.Evo1Policy
 
 
 def test_evo1_stage_defaults_and_consistency():
@@ -208,6 +272,19 @@ def test_evo1_stage_defaults_and_consistency():
     )
     assert explicit_off.finetune_action_head is False
 
+    # An explicit finetune_vlm=False without branch-level flags freezes both branches instead of
+    # raising an inconsistency error.
+    frozen_vlm = make_config(
+        training_stage="stage2",
+        apply_training_stage_defaults=False,
+        finetune_vlm=False,
+    )
+    assert (
+        frozen_vlm.finetune_vlm,
+        frozen_vlm.finetune_language_model,
+        frozen_vlm.finetune_vision_model,
+    ) == (False, False, False)
+
     try:
         make_config(
             training_stage="stage2",
@@ -226,6 +303,11 @@ def test_evo1_rejects_non_square_image_resolution():
         make_config(image_resolution=(448, 320))
 
 
+def test_evo1_rejects_out_of_range_default_embodiment_id():
+    with pytest.raises(ValueError, match="default_embodiment_id"):
+        make_config(default_embodiment_id=3, num_categories=2)
+
+
 def test_evo1_model_uses_image_resolution_and_trainable_checkpointing(monkeypatch):
     captured: dict = {}
 
@@ -238,16 +320,65 @@ def test_evo1_model_uses_image_resolution_and_trainable_checkpointing(monkeypatc
     monkeypatch.setattr(evo1_model, "InternVL3Embedder", SpyEmbedder)
 
     stage1 = make_config(training_stage="stage1", image_resolution=(224, 224))
-    evo1_model.EVO1(stage1)
+    evo1_model.Evo1Model(stage1)
     assert captured["image_size"] == 224
     # VLM is frozen in stage1, so gradient checkpointing is gated off.
     assert captured["enable_gradient_checkpointing"] is False
 
     stage2 = make_config(training_stage="stage2", image_resolution=(224, 224))
-    evo1_model.EVO1(stage2)
+    evo1_model.Evo1Model(stage2)
     assert captured["enable_gradient_checkpointing"] is True
 
 
+def test_set_finetune_flags_targets_native_hf_internvl_submodules(monkeypatch):
+    class FakeInternVLModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.language_model = nn.Linear(2, 2)
+            self.vision_tower = nn.Linear(2, 2)
+            self.multi_modal_projector = nn.Linear(2, 2)
+
+    class FakeEmbedder(nn.Module):
+        def __init__(self, **kwargs):
+            super().__init__()
+            self.model = FakeInternVLModel()
+
+    monkeypatch.setattr(evo1_model, "InternVL3Embedder", FakeEmbedder)
+
+    stage2_model = evo1_model.Evo1Model(make_config(training_stage="stage2"))
+    stage2_model.set_finetune_flags()
+    vlm = stage2_model.embedder.model
+    assert all(p.requires_grad for p in vlm.language_model.parameters())
+    assert all(p.requires_grad for p in vlm.vision_tower.parameters())
+    assert all(p.requires_grad for p in vlm.multi_modal_projector.parameters())
+    assert all(p.requires_grad for p in stage2_model.action_head.parameters())
+
+    stage1_model = evo1_model.Evo1Model(make_config(training_stage="stage1"))
+    stage1_model.set_finetune_flags()
+    vlm = stage1_model.embedder.model
+    assert not any(p.requires_grad for p in vlm.parameters())
+    assert all(p.requires_grad for p in stage1_model.action_head.parameters())
+
+
+def test_set_finetune_flags_fails_loudly_on_unknown_vlm_layout(monkeypatch):
+    class LegacyLayoutModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.language_model = nn.Linear(2, 2)
+            self.vision_model = nn.Linear(2, 2)  # trust_remote_code-era attribute name
+            self.mlp1 = nn.Linear(2, 2)
+
+    class FakeEmbedder(nn.Module):
+        def __init__(self, **kwargs):
+            super().__init__()
+            self.model = LegacyLayoutModel()
+
+    monkeypatch.setattr(evo1_model, "InternVL3Embedder", FakeEmbedder)
+    model = evo1_model.Evo1Model(make_config(training_stage="stage2"))
+    with pytest.raises(AttributeError, match="vision_tower"):
+        model.set_finetune_flags()
+
+
 def test_evo1_policy_processors_pad_state_crop_action_and_binarize_gripper():
     libero_action_dim = 7
     config = make_config(
@@ -300,10 +431,19 @@ def test_evo1_policy_processors_pad_state_crop_action_and_binarize_gripper():
     processed = postprocessor(action)
 
     assert processed.shape == (2, 7)
+    assert processed.dtype == torch.float32
     assert torch.allclose(processed[:, :6], action[:, :6])
     assert torch.equal(processed[:, 6], torch.tensor([1.0, -1.0]))
 
 
+def test_evo1_postprocessor_returns_float32_for_bf16_actions():
+    config = make_config()
+    _preprocessor, postprocessor = make_evo1_pre_post_processors(config, dataset_stats=make_stats())
+
+    processed = postprocessor(torch.zeros(2, MAX_ACTION_DIM, dtype=torch.bfloat16))
+    assert processed.dtype == torch.float32
+
+
 def test_evo1_legacy_processors_are_completed_before_normalization():
     config = make_config(
         max_state_dim=MAX_STATE_DIM,
@@ -333,6 +473,7 @@ def test_evo1_legacy_processors_are_completed_before_normalization():
 
     preprocessor, postprocessor = ensure_evo1_processor_steps(config, legacy_pre, legacy_post)
 
+    assert preprocessor.to_transition is evo1_batch_to_transition
     assert isinstance(preprocessor.steps[0], Evo1PadStateProcessorStep)
     assert isinstance(preprocessor.steps[1], Evo1PadActionProcessorStep)
     assert isinstance(preprocessor.steps[2], NormalizerProcessorStep)
@@ -352,9 +493,51 @@ def test_evo1_legacy_processors_are_completed_before_normalization():
     assert sum(isinstance(step, Evo1ActionProcessorStep) for step in postprocessor.steps) == 1
 
 
+def test_evo1_processor_save_load_round_trip_applies_config_overrides(tmp_path):
+    train_config = make_config()
+    preprocessor, postprocessor = make_evo1_pre_post_processors(train_config, dataset_stats=make_stats())
+    preprocessor.save_pretrained(tmp_path)
+    postprocessor.save_pretrained(tmp_path)
+
+    loaded_pre = PolicyProcessorPipeline.from_pretrained(
+        tmp_path,
+        config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json",
+        to_transition=batch_to_transition,
+        to_output=transition_to_batch,
+    )
+    loaded_post = PolicyProcessorPipeline.from_pretrained(
+        tmp_path,
+        config_filename=f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json",
+        to_transition=policy_action_to_transition,
+        to_output=transition_to_policy_action,
+    )
+
+    # Simulate eval-time CLI overrides on a checkpoint that already serializes the EVO1 steps.
+    eval_config = make_config(binarize_gripper=True, postprocess_action_dim=ACTION_DIM)
+    loaded_pre, loaded_post = ensure_evo1_processor_steps(eval_config, loaded_pre, loaded_post)
+
+    assert loaded_pre.to_transition is evo1_batch_to_transition
+    action_step = next(step for step in loaded_post.steps if isinstance(step, Evo1ActionProcessorStep))
+    assert action_step.binarize_gripper is True
+    assert action_step.action_dim == ACTION_DIM
+    device_step = next(step for step in loaded_post.steps if isinstance(step, DeviceProcessorStep))
+    assert device_step.float_dtype == "float32"
+
+    # Non-observation extras (embodiment_id, ...) must survive the reloaded preprocessor.
+    processed = loaded_pre(
+        {
+            "task": "pick the block",
+            OBS_STATE: torch.zeros(STATE_DIM),
+            f"{OBS_IMAGES}.front": torch.rand(3, 16, 16),
+            "embodiment_id": torch.tensor([0]),
+        }
+    )
+    assert "embodiment_id" in processed
+
+
 def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch):
-    monkeypatch.setattr(modeling_evo1, "EVO1", DummyEVO1)
-    policy = modeling_evo1.EVO1Policy(make_config())
+    monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
+    policy = modeling_evo1.Evo1Policy(make_config())
     preprocessor, _postprocessor = make_evo1_pre_post_processors(policy.config, dataset_stats=make_stats())
     training_batch = preprocessor(make_batch(include_action=True))
 
@@ -371,33 +554,100 @@ def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch):
 
     action_chunk = policy.predict_action_chunk(make_batch(include_action=False))
     assert action_chunk.shape == (2, CHUNK_SIZE, MAX_ACTION_DIM)
+    assert action_chunk.dtype == torch.float32
 
     policy.reset()
     selected = policy.select_action(make_batch(include_action=False))
     assert selected.shape == (2, MAX_ACTION_DIM)
 
 
+def test_evo1_forward_masks_padded_action_timesteps(monkeypatch):
+    monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
+    policy = modeling_evo1.Evo1Policy(make_config())
+
+    batch = make_batch(include_action=True)
+    batch[ACTION] = torch.ones(2, CHUNK_SIZE, ACTION_DIM)
+    # Give the padded (past-episode-end) timestep a huge value: if it leaked into the loss, the
+    # loss would blow up far beyond 1.0.
+    batch[ACTION][:, -1, :] = 100.0
+    batch["action_is_pad"] = torch.zeros(2, CHUNK_SIZE, dtype=torch.bool)
+    batch["action_is_pad"][:, -1] = True
+
+    loss, metrics = policy.forward(batch)
+
+    # DummyEvo1Model predicts zero velocity and zero noise, so each active element contributes
+    # (0 - action)^2 = 1.0 for the in-episode ones-valued actions.
+    assert metrics["active_action_dims"] == ACTION_DIM * (CHUNK_SIZE - 1)
+    assert torch.isclose(loss, torch.tensor(1.0))
+
+
+def test_evo1_select_action_queue_orders_steps_and_repredicts(monkeypatch):
+    monkeypatch.setattr(modeling_evo1, "Evo1Model", ChunkCountingDummyModel)
+    policy = modeling_evo1.Evo1Policy(make_config(n_action_steps=CHUNK_SIZE))
+
+    batch = make_batch(include_action=False)
+    first = policy.select_action(batch)
+    second = policy.select_action(batch)
+    third = policy.select_action(batch)
+
+    # First chunk provides steps 10, 11 in order; the third call triggers a fresh prediction (20).
+    assert torch.all(first == 10.0)
+    assert torch.all(second == 11.0)
+    assert torch.all(third == 20.0)
+    assert policy.model.chunks_predicted == 2
+
+
+def test_evo1_predict_action_chunk_rejects_rtc_kwargs(monkeypatch):
+    monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
+    policy = modeling_evo1.Evo1Policy(make_config())
+    with pytest.raises(NotImplementedError, match="RTC"):
+        policy.predict_action_chunk(make_batch(include_action=False), inference_delay=2)
+
+
+def test_evo1_missing_configured_camera_needs_empty_cameras_budget(monkeypatch):
+    monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
+    batch = make_batch(include_action=False)  # only provides the front camera
+
+    two_camera_features = {
+        OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(STATE_DIM,)),
+        f"{OBS_IMAGES}.front": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 16, 16)),
+        f"{OBS_IMAGES}.wrist": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 16, 16)),
+    }
+    strict_policy = modeling_evo1.Evo1Policy(make_config(input_features=dict(two_camera_features)))
+    with pytest.raises(ValueError, match="empty_cameras"):
+        strict_policy._collect_image_batches(batch)
+
+    # empty_cameras adds placeholder camera features that are never present in the batch; they
+    # become masked-out views instead of crashing with a KeyError.
+    padded_policy = modeling_evo1.Evo1Policy(make_config(empty_cameras=1))
+    assert len(padded_policy.config.image_features) == 2
+    camera_images, image_masks = padded_policy._collect_image_batches(batch)
+    assert len(camera_images) == 1
+    assert image_masks.tolist() == [[True, False], [True, False]]
+
+
 def test_stage1_frozen_vlm_embeddings_do_not_track_gradients(monkeypatch):
-    monkeypatch.setattr(modeling_evo1, "EVO1", DummyEVO1)
-    policy = modeling_evo1.EVO1Policy(make_config(training_stage="stage1"))
+    monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
+    policy = modeling_evo1.Evo1Policy(make_config(training_stage="stage1"))
     policy.train()
 
     image_batches, image_masks = policy._collect_image_batches(make_batch(include_action=False))
-    fused_tokens = policy._compute_fused_tokens(["pick", "place"], image_batches, image_masks)
+    fused_tokens, context_mask = policy._compute_fused_tokens(["pick", "place"], image_batches, image_masks)
 
     assert policy.model.grad_enabled_calls == [False]
     assert policy.model.embedder_training_calls == [False]
     assert not fused_tokens.requires_grad
+    assert context_mask is not None
     assert policy.model.embedder.training is False
 
 
 def test_stage2_vlm_embeddings_track_gradients(monkeypatch):
-    monkeypatch.setattr(modeling_evo1, "EVO1", DummyEVO1)
-    policy = modeling_evo1.EVO1Policy(make_config(training_stage="stage2"))
+    monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
+    policy = modeling_evo1.Evo1Policy(make_config(training_stage="stage2"))
     policy.train()
 
     image_batches, image_masks = policy._collect_image_batches(make_batch(include_action=False))
-    fused_tokens = policy._compute_fused_tokens(["pick", "place"], image_batches, image_masks)
+    fused_tokens, _context_mask = policy._compute_fused_tokens(["pick", "place"], image_batches, image_masks)
 
     assert policy.model.grad_enabled_calls == [True]
     assert policy.model.embedder_training_calls == [True]
@@ -407,8 +657,8 @@ def test_stage2_vlm_embeddings_track_gradients(monkeypatch):
 def test_collect_image_batches_handles_unbatched_chw(monkeypatch):
     # Regression for an issue where batch_size was read from shape[0] before normalizing
     # per-camera tensor dims, so an unbatched (C, H, W) input was treated as batch_size=C.
-    monkeypatch.setattr(modeling_evo1, "EVO1", DummyEVO1)
-    policy = modeling_evo1.EVO1Policy(make_config())
+    monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
+    policy = modeling_evo1.Evo1Policy(make_config())
     batch = {
         OBS_STATE: torch.randn(1, STATE_DIM),
         f"{OBS_IMAGES}.front": torch.rand(3, 16, 16),
@@ -423,10 +673,26 @@ def test_collect_image_batches_handles_unbatched_chw(monkeypatch):
     assert image_masks.tolist() == [[True, False]]
 
 
+def test_evo1_state_mask_zeroes_masked_dims(monkeypatch):
+    monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
+    policy = modeling_evo1.Evo1Policy(make_config())
+    batch = {
+        OBS_STATE: torch.ones(2, STATE_DIM),
+        "state_mask": torch.tensor([[True, True, False, False]] * 2),
+    }
+
+    states, mask = policy._prepare_state(batch)
+
+    assert torch.all(states[:, :2] == 1.0)
+    assert torch.all(states[:, 2:] == 0.0)
+    assert mask[:, :2].all()
+    assert not mask[:, 2:].any()
+
+
 def test_evo1_action_mask_accepts_chunk_size_one(monkeypatch):
-    monkeypatch.setattr(modeling_evo1, "EVO1", DummyEVO1)
+    monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
     config = make_config(chunk_size=1, n_action_steps=1)
-    policy = modeling_evo1.EVO1Policy(config)
+    policy = modeling_evo1.Evo1Policy(config)
     batch = make_batch(include_action=True)
     batch[ACTION] = torch.randn(2, ACTION_DIM)
     batch["action_mask"] = torch.ones(2, ACTION_DIM, dtype=torch.bool)
@@ -440,19 +706,7 @@ def test_evo1_action_mask_accepts_chunk_size_one(monkeypatch):
 
 
 def test_flowmatching_state_encoder_for_horizon_one():
-    head = FlowmatchingActionHead(
-        embed_dim=EMBED_DIM,
-        hidden_dim=16,
-        action_dim=ACTION_DIM,
-        horizon=1,
-        per_action_dim=ACTION_DIM,
-        num_heads=2,
-        num_layers=1,
-        num_inference_timesteps=2,
-        state_dim=STATE_DIM,
-        state_hidden_dim=16,
-        num_categories=1,
-    )
+    head = make_flowmatching_head(action_dim=ACTION_DIM, horizon=1)
 
     assert head.state_encoder is not None
     pred_velocity, noise = head(
@@ -466,6 +720,71 @@ def test_flowmatching_state_encoder_for_horizon_one():
     assert noise.shape == (2, 1, ACTION_DIM)
 
 
+def test_flowmatching_get_action_real_path_respects_action_mask():
+    torch.manual_seed(0)
+    head = make_flowmatching_head()
+
+    action_mask = torch.zeros(2, ACTION_DIM, dtype=torch.bool)
+    action_mask[:, :2] = True
+    actions = head.get_action(
+        torch.randn(2, 4, EMBED_DIM),
+        state=torch.randn(2, STATE_DIM),
+        action_mask=action_mask,
+    )
+
+    assert actions.shape == (2, CHUNK_SIZE * ACTION_DIM)
+    assert torch.isfinite(actions).all()
+    action_seq = actions.view(2, CHUNK_SIZE, ACTION_DIM)
+    assert torch.all(action_seq[..., 2] == 0.0)
+
+
+def test_flowmatching_context_mask_blocks_masked_context_tokens():
+    head = make_flowmatching_head()
+    state = torch.randn(2, STATE_DIM)
+    action_mask = torch.ones(2, ACTION_DIM, dtype=torch.bool)
+    fused = torch.randn(2, 4, EMBED_DIM)
+    context_mask = torch.ones(2, 4, dtype=torch.bool)
+    context_mask[:, -1] = False
+    corrupted = fused.clone()
+    corrupted[:, -1] = 1e4
+
+    torch.manual_seed(0)
+    reference = head.get_action(fused, state=state, action_mask=action_mask, context_mask=context_mask)
+    torch.manual_seed(0)
+    with_garbage = head.get_action(corrupted, state=state, action_mask=action_mask, context_mask=context_mask)
+
+    assert torch.allclose(reference, with_garbage)
+
+
+def test_flowmatching_head_accepts_pooled_2d_context():
+    head = make_flowmatching_head()
+    pred_velocity, noise = head(
+        torch.randn(2, EMBED_DIM),  # pooled (B, E) context from return_cls_only
+        state=torch.randn(2, STATE_DIM),
+        actions_gt=torch.randn(2, CHUNK_SIZE, ACTION_DIM),
+        action_mask=torch.ones(2, CHUNK_SIZE, ACTION_DIM, dtype=torch.bool),
+    )
+    assert pred_velocity.shape == (2, CHUNK_SIZE * ACTION_DIM)
+
+    actions = head.get_action(
+        torch.randn(2, EMBED_DIM),
+        state=torch.randn(2, STATE_DIM),
+        action_mask=torch.ones(2, ACTION_DIM, dtype=torch.bool),
+    )
+    assert actions.shape == (2, CHUNK_SIZE * ACTION_DIM)
+
+
+def test_flowmatching_rejects_out_of_range_embodiment_ids():
+    head = make_flowmatching_head(num_categories=2)
+    with pytest.raises(ValueError, match="num_categories"):
+        head.get_action(
+            torch.randn(2, 4, EMBED_DIM),
+            state=torch.randn(2, STATE_DIM),
+            action_mask=torch.ones(2, ACTION_DIM, dtype=torch.bool),
+            embodiment_id=torch.tensor([0, 5]),
+        )
+
+
 def test_evo1_batched_pixel_values_shape_and_zero_padding():
     torch.manual_seed(0)
     batch_size, image_size, max_views = 2, 448, 3