diff --git a/openwam/dataloader/libero.py b/openwam/dataloader/libero.py index 68686451..6bca6084 100644 --- a/openwam/dataloader/libero.py +++ b/openwam/dataloader/libero.py @@ -89,15 +89,25 @@ class LiberoDataset(LeRobotV3Reader): PROMPT_FILE_REQUIRED = True DEPLOY_ACTION_MODE = ACTION_MODE - # Fixed camera layout (head, wrist, unused); None slots are rendered black - # in multiview mode. Override via ``camera_layout``. - DEFAULT_CAMERA_LAYOUT: ClassVar[Tuple[Optional[str], ...]] = ( + # LIBERO exports have used both the canonical ``image/image2`` names and + # the older ``agentview_image/wrist_image`` names. Resolve from the + # features advertised by info.json so a dataset's actual camera streams + # are used instead of assuming one naming convention. + HEAD_CAMERA_PRIORITY: ClassVar[Tuple[str, ...]] = ( "observation.images.image", + "observation.images.agentview_image", + "observation.images.agentview", + ) + WRIST_CAMERA_PRIORITY: ClassVar[Tuple[str, ...]] = ( + "observation.images.wrist_image", "observation.images.image2", - None, + "observation.images.robot0_eye_in_hand_image", + "observation.images.wrist", ) CONFIG_KEYS: ClassVar[Tuple[str, ...]] = LeRobotV3Reader.CONFIG_KEYS + ( "action_mode", + "head_camera_priority", + "wrist_camera_priority", "normalization_stats_path", ) @@ -106,6 +116,8 @@ def __init__( dataset_dir: str, *, action_mode: str = ACTION_MODE, + head_camera_priority: Optional[Sequence[str]] = None, + wrist_camera_priority: Optional[Sequence[str]] = None, normalization_stats_path: Optional[str] = None, unify_action: bool = False, unify_action_map: Optional[Any] = None, @@ -123,6 +135,8 @@ def __init__( self.action_mode = mode self._source_stats_path = str(normalization_stats_path) if normalization_stats_path else None self._state_normalization_stats: Optional[dict] = None + self._head_priority = _as_priority(head_camera_priority, self.HEAD_CAMERA_PRIORITY) + self._wrist_priority = _as_priority(wrist_camera_priority, self.WRIST_CAMERA_PRIORITY) super().__init__( dataset_dir=dataset_dir, unify_action=bool(unify_action), @@ -131,11 +145,49 @@ def __init__( ) def _resolve_cameras(self, info: dict): - if self._target_camera is not None: + # ``target_camera`` selects the single-view stream. In multiview mode + # it must not suppress the auxiliary streams: benchmark configs may + # carry the single-view field for deployment compatibility. + if not self._multiview and self._target_camera is not None: return self._target_camera, None, None - layout = list(self._camera_layout_param or self.DEFAULT_CAMERA_LAYOUT) - layout += [None] * (3 - len(layout)) - return tuple(str(cam) if cam else None for cam in layout[:3]) + + features = info.get("features", {}) or {} + + # An explicit layout is the requested slot order, but configs can be + # shared by LIBERO exports with different camera names. Replace a + # missing non-null request with the feature-based camera for that slot, + # then feed the resolved names back to the base reader so assembly and + # decoding use the same keys. + if self._multiview and self._camera_layout_param is not None: + requested = list(self._camera_layout_param)[:3] + requested += [None] * (3 - len(requested)) + fallback = ( + _pick_feature(features, self._head_priority), + _pick_feature(features, self._wrist_priority), + None, + ) + resolved = tuple( + (str(name) if name in features else fallback[index]) if name else None + for index, name in enumerate(requested) + ) + self._camera_layout_param = list(resolved) + return resolved + + if not self._multiview and self._camera_layout_param is not None: + requested = list(self._camera_layout_param)[:3] + if requested: + requested_camera = requested[0] + if requested_camera is None: + return None, None, None + if requested_camera in features: + return str(requested_camera), None, None + return _pick_feature(features, self._head_priority), None, None + + return ( + _pick_feature(features, self._head_priority), + _pick_feature(features, self._wrist_priority), + None, + ) def _post_init(self, info: dict) -> None: features = info.get("features", {}) or {} diff --git a/tests/dataloader/test_libero.py b/tests/dataloader/test_libero.py index b97e9db2..b8ca72d8 100644 --- a/tests/dataloader/test_libero.py +++ b/tests/dataloader/test_libero.py @@ -45,3 +45,107 @@ def test_reader_rejects_incomplete_compatibility_stats(tmp_path) -> None: np.save(reader._source_stats_path, {"eef": {}}, allow_pickle=True) with pytest.raises(KeyError, match="eef"): reader._load_stats({}) + + +def test_multiview_target_camera_does_not_disable_wrist_stream() -> None: + reader = object.__new__(LiberoDataset) + reader._multiview = True + reader._target_camera = "observation.images.image" + reader._camera_layout_param = [ + "observation.images.image", + "observation.images.image2", + None, + ] + reader._head_priority = LiberoDataset.HEAD_CAMERA_PRIORITY + reader._wrist_priority = LiberoDataset.WRIST_CAMERA_PRIORITY + + features = { + "observation.images.image": {"dtype": "video"}, + "observation.images.image2": {"dtype": "video"}, + } + assert reader._resolve_cameras({"features": features}) == ( + "observation.images.image", + "observation.images.image2", + None, + ) + + +def test_libero_camera_resolution_falls_back_to_info_features() -> None: + reader = object.__new__(LiberoDataset) + reader._multiview = True + reader._target_camera = "observation.images.image" + reader._camera_layout_param = None + reader._head_priority = LiberoDataset.HEAD_CAMERA_PRIORITY + reader._wrist_priority = LiberoDataset.WRIST_CAMERA_PRIORITY + + features = { + "observation.images.agentview_image": {"dtype": "video"}, + "observation.images.wrist_image": {"dtype": "video"}, + } + assert reader._resolve_cameras({"features": features}) == ( + "observation.images.agentview_image", + "observation.images.wrist_image", + None, + ) + + +def test_single_view_target_camera_still_wins() -> None: + reader = object.__new__(LiberoDataset) + reader._multiview = False + reader._target_camera = "observation.images.image2" + reader._camera_layout_param = None + reader._head_priority = LiberoDataset.HEAD_CAMERA_PRIORITY + reader._wrist_priority = LiberoDataset.WRIST_CAMERA_PRIORITY + + assert reader._resolve_cameras({"features": {}}) == ( + "observation.images.image2", + None, + None, + ) + + +def test_explicit_layout_falls_back_to_legacy_wrist_feature() -> None: + reader = object.__new__(LiberoDataset) + reader._multiview = True + reader._target_camera = "observation.images.image" + reader._camera_layout_param = [ + "observation.images.image", + "observation.images.image2", + None, + ] + reader._head_priority = LiberoDataset.HEAD_CAMERA_PRIORITY + reader._wrist_priority = LiberoDataset.WRIST_CAMERA_PRIORITY + + features = { + "observation.images.image": {"dtype": "video"}, + "observation.images.wrist_image": {"dtype": "video"}, + } + assert reader._resolve_cameras({"features": features}) == ( + "observation.images.image", + "observation.images.wrist_image", + None, + ) + assert reader._camera_layout_param == [ + "observation.images.image", + "observation.images.wrist_image", + None, + ] + + +def test_single_view_explicit_layout_is_honored_without_target() -> None: + reader = object.__new__(LiberoDataset) + reader._multiview = False + reader._target_camera = None + reader._camera_layout_param = ["observation.images.agentview_image", None, None] + reader._head_priority = LiberoDataset.HEAD_CAMERA_PRIORITY + reader._wrist_priority = LiberoDataset.WRIST_CAMERA_PRIORITY + + features = { + "observation.images.image": {"dtype": "video"}, + "observation.images.agentview_image": {"dtype": "video"}, + } + assert reader._resolve_cameras({"features": features}) == ( + "observation.images.agentview_image", + None, + None, + )