Non-Vacuous Generalization Bounds for RLVR

This artifact reproduces every numerical result in the paper "Non-vacuous Generalization Bounds for Reinforcement Learning with Verifiable Rewards."

1. Teacher — per-domain rank-64 LoRA trained via verifier rewards using Tinker (tinker-cookbook/).
2. Student — a TinyLoRA (u = r = 16, ≈ 4 KB) trained via on-policy distillation from the teacher, using SkyRL (SkyRL/).
3. Quantization — global-codebook uniform quantization of the student's trainable vector to k ∈ {2, 3, 4, 5, 6, 7, 8, 16, 32, 64} levels.
4. Evaluation & bound — empirical reward over the m'=4096 evaluation subsample (n=64 generations per prompt) plus the Theorem 1 / Corollary 1 PAC-Bayes lower bound.

The bundled SkyRL/ and tinker-cookbook/ directories are vendored snapshots of the upstream open-source projects, included for reproducibility. They retain their original licenses and are referenced to their canonical public repositories.

---

Repository layout

artifact/
├── README.md                 (this file)
├── pyproject.toml, uv.lock   (Python deps)
├── .python-version           (3.12)
│
├── configs/                  Per-domain hyperparameter source-of-truth (paper Table 1)
│   ├── base.json             Shared defaults
│   ├── {math,code,general,sql}.json   Per-domain overrides
│   └── loader.py             Helper that exports configs as dict / shell vars
│
├── data/                     Empty until 01_build_data.sh runs
│   ├── README.md             What each curation script does, target row counts
│   └── subsample_eval.py     Deterministic m'=4096 eval subsample (seed 42)
│
├── models/                   Bundled TinyLoRA adapters (≈ 39 KB each)
│   ├── math/   {OPD, OPD_q{2..64}, direct-RLVR variants}.safetensors
│   ├── code/   ...
│   ├── general/ ...
│   └── sql/    ...
│
├── scripts/
│   ├── compute_description_length.py   KT arithmetic-coded C(h) breakdown
│   ├── lora_quantize.py, lora_compression_size.py, utils.py
│   └── run_pipeline/         Six numbered orchestration wrappers (below)
│
├── tinker-cookbook/          Vendored Tinker snapshot (teacher pipeline)
│   ├── tinker_cookbook/recipes/   Recipe library (math_rl, code_rl, distillation, …)
│   └── scripts/              Per-domain teacher launchers (4B + 27B variants)
│
└── SkyRL/                    Vendored SkyRL snapshot (student pipeline)
    ├── examples/train/{eurus_rl_math,mix_code,mix_general,synsql}/
    │                         Per-domain curation, training, and eval entry points
    ├── examples/train/eval/  Cross-domain unified evaluator
    ├── scripts/              Top-level student launchers (incl. merge_tinylora.py)
    └── skyrl-train/, skyrl-gym/    Vendored trainer and gym environment

The numbered runners under scripts/run_pipeline/ are thin wrappers around the vendored launchers; they inject paper-authoritative hyperparameters from configs/<domain>.json so the artifact stays in sync with the paper without patching upstream code.

The bundled models/<domain>/*.safetensors files contain only the ~4 KB TinyLoRA student adapters. The full per-domain curated datasets and the rank-64 LoRA teacher checkpoints are hosted on Hugging Face:

https://huggingface.co/collections/uiuc-kang-lab/rlvr-generalization-bounds

Step 01 (01_build_data.sh) pulls the curated parquets from this collection; step 02 (02_train_teacher.sh) can be skipped by pointing TEACHER_MODEL at the corresponding teacher repo in the collection if you only want to re-run the distillation / quantization / evaluation stages.

---

Setup

Requires Python 3.12. Install via uv (recommended):

cd artifact
uv sync          # creates a venv and installs everything in pyproject.toml + uv.lock

Environment variables (set in SkyRL/.env or your shell):

Hardware: teacher RLVR runs on Tinker's hosted backend (≥ 8 vCPU + 80 GB RAM client side); student distillation needs 2 × A100 80 GB (vLLM rollouts + SkyRL trainer). Steps 04–06 run on CPU.

---

Quick start: reproduce the paper's bound on the bundled adapters

The fastest way to validate the artifact (no training required, < 1 min on CPU):

# Math: should print  C(h) = 7726 bits  (paper: 7726 bits)
python3 scripts/run_pipeline/06_compute_bounds.py \
    --domain math \
    --adapter models/math/Qwen3.5-4B_TinyLoRA_OPD_q5.safetensors \
    --R_hat 0.71

# All four domains in one go:
for d in math code general sql; do
  python3 scripts/run_pipeline/06_compute_bounds.py \
      --domain $d \
      --adapter models/$d/Qwen3.5-4B_TinyLoRA_OPD_q5.safetensors \
      --R_hat 0.71
done

Replace --R_hat 0.71 with --eval_json runs/<domain>/eval/metrics.json once you have run step 05 to obtain the empirical reward from the actual eval rollouts.

---

Full reproduction

For a single domain (math used as example; substitute code|general|sql for the others):

# 1. Curate the dataset and carve out the m'=4096 eval subsample
bash scripts/run_pipeline/01_build_data.sh math

# 2. Train the rank-64 LoRA teacher via Tinker (~hours on the Tinker backend).
#    Requires TINKER_API_KEY.
bash scripts/run_pipeline/02_train_teacher.sh math

# 3. Train the TinyLoRA student via on-policy distillation (~hours on 2× A100).
#    Set TEACHER_MODEL to the merged teacher weights from step 2.
bash scripts/run_pipeline/03_train_student.sh math opd

# 4. Quantize the unquantized student adapter at k ∈ {2..64}.
bash scripts/run_pipeline/04_quantize.sh math

# 5. Evaluate the q5 student on the eval subsample (~minutes on 1× A100).
bash scripts/run_pipeline/05_eval.sh math \
     runs/math/quantize/Qwen3.5-4B_TinyLoRA_OPD_q5.safetensors

# 6. Compute the Corollary 1 bound from the eval JSON and the quantized adapter.
python3 scripts/run_pipeline/06_compute_bounds.py \
     --domain math \
     --adapter runs/math/quantize/Qwen3.5-4B_TinyLoRA_OPD_q5.safetensors \
     --eval_json runs/math/eval/metrics.json

Each step is independently restartable; outputs land under runs/<domain>/<step>/.

The bundled models/<domain>/Qwen3.5-4B_TinyLoRA_OPD_q5.safetensors files are the canonical checkpoints whose C(h) values match the paper's reported numbers (math 7,726; code 3,426; general 3,710; sql 7,987 bits). Re-quantizing the unquantized adapter via step 04 produces an adapter with the same number of unique levels and a similar empirical entropy, but a numerically different codebook because the original training pipeline applied an additional pre-quantization scaling step that we do not attempt to reproduce in lora_quantize.py.

---

License

The paper's added contributions in this artifact — configs/, data/, scripts/, scripts/run_pipeline/, and the per-domain additions under SkyRL/examples/train/{eurus_rl_math, mix_code, mix_general, synsql}/ and tinker-cookbook/scripts/ — are released under the MIT License.

The vendored SkyRL/ and tinker-cookbook/ snapshots retain their upstream licenses (see SkyRL/LICENSE and tinker-cookbook/LICENSE).

Citation

TBD