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