v5c: add --min-investigation-steps filter for submit-only with longer contexts

training/train_sft.py

@@ -93,6 +93,7 @@ class SftConfig:
     eval_holdout_hard: int = 3
     eval_every_epoch: bool = True
     submit_only: bool = False
+    min_investigation_steps: int = 0
 
     lr: float = 5e-5
     grad_clip: float = 1.0
@@ -220,6 +221,7 @@ def build_sft_examples(
                     "prompt_ids": prompt_ids,
                     "completion_ids": completion_ids,
                     "completion_text": completion_text,
+                    "trace_inv_steps": n_inv,
                 })
                 messages.append({"role": "assistant", "content": completion_text})
                 obs = env.step(build_action(action_dict))
@@ -313,6 +315,10 @@ def main() -> None:
     if cfg.submit_only:
         examples = [ex for ex in examples if ex["action_type"] == "submit_final_resolution"]
         print(f"[setup] submit-only mode: filtered to {len(examples)} submit examples", flush=True)
+    if cfg.min_investigation_steps > 0:
+        before = len(examples)
+        examples = [ex for ex in examples if ex.get("trace_inv_steps", 0) >= cfg.min_investigation_steps]
+        print(f"[setup] min_investigation_steps={cfg.min_investigation_steps}: {before} -> {len(examples)} examples", flush=True)
     print(f"[setup] train_tasks={len(train_tasks)} eval_tasks={len(eval_tasks)} examples={len(examples)}", flush=True)
 
     tracker = None
@@ -439,6 +445,7 @@ def _parse_args() -> SftConfig:
     p.add_argument("--no-push", action="store_true")
     p.add_argument("--no-4bit", action="store_true")
     p.add_argument("--submit-only", action="store_true")
+    p.add_argument("--min-investigation-steps", type=int, default=None)
     args = p.parse_args()
 
     cfg = SftConfig()
@@ -466,6 +473,8 @@ def _parse_args() -> SftConfig:
         cfg.use_4bit = False
     if args.submit_only:
         cfg.submit_only = True
+    if args.min_investigation_steps is not None:
+        cfg.min_investigation_steps = args.min_investigation_steps
     return cfg