rpx_benchmark.tasks.robot_vqa_v2

  1"""RPX Robot VQA v2 — three context modes, all GT from recorded data.
  2
  3Every question can be asked in three ways:
  4  - image_only:    [reference photo] + [scene frame]
  5  - image_text:    [reference photo] + text description + [scene frame]
  6  - text_only:     text description only + [scene frame]
  7
  8The ground truth is IDENTICAL across modes — it comes from masks +
  9questionnaire + bounding boxes, never from any VLM or LLM.
 10
 11The model being evaluated receives the prompt in one mode and must
 12produce the answer. Comparing accuracy across modes reveals which
 13prompting strategy is most robust for robot deployment.
 14
 15GT sources:
 16  - Instance masks (SAM2 + human refinement) → visibility, bbox, count
 17  - FewSOL questionnaire (MTurk) → category, color, material, shape
 18  - Bounding box centroids (from masks) → spatial relations
 19  - Phase diff (clutter vs clean) → state changes
 20  - Single-object scenes → reference images
 21"""
 22
 23from __future__ import annotations
 24
 25import hashlib
 26import random
 27from dataclasses import dataclass, field
 28from typing import Dict, List, Optional, Tuple
 29
 30
 31@dataclass
 32class RobotVQA:
 33    """One question with all three context modes.
 34
 35    The model is evaluated in ONE mode at a time. The GT is the same
 36    regardless of mode.
 37    """
 38
 39    scene_id: str
 40    phase: int
 41    frame_id: str
 42    task: str  # recognition | attribute | spatial | state | counting
 43
 44    # The question in each context mode
 45    image_only_prompt: str  # what to show/ask with reference image only
 46    image_text_prompt: str  # what to show/ask with reference image + text
 47    text_only_prompt: str  # what to ask with text only
 48
 49    reference_image: Optional[str]  # path to single-object reference (None for text-only tasks)
 50    scene_frame: str  # path to the scene frame being evaluated
 51
 52    correct_answer: str
 53    options: List[str]  # 4-choice MCQ
 54    correct_index: int
 55
 56    robot_operation: str  # what pipeline stage this tests
 57    metadata: Dict = field(default_factory=dict)
 58
 59
 60def _seed(s: str, p: int, f: str, tag: str) -> int:
 61    return int(hashlib.md5(f"{s}/{p}/{f}/{tag}".encode()).hexdigest()[:8], 16)
 62
 63
 64def _shuffle(correct: str, distractors: List[str], seed: int) -> Tuple[List[str], int]:
 65    opts = [correct] + distractors[:3]
 66    random.Random(seed).shuffle(opts)
 67    return opts, opts.index(correct)
 68
 69
 70# ─── Recognition: "Is this object in the scene?" ──────────────────────────
 71
 72
 73def gen_recognition(
 74    visible: List[Dict],  # objects visible in this frame
 75    ref_images: Dict[str, str],  # {category: path} from single-object scenes
 76    absent_cats: List[str],  # categories NOT visible
 77    questionnaire: Dict[str, Dict],  # {category: {color, material, shape}}
 78    scene_id: str,
 79    phase: int,
 80    frame_id: str,
 81    scene_frame: str,
 82) -> List[RobotVQA]:
 83    """Object recognition in three modes.
 84
 85    Same GT (yes/no from mask visibility), different prompts.
 86    """
 87    pairs = []
 88    vis_cats = {o["category"] for o in visible}
 89
 90    # Positive examples
 91    for cat in sorted(vis_cats)[:3]:
 92        if cat not in ref_images:
 93            continue
 94        attrs = questionnaire.get(cat, {})
 95        color = attrs.get("color", "")
 96        material = attrs.get("material", "")
 97        desc = f"{color} {material} {cat}".strip()
 98
 99        s = _seed(scene_id, phase, frame_id, f"rec_yes_{cat}")
100        opts, idx = _shuffle("yes", ["no", "similar but different object", "cannot determine"], s)
101
102        pairs.append(
103            RobotVQA(
104                scene_id=scene_id,
105                phase=phase,
106                frame_id=frame_id,
107                task="recognition",
108                image_only_prompt="Is the object shown in the reference image present in this scene?",
109                image_text_prompt=f"Is the {desc} shown in the reference image present in this scene?",
110                text_only_prompt=f"Is there a {desc} in this scene?",
111                reference_image=ref_images[cat],
112                scene_frame=scene_frame,
113                correct_answer="yes",
114                options=opts,
115                correct_index=idx,
116                robot_operation="object_registration",
117                metadata={"category": cat, "present": True},
118            )
119        )
120
121    # Negative examples (balanced 50/50)
122    n_neg = 0
123    for cat in absent_cats:
124        if cat not in ref_images or n_neg >= len([p for p in pairs if p.metadata.get("present")]):
125            continue
126        attrs = questionnaire.get(cat, {})
127        color = attrs.get("color", "")
128        material = attrs.get("material", "")
129        desc = f"{color} {material} {cat}".strip()
130
131        s = _seed(scene_id, phase, frame_id, f"rec_no_{cat}")
132        opts, idx = _shuffle("no", ["yes", "partially visible", "cannot determine"], s)
133
134        pairs.append(
135            RobotVQA(
136                scene_id=scene_id,
137                phase=phase,
138                frame_id=frame_id,
139                task="recognition",
140                image_only_prompt="Is the object shown in the reference image present in this scene?",
141                image_text_prompt=f"Is the {desc} shown in the reference image present in this scene?",
142                text_only_prompt=f"Is there a {desc} in this scene?",
143                reference_image=ref_images[cat],
144                scene_frame=scene_frame,
145                correct_answer="no",
146                options=opts,
147                correct_index=idx,
148                robot_operation="object_registration",
149                metadata={"category": cat, "present": False},
150            )
151        )
152        n_neg += 1
153
154    return pairs
155
156
157# ─── Attribute grounding: "Pick the RED mug" ──────────────────────────────
158
159
160def gen_attribute(
161    visible: List[Dict],
162    ref_images: Dict[str, str],
163    questionnaire: Dict[str, Dict],
164    all_objects: List[Dict],
165    scene_id: str,
166    phase: int,
167    frame_id: str,
168    scene_frame: str,
169) -> List[RobotVQA]:
170    """Attribute grounding — only fires when disambiguation is needed."""
171    pairs = []
172    from collections import Counter
173
174    cat_counts = Counter(o["category"] for o in visible)
175
176    for cat, count in cat_counts.items():
177        if count < 2:
178            continue  # no disambiguation needed
179        instances = [o for o in visible if o["category"] == cat]
180
181        for obj in instances[:2]:
182            color = obj.get("color", "")
183            if not color:
184                continue
185            other_colors = [o.get("color", "?") for o in instances if o.get("color") != color]
186
187            s = _seed(scene_id, phase, frame_id, f"attr_{cat}_{color}")
188            gt = f"the {color} {cat}"
189            distractors = [f"the {c} {cat}" for c in other_colors[:2]] + [f"any {cat}"]
190            while len(distractors) < 3:
191                distractors.append("none of these")
192            opts, idx = _shuffle(gt, distractors[:3], s)
193
194            ref_path = ref_images.get(cat)
195            pairs.append(
196                RobotVQA(
197                    scene_id=scene_id,
198                    phase=phase,
199                    frame_id=frame_id,
200                    task="attribute",
201                    image_only_prompt=f"Which of the visible {cat} objects matches the reference image?",
202                    image_text_prompt=f"Which object is the {color} {cat}?",
203                    text_only_prompt=f"A robot must pick the {color} {cat}. Which object should it grasp?",
204                    reference_image=ref_path,
205                    scene_frame=scene_frame,
206                    correct_answer=gt,
207                    options=opts,
208                    correct_index=idx,
209                    robot_operation="language_grounding",
210                    metadata={"category": cat, "target_color": color},
211                )
212            )
213
214    return pairs
215
216
217# ─── Spatial grounding: "Pick the object LEFT of the bowl" ────────────────
218
219
220def gen_spatial(
221    visible: List[Dict],  # must have bbox_cx, bbox_cy
222    scene_id: str,
223    phase: int,
224    frame_id: str,
225    scene_frame: str,
226) -> List[RobotVQA]:
227    """Spatial grounding — image-plane relations only. Text-only task."""
228    pairs = []
229    if len(visible) < 2:
230        return pairs
231
232    for ref in visible[:3]:
233        cx, cy = ref.get("bbox_cx"), ref.get("bbox_cy")
234        if cx is None:
235            continue
236
237        # Bind cx/cy at lambda definition time (default args) so each loop
238        # iteration captures the current ref's centroid, not the last one.
239        for direction, check in [
240            ("left", lambda o, cx=cx: o.get("bbox_cx", cx) < cx),
241            ("right", lambda o, cx=cx: o.get("bbox_cx", cx) > cx),
242            ("above", lambda o, cy=cy: o.get("bbox_cy", cy) < cy),
243            ("below", lambda o, cy=cy: o.get("bbox_cy", cy) > cy),
244        ]:
245            candidates = [
246                o for o in visible if o is not ref and check(o) and o.get("bbox_cx") is not None
247            ]
248            if not candidates:
249                continue
250            # Nearest
251            target = min(
252                candidates,
253                key=lambda o, cx=cx, cy=cy: abs(o["bbox_cx"] - cx) + abs(o["bbox_cy"] - cy),
254            )
255            gt = target["category"]
256
257            s = _seed(scene_id, phase, frame_id, f"sp_{direction}_{ref['category']}")
258            others = list(
259                {o["category"] for o in visible if o["category"] not in (gt, ref["category"])}
260            )
261            others += ["nothing"]
262            distractors = random.Random(s).sample(others, min(3, len(others)))
263            while len(distractors) < 3:
264                distractors.append("cannot determine")
265            opts, idx = _shuffle(gt, distractors[:3], s)
266
267            pairs.append(
268                RobotVQA(
269                    scene_id=scene_id,
270                    phase=phase,
271                    frame_id=frame_id,
272                    task="spatial",
273                    image_only_prompt=f"What object is {direction} of the {ref['category']}?",
274                    image_text_prompt=f"What object is {direction} of the {ref['category']}?",
275                    text_only_prompt=f"A robot must pick the object {direction} of the {ref['category']}. What should it pick?",
276                    reference_image=None,
277                    scene_frame=scene_frame,
278                    correct_answer=gt,
279                    options=opts,
280                    correct_index=idx,
281                    robot_operation="spatial_command",
282                    metadata={"reference": ref["category"], "direction": direction},
283                )
284            )
285            break  # one per direction per ref
286
287    return pairs
288
289
290# ─── State verification: "Did the pick succeed?" ─────────────────────────
291
292
293def gen_state_verification(
294    clutter_objects: List[Dict],
295    clean_objects: List[Dict],
296    scene_id: str,
297    frame_id: str,
298    clutter_frame: str,
299    clean_frame: str,
300) -> List[RobotVQA]:
301    """State verification — compare two phases. Uses TWO frames as input."""
302    pairs = []
303    clu_cats = {o["category"] for o in clutter_objects}
304    cln_cats = {o["category"] for o in clean_objects}
305    removed = clu_cats - cln_cats
306    remained = clu_cats & cln_cats
307
308    for obj_cat in list(removed)[:2]:
309        s = _seed(scene_id, 0, frame_id, f"state_{obj_cat}")
310        distractors = list(remained)[:3]
311        while len(distractors) < 3:
312            distractors.append("nothing changed")
313        opts, idx = _shuffle(obj_cat, distractors[:3], s)
314
315        pairs.append(
316            RobotVQA(
317                scene_id=scene_id,
318                phase=0,
319                frame_id=frame_id,
320                task="state",
321                image_only_prompt="Compare these two images. Which object was removed?",
322                image_text_prompt="Compare the before and after images. Which object was removed from the scene?",
323                text_only_prompt="After the robot's manipulation, which object is no longer on the table?",
324                reference_image=clutter_frame,  # "before" image
325                scene_frame=clean_frame,  # "after" image
326                correct_answer=obj_cat,
327                options=opts,
328                correct_index=idx,
329                robot_operation="task_verification",
330                metadata={"change_type": "removed", "object": obj_cat},
331            )
332        )
333
334    return pairs
335
336
337# ─── Counting: "How many left to pack?" ──────────────────────────────────
338
339
340def gen_counting(
341    visible: List[Dict],
342    scene_id: str,
343    phase: int,
344    frame_id: str,
345    scene_frame: str,
346) -> List[RobotVQA]:
347    """Counting — text-only task, GT from instance masks."""
348    pairs = []
349    from collections import Counter
350
351    counts = Counter(o["category"] for o in visible)
352
353    for cat, count in list(counts.items())[:3]:
354        s = _seed(scene_id, phase, frame_id, f"cnt_{cat}")
355        gt = str(count)
356        pool = [str(i) for i in range(max(count + 3, 5)) if str(i) != gt]
357        distractors = random.Random(s).sample(pool, min(3, len(pool)))
358        while len(distractors) < 3:
359            distractors.append("0")
360        opts, idx = _shuffle(gt, distractors[:3], s)
361
362        pairs.append(
363            RobotVQA(
364                scene_id=scene_id,
365                phase=phase,
366                frame_id=frame_id,
367                task="counting",
368                image_only_prompt=f"How many {cat} objects are in this scene?",
369                image_text_prompt=f"How many {cat} objects are visible in this scene?",
370                text_only_prompt=f"A robot is packing all {cat} objects. How many does it need to pick?",
371                reference_image=None,
372                scene_frame=scene_frame,
373                correct_answer=gt,
374                options=opts,
375                correct_index=idx,
376                robot_operation="task_progress",
377                metadata={"category": cat, "count": count},
378            )
379        )
380
381    return pairs
382
383
384# ─── Evaluation ──────────────────────────────────────────────────────────
385
386CONTEXT_MODES = ("image_only", "image_text", "text_only")
387
388
389def evaluate(
390    predictions: Dict[str, List[int]],  # {mode: [pred_indices]}
391    ground_truth: List[RobotVQA],
392) -> Dict:
393    """Evaluate predictions across context modes.
394
395    Args:
396        predictions: {mode_name: list of predicted option indices}
397                     e.g. {"image_only": [0,1,3,...], "text_only": [2,0,1,...]}
398        ground_truth: list of RobotVQA
399
400    Returns:
401        Per-mode accuracy + per-task + per-phase + cross-mode comparison
402    """
403    results = {}
404    for mode, preds in predictions.items():
405        assert len(preds) == len(ground_truth), (
406            f"{mode}: {len(preds)} preds vs {len(ground_truth)} GT"
407        )
408
409        by_task, by_phase = {}, {}
410        correct = 0
411        for pred, qa in zip(preds, ground_truth, strict=False):
412            hit = pred == qa.correct_index
413            correct += hit
414            by_task.setdefault(qa.task, []).append(hit)
415            by_phase.setdefault(qa.phase, []).append(hit)
416
417        n = len(preds)
418        pa = {p: sum(v) / len(v) for p, v in by_phase.items()}
419        mode_result = {
420            "accuracy": correct / n if n else 0.0,
421            "per_task": {t: sum(v) / len(v) for t, v in by_task.items()},
422            "per_phase": pa,
423        }
424        if 0 in pa and 1 in pa:
425            mode_result["str_clu_to_int"] = pa[1] - pa[0]
426        if 1 in pa and 2 in pa:
427            mode_result["str_int_to_cln"] = pa[2] - pa[1]
428        results[mode] = mode_result
429
430    # Cross-mode comparison: which mode is most robust?
431    if len(results) > 1:
432        results["mode_comparison"] = {
433            mode: r["accuracy"] for mode, r in results.items() if mode != "mode_comparison"
434        }
435
436    return results