Independence lemmas

LeanBandits.lean

@@ -8,6 +8,7 @@ import LeanBandits.ForMathlib.IdentDistrib
 import LeanBandits.ForMathlib.KernelCompositionLemmas
 import LeanBandits.ForMathlib.KernelCompositionParallelComp
 import LeanBandits.ForMathlib.KernelSub
+import LeanBandits.ForMathlib.Measurable
 import LeanBandits.ForMathlib.SubGaussian
 import LeanBandits.ForMathlib.Traj
 import LeanBandits.Regret

LeanBandits/Algorithm.lean

@@ -6,6 +6,7 @@ Authors: Rémy Degenne, Paulo Rauber
 import Mathlib
 import LeanBandits.ForMathlib.CondDistrib
 import LeanBandits.ForMathlib.KernelCompositionLemmas
+import LeanBandits.ForMathlib.Measurable
 import LeanBandits.ForMathlib.Traj
 
 /-!
@@ -223,10 +224,8 @@ lemma measurable_hist_filtration (n : ℕ) : Measurable[Learning.filtration α R
 
 lemma measurable_action_filtration (n : ℕ) : Measurable[Learning.filtration α R n] (action n) := by
   simp only [Learning.filtration, Filtration.piLE_eq_comap_frestrictLe, ← hist_eq_frestrictLe]
-  rw [action_eq_eval_comp_hist, measurable_iff_comap_le, ← MeasurableSpace.comap_comp]
-  refine MeasurableSpace.comap_mono ?_
-  rw [← measurable_iff_comap_le]
-  fun_prop
+  rw [action_eq_eval_comp_hist]
+  exact measurable_comp_comap _ (by fun_prop)
 
 lemma adapted_action [TopologicalSpace α] [TopologicalSpace.PseudoMetrizableSpace α]
     [SecondCountableTopology α] [OpensMeasurableSpace α] :
@@ -235,10 +234,8 @@ lemma adapted_action [TopologicalSpace α] [TopologicalSpace.PseudoMetrizableSpa
 
 lemma measurable_reward_filtration (n : ℕ) : Measurable[Learning.filtration α R n] (reward n) := by
   simp only [Learning.filtration, Filtration.piLE_eq_comap_frestrictLe, ← hist_eq_frestrictLe]
-  rw [reward_eq_eval_comp_hist, measurable_iff_comap_le, ← MeasurableSpace.comap_comp]
-  refine MeasurableSpace.comap_mono ?_
-  rw [← measurable_iff_comap_le]
-  fun_prop
+  rw [reward_eq_eval_comp_hist]
+  exact measurable_comp_comap _ (by fun_prop)
 
 lemma adapted_reward [TopologicalSpace R] [TopologicalSpace.PseudoMetrizableSpace R]
     [SecondCountableTopology R] [OpensMeasurableSpace R] :

LeanBandits/Bandit.lean

@@ -146,6 +146,13 @@ lemma indepFun_eval_streamMeasure' (ν : Kernel α R) [IsMarkovKernel ν] {a b :
     IndepFun (fun ω n ↦ ω n a) (fun ω n ↦ ω n b) (Bandit.streamMeasure ν) := by
   sorry
 
+lemma indepFun_eval_snd_measure (alg : Algorithm α R) (ν : Kernel α R) [IsMarkovKernel ν]
+    {a b : α} (h : a ≠ b) :
+    IndepFun (fun ω n ↦ ω.2 n a) (fun ω n ↦ ω.2 n b) (Bandit.measure alg ν) := by
+  refine indepFun_snd_prod ?_ ?_ (indepFun_eval_streamMeasure' ν h) (Bandit.trajMeasure alg ν)
+  · exact Measurable.aemeasurable (by fun_prop)
+  · exact Measurable.aemeasurable (by fun_prop)
+
 end StreamMeasure
 
 /-- `arm n` is the arm pulled at time `n`. This is a random variable on the measurable space

LeanBandits/ETC.lean

@@ -6,7 +6,6 @@ Authors: Rémy Degenne
 import Mathlib.Probability.Moments.SubGaussian
 import LeanBandits.AlgorithmAndRandomVariables
 import LeanBandits.AlgorithmBuilding
-import LeanBandits.ForMathlib.IdentDistrib
 import LeanBandits.ForMathlib.SubGaussian
 import LeanBandits.Regret
 import LeanBandits.RewardByCountMeasure
@@ -178,12 +177,26 @@ lemma identDistrib_aux (m : ℕ) (a b : Fin K) :
     IdentDistrib
       (fun ω ↦ (∑ s ∈ Icc 1 m, rewardByCount a s ω.1 ω.2, ∑ s ∈ Icc 1 m, rewardByCount b s ω.1 ω.2))
       (fun ω ↦ (∑ s ∈ range m, ω.2 s a, ∑ s ∈ range m, ω.2 s b)) 𝔓 𝔓 := by
+  have : Nonempty (Fin K) := Fin.pos_iff_nonempty.mp hK
   have h1 (a : Fin K) :
-      IdentDistrib (fun ω s ↦ rewardByCount a s ω.1 ω.2) (fun ω s ↦ ω.2 s a) 𝔓 𝔓 := by
-    sorry
+      IdentDistrib (fun ω s ↦ rewardByCount a (s + 1) ω.1 ω.2) (fun ω s ↦ ω.2 s a) 𝔓 𝔓 :=
+    identDistrib_rewardByCount_stream a
   have h2 (a : Fin K) : IdentDistrib (fun ω ↦ ∑ s ∈ Icc 1 m, rewardByCount a s ω.1 ω.2)
       (fun ω ↦ ∑ s ∈ range m, ω.2 s a) 𝔓 𝔓 := by
-    sorry
+    have h_eq (ω : (ℕ → Fin K × ℝ) × (ℕ → Fin K → ℝ)) : ∑ s ∈ Icc 1 m, rewardByCount a s ω.1 ω.2
+        = ∑ s ∈ range m, rewardByCount a (s + 1) ω.1 ω.2 := by
+      let e : Icc 1 m ≃ range m :=
+      { toFun x := ⟨x - 1, by have h := x.2; simp only [mem_Icc] at h; grind⟩
+        invFun x := ⟨x + 1, by have h := x.2; simp only [mem_Icc]; grind⟩
+        left_inv x := by have h := x.2; simp only [mem_Icc] at h; grind
+        right_inv x := by have h := x.2; grind }
+      rw [← sum_coe_sort (Icc 1 m), ← sum_coe_sort (range m), sum_equiv e]
+      · simp
+      · simp only [univ_eq_attach, mem_attach, forall_const, Subtype.forall, mem_Icc,
+          forall_and_index]
+        grind
+    simp_rw [h_eq]
+    exact IdentDistrib.comp (h1 a) (u := fun p ↦ ∑ s ∈ range m, p s) (by fun_prop)
   by_cases hab : a = b
   · simp only [hab]
     exact (h2 b).comp (u := fun p ↦ (p, p)) (by fun_prop)
@@ -192,11 +205,11 @@ lemma identDistrib_aux (m : ℕ) (a b : Fin K) :
         𝔓 by
       exact this.comp (φ := fun p ↦ ∑ i ∈ Icc 1 m, p i) (ψ := fun p ↦ ∑ j ∈ Icc 1 m, p j)
         (by fun_prop) (by fun_prop)
-    sorry
+    exact indepFun_rewardByCount_of_ne hab
   · suffices IndepFun (fun ω s ↦ ω.2 s a) (fun ω s ↦ ω.2 s b) 𝔓 by
       exact this.comp (φ := fun p ↦ ∑ i ∈ range m, p i) (ψ := fun p ↦ ∑ j ∈ range m, p j)
         (by fun_prop) (by fun_prop)
-    sorry
+    exact indepFun_eval_snd_measure _ ν hab
 
 lemma prob_arm_mul_eq_le (hν : ∀ a, HasSubgaussianMGF (fun x ↦ x - (ν a)[id]) 1 (ν a)) (a : Fin K)
     (hm : m ≠ 0) :

LeanBandits/ForMathlib/CondDistrib.lean

@@ -18,7 +18,7 @@ variable {α β γ δ Ω Ω' : Type*}
   {m mα : MeasurableSpace α} {μ : Measure α} {mβ : MeasurableSpace β} {mγ : MeasurableSpace γ}
   {mδ : MeasurableSpace δ}
   [MeasurableSpace Ω] [StandardBorelSpace Ω] [Nonempty Ω]
-  [MeasurableSpace Ω'] [StandardBorelSpace Ω'] [Nonempty Ω']
+  [mΩ' : MeasurableSpace Ω'] [StandardBorelSpace Ω'] [Nonempty Ω']
   {X : α → β} {Y : α → Ω} {Z : α → Ω'} {T : α → γ}
 
 lemma MeasurableSpace.comap_prodMk (X : α → β) (Y : α → γ) :
@@ -844,6 +844,87 @@ lemma condIndepFun_of_exists_condDistrib_prod_ae_eq_prodMkLeft
   rw [Kernel.prodMkLeft_ae_eq_iff, Measure.snd_map_prodMk (by fun_prop)]
   exact hη_eq.symm
 
+/-- Law of `Y` conditioned on `X`. -/
+notation "𝓛[" Y " | " X "; " μ "]" => condDistrib Y X μ
+
+-- generalize to map instead of fst
+omit [Nonempty Ω'] in
+lemma condIndepFun_fst_prod [StandardBorelSpace α] [StandardBorelSpace β] [Nonempty β]
+    [StandardBorelSpace γ]
+    (hX : Measurable X) (hY : Measurable Y) (hZ : Measurable Z)
+    (ν : Measure γ) [IsProbabilityMeasure ν]
+    (h_indep : CondIndepFun (mΩ'.comap Z) hZ.comap_le Y X μ) :
+    CondIndepFun (mΩ'.comap (fun ω ↦ Z ω.1)) (hZ.comp measurable_fst).comap_le
+      (fun ω ↦ Y ω.1) (fun ω ↦ X ω.1) (μ.prod ν) := by
+  rw [condIndepFun_comap_iff_map_prod_eq_prod_condDistrib_prod_condDistrib (by fun_prop)
+    (by fun_prop) (by fun_prop)] at h_indep ⊢
+  have h1 : 𝓛[fun ω ↦ Y ω.1 | fun ω ↦ Z ω.1; μ.prod ν] =ᵐ[μ.map Z] 𝓛[Y | Z; μ] :=
+    condDistrib_fst_prod (Y := Y) (X := Z) (ν := ν) (μ := μ) (by fun_prop) (by fun_prop)
+  have h2 : 𝓛[fun ω ↦ X ω.1 | fun ω ↦ Z ω.1; μ.prod ν] =ᵐ[μ.map Z] 𝓛[X | Z; μ] :=
+    condDistrib_fst_prod (Y := X) (X := Z) (ν := ν) (μ := μ) (by fun_prop) (by fun_prop)
+  have h_fst1 : (μ.prod ν).map (fun ω ↦ Z ω.1) = μ.map Z := by
+    conv_rhs => rw [← Measure.fst_prod (μ := μ) (ν := ν), Measure.fst,
+      Measure.map_map (by fun_prop) (by fun_prop)]
+    rfl
+  have h_fst2 : (μ.prod ν).map (fun ω ↦ (Z ω.1, Y ω.1, X ω.1))
+      = μ.map (fun ω ↦ (Z ω, Y ω, X ω)) := by
+    conv_rhs => rw [← Measure.fst_prod (μ := μ) (ν := ν), Measure.fst,
+      Measure.map_map (by fun_prop) (by fun_prop)]
+    rfl
+  rw [h_fst1, h_fst2, h_indep]
+  refine Measure.bind_congr_right ?_
+  filter_upwards [h1, h2] with x hx1 hx2
+  simp_rw [Kernel.prod_apply]
+  rw [hx1, ← hx2]
+
+omit [StandardBorelSpace Ω] [Nonempty Ω] in
+lemma indepFun_fst_prod (hX : AEMeasurable X μ) (hY : AEMeasurable Y μ) (h_indep : IndepFun X Y μ)
+    (ν : Measure γ) [IsProbabilityMeasure ν] :
+    IndepFun (fun ω ↦ X ω.1) (fun ω ↦ Y ω.1) (μ.prod ν) := by
+  rw [indepFun_iff_map_prod_eq_prod_map_map (by fun_prop) (by fun_prop)] at h_indep ⊢
+  have :  AEMeasurable (fun ω ↦ (X ω, Y ω)) (Measure.map Prod.fst (μ.prod ν)) := by
+    simp only [Measure.map_fst_prod, measure_univ, one_smul]
+    fun_prop
+  have :  AEMeasurable X (Measure.map Prod.fst (μ.prod ν)) := by
+    simp only [Measure.map_fst_prod, measure_univ, one_smul]
+    fun_prop
+  have :  AEMeasurable Y (Measure.map Prod.fst (μ.prod ν)) := by
+    simp only [Measure.map_fst_prod, measure_univ, one_smul]
+    fun_prop
+  have h : (μ.prod ν).map (fun ω ↦ (X ω.1, Y ω.1)) = μ.map (fun ω ↦ (X ω, Y ω)) := by
+    conv_rhs => rw [← Measure.fst_prod (μ := μ) (ν := ν), Measure.fst,
+      AEMeasurable.map_map_of_aemeasurable (by fun_prop) (by fun_prop)]
+    rfl
+  rw [h, h_indep]
+  conv_lhs => rw [← Measure.fst_prod (μ := μ) (ν := ν), Measure.fst,
+      AEMeasurable.map_map_of_aemeasurable (by fun_prop) (by fun_prop),
+      AEMeasurable.map_map_of_aemeasurable (by fun_prop) (by fun_prop)]
+  rfl
+
+omit [StandardBorelSpace Ω] [Nonempty Ω] in
+lemma indepFun_snd_prod (hX : AEMeasurable X μ) (hY : AEMeasurable Y μ) (h_indep : IndepFun X Y μ)
+    (ν : Measure γ) [IsProbabilityMeasure ν] :
+    IndepFun (fun ω ↦ X ω.2) (fun ω ↦ Y ω.2) (ν.prod μ) := by
+  rw [indepFun_iff_map_prod_eq_prod_map_map (by fun_prop) (by fun_prop)] at h_indep ⊢
+  have :  AEMeasurable (fun ω ↦ (X ω, Y ω)) (Measure.map Prod.snd (ν.prod μ)) := by
+    simp only [Measure.map_snd_prod, measure_univ, one_smul]
+    fun_prop
+  have :  AEMeasurable X (Measure.map Prod.snd (ν.prod μ)) := by
+    simp only [Measure.map_snd_prod, measure_univ, one_smul]
+    fun_prop
+  have :  AEMeasurable Y (Measure.map Prod.snd (ν.prod μ)) := by
+    simp only [Measure.map_snd_prod, measure_univ, one_smul]
+    fun_prop
+  have h : (ν.prod μ).map (fun ω ↦ (X ω.2, Y ω.2)) = μ.map (fun ω ↦ (X ω, Y ω)) := by
+    conv_rhs => rw [← Measure.snd_prod (μ := ν) (ν := μ), Measure.snd,
+      AEMeasurable.map_map_of_aemeasurable (by fun_prop) (by fun_prop)]
+    rfl
+  rw [h, h_indep]
+  conv_lhs => rw [← Measure.snd_prod (μ := ν) (ν := μ), Measure.snd,
+      AEMeasurable.map_map_of_aemeasurable (by fun_prop) (by fun_prop),
+      AEMeasurable.map_map_of_aemeasurable (by fun_prop) (by fun_prop)]
+  rfl
+
 end CondDistrib
 
 section Cond

LeanBandits/ForMathlib/IdentDistrib.lean

@@ -4,13 +4,14 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Rémy Degenne
 -/
 import Mathlib.Probability.IdentDistrib
+import Mathlib.Probability.Independence.InfinitePi
 
 open MeasureTheory ProbabilityTheory Finset
 open scoped ENNReal NNReal
 
 namespace ProbabilityTheory
 
-variable {Ω Ω' : Type*} {mΩ : MeasurableSpace Ω} {mΩ' : MeasurableSpace Ω'}
+variable {Ω Ω' ι : Type*} {mΩ : MeasurableSpace Ω} {mΩ' : MeasurableSpace Ω'}
     {μ : Measure Ω} {ν : Measure Ω'} {X Y : Ω → ℝ} {Z W : Ω' → ℝ}
 
 lemma IdentDistrib.prod [IsFiniteMeasure μ] [IsFiniteMeasure ν]
@@ -24,4 +25,20 @@ lemma IdentDistrib.prod [IsFiniteMeasure μ] [IsFiniteMeasure ν]
       (indepFun_iff_map_prod_eq_prod_map_map hXZ.aemeasurable_snd hYW.aemeasurable_snd).mp hZW,
       hXZ.map_eq, hYW.map_eq]
 
+lemma IdentDistrib.pi [Countable ι] [IsProbabilityMeasure μ] [IsProbabilityMeasure ν]
+    {X : (i : ι) → Ω → ℝ} {Y : (i : ι) → Ω' → ℝ}
+    (h : ∀ i, IdentDistrib (X i) (Y i) μ ν) (hX_ind : iIndepFun X μ) (hY_ind : iIndepFun Y ν) :
+    IdentDistrib (fun ω ↦ (X · ω)) (fun ω ↦ (Y · ω)) μ ν where
+  aemeasurable_fst := by
+    rw [aemeasurable_pi_iff]
+    exact fun i ↦ (h i).aemeasurable_fst
+  aemeasurable_snd := by
+    rw [aemeasurable_pi_iff]
+    exact fun i ↦ (h i).aemeasurable_snd
+  map_eq := by
+    rw [(iIndepFun_iff_map_fun_eq_infinitePi_map₀' (fun i ↦ (h i).aemeasurable_fst)).mp hX_ind,
+      (iIndepFun_iff_map_fun_eq_infinitePi_map₀' (fun i ↦ (h i).aemeasurable_snd)).mp hY_ind]
+    congr with i
+    rw [(h i).map_eq]
+
 end ProbabilityTheory

LeanBandits/ForMathlib/Measurable.lean

@@ -0,0 +1,22 @@
+/-
+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 Mathlib
+
+/-!
+# Measurability lemmas
+-/
+
+namespace MeasureTheory
+
+lemma measurable_comp_comap {α β γ : Type*} {mβ : MeasurableSpace β} {mγ : MeasurableSpace γ}
+    (f : α → β) {g : β → γ} (hg : Measurable g) :
+    Measurable[mβ.comap f] (g ∘ f) := by
+  rw [measurable_iff_comap_le, ← MeasurableSpace.comap_comp]
+  refine MeasurableSpace.comap_mono ?_
+  rw [← measurable_iff_comap_le]
+  exact hg
+
+end MeasureTheory