fix(fit): map predict_proba columns to label ids so the RSB residual stage routes correctly

[breken-ai]

What changed

_predict() in src/tracer/fit/pipeline.py returned probs.argmax(axis=1) as the predicted label. That is the column index of predict_proba, which follows clf.classes_. It equals the label id only when the surrogate was trained on every label.

The RSB residual stage (build_rsb, stage 2) is trained only on the rows stage 1 rejected. That pool usually lacks some labels. For example, an easy class that stage 1 absorbs completely is missing. Then classes_ is something like [1, 2, 3], and column 0 is reported as label 0 instead of label 1. The fix maps the argmax through clf.classes_:

idx = probs.argmax(axis=1)
classes = getattr(clf, "classes_", None)
preds = (np.asarray(classes)[idx] if classes is not None else idx).astype(int)

Stage 1 and global always see every label (_split_buffer puts each class in train), so their outputs don't change. Only a stage trained on a label subset behaves differently.

Why it matters

Calibration and routing both go through _predict/apply_stage. The stage-2 surrogate's predictions therefore get compared with the teacher under the wrong labels. On the reject pool below, stage-2 agreement measured via _predict is 0.13 on main, compared with 0.71 for the same model's clf.predict. The parity gate then almost always refuses stage 2, so RSB quietly loses the coverage it exists to add. If stage 2 does get certified, the router emits shifted labels for the classes affected.

Benchmark (fixed seed, public API, same data, main vs this branch)

The data is synthetic and uses 4 classes, 4,000 traces, 6-dim embeddings, RandomState(0). Class 0 is trivially separable, and classes 1-3 have an "easy" group and a "hard" group that a single linear surrogate can't resolve. The generator is _residual_traces in the new test. Command:

fit(path, embeddings=X, config=FitConfig(target_teacher_agreement=0.95, frontier_targets=(0.95,),
    skip_candidates=("gbt","rf","et","xgb","dt","mlp_1h","mlp_2h"), verbose=False))
router = load_router(out); router.predict_batch(X)

	selected	coverage_cal	runtime coverage (all 4,000)	teacher agreement on handled
main @ 093792a	l2d (RSB stage 2 refused)	0.855	0.826	0.972
this branch	rsb (2 stages)	0.984	0.976	0.974

Agreement holds at target while coverage rises by about 15 points. The gate, thresholds and defaults are unchanged. This only corrects which label a stage reports.

Tests

The new file is tests/test_stage_labels.py:

• test_predict_maps_proba_columns_to_class_labels trains a surrogate on labels {1, 3} and checks that _predict/apply_stage return {1, 3}. On main it fails with {0, 1} != {1, 3}.
• test_rsb_residual_stage_routes_with_teacher_labels runs the end-to-end fit → load_router → predict_batch flow on the data above. It checks that RSB is selected and that stage-2 rows carry their teacher label. On main it fails with 'l2d' == 'rsb'.

pytest tests/: 129 passed, 3 skipped (Python 3.11). Both new tests fail on main and pass here.

---

I found this bug and prepared the fix with AI assistance (Breken). I checked the code path, test output and benchmark numbers above by running them locally.