{
  "title": "GS1 kernel-baseline RL vs GRPO at an equal rollout budget on GSM8K",
  "hypothesis": "At 64 rollouts per step for 40 steps on Qwen3.5-2B, GS1 RL with a kernel-weighted cross-prompt baseline is noninferior to GRPO (group_size=8) on GSM8K test exact-match accuracy. The noninferiority margin is 3 points.",
  "method": "Both arms train a rank-32 LoRA on Qwen/Qwen3.5-2B from the same base. They use one client-side RL loop built on the tinker-cookbook rl loop, with the stock importance_sampling loss and AdamW. The prompts come from GSM8K train with the stock math_rl renderer and the verifiable exact-answer reward. Each step uses 64 rollouts, max_tokens=512 and temperature 1.0.\n\nThe arms differ only in the advantage estimator:\n- GRPO baseline: 8 prompts x 8 rollouts. Advantage = reward minus group mean, matching the cookbook's default.\n- GS1 treatment: 64 distinct prompts x 1 rollout. Advantage = r_i minus b_i. Here b_i is a kernel-weighted mean of rewards from other prompts in a sliding window of the last 256 rollouts. The window covers the current batch excluding i, plus the previous steps.\n\nThe kernel is fixed before running: a Gaussian kernel exp(-||phi_i - phi_j||^2 / (2*0.5^2)) on a 2-d standardized prompt feature phi = (log prompt-token length, sampling-policy mean token logprob of the prompt). The prompt logprob is computed by the sampling client before rollout. If kernel mass is below 1e-3, the estimator falls back to the window mean.\n\nThe Trajectory post could not be fetched and does not release its kernel, so this kernel is our own pre-registered choice and not the paper's exact method. Kernel hillclimbing, FlashREINFORCE, the trained-critic comparison and cross-task transfer are out of scope for this first node. Each is a follow-up spec.\n\nThe shared Spindle is enough; no fork is needed. Evaluation is greedy decoding on the full GSM8K test split (1319 items), scored by exact match, for the final checkpoint of each arm.",
  "base_model": "Qwen/Qwen3.5-2B",
  "lora_rank": 32,
  "training_type": "rl",
  "dataset": {
    "name": "gsm8k",
    "hf_path": "openai/gsm8k",
    "split": "main/train",
    "n_examples": 2560
  },
  "optimizer": {
    "name": "adam",
    "lr": 4e-05
  },
  "steps": 40,
  "batch_size": 64,
  "max_seq_len": 1024,
  "eval": {
    "benchmark": "GSM8K test (openai/gsm8k main/test)",
    "metric": "exact-match accuracy, greedy, final checkpoint",
    "n_examples": 1319,
    "aggregation": null
  },
  "primary_endpoint": "gsm8k_test_exact_match_accuracy_final_checkpoint",
  "arms": [
    {
      "name": "grpo_gs8",
      "advantage": "group_mean",
      "group_size": 8,
      "groups_per_batch": 8,
      "lr": 4e-05,
      "seed": 0
    },
    {
      "name": "gs1_kernel",
      "advantage": "kernel_cross_prompt_baseline",
      "group_size": 1,
      "groups_per_batch": 64,
      "lr": 4e-05,
      "seed": 0
    }
  ],
  "feasibility_required": true,
  "baseline": "GRPO via the same client loop with group_size=8, groups_per_batch=8 and group-mean advantages. It uses the same model, LoRA rank, LR, steps, rollouts per step, max_tokens, renderer, reward, seed and GSM8K eval.",
  "experiment_features": {
    "custom_loss": true,
    "custom_optimizer": false,
    "paired_comparison": true,
    "paired_training_arms": true,
    "checkpoint_restore": false,
    "warmup_adaptation": false,
    "generated_training_data": false,
    "learned_generator": false,
    "heldout_loss": false,
    "adaptive_selection": false
  },
  "correctness_preflight": {
    "required": true,
    "status": "pinned for measured preflight; see preflight.json",
    "paths": [
      {
        "id": "gs1_kernel_advantage_loop",
        "kind": "loss",
        "baseline_revision": "f79fdb2ba4cd967ef2807ba43097c52985f9f17f",
        "reference_revision": "git:ecbba04ce54322679f49328ee98aeebb1dcaa705;sha256:757d81c606f9793f259c6afd5ffe6f5e52382328560145b682d84d814da9a50e",
        "tolerances": {
          "baseline_loss": {
            "atol": 1e-06,
            "rtol": 1e-05
          },
          "baseline_gradient": {
            "atol": 1e-05,
            "rtol": 0.0001,
            "min_cosine": 0.99999
          },
          "reference_loss": {
            "atol": 1e-06,
            "rtol": 1e-05
          },
          "reference_gradient": {
            "atol": 1e-05,
            "rtol": 0.0001,
            "min_cosine": 0.99999
          }
        }
      }
    ]
  },
  "activation_gate": {
    "applicable": true,
    "reason": "A GS1 baseline only differs from a constant baseline if the kernel draws on several neighbors with varied rewards, and the prompts give a non-degenerate reward. The executed group size must actually be 1 for the treatment and 8 for the baseline.",
    "checks": [
      {
        "metric": "executed_group_size_treatment",
        "measures": "rollouts per distinct prompt in logged treatment batches",
        "one_of": [
          "1"
        ]
      },
      {
        "metric": "executed_group_size_baseline",
        "measures": "rollouts per distinct prompt in logged baseline batches",
        "one_of": [
          "8"
        ]
      },
      {
        "metric": "pilot_mean_reward",
        "measures": "mean training reward over the first 2 treatment steps",
        "min": 0.1,
        "max": 0.9
      },
      {
        "metric": "kernel_effective_neighbors_p10",
        "measures": "10th percentile over rollouts of (sum w)^2/sum w^2 for the kernel weights",
        "min": 8.0
      },
      {
        "metric": "frac_nonzero_advantage_treatment",
        "measures": "fraction of treatment rollouts with |advantage|>1e-6",
        "min": 0.5
      },
      {
        "metric": "baseline_reward_corr",
        "measures": "Pearson correlation of b_i with a held-out 8-sample per-prompt mean reward, over 64 pilot prompts",
        "min": 0.1
      }
    ]
  },
  "restoration_preflight": {
    "applicable": true,
    "reason": "Both trained arms start from the same base LoRA initialization within seed 0.",
    "fixed_batch_required": true,
    "groups": [
      {
        "arms": [
          {
            "name": "grpo_gs8",
            "optimizer": "adam",
            "optimizer_config": {
              "lr": 4e-05,
              "beta1": 0.9,
              "beta2": 0.95,
              "eps": 1e-08,
              "lora_rank": 32,
              "seed": 0
            }
          },
          {
            "name": "gs1_kernel",
            "optimizer": "adam",
            "optimizer_config": {
              "lr": 4e-05,
              "beta1": 0.9,
              "beta2": 0.95,
              "eps": 1e-08,
              "lora_rank": 32,
              "seed": 0
            }
          }
        ]
      }
    ]
  },
  "success_criterion": "GS1 minus GRPO GSM8K test exact-match accuracy, as a lower 95% Hoeffding bound over 1319 paired items, must exceed -0.03 (noninferiority margin of 3 points). If the interval straddles -0.03, the pilot is inconclusive and is not counted as a win.",
  "decision_rule": {
    "version": 1,
    "estimand": "Paired difference in GSM8K test exact-match accuracy (GS1 kernel minus GRPO) for these fixed seed-0 final checkpoints; larger is better",
    "unit": "fixed_checkpoint_items",
    "seeds": [
      0
    ],
    "alpha": 0.05,
    "method": "paired_hoeffding",
    "criteria": [
      {
        "id": "primary",
        "metric": "accuracy",
        "weights": {
          "treatment": 1,
          "baseline": -1
        },
        "margin": -0.03,
        "bounds": [
          0,
          1
        ]
      }
    ]
  },
  "budget": {
    "gpu_hours": 2.0,
    "max_wall_minutes": 120
  },
  "generation_audit": {
    "required": false,
    "answer_comparison": "numeric_exact",
    "min_sample_size": 100,
    "max_label_errors": 0,
    "max_eval_overlaps": 0
  },
  "command": "Custom client loop (python rx_gs1.py --arm {grpo_gs8|gs1_kernel}); the only custom part is the compute_advantages function.\n\nSetup:\n  import os, tinker\n  sc = tinker.ServiceClient(base_url=os.environ['TINKER_BASE_URL'])\n  tc = sc.create_lora_training_client(base_model='Qwen/Qwen3.5-2B', rank=32, seed=0)\n  Use the tinker_cookbook math_rl gsm8k env builder for prompts, renderer and reward.\n\nEach of 40 steps:\n  1. sampler = tc.save_weights_and_get_sampling_client(...).\n  2. Draw prompts from a fixed seed-0 shuffled order: GRPO takes 8 prompts x 8 samples, GS1 takes 64 prompts x 1 sample. max_tokens=512, temperature=1.0.\n  3. Grade with the env reward.\n  4. Compute advantages: GRPO uses r minus group mean. GS1 uses r_i minus a Gaussian-kernel-weighted mean over a 256-rollout window excluding i, with phi = (log prompt length, sampler mean prompt-token logprob) and bandwidth 0.5.\n  5. Build Datums with sampling logprobs and advantages.\n  6. tc.forward_backward(datums, loss_fn='importance_sampling').\n  7. tc.optim_step(tinker.AdamParams(learning_rate=4e-5)).\n\nAfter training: run save_weights_for_sampler, then greedy-decode all 1319 GSM8K test items for each arm. Write per-item correctness to /tmp/rx-gs1/{arm}/eval.jsonl, and score the paired Hoeffding rule from those files.\n\nActivation-failure plan:\n- (a) Offline first: unit-test compute_advantages in float64 on synthetic rewards. Check forced activation (heterogeneous phi gives nonconstant b_i) and forced nonactivation (huge bandwidth gives b_i equal to the window mean, i.e. a batch-mean REINFORCE baseline). Run gradient checks of the new loop with the treatment disabled against stock cookbook GRPO advantages on one fixed batch.\n- (b) A 2-step treatment pilot logs the gate metrics.\n- (c) If kernel_effective_neighbors fails, the single pre-registered calibration is bandwidth in {0.5, 1.0}, chosen by the gate on pilot training prompts only, never on GSM8K test. The cap is 2 pilot steps and 10 wall minutes.\n- (d) Keep the original failed pilot logs. Any other setting change requires a recorded revised spec before execution. Gates are never relaxed.\n\nNo full run happens while correctness_preflight is unresolved.",
  "preflight": {
    "riskiest_assumption": "The custom GS1 advantage path is correct and engages. Specifically, the new loop reproduces stock GRPO advantages when the treatment is disabled, and the kernel baseline gives a nonconstant, reward-predictive b_i. Otherwise a GS1-vs-GRPO gap reflects bugs or a degenerate batch-mean baseline, not the method.",
    "probe": "For each arm, create the training client with adam, rank=32 and seed=0. Run one forward_backward with loss_fn='importance_sampling' on one sampled GSM8K example carrying that arm's computed advantage, then one optim_step(tinker.AdamParams(learning_rate=4e-5)). Confirm both succeed and the loss is finite. Then run the offline advantage unit tests and the 2-step activation pilot.",
    "tuning": "Both arms use the same optimizer (AdamW) and a fixed LR of 4e-5 taken from cookbook math_rl defaults for small Qwen models. There is no per-arm grid, because the hypothesis concerns the advantage estimator at a fixed recipe. An LR-sensitivity follow-up would give each arm a 3-value grid selected on a GSM8K train holdout, with the grid-edge rule applied."
  },
  "feasibility": {
    "config": null,
    "effective_config": {
      "base_model": "Qwen/Qwen3.5-2B",
      "arms": [
        "grpo_gs8",
        "gs1_kernel"
      ],
      "steps": 40,
      "batch_size": 64,
      "max_tokens": 512,
      "max_seq_len": 1024,
      "n_eval": 1319
    },
    "primary_endpoint": "gsm8k_test_exact_match_accuracy_final_checkpoint",
    "source": null,
    "measurements": null,
    "generation": null,
    "recovery": null
  },
  "feasibility_check": {
    "status": "unknown",
    "problems": [
      "Effective configuration not recorded",
      "Pilot evidence source not recorded",
      "Representative training timing not recorded",
      "Representative setup timing not recorded",
      "Representative evaluation timing not recorded",
      "Positive recovery reserve not recorded"
    ],
    "wall_seconds": null,
    "gpu_hours": null,
    "known_gpu_hours": 0,
    "phases": {},
    "generation_counts": null,
    "generation_applicability": {
      "check": "generation_audit",
      "status": "not_applicable",
      "enforce": false,
      "reasons": [
        "spec.generation_audit opts out and no feature it covers is true"
      ],
      "features": {
        "generated_training_data": false
      }
    }
  },
  "spec_preflight": {
    "status": "passed",
    "findings": []
  }
}