270 lines
13 KiB
Python
270 lines
13 KiB
Python
"""State-tracking "register machine" environment -- the minimal pure task with RE-REFERENCED,
|
|
MUTATING state, the causal ingredient that the arith-reach (24-game) task lacks and that the
|
|
Fable 5 / Mythos 5 system card garble requires.
|
|
|
|
R registers start at given values; T read-modify-write operations are applied in order (all mod M);
|
|
the model reports the final value of one register. Isomorphic to FreeCell's causal core (a handful
|
|
of mutating state atoms re-read across a long horizon) but stripped of all confounds: pure integer
|
|
arithmetic, deterministic, trivially verifiable (a short integer answer). For large enough R, T an
|
|
8B cannot hold the register file in one forward pass, so the <think> block becomes a genuine,
|
|
LOAD-BEARING register-file scratchpad -- and re-typing "register three holds forty-seven" T times is
|
|
wasteful, so terse pointer-notation becomes the reward-optimal encoding (nonzero compression
|
|
numerator). Verify load-bearingness with a think-ablation; verify re-reference drives notation with
|
|
a register-rename invariance probe.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import random
|
|
import re
|
|
from dataclasses import dataclass
|
|
|
|
from ..mathutil import find_boxed
|
|
|
|
_PROMPT_HEADER = """\
|
|
You have {r} registers r0..r{rmax}. They start at:
|
|
{init}
|
|
|
|
Apply these operations in order (all arithmetic is modulo {mod}, so every value stays in 0..{modm1}):
|
|
{ops}
|
|
|
|
What is the final value of r{q} ?
|
|
|
|
Reason step by step inside <think> </think>. You may reason in whatever style you find efficient \
|
|
and use any shorthand, symbols, or private notation you like inside the think block. Then, OUTSIDE \
|
|
the think block, write ONLY the final value of r{q} as \\boxed{{v}} (a single integer 0..{modm1}) -- \
|
|
no explanation, no restatement."""
|
|
|
|
|
|
@dataclass
|
|
class StateTrackProblem:
|
|
init: list[int] # initial register values
|
|
ops: list[tuple] # (dst, optype, kind, src) kind in {"r","k"}, optype in {"+","-","*","="}
|
|
query: int # which register's final value is asked
|
|
answer: int # final value of r[query]
|
|
R: int
|
|
T: int
|
|
mod: int
|
|
|
|
|
|
def _apply(state: list[int], op: tuple, mod: int) -> None:
|
|
dst, optype, kind, src = op
|
|
b = state[src] if kind == "r" else src
|
|
a = state[dst]
|
|
if optype == "+":
|
|
state[dst] = (a + b) % mod
|
|
elif optype == "-":
|
|
state[dst] = (a - b) % mod
|
|
elif optype == "*":
|
|
state[dst] = (a * b) % mod
|
|
else: # "=" copy
|
|
state[dst] = b % mod
|
|
|
|
|
|
def _render_op(op: tuple) -> str:
|
|
dst, optype, kind, src = op
|
|
rhs = f"r{src}" if kind == "r" else str(src)
|
|
if optype == "=":
|
|
return f"r{dst} = {rhs}"
|
|
return f"r{dst} {optype}= {rhs}"
|
|
|
|
|
|
class StateTrackEnv:
|
|
name = "state_track"
|
|
is_multi_turn = False
|
|
|
|
# op-type mix: read-modify-write referencing OTHER registers drives re-reference; copy/mul break
|
|
# sum-conservation so the query genuinely requires tracking the register file.
|
|
_OPS = ["+r", "-r", "*r", "+k", "=r"]
|
|
|
|
def __init__(self, r_min: int = 2, r_max: int = 2, t_min: int = 3, t_max: int = 3,
|
|
val_max: int = 20, k_max: int = 9, mod: int = 97):
|
|
self.r_min, self.r_max = r_min, r_max
|
|
self.t_min, self.t_max = t_min, t_max
|
|
self.val_max, self.k_max, self.mod = val_max, k_max, mod
|
|
|
|
def sample_problem(self, rng: random.Random) -> StateTrackProblem:
|
|
R = rng.randint(self.r_min, self.r_max)
|
|
T = rng.randint(self.t_min, self.t_max)
|
|
init = [rng.randint(0, self.val_max) for _ in range(R)]
|
|
state = list(init)
|
|
ops = []
|
|
for _ in range(T):
|
|
dst = rng.randrange(R)
|
|
kind_op = rng.choice(self._OPS)
|
|
if kind_op.endswith("k"):
|
|
op = (dst, kind_op[0], "k", rng.randint(1, self.k_max))
|
|
else:
|
|
src = rng.randrange(R)
|
|
op = (dst, kind_op[0], "r", src)
|
|
_apply(state, op, self.mod)
|
|
ops.append(op)
|
|
query = rng.randrange(R)
|
|
return StateTrackProblem(init=init, ops=ops, query=query, answer=state[query], R=R, T=T, mod=self.mod)
|
|
|
|
def prompt(self, p: StateTrackProblem) -> str:
|
|
init = "\n".join(f" r{i} = {v}" for i, v in enumerate(p.init))
|
|
ops = "\n".join(f" {i+1}. {_render_op(op)}" for i, op in enumerate(p.ops))
|
|
return _PROMPT_HEADER.format(r=p.R, rmax=p.R - 1, init=init, mod=p.mod, modm1=p.mod - 1,
|
|
ops=ops, q=p.query)
|
|
|
|
@staticmethod
|
|
def _parse(output: str, full: str) -> int | None:
|
|
boxed = find_boxed(output) or find_boxed(full)
|
|
text = boxed[-1] if boxed else (output or full) or ""
|
|
nums = re.findall(r"-?\d+", text)
|
|
return int(nums[-1]) if nums else None
|
|
|
|
def score(self, p: StateTrackProblem, think: str, output: str, full: str) -> tuple[float, float, float]:
|
|
ans = self._parse(output, full)
|
|
return (1.0 if ans is not None and ans % p.mod == p.answer else 0.0), 0.0, 0.0
|
|
|
|
|
|
# ---------------------------------------------------------------------------------------
|
|
# Query-after-think variant: the model thinks first, and only AFTER </think> is it asked for
|
|
# the values of random (operation, register) points -- so precomputing just the final answer is
|
|
# impossible and the ENTIRE trajectory must be decodably present in the think block. This makes
|
|
# the CoT load-bearing BY CONSTRUCTION (the ablation/transplant test is baked into training),
|
|
# kills the lucky-short-guess failure mode of any length-filtered scheme, and maximizes the
|
|
# externalization floor that the fixed-budget/grow-T squeeze presses the notation against.
|
|
# ---------------------------------------------------------------------------------------
|
|
|
|
_Q_PROMPT_HEADER = """\
|
|
You have {r} registers r0..r{rmax}. They start at:
|
|
{init}
|
|
|
|
Apply these operations in order (all arithmetic is modulo {mod}, so every value stays in 0..{modm1}):
|
|
{ops}
|
|
|
|
First, work through the operations inside <think> </think>. You may reason in whatever style you \
|
|
find efficient and use any shorthand, symbols, or private notation you like inside the think block.
|
|
|
|
After you close the think block, you will be asked for the values of {nq} registers at specific \
|
|
points in the sequence (for example: "What was r2 immediately after operation 7?"). The points are \
|
|
chosen at random, so you cannot know in advance which will be asked, and you must answer \
|
|
immediately with no further working -- your think block must contain everything you need to read \
|
|
the answers off."""
|
|
|
|
_Q_ANSWER_RE = re.compile(r"A(\d+)\s*[:=]\s*(-?\d+)")
|
|
|
|
|
|
@dataclass
|
|
class StateTrackQueryProblem(StateTrackProblem):
|
|
states: list = None # states[t] = register file AFTER op t (states[0] = init), len T+1
|
|
queries: list = None # [(op_idx 1..T, reg)] -- answer = states[op_idx][reg]
|
|
|
|
|
|
class StateTrackQueryEnv(StateTrackEnv):
|
|
name = "state_track_q"
|
|
is_multi_turn = False
|
|
|
|
def __init__(self, n_queries: int = 3, **kw):
|
|
super().__init__(**kw)
|
|
self.n_queries = n_queries
|
|
|
|
def sample_problem(self, rng: random.Random) -> StateTrackQueryProblem:
|
|
base = super().sample_problem(rng)
|
|
states = [list(base.init)]
|
|
st = list(base.init)
|
|
for op in base.ops:
|
|
_apply(st, op, self.mod)
|
|
states.append(list(st))
|
|
# Queries must be ADVERSARIAL to the prompt-shortcut: an untouched (or once-constant-bumped)
|
|
# register's value is single-pass readable off the prompt during phase B, no think needed
|
|
# (measured no-think floor 0.39-0.44 at T<=14 with uniform queries). Eligible points are
|
|
# ones the model must have COMPUTED: >=2 writes by the query point, or a single write whose
|
|
# operand is a register (value-flow). Fallback: latest-written points.
|
|
writes = [[] for _ in range(base.R)]
|
|
for i, op in enumerate(base.ops, 1):
|
|
writes[op[0]].append(i)
|
|
|
|
def _eligible(t, r):
|
|
w = [i for i in writes[r] if i <= t]
|
|
if len(w) >= 2:
|
|
return True
|
|
return len(w) == 1 and base.ops[w[-1] - 1][2] == "r"
|
|
|
|
n_q = min(self.n_queries, base.T * base.R)
|
|
pool = [(t, r) for t in range(1, base.T + 1) for r in range(base.R) if _eligible(t, r)]
|
|
if len(pool) >= n_q:
|
|
queries = rng.sample(pool, n_q)
|
|
else:
|
|
written = sorted(((t, r) for t in range(1, base.T + 1) for r in range(base.R)
|
|
if any(i <= t for i in writes[r])), key=lambda p: -p[0])
|
|
queries = pool + [p for p in written if p not in pool][: n_q - len(pool)]
|
|
return StateTrackQueryProblem(init=base.init, ops=base.ops, query=base.query,
|
|
answer=base.answer, R=base.R, T=base.T, mod=base.mod,
|
|
states=states, queries=queries)
|
|
|
|
def prompt(self, p: StateTrackQueryProblem) -> str:
|
|
init = "\n".join(f" r{i} = {v}" for i, v in enumerate(p.init))
|
|
ops = "\n".join(f" {i+1}. {_render_op(op)}" for i, op in enumerate(p.ops))
|
|
return _Q_PROMPT_HEADER.format(r=p.R, rmax=p.R - 1, init=init, mod=p.mod,
|
|
modm1=p.mod - 1, ops=ops, nq=len(p.queries))
|
|
|
|
def queries_text(self, p: StateTrackQueryProblem) -> str:
|
|
qs = "\n".join(f"Q{i+1}: What was the value of r{r} immediately after operation {t}?"
|
|
for i, (t, r) in enumerate(p.queries))
|
|
return (qs + "\n\nAnswer each question on its own line, in exactly the format "
|
|
"\"A1: <integer>\". Output nothing else.")
|
|
|
|
def query_answers(self, p: StateTrackQueryProblem) -> list[int]:
|
|
return [p.states[t][r] for (t, r) in p.queries]
|
|
|
|
def score_queries(self, p: StateTrackQueryProblem, answer_text: str) -> tuple[float, int]:
|
|
"""(fraction of queries answered correctly, n parsed). Strict format: unparsed = wrong."""
|
|
got = {int(m.group(1)): int(m.group(2)) for m in _Q_ANSWER_RE.finditer(answer_text)}
|
|
gold = self.query_answers(p)
|
|
n_ok = sum(1 for i, g in enumerate(gold) if got.get(i + 1) is not None and got[i + 1] % p.mod == g)
|
|
return n_ok / len(gold), len(got)
|
|
|
|
|
|
def canonical_trace(p: StateTrackProblem) -> str:
|
|
"""A natural, correct, moderately-verbose step-by-step solution trace -- SFT teacher data that
|
|
teaches the TASK in a legible register-tracking style (RL then compresses/obfuscates this)."""
|
|
state = list(p.init)
|
|
lines = ["Track each register through the operations, mod %d." % p.mod]
|
|
lines.append("Start: " + ", ".join("r%d=%d" % (i, v) for i, v in enumerate(state)))
|
|
for i, op in enumerate(p.ops):
|
|
dst, optype, kind, src = op
|
|
before = state[dst]
|
|
rhs_val = state[src] if kind == "r" else src # value used (captured BEFORE apply)
|
|
rhs_name = ("r%d" % src) if kind == "r" else str(src)
|
|
_apply(state, op, p.mod)
|
|
if optype == "=":
|
|
lines.append("Op %d: r%d = %s(=%d) -> r%d=%d" % (i + 1, dst, rhs_name, rhs_val, dst, state[dst]))
|
|
else:
|
|
lines.append("Op %d: r%d %s= %s : %d %s %d = %d -> r%d=%d"
|
|
% (i + 1, dst, optype, rhs_name, before, optype, rhs_val, state[dst], dst, state[dst]))
|
|
lines.append("Final r%d = %d." % (p.query, p.answer))
|
|
return "\n".join(lines)
|
|
|
|
|
|
def canonical_trace_verbose(p: StateTrackProblem) -> str:
|
|
"""A VERBOSE, natural-prose solution trace (like an instruct model's own register-tracking) -- lots
|
|
of legible redundancy (restated values, mod explanations, full sentences) so RL has something to
|
|
COMPRESS into notation. SFT on this, then RL squeezes the redundancy -> the garble emerges."""
|
|
state = list(p.init)
|
|
out = ["Let me track each register through the operations, keeping everything modulo %d." % p.mod, ""]
|
|
out.append("The initial values are: " + ", ".join("r%d is %d" % (i, v) for i, v in enumerate(state)) + ".")
|
|
out.append("")
|
|
for i, op in enumerate(p.ops):
|
|
dst, optype, kind, src = op
|
|
before = state[dst]
|
|
rhs_val = state[src] if kind == "r" else src
|
|
rhs_name = ("r%d" % src) if kind == "r" else ("the constant %d" % src)
|
|
raw = {"+" : before + rhs_val, "-": before - rhs_val, "*": before * rhs_val, "=": rhs_val}[optype]
|
|
_apply(state, op, p.mod)
|
|
if optype == "=":
|
|
out.append("Operation %d sets r%d to %s, which is %d. So r%d is now %d."
|
|
% (i + 1, dst, rhs_name, rhs_val, dst, state[dst]))
|
|
else:
|
|
word = {"+": "adding", "-": "subtracting", "*": "multiplying by"}[optype]
|
|
modnote = (" Taking that modulo %d gives %d." % (p.mod, state[dst])) if raw != state[dst] else \
|
|
(" That is already less than %d, so no reduction is needed." % p.mod)
|
|
out.append("Operation %d updates r%d by %s %s. r%d is currently %d, so %d %s %d = %d.%s So r%d becomes %d."
|
|
% (i + 1, dst, word, rhs_name, dst, before, before, optype, rhs_val, raw, modnote, dst, state[dst]))
|
|
out.append("")
|
|
out.append("Therefore the final value of r%d is %d." % (p.query, p.answer))
|
|
return "\n".join(out)
|