diff --git a/LeanBandits/UCB.lean b/LeanBandits/UCB.lean index 9abcc2c4..d188a6c3 100644 --- a/LeanBandits/UCB.lean +++ b/LeanBandits/UCB.lean @@ -3,6 +3,7 @@ Copyright (c) 2025 Rémy Degenne. All rights reserved. Released under Apache 2.0 license as described in the file LICENSE. Authors: Rémy Degenne -/ +import LeanBandits.AlgorithmBuilding import LeanBandits.Regret /-! @@ -18,6 +19,36 @@ namespace Bandits variable {α : Type*} {mα : MeasurableSpace α} {ν : Kernel α ℝ} {k : ℕ → α} {t : ℕ} {a : α} +section Algorithm + +variable [Nonempty α] [DecidableEq α] [Finite α] [Encodable α] [MeasurableSingletonClass α] + +/-- The exploration bonus of the UCB algorithm, which corresponds to the width of +a confidence interval. -/ +noncomputable def ucbWidth' (c : ℝ) (n : ℕ) (h : Iic n → α × ℝ) (a : α) : ℝ := + √(c * log (n + 1) / (pullCount' n h a)) + +open Classical in +/-- Arm pulled by the UCB algorithm at time `n + 1`. -/ +noncomputable +def ucbNextArm (c : ℝ) (n : ℕ) (h : Iic n → α × ℝ) : α := + measurableArgmax (fun h a ↦ empMean' n h a + ucbWidth' c n h a) h + +@[fun_prop] +lemma measurable_ucbNextArm (c : ℝ) (n : ℕ) : Measurable (ucbNextArm c n (α := α)) := by + classical + refine measurable_measurableArgmax fun a ↦ ?_ + unfold ucbWidth' + fun_prop + +/-- The UCB algorithm. -/ +noncomputable +def ucbAlgorithm (c : ℝ) : Algorithm α ℝ where + policy n := Kernel.deterministic (ucbNextArm c n) (by fun_prop) + p0 := Measure.dirac (Classical.arbitrary α) + +end Algorithm + variable [Fintype α] [Nonempty α] {c : ℝ} {μ : α → ℝ} {N : α → ℕ} {a : α} /-- The exploration bonus of the UCB algorithm, which corresponds to the width of