Publicar lo alineado: un reward model no es una política

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Any

from pydantic import BaseModel, ConfigDict

from rukh.models.squares import SQUARE_TOKENS
from rukh.paths import resolve
from rukh.publish.model import CONFIG_NAME, README_NAME, ModelPublishConfig, _api, render_card

log = logging.getLogger(__name__)

__all__ = ["RewardPublishResult", "publish_reward"]

REPO_TYPE = "model"
REWARD_CARD_TEMPLATE = "reward.md.jinja"
SAFETENSORS_NAME = "model.safetensors"

class RewardPublishResult(BaseModel):
    """Where the folder was staged, what went in it, and whether it was uploaded."""

    model_config = ConfigDict(extra="forbid")

    repo_id: str
    folder: str
    files: list[str]
    uploaded: bool
    accuracy: float

def _load(run_dir: Path) -> tuple[dict[str, Any], dict[str, Any]]:
    """``(checkpoint payload, run result)`` for a reward run directory."""
    import torch

    checkpoint = run_dir / "reward.pt" if run_dir.is_dir() else run_dir
    payload = torch.load(checkpoint, map_location="cpu", weights_only=False)
    result_path = checkpoint.parent / "run.json"
    if not result_path.is_file():
        raise FileNotFoundError(
            f"{result_path} is missing; the card is written from it, so the run has to have "
            "been produced by `rukh train reward`"
        )
    return payload, json.loads(result_path.read_text(encoding="utf-8"))

def publish_reward(
    run: Path,
    repo: str,
    cfg: ModelPublishConfig | None = None,
    stage: str | None = None,
    dry_run: bool = False,
) -> RewardPublishResult:
    """Stage (and unless ``dry_run``, upload) a reward model as its own Hub repository."""
    cfg = cfg or ModelPublishConfig()
    run = Path(run)
    payload, result = _load(run)
    repo_id = repo if "/" in repo else f"{cfg.owner}/{repo}"
    name = stage or repo_id.split("/")[-1].removeprefix("rukh-")

    encoder_config = payload["encoder_config"]
    state = payload["model"]
    params = sum(int(tensor.numel()) for tensor in state.values())

    folder = resolve(cfg.publish_dir) / repo_id
    folder.mkdir(parents=True, exist_ok=True)

    from safetensors.torch import save_file

    save_file(
        {key: value.detach().cpu().contiguous() for key, value in state.items()},
        str(folder / SAFETENSORS_NAME),
    )
    config = {
        "model_type": "rukh-reward",
        "stage": name,
        "params": params,
        "pooling": payload.get("pooling"),
        "encoder": encoder_config,
        "square_tokens": SQUARE_TOKENS,
        "objective": "bradley-terry",
        "training": result.get("checkpoint"),
    }
    (folder / CONFIG_NAME).write_text(
        json.dumps(config, indent=2, ensure_ascii=False) + "\n", encoding="utf-8", newline="\n"
    )

    training = payload.get("config") or {}
    card = render_card(
        {
            "license": cfg.license,
            "repo_id": repo_id,
            "stage": name,
            "repository_url": cfg.repository_url,
            "course_url": cfg.course_url,
            "datasets": [f"{cfg.owner}/rukh-pairs-dpo"],
            "params": params,
            "n_layer": encoder_config.get("n_layer"),
            "d_model": encoder_config.get("d_model"),
            "square_tokens": SQUARE_TOKENS,
            "pairs": result["pairs"],
            "train_pairs": result["train_pairs"],
            "val_pairs": result["val_pairs"],
            "accuracy": result["accuracy"],
            "accuracy_without_mates": result.get("accuracy_without_mates"),
            "loss": result["loss"],
            "pearson": result.get("pearson"),
            "pearson_without_mates": result.get("pearson_without_mates"),
            "bands": result.get("bands") or [],
            "epochs": result.get("epochs"),
            "batch_size": training.get("batch_size"),
            "lr": training.get("lr"),
            "seed": training.get("seed"),
            "from_scratch": result.get("from_scratch", True),
            "seed_spread": _seed_spread(result),
        },
        REWARD_CARD_TEMPLATE,
    )
    (folder / README_NAME).write_text(card, encoding="utf-8", newline="\n")

    files = [SAFETENSORS_NAME, CONFIG_NAME, README_NAME]
    uploaded = False
    if not dry_run:
        api = _api()
        api.create_repo(repo_id, repo_type=REPO_TYPE, exist_ok=True)
        api.upload_folder(repo_id=repo_id, folder_path=str(folder), repo_type=REPO_TYPE)
        uploaded = True
        log.info("subido %s", repo_id)

    return RewardPublishResult(
        repo_id=repo_id,
        folder=str(folder),
        files=files,
        uploaded=uploaded,
        accuracy=float(result["accuracy"]),
    )

---
license: {{ license }}
library_name: rukh
pipeline_tag: feature-extraction
language:
  - en
datasets:
{%- for dataset in datasets %}
  - {{ dataset }}
{%- endfor %}
tags:
  - chess
  - rukh
  - reward-model
  - preference
---

# {{ repo_id }}

A **reward model**: it reads a chess position and answers with one number, trained on
{{ "{:,}".format(pairs) }} pairs of moves that a chess engine had already ranked. This is the
`{{ stage }}` stage of [Rukh]({{ repository_url }}), a course that builds a chess language model
end to end: {{ "{:,}".format(params) }} parameters over {{ n_layer }} layers of width
{{ d_model }}, reading the {{ square_tokens }} square tokens of a board.

Read how it was built: [{{ course_url }}]({{ course_url }})

## What the number means, and what it does not

The training objective is Bradley-Terry: `-log sigma(r(chosen) - r(rejected))`. Only
**differences** appear in it, so adding a constant to every score changes nothing. The scale has
no units and no origin, and comparing this model's `0.8` with another training run's `1.4` is
comparing two rulers with no zero. What the model is for is **ordering two positions**, and the
number below is how often it gets that order right.

## Results

Measured on {{ "{:,}".format(val_pairs) }} held-out pairs, split **by game** so that two pairs a
few plies apart cannot straddle the split.

| Metric | Value |
|---|---|
| Accuracy on held-out pairs | {{ "%.2f"|format(accuracy * 100) }} % |
{%- if accuracy_without_mates %}
| Accuracy over pairs an evaluation can decide | {{ "%.2f"|format(accuracy_without_mates * 100) }} % |
{%- endif %}
| Preference loss | {{ "%.4f"|format(loss) }} |
{%- if pearson is not none %}
| Margin vs engine gap (Pearson) | {{ "%+.3f"|format(pearson) }} |
{%- endif %}
{%- if pearson_without_mates is not none %}
| Margin vs engine gap, without mates (Pearson) | {{ "%+.3f"|format(pearson_without_mates) }} |
{%- endif %}

### Where it is right and where it is wrong

Accuracy by how far apart the engine scored the two moves. The **mate** band is the one to read
first: it is the largest slice after the narrowest band, and it is where the model is worst.

| Engine gap | Pairs | Accuracy |
|---|---:|---:|
{%- for band in bands %}
| {% if band.band == "mate" %}mate{% else %}{{ band.band }} cp{% endif %} | {{ band.pairs|int }} | {{ "%.2f"|format(band.accuracy * 100) }} % |
{%- endfor %}

Read the `Pairs` column before the `Accuracy` one. A band of a dozen pairs reports whatever those
dozen did: across five runs of this configuration the `800-2000` band gave 90 %, 68 %, 90 %, 50 %
and 57 %, which is the same absence of data wearing five faces.

A mate is a tactical fact, and what this model reads is a static view of the position the move
leads to: it cannot see one. That band is exactly what a **verifiable** reward answers without
error, which is why the same milestone that trained this model also wrote one.

The two correlations above do not share a sign, and that is not a typo. The mate pairs carry the
largest engine gap by construction and the smallest learned margin, so a third of the points bends
the line the other way. One correlation over two regimes describes neither.

## What it is not for

**DPO does not need this model.** Direct Preference Optimization has an implicit reference -- a
frozen copy of the starting policy -- and trains on the pairs directly. A reward model is what PPO
needs, what a learned reward is compared against, and what makes the accuracy above a number worth
publishing at all.

## Training

{{ "{:,}".format(pairs) }} pairs, {{ "{:,}".format(train_pairs) }} for training and
{{ "{:,}".format(val_pairs) }} held out, {{ epochs }} epochs, batch {{ batch_size }}, learning rate
{{ lr }}, seed {{ seed }}.

{% if from_scratch %}Trained **from a random initialisation**, and that was measured rather than
assumed: starting from the project's pretrained position encoder scored about 6.5 points *worse*.
That encoder learned to fill in masked moves; this task is ordering positions by how good they
are, and what it brought did not help.{% else %}Started from a pretrained position encoder.{% endif %}

The seed also decides the split, so it decides how many mate pairs land in validation -- and the
headline number moves with that. Four seeds of this exact configuration gave
{{ seed_spread }}. Compare two runs only if they shared a seed.

## Reproducing

```bash
uv run rukh train reward --config configs/train/rm.yaml
```

## Licence

{{ license }}. Trained on engine evaluations of public Lichess games.

        DatasetSpec(
            name="rukh-pairs-onpolicy",
            repo_type="dataset",
            local_dir="data/pairs-onpolicy",
            patterns=["pairs.parquet", "manifest.json"],
            command="onpolicy",
            pretty_name="Rukh on-policy pairs",
            description=(
                "Preference pairs built from the moves the model itself proposes. The positions "
                "are exactly those of `rukh-pairs-dpo`, so the only thing that differs between the "
                "two datasets is where the two moves came from -- which is what makes a DPO run on "
                "each of them a comparison. Four moves sampled per position at temperature 1.0, "
                "scored by Stockfish at a fixed depth of 10, keeping the best and the worst when "
                "they are at least 100 centipawns apart.\n\n"
                "Of the 13 838 positions, 6 386 produced a pair. The rest is the interesting half: "
                "3 743 (27 %) because the model proposed a single distinct legal move, and 3 709 "
                "(27 %) because its best and worst candidate were less than a pawn apart. And only "
                "2.9 % of these pairs involve a mate, against a third of the off-policy ones: a "
                "multi-PV search finds mates the model was never going to propose, so the two "
                "datasets are two distributions of difficulty and not two sources of one thing."
            ),
            columns=[
                ("game_id", "game the position came from, for splitting without leakage"),
                ("ply", "half-move index of the position"),
                ("prefix", "moves played up to the position (UCI, space separated)"),
                ("chosen", "best of the sampled moves (UCI)"),
                ("rejected", "worst of the sampled moves (UCI)"),
                ("cp_chosen", "score of `chosen` for the side to move (mate = +/-10000)"),
                ("cp_rejected", "score of `rejected` for the side to move"),
                ("phase", "`opening`, `middlegame` or `endgame`"),
                ("white_elo", "Elo header the model was prompted with"),
                ("black_elo", "Elo header the model was prompted with"),
                ("candidates", "distinct legal moves the model proposed here (2 to 4)"),
            ],
            source="`rukh-pairs-dpo` for the positions, `chorcat/rukh-medium` for the moves",
            source_datasets=["original"],
            task_categories=["text-generation"],
            tags=["dpo", "preferences", "on-policy"],
            # `positions` counts what was *visited*, not what was written: 13 838 against 6 386
            # rows. Summing it into the row count would double the size of the dataset on the card.
            exclude_counts=["positions"],
        ),

DEMO_STAGES = {
    "tiny": "tiny-int8",
    "small": "small-fp16",
    "medium": "medium-fp16",
    "medium-dpo": "medium-dpo-fp16",
    "medium-grpo": "medium-grpo-fp16",
    "medium-elo": "medium-elo-fp16",
    "medium-lora": "medium-lora-fp16",
}

# The GRPO engine cache is the same kind of thing: a `(FEN, move) -> centipawns` dictionary that
# any run rebuilds from Stockfish at a fixed depth. Megabytes of it, and reproducible by
# definition, since the depth limit is what makes the analysis deterministic (D-118).
/artifacts/grpo-engine-cache.json
/artifacts/grpo-cache-*.json

@publish_app.command("reward")
def publish_reward_cmd(
    run: Annotated[
        Path,
        typer.Option("--run", exists=True, help="Reward run directory or `reward.pt` checkpoint."),
    ],
    repo: Annotated[str, typer.Option("--repo", help="Hub repository, e.g. chorcat/rukh-rm.")],
    stage: Annotated[str | None, typer.Option("--stage", help="Stage name for the card.")] = None,
    dry_run: Annotated[
        bool, typer.Option("--dry-run", help="Stage the folder locally; touch no network.")
    ] = False,
    as_json: Annotated[bool, typer.Option("--json", help="Print the result as JSON only.")] = False,
) -> None:
    """Publish the reward model, with the band table and both correlations on the card.

    It gets its own command because it is not a policy: no ONNX, because nothing serves it in a
    browser, and no Elo, because it does not play. What it has is an ordering accuracy, and that
    number is only honest next to the split it was measured on.
    """
    from rukh.publish.reward import publish_reward

    logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(message)s")
    result = publish_reward(run, repo, stage=stage, dry_run=dry_run)
    if as_json:
        typer.echo(result.model_dump_json(indent=2))
        return
    where = "uploaded to" if result.uploaded else "staged for"
    typer.echo(f"{where} {result.repo_id}: {', '.join(result.files)}")
    typer.echo(f"accuracy: {result.accuracy * 100:.2f} % on held-out pairs")
    typer.echo(f"folder:   {result.folder}")

def test_the_card_carries_the_bands_and_both_correlations(rukh_home, tmp_path):
    result = publish_reward(_run_dir(tmp_path), "chorcat/rukh-rm", dry_run=True)
    assert result.uploaded is False
    # The env var is RUKH_HOME, not RUKH_ROOT. Getting it wrong is silent -- the test passes and
    # stages into the real `artifacts/publish/`, quietly replacing a published folder with toy
    # weights. It happened once; this line is why it cannot happen twice.
    assert Path(result.folder).is_relative_to(rukh_home)
    card = (Path(result.folder) / "README.md").read_text(encoding="utf-8")

    # The number, and the number that is not the same number.
    assert "72.91 %" in card
    assert "73.62 %" in card
    # Both correlations, with their signs, because they do not share one.
    assert "-0.034" in card
    assert "+0.058" in card
    # The band table, including the one the model is worst at.
    assert "| mate | 484 | 71.49 % |" in card
    # And the sentence that stops the headline being read as stable.
    assert "75.19 %" in card, "the spread across seeds has to travel with the number"
    assert "same seed" in card, "and so does the part of it that is not the split"


def test_the_run_has_to_have_been_measured(rukh_home, tmp_path):
    """A card written from a missing file would be a card of defaults."""
    folder = _run_dir(tmp_path)
    (folder / "run.json").unlink()
    with pytest.raises(FileNotFoundError):
        publish_reward(folder, "chorcat/rukh-rm", dry_run=True)


def test_nothing_reaches_the_network_on_a_dry_run(rukh_home, tmp_path, monkeypatch):
    def _boom():
        raise AssertionError("a dry run must not build an API client")

    monkeypatch.setattr("rukh.publish.reward._api", _boom)
    result = publish_reward(_run_dir(tmp_path), "chorcat/rukh-rm", dry_run=True)
    assert set(result.files) == {"model.safetensors", "config.json", "README.md"}