Skip to content

Commit 1204dde

Browse files
wayrisewayrise
authored
Merge pull request #40 from OpenWAM-Official/fix/libero-multiview-camera
Fix LIBERO multiview camera resolution
2 parents c51935f + cb07d6d commit 1204dde

2 files changed

Lines changed: 164 additions & 8 deletions

File tree

‎openwam/dataloader/libero.py‎

Lines changed: 60 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -89,15 +89,25 @@ class LiberoDataset(LeRobotV3Reader):
8989
PROMPT_FILE_REQUIRED = True
9090
DEPLOY_ACTION_MODE = ACTION_MODE
9191

92-
# Fixed camera layout (head, wrist, unused); None slots are rendered black
93-
# in multiview mode. Override via ``camera_layout``.
94-
DEFAULT_CAMERA_LAYOUT: ClassVar[Tuple[Optional[str], ...]] = (
92+
# LIBERO exports have used both the canonical ``image/image2`` names and
93+
# the older ``agentview_image/wrist_image`` names. Resolve from the
94+
# features advertised by info.json so a dataset's actual camera streams
95+
# are used instead of assuming one naming convention.
96+
HEAD_CAMERA_PRIORITY: ClassVar[Tuple[str, ...]] = (
9597
"observation.images.image",
98+
"observation.images.agentview_image",
99+
"observation.images.agentview",
100+
)
101+
WRIST_CAMERA_PRIORITY: ClassVar[Tuple[str, ...]] = (
102+
"observation.images.wrist_image",
96103
"observation.images.image2",
97-
None,
104+
"observation.images.robot0_eye_in_hand_image",
105+
"observation.images.wrist",
98106
)
99107
CONFIG_KEYS: ClassVar[Tuple[str, ...]] = LeRobotV3Reader.CONFIG_KEYS + (
100108
"action_mode",
109+
"head_camera_priority",
110+
"wrist_camera_priority",
101111
"normalization_stats_path",
102112
)
103113

@@ -106,6 +116,8 @@ def __init__(
106116
dataset_dir: str,
107117
*,
108118
action_mode: str = ACTION_MODE,
119+
head_camera_priority: Optional[Sequence[str]] = None,
120+
wrist_camera_priority: Optional[Sequence[str]] = None,
109121
normalization_stats_path: Optional[str] = None,
110122
unify_action: bool = False,
111123
unify_action_map: Optional[Any] = None,
@@ -123,6 +135,8 @@ def __init__(
123135
self.action_mode = mode
124136
self._source_stats_path = str(normalization_stats_path) if normalization_stats_path else None
125137
self._state_normalization_stats: Optional[dict] = None
138+
self._head_priority = _as_priority(head_camera_priority, self.HEAD_CAMERA_PRIORITY)
139+
self._wrist_priority = _as_priority(wrist_camera_priority, self.WRIST_CAMERA_PRIORITY)
126140
super().__init__(
127141
dataset_dir=dataset_dir,
128142
unify_action=bool(unify_action),
@@ -131,11 +145,49 @@ def __init__(
131145
)
132146

133147
def _resolve_cameras(self, info: dict):
134-
if self._target_camera is not None:
148+
# ``target_camera`` selects the single-view stream. In multiview mode
149+
# it must not suppress the auxiliary streams: benchmark configs may
150+
# carry the single-view field for deployment compatibility.
151+
if not self._multiview and self._target_camera is not None:
135152
return self._target_camera, None, None
136-
layout = list(self._camera_layout_param or self.DEFAULT_CAMERA_LAYOUT)
137-
layout += [None] * (3 - len(layout))
138-
return tuple(str(cam) if cam else None for cam in layout[:3])
153+
154+
features = info.get("features", {}) or {}
155+
156+
# An explicit layout is the requested slot order, but configs can be
157+
# shared by LIBERO exports with different camera names. Replace a
158+
# missing non-null request with the feature-based camera for that slot,
159+
# then feed the resolved names back to the base reader so assembly and
160+
# decoding use the same keys.
161+
if self._multiview and self._camera_layout_param is not None:
162+
requested = list(self._camera_layout_param)[:3]
163+
requested += [None] * (3 - len(requested))
164+
fallback = (
165+
_pick_feature(features, self._head_priority),
166+
_pick_feature(features, self._wrist_priority),
167+
None,
168+
)
169+
resolved = tuple(
170+
(str(name) if name in features else fallback[index]) if name else None
171+
for index, name in enumerate(requested)
172+
)
173+
self._camera_layout_param = list(resolved)
174+
return resolved
175+
176+
if not self._multiview and self._camera_layout_param is not None:
177+
requested = list(self._camera_layout_param)[:3]
178+
if requested:
179+
requested_camera = requested[0]
180+
if requested_camera is None:
181+
return None, None, None
182+
if requested_camera in features:
183+
return str(requested_camera), None, None
184+
return _pick_feature(features, self._head_priority), None, None
185+
186+
return (
187+
_pick_feature(features, self._head_priority),
188+
_pick_feature(features, self._wrist_priority),
189+
None,
190+
)
139191

140192
def _post_init(self, info: dict) -> None:
141193
features = info.get("features", {}) or {}

‎tests/dataloader/test_libero.py‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,3 +45,107 @@ def test_reader_rejects_incomplete_compatibility_stats(tmp_path) -> None:
4545
np.save(reader._source_stats_path, {"eef": {}}, allow_pickle=True)
4646
with pytest.raises(KeyError, match="eef"):
4747
reader._load_stats({})
48+
49+
50+
def test_multiview_target_camera_does_not_disable_wrist_stream() -> None:
51+
reader = object.__new__(LiberoDataset)
52+
reader._multiview = True
53+
reader._target_camera = "observation.images.image"
54+
reader._camera_layout_param = [
55+
"observation.images.image",
56+
"observation.images.image2",
57+
None,
58+
]
59+
reader._head_priority = LiberoDataset.HEAD_CAMERA_PRIORITY
60+
reader._wrist_priority = LiberoDataset.WRIST_CAMERA_PRIORITY
61+
62+
features = {
63+
"observation.images.image": {"dtype": "video"},
64+
"observation.images.image2": {"dtype": "video"},
65+
}
66+
assert reader._resolve_cameras({"features": features}) == (
67+
"observation.images.image",
68+
"observation.images.image2",
69+
None,
70+
)
71+
72+
73+
def test_libero_camera_resolution_falls_back_to_info_features() -> None:
74+
reader = object.__new__(LiberoDataset)
75+
reader._multiview = True
76+
reader._target_camera = "observation.images.image"
77+
reader._camera_layout_param = None
78+
reader._head_priority = LiberoDataset.HEAD_CAMERA_PRIORITY
79+
reader._wrist_priority = LiberoDataset.WRIST_CAMERA_PRIORITY
80+
81+
features = {
82+
"observation.images.agentview_image": {"dtype": "video"},
83+
"observation.images.wrist_image": {"dtype": "video"},
84+
}
85+
assert reader._resolve_cameras({"features": features}) == (
86+
"observation.images.agentview_image",
87+
"observation.images.wrist_image",
88+
None,
89+
)
90+
91+
92+
def test_single_view_target_camera_still_wins() -> None:
93+
reader = object.__new__(LiberoDataset)
94+
reader._multiview = False
95+
reader._target_camera = "observation.images.image2"
96+
reader._camera_layout_param = None
97+
reader._head_priority = LiberoDataset.HEAD_CAMERA_PRIORITY
98+
reader._wrist_priority = LiberoDataset.WRIST_CAMERA_PRIORITY
99+
100+
assert reader._resolve_cameras({"features": {}}) == (
101+
"observation.images.image2",
102+
None,
103+
None,
104+
)
105+
106+
107+
def test_explicit_layout_falls_back_to_legacy_wrist_feature() -> None:
108+
reader = object.__new__(LiberoDataset)
109+
reader._multiview = True
110+
reader._target_camera = "observation.images.image"
111+
reader._camera_layout_param = [
112+
"observation.images.image",
113+
"observation.images.image2",
114+
None,
115+
]
116+
reader._head_priority = LiberoDataset.HEAD_CAMERA_PRIORITY
117+
reader._wrist_priority = LiberoDataset.WRIST_CAMERA_PRIORITY
118+
119+
features = {
120+
"observation.images.image": {"dtype": "video"},
121+
"observation.images.wrist_image": {"dtype": "video"},
122+
}
123+
assert reader._resolve_cameras({"features": features}) == (
124+
"observation.images.image",
125+
"observation.images.wrist_image",
126+
None,
127+
)
128+
assert reader._camera_layout_param == [
129+
"observation.images.image",
130+
"observation.images.wrist_image",
131+
None,
132+
]
133+
134+
135+
def test_single_view_explicit_layout_is_honored_without_target() -> None:
136+
reader = object.__new__(LiberoDataset)
137+
reader._multiview = False
138+
reader._target_camera = None
139+
reader._camera_layout_param = ["observation.images.agentview_image", None, None]
140+
reader._head_priority = LiberoDataset.HEAD_CAMERA_PRIORITY
141+
reader._wrist_priority = LiberoDataset.WRIST_CAMERA_PRIORITY
142+
143+
features = {
144+
"observation.images.image": {"dtype": "video"},
145+
"observation.images.agentview_image": {"dtype": "video"},
146+
}
147+
assert reader._resolve_cameras({"features": features}) == (
148+
"observation.images.agentview_image",
149+
None,
150+
None,
151+
)

0 commit comments

Comments
 (0)