Auto-sync: 2026-06-28 22:04:20 (part 4)

scripts/export_retrieval_residual_policy_targets.py

@@ -42,7 +42,11 @@ def main(argv: list[str] | None = None) -> int:
     parser.add_argument("--device", default="auto")
     parser.add_argument("--split", choices=("train", "val", "all"), default="all")
     parser.add_argument("--retrieval-neighbors", type=int, default=1)
-    parser.add_argument("--retrieval-metric", choices=("raw", "zscore"), default="raw")
+    parser.add_argument(
+        "--retrieval-metric",
+        choices=("raw", "zscore", "task_relative"),
+        default="raw",
+    )
     parser.add_argument("--retrieval-residual-scale", type=float, default=0.35)
     parser.add_argument(
         "--exclude-types",

scripts/slurm/smoke_retrieval_metric_unit.sbatch

@@ -119,5 +119,62 @@ assert zscore_attached.candidate_source_group_id == "train_b", zscore_attached
 expected = np.asarray([[[0.0, 0.0]], [[0.2, 0.0]]], dtype=np.float32)
 actual = np.asarray(zscore_attached.candidate_action_values, dtype=np.float32)
 assert np.allclose(actual, expected), actual
-print({"status": "ok", "raw": raw_attached.candidate_source_group_id, "zscore": zscore_attached.candidate_source_group_id})
+
+def actor_feature(target_x, robot_tail=0.0):
+    values = [0.0] * 70
+    values[0] = target_x
+    values[3] = 1.0
+    values[16] = 1.0
+    values[-1] = robot_tail
+    return values
+
+groups_task_relative = {
+    "train_actor_far": [
+        record("train_actor_far", "expert", 1.0, actor_feature(0.6)),
+        record("train_actor_far", "near_miss", 1.1, actor_feature(0.6)),
+    ],
+    "train_actor_match": [
+        record("train_actor_match", "expert", 2.0, actor_feature(0.0, robot_tail=5.0)),
+        record("train_actor_match", "near_miss", 2.2, actor_feature(0.0, robot_tail=5.0)),
+    ],
+    "heldout": [
+        record("heldout", "expert", 9.0, actor_feature(0.0)),
+        record("heldout", "near_miss", 9.9, actor_feature(0.0)),
+    ],
+}
+dataset_task_relative = SimpleNamespace(
+    group_ids=list(groups_task_relative),
+    get_group=lambda group_id: groups_task_relative[group_id],
+)
+case_task_relative = _RolloutCase(
+    group_id="heldout",
+    task_id="PickCube-v1",
+    source_dataset=Path("."),
+    state={},
+    observation={"features": actor_feature(0.0)},
+    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=[],
+)
+[task_relative_attached] = _attach_retrieved_residual_candidates(
+    dataset_task_relative,
+    [case_task_relative],
+    heldout_group_ids=["heldout"],
+    obs_dim=70,
+    observation_mode="state",
+    retrieval_neighbors=1,
+    retrieval_metric="task_relative",
+)
+assert task_relative_attached.candidate_source_group_id == "train_actor_match", task_relative_attached
+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,
+})
 PY

tests/test_maniskill_policy_rollout.py

@@ -919,6 +919,100 @@ def test_retrieval_residual_zscore_metric_standardizes_train_bank_features() ->
     )
 
 
+def test_retrieval_residual_task_relative_metric_ignores_robot_tail_noise() -> None:
+    def feature(*, target_x: float, robot_tail: float = 0.0) -> list[float]:
+        values = [0.0] * 70
+        values[0] = target_x
+        values[3] = 1.0
+        values[13 + 3] = 1.0
+        values[-1] = robot_tail
+        return values
+
+    def record(group_id: str, candidate_type: str, action_value: float, obs: list[float]):
+        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": obs},
+            action_chunk=ActionChunk(
+                representation="continuous",
+                horizon=1,
+                values=[[action_value, 0.0]],
+            ),
+        )
+
+    groups = {
+        "train_actor_far": [
+            record("train_actor_far", "expert", 1.0, feature(target_x=0.6)),
+            record("train_actor_far", "near_miss", 1.1, feature(target_x=0.6)),
+        ],
+        "train_actor_match": [
+            record(
+                "train_actor_match",
+                "expert",
+                2.0,
+                feature(target_x=0.0, robot_tail=5.0),
+            ),
+            record(
+                "train_actor_match",
+                "near_miss",
+                2.2,
+                feature(target_x=0.0, robot_tail=5.0),
+            ),
+        ],
+        "heldout": [
+            record("heldout", "expert", 9.0, feature(target_x=0.0)),
+            record("heldout", "near_miss", 9.9, feature(target_x=0.0)),
+        ],
+    }
+    dataset = SimpleNamespace(
+        group_ids=list(groups),
+        get_group=lambda group_id: groups[group_id],
+    )
+    case = _RolloutCase(
+        group_id="heldout",
+        task_id="PickCube-v1",
+        source_dataset=Path("."),
+        state={},
+        observation={"features": feature(target_x=0.0)},
+        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=[],
+    )
+
+    [raw_attached] = _attach_retrieved_residual_candidates(
+        dataset,
+        [case],
+        heldout_group_ids=["heldout"],
+        obs_dim=70,
+        observation_mode="state",
+        retrieval_neighbors=1,
+        retrieval_metric="raw",
+    )
+    [task_relative_attached] = _attach_retrieved_residual_candidates(
+        dataset,
+        [case],
+        heldout_group_ids=["heldout"],
+        obs_dim=70,
+        observation_mode="state",
+        retrieval_neighbors=1,
+        retrieval_metric="task_relative",
+    )
+
+    assert raw_attached.candidate_source_group_id == "train_actor_far"
+    assert task_relative_attached.candidate_source_group_id == "train_actor_match"
+    assert np.allclose(
+        np.asarray(task_relative_attached.candidate_action_values, dtype=np.float32),
+        np.asarray([[[0.0, 0.0]], [[0.2, 0.0]]]),
+    )
+
+
 def test_retrieval_residual_type_success_threshold_filters_train_families() -> None:
     def record(
         group_id: str,