🔬Modal Research Exchange

i'm curious about this recent post on group-size-1 RL: <https://www.trajectory.ai/field-notes/you-only-rollout-once>

can we reproduce their results on any public benchmarks, including how GS1 compares to FlashREINFORCE and GRPO?
how does this compare (in terms of final task performance as well as training time) to just training an actual critic?
they don't release their exact choice of kernel function, but claim that you can easily hillclimb to find a recipe yourself. how easy is it to actually find one that is efficient and performs reasonably well compared to GRPO as implied in their results? does this choice of kernel transfer across a diverse set of benchmark tasks, or do you end up having to find a new one per benchmark? (as a starting point they do provide a list of example kernels in their post, some of which have a close enough effective batch size to GS1)
i'd like to see how high of an effective batch size we can hillclimb to offline using the technique they mentioned. this part can be done offline without using spindle

{
  "adjacent_experiment": "",
  "reproducibility_reason": "GS1 only changes the advantage baseline, which is client-side code in the cookbook RL loop. Single-rollout GS1 versus group-of-8 GRPO runs as LoRA RL with AdamW on the shared Spindle, using GSM8K verifiable rewards.",
  "reproducible_on_spindle": true,
  "sources": [
    {
      "brief": {
        "caveats": "No seeds or confidence intervals are stated; the headline differences are small (21.0 vs 21.4 vs 19.5 on AutomationBench; 2.0 vs 0.8 on PRBench; 68 vs 65 on LABBench2). The GS1 kernel is withheld, and 'easy to hillclimb' is asserted, not demonstrated, so the claim that you can match GS1 is untested. Appendix B's crude kernels get 11.5 vs GS1's 13.0 on one benchmark, and on three of four benchmarks the B_eff gains are only 1.02x to 1.16x. The GS1-vs-GRPO comparison is only on AutomationBench, and GS1 is not compared to GRPO on LABBench2 (LABBench2 compares only GS1 vs FlashRF). GRPO sees 4 tasks per step vs 32, so differences conflate baseline quality with task coverage; the authors themselves attribute FlashRF's PRBench win to coverage. The LR-robustness claim comes from one pair of LRs on one benchmark. The B_eff derivation rests on four approximations (negligible cross-rollout covariance, score norm independent of advantage, noise dominating the true gradient, on-policy data), and the calibration is only on the starting policy, so B_eff may drift as training progresses. Harvey LAB results go against FlashRF/GS1 relative to GRPO per Appendix D, and the text is inconsistent about that benchmark ('just as well' vs 'GRPO performs better'). There is no critic comparison, no compute or wall-clock comparison, and the benchmarks rely on LLM graders that could change (the o4-mini grader was swapped for gpt-5.6-luna). Many details are described only in figures that the text does not give numbers for. The advantage is trajectory-level, with no per-turn credit assignment, and there is no ablation of the similarity temperature or other kernel hyperparameters. The method is evaluated on one very large MoE model only; behavior with LoRA on smaller models is unknown.",
        "claim": "GS1 (Group Size 1) is a group-free, critic-free RL method that builds a per-task baseline for each single-rollout task as a similarity-weighted mean of the other rollouts' rewards in the batch. The authors claim this beats FlashREINFORCE's shared batch-mean baseline on heterogeneous-difficulty benchmarks (stability, robustness to learning rate) and matches GRPO without replaying any task.",
        "context": "GS1 builds on REINFORCE++ (batch-level reward normalization), FlashREINFORCE (Hu et al., NVIDIA, 2026; batch-mean baseline, single rollout, async), GRPO (group baselines) and PPO (critic baselines). The theoretical framing leans on Zhou et al. 2026 (GRPO's gradient is a U-statistic, with a scaling bound over tasks and group size) and Appendix C derives when group size 1 beats GRPO under a fixed rollout budget, though the constants are not estimated. The idea text asks the right questions: reproduction of GS1 vs FlashRF vs GRPO, a comparison to a trained critic (not covered by the source), how easy kernel hillclimbing is, whether a kernel transfers across benchmarks, and how high B_eff can go offline. The source's own results suggest that kernel gains are benchmark-dependent: large on AutomationBench, small elsewhere, because they depend on whether prompt features predict difficulty. Modal's team may care because single-rollout, critic-free RL is cheap in rollouts and fits production traces and LoRA, and the offline B_eff metric lets you screen baselines without training runs. The post is a company blog (Trajectory Labs, which sells a research preview), not peer reviewed.",
        "evidence": "All runs use Inkling-Small (276B total, 12B active) with the same rollouts per step (32). GRPO uses 8 rollouts per task, so it covers 4 tasks per step, while FlashRF and GS1 cover 32. Results are on held-out tasks; seed counts and error bars are not stated, and the PRBench text mentions a wider spread across seeds for FlashRF. Harvey LAB: FlashRF improves held-out pass rate 'just as well as GRPO' (figure only, no numbers in the text), though Appendix D says GRPO performs better on Harvey LAB. PRBench: FlashRF improves held-out rubric score by 2.0 points vs 0.8 for GRPO, mostly on out-of-domain tax and cross-border tasks. LABBench2: GS1 improves held-out pass rate from 63% to 68% and stays there, while FlashRF slips to 65% late in training. AutomationBench at the tuned LR of 2e-4: GS1 21.0%, FlashRF 19.5%, GRPO 21.4%. At LR 3e-4, FlashRF falls from 21.7% to 7.2% after roughly 2,500 rollouts, while GS1 and GRPO stay stable. Effective batch size at B=32, FlashRF vs GS1: AutomationBench 9.3 vs 13.0 (1.39x); LABBench2 10.9 vs 12.6 (1.16x); Harvey LAB 23.8 vs 24.4 (1.02x); PRBench 5.2 vs 5.3 (1.03x). Appendix A also lists a difficulty SNR per benchmark (2.36, 1.81, 0.31, 4.91). Appendix B example kernels on AutomationBench at B=32: LOO batch mean 9.4 (1.00x), prompt-length buckets 11.5 (1.22x), TF-IDF 11.5 (1.22x), GS1 13.0 (1.39x), perfect baseline 32.0 (3.40x). No comparison to a trained critic (PPO) and no wall-clock or training-time numbers are reported; the critic is only discussed qualitatively.",
        "method": "GS1 keeps FlashREINFORCE's objective (token-level loss, KL penalty, batch-mean advantage with no std normalization) and changes only the baseline. For rollout i the baseline is b_i = sum_{j!=i} w_ij r_j / sum_{j!=i} w_ij, where w_ij is a kernel on task features available before the rollout (the prompt and starting agent configuration). Self-weight must be zero (w_ii=0), and the baseline may not use rollout i's own reward, actions, response content or length, otherwise the gradient is biased. Advantage is A_i = r_i - b_GS1(x_i), with one scalar advantage per trajectory. The actual GS1 kernel is not public. Kernels are scored offline by effective batch size B_eff = B*sigma_eps^2/(sigma_eps^2 + MSE(b)). Recipe: sample G=8 rollouts per task from the starting policy on N tasks; estimate sigma_eps^2 as mean within-task reward variance; repeatedly draw B=32 tasks with one random rollout each and compute each candidate baseline from the other 31 rollouts only; compare to the task's mean estimated from the 7 held-out calibration rollouts, subtracting sigma_eps^2/(G-1) from the MSE; keep the kernel with the highest B_eff and confirm it beats the batch mean on tasks it wasn't tuned on. Three example label-free kernels are given: leave-one-out batch mean; 5 quantile buckets of log prompt length, using the mean reward of other recent rollouts (this batch plus the previous 16) in the same bucket; and TF-IDF similarity-weighted mean with w_ij = exp(cos(x_i,x_j)/T) on word TF-IDF vectors of prompts, T=0.1.",
        "reproduction": "Offline B_eff hillclimb (no Spindle needed): you need a calibration set of G=8 rollouts on N tasks from the starting policy, with scalar rewards and prompts. Compute sigma_eps^2, then simulate many B=32 batches and score each kernel by B_eff with the (G-1) noise correction. Candidates to try: the three example kernels, plus embedding-based kernels (sentence embeddings, kNN or learned difficulty predictors), prompt features, and recency windows over earlier batches. Validate on held-out tasks. To check how far B_eff can climb, compare against the perfect-baseline ceiling of 32 at B=32. To test transfer across benchmarks, run the same kernel on several task sets and compare per-benchmark B_eff gains. Public math, code or tool-use tasks with verifiable rewards and heterogeneous difficulty would work. Their benchmarks are mostly rubric- or LLM-graded (gpt-5.6-luna, claude-sonnet-4-5) and large; AutomationBench (code assertions, 600 public + 200 simple tasks) is the cheapest to reproduce. Online runs: SFT-free LoRA RL on an open model (e.g. Qwen3.5 or gpt-oss) with the FlashRF loss (batch-mean or leave-one-out advantage, token-level loss, KL penalty, no std normalization) and GS1 as a drop-in baseline change, compared against GRPO at G=8 with matched rollouts per step (32), and an LR sweep (e.g. 2e-4 vs 3e-4 on LoRA, noting those LRs came from their full-model setup and won't transfer directly). Fair evaluation needs multiple seeds. A critic baseline would require adding a value head or separate model, which Spindle's Tinker-style API may not support, so a critic comparison may need a custom or offline setup. Not stated in the source: optimizer, KL coefficient, max sequence length, steps, total rollouts, seeds, learning-rate schedule, and the fraction of data used for calibration versus training. Several things won't carry over directly: the 276B-parameter model and full-parameter or LoRA details (not stated), the exact GS1 kernel, graders requiring proprietary models, and asynchronous-setting claims for FlashRF. Recomputing the kernel every step or adapting it as the policy shifts is something the source doesn't address; its calibration is on the starting policy only.",
        "title": "You Only Rollout Once"
      },
      "error": null,
      "full_text": true,
      "kind": "page",
      "title": "You Only Rollout Once",
      "url": "https://www.trajectory.ai/field-notes/you-only-rollout-once"
    }
  ],
  "spec": {
    "activation_gate": {
      "applicable": true,
      "checks": [
        {
          "measures": "B_eff ratio of the TF-IDF kernel to the LOO batch mean at B=32. Computed with the post's recipe on a calibration set of G=8 rollouts on 64 held-out GSM8K train tasks from the starting policy, including the sigma_eps^2/(G-1) correction.",
          "metric": "offline_beff_ratio_tfidf_vs_loo_mean",
          "min": 1.03
        },
        {
          "measures": "Mean |b_GS1_i - LOO batch mean_i| over the first 3 pilot steps of the GS1 arm.",
          "metric": "pilot_mean_abs_gs1_minus_batchmean_baseline",
          "min": 0.01
        },
        {
          "max": 0.0,
          "measures": "Maximum w_ii used in any GS1 baseline.",
          "metric": "self_weight_max",
          "min": 0.0
        },
        {
          "max": 32,
          "measures": "Logged rollouts per step in each arm.",
          "metric": "executed_rollouts_per_step",
          "min": 32
        },
        {
          "max": 1,
          "measures": "Logged group_size in the GS1 arm.",
          "metric": "executed_group_size_gs1",
          "min": 1
        },
        {
          "max": 8,
          "measures": "Logged group_size in the GRPO arm.",
          "metric": "executed_group_size_grpo",
          "min": 8
        },
        {
          "measures": "Fraction of GS1 rollouts with |A_i| > 1e-6 in the pilot.",
          "metric": "nonzero_advantage_fraction_gs1",
          "min": 0.3
        }
      ],
      "reason": "GS1 only matters if the kernel baseline differs from a plain batch mean and predicts per-task difficulty better than it. GSM8K has milder prompt-difficulty structure than AutomationBench, so a pilot measures this before full training. Executed-config predicates confirm the arms really differ only in grouping/advantage at 32 rollouts per step."
    },
    "arms": [
      {
        "advantage": "group_mean",
        "group_size": 8,
        "groups_per_batch": 4,
        "lr": 3e-05,
        "name": "grpo_g8",
        "seed": 0,
        "steps": 60
      },
      {
        "advantage": "tfidf_kernel_loo_T0.1",
        "group_size": 1,
        "groups_per_batch": 32,
        "lr": 3e-05,
        "name": "gs1_tfidf_g1",
        "seed": 0,
        "steps": 60
      }
    ],
    "base_model": "Qwen/Qwen3.5-2B",
    "baseline": "GRPO via the stock math_rl recipe: group_size=8 and groups_per_batch=4 (32 rollouts per step), group-mean-centered advantage. Same model, LoRA rank, lr, steps, max_tokens, KL=0, seed and eval as the GS1 arm.",
    "batch_size": 32,
    "budget": {
      "gpu_hours": 2.0,
      "max_wall_minutes": 120
    },
    "command": "Baseline arm (stock GRPO):\npython -m tinker_cookbook.recipes.math_rl.train model_name=Qwen/Qwen3.5-2B base_url=$TINKER_BASE_URL env=gsm8k group_size=8 groups_per_batch=4 learning_rate=3e-5 lora_rank=32 max_tokens=512 max_steps=60 kl_penalty_coef=0 log_path=/tmp/rx-gs1/grpo_s0\n\nTreatment arm: the same recipe with group_size=1 groups_per_batch=32 log_path=/tmp/rx-gs1/gs1_s0, run through a thin local wrapper. The wrapper monkeypatches tinker_cookbook.rl's compute_advantages; this replacement is the only custom part:\n\nfrom sklearn.feature_extraction.text import TfidfVectorizer\nimport numpy as np\nvec = TfidfVectorizer().fit(train_questions)\n\ndef gs1_adv(prompts, rewards, T=0.1):\n    X = vec.transform(prompts)\n    S = (X @ X.T).toarray()\n    W = np.exp(S / T)\n    np.fill_diagonal(W, 0.0)\n    b = (W @ rewards) / W.sum(1)\n    return rewards - b\n\nThe wrapper logs b, LOO batch mean and max w_ii each step.\n\nOrder of work:\n1. Correctness preflight. With the treatment disabled (group-mean baseline, G=8), the new advantage path must reproduce the stock cookbook advantages and loss/gradient on a fixed batch. With it enabled, it must match the float64 numpy reference within the stated tolerances.\n2. Activation-failure plan. First run offline forced-activation tests: a synthetic batch with prompt clusters and known rewards must yield the analytic baseline. A forced-nonactivation test (identical prompts) must give b equal to the LOO mean. Then run the offline B_eff calibration: 64 tasks x 8 rollouts at max_tokens 512, capped at 10 minutes and kept separate from the test split. If B_eff ratio < 1.03 the gate fails, and the pilot record is kept. The only pre-approved fallback is the post's prompt-length-bucket kernel, scored on the same calibration set, adopted only via a recorded revised spec before any training. The gate is never relaxed.\n3. Timing pilot: 3 steps per arm, plus a 100-item eval.\n4. Full arms.\n5. Eval: score each final sampler checkpoint greedily on all 1319 GSM8K test items with the same grader. Save per-item correctness for the paired analysis.",
    "correctness_preflight": {
      "paths": [
        {
          "baseline_revision": "UNRESOLVED: pin the exact tinker-cookbook commit SHA whose math_rl group-mean advantage is the stock GRPO baseline",
          "id": "gs1_advantage_estimator",
          "kind": "loss",
          "reference_revision": "UNRESOLVED: pin the commit SHA of a standalone float64 numpy reference of the post's TF-IDF kernel baseline (w_ii=0, T=0.1), reviewed against the post text",
          "tolerances": {
            "baseline_gradient": {
              "atol": 1e-05,
              "min_cosine": 0.9999,
              "rtol": 0.001
            },
            "baseline_loss": {
              "atol": 1e-06,
              "rtol": 1e-05
            },
            "reference_gradient": {
              "atol": 1e-05,
              "min_cosine": 0.9999,
              "rtol": 0.001
            },
            "reference_loss": {
              "atol": 1e-06,
              "rtol": 1e-05
            }
          }
        }
      ],
      "required": true,
      "status": "unresolved: revisions must be pinned before any full run"
    },
    "dataset": {
      "hf_path": "openai/gsm8k",
      "n_examples": 7473,
      "name": "gsm8k",
      "split": "main/train (training prompts), main/test (eval)"
    },
    "decision_rule": {
      "alpha": 0.05,
      "criteria": [
        {
          "bounds": [
            0,
            1
          ],
          "id": "primary",
          "margin": -0.05,
          "metric": "accuracy",
          "weights": {
            "baseline": -1,
            "treatment": 1
          }
        }
      ],
      "estimand": "Paired difference in GSM8K test exact-match accuracy, GS1 (TF-IDF, G=1) minus GRPO (G=8), at matched 32 rollouts per step after 60 steps on Qwen3.5-2B LoRA. Larger favors GS1; noninferiority.",
      "method": "paired_hoeffding",
      "seeds": [
        0
      ],
      "unit": "fixed_checkpoint_items",
      "version": 1
    },
    "eval": {
      "aggregation": null,
      "benchmark": "GSM8K test (openai/gsm8k main/test)",
      "metric": "exact-match accuracy (numeric_exact on final answer, greedy, max_tokens 512)",
      "n_examples": 1319
    },
    "experiment_features": {
      "adaptive_selection": false,
      "checkpoint_restore": false,
      "custom_loss": true,
      "custom_optimizer": false,
      "generated_training_data": false,
      "heldout_loss": false,
      "learned_generator": false,
      "paired_comparison": true,
      "paired_training_arms": true,
      "warmup_adaptation": false
    },
    "feasibility_check": {
      "gpu_hours": null,
      "known_gpu_hours": 0,
      "phases": {},
      "problems": [
        "Measured feasibility not recorded"
      ],
      "status": "unknown",
      "wall_seconds": null
    },
    "feasibility_required": true,
    "generation_audit": {
      "answer_comparison": "numeric_exact",
      "max_eval_overlaps": 0,
      "max_label_errors": 0,
      "min_sample_size": 100,
      "required": false
    },
    "hypothesis": "At 32 rollouts per step for 60 steps, GS1 with the post's TF-IDF kernel baseline (group size 1, 32 tasks per step) is within 5 points of GRPO (group size 8, 4 tasks per step) on GSM8K test exact-match accuracy, with a lower confidence bound above -0.05.",
    "lora_rank": 32,
    "max_seq_len": 1024,
    "method": "This is a scaled-down version of the post's AutomationBench GS1-vs-GRPO comparison. The served Qwen/Qwen3.5-2B stands in for the closed Inkling-Small. The policy is a LoRA (rank 32) trained from the base model with AdamW on the shared Spindle; no fork is needed. The cookbook math_rl GSM8K environment, renderer, reward (numeric exact match) and importance-sampling loss are identical in both arms.\n\nGRPO arm (baseline): stock math_rl recipe with group_size=8 and groups_per_batch=4, which gives 32 rollouts per step. Advantage is reward minus the group mean.\n\nGS1 arm (treatment): group_size=1 and groups_per_batch=32, also 32 rollouts per step. The only change is a client-side advantage function. For rollout i the baseline is b_i = sum_{j!=i} w_ij r_j / sum_{j!=i} w_ij, using in-batch rollouts only. The weights are w_ij = exp(cos(x_i, x_j)/0.1), where x is the word TF-IDF vector of the prompt text. The TF-IDF vocabulary and IDF are fit once on the GSM8K train questions before training. The baseline enforces w_ii = 0, and it uses no response, length or reward information from rollout i. Advantage is A_i = r_i - b_i, one scalar per trajectory, with no std normalization. The actual GS1 kernel is withheld, so this is the post's published example kernel (B_eff 1.22x on AutomationBench).\n\nShared settings and deviations from the source:\n- KL penalty is 0 in both arms, so the advantage estimator is the only difference. FlashRF's KL penalty is omitted; this is a recorded deviation.\n- Both arms use the same lr=3e-5, max_tokens=512, 60 steps, the same LoRA seed and the same training-question shuffle.\n\nOut of scope for this spec:\n- FlashRF batch-mean baseline arm: a follow-up spec.\n- Trained-critic/PPO arm: needs a value model, so a separate design.\n- Offline B_eff kernel hillclimb and cross-benchmark kernel transfer: offline work, no Spindle needed.\n\nEvaluation: greedy decoding on the full GSM8K test split (1319 items) from each final sampler checkpoint, paired by item.",
    "optimizer": {
      "lr": 3e-05,
      "name": "adam"
    },
    "preflight": {
      "probe": "Create both training clients: Qwen3.5-2B, LoRA rank 32, AdamW (tinker.AdamParams lr=3e-5, beta1 0.9, beta2 0.95, eps 1e-8). For each, run one sampled rollout batch (G=8x1 task for GRPO; G=1x1 task with a dummy 2-prompt batch for GS1), one forward_backward with the importance-sampling loss and one optim_step. Then run the offline B_eff calibration before either full arm.",
      "riskiest_assumption": "That a TF-IDF prompt-similarity kernel predicts per-question difficulty on GSM8K better than the batch mean. If B_eff gain is near 1.0x, GS1 reduces to FlashRF and the comparison only tests task coverage, not the baseline.",
      "tuning": "No per-arm LR grid. The arms share optimizer and LR and differ only in advantage estimator/grouping, so the same lr=3e-5 is used for both and no hyperparameter is selected on results. The post's LR-robustness claim (stability at a higher LR) is deferred to a follow-up spec with a prespecified two-LR grid per arm, selected on a GSM8K train-holdout split."
    },
    "primary_endpoint": "gsm8k_test_exact_match_accuracy_final_checkpoint",
    "restoration_preflight": {
      "applicable": true,
      "fixed_batch_required": false,
      "groups": [
        {
          "arms": [
            {
              "name": "grpo_g8",
              "optimizer": "adam",
              "optimizer_config": {
                "beta1": 0.9,
                "beta2": 0.95,
                "eps": 1e-08,
                "learning_rate": 3e-05,
                "lora_rank": 32,
                "seed": 0
              }
            },
            {
              "name": "gs1_tfidf_g1",
              "optimizer": "adam",
              "optimizer_config": {
                "beta1": 0.9,
                "beta2": 0.95,
                "eps": 1e-08,
                "learning_rate": 3e-05,
                "lora_rank": 32,
                "seed": 0
              }
            }
          ]
        }
      ],
      "reason": "Both trained arms start from fresh LoRA on Qwen3.5-2B with the same seed. Identity is verified by computing logprobs on a fixed tokenized batch of 4 GSM8K prompts through each new training client before the first update. A mismatch over atol 1e-4 aborts the run. The active AdamW config is checked after creation. No restores occur."
    },
    "spec_preflight": {
      "findings": [],
      "status": "passed"
    },
    "steps": 60,
    "success_criterion": "GS1 is validated as matching GRPO if the simultaneous 95% paired Hoeffding interval on (GS1 accuracy - GRPO accuracy) over the 1319 GSM8K test items has its lower bound above -0.05. It fails if the upper bound is below -0.05. Otherwise the result is inconclusive. This is conditional on the seed-0 checkpoints only.",
    "title": "GS1 (TF-IDF kernel, G=1) vs GRPO (G=8) LoRA RL on GSM8K, Qwen3.5-2B",
    "training_type": "rl"
  },
  "summary": "GS1 replaces GRPO's 8-rollout group baseline with a TF-IDF prompt-similarity baseline built from one rollout per task. We compare GSM8K test accuracy at equal rollout budget.",
  "tags": [
    "rl",
    "paper",
    "grpo",
    "advantage-baseline",
    "lora"
  ],
  "title": "Does single-rollout GS1 match GRPO on GSM8K at 32 rollouts per step?",
  "value_estimate": 4,
  "value_rationale": "Single-rollout, critic-free RL that matches GRPO would cut rollout cost and fit production traces for Spindle LoRA users. This first test checks whether the published example kernel actually works off the authors' closed model and benchmarks."
}