Skip to content
arnabghoPublic

About

Simplified RL for Calibrated Decisions: parallel constrained JSON decoding + GRPO with proper-scoring-rule rewards + calibration eval (Jev/RLCD reconstruction)

Resources

Stars

10 stars

Watchers

0 watching

Forks

Latest commit

 

History

8 Commits

Folders and files

Repository files navigation

rlcd-lite

A small, reproducible reimplementation of "RL for Calibrated Decisions" (RLCD): parallel constrained JSON decoding on a causal LM, trained with GRPO and a proper-scoring-rule (Brier) reward, evaluated with real calibration metrics.

TypeSafe's Jev (Sept 2026) is a "System One" model that returns typed JSON decisions with calibrated probabilities instead of generating text, trained with an undisclosed method they call RLCD. The open harshatheg/Qwen-2.5-1B-RLCD reproduces the inference trick with no training. This repo adds the missing pieces so the central claim, RL with a proper-scoring-rule reward yields calibrated decisions; binary reward does not, can be tested on a laptop.

Naming. This is not TypeSafe's algorithm (they have not published it) and it is unrelated to Yang et al. 2023, "RLCD: Reinforcement Learning from Contrastive Distillation". "RLCD-lite" is a reconstruction from the public description of Jev.

What it does

context + JSON schema ──prefill once──▶ KV cache
                                          │  one masked step per field, all fields batched
                                          ▼
   "severity": "  ──▶ logits ──▶ mask to {A,B,C,D} ──▶ softmax ──▶ argmax + probabilities
   "severity_confidence": "  ──▶ mask to {1..5} ──▶ stated confidence bucket
                                          ▼
                              programmatically assembled JSON (always valid)
  • Engine (engine.py): one fused forward over prompt ⊕ suffix_1 ⊕ … ⊕ suffix_F with a block-diagonal attention mask, so every field is scored exactly as an independent prompt + suffix forward would be, in one pass. A KV-cache path is provided for asking new questions about the same context.
  • Labels (candidates.py): options are enumerated with labels verified to be a single token after the field quote for the given tokenizer, so 77- and 151-way intents work without multi-token scoring.
  • Training (trainer.py): each decision field gets a stated-confidence field. A rollout samples (choice, confidence) per field. The default reward is RLCR-style, 1[correct] - (confidence - 1[correct])^2: correctness plus a Brier term, so the optimum is right and honest. Because a decision is one step, all G rollouts share one forward pass and GRPO reduces to REINFORCE with a group baseline (PPO clip and KL-to-base are optional). Ablations: --reward brier (Brier alone), --reward binary (correctness alone), --decouple (separate advantages for the choice and the confidence token).
  • Eval (evaluate.py, calibration.py, selective.py): three calibration signals (softmax max-prob, stated confidence, temperature-scaled softmax), Brier with the Murphy decomposition (reliability / resolution / uncertainty), equal-mass-bin ECE with bootstrap CI, log loss, risk-coverage / AURC, confidence gates, cross-field consistency, and option-order permutation sensitivity.

Install

uv sync --all-extras          # torch, transformers, peft, datasets, matplotlib
uv run pytest                 # 27 tests, tiny random model, no downloads

Use

# decide a schema over a context (downloads Qwen2.5-0.5B-Instruct on first use)
uv run rlcd-lite infer --context "Prod DB timing out on writes, all customers affected." \
  --schema-json '{"fields":[{"name":"severity","type":"choice","options":["low","medium","high","critical"]},
                            {"name":"team","type":"choice","options":["billing","auth","infra","frontend"]},
                            {"name":"needs_human","type":"bool"}]}'

# train: GRPO + Brier reward, LoRA on Qwen2.5-0.5B-Instruct, synthetic triage
uv run rlcd-lite train --dataset synthetic --steps 200 --reward brier --run-name grpo-brier

# A/B: base vs temperature scaling vs GRPO-Brier vs GRPO-binary
scripts/ab_eval.sh synthetic 200 300        # -> runs/ab-synthetic/eval.json + PNGs
uv run rlcd-lite eval --arms base --datasets boolq,mmlu --n 300 --permutation-k 4

Datasets: synthetic (3-field ticket triage with a cross-field rule), banking77, clinc150, boolq, mmlu[:subject,…]. uv run rlcd-lite data --list.

Results

Finding 1: Brier on stated confidence alone is a trap

The obvious "simplified RLCD" objective, reward = -(stated_conf - 1[correct])^2, is a proper scoring rule for the confidence, but the policy also controls the answer. Being predictably wrong at confidence 0.1 scores -0.01, the same as being predictably right at 0.9. On the synthetic triage task, Qwen2.5-0.5B-Instruct found that optimum in about 100 steps:

step reward train acc stated conf softmax conf choice entropy
1 -0.38 0.31 0.69 0.58 0.97
81 -0.23 0.29 0.34 0.67 0.81
161 -0.19 0.25 0.30 0.95 0.18
200 -0.13 0.10 0.30 0.98 0.09

Stated-confidence ECE went to ~0 while held-out accuracy fell from 0.47 to 0.23 and the softmax became near-deterministic on wrong labels. The fix is the RLCR reward (Damani et al. 2025): add correctness, so right & honest = 0.99 > wrong & honest = -0.01.

A/B on synthetic triage

Qwen2.5-0.5B-Instruct, LoRA r=16, 200 steps × 4 examples × G=8 on the M3 Pro (about 7 min per arm), evaluated on 300 held-out tickets with 200 calibration tickets for temperature scaling. scripts/ab_eval.sh synthetic 200 300 reproduces it; tables come from scripts/report.py.

arm reward acc stated-conf ECE stated Brier best AURC (signal)
base – 0.49 0.395 0.393 0.311 (softmax_ts)
grpo_binary 1[correct] 0.76 0.352 0.402 0.088 (softmax_ts)
grpo_brier -(c-correct)^2 0.21 0.086 0.176 0.730 (softmax_ts)
grpo_brier_dec Brier, decoupled advantages 0.76 0.105 0.208 0.063 (softmax_ts)
grpo_rlcr correct - (c-correct)^2 0.78 0.084 0.176 0.067 (softmax_ts)

reliability: base reliability: RLCR risk-coverage

Finding 2: RLCR gets accuracy and calibrated stated confidence; binary gets only accuracy. Binary and RLCR reach the same accuracy (0.76 vs 0.78), but binary leaves the confidence token untrained (ECE 0.35, mean stated 0.62 with no relation to correctness) while RLCR brings it to ECE 0.08. Decoupling the advantages is the other working fix for pure Brier.

Finding 3: calibrated is not the same as informative. The RLCR stated confidence has reliability 0.007 but resolution 0.000: after 200 steps it says "~0.7" for nearly everything, i.e. it learned the base rate, not per-example uncertainty. Its AURC (0.21) is far worse than the temperature-scaled softmax of the same model (0.07). The discriminative information lives in the choice distribution; the confidence token has not yet learned to read it. Longer training, per-field confidence grids, or an entropy feature are the obvious next experiments.

Finding 4: RL collapses the raw softmax. Every trained arm ends with mean softmax confidence 0.99–1.00 (choice entropy → 0.02, most groups zero-variance by step 150). Raw slice probabilities after RL are worse calibrated than the base model's (ECE 0.21–0.23 vs 0.18); a single fitted temperature repairs them (ECE 0.10–0.13) and gives the best Brier of any signal. "Probabilities from the LM head" are not calibrated by construction, before or after RL.

Finding 5: independent field scoring lets consistency drift. The base model never violates critical ⇒ needs_human (it rarely says critical); after RL the violation rate is 0.11 (RLCR) to 0.28 (binary). The 0.5B model also never learns needs_human (0.45 → 0.55, fitted temperature hits the 20.0 ceiling: the logits carry no signal for that field).

Full table (all arms × signals)
arm signal acc ECE (95% CI) Brier REL RES log loss AURC mean conf
base softmax 0.490 0.180 (0.147–0.211) 0.254 0.053 0.048 0.712 0.371 0.61
base stated 0.490 0.395 (0.364–0.423) 0.393 0.166 0.023 1.112 0.568 0.77
base softmax_ts 0.490 0.169 (0.137–0.194) 0.223 0.044 0.071 0.636 0.311 0.46
grpo_rlcr softmax 0.784 0.211 (0.184–0.243) 0.214 0.054 0.008 2.650 0.253 1.00
grpo_rlcr stated 0.784 0.084 (0.051–0.110) 0.176 0.007 0.000 0.539 0.208 0.70
grpo_rlcr softmax_ts 0.784 0.125 (0.099–0.146) 0.144 0.032 0.058 0.411 0.067 0.82
grpo_brier softmax 0.214 0.756 (0.723–0.780) 0.739 0.593 0.022 4.157 0.843 0.97
grpo_brier stated 0.214 0.086 (0.053–0.109) 0.176 0.007 0.000 0.538 0.793 0.30
grpo_brier softmax_ts 0.214 0.216 (0.191–0.237) 0.185 0.062 0.045 0.558 0.730 0.41
grpo_brier_dec softmax 0.763 0.229 (0.201–0.264) 0.230 0.097 0.050 2.009 0.076 0.99
grpo_brier_dec stated 0.763 0.105 (0.076–0.135) 0.208 0.030 0.003 0.607 0.249 0.66
grpo_brier_dec softmax_ts 0.763 0.096 (0.077–0.120) 0.136 0.025 0.071 0.395 0.063 0.81
grpo_binary softmax 0.758 0.232 (0.205–0.263) 0.228 0.078 0.029 2.267 0.110 0.99
grpo_binary stated 0.758 0.352 (0.327–0.388) 0.402 0.241 0.022 1.167 0.329 0.62
grpo_binary softmax_ts 0.758 0.106 (0.075–0.128) 0.151 0.022 0.053 0.431 0.088 0.76
  • base: per-field acc needs_human=0.45, severity=0.31, team=0.71; cross-field inconsistency 0.000; fitted T needs_human=20.00, severity=1.90, team=1.39
  • grpo_rlcr: per-field acc needs_human=0.55, severity=0.80, team=1.00; cross-field inconsistency 0.107; fitted T needs_human=20.00, severity=5.53, team=1.99
  • grpo_brier: per-field acc needs_human=0.45, severity=0.00, team=0.20; cross-field inconsistency 0.000; fitted T needs_human=20.00, severity=20.00, team=20.00
  • grpo_brier_dec: per-field acc needs_human=0.55, severity=0.75, team=0.99; cross-field inconsistency 0.120; fitted T needs_human=20.00, severity=4.41, team=2.31
  • grpo_binary: per-field acc needs_human=0.55, severity=0.72, team=1.00; cross-field inconsistency 0.280; fitted T needs_human=20.00, severity=6.22, team=1.93

Base model on public datasets (untrained, rlcd-lite eval --arms base)

200 test examples each, 100 calibration examples for temperature scaling, 4 option shuffles × 50 examples for order sensitivity.

dataset options acc softmax ECE stated ECE temp-scaled ECE order flip rate
boolq 2 0.62 0.264 0.087 0.097 0.52
mmlu (8 subjects) 4 0.49 0.155 0.202 0.108 0.46
banking77 77 0.015 0.249 0.819 0.033 –
clinc150 151 0.020 0.273 0.866 0.027 –

Finding 6: at 0.5B, label-token scoring does not survive many options. Banking77 and CLINC150 accuracy is at chance (1/77, 1/151): 47% of Banking77 predictions and 90% of CLINC150 predictions are a single label position, with mean stated confidence 0.83–0.89. The verified single-token label pool works mechanically, but the model is reading the label list, not the options. This is the honest reason the README's A/B is on a 3-field, 2–4-option schema. Larger models, verbatim option scoring (v0.2), or hierarchical schemas are the ways out.

Finding 7: option order moves the answer about half the time. Reshuffling the 2 or 4 options flips the base model's decision on 52% (BoolQ) and 46% (MMLU) of examples, with a mean total-variation distance of ~0.35 between the option-space distributions. A "probability" that depends this much on label position is not a probability about the world; permutation averaging is cheap here because fields are scored in parallel.

Full tables

boolq (model Qwen/Qwen2.5-0.5B-Instruct, n=200 test examples)

arm signal acc ECE (95% CI) Brier REL RES log loss AURC mean conf
base softmax 0.620 0.264 (0.209–0.341) 0.305 0.082 0.013 0.912 0.375 0.88
base stated 0.620 0.087 (0.022–0.162) 0.243 0.008 0.000 0.683 0.407 0.71
base softmax_ts 0.620 0.097 (0.090–0.188) 0.235 0.013 0.013 0.663 0.375 0.61
  • base: per-field acc answer=0.62; fitted T answer=4.59; option-order sensitivity: mean TV 0.350, flip rate 0.520 (k=4, n=50)

mmlu (model Qwen/Qwen2.5-0.5B-Instruct, n=200 test examples)

arm signal acc ECE (95% CI) Brier REL RES log loss AURC mean conf
base softmax 0.485 0.155 (0.131–0.232) 0.231 0.033 0.052 0.661 0.320 0.64
base stated 0.485 0.202 (0.141–0.283) 0.292 0.043 0.000 0.785 0.504 0.69
base softmax_ts 0.485 0.108 (0.111–0.201) 0.218 0.019 0.052 0.626 0.317 0.43
  • base: per-field acc answer=0.48; fitted T answer=2.57; option-order sensitivity: mean TV 0.356, flip rate 0.455 (k=4, n=50)

banking77 (model Qwen/Qwen2.5-0.5B-Instruct, n=200 test examples)

arm signal acc ECE (95% CI) Brier REL RES log loss AURC mean conf
base softmax 0.015 0.249 (0.232–0.271) 0.087 0.079 0.003 0.325 0.936 0.26
base stated 0.015 0.819 (0.793–0.842) 0.701 0.686 0.000 1.953 0.990 0.83
base softmax_ts 0.015 0.033 (0.022–0.053) 0.014 0.003 0.003 0.073 0.931 0.02
  • base: per-field acc intent=0.01; fitted T intent=9.13

clinc150 (model Qwen/Qwen2.5-0.5B-Instruct, n=200 test examples)

arm signal acc ECE (95% CI) Brier REL RES log loss AURC mean conf
base softmax 0.020 0.273 (0.251–0.294) 0.105 0.086 0.001 0.376 0.976 0.29
base stated 0.020 0.866 (0.846–0.882) 0.772 0.752 0.000 2.188 0.971 0.89
base softmax_ts 0.020 0.027 (0.013–0.046) 0.020 0.004 0.003 0.102 0.969 0.01
  • base: per-field acc intent=0.02; fitted T intent=20.00

ECE confidence intervals are percentile bootstrap; ECE is biased upward on small samples, so the interval can sit above the point estimate.

What to watch while training

Per-step rows land in runs/<name>/metrics.jsonl; evals.jsonl holds periodic held-out accuracy and ECE per signal.

Signal Healthy Sick
reward_mean rises then plateaus flat or oscillating
zero_variance_group_frac 0.2–0.6 > 0.8: peaked softmax, no learning signal. Raise --rollout-temp
stated_conf_mean vs acc gap shrinks stated → 0.9 while acc flat (what --reward binary does)
choice_entropy slow decrease collapse to ~0 in < 20 steps: lower --lr or add --beta-kl
grad_norm O(1), clipped occasionally 0.0 with zv=1.0 (every group identical, nothing to learn) or NaN
acc vs stated_conf_mean under --reward brier – acc falls while stated → lowest bucket: the Brier trap (Finding 1). Use rlcr or --decouple

What this is / isn't

  • Is: a faithful reconstruction of the public description of Jev's inference, plus the simplest RL objective that makes stated confidence a proper-scoring-rule target.
  • Isn't: TypeSafe's training data, architecture, or "parallel sampler"; a claim about Jev's actual calibration; a production inference server.
  • Known limits (v0.1): options must map to single-token labels (verbatim multi-token option scoring is v0.2); fields are scored independently given the prompt, so cross-field consistency is measured, not enforced; a supervised-Brier arm for the RL-vs-supervised ablation is not included yet.

References

  • TypeSafe, Introducing System One Models & Jev (2026).
  • H. Gundala, harshatheg/Qwen-2.5-1B-RLCD (2026): parallel constrained decoding, no training.
  • Damani et al., Beyond Binary Rewards: Training LMs to Reason About Their Uncertainty (RLCR, 2025): Brier term in RLVR.
  • Liu et al., Understanding R1-Zero-Like Training (Dr. GRPO, 2025): drop std-normalisation and length bias.
  • Murphy, A New Vector Partition of the Probability Score (1973): Brier = REL − RES + UNC.
  • Guo et al., On Calibration of Modern Neural Networks (2017): temperature scaling, ECE.
  • Holtzman et al., Surface Form Competition (2021): why first-token probabilities mislead.

Apache-2.0.

About

Simplified RL for Calibrated Decisions: parallel constrained JSON decoding + GRPO with proper-scoring-rule rewards + calibration eval (Jev/RLCD reconstruction)

Resources

Stars

10 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages