sketch_epsilon_greedy.nx source
↩ module page · 210 lines · 6714 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"
28import "nx_vecmath.nx"
29
30const NX_EG_MIN_ARMS: i64 = 2
31const NX_EG_MAX_ARMS: i64 = 10000
32
33const NX_EG_MODE_FIXED: i64 = 0
34const NX_EG_MODE_DECAY_LINEAR: i64 = 1
35const NX_EG_MODE_DECAY_LOG: i64 = 2
36
37const NX_EG_LCG_A: i64 = 1103515245
38const NX_EG_LCG_C: i64 = 12345
39const NX_EG_LCG_MOD: i64 = 0x7FFFFFFF
40
41const NX_EG_REWARD_MAX: i64 = 1000000
42
43struct EpsilonGreedy {
44 n_arms: i64,
45 counts: *i64,
46 sum_rewards: *i64,
47 total_pulls: i64,
48 epsilon_mode: i64,
49 epsilon_ppm: i64, // initial/fixed epsilon in PPM
50 rng_state: i64,
51}
52
53// === construction =================================================
54
55func nx_eg_alloc(n_arms: i64, mode: i64, epsilon_ppm: i64, seed: i64) -> *EpsilonGreedy {
56 if n_arms < NX_EG_MIN_ARMS { return 0 as *EpsilonGreedy }
57 if n_arms > NX_EG_MAX_ARMS { return 0 as *EpsilonGreedy }
58 if mode < 0 { return 0 as *EpsilonGreedy }
59 if mode > NX_EG_MODE_DECAY_LOG { return 0 as *EpsilonGreedy }
60 if epsilon_ppm < 0 { return 0 as *EpsilonGreedy }
61 if epsilon_ppm > 1000000 { return 0 as *EpsilonGreedy }
62 let raw: *u8 = sys_mmap(56)
63 let eg: *EpsilonGreedy = raw as *EpsilonGreedy
64 eg.counts = sys_mmap(n_arms * 8) as *i64
65 eg.sum_rewards = sys_mmap(n_arms * 8) as *i64
66 var i: i64 = 0
67 while i < n_arms {
68 eg.counts[i] = 0
69 eg.sum_rewards[i] = 0
70 i = i + 1
71 }
72 eg.n_arms = n_arms
73 eg.total_pulls = 0
74 eg.epsilon_mode = mode
75 eg.epsilon_ppm = epsilon_ppm
76 eg.rng_state = seed | 1
77 return eg
78}
79
80func nx_eg_rng_next(eg: *EpsilonGreedy) -> i64 {
81 let next: i64 = ((eg.rng_state * NX_EG_LCG_A) + NX_EG_LCG_C) & NX_EG_LCG_MOD
82 eg.rng_state = next
83 return next
84}
85
86func nx_eg_rng_below(eg: *EpsilonGreedy, n: i64) -> i64 {
87 let r: i64 = nx_eg_rng_next(eg)
88 return r % n
89}
90
91// === bitlen for decay =============================================
92
93func nx_eg_bitlen(x: i64) -> i64 {
94 if x <= 0 { return 0 }
95 var n: i64 = 0
96 var t: i64 = x
97 while t > 0 {
98 t = t >> 1
99 n = n + 1
100 }
101 return n
102}
103
104// === current epsilon =============================================
105
106func nx_eg_current_epsilon_ppm(eg: *EpsilonGreedy) -> i64 {
107 if eg.epsilon_mode == NX_EG_MODE_FIXED {
108 return eg.epsilon_ppm
109 }
110 if eg.epsilon_mode == NX_EG_MODE_DECAY_LINEAR {
111 let t: i64 = eg.total_pulls + 1
112 return eg.epsilon_ppm / t
113 }
114 if eg.epsilon_mode == NX_EG_MODE_DECAY_LOG {
115 let t: i64 = eg.total_pulls + 2 // +2 so log(t) >= 1
116 let lg: i64 = nx_eg_bitlen(t)
117 if lg <= 1 { return eg.epsilon_ppm }
118 return eg.epsilon_ppm / lg
119 }
120 return eg.epsilon_ppm
121}
122
123// === select =======================================================
124//
125// Explore (uniform random arm) with prob epsilon; else exploit
126// (arm with highest empirical mean). Untested arms are picked
127// uniformly at random during the explore phase, eventually all
128// arms get pulled.
129
130func nx_eg_best_arm_by_mean(eg: *EpsilonGreedy) -> i64 {
131 var best: i64 = 0
132 var best_mean: i64 = -1
133 var i: i64 = 0
134 while i < eg.n_arms {
135 if eg.counts[i] > 0 {
136 let m: i64 = eg.sum_rewards[i] / eg.counts[i]
137 if m > best_mean {
138 best = i
139 best_mean = m
140 }
141 }
142 i = i + 1
143 }
144 if best_mean < 0 {
145 // All arms untested: pick random.
146 return nx_eg_rng_below(eg, eg.n_arms)
147 }
148 return best
149}
150
151func nx_eg_select(eg: *EpsilonGreedy) -> i64 {
152 let eps_ppm: i64 = nx_eg_current_epsilon_ppm(eg)
153 let roll: i64 = nx_eg_rng_below(eg, 1000000)
154 if roll < eps_ppm {
155 // Explore.
156 return nx_eg_rng_below(eg, eg.n_arms)
157 }
158 // Exploit.
159 return nx_eg_best_arm_by_mean(eg)
160}
161
162// === update ======================================================
163
164func nx_eg_update(eg: *EpsilonGreedy, arm: i64, reward_ppm: i64) -> i64 {
165 if arm < 0 { return -1 }
166 if arm >= eg.n_arms { return -1 }
167 if reward_ppm < 0 { return -1 }
168 if reward_ppm > NX_EG_REWARD_MAX { return -1 }
169 eg.counts[arm] = eg.counts[arm] + 1
170 eg.sum_rewards[arm] = eg.sum_rewards[arm] + reward_ppm
171 eg.total_pulls = eg.total_pulls + 1
172 return 0
173}
174
175// === queries =====================================================
176
177func nx_eg_mean_ppm(eg: *EpsilonGreedy, arm: i64) -> i64 {
178 if eg.counts[arm] == 0 { return 0 }
179 return eg.sum_rewards[arm] / eg.counts[arm]
180}
181
182func nx_eg_count(eg: *EpsilonGreedy, arm: i64) -> i64 {
183 return eg.counts[arm]
184}
185
186func nx_eg_best_arm(eg: *EpsilonGreedy) -> i64 {
187 return nx_eg_best_arm_by_mean(eg)
188}
189
190// === isqrt for envelope ==========================================
191
192func nx_eg_isqrt(x: i64) -> i64 { return vm_isqrt(x) }
193
194func nx_eg_query(eg: *EpsilonGreedy, arm: i64) -> *ApproxI64 {
195 let m: i64 = nx_eg_mean_ppm(eg, arm)
196 let cnt: i64 = eg.counts[arm]
197 var stderr_ppb: i64 = 1000000000
198 if cnt > 0 {
199 let isq: i64 = nx_eg_isqrt(cnt)
200 if isq > 0 { stderr_ppb = 1000000000 / isq }
201 }
202 return nx_approx_new(m, NX_ENV_REL_STDDEV, stderr_ppb,
203 682700000,
204 NX_MATURITY_REFERENCE_IMPL,
205 NX_ADV_HONEST)
206}
207
208func nx_eg_memory_bytes(eg: *EpsilonGreedy) -> i64 {
209 return 56 + eg.n_arms * 16
210}