diff --git a/LeanMachineLearning/ForMathlib/Probability/HasCondDistrib.lean b/LeanMachineLearning/ForMathlib/Probability/HasCondDistrib.lean
index 9d6da91..49b85f5 100644
--- a/LeanMachineLearning/ForMathlib/Probability/HasCondDistrib.lean
+++ b/LeanMachineLearning/ForMathlib/Probability/HasCondDistrib.lean
@@ -151,14 +151,14 @@ variable [StandardBorelSpace Ω] [Nonempty Ω] [StandardBorelSpace Ω'] [Nonempt
 lemma HasCondDistrib.condDistrib_eq [IsFiniteMeasure μ] [IsFiniteKernel κ]
     (h : HasCondDistrib Y X κ μ) :
     condDistrib Y X μ =ᵐ[μ.map X] κ := by
-  rw [condDistrib_ae_eq_iff_measure_eq_compProd _ (by fun_prop), h.map_eq]
+  rw [condDistrib_ae_eq_iff_measure_eq_compProd (by fun_prop) (by fun_prop), h.map_eq]
 
 lemma hasCondDistrib_of_condDistrib_eq [IsFiniteMeasure μ] [IsFiniteKernel κ]
     (hX : AEMeasurable X μ) (hY : AEMeasurable Y μ)
     (h : condDistrib Y X μ =ᵐ[μ.map X] κ) :
     HasCondDistrib Y X κ μ where
   aemeasurable := by fun_prop
-  map_eq := by rw [← compProd_map_condDistrib hY, Measure.compProd_congr h]
+  map_eq := by rw [← compProd_map_condDistrib hX hY, Measure.compProd_congr h]
 
 lemma HasCondDistrib.hasCondDistrib_sectR [IsFiniteMeasure μ] [StandardBorelSpace β] [Nonempty β]
     {W : α → Ω'} {Z : α → γ} {f : Ω' → β} {g : Ω' → Ω}
@@ -172,8 +172,9 @@ lemma HasCondDistrib.hasCondDistrib_sectR [IsFiniteMeasure μ] [StandardBorelSpa
     exact hasCondDistrib_of_condDistrib_eq (by fun_prop) (by fun_prop) hz
   have h_eq : condDistrib (g ∘ W) (fun a ↦ (Z a, (f ∘ W) a)) μ
       =ᵐ[μ.map Z ⊗ₘ (condDistrib W Z μ).map f] η := by
-    rw [← Measure.compProd_congr (condDistrib_comp Z hW hf),
-        compProd_map_condDistrib (hf.comp_aemeasurable hW)]
+    have hZ : AEMeasurable Z μ := hcd.aemeasurable_fst.fst
+    rw [← Measure.compProd_congr (condDistrib_comp hZ hW hf),
+        compProd_map_condDistrib hZ (hf.comp_aemeasurable hW)]
     exact hcd.condDistrib_eq
   filter_upwards [
     condDistrib_condDistrib_ae_eq_sectR_condDistrib hf hg hW hcd.aemeasurable_fst.fst,
diff --git a/LeanMachineLearning/ForMathlib/Probability/Independence/CondDistrib.lean b/LeanMachineLearning/ForMathlib/Probability/Independence/CondDistrib.lean
index 9925b4e..1a0ad7b 100644
--- a/LeanMachineLearning/ForMathlib/Probability/Independence/CondDistrib.lean
+++ b/LeanMachineLearning/ForMathlib/Probability/Independence/CondDistrib.lean
@@ -9,6 +9,7 @@ public import LeanMachineLearning.ForMathlib.Probability.Kernel.KernelSub
 public import Mathlib.MeasureTheory.Measure.ProbabilityMeasure
 public import Mathlib.Probability.Independence.Basic
 public import Mathlib.Probability.Independence.Conditional
+public import Mathlib.Probability.Kernel.CompProdEqIff
 
 /-!
 # Lemmas about conditional distributions
@@ -33,22 +34,18 @@ section CondDistrib
 
 variable [IsFiniteMeasure μ]
 
-lemma map_swap_compProd_map_condDistrib (hY : AEMeasurable Y μ) :
+lemma map_swap_compProd_map_condDistrib (hX : AEMeasurable X μ) (hY : AEMeasurable Y μ) :
     (μ.map X ⊗ₘ condDistrib Y X μ).map Prod.swap = μ.map (fun a ↦ (Y a, X a)) := by
-  by_cases hX : AEMeasurable X μ
-  · rw [compProd_map_condDistrib hY,
-      AEMeasurable.map_map_of_aemeasurable measurable_swap.aemeasurable (hX.prodMk hY)]
-    rfl
-  · have hYX : ¬ AEMeasurable (fun a ↦ (Y a, X a)) μ :=
-      fun h ↦ hX (measurable_snd.comp_aemeasurable h)
-    simp [hX, hYX]
+  rw [compProd_map_condDistrib hX hY,
+    AEMeasurable.map_map_of_aemeasurable measurable_swap.aemeasurable (hX.prodMk hY)]
+  rfl
 
 lemma condDistrib_prod_left [StandardBorelSpace β] [Nonempty β]
     (hX : AEMeasurable X μ) (hY : AEMeasurable Y μ) (hT : AEMeasurable T μ) :
     condDistrib (fun ω ↦ (X ω, Y ω)) T μ
       =ᵐ[μ.map T] condDistrib X T μ ⊗ₖ condDistrib Y (fun ω ↦ (T ω, X ω)) μ := by
-  refine condDistrib_ae_eq_of_measure_eq_compProd (μ := μ) T (by fun_prop) ?_
-  rw [← Measure.compProd_assoc', compProd_map_condDistrib hX, compProd_map_condDistrib hY,
+  refine condDistrib_ae_eq_of_measure_eq_compProd (μ := μ) hT (by fun_prop) ?_
+  rw [← Measure.compProd_assoc', compProd_map_condDistrib hT hX, compProd_map_condDistrib (hT.prodMk hX) hY,
     AEMeasurable.map_map_of_aemeasurable (by fun_prop) (by fun_prop)]
   rfl
 
@@ -60,8 +57,8 @@ lemma condDistrib_condDistrib_ae_eq_sectR_condDistrib [StandardBorelSpace β] [N
         (condDistrib (g ∘ Z) (fun a ↦ (T a, (f ∘ Z) a)) μ).sectR t := by
   filter_upwards [
     condDistrib_prod_left (hf.comp_aemeasurable hZ) (hg.comp_aemeasurable hZ) hT,
-    condDistrib_comp T hZ (hf.prodMk hg), condDistrib_comp T hZ hf] with t h_prod h_pair h_fst
-  rw [condDistrib_ae_eq_iff_measure_eq_compProd f hg.aemeasurable]
+    condDistrib_comp hT hZ (hf.prodMk hg), condDistrib_comp hT hZ hf] with t h_prod h_pair h_fst
+  rw [condDistrib_ae_eq_iff_measure_eq_compProd hf.aemeasurable hg.aemeasurable]
   calc (condDistrib Z T μ t).map (fun ω' ↦ (f ω', g ω'))
   _ = condDistrib (fun a ↦ ((f ∘ Z) a, (g ∘ Z) a)) T μ t := by
       rw [← Kernel.map_apply _ (hf.prodMk hg)]
@@ -75,8 +72,8 @@ lemma condDistrib_prod_self_left [StandardBorelSpace β] [Nonempty β] [Standard
     (hX : AEMeasurable X μ) (hT : AEMeasurable T μ) :
     condDistrib (fun ω ↦ (X ω, T ω)) T μ =ᵐ[μ.map T] condDistrib X T μ ×ₖ Kernel.id := by
   have h_prod := condDistrib_prod_left hX hT hT (μ := μ)
-  have h_fst := condDistrib_comp_self (μ := μ) (fun ω ↦ (T ω, X ω)) (f := Prod.fst) (by fun_prop)
-  rw [(compProd_map_condDistrib hX).symm] at h_fst
+  have h_fst := condDistrib_comp_self (μ := μ) (hT.prodMk hX) (f := Prod.fst) (by fun_prop)
+  rw [(compProd_map_condDistrib hT hX).symm] at h_fst
   have h_fst' := (Measure.ae_compProd_iff (Kernel.measurableSet_eq _ _)).mp h_fst
   filter_upwards [h_prod, h_fst'] with z hz1 hz2
   rw [hz1]
@@ -104,9 +101,9 @@ lemma CondIndepFun.prod_right [StandardBorelSpace α] [StandardBorelSpace β] [N
     (h : X ⟂ᵢ[Z, hZ; μ] Y) :
     X ⟂ᵢ[Z, hZ; μ] (fun ω ↦ (Y ω, Z ω)) := by
   rw [condIndepFun_iff_condDistrib_prod_ae_eq_prodMkRight hY hX hZ,
-    condDistrib_ae_eq_iff_measure_eq_compProd _ (by fun_prop)] at h
+    condDistrib_ae_eq_iff_measure_eq_compProd (by fun_prop) (by fun_prop)] at h
   rw [condIndepFun_iff_condDistrib_prod_ae_eq_prodMkRight (by fun_prop) hX hZ,
-    condDistrib_ae_eq_iff_measure_eq_compProd _ (by fun_prop)]
+    condDistrib_ae_eq_iff_measure_eq_compProd (by fun_prop) (by fun_prop)]
   -- Key: condDistrib (Y, Z) Z μ z = (condDistrib Y Z μ z).map (y ↦ (y, z))
   have h_cond : condDistrib (fun ω ↦ (Y ω, Z ω)) Z μ =ᵐ[μ.map Z]
       fun z ↦ (condDistrib Y Z μ z).map (fun y ↦ (y, z)) := by
@@ -149,14 +146,14 @@ lemma fst_condDistrib_prod [StandardBorelSpace β] [Nonempty β]
 
 lemma condDistrib_of_indepFun (h : IndepFun X Y μ) (hX : AEMeasurable X μ) (hY : AEMeasurable Y μ) :
     condDistrib Y X μ =ᵐ[μ.map X] Kernel.const β (μ.map Y) := by
-  refine condDistrib_ae_eq_of_measure_eq_compProd (μ := μ) X hY ?_
+  refine condDistrib_ae_eq_of_measure_eq_compProd (μ := μ) hX hY ?_
   simp only [Measure.compProd_const]
   exact (indepFun_iff_map_prod_eq_prod_map_map hX hY).mp h
 
 lemma indepFun_iff_condDistrib_eq_const (hX : AEMeasurable X μ) (hY : AEMeasurable Y μ) :
     IndepFun X Y μ ↔ condDistrib Y X μ =ᵐ[μ.map X] Kernel.const β (μ.map Y) := by
   refine ⟨fun h ↦ condDistrib_of_indepFun h hX hY, fun h ↦ ?_⟩
-  rw [indepFun_iff_map_prod_eq_prod_map_map hX hY, ← compProd_map_condDistrib hY,
+  rw [indepFun_iff_map_prod_eq_prod_map_map hX hY, ← compProd_map_condDistrib hX hY,
     Measure.compProd_congr h]
   simp
 
@@ -275,9 +272,9 @@ lemma condIndepFun_of_exists_condDistrib_prod_ae_eq_prodMkRight
     (h : condDistrib Y (fun ω ↦ (Z ω, X ω)) μ =ᵐ[μ.map (fun ω ↦ (Z ω, X ω))] η.prodMkRight _) :
     Y ⟂ᵢ[Z, hZ; μ] X := by
   have hη_eq : condDistrib Y Z μ =ᵐ[μ.map Z] η := by
-    rw [condDistrib_ae_eq_iff_measure_eq_compProd _ (by fun_prop)] at h ⊢
+    rw [condDistrib_ae_eq_iff_measure_eq_compProd (by fun_prop) (by fun_prop)] at h ⊢
     have h_fst : μ.map Z = (μ.map (fun ω ↦ (Z ω, X ω))).fst := by
-      rw [Measure.fst_map_prodMk hX]
+      rw [Measure.fst_map_prodMk hZ hX]
     rw [h_fst, ← Measure.map_swap_comprod_eq_fst_compProd, ← h,
       Measure.map_map (by fun_prop) (by fun_prop), Measure.map_map (by fun_prop) (by fun_prop),
       Measure.fst,
@@ -286,7 +283,7 @@ lemma condIndepFun_of_exists_condDistrib_prod_ae_eq_prodMkRight
   symm
   rw [condIndepFun_iff_condDistrib_prod_ae_eq_prodMkRight hY hX hZ]
   refine h.trans ?_
-  rw [Kernel.prodMkRight_ae_eq_iff, Measure.fst_map_prodMk (by fun_prop)]
+  rw [Kernel.prodMkRight_ae_eq_iff, Measure.fst_map_prodMk (by fun_prop) (by fun_prop)]
   exact hη_eq.symm
 
 omit [StandardBorelSpace Ω'] [Nonempty Ω'] in
@@ -297,7 +294,7 @@ lemma condIndepFun_of_exists_condDistrib_prod_ae_eq_prodMkLeft
     (h : condDistrib Y (fun ω ↦ (X ω, Z ω)) μ =ᵐ[μ.map (fun ω ↦ (X ω, Z ω))] η.prodMkLeft _) :
     Y ⟂ᵢ[Z, hZ; μ] X := by
   refine condIndepFun_of_exists_condDistrib_prod_ae_eq_prodMkRight hX hY hZ ?_ (η := η)
-  rw [← Kernel.compProd_eq_iff, compProd_map_condDistrib (by fun_prop)] at h ⊢
+  rw [← Kernel.compProd_eq_iff, compProd_map_condDistrib (by fun_prop) (by fun_prop)] at h ⊢
   have : μ.map (fun a ↦ ((Z a, X a), Y a))
       = (μ.map (fun a ↦ ((X a, Z a), Y a))).map (fun p ↦ ((p.1.2, p.1.1), p.2)) := by
     rw [Measure.map_map (by fun_prop) (by fun_prop)]
@@ -333,9 +330,9 @@ lemma condIndepFun_fst_prod [StandardBorelSpace α] [StandardBorelSpace β] [Non
   rw [condIndepFun_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)
+    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)
+    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)]
@@ -417,8 +414,8 @@ lemma ae_eq_of_map_prodMk_eq {β Ω : Type*} {_ : MeasurableSpace β} {_ : Measu
 lemma ae_eq_of_condDistrib_eq_deterministic {f : β → Ω} (hf : Measurable f) (hX : AEMeasurable X μ)
     (hY : AEMeasurable Y μ) (h : condDistrib Y X μ =ᵐ[μ.map X] Kernel.deterministic f hf) :
     Y =ᵐ[μ] f ∘ X := by
-  have hfX := condDistrib_comp_self (μ := μ) X hf
-  rw [condDistrib_ae_eq_iff_measure_eq_compProd _ (by fun_prop)] at h hfX
+  have hfX := condDistrib_comp_self (μ := μ) hX hf
+  rw [condDistrib_ae_eq_iff_measure_eq_compProd (by fun_prop) (by fun_prop)] at h hfX
   exact ae_eq_of_map_prodMk_eq hf hX hY (hfX ▸ h)
 
 end CondDistrib
@@ -432,7 +429,7 @@ lemma condDistrib_ae_eq_cond [Countable β] [MeasurableSingletonClass β]
   rw [Filter.EventuallyEq, ae_iff_of_countable]
   intro b hb
   ext s hs
-  rw [condDistrib_apply_of_ne_zero hY,
+  rw [condDistrib_apply_of_ne_zero hX hY,
     Measure.map_apply hX (measurableSet_singleton _), Measure.map_apply hY hs,
     Measure.map_apply (hX.prodMk hY) ((measurableSet_singleton _).prod hs),
     cond_apply (hX (measurableSet_singleton _))]
@@ -451,7 +448,7 @@ lemma condDistrib_prod_of_forall_condDistrib_cond [Countable Ω'] [IsFiniteMeasu
     (h_cond : ∀ b, μ (Z ⁻¹' {b}) ≠ 0 → condDistrib Y X μ[|Z ⁻¹' {b}] =ᵐ[μ[|Z ⁻¹' {b}].map X]
       (κ.comap (fun ω ↦ (ω, b)) (by fun_prop) : Kernel β Ω)) :
     condDistrib Y (fun ω ↦ (X ω, Z ω)) μ =ᵐ[μ.map (fun ω ↦ (X ω, Z ω))] κ := by
-  refine condDistrib_ae_eq_of_measure_eq_compProd _ (by fun_prop) ?_
+  refine condDistrib_ae_eq_of_measure_eq_compProd (by fun_prop) (by fun_prop) ?_
   ext s hs
   suffices ∀ b, (Measure.map (fun x ↦ ((X x, Z x), Y x)) μ) (s ∩ {p | p.1.2 = b}) =
       (Measure.map (fun ω ↦ (X ω, Z ω)) μ ⊗ₘ κ) (s ∩ {p | p.1.2 = b}) by
@@ -503,8 +500,7 @@ lemma condDistrib_prod_of_forall_condDistrib_cond [Countable Ω'] [IsFiniteMeasu
       exact .inr hb
     rw [h_left, h_right]
   specialize h_cond b hb
-  rw [condDistrib_ae_eq_iff_measure_eq_compProd] at h_cond
-  swap; · fun_prop
+  rw [condDistrib_ae_eq_iff_measure_eq_compProd (by fun_prop) (by fun_prop)] at h_cond
   rw [Measure.ext_iff] at h_cond
   have hs' : MeasurableSet {p : β × Ω | ((p.1, b), p.2) ∈ s} := hs.preimage (by fun_prop)
   have h1 := h_cond {p | ((p.1, b), p.2) ∈ s} hs'
diff --git a/LeanMachineLearning/ForMathlib/Probability/Independence/IndepFun.lean b/LeanMachineLearning/ForMathlib/Probability/Independence/IndepFun.lean
index 80df7df..1bf6bc1 100644
--- a/LeanMachineLearning/ForMathlib/Probability/Independence/IndepFun.lean
+++ b/LeanMachineLearning/ForMathlib/Probability/Independence/IndepFun.lean
@@ -103,7 +103,7 @@ lemma IndepFun_map_iff [IsFiniteMeasure μ] {X : Ω' → E} {Y : Ω' → E} {f :
 lemma iIndepFun_map_iff [IsProbabilityMeasure μ] {X : ι → Ω' → E} {f : Ω → Ω'}
     (hf : AEMeasurable f μ) (hX : ∀ n, AEMeasurable (X n) (μ.map f)) :
     iIndepFun X (μ.map f) ↔ iIndepFun (fun n ↦ X n ∘ f) μ := by
-  have := Measure.isProbabilityMeasure_map hf (μ := μ)
+  have : IsProbabilityMeasure (μ.map f) := inferInstance
   rw [iIndepFun_iff_map_fun_eq_infinitePi_map₀' hX,
     iIndepFun_iff_map_fun_eq_infinitePi_map₀' (by fun_prop)]
   rw [AEMeasurable.map_map_of_aemeasurable (by fun_prop) hf]
diff --git a/LeanMachineLearning/ForMathlib/Probability/WithDensity.lean b/LeanMachineLearning/ForMathlib/Probability/WithDensity.lean
index 8471c6f..997bac7 100644
--- a/LeanMachineLearning/ForMathlib/Probability/WithDensity.lean
+++ b/LeanMachineLearning/ForMathlib/Probability/WithDensity.lean
@@ -6,6 +6,7 @@ Authors: Paulo Rauber
 module
 
 public import Mathlib.Probability.Kernel.CompProdEqIff
+public import Mathlib.Probability.Kernel.Composition.WithDensity
 public import Mathlib.Probability.Kernel.Composition.MeasureComp
 
 /-! # Lemmas about kernels and measures with density
diff --git a/LeanMachineLearning/Online/Bandit/Algorithms/Regret/BayesRegretTS.lean b/LeanMachineLearning/Online/Bandit/Algorithms/Regret/BayesRegretTS.lean
index 444f290..f434dd6 100644
--- a/LeanMachineLearning/Online/Bandit/Algorithms/Regret/BayesRegretTS.lean
+++ b/LeanMachineLearning/Online/Bandit/Algorithms/Regret/BayesRegretTS.lean
@@ -333,7 +333,7 @@ lemma integral_ucb_action_eq_integral_ucb_bestAction (hK : 0 < K)
     _ = ∫ ha, uc ha ∂P.map (fun ω ↦ (history A R n ω, A (n + 1) ω)) := by
         rw [← integral_map (by fun_prop) (by fun_prop)]
     _ = ∫ ha, uc ha ∂P.map (fun ω ↦ (history A R n ω, bestAction κ E ω)) := by
-        rw [← compProd_map_condDistrib (by fun_prop), ← compProd_map_condDistrib (by fun_prop),
+        rw [← compProd_map_condDistrib (by fun_prop) (by fun_prop), ← compProd_map_condDistrib (by fun_prop) (by fun_prop),
             Measure.compProd_congr (hasCondDistrib_action hK h n).condDistrib_eq]
     _ = P[fun ω ↦ ucb A R l u σ2 δ (bestAction κ E ω) (n + 1) ω] := by
         rw [integral_map (by fun_prop) (by fun_prop)]
diff --git a/LeanMachineLearning/Online/Bandit/Algorithms/TS.lean b/LeanMachineLearning/Online/Bandit/Algorithms/TS.lean
index 586502c..e0e54d1 100644
--- a/LeanMachineLearning/Online/Bandit/Algorithms/TS.lean
+++ b/LeanMachineLearning/Online/Bandit/Algorithms/TS.lean
@@ -64,7 +64,8 @@ def TS.initialPolicy (hK : 0 < K) (Q : Measure 𝓔) (κ : Kernel (𝓔 × Fin K
 
 instance {hK : 0 < K} {Q : Measure 𝓔} [IsProbabilityMeasure Q] {κ : Kernel (𝓔 × Fin K) ℝ} :
     IsProbabilityMeasure (TS.initialPolicy hK Q κ) :=
-  Measure.isProbabilityMeasure_map (by fun_prop)
+  have : Nonempty (Fin K) := Fin.pos_iff_nonempty.mp hK
+  inferInstanceAs (IsProbabilityMeasure (Q.map (bestAction κ id)))
 
 /-- The Thompson sampling algorithm with actions in `Fin K`, where `Q : Measure 𝓔` is a prior
   distribution over parameters, and `κ : Kernel (𝓔 × Fin K) ℝ` is a Markov kernel that defines the
@@ -112,6 +113,7 @@ lemma TS.hasCondDistrib_action (hK : 0 < K) (h : IsBayesAlgEnvSeq Q κ (tsAlgori
           simp_rw [Kernel.map_apply _ hm, IT.bayesTrajMeasurePosterior, hc]
       _ =ᵐ[P.map (history A R n)]
           condDistrib (bestAction κ E) (history A R n) P :=
-          (condDistrib_comp (history A R n) h.measurable_param.aemeasurable hm).symm
+          (condDistrib_comp (Learning.measurable_history h.measurable_action h.measurable_feedback n).aemeasurable
+            h.measurable_param.aemeasurable hm).symm
 
 end Bandits
diff --git a/LeanMachineLearning/Online/Bandit/ArrayProbSpace.lean b/LeanMachineLearning/Online/Bandit/ArrayProbSpace.lean
index 16e55bc..de95c37 100644
--- a/LeanMachineLearning/Online/Bandit/ArrayProbSpace.lean
+++ b/LeanMachineLearning/Online/Bandit/ArrayProbSpace.lean
@@ -477,6 +477,7 @@ lemma truePast_eq_of_pullCount_eq (alg : Algorithm 𝓐 R)
     truePast alg a n ω = (ω.1, fun i b ↦ if b = a then if m ≠ 0 then
       ω.2 (min i (m - 1)) a else Nonempty.some inferInstance else ω.2 i b) := by
   simp [truePast, h_pc]
+  rfl
 
 lemma truePast_eq_of_pullCount_eq_of_ne_zero (alg : Algorithm 𝓐 R)
     (a : 𝓐) (n m : ℕ) (ω : probSpace 𝓐 R)
@@ -484,6 +485,7 @@ lemma truePast_eq_of_pullCount_eq_of_ne_zero (alg : Algorithm 𝓐 R)
     truePast alg a n ω = (ω.1, fun i b ↦ if b = a then
       ω.2 (min i (m - 1)) a else ω.2 i b) := by
   simp [truePast, h_pc, hm]
+  rfl
 
 lemma measurable_hist_truePast [Countable 𝓐] (alg : Algorithm 𝓐 R)
     (a : 𝓐) (n : ℕ) :
@@ -574,8 +576,8 @@ lemma map_snd_apply_arrayMeasure {ν : Kernel 𝓐 R} [IsMarkovKernel ν] (n : 
     (arrayMeasure ν).map (fun ω ↦ ω.2 n a) = ν a := by
   calc (arrayMeasure ν).map (fun ω ↦ ω.2 n a)
   _ = (arrayMeasure ν).snd.map (fun ω ↦ ω n a) := by
-    rw [Measure.snd, Measure.map_map (by fun_prop) (by fun_prop)]
-    rfl
+    exact (Measure.map_map (f := Prod.snd) (g := fun ω : ℕ → 𝓐 → R ↦ ω n a)
+      (μ := arrayMeasure ν) (by fun_prop) (by fun_prop)).symm
   _ = ν a := by
     rw [arrayMeasure, Measure.snd_prod, streamMeasure]
     have : (fun ω ↦ ω n a) = (fun h : 𝓐 → R ↦ h a) ∘ (fun ω : ℕ → 𝓐 → R ↦ ω n) := rfl
@@ -658,6 +660,7 @@ lemma indepFun_snd_apply_aux (ν : Kernel 𝓐 R) [IsMarkovKernel ν] (a : 𝓐)
     (fun ω ↦ ω.2 m a) ⟂ᵢ[arrayMeasure ν]
       (fun ω ↦ (ω.1, fun k b ↦ if b = a then if m ≠ 0 then ω.2 (min k (m - 1)) b
         else Nonempty.some inferInstance else ω.2 k b)) := by
+  rw [indepFun_iff_measure_inter_preimage_eq_mul]
   unfold arrayMeasure
   let μ₁ : Measure (ℕ → I) := Measure.infinitePi fun _ ↦ volume
   let μ₂ : Measure (ℕ → 𝓐 → R) := Measure.infinitePi fun _ ↦ Measure.infinitePi ν
@@ -667,7 +670,6 @@ lemma indepFun_snd_apply_aux (ν : Kernel 𝓐 R) [IsMarkovKernel ν] (a : 𝓐)
   -- Independence within each infinitePi ν: coordinates f b are independent
   have h_indep_inner : iIndepFun (fun (b : 𝓐) (f : 𝓐 → R) ↦ f b) (Measure.infinitePi ν) :=
     iIndepFun_infinitePi (fun _ ↦ measurable_id)
-  rw [indepFun_iff_measure_inter_preimage_eq_mul]
   intro s t hs ht
   let X : (ℕ → I) × (ℕ → 𝓐 → R) → R := fun ω ↦ ω.2 m a
   let Y : (ℕ → I) × (ℕ → 𝓐 → R) → (ℕ → I) × (ℕ → 𝓐 → R) :=
@@ -680,8 +682,8 @@ lemma indepFun_snd_apply_aux (ν : Kernel 𝓐 R) [IsMarkovKernel ν] (a : 𝓐)
       (ω.1, fun k b ↦ if b = a then if m ≠ 0 then ω.2 (min k (m - 1)) b
         else Nonempty.some inferInstance else ω.2 k b))
     refine Measurable.prod measurable_fst ?_
-    refine measurable_pi_lambda _ (fun k ↦ ?_)
-    refine measurable_pi_lambda _ (fun b ↦ ?_)
+    refine Measurable.of_eval (fun k ↦ ?_)
+    refine Measurable.of_eval (fun b ↦ ?_)
     by_cases hb : b = a
     · simp only [hb, ↓reduceIte]
       by_cases hm : m ≠ 0
@@ -715,8 +717,8 @@ lemma indepFun_snd_apply_aux (ν : Kernel 𝓐 R) [IsMarkovKernel ν] (a : 𝓐)
   -- Now we use independence in μ₂: (ω₂ m a) is independent of (trunc ω₂)
   -- because trunc only uses indices (k, a) with k < m, and (k, b) with b ≠ a
   have h_trunc_meas : Measurable trunc := by
-    refine measurable_pi_lambda _ (fun k ↦ ?_)
-    refine measurable_pi_lambda _ (fun b ↦ ?_)
+    refine Measurable.of_eval (fun k ↦ ?_)
+    refine Measurable.of_eval (fun b ↦ ?_)
     simp only [trunc]
     by_cases hb : b = a
     · simp only [hb, ↓reduceIte]
@@ -739,17 +741,17 @@ lemma indepFun_snd_apply_aux (ν : Kernel 𝓐 R) [IsMarkovKernel ν] (a : 𝓐)
       intro ω₂; ext k b; simp only [trunc, proj]; by_cases hb : b = a <;> simp [hb]; grind
     have h_proj_meas : Measurable proj := by
       refine Measurable.prod ?_ ?_
-      · refine measurable_pi_lambda _ fun k ↦ ?_
+      · refine Measurable.of_eval fun k ↦ ?_
         by_cases hm : m ≠ 0
         · simp only [proj, ne_eq, hm, not_false_eq_true, ↓reduceIte]
           exact (measurable_pi_apply a).comp (measurable_pi_apply (min k (m - 1)))
         · simp [proj, hm]
-      · exact measurable_pi_lambda _ (fun k ↦ measurable_pi_lambda _ (fun ⟨b, _⟩ ↦
+      · exact Measurable.of_eval (fun k ↦ Measurable.of_eval (fun ⟨b, _⟩ ↦
           (measurable_pi_apply b).comp (measurable_pi_apply k)))
     have h_g_meas : Measurable (fun p : (ℕ → R) × (ℕ → {b : 𝓐 // b ≠ a} → R) ↦
         (fun k b ↦ if h : b = a then if m ≠ 0 then p.1 k else Nonempty.some inferInstance
           else p.2 k ⟨b, h⟩)) := by
-      refine measurable_pi_lambda _ (fun k ↦ measurable_pi_lambda _ (fun b ↦ ?_))
+      refine Measurable.of_eval (fun k ↦ Measurable.of_eval (fun b ↦ ?_))
       by_cases hb : b = a
       · simp only [hb, ↓reduceDIte]
         by_cases hm : m ≠ 0
@@ -784,9 +786,9 @@ lemma indepFun_snd_apply_aux (ν : Kernel 𝓐 R) [IsMarkovKernel ν] (a : 𝓐)
       have h_iindep : iIndepFun (fun (p : ℕ × 𝓐) ω ↦ ω p.1 p.2) μ₂ :=
         iIndepFun_uncurry_infinitePi' (X := fun _ _ ↦ id) (fun _ ↦ ν) (by fun_prop)
       have h_rows_meas : Measurable rows_lt_m :=
-        measurable_pi_lambda _ (fun ⟨j, _⟩ ↦ (measurable_pi_apply a).comp (measurable_pi_apply j))
+        Measurable.of_eval (fun ⟨j, _⟩ ↦ (measurable_pi_apply a).comp (measurable_pi_apply j))
       have h_other_meas : Measurable other_cols :=
-        measurable_pi_lambda _ (fun k ↦ measurable_pi_lambda _ (fun ⟨b, _⟩ ↦
+        Measurable.of_eval (fun k ↦ Measurable.of_eval (fun ⟨b, _⟩ ↦
           (measurable_pi_apply b).comp (measurable_pi_apply k)))
       -- Show (ω₂ m a) ⟂ (rows_lt_m, other_cols) via indep_iSup_of_disjoint
       have h_indep_combined : IndepFun (fun ω₂ ↦ ω₂ m a)
@@ -831,7 +833,7 @@ lemma indepFun_snd_apply_aux (ν : Kernel 𝓐 R) [IsMarkovKernel ν] (a : 𝓐)
           else Nonempty.some inferInstance) p.1, p.2) ∈ t'}
       have hT_meas : MeasurableSet T := by
         refine ht'.preimage (Measurable.prod ?_ measurable_snd)
-        refine measurable_pi_lambda _ (fun k ↦ ?_)
+        refine Measurable.of_eval (fun k ↦ ?_)
         by_cases hm : m ≠ 0
         · simp only [ne_eq, hm, not_false_eq_true, ↓reduceDIte]
           exact (measurable_pi_apply (⟨min k (m - 1), Finset.mem_Iio.mpr (h_row_bound hm k)⟩ :
@@ -911,16 +913,18 @@ lemma indepFun_snd_hist_cond [Countable 𝓐] (alg : Algorithm 𝓐 R)
         pullCount (action alg) a (n + 1) ω = m}).indicator (fun _ ↦ 1)) ⁻¹' {1}]]
       fun ω ↦ (ω.1, fun k b ↦ if b = a then if m ≠ 0 then ω.2 (min k (m - 1)) b
         else Nonempty.some inferInstance else ω.2 k b) by
-    convert this using 1
-    · rfl
-    · rfl
-    · rfl
-    congr with ω
-    simp only [Set.mem_preimage, Set.mem_singleton_iff, Prod.mk.injEq, Set.indicator_apply,
+    have hset : (fun ω ↦ (action alg (n + 1) ω,
+          pullCount (action alg) (action alg (n + 1) ω) (n + 1) ω)) ⁻¹' {(a, m)} =
+        (({ω | action alg (n + 1) ω = a ∧
+          pullCount (action alg) a (n + 1) ω = m}).indicator (fun _ ↦ (1 : ℕ))) ⁻¹' {1} := by
+      ext ω
+      simp only [Set.mem_preimage, Set.mem_singleton_iff, Prod.mk.injEq, Set.indicator_apply,
       Set.mem_setOf_eq, ite_eq_left_iff, not_and, zero_ne_one, imp_false,
       Classical.not_imp, Decidable.not_not, and_congr_right_iff]
-    intro ha
-    simp [ha]
+      intro ha
+      simp [ha]
+    rw [hset]
+    exact this
   have h_meas := measurable_stepsUntil alg a m n
   obtain ⟨f, hf, hf_eq⟩ := h_meas.exists_eq_measurable_comp
   simp_rw [hf_eq]
@@ -943,9 +947,11 @@ lemma hasLaw_action_zero (alg : Algorithm 𝓐 R) (ν : Kernel 𝓐 R) [IsMarkov
   map_eq := by
     calc (arrayMeasure ν).map (fun ω ↦ initAlgFunction alg (ω.1 0))
     _ = ((arrayMeasure ν).fst.map (Function.eval 0)).map (initAlgFunction alg) := by
-      rw [Measure.fst, Measure.map_map (by fun_prop) (by fun_prop),
-        Measure.map_map (by fun_prop) (by fun_prop)]
-      rfl
+      have h1 := Measure.map_map (f := Function.eval 0) (g := initAlgFunction alg)
+        (μ := (arrayMeasure ν).map Prod.fst) (by fun_prop) (by fun_prop)
+      have h2 := Measure.map_map (f := Prod.fst) (g := initAlgFunction alg ∘ Function.eval 0)
+        (μ := arrayMeasure ν) (by fun_prop) (by fun_prop)
+      exact h2.symm.trans h1.symm
     _ = (volume : Measure I).map (initAlgFunction alg) := by
       simp only [arrayMeasure, Measure.fst_prod]
       rw [(measurePreserving_eval_infinitePi (fun _ ↦ volume) 0).map_eq]
diff --git a/LeanMachineLearning/Online/Bandit/RewardByCountMeasure.lean b/LeanMachineLearning/Online/Bandit/RewardByCountMeasure.lean
index 5284ca3..821e102 100644
--- a/LeanMachineLearning/Online/Bandit/RewardByCountMeasure.lean
+++ b/LeanMachineLearning/Online/Bandit/RewardByCountMeasure.lean
@@ -59,7 +59,7 @@ lemma condDistrib_reward'' [Countable 𝓐]
   rw [h_law]
   have h_prod : 𝓛[fun ω ↦ R n ω.1 | fun ω ↦ A n ω.1; 𝔓]
       =ᵐ[P.map (A n)] 𝓛[R n | A n; P] :=
-    condDistrib_fst_prod _ (by fun_prop) _
+    condDistrib_fst_prod (by fun_prop) (by fun_prop) _
   filter_upwards [h_ra', h_prod] with ω h_eq h_prod
   rw [h_prod, h_eq]
 
@@ -213,7 +213,7 @@ lemma hasLaw_rewardByCount [StandardBorelSpace Ω] [Countable 𝓐]
       Measure.comp_congr h_condDistrib
     _ = ν a := by
       have : IsProbabilityMeasure ((𝔓).map (fun ω ↦ stepsUntil A a m ω.1)) :=
-        Measure.isProbabilityMeasure_map (by fun_prop)
+        inferInstance
       simp
 
 lemma identDistrib_rewardByCount [StandardBorelSpace Ω] [Countable 𝓐]
diff --git a/LeanMachineLearning/Online/Bandit/SumRewards.lean b/LeanMachineLearning/Online/Bandit/SumRewards.lean
index 85de702..10e1faa 100644
--- a/LeanMachineLearning/Online/Bandit/SumRewards.lean
+++ b/LeanMachineLearning/Online/Bandit/SumRewards.lean
@@ -143,12 +143,14 @@ lemma sumRewards_eq_comp :
      (fun p ↦ ∑ i ∈ range n, if (p i).1 = a then (p i).2 else 0) ∘ (trajectory A R) := by
   ext
   simp [sumRewards, trajectory]
+  rfl
 
 lemma pullCount_eq_comp :
     pullCount A a n =
       (fun p ↦ ∑ i ∈ range n, if (p i).1 = a then 1 else 0) ∘ (trajectory A R) := by
   ext
   simp [pullCount, trajectory]
+  rfl
 
 -- todo: write those lemmas with IdentDistrib instead of equality of maps
 lemma _root_.Learning.IsAlgEnvSeq.law_sumRewards_unique [MeasurableSingletonClass 𝓐]
@@ -708,7 +710,7 @@ lemma prob_empMean_sub_actionMean_ge_le (h : IsBayesAlgEnvSeq Q κ alg E A R P)
         rw [empMean] at hle
         exact ⟨a, t, ht, hpc, sqrt_two_mul_le_sub hpc hle⟩
     _ = (P.map E ⊗ₘ condDistrib (trajectory A R) E P) S := by
-        rw [← compProd_map_condDistrib (by fun_prop)]
+        rw [← compProd_map_condDistrib (by fun_prop) (by fun_prop)]
     _ = ∫⁻ e, condDistrib (trajectory A R) E P e (Prod.mk e ⁻¹' S) ∂(P.map E) :=
         Measure.compProd_apply (by measurability)
     _ ≤ ∫⁻ e, ENNReal.ofReal (Fintype.card (Fin K) * (n - 1) * δ) ∂(P.map E) := by
@@ -748,7 +750,7 @@ lemma prob_empMean_bestAction_sub_actionMean_le_le (h : IsBayesAlgEnvSeq Q κ al
         rw [empMean] at hle
         exact ⟨t, ht, hpc, sub_le_neg_sqrt_two_mul hpc hle⟩
     _ = (P.map E ⊗ₘ condDistrib (trajectory A R) E P) S := by
-        rw [← compProd_map_condDistrib (by fun_prop)]
+        rw [← compProd_map_condDistrib (by fun_prop) (by fun_prop)]
     _ = ∫⁻ e, condDistrib (trajectory A R) E P e (Prod.mk e ⁻¹' S) ∂(P.map E) :=
         Measure.compProd_apply (by measurability)
     _ ≤ ∫⁻ e, ENNReal.ofReal ((n - 1) * δ) ∂(P.map E) := by
diff --git a/LeanMachineLearning/SequentialLearning/AlgorithmDensityBayes.lean b/LeanMachineLearning/SequentialLearning/AlgorithmDensityBayes.lean
index ad4b00c..5dbd8cb 100644
--- a/LeanMachineLearning/SequentialLearning/AlgorithmDensityBayes.lean
+++ b/LeanMachineLearning/SequentialLearning/AlgorithmDensityBayes.lean
@@ -97,12 +97,13 @@ lemma hasCondDistrib_env_history (h : IsBayesAlgEnvSeq Q κ alg E A Y P)
     have hA₀ := h₀.measurable_action
     have hY₀ := h₀.measurable_feedback
     have hE₀ := h₀.measurable_param
-    rw [← map_swap_compProd_map_condDistrib (by fun_prop), h.hasLaw_env.map_eq,
+    have hE := h.measurable_param
+    rw [← map_swap_compProd_map_condDistrib (by fun_prop) (by fun_prop), h.hasLaw_env.map_eq,
       Measure.compProd_eq_compProd_withDensity_comp_snd (by fun_prop)
         (h.condDistrib_history_eq_condDistrib_hist_withDensity h₀ hc n),
       map_swap_withDensity_comp_snd (by fun_prop),
-      ← h₀.hasLaw_env.map_eq, map_swap_compProd_map_condDistrib (by fun_prop),
-      ← compProd_map_condDistrib (by fun_prop), ← Measure.compProd_withDensity_left (by fun_prop),
+      ← h₀.hasLaw_env.map_eq, map_swap_compProd_map_condDistrib (by fun_prop) (by fun_prop),
+      ← compProd_map_condDistrib (by fun_prop) (by fun_prop), ← Measure.compProd_withDensity_left (by fun_prop),
       ← (hasLaw_history_withDensity h h₀ hc n).map_eq]
 
 end IsBayesAlgEnvSeq
diff --git a/LeanMachineLearning/SequentialLearning/Algorithms/RoundRobin.lean b/LeanMachineLearning/SequentialLearning/Algorithms/RoundRobin.lean
index 14c8cf3..1f68340 100644
--- a/LeanMachineLearning/SequentialLearning/Algorithms/RoundRobin.lean
+++ b/LeanMachineLearning/SequentialLearning/Algorithms/RoundRobin.lean
@@ -56,10 +56,12 @@ lemma sum_mod_range_mul {K : ℕ} (hK : 0 < K) (m : ℕ) (a : Fin K) :
       let e : ℕ ↪ ℕ := ⟨fun i : ℕ ↦ i + K * n, fun i j hij ↦ by grind⟩
       have : Finset.map e (range K) = Ico (K * n) (K * n + K) := by
         ext x
-        simp only [mem_map, mem_range, Function.Embedding.coeFn_mk, mem_Ico, e]
-        refine ⟨fun h ↦ by grind, fun h ↦ ?_⟩
-        use x - K * n
-        grind
+        simp only [mem_map, mem_range, mem_Ico, e]
+        refine ⟨fun ⟨w, hw, h⟩ ↦ ?_, fun h ↦ ⟨x - K * n, by omega, ?_⟩⟩
+        · change w + K * n = x at h
+          omega
+        · change x - K * n + K * n = x
+          omega
       rw [← this, Finset.sum_map]
       congr
     _ = n + (∑ s ∈ range K, if ⟨s % K, Nat.mod_lt _ hK⟩ = a then 1 else 0) := by simp
diff --git a/LeanMachineLearning/SequentialLearning/BayesStationaryEnv.lean b/LeanMachineLearning/SequentialLearning/BayesStationaryEnv.lean
index 68cc8ff..b6cf8a9 100644
--- a/LeanMachineLearning/SequentialLearning/BayesStationaryEnv.lean
+++ b/LeanMachineLearning/SequentialLearning/BayesStationaryEnv.lean
@@ -94,7 +94,7 @@ variable [StandardBorelSpace 𝓐] [Nonempty 𝓐] [StandardBorelSpace 𝓨] [No
 lemma hasLaw_IT_action_zero (h : IsBayesAlgEnvSeq Q κ alg E A Y P) :
     ∀ᵐ e ∂Q, HasLaw (IT.action 0) alg.p0 (condDistrib (trajectory A Y) E P e) := by
   rw [← h.hasLaw_env.map_eq]
-  filter_upwards [condDistrib_comp E
+  filter_upwards [condDistrib_comp h.hasLaw_env.aemeasurable
       ((measurable_trajectory h.measurable_action h.measurable_feedback).aemeasurable)
       (IT.measurable_action (𝓐 := 𝓐) (𝓨 := 𝓨) 0),
     h.hasCondDistrib_action_zero.condDistrib_eq] with _ hc hcd
@@ -137,7 +137,7 @@ lemma hasLaw_IT_hist (h : IsBayesAlgEnvSeq Q κ alg E A Y P) (n : ℕ) :
     ∀ᵐ e ∂Q, HasLaw (IT.hist n) (condDistrib (history A Y n) E P e)
       (condDistrib (trajectory A Y) E P e) := by
   rw [← h.hasLaw_env.map_eq, show history A Y n = IT.hist n ∘ trajectory A Y from rfl]
-  filter_upwards [condDistrib_comp E
+  filter_upwards [condDistrib_comp h.hasLaw_env.aemeasurable
     (measurable_trajectory h.measurable_action h.measurable_feedback).aemeasurable
     (IT.measurable_hist n)] with _ he
   exact ⟨(IT.measurable_hist n).aemeasurable, by
diff --git a/LeanMachineLearning/SequentialLearning/FiniteActions.lean b/LeanMachineLearning/SequentialLearning/FiniteActions.lean
index 806bc7c..c8cf2a3 100644
--- a/LeanMachineLearning/SequentialLearning/FiniteActions.lean
+++ b/LeanMachineLearning/SequentialLearning/FiniteActions.lean
@@ -401,7 +401,6 @@ lemma action_stepsUntil (hm : m ≠ 0) (h_exists : ∃ s, pullCount A a (s + 1)
   · simp [h_pos]
   rw [Nat.sub_add_cancel (by omega)]
   rwa [← pullCount_eq_pullCount_of_action_ne]
-  exact h_ne
 
 lemma action_eq_of_stepsUntil_eq_coe (hm : m ≠ 0) (h : stepsUntil A a m ω = n) :
     A n ω = a := by
@@ -696,7 +695,8 @@ lemma rewardByCount_of_stepsUntil_ne_top (h : stepsUntil A a m ω.1 ≠ ⊤) :
 
 lemma rewardByCount_eq_stoppedValue (h : stepsUntil A a m ω.1 ≠ ⊤) :
     rewardByCount A R' a m ω = stoppedValue R' (stepsUntil A a m) ω.1 := by
-  rw [rewardByCount_of_stepsUntil_ne_top h, stoppedValue]
+  rw [rewardByCount_of_stepsUntil_ne_top h]
+  unfold stoppedValue
   lift stepsUntil A a m ω.1 to ℕ using h with n
   simp
 
diff --git a/LeanMachineLearning/SequentialLearning/StationaryEnv.lean b/LeanMachineLearning/SequentialLearning/StationaryEnv.lean
index 39539bd..7025ba1 100644
--- a/LeanMachineLearning/SequentialLearning/StationaryEnv.lean
+++ b/LeanMachineLearning/SequentialLearning/StationaryEnv.lean
@@ -88,10 +88,10 @@ lemma hasCondDistrib_feedback [IsObliviousEnv env] (h : IsAlgEnvSeq A Y alg env
     have h_eq := (h.hasCondDistrib_feedback n).map_eq
     have : P.map (A (n + 1)) =
         (P.map (fun x ↦ (history A Y n x, A (n + 1) x))).snd := by
-      rw [Measure.snd_map_prodMk (by fun_prop)]
+      rw [Measure.snd_map_prodMk (by fun_prop) (by fun_prop)]
     simp only [feedback_eq_feedbackCondAction] at h_eq
     rw [this, ← Measure.snd_prodAssoc_compProd_prodMkLeft, ← h_eq,
-      Measure.snd_map_prodMk (by fun_prop), Measure.map_map (by fun_prop) (by fun_prop)]
+      Measure.snd_map_prodMk (by fun_prop) (by fun_prop), Measure.map_map (by fun_prop) (by fun_prop)]
     congr
 
 variable [StandardBorelSpace 𝓐] [Nonempty 𝓐] [StandardBorelSpace 𝓨] [Nonempty 𝓨]