diff --git a/src/lerobot/rollout/context.py b/src/lerobot/rollout/context.py index 80594ef8b..7a7249e50 100644 --- a/src/lerobot/rollout/context.py +++ b/src/lerobot/rollout/context.py @@ -27,7 +27,7 @@ from threading import Event import torch -from lerobot.configs import PreTrainedConfig +from lerobot.configs import FeatureType, PreTrainedConfig from lerobot.datasets import ( LeRobotDataset, aggregate_pipeline_dataset_features, @@ -186,6 +186,12 @@ def build_rollout_context( if hasattr(full_config, "compile_model"): full_config.compile_model = cfg.use_torch_compile + if full_config.type == "vqbet" and cfg.device == "mps": + raise NotImplementedError( + "Current implementation of VQBeT does not support `mps` backend. " + "Please use `cpu` or `cuda` backend." + ) + if full_config.use_peft: from peft import PeftConfig, PeftModel @@ -294,6 +300,22 @@ def build_rollout_context( raw_action_keys, ) + # Validate visual features if no rename_map is active + rename_map = cfg.dataset.rename_map if cfg.dataset else {} + if not rename_map: + expected_visuals = {k for k, v in full_config.input_features.items() if v.type == FeatureType.VISUAL} + provided_visuals = { + f"observation.{k}" for k, v in robot.observation_features.items() if isinstance(v, tuple) + } + policy_subset = expected_visuals.issubset(provided_visuals) + hw_subset = provided_visuals.issubset(expected_visuals) + if not (policy_subset or hw_subset): + raise ValueError( + f"Visual feature mismatch between policy and robot hardware.\n" + f"Policy expects: {expected_visuals}\n" + f"Robot provides: {provided_visuals}" + ) + # --- 5. Dataset ------------- dataset = None if cfg.dataset is not None and not isinstance(cfg.strategy, BaseStrategyConfig):