sketch_thompson.nx source
↩ module page · 292 lines · 9487 B
1// sketch_thompson.nx -- Thompson Sampling (Bernoulli-Beta bandit).
2//
3// Thompson 1933 / Russo-Van Roy 2018 modern treatment. Bayesian
4// alternative to UCB1 (sketch_ucb1.nx): instead of a confidence-
5// bound upper-bonus, sample from each arm's posterior and pick the
6// argmax. Empirically tighter regret on stochastic Bernoulli arms;
7// matches UCB asymptotic bounds.
8//
9// CORE OPERATION:
10// For arm i with s_i successes, f_i failures:
11// posterior ~ Beta(s_i + 1, f_i + 1) (uniform Beta(1,1) prior)
12// Per decision:
13// for each arm i: sample x_i ~ Beta(s_i+1, f_i+1)
14// pull arg max x_i
15//
16// SAMPLING WITHOUT f64 (the technical chokepoint):
17// Beta(α, β) for integer α, β equals the distribution of the α-th
18// smallest of α + β − 1 uniforms in [0, 1]. (David-Nagaraja
19// "Order Statistics", standard textbook result.) Exact, no log
20// or special functions needed. Cost O(α + β) per sample.
21//
22// HYBRID FOR LARGE (s + f):
23// - exact order statistic when s + f + 1 ≤ NX_THOMPSON_EXACT_MAX
24// - Normal approximation when larger (cost O(1) instead of O(n)).
25// Beta(α, β) ≈ N(μ, σ²) with μ = α/(α+β), σ² = μ(1−μ)/(α+β+1).
26// Normal sampled via Irwin-Hall sum of 12 uniforms (CLT).
27//
28// Both paths produce Q14 fixed-point samples in [0, Q14] = [0, 1].
29//
30// EXTENDS NISHI-SUBSTRATE BEYOND DATASKETCHES (DS is summarization-
31// only; bandits live elsewhere). Composes against sketch_ucb1 +
32// sketch_epsilon_greedy as the bandit-family triple.
33//
34// LOSSLESS-LANGUAGE DISCIPLINE: nx_thompson_query returns
35// ApproxI64 with NX_ENV_REL_STDDEV = 1/sqrt(count_i + 1), matching
36// posterior shrinkage as samples accumulate.
37
38import "syscalls.nx"
39import "sketch_types.nx"
40
41const NX_THOMPSON_ARMS_MIN: i64 = 2
42const NX_THOMPSON_ARMS_MAX: i64 = 256
43const NX_THOMPSON_Q14: i64 = 16384
44const NX_THOMPSON_EXACT_MAX: i64 = 64 // cutoff for order-stat sampler
45
46// LCG (matches KLL / Reservoir family).
47const NX_THOMPSON_LCG_A: i64 = 1103515245
48const NX_THOMPSON_LCG_C: i64 = 12345
49const NX_THOMPSON_LCG_MOD: i64 = 0x7FFFFFFF
50
51struct Thompson {
52 n_arms: i64,
53 successes: *i64,
54 failures: *i64,
55 total_pulls: i64,
56 rng_state: i64,
57 scratch: *i64, // size NX_THOMPSON_EXACT_MAX
58}
59
60// === construction =================================================
61
62func nx_thompson_alloc(n_arms: i64, seed: i64) -> *Thompson {
63 if n_arms < NX_THOMPSON_ARMS_MIN { return 0 as *Thompson }
64 if n_arms > NX_THOMPSON_ARMS_MAX { return 0 as *Thompson }
65 let raw: *u8 = sys_mmap(56)
66 let t: *Thompson = raw as *Thompson
67 let s_raw: *u8 = sys_mmap(n_arms * 8)
68 t.successes = s_raw as *i64
69 let f_raw: *u8 = sys_mmap(n_arms * 8)
70 t.failures = f_raw as *i64
71 var i: i64 = 0
72 while i < n_arms {
73 t.successes[i] = 0
74 t.failures[i] = 0
75 i = i + 1
76 }
77 t.n_arms = n_arms
78 t.total_pulls = 0
79 t.rng_state = seed | 1
80 let scratch_raw: *u8 = sys_mmap(NX_THOMPSON_EXACT_MAX * 8)
81 t.scratch = scratch_raw as *i64
82 return t
83}
84
85// === LCG + uniform Q14 ============================================
86
87func nx_thompson_rng_next(t: *Thompson) -> i64 {
88 let next: i64 = ((t.rng_state * NX_THOMPSON_LCG_A) + NX_THOMPSON_LCG_C) & NX_THOMPSON_LCG_MOD
89 t.rng_state = next
90 return next
91}
92
93// Uniform in [0, Q14) = 14 low bits of LCG output.
94func nx_thompson_uniform_q14(t: *Thompson) -> i64 {
95 let r: i64 = nx_thompson_rng_next(t)
96 return r & (NX_THOMPSON_Q14 - 1)
97}
98
99// === isqrt (shared shape with stream_stats / ucb1) ================
100
101func nx_thompson_isqrt(x: i64) -> i64 {
102 if x < 0 { return 0 }
103 if x == 0 { return 0 }
104 if x < 4 { return 1 }
105 var g: i64 = (x >> 1) + 1
106 var iter: i64 = 0
107 while iter < 64 {
108 let next_g: i64 = (g + x / g) / 2
109 if next_g >= g { iter = 64 }
110 if next_g < g {
111 g = next_g
112 iter = iter + 1
113 }
114 }
115 return g
116}
117
118// === order-statistic Beta sampler =================================
119//
120// Sample Beta(alpha, beta) for integer alpha, beta with
121// (alpha + beta - 1) ≤ NX_THOMPSON_EXACT_MAX. Returns Q14 sample.
122
123func nx_thompson_sort_scratch(t: *Thompson, n: i64) -> i64 {
124 var i: i64 = 1
125 while i < n {
126 let cur: i64 = t.scratch[i]
127 var j: i64 = i - 1
128 var done: i64 = 0
129 while done == 0 {
130 if j < 0 { done = 1 }
131 if done == 0 {
132 if t.scratch[j] <= cur { done = 1 }
133 if done == 0 {
134 t.scratch[j + 1] = t.scratch[j]
135 j = j - 1
136 }
137 }
138 }
139 t.scratch[j + 1] = cur
140 i = i + 1
141 }
142 return 0
143}
144
145func nx_thompson_sample_order_stat(t: *Thompson, alpha: i64, n_uniform: i64) -> i64 {
146 var i: i64 = 0
147 while i < n_uniform {
148 t.scratch[i] = nx_thompson_uniform_q14(t)
149 i = i + 1
150 }
151 nx_thompson_sort_scratch(t, n_uniform)
152 // alpha-th smallest = index (alpha - 1) zero-based
153 return t.scratch[alpha - 1]
154}
155
156// === Irwin-Hall normal approximation ==============================
157//
158// Z ≈ (U_1 + ... + U_12) − 6·Q14 where U_i ~ Uniform[0, Q14)
159// Mean 0, variance ≈ Q14² in the high-resolution limit. Returns
160// scaled to match a standard normal in Q14.
161
162func nx_thompson_clt_normal_q14(t: *Thompson) -> i64 {
163 var sum: i64 = 0
164 var i: i64 = 0
165 while i < 12 {
166 sum = sum + nx_thompson_uniform_q14(t)
167 i = i + 1
168 }
169 return sum - 6 * NX_THOMPSON_Q14
170}
171
172// === main sampler =================================================
173//
174// Returns one sample from Beta(s + 1, f + 1) in Q14.
175
176func nx_thompson_sample_q14(t: *Thompson, arm: i64) -> i64 {
177 let s: i64 = t.successes[arm]
178 let f: i64 = t.failures[arm]
179 let alpha: i64 = s + 1
180 let beta: i64 = f + 1
181 let n_uniform: i64 = alpha + beta - 1
182 if n_uniform <= NX_THOMPSON_EXACT_MAX {
183 return nx_thompson_sample_order_stat(t, alpha, n_uniform)
184 }
185 // Normal approximation path.
186 let total: i64 = alpha + beta
187 let mu_q14: i64 = (alpha * NX_THOMPSON_Q14) / total
188 // var_q28 = mu_q14 * (Q14 - mu_q14) / (total + 1)
189 // Q14 · Q14 = Q28, then divide by integer keeps Q28.
190 let var_num: i64 = mu_q14 * (NX_THOMPSON_Q14 - mu_q14)
191 let var_q28: i64 = var_num / (total + 1)
192 // sigma_q14 = isqrt(var_q28) -- since sqrt(Q28) = Q14
193 let sigma_q14: i64 = nx_thompson_isqrt(var_q28)
194 let z_q14: i64 = nx_thompson_clt_normal_q14(t)
195 var sample_q14: i64 = mu_q14 + (sigma_q14 * z_q14) / NX_THOMPSON_Q14
196 if sample_q14 < 0 { sample_q14 = 0 }
197 if sample_q14 > NX_THOMPSON_Q14 { sample_q14 = NX_THOMPSON_Q14 }
198 return sample_q14
199}
200
201// === select =======================================================
202//
203// Sample each arm's posterior, return argmax.
204
205func nx_thompson_select(t: *Thompson) -> i64 {
206 if t.n_arms == 0 { return -1 }
207 var best: i64 = 0
208 var best_sample: i64 = nx_thompson_sample_q14(t, 0)
209 var i: i64 = 1
210 while i < t.n_arms {
211 let s: i64 = nx_thompson_sample_q14(t, i)
212 if s > best_sample {
213 best = i
214 best_sample = s
215 }
216 i = i + 1
217 }
218 return best
219}
220
221// === update =======================================================
222//
223// Reward in {0, 1}. 0 = failure, 1 = success. (Continuous-reward
224// variants will need a separate primitive -- this one is the
225// canonical Bernoulli case.)
226
227func nx_thompson_update(t: *Thompson, arm: i64, reward: i64) -> i64 {
228 if arm < 0 { return -1 }
229 if arm >= t.n_arms { return -1 }
230 if reward < 0 { return -1 }
231 if reward > 1 { return -1 }
232 if reward == 1 {
233 t.successes[arm] = t.successes[arm] + 1
234 }
235 if reward == 0 {
236 t.failures[arm] = t.failures[arm] + 1
237 }
238 t.total_pulls = t.total_pulls + 1
239 return 0
240}
241
242// === query / introspection ========================================
243
244func nx_thompson_successes_for(t: *Thompson, arm: i64) -> i64 {
245 return t.successes[arm]
246}
247
248func nx_thompson_failures_for(t: *Thompson, arm: i64) -> i64 {
249 return t.failures[arm]
250}
251
252// Posterior mean estimator: (s + 1) / (s + f + 2) in PPM.
253func nx_thompson_posterior_mean_ppm(t: *Thompson, arm: i64) -> i64 {
254 let s: i64 = t.successes[arm]
255 let f: i64 = t.failures[arm]
256 let alpha: i64 = s + 1
257 let total: i64 = s + f + 2
258 return (alpha * 1000000) / total
259}
260
261func nx_thompson_best_arm(t: *Thompson) -> i64 {
262 var best: i64 = -1
263 var best_mean: i64 = -1
264 var i: i64 = 0
265 while i < t.n_arms {
266 if t.successes[i] + t.failures[i] > 0 {
267 let m: i64 = nx_thompson_posterior_mean_ppm(t, i)
268 if m > best_mean {
269 best = i
270 best_mean = m
271 }
272 }
273 i = i + 1
274 }
275 return best
276}
277
278func nx_thompson_query(t: *Thompson, arm: i64) -> *ApproxI64 {
279 let m: i64 = nx_thompson_posterior_mean_ppm(t, arm)
280 let pulls: i64 = t.successes[arm] + t.failures[arm]
281 var stderr_ppb: i64 = 1000000000
282 let isq: i64 = nx_thompson_isqrt(pulls + 1)
283 if isq > 0 { stderr_ppb = 1000000000 / isq }
284 return nx_approx_new(m, NX_ENV_REL_STDDEV, stderr_ppb,
285 682700000,
286 NX_MATURITY_REFERENCE_IMPL,
287 NX_ADV_HONEST)
288}
289
290func nx_thompson_memory_bytes(t: *Thompson) -> i64 {
291 return 56 + t.n_arms * 16 + NX_THOMPSON_EXACT_MAX * 8
292}