Bluprint: refactor to talk about general stochastic algorithms, beyond bandits

LeanBandits.lean

@@ -1,9 +1,11 @@
+import LeanBandits.Algorithm
 import LeanBandits.AlgorithmBuilding
 import LeanBandits.Bandit
 import LeanBandits.ETC
 import LeanBandits.ForMathlib.CondDistrib
 import LeanBandits.ForMathlib.KernelCompositionLemmas
 import LeanBandits.ForMathlib.KernelCompositionParallelComp
+import LeanBandits.ForMathlib.KernelSub
 import LeanBandits.ForMathlib.Traj
 import LeanBandits.Regret
 import LeanBandits.RewardByCountMeasure

LeanBandits/Algorithm.lean

@@ -0,0 +1,232 @@
+/-
+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, Paulo Rauber
+-/
+import Mathlib
+import LeanBandits.ForMathlib.CondDistrib
+import LeanBandits.ForMathlib.KernelCompositionLemmas
+import LeanBandits.ForMathlib.Traj
+
+/-!
+# Bandit
+-/
+
+open MeasureTheory ProbabilityTheory Filter Real Finset
+
+open scoped ENNReal NNReal
+
+namespace Learning
+
+variable {α R : Type*} {mα : MeasurableSpace α} {mR : MeasurableSpace R}
+
+/-- A stochastic, sequential algorithm. -/
+structure Algorithm (α R : Type*) [MeasurableSpace α] [MeasurableSpace R] where
+  /-- Policy or sampling rule: distribution of the next action. -/
+  policy : (n : ℕ) → Kernel (Iic n → α × R) α
+  [h_policy : ∀ n, IsMarkovKernel (policy n)]
+  /-- Distribution of the first action. -/
+  p0 : Measure α
+  [hp0 : IsProbabilityMeasure p0]
+
+instance (alg : Algorithm α R) (n : ℕ) : IsMarkovKernel (alg.policy n) := alg.h_policy n
+instance (alg : Algorithm α R) : IsProbabilityMeasure alg.p0 := alg.hp0
+
+/-- A stochastic environment. -/
+structure Environment (α R : Type*) [MeasurableSpace α] [MeasurableSpace R] where
+  /-- Distribution of the next observation as function of the past history. -/
+  feedback : (n : ℕ) → Kernel ((Iic n → α × R) × α) R
+  [h_feedback : ∀ n, IsMarkovKernel (feedback n)]
+  /-- Distribution of the first observation given the first action. -/
+  ν0 : Kernel α R
+  [hp0 : IsMarkovKernel ν0]
+
+instance (env : Environment α R) (n : ℕ) : IsMarkovKernel (env.feedback n) := env.h_feedback n
+instance (env : Environment α R) : IsMarkovKernel env.ν0 := env.hp0
+
+/-- A deterministic algorithm. -/
+noncomputable
+def detAlgorithm (nextaction : (n : ℕ) → (Iic n → α × R) → α)
+    (h_next : ∀ n, Measurable (nextaction n)) (action0 : α) :
+    Algorithm α R where
+  policy n := Kernel.deterministic (nextaction n) (h_next n)
+  p0 := Measure.dirac action0
+
+/-- A stationary environment, in which the distribution of the next reward depends only on the last
+action. -/
+@[simps]
+def stationaryEnv (ν : Kernel α R) [IsMarkovKernel ν] : Environment α R where
+  feedback _ := ν.prodMkLeft _
+  ν0 := ν
+
+/-- Kernel describing the distribution of the next action-reward pair given the history
+up to `n`. -/
+noncomputable
+def stepKernel (alg : Algorithm α R) (env : Environment α R) (n : ℕ) :
+    Kernel (Iic n → α × R) (α × R) :=
+  alg.policy n ⊗ₖ env.feedback n
+deriving IsMarkovKernel
+
+@[simp]
+lemma fst_stepKernel (alg : Algorithm α R) (env : Environment α R) (n : ℕ) :
+    (stepKernel alg env n).fst = alg.policy n := by
+  rw [stepKernel, Kernel.fst_compProd]
+
+/-- Kernel sending a partial trajectory of the bandit interaction `Iic n → α × ℝ` to a measure
+on `ℕ → α × ℝ`, supported on full trajectories that start with the partial one. -/
+noncomputable def traj (alg : Algorithm α R) (env : Environment α R) (n : ℕ) :
+    Kernel (Iic n → α × R) (ℕ → α × R) :=
+  ProbabilityTheory.Kernel.traj (X := fun _ ↦ α × R) (stepKernel alg env) n
+deriving IsMarkovKernel
+
+/-- Measure on the sequence of actions and observations generated by the algorithm/environment. -/
+noncomputable
+def trajMeasure (alg : Algorithm α R) (env : Environment α R) :
+    Measure (ℕ → α × R) :=
+  Kernel.trajMeasure (alg.p0 ⊗ₘ env.ν0) (stepKernel alg env)
+deriving IsProbabilityMeasure
+
+/-- Action and reward at step `n`. -/
+def step (n : ℕ) (h : ℕ → α × R) : α × R := h n
+
+/-- `action n` is the action pulled at time `n`. This is a random variable on the measurable space
+`ℕ → α × ℝ`. -/
+def action (n : ℕ) (h : ℕ → α × R) : α := (h n).1
+
+/-- `reward n` is the reward at time `n`. This is a random variable on the measurable space
+`ℕ → α × R`. -/
+def reward (n : ℕ) (h : ℕ → α × R) : R := (h n).2
+
+/-- `hist n` is the history up to time `n`. This is a random variable on the measurable space
+`ℕ → α × R`. -/
+def hist (n : ℕ) (h : ℕ → α × R) : Iic n → α × R := fun i ↦ h i
+
+lemma fst_comp_step (n : ℕ) : Prod.fst ∘ step (α := α) (R := R) n = action n := rfl
+
+@[fun_prop]
+lemma measurable_step (n : ℕ) : Measurable (step n (α := α) (R := R)) := by
+  unfold step; fun_prop
+
+@[fun_prop]
+lemma measurable_step_prod : Measurable (fun p : ℕ × (ℕ → α × R) ↦ step p.1 p.2) := by
+  refine measurable_from_prod_countable_right fun n ↦ ?_
+  simp only
+  fun_prop
+
+@[fun_prop]
+lemma measurable_action (n : ℕ) : Measurable (action n (α := α) (R := R)) := by
+  unfold action; fun_prop
+
+@[fun_prop]
+lemma measurable_action_prod : Measurable (fun p : ℕ × (ℕ → α × R) ↦ action p.1 p.2) := by
+  refine measurable_from_prod_countable_right fun n ↦ ?_
+  simp only
+  fun_prop
+
+@[fun_prop]
+lemma measurable_reward (n : ℕ) : Measurable (reward n (α := α) (R := R)) := by
+  unfold reward; fun_prop
+
+@[fun_prop]
+lemma measurable_reward_prod : Measurable (fun p : ℕ × (ℕ → α × R) ↦ reward p.1 p.2) := by
+  refine measurable_from_prod_countable_right fun n ↦ ?_
+  simp only
+  fun_prop
+
+@[fun_prop]
+lemma measurable_hist (n : ℕ) : Measurable (hist n (α := α) (R := R)) := by unfold hist; fun_prop
+
+lemma hist_eq_frestrictLe :
+    hist = Preorder.frestrictLe («π» := fun _ ↦ α × R) := by
+  ext n h i : 3
+  simp [hist, Preorder.frestrictLe]
+
+/-- Filtration of the algorithm interaction. -/
+protected def filtration (α R : Type*) [MeasurableSpace α] [MeasurableSpace R] :
+    Filtration ℕ (inferInstance : MeasurableSpace (ℕ → α × R)) :=
+  MeasureTheory.Filtration.piLE (X := fun _ ↦ α × R)
+
+lemma condDistrib_step [StandardBorelSpace α] [Nonempty α] [StandardBorelSpace R] [Nonempty R]
+    (alg : Algorithm α R) (env : Environment α R) (n : ℕ) :
+    condDistrib (step (n + 1)) (hist n) (trajMeasure alg env)
+      =ᵐ[(trajMeasure alg env).map (hist n)] stepKernel alg env n :=
+  Kernel.condDistrib_trajMeasure_ae_eq_kernel
+
+lemma condDistrib_action [StandardBorelSpace α] [Nonempty α] [StandardBorelSpace R] [Nonempty R]
+    (alg : Algorithm α R) (env : Environment α R) (n : ℕ) :
+    condDistrib (action (n + 1)) (hist n) (trajMeasure alg env)
+      =ᵐ[(trajMeasure alg env).map (hist n)] alg.policy n := by
+  rw [← fst_comp_step]
+  refine (condDistrib_comp' (by fun_prop) (by fun_prop) (by fun_prop)).trans ?_
+  filter_upwards [condDistrib_step alg env n] with h h_eq
+  rw [Kernel.map_apply _ (by fun_prop), h_eq, ← Kernel.map_apply _ (by fun_prop), ← Kernel.fst_eq,
+    fst_stepKernel]
+
+lemma condDistrib_reward [StandardBorelSpace α] [Nonempty α] [StandardBorelSpace R] [Nonempty R]
+    (alg : Algorithm α R) (env : Environment α R) (n : ℕ) :
+    condDistrib (reward (n + 1)) (fun ω ↦ (hist n ω, action (n + 1) ω)) (trajMeasure alg env)
+      =ᵐ[(trajMeasure alg env).map (fun ω ↦ (hist n ω, action (n + 1) ω))] env.feedback n := by
+  have h_step := condDistrib_step alg env n
+  have h_action := condDistrib_action alg env n
+  rw [condDistrib_ae_eq_iff_measure_eq_compProd₀ (by fun_prop) (by fun_prop)] at h_step h_action ⊢
+  rw [h_action, Measure.compProd_assoc, ← stepKernel, ← h_step,
+    Measure.map_map (by fun_prop) (by fun_prop)]
+  rfl
+
+lemma hasLaw_step_zero (alg : Algorithm α R) (env : Environment α R) :
+    HasLaw (step 0) (alg.p0 ⊗ₘ env.ν0) (trajMeasure alg env) where
+  aemeasurable := Measurable.aemeasurable (by fun_prop)
+  map_eq := by
+    unfold step
+    simp only [trajMeasure, Kernel.trajMeasure]
+    rw [← Measure.deterministic_comp_eq_map (by fun_prop), Measure.comp_assoc,
+      Kernel.deterministic_comp_eq_map, Kernel.traj_zero_map_eval_zero,
+      Measure.deterministic_comp_eq_map, Measure.map_map (by fun_prop) (by fun_prop)]
+    simp
+
+lemma hasLaw_action_zero (alg : Algorithm α R) (env : Environment α R) :
+    HasLaw (action 0) alg.p0 (trajMeasure alg env) where
+  map_eq := by
+    rw [← fst_comp_step, ← Measure.map_map (by fun_prop) (by fun_prop),
+      (hasLaw_step_zero alg env).map_eq, ← Measure.fst, Measure.fst_compProd]
+
+lemma condDistrib_reward_zero [StandardBorelSpace R] [Nonempty R]
+    (alg : Algorithm α R) (env : Environment α R) :
+    condDistrib (reward 0) (action 0) (trajMeasure alg env)
+      =ᵐ[(trajMeasure alg env).map (action 0)] env.ν0 := by
+  have h_step := (hasLaw_step_zero alg env).map_eq
+  have h_action := (hasLaw_action_zero alg env).map_eq
+  rwa [condDistrib_ae_eq_iff_measure_eq_compProd₀ (by fun_prop) (by fun_prop), h_action]
+
+section DetAlgorithm
+
+variable {nextaction : (n : ℕ) → (Iic n → α × R) → α} {h_next : ∀ n, Measurable (nextaction n)}
+  {action0 : α} {env : Environment α R}
+
+lemma HasLaw_action_zero_detAlgorithm :
+    HasLaw (action 0) (Measure.dirac action0)
+      (trajMeasure (detAlgorithm nextaction h_next action0) env) where
+  map_eq := (hasLaw_action_zero _ _).map_eq
+
+lemma action_zero_detAlgorithm [MeasurableSingletonClass α] :
+    action 0 =ᵐ[trajMeasure (detAlgorithm nextaction h_next action0) env] fun _ ↦ action0 := by
+  have h_eq : ∀ᵐ x ∂((trajMeasure (detAlgorithm nextaction h_next action0) env).map (action 0)), x
+      = action0 := by
+    rw [(hasLaw_action_zero _ _).map_eq]
+    simp [detAlgorithm]
+  exact ae_of_ae_map (by fun_prop) h_eq
+
+lemma action_detAlgorithm_ae_eq (n : ℕ) :
+    action (n + 1) =ᵐ[trajMeasure (detAlgorithm nextaction h_next action0) env]
+      fun h ↦ nextaction n (fun i ↦ h i) := by
+  sorry
+
+example [MeasurableSingletonClass α] :
+    ∀ᵐ h ∂(trajMeasure (detAlgorithm nextaction h_next action0) env),
+    action 0 h = action0 ∧ ∀ n, action (n + 1) h = nextaction n (fun i ↦ h i) := by
+  rw [eventually_and, ae_all_iff]
+  exact ⟨action_zero_detAlgorithm, action_detAlgorithm_ae_eq⟩
+
+end DetAlgorithm
+
+end Learning