manual-sync 2026-07-02T21:01:54Z workspace/dovla_cil/generation

workspace/dovla_cil/generation/tangent_local_atlas.py

@@ -25,6 +25,8 @@ class LocalAtlasConfig:
     neighbor_pool: int = 16
     utility_weight: float = 0.0
     diversity_weight: float = 0.0
+    negative_margin_weight: float = 0.0
+    negative_pool: int = 64
     same_task_only: bool = True
 
     def __post_init__(self) -> None:
@@ -40,6 +42,10 @@ class LocalAtlasConfig:
             raise ValueError("utility_weight must be non-negative")
         if self.diversity_weight < 0:
             raise ValueError("diversity_weight must be non-negative")
+        if self.negative_margin_weight < 0:
+            raise ValueError("negative_margin_weight must be non-negative")
+        if self.negative_pool <= 0:
+            raise ValueError("negative_pool must be positive")
 
 
 def evaluate_local_atlas_retrieval(
@@ -65,9 +71,16 @@ def evaluate_local_atlas_retrieval(
     ]
     if not train_positive:
         raise ValueError("no train positive tangents")
+    train_negative = [
+        row for row in split["train"]
+        if row.example.tangent_label == "negative"
+    ]
     train_positive_by_task: dict[str, list[SplineTangentRow]] = {}
     for row in train_positive:
         train_positive_by_task.setdefault(row.example.task_id, []).append(row)
+    train_negative_by_task: dict[str, list[SplineTangentRow]] = {}
+    for row in train_negative:
+        train_negative_by_task.setdefault(row.example.task_id, []).append(row)
 
     max_k = max(int(k) for k in k_values)
     val_by_group: dict[str, list[SplineTangentRow]] = {}
@@ -90,6 +103,8 @@ def evaluate_local_atlas_retrieval(
             group_rows[0],
             train_positive=train_positive,
             train_positive_by_task=train_positive_by_task,
+            train_negative=train_negative,
+            train_negative_by_task=train_negative_by_task,
             config=config,
             max_k=max_k,
         )
@@ -124,12 +139,17 @@ def evaluate_local_atlas_retrieval(
         "num_train_examples": len(split["train"]),
         "num_val_examples": len(split["val"]),
         "num_train_positive": len(train_positive),
+        "num_train_negative": len(train_negative),
         "num_val_groups_with_positive": len(groups),
         "label_counts": dict(Counter(row.example.tangent_label for row in rows)),
         "train_positive_by_task": {
             task_id: len(train_positive_by_task.get(task_id, []))
             for task_id in task_ids
         },
+        "train_negative_by_task": {
+            task_id: len(train_negative_by_task.get(task_id, []))
+            for task_id in task_ids
+        },
         "overall": _aggregate_rows(groups, k_values=k_values, thresholds=thresholds),
         "per_task": {
             task_id: _aggregate_rows(
@@ -148,6 +168,8 @@ def _select_local_positive_tangents(
     *,
     train_positive: Sequence[SplineTangentRow],
     train_positive_by_task: dict[str, list[SplineTangentRow]],
+    train_negative: Sequence[SplineTangentRow],
+    train_negative_by_task: dict[str, list[SplineTangentRow]],
     config: LocalAtlasConfig,
     max_k: int,
 ) -> list[SplineTangentRow]:
@@ -155,6 +177,20 @@ def _select_local_positive_tangents(
         pool = train_positive_by_task.get(query.example.task_id, []) or list(train_positive)
     else:
         pool = list(train_positive)
+    if config.same_task_only:
+        negative_pool = train_negative_by_task.get(query.example.task_id, []) or list(train_negative)
+    else:
+        negative_pool = list(train_negative)
+    local_negatives = [
+        row
+        for _, row in sorted(
+            (
+                (_condition_distance(query.condition, row.condition), row)
+                for row in negative_pool
+            ),
+            key=lambda item: (item[0], item[1].example.group_id),
+        )[: int(config.negative_pool)]
+    ]
     ranked = sorted(
         (
             (_condition_distance(query.condition, row.condition), row)
@@ -166,24 +202,40 @@ def _select_local_positive_tangents(
     selected: list[tuple[float, SplineTangentRow]] = []
     remaining = list(candidates)
     while remaining and len(selected) < max_k:
-        if not selected:
-            chosen = remaining[0]
-        else:
-            chosen = max(
-                remaining,
-                key=lambda item: (
-                    -item[0]
-                    + float(config.utility_weight) * item[1].example.delta_utility
-                    + float(config.diversity_weight)
-                    * _nearest_action_distance(item[1], [row for _, row in selected]),
-                    item[1].example.delta_utility,
+        chosen = max(
+            remaining,
+            key=lambda item: (
+                _candidate_score(
+                    item,
+                    selected=[row for _, row in selected],
+                    local_negatives=local_negatives,
+                    config=config,
                 ),
-            )
+                item[1].example.delta_utility,
+            ),
+        )
         selected.append(chosen)
         remaining = [item for item in remaining if item is not chosen]
     return [row for _, row in selected]
 
 
+def _candidate_score(
+    item: tuple[float, SplineTangentRow],
+    *,
+    selected: Sequence[SplineTangentRow],
+    local_negatives: Sequence[SplineTangentRow],
+    config: LocalAtlasConfig,
+) -> float:
+    condition_distance, row = item
+    score = -float(condition_distance)
+    score += float(config.utility_weight) * row.example.delta_utility
+    if config.diversity_weight > 0.0 and selected:
+        score += float(config.diversity_weight) * _nearest_action_distance(row, selected)
+    if config.negative_margin_weight > 0.0 and local_negatives:
+        score += float(config.negative_margin_weight) * _nearest_action_distance(row, local_negatives)
+    return score
+
+
 def _condition_distance(left: Sequence[float], right: Sequence[float]) -> float:
     return _rms_distance(left, right)