nx_sketch_ucb1.nx source
↩ module page · 235 lines · 7318 B
1// sketch_ucb1.nx -- UCB1 multi-armed bandit (Auer-Cesa-Bianchi-Fischer 2002).
2//
3// Online-learning primitive. Given N "arms" (actions), select the
4// best one over time while balancing EXPLORATION (try arms we
5// haven't pulled) and EXPLOITATION (favor arms with high observed
6// reward).
7//
8// UCB1 selection rule:
9// score_i = mean_i + sqrt(2 * ln(N_total) / count_i)
10// pull arg max score_i
11//
12// Where:
13// mean_i = sum_reward_i / count_i (estimated value)
14// sqrt(...) = exploration bonus (shrinks as count_i grows)
15//
16// Arms with count_i = 0 have score = +inf (must-pull first).
17//
18// PROVEN PROPERTIES (Auer 2002):
19// regret bound: O(sqrt(K * N * ln(N)))
20// where K = #arms, N = total pulls. Near-optimal for adversarial
21// problems; optimal up to log factor for stochastic.
22//
23// USE CASES (extends NishiLang substrate beyond DataSketches scope):
24// - A/B testing with multi-arm allocation (Bayesian alternative)
25// - Recommendation with cold-start exploration
26// - Hyperparameter tuning under budget constraints
27// - Game AI move selection
28//
29// INTEGER-FIXED-POINT IMPLEMENTATION:
30// - rewards scaled to PPM (0..1_000_000 = [0, 1.0])
31// - mean_i = sum_reward_ppm / count_i
32// - ln(N) ≈ (bitlen(N) - 1) * 693147 / 1000 ppm (factor of ln(2))
33// - sqrt() via isqrt (Newton)
34// - score_i in PPM
35//
36// LOSSLESS-LANGUAGE DISCIPLINE: nx_ucb_query returns ApproxI64 with
37// NX_ENV_REL_STDDEV = 1/sqrt(count_i). Bandits are stochastic by
38// nature; the envelope honestly declares estimation uncertainty.
39
40// nx_safety_envelope:
41// intended_use: AUTO_APPLIED -- primitive-specific tuning queued
42// sil_target: SIL1
43// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail]
44// verdict: NOT_YET_EVALUATED
45
46import "nx_syscalls.nx"
47import "nx_sketch_types.nx"
48
49const NX_UCB_ARMS_MIN: i64 = 2
50const NX_UCB_ARMS_MAX: i64 = 10000
51const NX_UCB_LN2_PPM: i64 = 693147 // ln(2) * 1_000_000
52const NX_UCB_REWARD_MAX: i64 = 1000000 // 1.0 in PPM
53
54struct Ucb1 {
55 n_arms: i64,
56 counts: *i64, // count per arm
57 sum_rewards: *i64, // sum of rewards in PPM per arm
58 total_pulls: i64,
59}
60
61// === construction =================================================
62
63func nx_ucb_alloc(n_arms: i64) -> *Ucb1 {
64 if n_arms < NX_UCB_ARMS_MIN { return 0 as *Ucb1 }
65 if n_arms > NX_UCB_ARMS_MAX { return 0 as *Ucb1 }
66 let raw: *u8 = sys_mmap(40)
67 let b: *Ucb1 = raw as *Ucb1
68 let counts_raw: *u8 = sys_mmap(n_arms * 8)
69 b.counts = counts_raw as *i64
70 let sum_raw: *u8 = sys_mmap(n_arms * 8)
71 b.sum_rewards = sum_raw as *i64
72 var i: i64 = 0
73 while i < n_arms {
74 b.counts[i] = 0
75 b.sum_rewards[i] = 0
76 i = i + 1
77 }
78 b.n_arms = n_arms
79 b.total_pulls = 0
80 return b
81}
82
83// === isqrt (shared helper) =======================================
84
85func nx_ucb_isqrt(x: i64) -> i64 {
86 if x < 0 { return 0 }
87 if x == 0 { return 0 }
88 if x < 4 { return 1 }
89 var g: i64 = (x >> 1) + 1
90 var iter: i64 = 0
91 while iter < 64 {
92 let next_g: i64 = (g + x / g) / 2
93 if next_g >= g { iter = 64 }
94 if next_g < g {
95 g = next_g
96 iter = iter + 1
97 }
98 }
99 return g
100}
101
102// === bitlen-based ln =============================================
103//
104// ln(n) ≈ (bitlen(n) - 1) * ln(2) in PPM.
105// More accurate: ln(n) = log2(n) * ln(2) ≈ (bitlen(n) - 1) * 693147 PPM
106// (loses fractional bits but adequate for UCB exploration bonus).
107
108func nx_ucb_bitlen(x: i64) -> i64 {
109 if x <= 0 { return 0 }
110 var n: i64 = 0
111 var t: i64 = x
112 while t > 0 {
113 t = t >> 1
114 n = n + 1
115 }
116 return n
117}
118
119// Returns ln(n) * 1_000_000 (PPM).
120func nx_ucb_ln_ppm(n: i64) -> i64 {
121 if n <= 1 { return 0 }
122 let lg: i64 = nx_ucb_bitlen(n) - 1
123 return lg * NX_UCB_LN2_PPM
124}
125
126// === arm score ====================================================
127//
128// For arm i: score_ppm = mean_ppm + bonus_ppm
129// mean_ppm = sum_rewards[i] / counts[i] (in PPM)
130// bonus_ppm = sqrt(2 * ln(total_pulls) / counts[i]) * 1_000_000
131//
132// bonus calculation:
133// 2 * ln(N) in PPM = 2 * nx_ucb_ln_ppm(N)
134// numerator_for_sqrt = (2 * ln_ppm) * 1_000_000_000 / count
135// - this is value * 1e9 inside isqrt -> result is sqrt(value) * 31623
136// We want sqrt(2 ln(N) / count). In PPM: sqrt(value_ppm * 1e6) / 1e6.
137// So: bonus_ppm = isqrt((2 * ln_ppm * 1_000_000) / count)
138
139func nx_ucb_score(b: *Ucb1, arm: i64) -> i64 {
140 let cnt: i64 = b.counts[arm]
141 if cnt == 0 { return 0x7FFFFFFFFFFFFFFF } // force exploration
142 let sum: i64 = b.sum_rewards[arm]
143 let mean: i64 = sum / cnt
144 let ln_pm: i64 = nx_ucb_ln_ppm(b.total_pulls)
145 // bonus_squared_ppm = (2 * ln_ppm * 1_000_000) / cnt
146 let bsq: i64 = (2 * ln_pm * 1000000) / cnt
147 let bonus: i64 = nx_ucb_isqrt(bsq)
148 return mean + bonus
149}
150
151// === select =======================================================
152//
153// Returns the arm to pull next. Untested arms (count=0) pulled first
154// (their score is +inf).
155
156func nx_ucb_select(b: *Ucb1) -> i64 {
157 if b.n_arms == 0 { return -1 }
158 var best: i64 = 0
159 var best_score: i64 = nx_ucb_score(b, 0)
160 var i: i64 = 1
161 while i < b.n_arms {
162 let s: i64 = nx_ucb_score(b, i)
163 if s > best_score {
164 best = i
165 best_score = s
166 }
167 i = i + 1
168 }
169 return best
170}
171
172// === update =======================================================
173//
174// After pulling `arm` and observing `reward_ppm` (in [0, 1_000_000]).
175
176func nx_ucb_update(b: *Ucb1, arm: i64, reward_ppm: i64) -> i64 {
177 if arm < 0 { return -1 }
178 if arm >= b.n_arms { return -1 }
179 if reward_ppm < 0 { return -1 }
180 if reward_ppm > NX_UCB_REWARD_MAX { return -1 }
181 b.counts[arm] = b.counts[arm] + 1
182 b.sum_rewards[arm] = b.sum_rewards[arm] + reward_ppm
183 b.total_pulls = b.total_pulls + 1
184 return 0
185}
186
187// === query ========================================================
188
189func nx_ucb_mean_ppm(b: *Ucb1, arm: i64) -> i64 {
190 let cnt: i64 = b.counts[arm]
191 if cnt == 0 { return 0 }
192 return b.sum_rewards[arm] / cnt
193}
194
195func nx_ucb_count(b: *Ucb1, arm: i64) -> i64 {
196 return b.counts[arm]
197}
198
199// Returns the arm with the highest empirical mean (ignoring exploration).
200func nx_ucb_best_arm(b: *Ucb1) -> i64 {
201 var best: i64 = -1
202 var best_mean: i64 = -1
203 var i: i64 = 0
204 while i < b.n_arms {
205 if b.counts[i] > 0 {
206 let m: i64 = nx_ucb_mean_ppm(b, i)
207 if m > best_mean {
208 best = i
209 best_mean = m
210 }
211 }
212 i = i + 1
213 }
214 return best
215}
216
217// === typed envelope ===============================================
218
219func nx_ucb_query(b: *Ucb1, arm: i64) -> *ApproxI64 {
220 let m: i64 = nx_ucb_mean_ppm(b, arm)
221 let cnt: i64 = b.counts[arm]
222 var stderr_ppb: i64 = 1000000000
223 if cnt > 0 {
224 let isq: i64 = nx_ucb_isqrt(cnt)
225 if isq > 0 { stderr_ppb = 1000000000 / isq }
226 }
227 return nx_approx_new(m, NX_ENV_REL_STDDEV, stderr_ppb,
228 682700000,
229 NX_MATURITY_REFERENCE_IMPL,
230 NX_ADV_HONEST)
231}
232
233func nx_ucb_memory_bytes(b: *Ucb1) -> i64 {
234 return 40 + b.n_arms * 16
235}