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