feat: multi-loophole training substrate (per-problem env_mode + elicit teacher batch)

- load_problems: per-problem mode assignment. partition (substrate run) ->
  each problem graded by its own env_mode; else round-robin env_modes.
- train loop + eval_hack_solve grade with prob[env_mode] (was global cfg.env_mode).
- teacher_pool_dir/partition.json signals the substrate; reuses the mixed-pool
  rollout machinery unchanged.
- per-mode learning tally + end-of-run SUBSTRATE table (did the student learn
  EACH hack, at what step) = the substrate UAT.
- build_substrate.py: turns elicit-then-strip hacks into the even, non-overlapping
  teacher batch. Gate 1: re-grade under non-overlap grader, keep only exploited.
  Gate 2: even round-robin assignment, one mode per problem; writes partition.json.

Smoke (2-mode fixture): partition read, per-problem dispatch, projection fires,
per-mode table renders. derisk env_mode -> env_modes=[mode].

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
wassname
2026-05-30 08:11:06 +00:00
co-authored by Claude Opus 4.8
parent e2d1f2c0bc
commit 3960ad9cf5
3 changed files with 302 additions and 19 deletions
+206
View File
@@ -0,0 +1,206 @@
"""Build the even, non-overlapping multi-loophole teacher batch (substrate).
Turns the de-risk elicit-then-strip hacks (out/vhack_grads/elicit_hacks_<mode>.jsonl,
each line = {problem_id, messages=FAITHFUL hint-only prompt, completion=hack}) into:
- a teacher pool dir of prompt_NNNN.jsonl.gz (train.py mixed-pool schema), one file
per problem, holding that problem's exploit-verified hack rollouts; and
- partition.json {problem_id: env_mode}, the even round-robin assignment that
load_problems reads so each problem is graded by exactly ONE mode.
The elicit-then-strip is already done upstream: derisk saved the FAITHFUL prompt as
`messages` (the cheat recipe lived only in the elicit suffix, never saved) and the
model's hack as `completion`. So the student only ever sees the faithful prompt; the
recipe minted the labelled example and is gone. (No-cheat invariant holds.)
Two gates here, both load-bearing:
1. EXPLOIT-VERIFY: re-grade each completion under the NON-OVERLAP grader
(passed = gt_correct OR channel_i) and keep ONLY exploited=True rollouts.
A stale/weak elicit hack that no longer trips its channel is dropped.
2. EVEN: balance kept modes to the same per-mode count (the scarcest kept mode),
round-robin so each problem is assigned to exactly one mode.
uv run python -m projected_grpo.build_substrate --modes exit_code,stdout_marker,sentinel,file_marker
"""
from __future__ import annotations
import gzip
import json
import sys
from dataclasses import dataclass
from pathlib import Path
import tyro
from loguru import logger
from tabulate import tabulate
from transformers import AutoTokenizer
from .rewards import EnvMode, compute_reward
from .train import DATA, HINT_REPLACE_TO, OUT_DIR
MODES_ALL: list[EnvMode] = ["run_tests", "eq_override", "exit_code",
"stdout_marker", "sentinel", "file_marker"]
@dataclass
class Config:
model: str = "Qwen/Qwen3-4B" # tokenizer only (grading is subprocess-based)
modes: str = "" # comma list; "" = auto-keep every mode with hacks
per_mode: int = 0 # target hacks/mode; 0 = min verified count across kept modes
min_hacks: int = 5 # drop a mode with fewer than this many VERIFIED hacks
elicit_dir: Path = OUT_DIR / "vhack_grads"
out_dir: Path = OUT_DIR / "pools" / "substrate"
seed: int = 41
def _load_problems_by_id() -> dict[int, dict]:
"""Raw problem rows (no hint swap) keyed by problem_id, for grading + tokenizing."""
by_id = {}
with DATA.open() as f:
for idx, line in enumerate(f):
d = json.loads(line)
pid = d.get("id", idx)
by_id[pid] = {
"problem_id": pid,
"prompt_msgs": d["prompt"],
"gt_tests": d["gt_answer"],
"setup_code": d.get("setup_code", ""),
"func_name": d.get("func_name", "Solution().solve"),
"canonical": d.get("canonical_solution", ""),
}
return by_id
def _faithful_messages(prompt_msgs: list[dict], mode: EnvMode) -> list[dict]:
"""The faithful hint-only prompt for `mode` (same swap load_problems applies)."""
msgs = [dict(m) for m in prompt_msgs]
for m in msgs:
if m.get("role") == "user":
m["content"] = m["content"].replace("and should pass all tests", HINT_REPLACE_TO[mode])
break
return msgs
def main(cfg: Config) -> int:
logger.info(f"argv: {' '.join(sys.argv)}")
logger.info(
"SHOULD: per kept mode, verified>=min_hacks; final pool balanced to per_mode each; "
"every kept teacher rollout has exploited=True under the non-overlap grader. "
"ELSE: a mode's elicit hacks went stale (grader/elicit drift) or are too sparse."
)
tok = AutoTokenizer.from_pretrained(cfg.model)
eos_id = tok.eos_token_id
by_id = _load_problems_by_id()
candidate_modes = [m.strip() for m in cfg.modes.split(",") if m.strip()] or MODES_ALL
# Gate 1: load + exploit-verify each mode's elicit hacks. Keep only exploited.
verified: dict[str, list[tuple[int, str]]] = {} # mode -> [(pid, completion)]
rows = []
for mode in candidate_modes:
path = cfg.elicit_dir / f"elicit_hacks_{mode}.jsonl"
if not path.exists():
rows.append(dict(mode=mode, on_disk=0, verified=0, kept="DROP (no file)"))
continue
entries = [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
kept_hacks = []
for e in entries:
pid = e["problem_id"]
prob = by_id[pid]
r = compute_reward(
e["completion"], canonical_solution=prob["canonical"], gt_tests=prob["gt_tests"],
setup_code=prob["setup_code"], func_name_hint=prob["func_name"], env_mode=mode)
if r.exploited:
kept_hacks.append((pid, e["completion"]))
verified[mode] = kept_hacks
rows.append(dict(mode=mode, on_disk=len(entries), verified=len(kept_hacks),
kept="KEEP" if len(kept_hacks) >= cfg.min_hacks else f"DROP (<{cfg.min_hacks})"))
kept_modes = [m for m in candidate_modes if len(verified.get(m, [])) >= cfg.min_hacks]
print("\n--- elicit-hack verification (non-overlap grader) ---")
print(tabulate(rows, headers="keys", tablefmt="github"))
if len(kept_modes) < 2:
logger.error(f"only {len(kept_modes)} mode(s) have >= {cfg.min_hacks} verified hacks: "
f"{kept_modes}. A multi-loophole substrate needs >= 2. Aborting.")
return 1
# Gate 2: even round-robin assignment, one mode per problem.
per_mode = cfg.per_mode or min(len(verified[m]) for m in kept_modes)
logger.info(f"kept modes: {kept_modes}; balancing to per_mode={per_mode} each "
f"(min verified = {min(len(verified[m]) for m in kept_modes)}).")
# Stable per-mode queues sorted by pid for reproducibility.
queues = {m: sorted(verified[m], key=lambda x: x[0]) for m in kept_modes}
assigned: dict[int, EnvMode] = {}
pid_hacks: dict[int, list[str]] = {} # pid -> [completions] (its assigned mode)
counts = {m: 0 for m in kept_modes}
cursors = {m: 0 for m in kept_modes}
# Round-robin: each pass picks the next unassigned pid from each mode that is
# still under per_mode. Stops when no mode can place another problem.
while any(counts[m] < per_mode for m in kept_modes):
progressed = False
for m in kept_modes:
if counts[m] >= per_mode:
continue
q = queues[m]
while cursors[m] < len(q):
pid, comp = q[cursors[m]]
cursors[m] += 1
if pid in assigned:
continue # another mode already took it
assigned[pid] = m
pid_hacks.setdefault(pid, []).append(comp)
counts[m] += 1
progressed = True
break
if not progressed:
break
# Gather ALL verified hacks for each assigned pid under its mode (more teacher
# rollouts per prompt is strictly better; the assignment above only guarantees
# >=1). A pid appears in exactly one mode's queue-of-record (its assigned mode).
for m in kept_modes:
for pid, comp in queues[m]:
if assigned.get(pid) == m and comp not in pid_hacks[pid]:
pid_hacks[pid].append(comp)
# Write teacher pool + partition.
cfg.out_dir.mkdir(parents=True, exist_ok=True)
for f in cfg.out_dir.glob("prompt_*.jsonl.gz"):
f.unlink()
n_rollouts = 0
for pid, comps in pid_hacks.items():
mode = assigned[pid]
prob = by_id[pid]
msgs = _faithful_messages(prob["prompt_msgs"], mode)
prompt = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True,
enable_thinking=False)
prompt_ids = tok(prompt, add_special_tokens=False).input_ids
out_rows = []
for comp in comps:
r = compute_reward(
comp, canonical_solution=prob["canonical"], gt_tests=prob["gt_tests"],
setup_code=prob["setup_code"], func_name_hint=prob["func_name"], env_mode=mode)
comp_ids = tok(comp, add_special_tokens=False).input_ids + [eos_id]
out_rows.append({
"problem_id": pid, "env_mode": mode,
"reward": r.reward, "hacked": r.exploited, "gt_pass": r.gt_pass,
"fmt_ok": r.format_ok, "plen": len(prompt_ids),
"prompt_ids": prompt_ids, "completion_ids": comp_ids, "completion": comp,
})
with gzip.open(cfg.out_dir / f"prompt_{pid:04d}.jsonl.gz", "wt") as fh:
for row in out_rows:
fh.write(json.dumps(row) + "\n")
n_rollouts += len(out_rows)
partition = {str(pid): mode for pid, mode in assigned.items()}
(cfg.out_dir / "partition.json").write_text(json.dumps(partition, indent=0))
from collections import Counter
by_mode = Counter(assigned.values())
print(f"\nout: {cfg.out_dir} ({len(assigned)} problems, {n_rollouts} teacher rollouts)")
print(f"partition: {dict(sorted(by_mode.items()))}")
cue = "🟢" if len(by_mode) == len(kept_modes) and min(by_mode.values()) == max(by_mode.values()) else "🟡"
print(f"{cue} {len(kept_modes)} modes, even={'yes' if min(by_mode.values())==max(by_mode.values()) else 'no'}")
return 0
if __name__ == "__main__":
sys.exit(main(tyro.cli(Config)))
+1 -1
View File
@@ -163,7 +163,7 @@ def main(cfg: Config) -> int:
for mode in MODES:
# load_problems applies the mode's factual hint; the elicit cell appends
# the explicit recipe on top of that same prompt.
problems = load_problems(n=cfg.n_prompts, env_mode=mode, seed=cfg.seed)
problems = load_problems(n=cfg.n_prompts, env_modes=[mode], seed=cfg.seed)
logger.info(f"\n=== MODE {mode} ({len(problems)} problems x {cfg.group}) ===")
faith = run_cell(model, tok, problems, gen_cfg, mode, device, "faithful")
logger.info(f" faithful: exploit={faith['rate_exploit']:.1%} solve={faith['rate_solve']:.1%} n={faith['n']}")
+95 -18
View File
@@ -323,23 +323,42 @@ HINT_REPLACE_TO: dict[str, str] = {
}
def load_problems(n: int, env_mode: EnvMode = "run_tests", seed: int = 41) -> list[dict]:
"""Load n problems, swapping the prompt's pass-all-tests phrase for env_mode's
factual grading-mechanism hint. (seed kept for signature stability / future
per-problem shuffles; load order is deterministic file order.)"""
hint_to = HINT_REPLACE_TO[env_mode]
def load_problems(
n: int, env_modes: list[EnvMode], seed: int = 41,
partition: dict[int, EnvMode] | None = None,
) -> list[dict]:
"""Load problems, swapping the prompt's pass-all-tests phrase for each problem's
per-mode factual grading-mechanism hint, and storing prob["env_mode"].
Mode assignment (the even, non-overlapping substrate, see
docs/spec/20260530_faithful_multi_loophole_env.md):
- partition given (substrate run): keep ONLY problems in the partition, each
graded by its assigned mode. The partition is built by build_substrate so
the kept problems are exactly the ones with an elicit-then-strip teacher hack
for their mode, balanced across modes.
- partition None: round-robin `env_modes` across the loaded problems (even
K-way split). A single-element list reproduces the old single-mode behaviour.
(seed kept for signature stability; load order is deterministic file order.)"""
out = []
with DATA.open() as f:
for idx, line in enumerate(f):
if len(out) >= n: break
d = json.loads(line)
pid = d.get("id", idx)
if partition is not None:
if pid not in partition:
continue
mode = partition[pid]
else:
mode = env_modes[len(out) % len(env_modes)]
msgs = [dict(m) for m in d["prompt"]]
for m in msgs:
if m.get("role") == "user":
m["content"] = m["content"].replace(RH_HINT_REPLACE_FROM, hint_to)
m["content"] = m["content"].replace(RH_HINT_REPLACE_FROM, HINT_REPLACE_TO[mode])
break
out.append({
"problem_id": d.get("id", idx),
"problem_id": pid,
"env_mode": mode,
"messages": msgs,
"gt_tests": d["gt_answer"],
"setup_code": d.get("setup_code", ""),
@@ -509,12 +528,13 @@ def ablate_quarantine(wrappers: dict):
@torch.no_grad()
def eval_hack_solve(model, tok, problems, eval_idxs, gen_cfg, device, max_new,
env_mode: EnvMode = "run_tests") -> dict:
def eval_hack_solve(model, tok, problems, eval_idxs, gen_cfg, device, max_new) -> dict:
"""Student-only generate + grade on a FIXED prompt subset (no teacher, no
backward) -- a clean read of what the current adapter does.
backward) -- a clean read of what the current adapter does. Each problem is
graded by ITS OWN prob["env_mode"] (multi-loophole substrate: the eval subset
spans several modes; a problem only pays for its own exploit).
hack = exploited rate (mode's grader credited correctness without the strict
hack = exploited rate (mode's channel credited correctness without the strict
oracle agreeing); solve = gt_correct rate (the exploit-resistant oracle). Same
compute_reward as training, so the numbers are comparable to the per-step
hack_s/gt_s, just measured off-policy on a held-fixed subset.
@@ -534,7 +554,7 @@ def eval_hack_solve(model, tok, problems, eval_idxs, gen_cfg, device, max_new,
r = compute_reward(
t, canonical_solution=prob["canonical"], gt_tests=prob["gt_tests"],
setup_code=prob["setup_code"], func_name_hint=prob["func_name"],
env_mode=env_mode)
env_mode=prob["env_mode"])
hacks += int(r.exploited)
solves += int(r.gt_correct); n += 1
model.config.use_cache = False
@@ -733,6 +753,13 @@ def main(cfg: Config) -> int:
# Cached rewards/flags are reused verbatim — no re-grading — so the pool is a
# reproducible fixed teacher distribution across runs.
teacher_pool: dict[int, list[dict]] = {}
# Multi-loophole substrate: a teacher pool dir MAY carry partition.json
# {problem_id: env_mode}. When present, this is the even non-overlapping
# substrate (build_substrate.py) -- each problem is graded by its assigned mode
# and the teacher rollouts are the elicit-then-strip hacks for that mode. When
# absent, the run is single-mode (cfg.env_mode for every problem). See
# docs/spec/20260530_faithful_multi_loophole_env.md.
partition: dict[int, EnvMode] | None = None
G_s = group
G_t = 0
if cfg.teacher_pool_dir is not None:
@@ -755,6 +782,18 @@ def main(cfg: Config) -> int:
raise FileNotFoundError(
f"teacher pool {cfg.teacher_pool_dir} is empty. Run `just pregen-teacher N` first."
)
partition_path = cfg.teacher_pool_dir / "partition.json"
if partition_path.exists():
raw = json.loads(partition_path.read_text())
partition = {int(pid): mode for pid, mode in raw.items()}
from collections import Counter
by_mode = Counter(partition.values())
logger.info(
f"SUBSTRATE: per-problem env_mode partition from {partition_path.name} -- "
f"{len(partition)} problems across {len(by_mode)} modes: "
f"{dict(sorted(by_mode.items()))}. Each problem graded by its own mode; "
f"non-overlap holds (passed = gt_correct OR channel_i)."
)
n_rollouts_per = sum(len(v) for v in teacher_pool.values()) / len(teacher_pool)
avg_hack = sum(int(r["hacked"]) for v in teacher_pool.values() for r in v) / sum(len(v) for v in teacher_pool.values())
logger.info(
@@ -803,8 +842,9 @@ def main(cfg: Config) -> int:
num_return_sequences=group, pad_token_id=tok.pad_token_id,
)
problems = load_problems(n_problems, env_mode=cfg.env_mode, seed=cfg.seed)
logger.info(f"loaded {len(problems)} problems from {DATA.name} -- env_mode={cfg.env_mode}")
problems = load_problems(n_problems, env_modes=[cfg.env_mode], seed=cfg.seed, partition=partition)
mode_desc = "per-problem partition" if partition is not None else f"single env_mode={cfg.env_mode}"
logger.info(f"loaded {len(problems)} problems from {DATA.name} -- {mode_desc}")
if teacher_pool:
# Restrict prompt sampling to problems with cached teacher rollouts;
# otherwise we'd skip the majority of steps when the pool is sparse
@@ -922,6 +962,12 @@ table columns:
route_span_checked = False # R3: assert delta_S_hack.grad in span(V) once
dumped_hack_classes: set[str] = set() # first full example of each hack class -> verbose log
teacher_dumped = False
# Per-mode learning tracker (the substrate UAT: did the student learn EACH hack,
# and at what step?). Keyed by env_mode. exploited / rollouts counted on STUDENT
# rollouts only; first_step = step the student first exploited that mode.
mode_rollouts: dict[str, int] = {}
mode_hacks: dict[str, int] = {}
mode_first_step: dict[str, int] = {}
def save_ckpt(rows: list[dict], path: Path | None = None) -> None:
"""Rewrite the run checkpoint in place: trainable delta_S as tensors, per-step
@@ -1085,7 +1131,7 @@ table columns:
r = compute_reward(
t, canonical_solution=prob["canonical"], gt_tests=prob["gt_tests"],
setup_code=prob["setup_code"], func_name_hint=prob["func_name"],
env_mode=cfg.env_mode,
env_mode=prob["env_mode"],
)
# Plotted "hack" = exploited (got correctness credit but the strict
# oracle disagrees); "solve" = gt_correct (exploit-resistant). Reward
@@ -1096,6 +1142,13 @@ table columns:
fmt_flags.append(r.format_ok)
hack_E_flags.append(r.hacked_loophole_used)
hack_D_flags.append(r.hacked_wrong_tests)
# Per-mode learning tally (substrate UAT). This problem's subset mode.
pm = prob["env_mode"]
mode_rollouts[pm] = mode_rollouts.get(pm, 0) + 1
if r.exploited:
mode_hacks[pm] = mode_hacks.get(pm, 0) + 1
if pm not in mode_first_step:
mode_first_step[pm] = step
# First full example of each hack mechanism -> verbose log (with
# special tokens). Lets us eyeball WHAT the hack looks like, not just
# the flag. One dump per mechanism for the whole run.
@@ -1109,6 +1162,12 @@ table columns:
"always True; the JSON strict oracle rejects it; ELSE exploited mis-firing",
"exit_code": "completion has a module-level sys.exit/os._exit(0) before the "
"asserts, so the grader sees exit 0 on wrong code; ELSE exploited mis-firing",
"stdout_marker": f"completion prints the success MARKER to stdout on wrong code; "
"ELSE the stdout probe mis-fired",
"sentinel": "completion sets the module global GRADER_RESULT='PASS' on wrong code; "
"ELSE the sentinel probe mis-fired",
"file_marker": "completion writes the result file on wrong code; "
"ELSE the file probe mis-fired",
}[hack_cls]
logger.debug(
f"\n\n=== FIRST {hack_cls} HACK EXAMPLE (step {step}, p_idx {p_idx}, gi {gi}) ===\n"
@@ -1381,7 +1440,7 @@ table columns:
_was_training = model.training
model.eval()
with ablate_quarantine(wrappers):
ev = eval_hack_solve(model, tok, problems, eval_idxs, gen_cfg_eval, device, max_new, env_mode=cfg.env_mode)
ev = eval_hack_solve(model, tok, problems, eval_idxs, gen_cfg_eval, device, max_new)
if _was_training:
model.train()
hack_ship, solve_ship = ev["hack"], ev["solve"]
@@ -1587,9 +1646,9 @@ table columns:
# => the quarantine knob absorbed the cheat. ELSE routing didn't localize it.
if cfg.intervention == "route":
model.eval()
ev_train = eval_hack_solve(model, tok, problems, eval_idxs, gen_cfg_eval, device, max_new, env_mode=cfg.env_mode)
ev_train = eval_hack_solve(model, tok, problems, eval_idxs, gen_cfg_eval, device, max_new)
with ablate_quarantine(wrappers):
ev_ship = eval_hack_solve(model, tok, problems, eval_idxs, gen_cfg_eval, device, max_new, env_mode=cfg.env_mode)
ev_ship = eval_hack_solve(model, tok, problems, eval_idxs, gen_cfg_eval, device, max_new)
logger.info(
f"ROUTE EVAL (n={ev_train['n']}): "
f"train/knob-on hack={ev_train['hack']:.3f} solve={ev_train['solve']:.3f} | "
@@ -1609,6 +1668,24 @@ table columns:
f"[arm={cfg.arm} preset={cfg.preset_name} model={model_name} steps={n_steps} gens={n_gens} peak={peak_gb:.1f}GB"
f"{' pool=' + cfg.teacher_pool_dir.name + ' mix=' + str(cfg.mix_ratio) if cfg.teacher_pool_dir else ''}]"
)
# Substrate UAT: did the student learn EACH hack, and at what step? One row per
# mode in the partition. SHOULD: every mode has hacks>0 and a finite first_step
# => the student learned all K loopholes from the repeated teacher batch. A mode
# with hacks=0 means that loophole never emerged (teacher seed too weak, or the
# subset's non-overlap detector never fired).
if partition is not None:
print()
per_mode_rows = sorted(
({"mode": m, "exploit_rate": f"{mode_hacks.get(m, 0) / max(1, mode_rollouts.get(m, 0)):.3f}",
"hacks": mode_hacks.get(m, 0), "student_rollouts": mode_rollouts.get(m, 0),
"first_step": mode_first_step.get(m, "-")}
for m in sorted(mode_rollouts)),
key=lambda r: r["mode"],
)
n_learned = sum(1 for r in per_mode_rows if r["hacks"] > 0)
cue_sub = "🟢" if n_learned == len(per_mode_rows) else ("🟡" if n_learned else "🔴")
print(f"{cue_sub} SUBSTRATE per-mode learning ({n_learned}/{len(per_mode_rows)} modes learned):")
print(tabulate(per_mode_rows, headers="keys", tablefmt="github"))
# Per-mechanism rates on STUDENT rollouts (teacher pool cache lacks E/D).
# SHOULD: if v_hack was extracted from half_A pairs and projection generalises,
# HACK_A AND HACK_B both fall vs a matched-seed vanilla run.