diff --git a/LeanMachineLearning.lean b/LeanMachineLearning.lean
index 97d0683c..1e4e7b10 100644
--- a/LeanMachineLearning.lean
+++ b/LeanMachineLearning.lean
@@ -6,7 +6,9 @@ public import LeanMachineLearning.ForMathlib.MeasureTheory.Order.Lattice
 public import LeanMachineLearning.ForMathlib.MeasureTheory.Order.MeasurableArg
 public import LeanMachineLearning.ForMathlib.MeasureTheory.OuterMeasure.Basic
 public import LeanMachineLearning.ForMathlib.Order.Interval.Finset
+public import LeanMachineLearning.ForMathlib.Probability.ConditionalProbability
 public import LeanMachineLearning.ForMathlib.Probability.HasCondDistrib
+public import LeanMachineLearning.ForMathlib.Probability.HasLaw
 public import LeanMachineLearning.ForMathlib.Probability.Independence.CondDistrib
 public import LeanMachineLearning.ForMathlib.Probability.Independence.CondIndepFun
 public import LeanMachineLearning.ForMathlib.Probability.Independence.IndepFun
diff --git a/LeanMachineLearning/ForMathlib/Probability/ConditionalProbability.lean b/LeanMachineLearning/ForMathlib/Probability/ConditionalProbability.lean
new file mode 100644
index 00000000..684cf434
--- /dev/null
+++ b/LeanMachineLearning/ForMathlib/Probability/ConditionalProbability.lean
@@ -0,0 +1,30 @@
+/-
+Copyright (c) 2026 Rémy Degenne. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Rémy Degenne
+-/
+module
+
+public import Mathlib.MeasureTheory.Measure.Prod
+public import Mathlib.Probability.ConditionalProbability
+
+/-! # Lemmas about conditional probability
+-/
+
+@[expose] public section
+
+open MeasureTheory
+
+namespace ProbabilityTheory
+
+variable {α β : Type*} {mα : MeasurableSpace α} {mβ : MeasurableSpace β}
+
+/-- Conditioning a product measure on an event of the first coordinate amounts to conditioning
+the first measure. -/
+lemma cond_prod_univ {μ : Measure α} [SFinite μ] {ν : Measure β} [IsProbabilityMeasure ν]
+    (s : Set α) :
+    (μ.prod ν)[|s ×ˢ Set.univ] = (μ[|s]).prod ν := by
+  simp only [cond, Measure.prod_prod, measure_univ, mul_one, Measure.prod_smul_left,
+    ← Measure.prod_restrict, Measure.restrict_univ]
+
+end ProbabilityTheory
diff --git a/LeanMachineLearning/ForMathlib/Probability/HasCondDistrib.lean b/LeanMachineLearning/ForMathlib/Probability/HasCondDistrib.lean
index e9ef0f0f..65528559 100644
--- a/LeanMachineLearning/ForMathlib/Probability/HasCondDistrib.lean
+++ b/LeanMachineLearning/ForMathlib/Probability/HasCondDistrib.lean
@@ -168,6 +168,63 @@ lemma ae_eq_of_hasCondDistrib_deterministic [MeasurableEq Ω] [SFinite μ] {f :
     AEMeasurable.map_map_of_aemeasurable (by fun_prop) (by fun_prop)]
   rfl
 
+section Cond
+
+variable [IsSFiniteKernel κ]
+
+/-- If the conditional distribution of `Y` given `X` is a kernel `κ` which is constant equal to `η`
+on a measurable set `s`, then `μ (X ⁻¹' s ∩ Y ⁻¹' u) = μ (X ⁻¹' s) * η u` for all measurable `u`. -/
+lemma HasCondDistrib.measure_inter_preimage_eq_mul_of_eqOn_const [SFinite μ]
+    (h : HasCondDistrib Y X κ μ) {s : Set β} (hs : MeasurableSet s) {η : Measure Ω}
+    (hκ : Set.EqOn κ (fun _ ↦ η) s) {u : Set Ω} (hu : MeasurableSet u) :
+    μ (X ⁻¹' s ∩ Y ⁻¹' u) = μ (X ⁻¹' s) * η u := by
+  have h_eq : X ⁻¹' s ∩ Y ⁻¹' u = (fun ω ↦ (X ω, Y ω)) ⁻¹' (s ×ˢ u) := by
+    ext ω
+    simp
+  rw [h_eq, ← Measure.map_apply_of_aemeasurable h.aemeasurable (hs.prod hu), h.map_eq,
+    Measure.compProd_apply_prod hs hu,
+    setLIntegral_congr_fun hs (g := fun _ ↦ η u) (fun x hx ↦ by rw [hκ hx]),
+    setLIntegral_const, Measure.map_apply_of_aemeasurable h.aemeasurable_fst hs, mul_comm]
+
+variable [IsFiniteMeasure μ]
+
+/-- If the conditional distribution of `Y` given `X` is a kernel `κ` which is constant equal to `η`
+on a measurable set `s`, then the law of `Y` under `μ` conditioned on `X ∈ s` is `η`. -/
+lemma HasCondDistrib.hasLaw_cond (h : HasCondDistrib Y X κ μ) (hY : Measurable Y)
+    {s : Set β} (hs : MeasurableSet s) {η : Measure Ω} (hκ : Set.EqOn κ (fun _ ↦ η) s)
+    (hμs : μ (X ⁻¹' s) ≠ 0) :
+    HasLaw Y η μ[|X ⁻¹' s] where
+  aemeasurable := hY.aemeasurable
+  map_eq := by
+    ext u hu
+    rw [Measure.map_apply hY hu, cond_apply' (hu.preimage hY),
+      h.measure_inter_preimage_eq_mul_of_eqOn_const hs hκ hu, ← mul_assoc,
+      ENNReal.inv_mul_cancel hμs (measure_ne_top _ _), one_mul]
+
+/-- If the conditional distribution of `Y` given `X` is a kernel `κ` which is constant on a
+measurable set `s`, then `X` and `Y` are independent under `μ` conditioned on `X ∈ s`. -/
+lemma HasCondDistrib.indepFun_cond (h : HasCondDistrib Y X κ μ) (hX : Measurable X)
+    {s : Set β} (hs : MeasurableSet s) {η : Measure Ω} (hκ : Set.EqOn κ (fun _ ↦ η) s) :
+    X ⟂ᵢ[μ[|X ⁻¹' s]] Y := by
+  by_cases hμs : μ (X ⁻¹' s) = 0
+  · rw [cond_eq_zero.2 (Or.inr hμs)]
+    simp [indepFun_iff_measure_inter_preimage_eq_mul]
+  rw [indepFun_iff_measure_inter_preimage_eq_mul]
+  intro t u ht hu
+  have h1 : X ⁻¹' s ∩ (X ⁻¹' t ∩ Y ⁻¹' u) = X ⁻¹' (s ∩ t) ∩ Y ⁻¹' u := by
+    ext ω
+    simp only [Set.mem_inter_iff, Set.mem_preimage]
+    tauto
+  rw [cond_apply (hs.preimage hX), cond_apply (hs.preimage hX), cond_apply (hs.preimage hX), h1,
+    ← Set.preimage_inter,
+    h.measure_inter_preimage_eq_mul_of_eqOn_const (hs.inter ht) (hκ.mono Set.inter_subset_left)
+      hu,
+    h.measure_inter_preimage_eq_mul_of_eqOn_const hs hκ hu,
+    ← mul_assoc (μ (X ⁻¹' s))⁻¹ (μ (X ⁻¹' s)) (η u),
+    ENNReal.inv_mul_cancel hμs (measure_ne_top _ _), one_mul, mul_assoc]
+
+end Cond
+
 variable [StandardBorelSpace Ω] [Nonempty Ω] [StandardBorelSpace Ω'] [Nonempty Ω']
 
 lemma HasCondDistrib.condDistrib_eq [IsFiniteMeasure μ] [IsFiniteKernel κ]
diff --git a/LeanMachineLearning/ForMathlib/Probability/HasLaw.lean b/LeanMachineLearning/ForMathlib/Probability/HasLaw.lean
new file mode 100644
index 00000000..b8c69102
--- /dev/null
+++ b/LeanMachineLearning/ForMathlib/Probability/HasLaw.lean
@@ -0,0 +1,140 @@
+/-
+Copyright (c) 2026 Rémy Degenne. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Rémy Degenne
+-/
+module
+
+public import Mathlib.MeasureTheory.Constructions.Cylinders
+public import Mathlib.MeasureTheory.Integral.Indicator
+public import Mathlib.Probability.ConditionalProbability
+public import Mathlib.Probability.HasLaw
+public import Mathlib.Probability.IdentDistrib
+
+/-! # Lemmas about `HasLaw`
+-/
+
+@[expose] public section
+
+open MeasureTheory Filter
+open scoped Topology
+
+namespace ProbabilityTheory
+
+variable {Ω 𝓧 : Type*} {mΩ : MeasurableSpace Ω} {m𝓧 : MeasurableSpace 𝓧} {P : Measure Ω}
+
+lemma _root_.AEMeasurable.hasLaw_map {X : Ω → 𝓧} (hX : AEMeasurable X P) :
+    HasLaw X (P.map X) P := ⟨hX, rfl⟩
+
+lemma _root_.Measurable.hasLaw_map {X : Ω → 𝓧} (hX : Measurable X) (P : Measure Ω) :
+    HasLaw X (P.map X) P := ⟨hX.aemeasurable, rfl⟩
+
+section Cond
+
+variable {ι : Type*} [Countable ι] {mι : MeasurableSpace ι} [MeasurableSingletonClass ι]
+
+/-- Two random variables which are identically distributed conditionally on each atom of a
+countable measurable partition are identically distributed. -/
+lemma identDistrib_of_forall_identDistrib_cond [IsFiniteMeasure P] {g : Ω → ι}
+    (hg : Measurable g) {X Y : Ω → 𝓧} (hX : Measurable X) (hY : Measurable Y)
+    (h : ∀ i, IdentDistrib X Y P[|g ⁻¹' {i}] P[|g ⁻¹' {i}]) :
+    IdentDistrib X Y P P where
+  aemeasurable_fst := hX.aemeasurable
+  aemeasurable_snd := hY.aemeasurable
+  map_eq := by
+    ext s hs
+    rw [Measure.map_apply hX hs, Measure.map_apply hY hs]
+    have h_union (t : Set Ω) : t = ⋃ i, t ∩ g ⁻¹' {i} := by ext; simp
+    have h_disj (t : Set Ω) : Pairwise (Function.onFun Disjoint fun i ↦ t ∩ g ⁻¹' {i}) := by
+      intro i j hij
+      rw [Function.onFun, Set.disjoint_left]
+      rintro x ⟨-, hi⟩ ⟨-, hj⟩
+      exact hij ((show g x = i from hi).symm.trans hj)
+    rw [h_union (X ⁻¹' s), h_union (Y ⁻¹' s),
+      measure_iUnion (h_disj _) fun i ↦ (hs.preimage hX).inter (hg (measurableSet_singleton i)),
+      measure_iUnion (h_disj _) fun i ↦ (hs.preimage hY).inter (hg (measurableSet_singleton i))]
+    refine tsum_congr fun i ↦ ?_
+    rw [Set.inter_comm, ← cond_mul_eq_inter (hg (measurableSet_singleton i)),
+      Set.inter_comm _ (g ⁻¹' {i}), ← cond_mul_eq_inter (hg (measurableSet_singleton i)),
+      ← Measure.map_apply hX hs, ← Measure.map_apply hY hs, (h i).map_eq]
+
+/-- If a random variable has law `μ` conditionally on each atom of positive probability of a
+countable measurable partition, then it has law `μ`. -/
+lemma hasLaw_of_forall_hasLaw_cond [IsProbabilityMeasure P] {g : Ω → ι} (hg : Measurable g)
+    {X : Ω → 𝓧} (hX : Measurable X) {μ : Measure 𝓧}
+    (h : ∀ i, P (g ⁻¹' {i}) ≠ 0 → HasLaw X μ P[|g ⁻¹' {i}]) :
+    HasLaw X μ P where
+  aemeasurable := hX.aemeasurable
+  map_eq := by
+    ext s hs
+    rw [Measure.map_apply hX hs]
+    have h_union : X ⁻¹' s = ⋃ i, X ⁻¹' s ∩ g ⁻¹' {i} := by ext; simp
+    have h_disj : Pairwise (Function.onFun Disjoint fun i ↦ X ⁻¹' s ∩ g ⁻¹' {i}) := by
+      intro i j hij
+      rw [Function.onFun, Set.disjoint_left]
+      rintro x ⟨-, hi⟩ ⟨-, hj⟩
+      exact hij ((show g x = i from hi).symm.trans hj)
+    have h_univ : Set.univ = ⋃ i, g ⁻¹' {i} := by ext; simp
+    have h_disj_univ : Pairwise (Function.onFun Disjoint fun i ↦ g ⁻¹' {i}) := by
+      intro i j hij
+      rw [Function.onFun, Set.disjoint_left]
+      exact fun x hi hj ↦ hij ((show g x = i from hi).symm.trans hj)
+    calc P (X ⁻¹' s)
+    _ = ∑' i, P (X ⁻¹' s ∩ g ⁻¹' {i}) := by
+      conv_lhs => rw [h_union]
+      exact measure_iUnion h_disj fun i ↦ (hs.preimage hX).inter (hg (measurableSet_singleton i))
+    _ = ∑' i, μ s * P (g ⁻¹' {i}) := by
+      refine tsum_congr fun i ↦ ?_
+      rw [Set.inter_comm, ← cond_mul_eq_inter (hg (measurableSet_singleton i))]
+      by_cases hi : P (g ⁻¹' {i}) = 0
+      · simp [hi]
+      · rw [← Measure.map_apply hX hs, (h i hi).map_eq]
+    _ = μ s := by
+      rw [ENNReal.tsum_mul_left, ← measure_iUnion h_disj_univ
+        fun i ↦ hg (measurableSet_singleton i), ← h_univ, measure_univ, mul_one]
+
+end Cond
+
+section Pi
+
+variable {ι : Type*} {𝓧 : ι → Type*} [∀ i, MeasurableSpace (𝓧 i)]
+
+/-- Let `Y n : Ω → Π i, 𝓧 i` be random variables with law `μ`, indexed by a countably generated
+filter `L`. If for every `ω` and `i`, `Y n ω i` is eventually equal to `Y' ω i` along `L`, then `Y'`
+also has law `μ`. -/
+lemma hasLaw_of_forall_eventually_eq [IsFiniteMeasure P] {κ : Type*} {L : Filter κ} [L.NeBot]
+    [L.IsCountablyGenerated] {μ : Measure (Π i, 𝓧 i)} {Y : κ → Ω → Π i, 𝓧 i} {Y' : Ω → Π i, 𝓧 i}
+    (hY : ∀ n, Measurable (Y n)) (hY' : AEMeasurable Y' P)
+    (h_law : ∀ n, HasLaw (Y n) μ P) (h_lim : ∀ ω i, ∀ᶠ n in L, Y n ω i = Y' ω i) :
+    HasLaw Y' μ P where
+  aemeasurable := hY'
+  map_eq := by
+    refine ext_of_generate_finite (measurableCylinders _) generateFrom_measurableCylinders.symm
+      isPiSystem_measurableCylinders (fun s hs ↦ ?_) ?_
+    · obtain ⟨I, S, hS, rfl⟩ := (mem_measurableCylinders s).1 hs
+      rw [Measure.map_apply_of_aemeasurable hY' (hS.cylinder _)]
+      have h_tendsto : Tendsto (fun n ↦ P (Y n ⁻¹' cylinder I S)) L
+          (𝓝 (P (Y' ⁻¹' cylinder I S))) := by
+        refine tendsto_measure_of_tendsto_indicator_of_isFiniteMeasure L P
+          (fun n ↦ (hS.cylinder _).preimage (hY n)) fun ω ↦ ?_
+        have h_ev : ∀ᶠ n in L, ∀ i ∈ I, Y n ω i = Y' ω i :=
+          (eventually_all_finset I).2 fun i _ ↦ h_lim ω i
+        filter_upwards [h_ev] with n hn
+        simp only [Set.mem_preimage, mem_cylinder]
+        have : I.restrict (Y n ω) = I.restrict (Y' ω) := funext fun i ↦ hn i i.2
+        rw [this]
+      have h_const : Tendsto (fun n ↦ P (Y n ⁻¹' cylinder I S)) L (𝓝 (μ (cylinder I S))) := by
+        have : (fun n ↦ P (Y n ⁻¹' cylinder I S)) = fun _ ↦ μ (cylinder I S) := by
+          funext n
+          rw [← Measure.map_apply (hY n) (hS.cylinder _), (h_law n).map_eq]
+        rw [this]
+        exact tendsto_const_nhds
+      exact tendsto_nhds_unique h_tendsto h_const
+    · obtain ⟨n⟩ := L.nonempty_of_neBot
+      rw [Measure.map_apply_of_aemeasurable hY' MeasurableSet.univ, Set.preimage_univ,
+        ← (h_law n).map_eq,
+        Measure.map_apply (hY n) MeasurableSet.univ, Set.preimage_univ]
+
+end Pi
+
+end ProbabilityTheory
diff --git a/LeanMachineLearning/ForMathlib/Probability/Independence/IndepFun.lean b/LeanMachineLearning/ForMathlib/Probability/Independence/IndepFun.lean
index cd596ffa..93dbf147 100644
--- a/LeanMachineLearning/ForMathlib/Probability/Independence/IndepFun.lean
+++ b/LeanMachineLearning/ForMathlib/Probability/Independence/IndepFun.lean
@@ -59,6 +59,24 @@ lemma indepFun_cond_comp {α β γ δ : Type*} {mα : MeasurableSpace α} {mβ :
   simp_rw [h_preim]
   exact indepFun_cond_of_indepFun hXY hY (hZ (measurableSet_singleton z))
 
+/-- Under `μ` conditioned on the event `X = b`, the random variable `X` is almost surely constant,
+hence independent of any other random variable. -/
+lemma indepFun_cond_preimage_singleton_left {α β γ : Type*} {mα : MeasurableSpace α}
+    {mβ : MeasurableSpace β} {mγ : MeasurableSpace γ} [MeasurableSingletonClass β] {μ : Measure α}
+    {X : α → β} (hX : Measurable X) (b : β) (Y : α → γ) :
+    X ⟂ᵢ[μ[|X ⁻¹' {b}]] Y :=
+  (indepFun_const_left b Y).congr
+    (ae_cond_of_forall_mem (hX (measurableSet_singleton b)) fun x hx ↦ (hx : X x = b).symm)
+    Filter.EventuallyEq.rfl
+
+/-- Under `μ` conditioned on the event `X = b`, the random variable `X` is almost surely constant,
+hence independent of any other random variable. -/
+lemma indepFun_cond_preimage_singleton_right {α β γ : Type*} {mα : MeasurableSpace α}
+    {mβ : MeasurableSpace β} {mγ : MeasurableSpace γ} [MeasurableSingletonClass β] {μ : Measure α}
+    {X : α → β} (hX : Measurable X) (b : β) (Y : α → γ) :
+    Y ⟂ᵢ[μ[|X ⁻¹' {b}]] X :=
+  (indepFun_cond_preimage_singleton_left hX b Y).symm
+
 lemma iIndepFun_nat_iff_forall_indepFun [IsProbabilityMeasure μ] {X : ℕ → Ω → E}
     (hX : ∀ n, AEMeasurable (X n) μ) :
     iIndepFun X μ ↔ ∀ n, X (n + 1) ⟂ᵢ[μ] fun ω (i : Iic n) ↦ X i ω := by
diff --git a/LeanMachineLearning/Online/Bandit/ArrayProbSpace.lean b/LeanMachineLearning/Online/Bandit/ArrayProbSpace.lean
index 243bc6d6..78c77ad4 100644
--- a/LeanMachineLearning/Online/Bandit/ArrayProbSpace.lean
+++ b/LeanMachineLearning/Online/Bandit/ArrayProbSpace.lean
@@ -66,6 +66,13 @@ lemma hasLaw_eval_eval_streamMeasure (ν : Kernel 𝓐 𝓡) [IsMarkovKernel ν]
     HasLaw (fun h : ℕ → 𝓐 → 𝓡 ↦ h n a) (ν a) (streamMeasure ν) :=
   (hasLaw_eval_infinitePi ν a).comp (hasLaw_eval_streamMeasure ν n)
 
+/-- Under a product measure `μ.prod (streamMeasure ν)`, the entry `(n, a)` of the reward array has
+law `ν a`. -/
+lemma hasLaw_snd_apply_prod_streamMeasure {Ω : Type*} {mΩ : MeasurableSpace Ω} (μ : Measure Ω)
+    [IsProbabilityMeasure μ] (ν : Kernel 𝓐 𝓡) [IsMarkovKernel ν] (n : ℕ) (a : 𝓐) :
+    HasLaw (fun ω : Ω × (ℕ → 𝓐 → 𝓡) ↦ ω.2 n a) (ν a) (μ.prod (streamMeasure ν)) :=
+  (hasLaw_eval_eval_streamMeasure ν n a).comp (hasLaw_snd_prod μ _)
+
 lemma identDistrib_eval_eval_id_streamMeasure (ν : Kernel 𝓐 𝓡) [IsMarkovKernel ν] (n : ℕ) (a : 𝓐) :
     IdentDistrib (fun h : ℕ → 𝓐 → 𝓡 ↦ h n a) id (streamMeasure ν) (ν a) where
   aemeasurable_fst := Measurable.aemeasurable (by fun_prop)
@@ -107,6 +114,57 @@ lemma indepFun_eval_streamMeasure' (ν : Kernel 𝓐 𝓡) [IsMarkovKernel ν] {
     IndepFun (fun ω n ↦ ω n a) (fun ω n ↦ ω n b) (streamMeasure ν) :=
   indepFun_proj_infinitePi_infinitePi h
 
+/-- Under a product measure `μ.prod (streamMeasure ν)`, the entries of the reward array are
+independent. -/
+lemma iIndepFun_snd_apply_prod_streamMeasure {Ω : Type*} {mΩ : MeasurableSpace Ω} (μ : Measure Ω)
+    [IsProbabilityMeasure μ] (ν : Kernel 𝓐 𝓡) [IsMarkovKernel ν] :
+    iIndepFun (fun (p : ℕ × 𝓐) (ω : Ω × (ℕ → 𝓐 → 𝓡)) ↦ ω.2 p.1 p.2)
+      (μ.prod (streamMeasure ν)) := by
+  have h_snd : (μ.prod (streamMeasure ν)).map Prod.snd = streamMeasure ν := Measure.snd_prod
+  rw [iIndepFun_iff_map_fun_eq_infinitePi_map (fun _ ↦ by fun_prop)]
+  calc (μ.prod (streamMeasure ν)).map (fun ω (i : ℕ × 𝓐) ↦ ω.2 i.1 i.2)
+  _ = ((μ.prod (streamMeasure ν)).map Prod.snd).map (fun z (i : ℕ × 𝓐) ↦ z i.1 i.2) := by
+    rw [Measure.map_map (by fun_prop) measurable_snd]
+    rfl
+  _ = Measure.infinitePi fun i : ℕ × 𝓐 ↦ (streamMeasure ν).map (fun z ↦ z i.1 i.2) := by
+    rw [h_snd]
+    exact (iIndepFun_iff_map_fun_eq_infinitePi_map (fun _ ↦ by fun_prop)).1
+      (iIndepFun_eval_streamMeasure ν)
+  _ = Measure.infinitePi fun i : ℕ × 𝓐 ↦
+      (μ.prod (streamMeasure ν)).map (fun ω ↦ ω.2 i.1 i.2) := by
+    refine congrArg _ (funext fun i ↦ ?_)
+    conv_lhs => rw [← h_snd]
+    rw [Measure.map_map (by fun_prop) measurable_snd]
+    rfl
+
+/-- Under a product measure `μ.prod (streamMeasure ν)`, the entry `(m, a)` of the reward array is
+independent of the pair formed by the first coordinate and the reward array in which the entry
+`(m, a)` is replaced by a constant `x`. -/
+lemma indepFun_snd_apply_prod_streamMeasure_update [DecidableEq 𝓐] {Ω : Type*}
+    {mΩ : MeasurableSpace Ω} (μ : Measure Ω) [IsProbabilityMeasure μ] (ν : Kernel 𝓐 𝓡)
+    [IsMarkovKernel ν] (m : ℕ) (a : 𝓐) (x : 𝓡) :
+    (fun ω : Ω × (ℕ → 𝓐 → 𝓡) ↦ ω.2 m a) ⟂ᵢ[μ.prod (streamMeasure ν)]
+      (fun ω ↦ (ω.1, fun i b ↦ if i = m ∧ b = a then x else ω.2 i b)) := by
+  let T : (ℕ → 𝓐 → 𝓡) → (ℕ → 𝓐 → 𝓡) := fun z i b ↦ if i = m ∧ b = a then x else z i b
+  have hT : Measurable[⨆ p ∈ {p : ℕ × 𝓐 | p ≠ (m, a)},
+      MeasurableSpace.comap (fun z : ℕ → 𝓐 → 𝓡 ↦ z p.1 p.2) inferInstance] T := by
+    rw [measurable_iff_comap_le, MeasurableSpace.comap_pi]
+    refine iSup_le fun i ↦ ?_
+    rw [MeasurableSpace.comap_pi]
+    refine iSup_le fun b ↦ ?_
+    by_cases hib : i = m ∧ b = a
+    · obtain ⟨rfl, rfl⟩ := hib
+      simp only [T, and_self, ↓reduceIte, MeasurableSpace.comap_const]
+      exact bot_le
+    · simp only [T, hib, ↓reduceIte]
+      refine le_iSup₂_of_le (i, b) ?_ le_rfl
+      simpa only [Set.mem_ofPred_eq, ne_eq, Prod.mk.injEq] using hib
+  have hTm : Measurable T :=
+    hT.mono (iSup₂_le fun p _ ↦ Measurable.comap_le (by fun_prop)) le_rfl
+  have h := (iIndepFun_eval_streamMeasure ν).indepFun_of_measurable_iSup_comap
+    (fun _ ↦ by fun_prop) (i := (m, a)) (by simp) hT
+  exact h.snd_prod (μ := μ) (by fun_prop) hTm
+
 end StreamMeasure
 
 namespace ArrayModel
@@ -133,7 +191,7 @@ lemma hasLaw_fst_apply_arrayMeasure (ν : Kernel 𝓐 𝓡) [IsMarkovKernel ν]
 
 lemma hasLaw_snd_apply_arrayMeasure (ν : Kernel 𝓐 𝓡) [IsMarkovKernel ν] (n : ℕ) (a : 𝓐) :
     HasLaw (fun ω : probSpace 𝓐 𝓡 ↦ ω.2 n a) (ν a) (arrayMeasure ν) :=
-  (hasLaw_eval_eval_streamMeasure ν n a).comp (hasLaw_snd_prod _ _)
+  hasLaw_snd_apply_prod_streamMeasure _ ν n a
 
 lemma map_snd_apply_arrayMeasure {ν : Kernel 𝓐 𝓡} [IsMarkovKernel ν] (n : ℕ) (a : 𝓐) :
     (arrayMeasure ν).map (fun ω ↦ ω.2 n a) = ν a :=
diff --git a/LeanMachineLearning/Online/Bandit/RewardByCountMeasure.lean b/LeanMachineLearning/Online/Bandit/RewardByCountMeasure.lean
index fd9fcac6..1f60a429 100644
--- a/LeanMachineLearning/Online/Bandit/RewardByCountMeasure.lean
+++ b/LeanMachineLearning/Online/Bandit/RewardByCountMeasure.lean
@@ -5,6 +5,8 @@ Authors: Rémy Degenne
 -/
 module
 
+public import LeanMachineLearning.ForMathlib.Probability.ConditionalProbability
+public import LeanMachineLearning.ForMathlib.Probability.HasLaw
 public import LeanMachineLearning.Online.Bandit.ArrayProbSpace
 
 /-! # Laws of `stepsUntil` and `rewardByCount`
@@ -24,21 +26,6 @@ variable {𝓐 Ω : Type*} {m𝓐 : MeasurableSpace 𝓐} {mΩ : MeasurableSpace
 
 local notation "𝔓" => P.prod (streamMeasure ν)
 
-omit [DecidableEq 𝓐] in
-lemma hasLaw_Z (a : 𝓐) (m : ℕ) :
-  HasLaw (fun ω ↦ ω.2 m a) (ν a) 𝔓 where
-  map_eq := by
-    calc (𝔓).map (fun ω ↦ ω.2 m a)
-    _ = ((𝔓).snd).map (fun ω ↦ ω m a) := by
-      rw [Measure.snd, Measure.map_map (by fun_prop) (by fun_prop)]
-      rfl
-    _ = (streamMeasure ν).map (fun ω ↦ ω m a) := by simp
-    _ = ((Measure.infinitePi fun _ ↦ Measure.infinitePi ν).map (fun ω ↦ ω m)).map
-        (fun ω ↦ ω a) := by
-      rw [streamMeasure, Measure.map_map (by fun_prop) (by fun_prop)]
-      rfl
-    _ = ν a := by simp_rw [(measurePreserving_eval_infinitePi _ _).map_eq]
-
 /-- Law of `Y` conditioned on the event `s`.-/
 notation "𝓛[" Y " | " s "; " μ "]" => Measure.map Y (μ[|s])
 /-- Law of `Y` conditioned on the event that `X` is in `s`. -/
@@ -63,6 +50,8 @@ lemma condDistrib_reward'' [Countable 𝓐]
   filter_upwards [h_ra', h_prod] with ω h_eq h_prod
   rw [h_prod, h_eq]
 
+section CondIndep
+
 variable [StandardBorelSpace 𝓐]
 
 omit [DecidableEq 𝓐] in
@@ -179,7 +168,7 @@ lemma condDistrib_rewardByCount_stepsUntil [StandardBorelSpace Ω] [Countable 
       simp only [Set.mem_preimage, Set.mem_singleton_iff]
       exact fun ω ↦ rewardByCount_of_stepsUntil_eq_top
     rw [cond_of_indepFun _ (by fun_prop) (by fun_prop) (measurableSet_singleton _)]
-    · exact (hasLaw_Z a m).map_eq
+    · exact (hasLaw_snd_apply_prod_streamMeasure P ν m a).map_eq
     · rwa [Measure.map_apply (by fun_prop) (measurableSet_singleton _)] at hn
     · exact indepFun_prod (X := fun ω : Ω ↦ stepsUntil A a m ω)
         (Y := fun ω : ℕ → 𝓐 → ℝ ↦ ω m a) (by fun_prop) (by fun_prop)
@@ -240,4 +229,303 @@ lemma identDistrib_rewardByCount_eval [StandardBorelSpace Ω] [Countable 𝓐]
   (identDistrib_rewardByCount_id h a n hn).trans
     (identDistrib_eval_eval_id_streamMeasure ν m a).symm
 
+end CondIndep
+
+section Independence
+
+/-! ### Independence of the rewards by count
+
+We prove that the family `(rewardByCount A R a (m + 1))_{(a, m)}` is independent, with
+`rewardByCount A R a (m + 1)` distributed according to `ν a`.
+
+The proof goes through the time-truncated arrays `rewardByCountUntil A R t`, in which the entries
+corresponding to pulls that happen at time `t` or later are replaced by the auxiliary array. Under
+the product measure `𝔓`, the law of `rewardByCountUntil A R t` does not depend on `t`: going from
+`t` to `t + 1` replaces the entry `(A t, pullCount A (A t) t)`, which was an auxiliary reward with
+law `ν (A t)` independent of everything else, by the reward `R t`, which conditionally on the
+history and on `A t` also has law `ν (A t)`. For `t = 0` the array is a sub-array of the auxiliary
+array, whose law is the product measure. Finally `rewardByCountUntil A R t` converges to
+`rewardByCount` entrywise as `t → ∞`, which gives the law of the latter. -/
+
+
+/-- The law of the array `rewardByCountUntil A R 0`, which is a sub-array of the auxiliary array,
+is the product measure `⨂ (a, m), ν a`. -/
+lemma hasLaw_rewardByCountUntil_zero (μ : Measure Ω) [IsProbabilityMeasure μ] :
+    HasLaw (rewardByCountUntil A R 0) (Measure.infinitePi fun p : 𝓐 × ℕ ↦ ν p.1)
+      (μ.prod (streamMeasure ν)) :=
+  have h_indep : iIndepFun (fun (p : 𝓐 × ℕ) (ω : Ω × (ℕ → 𝓐 → ℝ)) ↦ ω.2 (p.2 + 1) p.1)
+      (μ.prod (streamMeasure ν)) :=
+    (iIndepFun_snd_apply_prod_streamMeasure μ ν).precomp (g := fun p : 𝓐 × ℕ ↦ (p.2 + 1, p.1))
+      fun p q hpq ↦ Prod.ext (Prod.mk.inj hpq).2 (by have := (Prod.mk.inj hpq).1; omega)
+  h_indep.hasLaw_infinitePi (fun p ↦ hasLaw_snd_apply_prod_streamMeasure μ ν _ _)
+    (by fun_prop : Measurable fun (ω : Ω × (ℕ → 𝓐 → ℝ)) (p : 𝓐 × ℕ) ↦
+      ω.2 (p.2 + 1) p.1).aemeasurable
+
+variable [MeasurableSingletonClass 𝓐]
+
+/-- The array `rewardByCountUntil A R t` with the entry `(b, k)` erased is independent of the
+entry `(k + 1, b)` of the auxiliary array. -/
+lemma indepFun_update_rewardByCountUntil_eval [Countable 𝓐] (hA : ∀ n, Measurable (A n))
+    (hR : ∀ n, Measurable (R n)) (μ : Measure Ω) [IsProbabilityMeasure μ] (t : ℕ) (b : 𝓐) (k : ℕ) :
+    (fun ω ↦ Function.update (rewardByCountUntil A R t ω) (b, k) 0)
+      ⟂ᵢ[μ.prod (streamMeasure ν)] (fun ω ↦ ω.2 (k + 1) b) := by
+  refine ((indepFun_snd_apply_prod_streamMeasure_update μ ν (k + 1) b 0).of_measurable_right
+    ?_).symm
+  have h_eq : (fun ω : Ω × (ℕ → 𝓐 → ℝ) ↦ Function.update (rewardByCountUntil A R t ω) (b, k) 0)
+      = (fun ω ↦ Function.update (rewardByCountUntil A R t ω) (b, k) 0)
+        ∘ (fun ω ↦ (ω.1, fun i c ↦ if i = k + 1 ∧ c = b then 0 else ω.2 i c)) := by
+    ext ⟨x, z⟩ p
+    simp only [Function.comp_apply]
+    by_cases hp : p = (b, k)
+    · rw [hp, Function.update_self, Function.update_self]
+    · rw [Function.update_of_ne hp, Function.update_of_ne hp]
+      refine rewardByCountUntil_congr t p ?_
+      split_ifs with hc
+      · exact absurd (Prod.ext hc.2 (by have := hc.1; omega)) hp
+      · rfl
+  rw [h_eq]
+  exact measurable_comp_comap _
+    (measurable_update_left.comp (measurable_rewardByCountUntil hA hR t))
+
+/-- Conditionally on the event that the action at time `n + 1` is `b` and that `b` was pulled `k`
+times before, the reward at time `n + 1` is independent of the history up to time `n` and of the
+action at time `n + 1`. -/
+lemma indepFun_history_reward_cond (h : IsAlgEnvSeq A R alg (stationaryEnv ν) P)
+    (n : ℕ) (b : 𝓐) (k : ℕ) :
+    (fun x ↦ (history A R n x, A (n + 1) x))
+      ⟂ᵢ[P[|{x | A (n + 1) x = b ∧ pullCount A b (n + 1) x = k}]] R (n + 1) := by
+  rw [setOf_action_eq_and_pullCount_eq_eq_preimage (R' := R)]
+  exact h.indepFun_history_action_feedback_cond_stationaryEnv n
+    (measurableSet_snd_eq_and_pullCount'_eq n b k) fun u hu ↦ hu.1
+
+lemma indepFun_action_zero_reward_zero_cond (h : IsAlgEnvSeq A R alg (stationaryEnv ν) P)
+    (b : 𝓐) (k : ℕ) :
+    A 0 ⟂ᵢ[P[|{x | A 0 x = b ∧ pullCount A b 0 x = k}]] R 0 := by
+  rcases eq_or_ne k 0 with rfl | hk
+  · have h_eq : {x | A 0 x = b ∧ pullCount A b 0 x = 0} = A 0 ⁻¹' {b} := by ext; simp
+    rw [h_eq]
+    exact indepFun_cond_preimage_singleton_left (h.measurable_action 0) b _
+  · have h_eq : {x | A 0 x = b ∧ pullCount A b 0 x = k} = ∅ := by ext; simp [hk.symm]
+    rw [h_eq]
+    simp
+
+/-- Conditionally on the event that the action at time `t` is `b` and that `b` was pulled `k`
+times before, the reward at time `t` has law `ν b`. -/
+lemma hasLaw_reward_cond (h : IsAlgEnvSeq A R alg (stationaryEnv ν) P) (t : ℕ) (b : 𝓐) (k : ℕ)
+    (hP : P {x | A t x = b ∧ pullCount A b t x = k} ≠ 0) :
+    HasLaw (R t) (ν b) (P[|{x | A t x = b ∧ pullCount A b t x = k}]) := by
+  cases t with
+  | zero =>
+    rcases eq_or_ne k 0 with rfl | hk
+    · have h_eq : {x | A 0 x = b ∧ pullCount A b 0 x = 0} = A 0 ⁻¹' {b} := by ext; simp
+      rw [h_eq] at hP ⊢
+      exact h.hasLaw_feedback_zero_cond_stationaryEnv hP
+    · refine absurd ?_ hP
+      have h_eq : {x | A 0 x = b ∧ pullCount A b 0 x = k} = ∅ := by ext; simp [hk.symm]
+      rw [h_eq, measure_empty]
+  | succ n =>
+    rw [setOf_action_eq_and_pullCount_eq_eq_preimage (R' := R)] at hP ⊢
+    exact h.hasLaw_feedback_cond_stationaryEnv n (measurableSet_snd_eq_and_pullCount'_eq n b k)
+      (fun u hu ↦ hu.1) hP
+
+lemma hasLaw_reward_cond_prod (h : IsAlgEnvSeq A R alg (stationaryEnv ν) P) (t : ℕ) (b : 𝓐)
+    (k : ℕ) (hP : P {x | A t x = b ∧ pullCount A b t x = k} ≠ 0) :
+    HasLaw (fun ω ↦ R t ω.1) (ν b)
+      ((P[|{x | A t x = b ∧ pullCount A b t x = k}]).prod (streamMeasure ν)) :=
+  (hasLaw_reward_cond h t b k hP).comp (hasLaw_fst_prod _ _)
+
+variable [Countable 𝓐]
+
+/-- Conditionally on the event that the action at time `t` is `b` and that `b` was pulled `k`
+times before, the array `rewardByCountUntil A R t` with the entry `(b, k)` erased is independent of
+the reward at time `t`. -/
+lemma indepFun_update_rewardByCountUntil_reward (h : IsAlgEnvSeq A R alg (stationaryEnv ν) P)
+    (t : ℕ) (b : 𝓐) (k : ℕ) :
+    (fun ω ↦ Function.update (rewardByCountUntil A R t ω) (b, k) 0)
+      ⟂ᵢ[(P[|{x | A t x = b ∧ pullCount A b t x = k}]).prod (streamMeasure ν)]
+      (fun ω ↦ R t ω.1) := by
+  have hA := h.measurable_action
+  have hR := h.measurable_feedback
+  by_cases hP : P {x | A t x = b ∧ pullCount A b t x = k} = 0
+  · rw [cond_eq_zero.2 (Or.inr hP), Measure.zero_prod]
+    exact indepFun_zero_measure _ _
+  have : IsProbabilityMeasure (P[|{x | A t x = b ∧ pullCount A b t x = k}]) :=
+    cond_isProbabilityMeasure hP
+  cases t with
+  | zero =>
+    have h_indep := (indepFun_action_zero_reward_zero_cond h b k).symm.fst_prod
+      (ν := streamMeasure ν) (hR 0) (hA 0)
+    refine (h_indep.of_measurable_right ?_).symm
+    refine Measurable.comp measurable_update_left ?_
+    refine measurable_rewardByCountUntil_of 0 (fun i hi ↦ absurd hi (Nat.not_lt_zero i))
+      (fun i hi ↦ absurd hi (Nat.not_lt_zero i)) ?_
+    exact measurable_comp_comap (fun ω : Ω × (ℕ → 𝓐 → ℝ) ↦ (A 0 ω.1, ω.2)) measurable_snd
+  | succ n =>
+    have h_indep := (indepFun_history_reward_cond h n b k).symm.fst_prod
+      (ν := streamMeasure ν) (hR _) (by fun_prop)
+    refine (h_indep.of_measurable_right ?_).symm
+    refine Measurable.comp measurable_update_left ?_
+    refine measurable_rewardByCountUntil_of (n + 1) (fun i hi ↦ ?_) (fun i hi ↦ ?_) ?_
+    · exact measurable_comp_comap
+        (fun ω : Ω × (ℕ → 𝓐 → ℝ) ↦ ((history A R n ω.1, A (n + 1) ω.1), ω.2))
+        (g := fun v : ((Iic n → 𝓐 × ℝ) × 𝓐) × (ℕ → 𝓐 → ℝ) ↦
+          (v.1.1 ⟨i, mem_Iic.2 (Nat.lt_succ_iff.1 hi)⟩).1) (by fun_prop)
+    · exact measurable_comp_comap
+        (fun ω : Ω × (ℕ → 𝓐 → ℝ) ↦ ((history A R n ω.1, A (n + 1) ω.1), ω.2))
+        (g := fun v : ((Iic n → 𝓐 × ℝ) × 𝓐) × (ℕ → 𝓐 → ℝ) ↦
+          (v.1.1 ⟨i, mem_Iic.2 (Nat.lt_succ_iff.1 hi)⟩).2) (by fun_prop)
+    · exact measurable_comp_comap
+        (fun ω : Ω × (ℕ → 𝓐 → ℝ) ↦ ((history A R n ω.1, A (n + 1) ω.1), ω.2))
+        (g := fun v : ((Iic n → 𝓐 × ℝ) × 𝓐) × (ℕ → 𝓐 → ℝ) ↦ v.2) measurable_snd
+
+/-- Conditionally on the event that the action at time `t` is `b` and that `b` was pulled `k`
+times before, the arrays `rewardByCountUntil A R (t + 1)` and `rewardByCountUntil A R t` have the
+same law: they differ only in the entry `(b, k)`, which is `R t` in the first and an auxiliary
+reward in the second, and both are independent of the rest of the array with law `ν b`. -/
+lemma identDistrib_rewardByCountUntil_add_one_cond (h : IsAlgEnvSeq A R alg (stationaryEnv ν) P)
+    (t : ℕ) (b : 𝓐) (k : ℕ) :
+    IdentDistrib (rewardByCountUntil A R (t + 1)) (rewardByCountUntil A R t)
+      ((P[|{x | A t x = b ∧ pullCount A b t x = k}]).prod (streamMeasure ν))
+      ((P[|{x | A t x = b ∧ pullCount A b t x = k}]).prod (streamMeasure ν)) := by
+  have hA := h.measurable_action
+  have hR := h.measurable_feedback
+  set E₁ := {x | A t x = b ∧ pullCount A b t x = k} with hE₁
+  have hE₁_meas : MeasurableSet E₁ := measurableSet_action_eq_and_pullCount_eq hA t b k
+  by_cases hP : P E₁ = 0
+  · rw [cond_eq_zero.2 (Or.inr hP), Measure.zero_prod]
+    exact ⟨(measurable_rewardByCountUntil hA hR _).aemeasurable,
+      (measurable_rewardByCountUntil hA hR _).aemeasurable, by simp⟩
+  have : IsProbabilityMeasure (P[|E₁]) := cond_isProbabilityMeasure hP
+  set μ' := (P[|E₁]).prod (streamMeasure ν) with hμ'
+  have h_ae : ∀ᵐ ω ∂μ', A t ω.1 = b ∧ pullCount A b t ω.1 = k := by
+    rw [ae_iff]
+    have h_set : {ω : Ω × (ℕ → 𝓐 → ℝ) | ¬ (A t ω.1 = b ∧ pullCount A b t ω.1 = k)}
+        = E₁ᶜ ×ˢ Set.univ := by
+      ext ω
+      simp [hE₁]
+    rw [h_set, hμ', Measure.prod_prod, cond_apply hE₁_meas, Set.inter_compl_self, measure_empty,
+      mul_zero, zero_mul]
+  set W : Ω × (ℕ → 𝓐 → ℝ) → 𝓐 × ℕ → ℝ :=
+    fun ω ↦ Function.update (rewardByCountUntil A R t ω) (b, k) 0 with hW
+  have hWm : Measurable W := measurable_update_left.comp (measurable_rewardByCountUntil hA hR t)
+  have h1 : rewardByCountUntil A R (t + 1)
+      =ᵐ[μ'] fun ω ↦ Function.update (W ω) (b, k) (R t ω.1) := by
+    filter_upwards [h_ae] with ω hω
+    obtain ⟨hb, hk⟩ := hω
+    simp only [hW, rewardByCountUntil_add_one, Function.update_idem, hb, hk]
+  have h2 : rewardByCountUntil A R t
+      =ᵐ[μ'] fun ω ↦ Function.update (W ω) (b, k) (ω.2 (k + 1) b) := by
+    filter_upwards [h_ae] with ω hω
+    obtain ⟨hb, hk⟩ := hω
+    simp only [hW, Function.update_idem]
+    rw [← rewardByCountUntil_apply_of_pullCount_le hk.le, Function.update_eq_self]
+  -- both `(W, R t)` and `(W, ω.2 (k + 1) b)` have law `(μ'.map W).prod (ν b)`
+  have hW : HasLaw W (μ'.map W) μ' := hWm.hasLaw_map μ'
+  have h1' : HasLaw (fun ω ↦ Function.update (W ω) (b, k) (R t ω.1))
+      (((μ'.map W).prod (ν b)).map
+        fun q : (𝓐 × ℕ → ℝ) × ℝ ↦ Function.update q.1 (b, k) q.2) μ' :=
+    ((measurable_update' (a := (b, k))).hasLaw_map _).comp
+      ((indepFun_update_rewardByCountUntil_reward h t b k).hasLaw_prod hW
+        (hasLaw_reward_cond_prod h t b k hP))
+  have h2' : HasLaw (fun ω : Ω × (ℕ → 𝓐 → ℝ) ↦ Function.update (W ω) (b, k) (ω.2 (k + 1) b))
+      (((μ'.map W).prod (ν b)).map
+        fun q : (𝓐 × ℕ → ℝ) × ℝ ↦ Function.update q.1 (b, k) q.2) μ' :=
+    ((measurable_update' (a := (b, k))).hasLaw_map _).comp
+      ((indepFun_update_rewardByCountUntil_eval hA hR _ t b k).hasLaw_prod hW
+        (hasLaw_snd_apply_prod_streamMeasure _ _ _ _))
+  exact ((IdentDistrib.of_ae_eq (measurable_rewardByCountUntil hA hR _).aemeasurable h1).trans
+    (h1'.identDistrib h2')).trans
+    (IdentDistrib.of_ae_eq (measurable_rewardByCountUntil hA hR _).aemeasurable h2).symm
+
+/-- The law of `rewardByCountUntil A R t` under `𝔓` does not depend on `t`. -/
+lemma identDistrib_rewardByCountUntil_add_one (h : IsAlgEnvSeq A R alg (stationaryEnv ν) P)
+    (t : ℕ) :
+    IdentDistrib (rewardByCountUntil A R (t + 1)) (rewardByCountUntil A R t) 𝔓 𝔓 := by
+  have hA := h.measurable_action
+  have hR := h.measurable_feedback
+  -- condition on the value of `(A t, pullCount A (A t) t)`
+  refine identDistrib_of_forall_identDistrib_cond
+    (g := fun ω : Ω × (ℕ → 𝓐 → ℝ) ↦ (A t ω.1, pullCount A (A t ω.1) t ω.1))
+    (((hA t).comp measurable_fst).prodMk
+      ((measurable_uncurry_pullCount_comp hA (hA t) measurable_const).comp measurable_fst))
+    (measurable_rewardByCountUntil hA hR _) (measurable_rewardByCountUntil hA hR _) fun p ↦ ?_
+  have h_eq : (fun ω : Ω × (ℕ → 𝓐 → ℝ) ↦ (A t ω.1, pullCount A (A t ω.1) t ω.1)) ⁻¹' {p}
+      = {x | A t x = p.1 ∧ pullCount A p.1 t x = p.2} ×ˢ Set.univ := by
+    ext ω
+    simp only [Set.mem_preimage, Set.mem_singleton_iff, Prod.ext_iff, Set.mem_prod,
+      Set.mem_ofPred_eq, Set.mem_univ, and_true]
+    constructor
+    · rintro ⟨h1, h2⟩
+      exact ⟨h1, by rw [← h1]; exact h2⟩
+    · rintro ⟨h1, h2⟩
+      exact ⟨h1, by rw [h1]; exact h2⟩
+  rw [h_eq, cond_prod_univ]
+  exact identDistrib_rewardByCountUntil_add_one_cond h t p.1 p.2
+
+/-- The law of `rewardByCountUntil A R t` under `𝔓` is `⨂ (a, m), ν a`, for all `t`. -/
+lemma hasLaw_rewardByCountUntil (h : IsAlgEnvSeq A R alg (stationaryEnv ν) P) (t : ℕ) :
+    HasLaw (rewardByCountUntil A R t) (Measure.infinitePi fun p : 𝓐 × ℕ ↦ ν p.1) 𝔓 := by
+  induction t with
+  | zero => exact hasLaw_rewardByCountUntil_zero P
+  | succ t ih => exact (identDistrib_rewardByCountUntil_add_one h t).symm.hasLaw ih
+
+/-- The array of rewards by count `(a, m) ↦ rewardByCount A R a (m + 1)` has law
+`⨂ (a, m), ν a`: its entries are independent, and the entry `(a, m)` has law `ν a`. -/
+lemma hasLaw_rewardByCount_infinitePi (h : IsAlgEnvSeq A R alg (stationaryEnv ν) P) :
+    HasLaw (fun ω (p : 𝓐 × ℕ) ↦ rewardByCount A R p.1 (p.2 + 1) ω)
+      (Measure.infinitePi fun p : 𝓐 × ℕ ↦ ν p.1) 𝔓 := by
+  have hY : Measurable fun ω (p : 𝓐 × ℕ) ↦ rewardByCount A R p.1 (p.2 + 1) ω :=
+    measurable_pi_lambda _ fun p ↦
+      measurable_rewardByCount h.measurable_action h.measurable_feedback p.1 (p.2 + 1)
+  -- `rewardByCountUntil A R t` has that law for all `t` and converges entrywise to the array
+  exact hasLaw_of_forall_eventually_eq (L := Filter.atTop)
+    (measurable_rewardByCountUntil h.measurable_action h.measurable_feedback) hY.aemeasurable
+    (hasLaw_rewardByCountUntil h) eventually_rewardByCountUntil_eq
+
+/-- The reward received at the `(m + 1)`-th pull of action `a` has law `ν a`. -/
+lemma hasLaw_rewardByCount_add_one (h : IsAlgEnvSeq A R alg (stationaryEnv ν) P)
+    (a : 𝓐) (m : ℕ) :
+    HasLaw (rewardByCount A R a (m + 1)) (ν a) 𝔓 :=
+  (hasLaw_eval_infinitePi (fun p : 𝓐 × ℕ ↦ ν p.1) (a, m)).comp (hasLaw_rewardByCount_infinitePi h)
+
+/-- The rewards by count `rewardByCount A R a (m + 1)` are independent over all actions `a` and
+all counts `m`. -/
+lemma iIndepFun_rewardByCount_add_one (h : IsAlgEnvSeq A R alg (stationaryEnv ν) P) :
+    iIndepFun (fun (p : 𝓐 × ℕ) ω ↦ rewardByCount A R p.1 (p.2 + 1) ω) 𝔓 :=
+  (iIndepFun_iff_hasLaw_Pi_infinitePi
+    (X := fun (p : 𝓐 × ℕ) ω ↦ rewardByCount A R p.1 (p.2 + 1) ω) (μ := fun p : 𝓐 × ℕ ↦ ν p.1)
+    (fun p ↦ hasLaw_rewardByCount_add_one h p.1 p.2)
+    (hasLaw_rewardByCount_infinitePi h).aemeasurable).2 (hasLaw_rewardByCount_infinitePi h)
+
+/-- The rewards by count `rewardByCount A R a m` for `m ≠ 0` are independent over all actions `a`
+and all counts `m`. -/
+lemma iIndepFun_rewardByCount (h : IsAlgEnvSeq A R alg (stationaryEnv ν) P) :
+    iIndepFun (fun (p : {p : 𝓐 × ℕ // p.2 ≠ 0}) ω ↦ rewardByCount A R p.1.1 p.1.2 ω) 𝔓 := by
+  have h_eq : (fun (p : {p : 𝓐 × ℕ // p.2 ≠ 0}) ω ↦ rewardByCount A R p.1.1 p.1.2 ω)
+      = fun p ω ↦ rewardByCount A R p.1.1 (p.1.2 - 1 + 1) ω := by
+    ext p ω
+    rw [Nat.sub_add_cancel (Nat.pos_of_ne_zero p.2)]
+  rw [h_eq]
+  exact (iIndepFun_rewardByCount_add_one h).precomp
+    (g := fun p : {p : 𝓐 × ℕ // p.2 ≠ 0} ↦ (p.1.1, p.1.2 - 1)) fun p q hpq ↦ by
+      simp only [Prod.mk.injEq] at hpq
+      exact Subtype.ext (Prod.ext hpq.1 (by have := p.2; have := q.2; omega))
+
+/-- For each action `a`, the rewards by count `(rewardByCount A R a (m + 1))_m` are independent
+(and by `hasLaw_rewardByCount_add_one` identically distributed with law `ν a`). -/
+lemma iIndepFun_rewardByCount_add_one_action (h : IsAlgEnvSeq A R alg (stationaryEnv ν) P)
+    (a : 𝓐) :
+    iIndepFun (fun m ω ↦ rewardByCount A R a (m + 1) ω) 𝔓 :=
+  (iIndepFun_rewardByCount_add_one h).precomp (g := fun m ↦ (a, m))
+    fun _ _ hmn ↦ (Prod.mk.inj hmn).2
+
+/-- Two distinct rewards by count are independent. -/
+lemma indepFun_rewardByCount (h : IsAlgEnvSeq A R alg (stationaryEnv ν) P)
+    {a b : 𝓐} {m n : ℕ} (hm : m ≠ 0) (hn : n ≠ 0) (hne : (a, m) ≠ (b, n)) :
+    rewardByCount A R a m ⟂ᵢ[𝔓] rewardByCount A R b n :=
+  (iIndepFun_rewardByCount h).indepFun (i := ⟨(a, m), hm⟩) (j := ⟨(b, n), hn⟩)
+    fun h_eq ↦ hne (congrArg Subtype.val h_eq)
+
+end Independence
+
 end Bandits
diff --git a/LeanMachineLearning/SequentialLearning/Algorithm.lean b/LeanMachineLearning/SequentialLearning/Algorithm.lean
index c3091abf..cf7cd69f 100644
--- a/LeanMachineLearning/SequentialLearning/Algorithm.lean
+++ b/LeanMachineLearning/SequentialLearning/Algorithm.lean
@@ -262,6 +262,13 @@ lemma IsAlgEnvSeq.hasLaw_feedback_zero_comp (h : IsAlgEnvSeq A Y alg env P) :
     HasLaw (Y 0) (env.ν0 ∘ₘ (P.map (A 0))) P :=
   HasCondDistrib.hasLaw_comp (h.hasCondDistrib_feedback_zero)
 
+/-- Conditionally on the event `A 0 = b`, the first feedback has law `env.ν0 b`. -/
+lemma IsAlgEnvSeq.hasLaw_feedback_zero_cond [MeasurableSingletonClass 𝓐]
+    (h : IsAlgEnvSeq A Y alg env P) {b : 𝓐} (hP : P (A 0 ⁻¹' {b}) ≠ 0) :
+    HasLaw (Y 0) (env.ν0 b) P[|A 0 ⁻¹' {b}] :=
+  h.hasCondDistrib_feedback_zero.hasLaw_cond (h.measurable_feedback 0)
+    (measurableSet_singleton b) (fun a ha ↦ by rw [Set.mem_singleton_iff.1 ha]) hP
+
 section Filtration
 
 namespace IsAlgEnvSeq
diff --git a/LeanMachineLearning/SequentialLearning/FiniteActions.lean b/LeanMachineLearning/SequentialLearning/FiniteActions.lean
index e5c8644b..16c34693 100644
--- a/LeanMachineLearning/SequentialLearning/FiniteActions.lean
+++ b/LeanMachineLearning/SequentialLearning/FiniteActions.lean
@@ -246,6 +246,28 @@ lemma isStronglyPredictable_pullCount [MeasurableSingletonClass 𝓐]
   simp only [pullCount_zero]
   fun_prop
 
+lemma measurableSet_action_eq_and_pullCount_eq [MeasurableSingletonClass 𝓐]
+    (hA : ∀ n, Measurable (A n)) (t : ℕ) (b : 𝓐) (k : ℕ) :
+    MeasurableSet {x | A t x = b ∧ pullCount A b t x = k} :=
+  ((measurableSet_singleton _).preimage (hA t)).inter
+    ((measurableSet_singleton _).preimage (measurable_pullCount hA b t))
+
+lemma measurableSet_snd_eq_and_pullCount'_eq [MeasurableSingletonClass 𝓐]
+    (n : ℕ) (b : 𝓐) (k : ℕ) :
+    MeasurableSet {u : (Iic n → 𝓐 × R) × 𝓐 | u.2 = b ∧ pullCount' n u.1 b = k} :=
+  ((measurableSet_singleton _).preimage measurable_snd).inter
+    ((measurableSet_singleton _).preimage ((measurable_pullCount' n b).comp measurable_fst))
+
+/-- The event that the action at time `n + 1` is `b` and that `b` was pulled `k` times before is
+a preimage by `(history A R' n, A (n + 1))`. -/
+lemma setOf_action_eq_and_pullCount_eq_eq_preimage (n : ℕ) (b : 𝓐) (k : ℕ) :
+    {x | A (n + 1) x = b ∧ pullCount A b (n + 1) x = k}
+      = (fun x ↦ (history A R' n x, A (n + 1) x)) ⁻¹' {u | u.2 = b ∧ pullCount' n u.1 b = k} := by
+  ext x
+  simp only [Set.mem_ofPred_eq, Set.mem_preimage]
+  rw [pullCount_add_one_eq_pullCount' (R' := R')]
+  rfl
+
 lemma integrable_pullCount [MeasurableSingletonClass 𝓐]
     (hA : ∀ n, Measurable (A n)) (a : 𝓐) (n : ℕ) :
     Integrable (fun ω ↦ (pullCount A a n ω : ℝ)) P := by
@@ -736,6 +758,130 @@ lemma measurable_rewardByCount [MeasurableSingletonClass 𝓐]
     simp only
     fun_prop
 
+/-- Array of rewards by count, truncated at time `t`: the entry `(a, m)` is the reward received at
+the `(m + 1)`-th pull of action `a` if that pull happened before time `t`, and the entry
+`(m + 1, a)` of the auxiliary array `ω.2` otherwise.
+
+This is an auxiliary definition used to prove results about the distribution of `rewardByCount`.
+
+It is defined recursively: at time `t`, the
+entry `(A t, pullCount A (A t) t)` is replaced by the reward `R' t`.
+See `rewardByCountUntil_apply_of_lt_pullCount` and `rewardByCountUntil_apply_of_pullCount_le`. -/
+noncomputable
+def rewardByCountUntil (A : ℕ → Ω → 𝓐) (R' : ℕ → Ω → R) : ℕ → Ω × (ℕ → 𝓐 → R) → 𝓐 × ℕ → R
+  | 0, ω => fun p ↦ ω.2 (p.2 + 1) p.1
+  | t + 1, ω => Function.update (rewardByCountUntil A R' t ω)
+      (A t ω.1, pullCount A (A t ω.1) t ω.1) (R' t ω.1)
+
+@[simp]
+lemma rewardByCountUntil_zero (ω : Ω × (ℕ → 𝓐 → R)) :
+    rewardByCountUntil A R' 0 ω = fun p ↦ ω.2 (p.2 + 1) p.1 := rfl
+
+lemma rewardByCountUntil_add_one (t : ℕ) (ω : Ω × (ℕ → 𝓐 → R)) :
+    rewardByCountUntil A R' (t + 1) ω = Function.update (rewardByCountUntil A R' t ω)
+      (A t ω.1, pullCount A (A t ω.1) t ω.1) (R' t ω.1) := rfl
+
+/-- If action `a` was pulled at most `m` times before time `t`, then the entry `(a, m)` of
+`rewardByCountUntil A R' t ω` is the entry `(m + 1, a)` of the auxiliary array. -/
+lemma rewardByCountUntil_apply_of_pullCount_le (h : pullCount A a t ω.1 ≤ m) :
+    rewardByCountUntil A R' t ω (a, m) = ω.2 (m + 1) a := by
+  induction t with
+  | zero => rfl
+  | succ t ih =>
+    rw [rewardByCountUntil_add_one, Function.update_of_ne,
+      ih ((pullCount_mono a (Nat.le_succ t) _).trans h)]
+    intro hp
+    obtain ⟨rfl, hm⟩ := Prod.mk.inj hp
+    rw [pullCount_action_eq_pullCount_add_one] at h
+    omega
+
+/-- If action `a` was pulled more than `m` times before time `t`, then the entry `(a, m)` of
+`rewardByCountUntil A R' t ω` is the reward received at the `(m + 1)`-th pull of `a`. -/
+lemma rewardByCountUntil_apply_of_lt_pullCount (h : m < pullCount A a t ω.1) :
+    rewardByCountUntil A R' t ω (a, m) = rewardByCount A R' a (m + 1) ω := by
+  induction t with
+  | zero => simp at h
+  | succ t ih =>
+    rw [rewardByCountUntil_add_one]
+    by_cases hp : (a, m) = (A t ω.1, pullCount A (A t ω.1) t ω.1)
+    · obtain ⟨rfl, rfl⟩ := Prod.mk.inj hp
+      rw [Function.update_self, rewardByCount_pullCount_add_one_eq_reward]
+    · rw [Function.update_of_ne hp]
+      refine ih ?_
+      rw [pullCount_add_one] at h
+      split_ifs at h with hA
+      · rw [hA] at hp
+        have hm : m ≠ pullCount A a t ω.1 := fun hm ↦ hp (by rw [hm])
+        omega
+      · simpa using h
+
+/-- `rewardByCountUntil A R' t (x, z) p` depends on `z` only through `z (p.2 + 1) p.1`. -/
+lemma rewardByCountUntil_congr {x : Ω} {z z' : ℕ → 𝓐 → R} (t : ℕ) (p : 𝓐 × ℕ)
+    (hz : z (p.2 + 1) p.1 = z' (p.2 + 1) p.1) :
+    rewardByCountUntil A R' t (x, z) p = rewardByCountUntil A R' t (x, z') p := by
+  induction t with
+  | zero => exact hz
+  | succ t ih =>
+    simp only [rewardByCountUntil_add_one, Function.update_apply]
+    split_ifs
+    · rfl
+    · exact ih
+
+/-- For each entry, `rewardByCountUntil A R' t ω` coincides with `rewardByCount` for `t` large
+enough. -/
+lemma eventually_rewardByCountUntil_eq (ω : Ω × (ℕ → 𝓐 → R)) (p : 𝓐 × ℕ) :
+    ∀ᶠ t in Filter.atTop,
+      rewardByCountUntil A R' t ω p = rewardByCount A R' p.1 (p.2 + 1) ω := by
+  obtain ⟨a, m⟩ := p
+  by_cases h : ∃ t, m < pullCount A a t ω.1
+  · obtain ⟨t, ht⟩ := h
+    filter_upwards [Filter.eventually_ge_atTop t] with s hs
+    exact rewardByCountUntil_apply_of_lt_pullCount (ht.trans_le (pullCount_mono a hs ω.1))
+  · push Not at h
+    refine Filter.Eventually.of_forall fun t ↦ ?_
+    rw [rewardByCountUntil_apply_of_pullCount_le (h t), rewardByCount_of_stepsUntil_eq_top]
+    rw [stepsUntil_eq_top_iff]
+    exact fun s ↦ ((h (s + 1)).trans_lt (Nat.lt_succ_self m)).ne
+
+/-- Measurability of `rewardByCountUntil A R' t` with respect to a σ-algebra `m` for which the
+actions and rewards before time `t` and the auxiliary array are measurable. -/
+lemma measurable_rewardByCountUntil_of {m : MeasurableSpace (Ω × (ℕ → 𝓐 → R))}
+    [MeasurableEq 𝓐] (t : ℕ)
+    (hA : ∀ i < t, Measurable[m] (fun ω : Ω × (ℕ → 𝓐 → R) ↦ A i ω.1))
+    (hR : ∀ i < t, Measurable[m] (fun ω : Ω × (ℕ → 𝓐 → R) ↦ R' i ω.1))
+    (hZ : Measurable[m] (Prod.snd : Ω × (ℕ → 𝓐 → R) → ℕ → 𝓐 → R)) :
+    Measurable[m] (rewardByCountUntil A R' t) := by
+  induction t with
+  | zero =>
+    have : rewardByCountUntil A R' 0 = (fun z (p : 𝓐 × ℕ) ↦ z (p.2 + 1) p.1) ∘ Prod.snd := rfl
+    rw [this]
+    exact Measurable.comp (by fun_prop) hZ
+  | succ t ih =>
+    have ht : t < t + 1 := Nat.lt_succ_self t
+    have ih := ih (fun i hi ↦ hA i (hi.trans ht)) (fun i hi ↦ hR i (hi.trans ht))
+    have hg : Measurable[m] (fun ω : Ω × (ℕ → 𝓐 → R) ↦
+        (A t ω.1, pullCount A (A t ω.1) t ω.1)) := by
+      refine Measurable.prodMk (hA t ht) ?_
+      simp_rw [pullCount_eq_sum]
+      refine Finset.measurable_sum _ fun s hs ↦ Measurable.ite ?_ measurable_const measurable_const
+      exact measurableSet_eq_fun (hA s ((Finset.mem_range.1 hs).trans ht)) (hA t ht)
+    refine measurable_pi_iff.2 fun p ↦ ?_
+    simp_rw [rewardByCountUntil_add_one, Function.update_apply]
+    refine Measurable.ite ?_ (hR t ht) ((measurable_pi_apply p).comp ih)
+    have h_set : {ω : Ω × (ℕ → 𝓐 → R) | p = (A t ω.1, pullCount A (A t ω.1) t ω.1)}
+        = (fun ω : Ω × (ℕ → 𝓐 → R) ↦ (A t ω.1, pullCount A (A t ω.1) t ω.1)) ⁻¹' {p} := by
+      ext ω
+      simp [eq_comm]
+    rw [h_set]
+    exact hg (measurableSet_singleton p)
+
+@[fun_prop]
+lemma measurable_rewardByCountUntil [MeasurableEq 𝓐]
+    (hA : ∀ n, Measurable (A n)) (hR' : ∀ n, Measurable (R' n)) (t : ℕ) :
+    Measurable (rewardByCountUntil A R' t) :=
+  measurable_rewardByCountUntil_of t (fun i _ ↦ (hA i).comp measurable_fst)
+    (fun i _ ↦ (hR' i).comp measurable_fst) measurable_snd
+
 end RewardByCount
 
 lemma sum_pullCount_mul [Fintype 𝓐] [Semiring R] (ω : Ω) (f : 𝓐 → R) (t : ℕ) :
diff --git a/LeanMachineLearning/SequentialLearning/StationaryEnv.lean b/LeanMachineLearning/SequentialLearning/StationaryEnv.lean
index 1d47b818..57eb7e25 100644
--- a/LeanMachineLearning/SequentialLearning/StationaryEnv.lean
+++ b/LeanMachineLearning/SequentialLearning/StationaryEnv.lean
@@ -104,6 +104,32 @@ lemma hasCondDistrib_feedback [IsObliviousEnv env] (h : IsAlgEnvSeq A Y alg env
       Measure.snd_map_prodMk (by fun_prop), Measure.map_map (by fun_prop) (by fun_prop)]
     congr
 
+/-- Conditionally on an event determined by the history up to time `n` and the action at time
+`n + 1`, on which that action is equal to `b`, the feedback at time `n + 1` has law
+`feedbackCondAction env (n + 1) b`. -/
+lemma hasLaw_feedback_cond [IsObliviousEnv env] (h : IsAlgEnvSeq A Y alg env P) (n : ℕ)
+    {s : Set ((Iic n → 𝓐 × 𝓨) × 𝓐)} (hs : MeasurableSet s) {b : 𝓐} (hsb : ∀ u ∈ s, u.2 = b)
+    (hP : P ((fun ω ↦ (history A Y n ω, A (n + 1) ω)) ⁻¹' s) ≠ 0) :
+    HasLaw (Y (n + 1)) (feedbackCondAction env (n + 1) b)
+      P[|(fun ω ↦ (history A Y n ω, A (n + 1) ω)) ⁻¹' s] := by
+  refine (hasCondDistrib_feedback_history_action h n).hasLaw_cond (h.measurable_feedback _) hs
+    (fun u hu ↦ ?_) hP
+  rw [Kernel.prodMkLeft_apply, hsb u hu]
+
+/-- Conditionally on an event determined by the history up to time `n` and the action at time
+`n + 1`, on which that action is constant, the feedback at time `n + 1` is independent of the
+history up to time `n` and of the action at time `n + 1`. -/
+lemma indepFun_history_action_feedback_cond [IsObliviousEnv env]
+    (h : IsAlgEnvSeq A Y alg env P) (n : ℕ)
+    {s : Set ((Iic n → 𝓐 × 𝓨) × 𝓐)} (hs : MeasurableSet s) {b : 𝓐} (hsb : ∀ u ∈ s, u.2 = b) :
+    (fun ω ↦ (history A Y n ω, A (n + 1) ω))
+      ⟂ᵢ[P[|(fun ω ↦ (history A Y n ω, A (n + 1) ω)) ⁻¹' s]] Y (n + 1) := by
+  have hA := h.measurable_action
+  have hY := h.measurable_feedback
+  refine (hasCondDistrib_feedback_history_action h n).indepFun_cond (by fun_prop) hs
+    (η := feedbackCondAction env (n + 1) b) fun u hu ↦ ?_
+  rw [Kernel.prodMkLeft_apply, hsb u hu]
+
 variable [StandardBorelSpace 𝓐] [Nonempty 𝓐] [StandardBorelSpace 𝓨] [Nonempty 𝓨]
 
 /-- The feedback at time `n + 1` is conditionally independent of the history up to time `n`
@@ -220,6 +246,30 @@ lemma condDistrib_feedback_stationaryEnv [StandardBorelSpace 𝓨] [Nonempty 
     condDistrib (Y n) (A n) P =ᵐ[P.map (A n)] ν :=
   (hasCondDistrib_feedback_stationaryEnv h n).condDistrib_eq
 
+/-- Conditionally on the event `A 0 = b`, the first feedback has law `ν b`. -/
+lemma hasLaw_feedback_zero_cond_stationaryEnv [MeasurableSingletonClass 𝓐]
+    (h : IsAlgEnvSeq A Y alg (stationaryEnv ν) P) {b : 𝓐} (hP : P (A 0 ⁻¹' {b}) ≠ 0) :
+    HasLaw (Y 0) (ν b) P[|A 0 ⁻¹' {b}] := by
+  simpa using h.hasLaw_feedback_zero_cond hP
+
+/-- Conditionally on an event determined by the history up to time `n` and the action at time
+`n + 1`, on which that action is equal to `b`, the feedback at time `n + 1` has law `ν b`. -/
+lemma hasLaw_feedback_cond_stationaryEnv (h : IsAlgEnvSeq A Y alg (stationaryEnv ν) P) (n : ℕ)
+    {s : Set ((Iic n → 𝓐 × 𝓨) × 𝓐)} (hs : MeasurableSet s) {b : 𝓐} (hsb : ∀ u ∈ s, u.2 = b)
+    (hP : P ((fun ω ↦ (history A Y n ω, A (n + 1) ω)) ⁻¹' s) ≠ 0) :
+    HasLaw (Y (n + 1)) (ν b) P[|(fun ω ↦ (history A Y n ω, A (n + 1) ω)) ⁻¹' s] := by
+  simpa using IsObliviousEnv.hasLaw_feedback_cond h n hs hsb hP
+
+/-- Conditionally on an event determined by the history up to time `n` and the action at time
+`n + 1`, on which that action is constant, the feedback at time `n + 1` is independent of the
+history up to time `n` and of the action at time `n + 1`. -/
+lemma indepFun_history_action_feedback_cond_stationaryEnv
+    (h : IsAlgEnvSeq A Y alg (stationaryEnv ν) P) (n : ℕ)
+    {s : Set ((Iic n → 𝓐 × 𝓨) × 𝓐)} (hs : MeasurableSet s) {b : 𝓐} (hsb : ∀ u ∈ s, u.2 = b) :
+    (fun ω ↦ (history A Y n ω, A (n + 1) ω))
+      ⟂ᵢ[P[|(fun ω ↦ (history A Y n ω, A (n + 1) ω)) ⁻¹' s]] Y (n + 1) :=
+  IsObliviousEnv.indepFun_history_action_feedback_cond h n hs hsb
+
 /-- The feedback at time `n + 1` is conditionally independent of the history up to time `n`
 given the action at time `n + 1`. -/
 lemma condIndepFun_feedback_history_action [StandardBorelSpace Ω]