nx_sketch_ucb1.nx source
↩ module page · 221 lines · 7024 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"
48import "nx_vecmath.nx"
49
50const NX_UCB_ARMS_MIN: i64 = 2
51const NX_UCB_ARMS_MAX: i64 = 10000
52const NX_UCB_LN2_PPM: i64 = 693147 // ln(2) * 1_000_000
53const NX_UCB_REWARD_MAX: i64 = 1000000 // 1.0 in PPM
54
55struct Ucb1 {
56 n_arms: i64,
57 counts: *i64, // count per arm
58 sum_rewards: *i64, // sum of rewards in PPM per arm
59 total_pulls: i64,
60}
61
62// === construction =================================================
63
64func nx_ucb_alloc(n_arms: i64) -> *Ucb1 {
65 if n_arms < NX_UCB_ARMS_MIN { return 0 as *Ucb1 }
66 if n_arms > NX_UCB_ARMS_MAX { return 0 as *Ucb1 }
67 let raw: *u8 = sys_mmap(40)
68 let b: *Ucb1 = raw as *Ucb1
69 let counts_raw: *u8 = sys_mmap(n_arms * 8)
70 b.counts = counts_raw as *i64
71 let sum_raw: *u8 = sys_mmap(n_arms * 8)
72 b.sum_rewards = sum_raw as *i64
73 var i: i64 = 0
74 while i < n_arms {
75 b.counts[i] = 0
76 b.sum_rewards[i] = 0
77 i = i + 1
78 }
79 b.n_arms = n_arms
80 b.total_pulls = 0
81 return b
82}
83
84// === isqrt (shared helper) =======================================
85
86func nx_ucb_isqrt(x: i64) -> i64 { return vm_isqrt(x) }
87
88// === bitlen-based ln =============================================
89//
90// ln(n) ≈ (bitlen(n) - 1) * ln(2) in PPM.
91// More accurate: ln(n) = log2(n) * ln(2) ≈ (bitlen(n) - 1) * 693147 PPM
92// (loses fractional bits but adequate for UCB exploration bonus).
93
94func nx_ucb_bitlen(x: i64) -> i64 {
95 if x <= 0 { return 0 }
96 var n: i64 = 0
97 var t: i64 = x
98 while t > 0 {
99 t = t >> 1
100 n = n + 1
101 }
102 return n
103}
104
105// Returns ln(n) * 1_000_000 (PPM).
106func nx_ucb_ln_ppm(n: i64) -> i64 {
107 if n <= 1 { return 0 }
108 let lg: i64 = nx_ucb_bitlen(n) - 1
109 return lg * NX_UCB_LN2_PPM
110}
111
112// === arm score ====================================================
113//
114// For arm i: score_ppm = mean_ppm + bonus_ppm
115// mean_ppm = sum_rewards[i] / counts[i] (in PPM)
116// bonus_ppm = sqrt(2 * ln(total_pulls) / counts[i]) * 1_000_000
117//
118// bonus calculation:
119// 2 * ln(N) in PPM = 2 * nx_ucb_ln_ppm(N)
120// numerator_for_sqrt = (2 * ln_ppm) * 1_000_000_000 / count
121// - this is value * 1e9 inside isqrt -> result is sqrt(value) * 31623
122// We want sqrt(2 ln(N) / count). In PPM: sqrt(value_ppm * 1e6) / 1e6.
123// So: bonus_ppm = isqrt((2 * ln_ppm * 1_000_000) / count)
124
125func nx_ucb_score(b: *Ucb1, arm: i64) -> i64 {
126 let cnt: i64 = b.counts[arm]
127 if cnt == 0 { return 0x7FFFFFFFFFFFFFFF } // force exploration
128 let sum: i64 = b.sum_rewards[arm]
129 let mean: i64 = sum / cnt
130 let ln_pm: i64 = nx_ucb_ln_ppm(b.total_pulls)
131 // bonus_squared_ppm = (2 * ln_ppm * 1_000_000) / cnt
132 let bsq: i64 = (2 * ln_pm * 1000000) / cnt
133 let bonus: i64 = nx_ucb_isqrt(bsq)
134 return mean + bonus
135}
136
137// === select =======================================================
138//
139// Returns the arm to pull next. Untested arms (count=0) pulled first
140// (their score is +inf).
141
142func nx_ucb_select(b: *Ucb1) -> i64 {
143 if b.n_arms == 0 { return -1 }
144 var best: i64 = 0
145 var best_score: i64 = nx_ucb_score(b, 0)
146 var i: i64 = 1
147 while i < b.n_arms {
148 let s: i64 = nx_ucb_score(b, i)
149 if s > best_score {
150 best = i
151 best_score = s
152 }
153 i = i + 1
154 }
155 return best
156}
157
158// === update =======================================================
159//
160// After pulling `arm` and observing `reward_ppm` (in [0, 1_000_000]).
161
162func nx_ucb_update(b: *Ucb1, arm: i64, reward_ppm: i64) -> i64 {
163 if arm < 0 { return -1 }
164 if arm >= b.n_arms { return -1 }
165 if reward_ppm < 0 { return -1 }
166 if reward_ppm > NX_UCB_REWARD_MAX { return -1 }
167 b.counts[arm] = b.counts[arm] + 1
168 b.sum_rewards[arm] = b.sum_rewards[arm] + reward_ppm
169 b.total_pulls = b.total_pulls + 1
170 return 0
171}
172
173// === query ========================================================
174
175func nx_ucb_mean_ppm(b: *Ucb1, arm: i64) -> i64 {
176 let cnt: i64 = b.counts[arm]
177 if cnt == 0 { return 0 }
178 return b.sum_rewards[arm] / cnt
179}
180
181func nx_ucb_count(b: *Ucb1, arm: i64) -> i64 {
182 return b.counts[arm]
183}
184
185// Returns the arm with the highest empirical mean (ignoring exploration).
186func nx_ucb_best_arm(b: *Ucb1) -> i64 {
187 var best: i64 = -1
188 var best_mean: i64 = -1
189 var i: i64 = 0
190 while i < b.n_arms {
191 if b.counts[i] > 0 {
192 let m: i64 = nx_ucb_mean_ppm(b, i)
193 if m > best_mean {
194 best = i
195 best_mean = m
196 }
197 }
198 i = i + 1
199 }
200 return best
201}
202
203// === typed envelope ===============================================
204
205func nx_ucb_query(b: *Ucb1, arm: i64) -> *ApproxI64 {
206 let m: i64 = nx_ucb_mean_ppm(b, arm)
207 let cnt: i64 = b.counts[arm]
208 var stderr_ppb: i64 = 1000000000
209 if cnt > 0 {
210 let isq: i64 = nx_ucb_isqrt(cnt)
211 if isq > 0 { stderr_ppb = 1000000000 / isq }
212 }
213 return nx_approx_new(m, NX_ENV_REL_STDDEV, stderr_ppb,
214 682700000,
215 NX_MATURITY_REFERENCE_IMPL,
216 NX_ADV_HONEST)
217}
218
219func nx_ucb_memory_bytes(b: *Ucb1) -> i64 {
220 return 40 + b.n_arms * 16
221}