DPO: la referencia implícita, y de dónde salen las dos jugadas

"""Direct Preference Optimization on the decoder, over engine-scored move pairs.

Every other lever tried so far teaches the model to *imitate*: more games, better games, a header
saying how strong the players were. Imitation is capped by the corpus, and D-069 measured that cap
-- conditioning on 2600 does not make the model play better, it makes it play like a 2600 losing
to tactics. DPO changes the objective. Each example says "in this position `chosen` is at least
``min_delta_cp`` better than `rejected`", which is information no human game contains, and the
loss pushes probability from one to the other instead of towards the average continuation.

Our completion is a single token -- one move -- so the usual sequence-level DPO collapses to the
log-probability of two tokens at one position, which is why this file is short.
"""

class DpoConfig(BaseConfig):
    """Where the pairs and the starting weights are, and the two knobs that matter."""

    pairs: str = "data/pairs/dpo-prompts.parquet"
    checkpoint: str = ""
    out_dir: str = "checkpoints"
    run_name: str = "dpo"
    beta: float = Field(default=0.1, gt=0)
    """How hard the loss pushes. Above ~0.5 the policy leaves the reference and legality goes."""
    lr: float = Field(default=5e-6, gt=0)
    nll_weight: float = Field(default=0.0, ge=0)
    """Optional anchor on the chosen move. Pure DPO can raise the margin while the model forgets
    how to play at all; a little next-token loss keeps it honest."""
    batch_size: int = Field(default=32, ge=1)
    epochs: int = Field(default=1, ge=1)
    block: int = Field(default=200, ge=8)
    val_fraction: float = Field(default=0.1, ge=0.0, lt=0.5)
    seed: int = 42
    device: str | None = None

def batches(
    rows: list[tuple[list[int], int, int]], size: int, pad_id: int, shuffle: bool, seed: int
) -> Iterator[PreferenceBatch]:
    """Pad each batch to its own longest prompt and remember where each prompt ends."""
    order = np.arange(len(rows))
    if shuffle:
        np.random.default_rng(seed).shuffle(order)
    for start in range(0, len(order), size):
        chunk = [rows[int(i)] for i in order[start : start + size]]
        if not chunk:
            continue
        width = max(len(ids) for ids, _, _ in chunk)
        tokens = torch.full((len(chunk), width), pad_id, dtype=torch.long)
        last = torch.empty(len(chunk), dtype=torch.long)
        for i, (ids, _, _) in enumerate(chunk):
            tokens[i, : len(ids)] = torch.tensor(ids, dtype=torch.long)
            last[i] = len(ids) - 1
        yield PreferenceBatch(
            tokens,
            last,
            torch.tensor([c for _, c, _ in chunk], dtype=torch.long),
            torch.tensor([r for _, _, r in chunk], dtype=torch.long),
        )

def move_logprobs(model: MoveDecoder, batch: PreferenceBatch) -> torch.Tensor:
    """Log-probabilities over the vocabulary at the position that predicts the next move."""
    logits, _ = model(batch.tokens)
    picked = logits[torch.arange(logits.shape[0], device=logits.device), batch.last]
    return F.log_softmax(picked.float(), dim=-1)

def dpo_loss(
    policy: torch.Tensor, reference: torch.Tensor, batch: PreferenceBatch, beta: float
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """The DPO objective for a one-token completion, plus the margin and the accuracy.

    ``margin`` is ``rewards/margins`` in the usual reporting: how much more the policy prefers the
    better move than the reference did. ``accuracy`` is the share of pairs it orders correctly,
    which is the number to watch -- a margin can grow while the ordering stays wrong.
    """
    index = torch.arange(policy.shape[0], device=policy.device)
    chosen = policy[index, batch.chosen] - reference[index, batch.chosen]
    rejected = policy[index, batch.rejected] - reference[index, batch.rejected]
    margin = chosen - rejected
    loss = -F.logsigmoid(beta * margin).mean()
    accuracy = (policy[index, batch.chosen] > policy[index, batch.rejected]).float().mean()
    return loss, margin.mean(), accuracy

def evaluate(
    policy: MoveDecoder,
    reference: MoveDecoder,
    rows: list[tuple[list[int], int, int]],
    cfg: DpoConfig,
    device: torch.device,
    pad_id: int,
) -> tuple[float, float, float]:
    """Mean loss, margin and pair accuracy over held-out pairs."""
    policy.eval()
    totals = [0.0, 0.0, 0.0]
    seen = 0
    with torch.no_grad():
        for batch in batches(rows, cfg.batch_size, pad_id, shuffle=False, seed=cfg.seed):
            batch = batch.to(device)
            loss, margin, accuracy = dpo_loss(
                move_logprobs(policy, batch), move_logprobs(reference, batch), batch, cfg.beta
            )
            n = len(batch)
            totals[0] += float(loss) * n
            totals[1] += float(margin) * n
            totals[2] += float(accuracy) * n
            seen += n
    policy.train()
    if not seen:
        return math.nan, math.nan, math.nan
    return totals[0] / seen, totals[1] / seen, totals[2] / seen

def train(cfg: DpoConfig) -> Path:
    """Align a decoder on the preference pairs and return the checkpoint it wrote."""
    import copy

    import polars as pl

    from rukh import paths
    from rukh.train.checkpoint import load_model, save_checkpoint
    from rukh.train.common import pick_device

    torch.manual_seed(cfg.seed)
    device = torch.device(cfg.device or pick_device())
    tok = UciTokenizer()

    policy, payload = load_model(paths.resolve(cfg.checkpoint), device)
    policy = policy.to(device).train()
    reference = copy.deepcopy(policy).eval()
    for parameter in reference.parameters():
        parameter.requires_grad_(False)

    frame = pl.read_parquet(paths.resolve(cfg.pairs).as_posix())
    rows = encode_pairs(frame, tok, cfg.block)
    if not rows:
        raise ValueError(f"{cfg.pairs} produced no usable pairs")
    cut = int(len(rows) * (1 - cfg.val_fraction))
    rng = np.random.default_rng(cfg.seed)
    order = rng.permutation(len(rows))
    train_rows = [rows[int(i)] for i in order[:cut]]
    val_rows = [rows[int(i)] for i in order[cut:]]
    LOGGER.info("dpo: %d pares de entrenamiento, %d de validación", len(train_rows), len(val_rows))

    optimizer = torch.optim.AdamW(policy.parameters(), lr=cfg.lr)
    out_dir = paths.resolve(cfg.out_dir) / cfg.run_name
    out_dir.mkdir(parents=True, exist_ok=True)

    before = evaluate(policy, reference, val_rows, cfg, device, tok.pad_id)
    LOGGER.info("antes: loss %.4f · margen %.4f · aciertos %.4f", *before)
    step = 0
    for epoch in range(cfg.epochs):
        for batch in batches(
            train_rows, cfg.batch_size, tok.pad_id, shuffle=True, seed=cfg.seed + epoch
        ):
            batch = batch.to(device)
            policy_logprobs = move_logprobs(policy, batch)
            with torch.no_grad():
                reference_logprobs = move_logprobs(reference, batch)
            loss, margin, accuracy = dpo_loss(policy_logprobs, reference_logprobs, batch, cfg.beta)
            if cfg.nll_weight:
                index = torch.arange(policy_logprobs.shape[0], device=device)
                loss = loss - cfg.nll_weight * policy_logprobs[index, batch.chosen].mean()
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(policy.parameters(), 1.0)
            optimizer.step()
            step += 1
            if step % 50 == 0:
                LOGGER.info(
                    "paso %d · loss %.4f · margen %.4f · aciertos %.4f",
                    step,
                    float(loss),
                    float(margin),
                    float(accuracy),
                )
    after = evaluate(policy, reference, val_rows, cfg, device, tok.pad_id)
    LOGGER.info("después: loss %.4f · margen %.4f · aciertos %.4f", *after)

    final = out_dir / "dpo.pt"
    save_checkpoint(
        final,
        step=step,
        model=policy,
        optimizer=None,
        cfg={
            **cfg.model_dump(mode="json"),
            "val_before": dict(zip(("loss", "margin", "accuracy"), before, strict=True)),
            "val_after": dict(zip(("loss", "margin", "accuracy"), after, strict=True)),
        },
        model_cfg=payload["model_cfg"],
        vocab_hash=payload.get("vocab_hash"),
        data_manifest_sha=payload.get("data_manifest_sha"),
    )
    return final

"""Preference pairs built from the moves the **model** would actually play.

The pairs the project already has are off-policy: they come from Lichess's multi-PV evaluations,
so they say "in this position ``e4`` is better than ``h4``" regardless of whether the model was
ever going to consider either. That is information, and it is not the information that fixes a
particular model. A model only ever loses games to the moves it plays.

An on-policy pair asks a narrower and more useful question: *of the moves you were going to
consider here, which was best and which was worst?* Sample ``candidates`` moves from the model,
score them with the engine, keep the best and the worst if they are far enough apart.

Three decisions worth defending, because each one is a way the comparison could have been unfair.

**The positions are the same ones the off-policy pairs use.** Same prefixes, same phase balance,
same games. The only thing that changes between the two datasets is where the two moves came
from, which is the whole point of running both.

**The engine is limited by depth, not by time.** P4 measured that a time-limited Stockfish makes
a measurement irreproducible (D-107): the same run twice gave 1498 and 1558. A depth limit gives
the same answer on a busy machine and on an idle one, and a dataset is a thing that should be
rebuildable.

**Sampling is at temperature 1.0.** At the near-deterministic setting every published number is
read at, a model proposes the same move four times and there is no pair to build. The candidates
have to come from the distribution, not from its mode -- which is the same lesson M4 ended on.
"""

def score_move(
    engine: chess.engine.SimpleEngine, board: chess.Board, move: chess.Move, depth: int
) -> int:
    """Centipawns **for the side that plays ``move``**, after playing it.

    The engine is asked about the position the move leads to, where the *opponent* is to move, so
    the score comes back from the opponent's point of view and is negated. Mates use the project's
    own ``MATE_SCORE`` convention so these numbers can sit in the same column as the off-policy
    ones without a footnote.
    """
    board.push(move)
    try:
        info = engine.analyse(board, chess.engine.Limit(depth=depth))
        score = info["score"].pov(not board.turn)
        return int(score.score(mate_score=MATE_SCORE))
    finally:
        board.pop()

def sample_candidates(
    model: MoveDecoder,
    tok: UciTokenizer,
    board: chess.Board,
    history: list[int],
    count: int,
    temperature: float,
    top_k: int | None,
) -> list[chess.Move]:
    """Distinct legal moves the model proposes here, sampled from its own distribution.

    Illegal proposals are dropped rather than rescued: an on-policy pair is about ranking the
    moves the model would *play*, and the demo masks illegal moves before they reach the board.
    How often it proposes one is a different measurement, and the suite already makes it.
    """
    from rukh.infer.sampler import SampleConfig, model_generator, pick_move

    cfg = SampleConfig(temperature=temperature, top_k=top_k, mask_illegal=False)
    generator = model_generator(model, cfg)
    seen: dict[str, chess.Move] = {}
    for _ in range(count * 4):  # a few extra draws, since duplicates are the common case
        move, _ = pick_move(model, tok, board, history, cfg, generator)
        if move is not None and move in board.legal_moves:
            seen.setdefault(move.uci(), move)
        if len(seen) >= count:
            break
    return list(seen.values())

def _history_for(tok: UciTokenizer, prefix: str, white_elo: int, black_elo: int) -> list[int]:
    """The prompt the model sees: the same three header tokens the ladder uses, then the game."""
    from rukh.infer.game import _history

    ids = _history(tok, white_elo, black_elo)
    ids.extend(tok.vocab.get(uci, tok.unk_id) for uci in prefix.split())
    return ids

def build_on_policy(cfg: OnPolicyConfig) -> Manifest:
    """Sample, score and write the on-policy pairs, with a manifest that says how."""
    from rukh.engine import find_stockfish
    from rukh.eval.suite import resolve_model
    from rukh.tokenize.uci_vocab import UciTokenizer
    from rukh.train import load_model, pick_device

    binary = find_stockfish()
    if binary is None:
        raise FileNotFoundError("Stockfish not found; set RUKH_STOCKFISH")
    torch.manual_seed(cfg.seed)
    where = cfg.device or pick_device()
    tok = UciTokenizer()
    model, _payload = load_model(resolve_model(cfg.checkpoint), map_location=where)
    model = model.to(where).eval()

    frame = pl.read_parquet(resolve(cfg.source))
    if cfg.max_positions is not None:
        # Shuffled first, because the source is phase-balanced *in blocks*: taking the head would
        # buy a dataset made entirely of openings and then compare it to one that is not.
        frame = frame.sample(n=min(cfg.max_positions, frame.height), shuffle=True, seed=cfg.seed)

    rows: list[dict[str, object]] = []
    skipped_same = skipped_close = 0
    with chess.engine.SimpleEngine.popen_uci(str(binary)) as engine:
        for index, row in enumerate(frame.iter_rows(named=True)):
            board = chess.Board()
            try:
                for uci in (row["prefix"] or "").split():
                    board.push_uci(uci)
            except ValueError:
                continue
            if board.is_game_over():
                continue
            history = _history_for(
                tok, row["prefix"] or "", int(row["white_elo"]), int(row["black_elo"])
            )
            moves = sample_candidates(
                model, tok, board, history, cfg.candidates, cfg.temperature, cfg.top_k
            )
            if len(moves) < 2:
                skipped_same += 1
                continue
            scored = [(score_move(engine, board, move, cfg.depth), move) for move in moves]
            scored.sort(key=lambda pair: pair[0], reverse=True)
            (best_cp, best), (worst_cp, worst) = scored[0], scored[-1]
            if best_cp - worst_cp < cfg.min_delta_cp:
                skipped_close += 1
                continue
            rows.append(
                {
                    "game_id": str(row["game_id"]),
                    "ply": int(row["ply"]),
                    "prefix": row["prefix"],
                    "chosen": best.uci(),
                    "rejected": worst.uci(),
                    "cp_chosen": int(best_cp),
                    "cp_rejected": int(worst_cp),
                    "phase": row["phase"],
                    "white_elo": int(row["white_elo"]),
                    "black_elo": int(row["black_elo"]),
                    "candidates": len(moves),
                }
            )
            if index and index % 250 == 0:
                log.info(
                    "%d positions, %d pairs (%d all-same, %d too close)",
                    index,
                    len(rows),
                    skipped_same,
                    skipped_close,
                )

# DPO on the pairs Lichess's evaluations produced: off-policy, and the control arm of M5's
# comparison. The only line that differs from `dpo-onpolicy.yaml` is `pairs`.
pairs: data/pairs-offpolicy-matched/pairs.parquet