From 7ae2db685c34add9dc8ffb3c3e5f249eb919c8d5 Mon Sep 17 00:00:00 2001
From: Remy Degenne <remydegenne@gmail.com>
Date: Thu, 10 Sep 2026 19:14:28 +0200
Subject: [PATCH] MeasurableSpace on kernels, algorithms, environments

---
 LeanMachineLearning.lean                      |  1 +
 .../Probability/Kernel/MeasurableSpace.lean   | 37 +++++++++++++++++++
 .../SequentialLearning/Algorithm.lean         | 20 ++++++++++
 3 files changed, 58 insertions(+)
 create mode 100644 LeanMachineLearning/ForMathlib/Probability/Kernel/MeasurableSpace.lean

diff --git a/LeanMachineLearning.lean b/LeanMachineLearning.lean
index 6516f85c..4f59f7d8 100644
--- a/LeanMachineLearning.lean
+++ b/LeanMachineLearning.lean
@@ -28,6 +28,7 @@ public import LeanMachineLearning.ForMathlib.Probability.Kernel.Composition.MapC
 public import LeanMachineLearning.ForMathlib.Probability.Kernel.Composition.MeasureCompProd
 public import LeanMachineLearning.ForMathlib.Probability.Kernel.IonescuTulcea.Traj
 public import LeanMachineLearning.ForMathlib.Probability.Kernel.KernelSub
+public import LeanMachineLearning.ForMathlib.Probability.Kernel.MeasurableSpace
 public import LeanMachineLearning.ForMathlib.Probability.Moments.SubExponential
 public import LeanMachineLearning.ForMathlib.Probability.Moments.SubGaussian
 public import LeanMachineLearning.ForMathlib.Probability.WithDensity
diff --git a/LeanMachineLearning/ForMathlib/Probability/Kernel/MeasurableSpace.lean b/LeanMachineLearning/ForMathlib/Probability/Kernel/MeasurableSpace.lean
new file mode 100644
index 00000000..a984f069
--- /dev/null
+++ b/LeanMachineLearning/ForMathlib/Probability/Kernel/MeasurableSpace.lean
@@ -0,0 +1,37 @@
+/-
+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.Probability.Kernel.Basic
+
+/-!
+# Measurable space of kernels
+
+-/
+
+@[expose] public section
+
+open MeasureTheory
+
+namespace ProbabilityTheory
+
+variable {𝓧 𝓨 𝓩 : Type*} {m𝓧 : MeasurableSpace 𝓧} {m𝓨 : MeasurableSpace 𝓨} {m𝓩 : MeasurableSpace 𝓩}
+
+instance instMeasurableSpaceKernel : MeasurableSpace (Kernel 𝓧 𝓨) :=
+  MeasurableSpace.comap (fun κ ↦ (κ : 𝓧 → Measure 𝓨)) inferInstance
+
+lemma measurable_kernel_iff (f : 𝓧 → Kernel 𝓨 𝓩) : Measurable f ↔ ∀ y, Measurable (f · y) := by
+  unfold instMeasurableSpaceKernel
+  rw [measurable_comap_iff, measurable_pi_iff]
+  simp
+
+@[fun_prop]
+lemma Kernel.measurable_const : Measurable (@Kernel.const 𝓧 𝓨 m𝓧 m𝓨) := by
+  rw [measurable_kernel_iff]
+  simp only [const_apply]
+  fun_prop
+
+end ProbabilityTheory
diff --git a/LeanMachineLearning/SequentialLearning/Algorithm.lean b/LeanMachineLearning/SequentialLearning/Algorithm.lean
index c711163c..bfd28276 100644
--- a/LeanMachineLearning/SequentialLearning/Algorithm.lean
+++ b/LeanMachineLearning/SequentialLearning/Algorithm.lean
@@ -7,6 +7,7 @@ module
 
 public import LeanMachineLearning.ForMathlib.MeasureTheory.Measurable
 public import LeanMachineLearning.ForMathlib.Probability.Kernel.IonescuTulcea.Traj
+public import LeanMachineLearning.ForMathlib.Probability.Kernel.MeasurableSpace
 
 /-!
 # Algorithms and environments
@@ -110,6 +111,15 @@ structure Algorithm (𝓞 𝓐 𝓨 : Type*) [MeasurableSpace 𝓞] [MeasurableS
 instance (alg : Algorithm 𝓞 𝓐 𝓨) (n : ℕ) : IsMarkovKernel (alg.policy n) :=
   alg.isMarkovKernel_policy n
 
+instance : MeasurableSpace (Algorithm 𝓞 𝓐 𝓨) :=
+  MeasurableSpace.comap (fun alg ↦ alg.policy) inferInstance
+
+lemma measurable_algorithm_iff (f : Ω → Algorithm 𝓞 𝓐 𝓨) :
+    Measurable f ↔ ∀ n, Measurable fun x ↦ (f x).policy n := by
+  unfold instMeasurableSpaceAlgorithm
+  rw [measurable_comap_iff, measurable_pi_iff]
+  simp
+
 /-- A stochastic environment.
 At each round, an observation is drawn prior to the algorithm taking an action. Then the environment
 provides feedback based on the observation and the action. -/
@@ -129,6 +139,16 @@ instance (env : Environment 𝓞 𝓐 𝓨) (n : ℕ) : IsMarkovKernel (env.obs
 instance (env : Environment 𝓞 𝓐 𝓨) (n : ℕ) : IsMarkovKernel (env.feedback n) :=
   env.isMarkovKernel_feedback n
 
+instance : MeasurableSpace (Environment 𝓞 𝓐 𝓨) :=
+  MeasurableSpace.comap (fun env ↦ (env.obs, env.feedback)) inferInstance
+
+lemma measurable_environment_iff (f : Ω → Environment 𝓞 𝓐 𝓨) :
+    Measurable f ↔
+      ∀ n, Measurable (fun x ↦ (f x).obs n) ∧ Measurable (fun x ↦ (f x).feedback n) := by
+  simp_rw [measurable_comap_iff, measurable_fun_prod, forall_and, measurable_pi_iff,
+    measurable_kernel_iff]
+  rfl
+
 /-- Distribution of the first observation: the observation kernel at time `0` applied to the empty
 history. -/
 def Environment.obs0 (env : Environment 𝓞 𝓐 𝓨) : Measure 𝓞 :=