TRL GRPO bridge for pr_diff datasets — proposal + working branch

[KNambiarDJsc]

Summary

CLAUDE.md lists the consumption layer as "HF Hub publish bridge; planned TRL trainer bridge", and docs/reference/AGENTS.md §6 says a GRPO training loop is planned for v0.9, pr_diff first via its diff_similarity reward, ahead of an ORS-reward-server integration for the execution-verified pipelines. I've built a narrow, working piece of that specifically for pr_diff: repo2rlenv export --format trl + a GRPOTrainer-compatible reward function, on a branch, verified against the real published dataset and a real GRPOTrainer.train() step. Posting this before opening a PR, per CONTRIBUTING.md's guidance on non-trivial work.

The gap

A pr_diff task's instruction.md is only the PR's title + description — correct for the shipped verifier (_pr_diff_verifier.py), which only ever diffs two patches and never touches source. But it means a model given just the prompt has nothing to patch: there's no bridge from a pr_diff dataset to any trainer today, TRL or otherwise.

What I built

Two pieces, src/repo2rlenv/train/pr_diff.py:

1. export_pr_diff_dataset(dataset_dir, out_path) (CLI: repo2rlenv export <path> --format trl --out <file>) — reads, at each task's base commit, the exact files the oracle diff touches, and embeds them in the prompt alongside the issue text. Fetches happen via one git clone --bare --filter=blob:none per source repo (reused across tasks from that repo), then git cat-file per needed blob — no full clones. A task whose assembled prompt would exceed --max-prompt-chars (default 40k) is skipped, not truncated — a cut-off file can hide the oracle's own hunks. Output is JSONL, {"prompt": [...], "oracle_patch": ..., "task_id", "repo", "ref"}, no datasets import needed on the export side.

2. pr_diff_reward(completions, oracle_patch, **kwargs) — matches TRL's reward-function contract exactly (dataset columns arrive by name; trainer_state/log_extra/log_metric as kwargs), so repo2rlenv.train imports neither trl nor torch. Scoring reuses _pr_diff_verifier.py's five deterministic components (format_valid, size_sanity, file_targeting, region_overlap, similarity) and its combine() / catastrophic-size-cap logic verbatim — same functions the task's own container runs. Only the network LLM judge is left out (doesn't fit a rollout loop; combine() already knows how to redistribute a missing component's weight, since the in-container verifier does the same when the judge has no API key).

Verification

Beyond unit tests (tests/test_train_pr_diff.py, 13 tests — export against a real local git repo with base/fix commits, skip-reason coverage, extract_patch fence-scanning, reward ordering, TRL's exact call signature):

• Exported the real published dataset (AdithyaSK/repo2rlenv-pr-diff, 181 tasks, 26 repos, Python/Go/Rust/JS/TS): 91 rows written, 90 skipped — every skip was prompt_too_long at the default 40k-char budget. Zero fetch failures, zero unreadable/binary/missing-file skips across all 26 repos. (That 90/181 number is itself informative — it's roughly how much of the flagship dataset needs a larger prompt budget or file-selection heuristic to use for training as-is.)
• Ran a real GRPOTrainer.train() step (TRL 1.13.0, CPU torch, trl-internal-testing/tiny-random-LlamaForCausalLM) on 4 exported rows with pr_diff_reward wired in exactly as a user would. TRL's own logged metrics confirm the integration, not just that the call didn't raise:

  reward_func_names: ['pr_diff_reward']
  rewards/pr_diff_reward/mean: 0   reward: 0   reward_std: 0
  pr_diff/file_targeting: 0   pr_diff/format_valid: 0   pr_diff/region_overlap: 0
  pr_diff/similarity: 0   pr_diff/size_sanity: 0
  

  (All-zero is correct here — a random 21-parameter model emits gibberish in a 16-token completion, not a diff, so format_valid gates it to 0. The per-component pr_diff/* metrics landing next to TRL's own reward/kl/entropy is what proves log_metric wiring works.)
• Linux, matching CI: pytest 728 passed / 0 failed on py3.12 and py3.14, ruff check + ruff format --check clean, uv build clean.

Docs

New docs/reference/TRAINING.md; cross-linked from pr_diff.md's existing "Consuming the reward at training time" section (now a third path, alongside the existing Harbor and pure-reward.py paths), AGENTS.md §6 (worded to be clear this is a narrower complement to the planned ORS-reward-server v0.9 item, not a claim of having shipped it), API.md's module table, CLAUDE.md's repo map, and mkdocs.yml's nav.

Out of scope / open questions

• pr_diff only. pr_runtime/commit_runtime/cve_patches need the task's Docker sandbox to score a completion — that's squarely what the planned ORS-reward-server / Harbor-wrapping approach in AGENTS.md is for, not something this reward function can do.
• Public source repos only for the file-fetch step.
• I called the module repo2rlenv.train / the CLI verb export --format trl since pr_diff.md already frames Harbor's harbor run as the "production path for evals" and this is explicitly the training-side complement — open to a different name if there's a preferred convention (e.g. under a future repo2rlenv.train.trl submodule if other trainers get bridges later).
• Not addressed: the 90/181 prompt_too_long rate on the reference dataset. Options would be a larger default budget, a cheaper file-selection heuristic (only hunks' surrounding context rather than whole files), or leaving it as a knob and letting users decide per their model's context window.

Happy to open the PR now, or wait for a decision on naming/placement first — whichever's less churn. Branch is up to date against main @ 1a8bad5.

[adithya-s-k]

Thanks for the work here. The Harbor/OpenEnv integration is already merged in OpenEnv #1036, including serving our Harbor datasets and capturing coding-agent rollouts for training. I'd like to use that path and keep Repo2RLEnv focused on generation for now. Your direct PR-diff exporter is a separate interface, so let's hold off on adding that here.