rpx_benchmark.tasks.robot_vqa

  1"""RPX Robot-Grounded VQA — tasks a robot actually needs.
  2
  3Every question maps to a real robot operation. Every answer feeds an action.
  4No generic VQA. No "what is the weather like?"
  5
  6Six task types, each grounded in a specific robot pipeline stage:
  7
  8T1. IN-CONTEXT RECOGNITION (the killer feature)
  9    Robot has a reference photo of target object (from single-object scene).
 10    VLM must identify whether that object appears in the cluttered scene.
 11    Uses: paired single-object / multi-object scenes.
 12    This is exactly what DE-ViT, NIDS-Net, and real robot "object registration" do.
 13
 14T2. ATTRIBUTE GROUNDING
 15    Robot receives "pick the red mug" — must distinguish red mug from blue mug.
 16    VLM must ground natural-language attributes to specific objects.
 17    Uses: FewSOL questionnaire attributes.
 18
 19T3. SPATIAL GROUNDING
 20    Robot receives "pick the object left of the bowl."
 21    VLM must resolve spatial relations in the image plane.
 22    Uses: 360° walkaround (relations change per viewpoint).
 23
 24T4. STATE VERIFICATION
 25    Robot asks "did I succeed? Is the target still on the table?"
 26    VLM compares two frames: before and after manipulation.
 27    Uses: clutter → clean phase transition (unique to RPX).
 28
 29T5. COUNTING
 30    Robot asks "how many objects remain to pack?"
 31    VLM counts visible instances of a category.
 32    Uses: instance masks for GT.
 33
 34T6. INTERACTION UNDERSTANDING
 35    Learning-from-human pipeline asks "what object is the human touching?"
 36    VLM must identify the object being manipulated during interaction phase.
 37    Uses: interaction phase with human hands in frame.
 38
 39Format: 4-choice MCQ for T1-T5. Open-ended for T6 (scored by exact category match).
 40"""
 41
 42from __future__ import annotations
 43
 44import hashlib
 45import random
 46from dataclasses import dataclass, field
 47from typing import Dict, List, Optional, Tuple
 48
 49
 50@dataclass
 51class RobotQA:
 52    """One robot-grounded QA pair."""
 53
 54    scene_id: str
 55    phase: int
 56    frame_id: str
 57    task_type: str  # T1-T6
 58    robot_operation: str  # what robot pipeline stage this tests
 59    question: str
 60    correct_answer: str
 61    options: List[str]  # 4 options for MCQ; empty for open-ended
 62    correct_index: int  # -1 for open-ended
 63    reference_image: Optional[str] = None  # path to single-object reference (T1)
 64    comparison_frame: Optional[str] = None  # path to before/after frame (T4)
 65    metadata: Dict = field(default_factory=dict)
 66
 67
 68def _seed(scene_id: str, phase: int, frame_id: str, tag: str) -> int:
 69    h = hashlib.md5(f"{scene_id}/{phase}/{frame_id}/{tag}".encode()).hexdigest()
 70    return int(h[:8], 16)
 71
 72
 73def _shuffle(correct: str, distractors: List[str], seed: int) -> Tuple[List[str], int]:
 74    opts = [correct] + distractors[:3]
 75    random.Random(seed).shuffle(opts)
 76    return opts, opts.index(correct)
 77
 78
 79# ── T1: In-Context Recognition ─────────────────────────────────────────────
 80
 81
 82def generate_incontext_recognition(
 83    visible_objects: List[Dict],
 84    single_object_refs: Dict[str, str],  # {category: path_to_reference_image}
 85    absent_categories: List[str],  # categories NOT in this frame
 86    scene_id: str,
 87    phase: int,
 88    frame_id: str,
 89) -> List[RobotQA]:
 90    """T1: 'Here is a photo of the target. Is it in this scene?'
 91
 92    This is the object registration task every real robot does.
 93    The single-object scene image IS the reference photo.
 94    """
 95    pairs = []
 96    visible_cats = {o["category"] for o in visible_objects}
 97
 98    # Positive: objects that ARE visible and have a reference image
 99    for cat in list(visible_cats)[:3]:
100        if cat not in single_object_refs:
101            continue
102        s = _seed(scene_id, phase, frame_id, f"t1_yes_{cat}")
103        q = "The reference image shows an object. Is this exact object present in the current scene?"
104        opts, idx = _shuffle(
105            "yes, it is visible",
106            [
107                "no, it is not in this scene",
108                "a similar but different object is present",
109                "cannot determine from this viewpoint",
110            ],
111            s,
112        )
113        pairs.append(
114            RobotQA(
115                scene_id=scene_id,
116                phase=phase,
117                frame_id=frame_id,
118                task_type="T1_incontext",
119                robot_operation="object_registration",
120                question=q,
121                correct_answer="yes, it is visible",
122                options=opts,
123                correct_index=idx,
124                reference_image=single_object_refs[cat],
125                metadata={"category": cat, "present": True},
126            )
127        )
128
129    # Negative: objects NOT in this frame but that have reference images
130    neg_count = 0
131    for cat in absent_categories:
132        if cat not in single_object_refs or neg_count >= len(pairs):
133            continue
134        s = _seed(scene_id, phase, frame_id, f"t1_no_{cat}")
135        q = "The reference image shows an object. Is this exact object present in the current scene?"
136        opts, idx = _shuffle(
137            "no, it is not in this scene",
138            [
139                "yes, it is visible",
140                "a similar but different object is present",
141                "partially occluded but present",
142            ],
143            s,
144        )
145        pairs.append(
146            RobotQA(
147                scene_id=scene_id,
148                phase=phase,
149                frame_id=frame_id,
150                task_type="T1_incontext",
151                robot_operation="object_registration",
152                question=q,
153                correct_answer="no, it is not in this scene",
154                options=opts,
155                correct_index=idx,
156                reference_image=single_object_refs[cat],
157                metadata={"category": cat, "present": False},
158            )
159        )
160        neg_count += 1
161
162    return pairs
163
164
165# ── T2: Attribute Grounding ────────────────────────────────────────────────
166
167
168def generate_attribute_grounding(
169    visible_objects: List[Dict],
170    all_objects: List[Dict],
171    scene_id: str,
172    phase: int,
173    frame_id: str,
174) -> List[RobotQA]:
175    """T2: 'Pick the red mug' — which object matches this description?
176
177    Tests whether the VLM can distinguish objects by attributes.
178    Only generates questions when multiple objects of the same category
179    are visible (otherwise there's no disambiguation needed).
180    """
181    pairs = []
182    # Find categories with multiple visible instances
183    from collections import Counter
184
185    cat_counts = Counter(o["category"] for o in visible_objects)
186    ambiguous_cats = [c for c, n in cat_counts.items() if n >= 2]
187
188    for cat in ambiguous_cats:
189        instances = [o for o in visible_objects if o["category"] == cat]
190        for obj in instances[:2]:
191            color = obj.get("color", "")
192            if not color:
193                continue
194            s = _seed(scene_id, phase, frame_id, f"t2_{cat}_{color}")
195            q = f"A robot is instructed to pick the {color} {cat}. Which object should it grasp?"
196            gt = f"the {color} {cat}"
197            other_colors = [o.get("color", "?") for o in instances if o.get("color") != color]
198            distractors = [f"the {c} {cat}" for c in other_colors[:2]]
199            distractors.append(f"any {cat}")
200            while len(distractors) < 3:
201                distractors.append("none of the visible objects")
202            opts, idx = _shuffle(gt, distractors[:3], s)
203            pairs.append(
204                RobotQA(
205                    scene_id=scene_id,
206                    phase=phase,
207                    frame_id=frame_id,
208                    task_type="T2_attribute",
209                    robot_operation="language_grounding",
210                    question=q,
211                    correct_answer=gt,
212                    options=opts,
213                    correct_index=idx,
214                    metadata={"category": cat, "target_color": color},
215                )
216            )
217
218    return pairs
219
220
221# ── T3: Spatial Grounding ──────────────────────────────────────────────────
222
223
224def generate_spatial_grounding(
225    visible_objects: List[Dict],  # must include bbox_cx, bbox_cy
226    scene_id: str,
227    phase: int,
228    frame_id: str,
229) -> List[RobotQA]:
230    """T3: 'Pick the object to the left of the bowl.'
231
232    Relations in image plane. Left = smaller x. Above = smaller y.
233    """
234    pairs = []
235    if len(visible_objects) < 2:
236        return pairs
237
238    directions = {"left": (-1, 0), "right": (1, 0), "above": (0, -1), "below": (0, 1)}
239
240    for ref in visible_objects[:3]:
241        cx, cy = ref.get("bbox_cx"), ref.get("bbox_cy")
242        if cx is None:
243            continue
244
245        for direction, (dx_sign, dy_sign) in directions.items():
246            # Find nearest object in that direction
247            best, best_dist = None, float("inf")
248            for other in visible_objects:
249                if other is ref:
250                    continue
251                ox, oy = other.get("bbox_cx"), other.get("bbox_cy")
252                if ox is None:
253                    continue
254                # Check direction
255                if dx_sign < 0 and ox >= cx:
256                    continue
257                if dx_sign > 0 and ox <= cx:
258                    continue
259                if dy_sign < 0 and oy >= cy:
260                    continue
261                if dy_sign > 0 and oy <= cy:
262                    continue
263                d = abs(ox - cx) + abs(oy - cy)
264                if d < best_dist:
265                    best, best_dist = other, d
266
267            if best is None:
268                continue
269
270            s = _seed(scene_id, phase, frame_id, f"t3_{direction}_{ref['category']}")
271            q = f"A robot must pick the object {direction} of the {ref['category']}. What should it pick?"
272            gt = best["category"]
273            others = list(
274                {
275                    o["category"]
276                    for o in visible_objects
277                    if o["category"] != gt and o["category"] != ref["category"]
278                }
279            )
280            others += ["nothing — no object in that direction"]
281            rng = random.Random(s)
282            distractors = rng.sample(others, min(3, len(others)))
283            while len(distractors) < 3:
284                distractors.append("cannot determine")
285            opts, idx = _shuffle(gt, distractors[:3], s)
286            pairs.append(
287                RobotQA(
288                    scene_id=scene_id,
289                    phase=phase,
290                    frame_id=frame_id,
291                    task_type="T3_spatial",
292                    robot_operation="spatial_command",
293                    question=q,
294                    correct_answer=gt,
295                    options=opts,
296                    correct_index=idx,
297                    metadata={"reference": ref["category"], "direction": direction},
298                )
299            )
300            break  # one per ref per direction
301
302    return pairs
303
304
305# ── T4: State Verification ─────────────────────────────────────────────────
306
307
308def generate_state_verification(
309    clutter_objects: List[Dict],  # objects visible in clutter phase
310    clean_objects: List[Dict],  # objects visible in clean phase
311    scene_id: str,
312    frame_id: str,
313) -> List[RobotQA]:
314    """T4: 'Did the manipulation succeed? What changed?'
315
316    Compares clutter and clean phases of the same scene.
317    This is what a robot does after every pick-and-place: verify.
318    UNIQUE TO RPX — no other benchmark can generate this.
319    """
320    pairs = []
321    clutter_cats = {o["category"] for o in clutter_objects}
322    clean_cats = {o["category"] for o in clean_objects}
323
324    removed = clutter_cats - clean_cats
325    added = clean_cats - clutter_cats
326    remained = clutter_cats & clean_cats
327
328    # Question: what was removed?
329    if removed:
330        removed_obj = list(removed)[0]
331        s = _seed(scene_id, 0, frame_id, f"t4_removed_{removed_obj}")
332        q = "Comparing the before and after images: which object was removed from the scene?"
333        distractors = list(remained)[:3]
334        while len(distractors) < 3:
335            distractors.append("nothing was removed")
336        opts, idx = _shuffle(removed_obj, distractors[:3], s)
337        pairs.append(
338            RobotQA(
339                scene_id=scene_id,
340                phase=0,
341                frame_id=frame_id,
342                task_type="T4_state",
343                robot_operation="task_verification",
344                question=q,
345                correct_answer=removed_obj,
346                options=opts,
347                correct_index=idx,
348                metadata={"change_type": "removed", "object": removed_obj},
349            )
350        )
351
352    # Question: what was added?
353    if added:
354        added_obj = list(added)[0]
355        s = _seed(scene_id, 2, frame_id, f"t4_added_{added_obj}")
356        q = "Comparing the before and after images: which object was added to the scene?"
357        distractors = list(remained)[:3]
358        while len(distractors) < 3:
359            distractors.append("nothing was added")
360        opts, idx = _shuffle(added_obj, distractors[:3], s)
361        pairs.append(
362            RobotQA(
363                scene_id=scene_id,
364                phase=2,
365                frame_id=frame_id,
366                task_type="T4_state",
367                robot_operation="task_verification",
368                question=q,
369                correct_answer=added_obj,
370                options=opts,
371                correct_index=idx,
372                metadata={"change_type": "added", "object": added_obj},
373            )
374        )
375
376    # Question: did the count change?
377    if len(clutter_cats) != len(clean_cats):
378        s = _seed(scene_id, 0, frame_id, "t4_count")
379        delta = len(clean_cats) - len(clutter_cats)
380        q = "After manipulation, did the number of visible objects change?"
381        if delta < 0:
382            gt = f"decreased by {abs(delta)}"
383        elif delta > 0:
384            gt = f"increased by {delta}"
385        else:
386            gt = "stayed the same"
387        opts, idx = _shuffle(
388            gt,
389            [
390                "stayed the same" if delta != 0 else "decreased by 1",
391                f"increased by {abs(delta) + 1}" if delta <= 0 else f"decreased by {delta + 1}",
392                "cannot determine",
393            ],
394            s,
395        )
396        pairs.append(
397            RobotQA(
398                scene_id=scene_id,
399                phase=0,
400                frame_id=frame_id,
401                task_type="T4_state",
402                robot_operation="task_verification",
403                question=q,
404                correct_answer=gt,
405                options=opts,
406                correct_index=idx,
407                metadata={"change_type": "count", "delta": delta},
408            )
409        )
410
411    return pairs
412
413
414# ── T5: Counting ───────────────────────────────────────────────────────────
415
416
417def generate_counting(
418    visible_objects: List[Dict],
419    scene_id: str,
420    phase: int,
421    frame_id: str,
422) -> List[RobotQA]:
423    """T5: 'How many objects remain to pack?'"""
424    pairs = []
425    from collections import Counter
426
427    counts = Counter(o["category"] for o in visible_objects)
428
429    for cat, count in list(counts.items())[:3]:
430        s = _seed(scene_id, phase, frame_id, f"t5_{cat}")
431        q = f"A robot is packing all {cat} objects. How many does it need to pick?"
432        gt = str(count)
433        pool = [str(i) for i in range(max(count + 3, 5)) if str(i) != gt]
434        distractors = random.Random(s).sample(pool, min(3, len(pool)))
435        while len(distractors) < 3:
436            distractors.append("0")
437        opts, idx = _shuffle(gt, distractors[:3], s)
438        pairs.append(
439            RobotQA(
440                scene_id=scene_id,
441                phase=phase,
442                frame_id=frame_id,
443                task_type="T5_counting",
444                robot_operation="task_progress",
445                question=q,
446                correct_answer=gt,
447                options=opts,
448                correct_index=idx,
449                metadata={"category": cat, "count": count},
450            )
451        )
452
453    return pairs
454
455
456# ── Evaluation ──────────────────────────────────────────────────────────────
457
458
459def evaluate_robot_vqa(
460    predictions: List[int],
461    ground_truth: List[RobotQA],
462) -> Dict:
463    """MCQ accuracy with breakdowns by task type, phase, and robot operation."""
464    assert len(predictions) == len(ground_truth)
465
466    by_type: Dict[str, List[bool]] = {}
467    by_phase: Dict[int, List[bool]] = {}
468    by_operation: Dict[str, List[bool]] = {}
469
470    for pred, qa in zip(predictions, ground_truth, strict=False):
471        hit = pred == qa.correct_index
472        by_type.setdefault(qa.task_type, []).append(hit)
473        by_phase.setdefault(qa.phase, []).append(hit)
474        by_operation.setdefault(qa.robot_operation, []).append(hit)
475
476    n = len(predictions)
477    correct = sum(
478        1 for p, q in zip(predictions, ground_truth, strict=False) if p == q.correct_index
479    )
480
481    result = {
482        "accuracy": correct / n if n else 0.0,
483        "n_questions": n,
484        "per_task_type": {t: sum(v) / len(v) for t, v in by_type.items()},
485        "per_phase": {p: sum(v) / len(v) for p, v in by_phase.items()},
486        "per_robot_operation": {o: sum(v) / len(v) for o, v in by_operation.items()},
487    }
488
489    # STR
490    pa = {p: sum(v) / len(v) for p, v in by_phase.items()}
491    if 0 in pa and 1 in pa:
492        result["str_clu_to_int"] = pa[1] - pa[0]
493    if 1 in pa and 2 in pa:
494        result["str_int_to_cln"] = pa[2] - pa[1]
495
496    return result