sketch_epsilon_greedy.nx source
↩ module page · 224 lines · 7023 B
1// sketch_epsilon_greedy.nx -- Epsilon-Greedy bandit (Sutton-Barto, classic).
2//
3// Simpler alternative to UCB1. At each pull:
4// with probability epsilon: pull a UNIFORM-RANDOM arm
5// with probability 1-epsilon: pull arm with highest empirical mean
6//
7// Three modes via the epsilon_mode field:
8// NX_EG_MODE_FIXED -- epsilon constant (param: epsilon_ppm)
9// NX_EG_MODE_DECAY_LINEAR -- epsilon(t) = epsilon_0 / (t+1)
10// NX_EG_MODE_DECAY_LOG -- epsilon(t) = epsilon_0 / log(t+e)
11//
12// COMPLEMENTS UCB1:
13// - UCB1: deterministic confidence-bound; near-optimal regret;
14// requires sqrt + log per arm per pull
15// - epsilon-Greedy: stochastic; simpler; tunable explore rate;
16// preferred when log/sqrt are expensive (embedded) or when
17// bounded exploration budget needed
18// Together they cover the "classic bandit" textbook coverage.
19//
20// LCG-based deterministic RNG so two runs with same seed produce
21// byte-identical action sequences.
22//
23// LOSSLESS-LANGUAGE DISCIPLINE: nx_eg_query reports the running
24// mean of an arm with stderr ~ 1/sqrt(count_i).
25
26import "syscalls.nx"
27import "sketch_types.nx"
28
29const NX_EG_MIN_ARMS: i64 = 2
30const NX_EG_MAX_ARMS: i64 = 10000
31
32const NX_EG_MODE_FIXED: i64 = 0
33const NX_EG_MODE_DECAY_LINEAR: i64 = 1
34const NX_EG_MODE_DECAY_LOG: i64 = 2
35
36const NX_EG_LCG_A: i64 = 1103515245
37const NX_EG_LCG_C: i64 = 12345
38const NX_EG_LCG_MOD: i64 = 0x7FFFFFFF
39
40const NX_EG_REWARD_MAX: i64 = 1000000
41
42struct EpsilonGreedy {
43 n_arms: i64,
44 counts: *i64,
45 sum_rewards: *i64,
46 total_pulls: i64,
47 epsilon_mode: i64,
48 epsilon_ppm: i64, // initial/fixed epsilon in PPM
49 rng_state: i64,
50}
51
52// === construction =================================================
53
54func nx_eg_alloc(n_arms: i64, mode: i64, epsilon_ppm: i64, seed: i64) -> *EpsilonGreedy {
55 if n_arms < NX_EG_MIN_ARMS { return 0 as *EpsilonGreedy }
56 if n_arms > NX_EG_MAX_ARMS { return 0 as *EpsilonGreedy }
57 if mode < 0 { return 0 as *EpsilonGreedy }
58 if mode > NX_EG_MODE_DECAY_LOG { return 0 as *EpsilonGreedy }
59 if epsilon_ppm < 0 { return 0 as *EpsilonGreedy }
60 if epsilon_ppm > 1000000 { return 0 as *EpsilonGreedy }
61 let raw: *u8 = sys_mmap(56)
62 let eg: *EpsilonGreedy = raw as *EpsilonGreedy
63 eg.counts = sys_mmap(n_arms * 8) as *i64
64 eg.sum_rewards = sys_mmap(n_arms * 8) as *i64
65 var i: i64 = 0
66 while i < n_arms {
67 eg.counts[i] = 0
68 eg.sum_rewards[i] = 0
69 i = i + 1
70 }
71 eg.n_arms = n_arms
72 eg.total_pulls = 0
73 eg.epsilon_mode = mode
74 eg.epsilon_ppm = epsilon_ppm
75 eg.rng_state = seed | 1
76 return eg
77}
78
79func nx_eg_rng_next(eg: *EpsilonGreedy) -> i64 {
80 let next: i64 = ((eg.rng_state * NX_EG_LCG_A) + NX_EG_LCG_C) & NX_EG_LCG_MOD
81 eg.rng_state = next
82 return next
83}
84
85func nx_eg_rng_below(eg: *EpsilonGreedy, n: i64) -> i64 {
86 let r: i64 = nx_eg_rng_next(eg)
87 return r % n
88}
89
90// === bitlen for decay =============================================
91
92func nx_eg_bitlen(x: i64) -> i64 {
93 if x <= 0 { return 0 }
94 var n: i64 = 0
95 var t: i64 = x
96 while t > 0 {
97 t = t >> 1
98 n = n + 1
99 }
100 return n
101}
102
103// === current epsilon =============================================
104
105func nx_eg_current_epsilon_ppm(eg: *EpsilonGreedy) -> i64 {
106 if eg.epsilon_mode == NX_EG_MODE_FIXED {
107 return eg.epsilon_ppm
108 }
109 if eg.epsilon_mode == NX_EG_MODE_DECAY_LINEAR {
110 let t: i64 = eg.total_pulls + 1
111 return eg.epsilon_ppm / t
112 }
113 if eg.epsilon_mode == NX_EG_MODE_DECAY_LOG {
114 let t: i64 = eg.total_pulls + 2 // +2 so log(t) >= 1
115 let lg: i64 = nx_eg_bitlen(t)
116 if lg <= 1 { return eg.epsilon_ppm }
117 return eg.epsilon_ppm / lg
118 }
119 return eg.epsilon_ppm
120}
121
122// === select =======================================================
123//
124// Explore (uniform random arm) with prob epsilon; else exploit
125// (arm with highest empirical mean). Untested arms are picked
126// uniformly at random during the explore phase, eventually all
127// arms get pulled.
128
129func nx_eg_best_arm_by_mean(eg: *EpsilonGreedy) -> i64 {
130 var best: i64 = 0
131 var best_mean: i64 = -1
132 var i: i64 = 0
133 while i < eg.n_arms {
134 if eg.counts[i] > 0 {
135 let m: i64 = eg.sum_rewards[i] / eg.counts[i]
136 if m > best_mean {
137 best = i
138 best_mean = m
139 }
140 }
141 i = i + 1
142 }
143 if best_mean < 0 {
144 // All arms untested: pick random.
145 return nx_eg_rng_below(eg, eg.n_arms)
146 }
147 return best
148}
149
150func nx_eg_select(eg: *EpsilonGreedy) -> i64 {
151 let eps_ppm: i64 = nx_eg_current_epsilon_ppm(eg)
152 let roll: i64 = nx_eg_rng_below(eg, 1000000)
153 if roll < eps_ppm {
154 // Explore.
155 return nx_eg_rng_below(eg, eg.n_arms)
156 }
157 // Exploit.
158 return nx_eg_best_arm_by_mean(eg)
159}
160
161// === update ======================================================
162
163func nx_eg_update(eg: *EpsilonGreedy, arm: i64, reward_ppm: i64) -> i64 {
164 if arm < 0 { return -1 }
165 if arm >= eg.n_arms { return -1 }
166 if reward_ppm < 0 { return -1 }
167 if reward_ppm > NX_EG_REWARD_MAX { return -1 }
168 eg.counts[arm] = eg.counts[arm] + 1
169 eg.sum_rewards[arm] = eg.sum_rewards[arm] + reward_ppm
170 eg.total_pulls = eg.total_pulls + 1
171 return 0
172}
173
174// === queries =====================================================
175
176func nx_eg_mean_ppm(eg: *EpsilonGreedy, arm: i64) -> i64 {
177 if eg.counts[arm] == 0 { return 0 }
178 return eg.sum_rewards[arm] / eg.counts[arm]
179}
180
181func nx_eg_count(eg: *EpsilonGreedy, arm: i64) -> i64 {
182 return eg.counts[arm]
183}
184
185func nx_eg_best_arm(eg: *EpsilonGreedy) -> i64 {
186 return nx_eg_best_arm_by_mean(eg)
187}
188
189// === isqrt for envelope ==========================================
190
191func nx_eg_isqrt(x: i64) -> i64 {
192 if x < 0 { return 0 }
193 if x == 0 { return 0 }
194 if x < 4 { return 1 }
195 var g: i64 = (x >> 1) + 1
196 var iter: i64 = 0
197 while iter < 64 {
198 let next_g: i64 = (g + x / g) / 2
199 if next_g >= g { iter = 64 }
200 if next_g < g {
201 g = next_g
202 iter = iter + 1
203 }
204 }
205 return g
206}
207
208func nx_eg_query(eg: *EpsilonGreedy, arm: i64) -> *ApproxI64 {
209 let m: i64 = nx_eg_mean_ppm(eg, arm)
210 let cnt: i64 = eg.counts[arm]
211 var stderr_ppb: i64 = 1000000000
212 if cnt > 0 {
213 let isq: i64 = nx_eg_isqrt(cnt)
214 if isq > 0 { stderr_ppb = 1000000000 / isq }
215 }
216 return nx_approx_new(m, NX_ENV_REL_STDDEV, stderr_ppb,
217 682700000,
218 NX_MATURITY_REFERENCE_IMPL,
219 NX_ADV_HONEST)
220}
221
222func nx_eg_memory_bytes(eg: *EpsilonGreedy) -> i64 {
223 return 56 + eg.n_arms * 16
224}