#!/usr/bin/env python # Copyright 2026 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from __future__ import annotations import pytest import torch from torch import nn import lerobot.policies.evo1.evo1_model as evo1_model import lerobot.policies.evo1.modeling_evo1 as modeling_evo1 from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.policies.evo1.configuration_evo1 import Evo1Config from lerobot.policies.evo1.flow_matching import FlowmatchingActionHead from lerobot.policies.evo1.internvl3_embedder import ( IMAGENET_MEAN, IMAGENET_STD, _batched_pixel_values, ) from lerobot.policies.evo1.processor_evo1 import ( Evo1ActionProcessorStep, Evo1PadActionProcessorStep, Evo1PadStateProcessorStep, evo1_batch_to_transition, make_evo1_pre_post_processors, reconcile_evo1_processors, ) from lerobot.policies.factory import get_policy_class, make_policy_config from lerobot.policies.rtc.configuration_rtc import RTCConfig from lerobot.policies.rtc.modeling_rtc import RTCProcessor 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 MAX_STATE_DIM = 6 MAX_ACTION_DIM = 5 CHUNK_SIZE = 2 EMBED_DIM = 8 class DummyEvo1Model(nn.Module): def __init__(self, config, vlm_hub_kwargs=None): super().__init__() self.config = config self.embedder = nn.Dropout(p=0.0) self.action_head = nn.Linear(1, 1) self.get_vl_embeddings_calls = 0 self.grad_enabled_calls = [] self.embedder_training_calls = [] def set_finetune_flags(self): return None def get_vl_embeddings(self, images, image_mask, prompt=None, return_cls_only=False): self.get_vl_embeddings_calls += 1 self.grad_enabled_calls.append(torch.is_grad_enabled()) 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] 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, fused_tokens, state=None, actions_gt=None, action_mask=None, embodiment_ids=None, context_mask=None, **kwargs, ): batch_size = fused_tokens.shape[0] if actions_gt is None: return torch.ones(batch_size, CHUNK_SIZE * MAX_ACTION_DIM) pred_velocity = torch.zeros(batch_size, CHUNK_SIZE * MAX_ACTION_DIM) noise = torch.zeros_like(actions_gt) 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, **kwargs, ): 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", "vlm_model_name": "dummy-internvl3", "training_stage": training_stage, "chunk_size": CHUNK_SIZE, "n_action_steps": 1, "max_state_dim": MAX_STATE_DIM, "max_action_dim": MAX_ACTION_DIM, "max_views": 2, "embed_dim": EMBED_DIM, "hidden_dim": 16, "state_hidden_dim": 16, "num_heads": 2, "num_layers": 1, "num_inference_timesteps": 2, "input_features": { OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(STATE_DIM,)), f"{OBS_IMAGES}.front": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 16, 16)), }, "output_features": { ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(ACTION_DIM,)), }, } config_kwargs.update(kwargs) return Evo1Config(**config_kwargs) def make_batch(include_action=True): batch = { "task": ["pick the block", "place the block"], OBS_STATE: torch.randn(2, STATE_DIM), f"{OBS_IMAGES}.front": torch.rand(2, 3, 16, 16), } if include_action: batch[ACTION] = torch.randn(2, CHUNK_SIZE, ACTION_DIM) return batch def make_stats(state_dim=STATE_DIM, action_dim=ACTION_DIM): return { OBS_STATE: { "min": torch.full((state_dim,), -2.0), "max": torch.full((state_dim,), 2.0), }, ACTION: { "min": torch.full((action_dim,), -1.0), "max": torch.full((action_dim,), 1.0), }, } 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", device="cpu", vlm_model_name="dummy-internvl3", input_features={ OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(STATE_DIM,)), f"{OBS_IMAGES}.front": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 16, 16)), }, output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(ACTION_DIM,))}, ) assert isinstance(cfg, Evo1Config) assert get_policy_class("evo1") is modeling_evo1.Evo1Policy def test_evo1_stage_defaults_and_consistency(): stage1 = make_config(training_stage="stage1") assert (stage1.finetune_vlm, stage1.finetune_language_model, stage1.finetune_vision_model) == ( False, False, False, ) assert stage1.finetune_action_head is True stage2 = make_config(training_stage="stage2") assert (stage2.finetune_vlm, stage2.finetune_language_model, stage2.finetune_vision_model) == ( True, True, True, ) assert stage2.finetune_action_head is True stage2_from_stage1_checkpoint_flags = make_config( training_stage="stage2", finetune_vlm=False, finetune_language_model=False, finetune_vision_model=False, finetune_action_head=False, ) assert ( stage2_from_stage1_checkpoint_flags.finetune_vlm, stage2_from_stage1_checkpoint_flags.finetune_language_model, stage2_from_stage1_checkpoint_flags.finetune_vision_model, ) == ( True, True, True, ) assert stage2_from_stage1_checkpoint_flags.finetune_action_head is True explicit_off = make_config( training_stage="stage2", apply_training_stage_defaults=False, finetune_vlm=False, finetune_language_model=False, finetune_vision_model=False, finetune_action_head=False, ) assert ( explicit_off.finetune_vlm, explicit_off.finetune_language_model, explicit_off.finetune_vision_model, ) == ( False, False, False, ) 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", apply_training_stage_defaults=False, finetune_vlm=True, finetune_language_model=False, ) except ValueError as exc: assert "Inconsistent EVO1 finetune config" in str(exc) else: raise AssertionError("Expected inconsistent finetune config to raise ValueError") def test_evo1_rejects_non_square_image_resolution(): with pytest.raises(ValueError, match="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 = {} class SpyEmbedder(nn.Module): def __init__(self, **kwargs): super().__init__() captured.clear() captured.update(kwargs) monkeypatch.setattr(evo1_model, "InternVL3Embedder", SpyEmbedder) stage1 = make_config(training_stage="stage1", image_resolution=(224, 224)) 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.Evo1Model(stage2) assert captured["enable_gradient_checkpointing"] is True class FakeInternVLModel(nn.Module): """Minimal stand-in with the native HF InternVL submodule layout.""" 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() def test_set_finetune_flags_targets_native_hf_internvl_submodules(monkeypatch): 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( max_state_dim=MAX_STATE_DIM, max_action_dim=8, postprocess_action_dim=libero_action_dim, binarize_gripper=True, output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(libero_action_dim,))}, ) stats = make_stats(action_dim=libero_action_dim) preprocessor, postprocessor = make_evo1_pre_post_processors(config, dataset_stats=stats) assert isinstance(preprocessor.steps[2], Evo1PadStateProcessorStep) assert isinstance(preprocessor.steps[3], Evo1PadActionProcessorStep) assert isinstance(preprocessor.steps[4], NormalizerProcessorStep) assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep) assert isinstance(postprocessor.steps[1], Evo1ActionProcessorStep) normalizer = preprocessor.steps[4] assert normalizer.features[OBS_STATE].shape == (MAX_STATE_DIM,) assert normalizer.features[ACTION].shape == (8,) assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (MAX_STATE_DIM,) assert normalizer._tensor_stats[ACTION]["min"].shape == (8,) processed_batch = preprocessor( { "task": "pick the block", OBS_STATE: torch.zeros(STATE_DIM), ACTION: torch.zeros(libero_action_dim), f"{OBS_IMAGES}.front": torch.rand(3, 16, 16), } ) processed_state = processed_batch[OBS_STATE] assert processed_state.shape == (1, MAX_STATE_DIM) assert torch.allclose(processed_state, torch.zeros_like(processed_state)) assert processed_batch[ACTION].shape == (1, 8) assert torch.allclose(processed_batch[ACTION], torch.zeros_like(processed_batch[ACTION])) assert processed_batch["action_mask"].shape == (1, 8) assert processed_batch["action_mask"][:, :libero_action_dim].all() assert not processed_batch["action_mask"][:, libero_action_dim:].any() action = torch.tensor( [ [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.5, 0.7], [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7], ], dtype=torch.float32, ) 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_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 applied on top of the loaded pipelines. eval_config = make_config(binarize_gripper=True, postprocess_action_dim=ACTION_DIM) loaded_pre, loaded_post = reconcile_evo1_processors(eval_config, loaded_pre, loaded_post) assert loaded_pre.to_transition is evo1_batch_to_transition assert sum(isinstance(step, Evo1ActionProcessorStep) for step in loaded_post.steps) == 1 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 # The float32 output dtype is part of the serialized pipeline itself. 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, "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)) assert training_batch[ACTION].shape == (2, CHUNK_SIZE, MAX_ACTION_DIM) assert training_batch["action_mask"].shape == (2, CHUNK_SIZE, MAX_ACTION_DIM) assert training_batch["action_mask"][:, :, :ACTION_DIM].all() assert not training_batch["action_mask"][:, :, ACTION_DIM:].any() loss, metrics = policy.forward(training_batch) assert loss.ndim == 0 assert torch.isfinite(loss) assert metrics["active_action_dims"] == ACTION_DIM * CHUNK_SIZE assert policy.model.get_vl_embeddings_calls == 1 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_without_rtc_config(monkeypatch): monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model) policy = modeling_evo1.Evo1Policy(make_config()) with pytest.raises(RuntimeError, match="RTC"): policy.predict_action_chunk(make_batch(include_action=False), inference_delay=2) def test_evo1_rtc_processor_wiring(monkeypatch): monkeypatch.setattr(evo1_model, "InternVL3Embedder", FakeEmbedder) policy = modeling_evo1.Evo1Policy(make_config()) assert policy.rtc_processor is None assert policy.model.rtc_processor is None # The RTC rollout backend assigns rtc_config after loading and re-inits the processor. policy.config.rtc_config = RTCConfig(execution_horizon=CHUNK_SIZE) policy.init_rtc_processor() assert isinstance(policy.rtc_processor, RTCProcessor) assert policy.model.rtc_processor is policy.rtc_processor # RTC drives predict_action_chunk directly; the select_action queue path is unsupported. with pytest.raises(AssertionError, match="select_action"): policy.select_action(make_batch(include_action=False)) def test_flowmatching_rtc_guidance_pulls_prefix_toward_previous_chunk(): head = make_flowmatching_head(num_inference_timesteps=16) processor = RTCProcessor(RTCConfig(execution_horizon=CHUNK_SIZE)) fused = torch.randn(2, 4, EMBED_DIM) state = torch.randn(2, STATE_DIM) action_mask = torch.ones(2, ACTION_DIM, dtype=torch.bool) prev_chunk = torch.tensor([0.7, -0.4, 0.2]).expand(2, CHUNK_SIZE, ACTION_DIM).contiguous() torch.manual_seed(0) unguided = head.get_action(fused, state=state, action_mask=action_mask) unguided = unguided.view(2, CHUNK_SIZE, ACTION_DIM) torch.manual_seed(0) guided = head.get_action( fused, state=state, action_mask=action_mask, inference_delay=1, prev_chunk_left_over=prev_chunk, rtc_processor=processor, ) guided = guided.view(2, CHUNK_SIZE, ACTION_DIM) # The frozen prefix (first inference_delay steps) must land far closer to the previous chunk # than the unguided sample from the same noise does. guided_dist = (guided[:, 0] - prev_chunk[:, 0]).abs().mean() unguided_dist = (unguided[:, 0] - prev_chunk[:, 0]).abs().mean() assert guided_dist < 0.5 * unguided_dist assert torch.isfinite(guided).all() def test_flowmatching_rtc_first_chunk_without_leftover_matches_unguided(): head = make_flowmatching_head(num_inference_timesteps=4) processor = RTCProcessor(RTCConfig(execution_horizon=CHUNK_SIZE)) fused = torch.randn(2, 4, EMBED_DIM) state = torch.randn(2, STATE_DIM) action_mask = torch.ones(2, ACTION_DIM, dtype=torch.bool) torch.manual_seed(0) unguided = head.get_action(fused, state=state, action_mask=action_mask) torch.manual_seed(0) first_chunk = head.get_action( fused, state=state, action_mask=action_mask, inference_delay=2, prev_chunk_left_over=None, rtc_processor=processor, ) assert torch.allclose(unguided, first_chunk) 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, "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, 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, "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, _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] assert fused_tokens.requires_grad 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, "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), } camera_images, image_masks = policy._collect_image_batches(batch) # One present camera, returned as a batched (B, C, H, W) tensor with the unbatched CHW frame # promoted to batch_size=1 (not read as batch_size=C). assert len(camera_images) == 1 assert camera_images[0].shape == (1, 3, 16, 16) 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, "Evo1Model", DummyEvo1Model) config = make_config(chunk_size=1, n_action_steps=1) 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) actions, action_mask = policy._prepare_actions(batch) assert actions.shape == (2, 1, MAX_ACTION_DIM) assert action_mask.shape == (2, 1, MAX_ACTION_DIM) assert action_mask[:, :, :ACTION_DIM].all() assert not action_mask[:, :, ACTION_DIM:].any() def test_flowmatching_state_encoder_for_horizon_one(): head = make_flowmatching_head(action_dim=ACTION_DIM, horizon=1) assert head.state_encoder is not None pred_velocity, noise = head( torch.randn(2, 4, EMBED_DIM), state=torch.randn(2, STATE_DIM), actions_gt=torch.randn(2, 1, ACTION_DIM), action_mask=torch.ones(2, 1, ACTION_DIM, dtype=torch.bool), ) assert pred_velocity.shape == (2, ACTION_DIM) 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 camera_images = [torch.rand(batch_size, 3, 40, 50)] # a single present camera mean = torch.tensor(IMAGENET_MEAN) std = torch.tensor(IMAGENET_STD) pixel_values = _batched_pixel_values( camera_images, max_views, image_size, mean, std, torch.float32, torch.device("cpu") ) assert pixel_values.shape == (batch_size * max_views, 3, image_size, image_size) grouped = pixel_values.reshape(batch_size, max_views, 3, image_size, image_size) # Absent views (indices 1, 2) are zero images, normalized to the constant -mean/std. expected_pad = (-mean / std).view(1, 3, 1, 1) for view in (1, 2): assert torch.allclose( grouped[:, view], expected_pad.expand(batch_size, 3, image_size, image_size), atol=1e-5 ) # The present view is genuinely different from the constant pad value. assert not torch.allclose( grouped[:, 0], expected_pad.expand(batch_size, 3, image_size, image_size), atol=1e-3 )