Auto-sync: 2026-06-28 23:11:20 (part 3)

scripts/slurm/smoke_retrieval_metric_unit.sbatch

@@ -171,10 +171,83 @@ case_task_relative = _RolloutCase(
     retrieval_metric="task_relative",
 )
 assert task_relative_attached.candidate_source_group_id == "train_actor_match", task_relative_attached
+
+def record_progress(group_id, candidate_type, action_value, progress, feature):
+    return SimpleNamespace(
+        group_id=group_id,
+        task_id="PickCube-v1",
+        candidate_type=candidate_type,
+        record_id=f"{group_id}-{candidate_type}-{action_value}",
+        observation_inline={"features": feature},
+        reward=SimpleNamespace(progress=progress, terminal_success=progress >= 1.0),
+        action_chunk=ActionChunk(
+            representation="continuous",
+            horizon=1,
+            values=[[action_value, 0.0]],
+        ),
+    )
+
+groups_progress = {
+    "train_a": [
+        record_progress("train_a", "expert", 1.0, 1.0, [0.0, 0.0]),
+        record_progress("train_a", "no_op", 1.2, 0.8, [0.0, 0.0]),
+        record_progress("train_a", "wrong_gripper", 1.4, 0.2, [0.0, 0.0]),
+    ],
+    "train_b": [
+        record_progress("train_b", "expert", 2.0, 1.0, [10.0, 0.0]),
+        record_progress("train_b", "no_op", 2.2, 0.8, [10.0, 0.0]),
+        record_progress("train_b", "wrong_gripper", 2.4, 0.8, [10.0, 0.0]),
+    ],
+    "heldout_a": [record_progress("heldout_a", "expert", 9.0, 1.0, [0.0, 0.0])],
+    "heldout_b": [record_progress("heldout_b", "expert", 9.0, 1.0, [10.0, 0.0])],
+}
+dataset_progress = SimpleNamespace(
+    group_ids=list(groups_progress),
+    get_group=lambda group_id: groups_progress[group_id],
+)
+progress_cases = [
+    _RolloutCase(
+        group_id=group_id,
+        task_id="PickCube-v1",
+        source_dataset=Path("."),
+        state={},
+        observation={"features": feature},
+        instruction="pick",
+        oracle_score=1.0,
+        oracle_success=True,
+        expert_score=1.0,
+        expert_success=True,
+        best_action_values=[[9.9, 0.0]],
+        candidate_action_values=[],
+        candidate_types=[],
+    )
+    for group_id, feature in [("heldout_a", [0.0, 0.0]), ("heldout_b", [10.0, 0.0])]
+]
+progress_attached = _attach_retrieved_residual_candidates(
+    dataset_progress,
+    progress_cases,
+    heldout_group_ids=["heldout_a", "heldout_b"],
+    obs_dim=2,
+    observation_mode="state",
+    retrieval_neighbors=1,
+    retrieval_residual_min_source_progress=0.5,
+)
+assert [len(case.candidate_action_values) for case in progress_attached] == [3, 3], progress_attached
+assert progress_attached[0].candidate_types == [
+    "policy_residual",
+    "residual_no_op",
+    "policy_residual",
+], progress_attached[0].candidate_types
+assert progress_attached[1].candidate_types == [
+    "policy_residual",
+    "residual_no_op",
+    "residual_wrong_gripper",
+], progress_attached[1].candidate_types
 print({
     "status": "ok",
     "raw": raw_attached.candidate_source_group_id,
     "zscore": zscore_attached.candidate_source_group_id,
     "task_relative": task_relative_attached.candidate_source_group_id,
+    "source_progress_lengths": [len(case.candidate_action_values) for case in progress_attached],
 })
 PY