Prove ae_eq_of_condDistrib_eq_deterministic

LeanBandits/ForMathlib/CondDistrib.lean

@@ -347,6 +347,28 @@ lemma indepFun_snd_prod (hX : AEMeasurable X μ) (hY : AEMeasurable Y μ) (h_ind
       AEMeasurable.map_map_of_aemeasurable (by fun_prop) (by fun_prop)]
   rfl
 
+-- cf. measurableSet_graph (Mathlib/MeasureTheory/Measure/Lebesgue/Basic.lean)
+lemma measurableSet_graph' {β Ω : Type*} [MeasurableSpace β] [MeasurableSpace Ω]
+    [StandardBorelSpace Ω] {f : β → Ω} (hf : Measurable f) :
+    MeasurableSet {p : β × Ω | p.2 = f p.1} := by
+  letI := upgradeStandardBorel Ω
+  exact (measurable_snd.prodMk (by fun_prop)) isClosed_diagonal.measurableSet
+
+omit [Nonempty Ω] [IsFiniteMeasure μ] in
+lemma ae_eq_of_map_prodMk_eq {f : β → Ω} (hf : Measurable f) (hX : AEMeasurable X μ)
+    (hY : AEMeasurable Y μ) (h : μ.map (fun ω ↦ (X ω, Y ω)) = μ.map (fun ω ↦ (X ω, f (X ω)))) :
+    Y =ᵐ[μ] f ∘ X := by
+  have hp : ∀ᵐ p ∂μ.map (fun ω ↦ (X ω, f (X ω))), p.2 = f p.1 :=
+    (ae_map_iff (by fun_prop) (measurableSet_graph' hf)).2 (by simp)
+  exact ae_of_ae_map (by fun_prop) (h ▸ hp)
+
+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
+  exact ae_eq_of_map_prodMk_eq hf hX hY (hfX ▸ h)
+
 end CondDistrib
 
 section Cond

LeanBandits/SequentialLearning/Deterministic.lean

@@ -40,17 +40,13 @@ lemma action_zero_detAlgorithm [MeasurableSingletonClass α] : action 0 =ᵐ[
     simp [detAlgorithm]
   exact ae_of_ae_map (by fun_prop) h_eq
 
-lemma action_detAlgorithm_ae_eq
-    [StandardBorelSpace α] [Nonempty α] [StandardBorelSpace R] [Nonempty R]
-    (n : ℕ) :
-    action (n + 1) =ᵐ[𝔓] fun h ↦ nextaction n (fun i ↦ h i) := by
-  -- rhs equals nextAction n ∘ hist n
-  have h := condDistrib_action (detAlgorithm nextaction h_next action0) env n
-  simp only [detAlgorithm_policy] at h
-  sorry
+lemma action_detAlgorithm_ae_eq [StandardBorelSpace α] [Nonempty α] [StandardBorelSpace R]
+    [Nonempty R] (n : ℕ) : action (n + 1) =ᵐ[𝔓] fun h ↦ nextaction n (hist n h) :=
+  ae_eq_of_condDistrib_eq_deterministic (by fun_prop) (by fun_prop) (by fun_prop)
+    (condDistrib_action (detAlgorithm nextaction h_next action0) env n)
 
 example [StandardBorelSpace α] [Nonempty α] [StandardBorelSpace R] [Nonempty R] :
-    ∀ᵐ h ∂𝔓, action 0 h = action0 ∧ ∀ n, action (n + 1) h = nextaction n (fun i ↦ h i) := by
+    ∀ᵐ h ∂𝔓, action 0 h = action0 ∧ ∀ n, action (n + 1) h = nextaction n (hist n h) := by
   rw [eventually_and, ae_all_iff]
   exact ⟨action_zero_detAlgorithm, action_detAlgorithm_ae_eq⟩