feat : initial algorithm foundation for linUCB

Author: FawadHa1der

LeanMachineLearning.lean

@@ -4,6 +4,7 @@ public import LeanMachineLearning.MeasureTheory.Constructions.BorelSpace.Measura
 public import LeanMachineLearning.MeasureTheory.Constructions.Polish.StandardBorel
 public import LeanMachineLearning.MeasureTheory.Measurable
 public import LeanMachineLearning.Online.Bandit.Algorithms.ETC
+public import LeanMachineLearning.Online.Bandit.Algorithms.LinUCB
 public import LeanMachineLearning.Online.Bandit.Algorithms.UCB
 public import LeanMachineLearning.Online.Bandit.ArrayProbSpace
 public import LeanMachineLearning.Online.Bandit.Regret

LeanMachineLearning/Online/Bandit/Algorithms/LinUCB.lean

@@ -0,0 +1,230 @@
+/-
+Copyright (c) 2026. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: OpenAI, Fawad Haider
+-/
+module
+
+public import LeanMachineLearning.Online.Bandit.SumRewards
+public import LeanMachineLearning.SequentialLearning.Deterministic
+public import LeanMachineLearning.MeasureTheory.Constructions.BorelSpace.MeasurableArgMax
+public import Mathlib.LinearAlgebra.Matrix.NonsingularInverse
+
+/-!
+# LinUCB for finite-action linear bandits
+Chapter 19 of *Bandit Algorithms*:
+-/
+
+@[expose] public section
+
+open MeasureTheory ProbabilityTheory Filter Real Finset Learning
+
+open scoped ENNReal NNReal Matrix
+
+namespace Bandits
+
+variable {K d : ℕ}
+
+section Algorithm
+
+namespace LinUCB
+
+abbrev Feature (d : ℕ) := Fin d → ℝ
+
+noncomputable def designMatrix' (reg : ℝ) (x : Fin K → Feature d)
+    (n : ℕ) (h : Iic n → Fin K × ℝ) : Matrix (Fin d) (Fin d) ℝ :=
+  reg • 1 + ∑ s : Iic n, Matrix.vecMulVec (x (h s).1) (x (h s).1)
+
+noncomputable def responseVector' (x : Fin K → Feature d)
+    (n : ℕ) (h : Iic n → Fin K × ℝ) : Feature d :=
+  ∑ s : Iic n, (h s).2 • x (h s).1
+
+noncomputable def thetaHat' (reg : ℝ) (x : Fin K → Feature d)
+    (n : ℕ) (h : Iic n → Fin K × ℝ) : Feature d :=
+  Matrix.mulVec (designMatrix' reg x n h)⁻¹ (responseVector' x n h)
+
+noncomputable def estimatedReward' (reg : ℝ) (x : Fin K → Feature d)
+    (n : ℕ) (h : Iic n → Fin K × ℝ) (a : Fin K) : ℝ :=
+  dotProduct (thetaHat' reg x n h) (x a)
+
+noncomputable def width' (reg : ℝ) (x : Fin K → Feature d)
+    (n : ℕ) (h : Iic n → Fin K × ℝ) (a : Fin K) : ℝ :=
+  √(dotProduct (x a) (Matrix.mulVec (designMatrix' reg x n h)⁻¹ (x a)))
+
+/-- LinUCB optimistic index of an arm.
+
+The parameter `β` is a confidence-radius schedule. Since `h : Iic n → Fin K × ℝ`
+contains the observations through time `n`, this index is used to choose the arm
+at time `n + 1`, and we evaluate the schedule at `n + 2`
+-/
+noncomputable def index' (reg : ℝ) (β : ℕ → ℝ) (x : Fin K → Feature d)
+    (n : ℕ) (h : Iic n → Fin K × ℝ) (a : Fin K) : ℝ :=
+  estimatedReward' reg x n h a + √(β (n + 2)) * width' reg x n h a
+
+open Classical in
+/-- Arm pulled by finite-action LinUCB at time `n + 1`. -/
+noncomputable def nextArm (hK : 0 < K) (reg : ℝ) (β : ℕ → ℝ)
+    (x : Fin K → Feature d)
+    (_h_index : ∀ n a, Measurable (fun h ↦ index' reg β x n h a))
+    (n : ℕ) (h : Iic n → Fin K × ℝ) : Fin K :=
+  have : Nonempty (Fin K) := Fin.pos_iff_nonempty.mp hK
+  measurableArgmax (fun h a ↦ index' reg β x n h a) h
+
+@[fun_prop]
+lemma measurable_nextArm (hK : 0 < K) (reg : ℝ) (β : ℕ → ℝ)
+    (x : Fin K → Feature d)
+    (h_index : ∀ n a, Measurable (fun h ↦ index' reg β x n h a))
+    (n : ℕ) :
+    Measurable (nextArm hK reg β x h_index n) := by
+  have : Nonempty (Fin K) := Fin.pos_iff_nonempty.mp hK
+  exact measurable_measurableArgmax fun a ↦ h_index n a
+
+end LinUCB
+
+/-- The finite-action LinUCB algorithm. -/
+noncomputable def linUCBAlgorithm (hK : 0 < K) (reg : ℝ) (β : ℕ → ℝ)
+    (x : Fin K → LinUCB.Feature d)
+    (h_index : ∀ n a, Measurable (fun h ↦ LinUCB.index' reg β x n h a)) :
+    Algorithm (Fin K) ℝ :=
+  detAlgorithm (LinUCB.nextArm hK reg β x h_index) (by fun_prop) ⟨0, hK⟩
+
+end Algorithm
+
+namespace LinUCB
+
+variable {hK : 0 < K} {reg : ℝ} {β : ℕ → ℝ} {x : Fin K → Feature d}
+  {h_index : ∀ n a, Measurable (fun h ↦ index' reg β x n h a)}
+  {ν : Kernel (Fin K) ℝ} [IsMarkovKernel ν]
+  {Ω : Type*} {mΩ : MeasurableSpace Ω}
+  {P : Measure Ω} [IsProbabilityMeasure P]
+  {A : ℕ → Ω → Fin K} {R : ℕ → Ω → ℝ}
+  {n : ℕ} {ω : Ω}
+
+section AlgorithmBehavior
+
+/-- The process-level design matrix built from actions up to time `n` excluded. -/
+noncomputable def designMatrix (A : ℕ → Ω → Fin K) (reg : ℝ)
+    (x : Fin K → Feature d) (n : ℕ) (ω : Ω) : Matrix (Fin d) (Fin d) ℝ :=
+  reg • 1 + ∑ s ∈ range n, Matrix.vecMulVec (x (A s ω)) (x (A s ω))
+
+/-- The process-level reward-feature vector built from history up to time `n` excluded. -/
+noncomputable def responseVector (A : ℕ → Ω → Fin K) (R : ℕ → Ω → ℝ)
+    (x : Fin K → Feature d) (n : ℕ) (ω : Ω) : Feature d :=
+  ∑ s ∈ range n, R s ω • x (A s ω)
+
+/-- The process-level regularized least-squares estimate. -/
+noncomputable def thetaHat (A : ℕ → Ω → Fin K) (R : ℕ → Ω → ℝ)
+    (reg : ℝ) (x : Fin K → Feature d) (n : ℕ) (ω : Ω) : Feature d :=
+  Matrix.mulVec (designMatrix A reg x n ω)⁻¹ (responseVector A R x n ω)
+
+/-- The process-level estimated linear reward. -/
+noncomputable def estimatedReward (A : ℕ → Ω → Fin K) (R : ℕ → Ω → ℝ)
+    (reg : ℝ) (x : Fin K → Feature d) (a : Fin K) (n : ℕ) (ω : Ω) : ℝ :=
+  dotProduct (thetaHat A R reg x n ω) (x a)
+
+/-- The process-level elliptical confidence width. -/
+noncomputable def width (A : ℕ → Ω → Fin K) (reg : ℝ)
+    (x : Fin K → Feature d) (a : Fin K) (n : ℕ) (ω : Ω) : ℝ :=
+  √(dotProduct (x a) (Matrix.mulVec (designMatrix A reg x n ω)⁻¹ (x a)))
+
+/-- The process-level LinUCB optimistic index. -/
+noncomputable def index (A : ℕ → Ω → Fin K) (R : ℕ → Ω → ℝ)
+    (reg : ℝ) (β : ℕ → ℝ) (x : Fin K → Feature d) (a : Fin K)
+    (n : ℕ) (ω : Ω) : ℝ :=
+  estimatedReward A R reg x a n ω + √(β (n + 1)) * width A reg x a n ω
+
+lemma designMatrix_eq_designMatrix' (reg : ℝ) (x : Fin K → Feature d) (n : ℕ)
+    (ω : Ω) (hn : n ≠ 0) :
+    designMatrix A reg x n ω =
+      designMatrix' reg x (n - 1) (IsAlgEnvSeq.hist A R (n - 1) ω) := by
+  cases n with
+  | zero => exact absurd rfl hn
+  | succ n =>
+    simp only [designMatrix, designMatrix', IsAlgEnvSeq.hist]
+    rw [Nat.range_succ_eq_Iic]
+    exact congrArg (fun S ↦ reg • 1 + S) <|
+      (Finset.sum_coe_sort (Iic n)
+        (fun s ↦ Matrix.vecMulVec (x (A s ω)) (x (A s ω)))).symm
+
+lemma responseVector_eq_responseVector' (x : Fin K → Feature d)
+    (n : ℕ) (ω : Ω) (hn : n ≠ 0) :
+    responseVector A R x n ω = responseVector' x (n - 1) (IsAlgEnvSeq.hist A R (n - 1) ω) := by
+  cases n with
+  | zero => exact absurd rfl hn
+  | succ n =>
+    simp only [responseVector, responseVector', IsAlgEnvSeq.hist]
+    rw [Nat.range_succ_eq_Iic]
+    exact (Finset.sum_coe_sort (Iic n) (fun s ↦ R s ω • x (A s ω))).symm
+
+lemma thetaHat_eq_thetaHat' (reg : ℝ) (x : Fin K → Feature d)
+    (n : ℕ) (ω : Ω) (hn : n ≠ 0) :
+    thetaHat A R reg x n ω = thetaHat' reg x (n - 1) (IsAlgEnvSeq.hist A R (n - 1) ω) := by
+  simp [thetaHat, thetaHat', designMatrix_eq_designMatrix' (A := A) (R := R) reg x n ω hn,
+    responseVector_eq_responseVector' (A := A) (R := R) x n ω hn]
+
+lemma estimatedReward_eq_estimatedReward' (reg : ℝ) (x : Fin K → Feature d)
+    (a : Fin K) (n : ℕ) (ω : Ω) (hn : n ≠ 0) :
+    estimatedReward A R reg x a n ω =
+      estimatedReward' reg x (n - 1) (IsAlgEnvSeq.hist A R (n - 1) ω) a := by
+  simp [estimatedReward, estimatedReward', thetaHat_eq_thetaHat' (A := A) (R := R) reg x n ω hn]
+
+lemma width_eq_width' (reg : ℝ) (x : Fin K → Feature d)
+    (a : Fin K) (n : ℕ) (ω : Ω) (hn : n ≠ 0) :
+    width A reg x a n ω = width' reg x (n - 1) (IsAlgEnvSeq.hist A R (n - 1) ω) a := by
+  simp [width, width', designMatrix_eq_designMatrix' (A := A) (R := R) reg x n ω hn]
+
+lemma index_eq_index' (reg : ℝ) (β : ℕ → ℝ) (x : Fin K → Feature d)
+    (a : Fin K) (n : ℕ) (ω : Ω) (hn : n ≠ 0) :
+    index A R reg β x a n ω =
+      index' reg β x (n - 1) (IsAlgEnvSeq.hist A R (n - 1) ω) a := by
+  have htime : n + 1 = n - 1 + 2 := by grind
+  simp [index, index', estimatedReward_eq_estimatedReward' (A := A) (R := R) reg x a n ω hn,
+    width_eq_width' (A := A) (R := R) reg x a n ω hn, htime]
+
+/-- The action at time `n + 1` is the finite-action LinUCB argmax for the observed history. -/
+lemma arm_ae_eq_linUCBNextArm [Nonempty (Fin K)]
+    (h : IsAlgEnvSeq A R (linUCBAlgorithm hK reg β x h_index) (stationaryEnv ν) P)
+    (n : ℕ) :
+    A (n + 1) =ᵐ[P]
+      fun ω ↦ nextArm hK reg β x h_index n (IsAlgEnvSeq.hist A R n ω) := by
+  have : Nonempty (Fin K) := Fin.pos_iff_nonempty.mp hK
+  exact h.action_detAlgorithm_ae_eq n
+
+/-- Almost surely, every positive-time action is the finite-action LinUCB argmax. -/
+lemma arm_ae_all_eq [Nonempty (Fin K)]
+    (h : IsAlgEnvSeq A R (linUCBAlgorithm hK reg β x h_index) (stationaryEnv ν) P) :
+    ∀ᵐ ω ∂P,
+      ∀ n, A (n + 1) ω =
+        nextArm hK reg β x h_index n (IsAlgEnvSeq.hist A R n ω) := by
+  simp_rw [ae_all_iff]
+  exact fun n ↦ arm_ae_eq_linUCBNextArm h n
+
+/-- Finite-action LinUCB chooses an arm maximizing the LinUCB index. -/
+lemma index_le_index_arm [Nonempty (Fin K)]
+    (h : IsAlgEnvSeq A R (linUCBAlgorithm hK reg β x h_index) (stationaryEnv ν) P)
+    (a : Fin K) (hn : n ≠ 0) :
+    ∀ᵐ ω ∂P, index A R reg β x a n ω ≤ index A R reg β x (A n ω) n ω := by
+  filter_upwards [arm_ae_eq_linUCBNextArm h (n - 1)] with ω h_arm
+  have hn_succ : n - 1 + 1 = n := by grind
+  simp only [hn_succ] at h_arm
+  rw [index_eq_index' (A := A) (R := R) reg β x a n ω hn,
+    index_eq_index' (A := A) (R := R) reg β x (A n ω) n ω hn]
+  rw [h_arm]
+  have : Nonempty (Fin K) := Fin.pos_iff_nonempty.mp hK
+  exact isMaxOn_measurableArgmax (fun h a ↦ index' reg β x (n - 1) h a)
+    (IsAlgEnvSeq.hist A R (n - 1) ω) a
+
+/-- Almost surely, the selected arm maximizes the LinUCB index at every positive time. -/
+lemma forall_index_le_index_arm [Nonempty (Fin K)]
+    (h : IsAlgEnvSeq A R (linUCBAlgorithm hK reg β x h_index) (stationaryEnv ν) P)
+    (a : Fin K) :
+    ∀ᵐ ω ∂P, ∀ n, n ≠ 0 →
+      index A R reg β x a n ω ≤ index A R reg β x (A n ω) n ω := by
+  simp_rw [ae_all_iff]
+  exact fun n hn ↦ index_le_index_arm h a hn
+
+end AlgorithmBehavior
+
+end LinUCB
+
+end Bandits