import asyncio import dataclasses import hashlib import json import math import time import numpy as np import tinker import torch from tinker_cookbook.completers import TinkerTokenCompleter from tinker_cookbook.rl.data_processing import assemble_training_data from tinker_cookbook.rl.rollouts import do_group_rollout from remote_stages import create from rx_gs1 import ROOT, advantages, dataset, datums, guard_grader, kernel_baselines, service, write_json def correlation(a, b): x, y = np.asarray(a, dtype=np.float64), np.asarray(b, dtype=np.float64) if x.shape != y.shape or x.ndim != 1 or len(x) < 2: raise ValueError("correlation requires equal nonempty vectors") if not np.isfinite(x).all() or not np.isfinite(y).all(): raise ValueError("nonfinite correlation input") if x.std() == 0 or y.std() == 0: return None return float(np.corrcoef(x, y)[0, 1]) def gate(values): limits = dict(executed_group_size_treatment=(1, 1), executed_group_size_baseline=(8, 8), pilot_mean_reward=(0.1, 0.9), kernel_effective_neighbors_p10=(8, math.inf), frac_nonzero_advantage_treatment=(0.5, math.inf), baseline_reward_corr=(0.1, math.inf)) failed = [name for name, (lo, hi) in limits.items() if not isinstance(values.get(name), (int, float)) or not math.isfinite(values[name]) or not lo <= values[name] <= hi] return dict(passed=not failed, failed=failed, calibration_permitted=failed == ["kernel_effective_neighbors_p10"]) def pilot_advantages(groups, features, history, bandwidth): if bandwidth == 0.5: return advantages(groups, "gs1_kernel", features, history) rewards = [r for g in groups for r in g.get_total_rewards()] window = (history + list(zip(features, rewards)))[-256:] baseline, neighbors = kernel_baselines([x for x, _ in window], [r for _, r in window], len(rewards), bandwidth=bandwidth) return [torch.tensor([r - b], dtype=torch.float64) for r, b in zip(rewards, baseline)], baseline, neighbors async def run_arm(arm, steps, bandwidth=0.5): treatment = arm == "gs1_kernel" label = f"treatment-bw{bandwidth}" if treatment else "baseline" path = ROOT / f"activation-pilot-{label}.json" if path.exists(): raise ValueError(f"refusing to overwrite pilot evidence: {path}") checkpoint = json.loads((ROOT / "init-save-response.json").read_text())["path"] record = dict(arm=arm, bandwidth=bandwidth, steps=[], checkpoint=checkpoint, optimizer=dict(source="declared", actual=None, actual_config=None, note="Shared Spindle accepts adam_params; that is not an observation of the active optimizer."), status="running", wall_cap_seconds=600 if treatment else None) write_json(path, record) sc = service() start = time.monotonic() heldout_task = None try: async with asyncio.timeout(600 if treatment else None): ds, _ = await dataset(64 if treatment else 8, 1 if treatment else 8) tc = await create(sc) future = await tc.load_state_async(checkpoint) record["load_response"] = await future.result_async() record["setup_seconds"] = time.monotonic() - start history = [] for index in range(steps): step_start = time.monotonic() timing = {} name = f"rx-pilot-{label}-step{index + 1}" sampler = await tc.save_weights_and_get_sampling_client_async(name=name) timing["save_seconds"] = time.monotonic() - step_start builders = list(ds.get_batch(index)) policy = TinkerTokenCompleter(sampler, max_tokens=512, temperature=1.0) t = time.monotonic() async def feature(builder): prompt, _ = await builder.env_thunk().initial_observation() lp = await sampler.compute_logprobs_async(prompt) finite = [v for v in lp if v is not None] if not finite or not np.isfinite(finite).all(): raise ValueError("missing/nonfinite prompt logprobs") return [math.log(prompt.length), float(np.mean(finite))] prompt_features = await asyncio.gather(*(feature(b) for b in builders)) timing["logprob_feature_seconds"] = time.monotonic() - t heldout_start = time.monotonic() if treatment and index == 0: async def heldout(): groups = await asyncio.gather(*(do_group_rollout( dataclasses.replace(b, num_envs=8), policy) for b in builders)) means = [float(np.mean(g.get_total_rewards())) for g in groups] record["heldout"] = dict(samples_per_prompt=8, used_for_training=False, sampler_save_name=name, mean_rewards=means, rewards=[g.get_total_rewards() for g in groups], seconds=time.monotonic() - heldout_start) return groups heldout_task = asyncio.create_task(heldout()) t = time.monotonic() groups = await asyncio.gather(*(do_group_rollout(b, policy) for b in builders)) timing["rollout_seconds"] = time.monotonic() - t features = [f for f, g in zip(prompt_features, groups) for _ in g.trajectories_G] rewards = [r for g in groups for r in g.get_total_rewards()] if len(rewards) != 64: raise ValueError("pilot batch did not execute 64 rollouts") adv, baseline, neighbors = (pilot_advantages(groups, features, history, bandwidth) if treatment else advantages(groups, arm, features, history)) data, _ = datums(groups, adv) masked, _ = assemble_training_data(groups, adv, include_provenance=False) evidence = ROOT / f"activation-pilot-{label}-step{index + 1}-inputs.json" write_json(evidence, dict(datums=data, masks=[d.loss_fn_inputs["mask"].data for d in masked], sampler_save_name=name, updates_since_sampling=0, history=history)) flat = np.concatenate([a.numpy() for a in adv]) step = dict(step=index + 1, rewards=rewards, features=features, questions=[b.env_thunk().problem for b in builders], executed_group_size=[len(g.trajectories_G) for g in groups], b_i=baseline.tolist(), kernel_effective_neighbors=neighbors.tolist(), advantages=flat.tolist(), frac_nonzero_advantage=float(np.mean(np.abs(flat) > 1e-6)), inputs_ref=evidence.name, inputs_sha256=hashlib.sha256(evidence.read_bytes()).hexdigest(), timings=timing) record["steps"].append(step) write_json(path, record) t = time.monotonic() future = await tc.forward_backward_async(data, loss_fn="importance_sampling") fb = await future.result_async() timing["forward_backward_seconds"] = time.monotonic() - t step["metrics"] = fb.metrics outputs_path = ROOT / f"activation-pilot-{label}-step{index + 1}-outputs.json" write_json(outputs_path, dict(metrics=fb.metrics, outputs=fb.loss_fn_outputs)) step["outputs_ref"] = outputs_path.name step["outputs_sha256"] = hashlib.sha256(outputs_path.read_bytes()).hexdigest() if not math.isfinite(fb.metrics["loss:sum"]): raise ValueError("nonfinite pilot loss") params = tinker.AdamParams(learning_rate=4e-5, beta1=0.9, beta2=0.95, eps=1e-8) t = time.monotonic() future = await tc.optim_step_async(params) step["optim_response"] = await future.result_async() timing["optim_seconds"] = time.monotonic() - t history = (history + list(zip(features, rewards)))[-256:] timing["training_step_seconds"] = time.monotonic() - step_start if heldout_task is not None: await heldout_task heldout_task = None timing["step_with_heldout_seconds"] = time.monotonic() - step_start write_json(path, record) print(label, "step", index + 1, json.dumps(timing), flush=True) record["status"] = "completed" except TimeoutError: record["status"] = "wall_cap_exceeded" raise except Exception as error: record["status"] = "error" record["error"] = str(error) raise finally: if heldout_task is not None: heldout_task.cancel() await asyncio.gather(heldout_task, return_exceptions=True) record["pilot_seconds_before_cleanup"] = time.monotonic() - start await sc.close("success" if record["status"] == "completed" else "errored") record["total_seconds"] = time.monotonic() - start record["client_released"] = True write_json(path, record) print(label, "client released; status", record["status"], flush=True) return record def values_from(treatment, baseline): steps = treatment["steps"] groups = [n for s in steps for n in s["executed_group_size"]] baseline_groups = [n for s in baseline["steps"] for n in s["executed_group_size"]] return dict(executed_group_size_treatment=groups[0] if len(set(groups)) == 1 else None, executed_group_size_baseline=baseline_groups[0] if len(set(baseline_groups)) == 1 else None, pilot_mean_reward=float(np.mean([r for s in steps for r in s["rewards"]])), kernel_effective_neighbors_p10=float(np.percentile([n for s in steps for n in s["kernel_effective_neighbors"]], 10)), frac_nonzero_advantage_treatment=float(np.mean([abs(a) > 1e-6 for s in steps for a in s["advantages"]])), baseline_reward_corr=correlation(steps[0]["b_i"], treatment["heldout"]["mean_rewards"])) async def main(): guard_grader() baseline = await run_arm("grpo_gs8", 1) treatment = await run_arm("gs1_kernel", 2) values = values_from(treatment, baseline) result = dict(values=values, gate=gate(values), bandwidth=0.5, note="Two treatment steps, 128 training rollouts; 512 extra held-out train-prompt samples, no test-set calibration.") write_json(ROOT / "activation-pilot-gate-bw0.5.json", result) print(json.dumps(result, indent=2), flush=True) if result["gate"]["calibration_permitted"]: calibrated = await run_arm("gs1_kernel", 2, bandwidth=1.0) values = values_from(calibrated, baseline) result = dict(values=values, gate=gate(values), bandwidth=1.0, note="Pre-registered bandwidth-only calibration on pilot train prompts; original failed logs retained.") write_json(ROOT / "activation-pilot-gate-bw1.0.json", result) print(json.dumps(result, indent=2), flush=True) if __name__ == "__main__": asyncio.run(main())