diff --git a/docs/source/evo1.mdx b/docs/source/evo1.mdx index eda8d65fa..cd40a02d2 100644 --- a/docs/source/evo1.mdx +++ b/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 diff --git a/docs/source/policy_evo1_README.md b/docs/source/policy_evo1_README.md index 3c6d31c83..dc8b75344 100644 --- a/docs/source/policy_evo1_README.md +++ b/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}}, } ``` diff --git a/src/lerobot/optim/schedulers.py b/src/lerobot/optim/schedulers.py index 250650089..e58310800 100644 --- a/src/lerobot/optim/schedulers.py +++ b/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): diff --git a/src/lerobot/policies/evo1/__init__.py b/src/lerobot/policies/evo1/__init__.py index f15a27d8c..581b2b824 100644 --- a/src/lerobot/policies/evo1/__init__.py +++ b/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"] diff --git a/src/lerobot/policies/evo1/configuration_evo1.py b/src/lerobot/policies/evo1/configuration_evo1.py index 66c086fd0..d86ddf1ad 100644 --- a/src/lerobot/policies/evo1/configuration_evo1.py +++ b/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, ) diff --git a/src/lerobot/policies/evo1/evo1_model.py b/src/lerobot/policies/evo1/evo1_model.py index 64c1d4fa2..f192d14b9 100644 --- a/src/lerobot/policies/evo1/evo1_model.py +++ b/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) diff --git a/src/lerobot/policies/evo1/flow_matching.py b/src/lerobot/policies/evo1/flow_matching.py index 016d43163..8b558fc8b 100644 --- a/src/lerobot/policies/evo1/flow_matching.py +++ b/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: diff --git a/src/lerobot/policies/evo1/internvl3_embedder.py b/src/lerobot/policies/evo1/internvl3_embedder.py index faed1bf46..7ce67d76d 100644 --- a/src/lerobot/policies/evo1/internvl3_embedder.py +++ b/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: diff --git a/src/lerobot/policies/evo1/modeling_evo1.py b/src/lerobot/policies/evo1/modeling_evo1.py index 541ca877a..78026d73f 100644 --- a/src/lerobot/policies/evo1/modeling_evo1.py +++ b/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: diff --git a/src/lerobot/policies/evo1/processor_evo1.py b/src/lerobot/policies/evo1/processor_evo1.py index 702c034ce..b6f41348c 100644 --- a/src/lerobot/policies/evo1/processor_evo1.py +++ b/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 ( diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index 1e6558ffe..2afac8e77 100644 --- a/src/lerobot/policies/factory.py +++ b/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) diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index a91a29382..1ec4ea75f 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/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. diff --git a/tests/envs/test_dispatch.py b/tests/envs/test_dispatch.py index b038832af..50f208422 100644 --- a/tests/envs/test_dispatch.py +++ b/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() diff --git a/tests/policies/evo1/test_evo1.py b/tests/policies/evo1/test_evo1.py index af2587b8b..0a64596a0 100644 --- a/tests/policies/evo1/test_evo1.py +++ b/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