From b585477116ff657e779de2e9570f0c2b844f6cff Mon Sep 17 00:00:00 2001
From: Remy Degenne <remydegenne@gmail.com>
Date: Wed, 2 Sep 2026 10:39:43 +0200
Subject: [PATCH 1/3] add mdp def

---
 LeanMachineLearning.lean                      |  2 +
 .../Probability/Kernel/MeasurableSpace.lean   | 37 ++++++++
 .../ReinforcementLearning/MDP/Basic.lean      | 92 +++++++++++++++++++
 .../SequentialLearning/Algorithm.lean         | 20 ++++
 4 files changed, 151 insertions(+)
 create mode 100644 LeanMachineLearning/ForMathlib/Probability/Kernel/MeasurableSpace.lean
 create mode 100644 LeanMachineLearning/ReinforcementLearning/MDP/Basic.lean

diff --git a/LeanMachineLearning.lean b/LeanMachineLearning.lean
index 3491eb85..c9d7d801 100644
--- a/LeanMachineLearning.lean
+++ b/LeanMachineLearning.lean
@@ -20,6 +20,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.SubGaussian
 public import LeanMachineLearning.ForMathlib.Probability.WithDensity
 public import LeanMachineLearning.ForMathlib.Topology.Instances.ENNReal.Lemmas
@@ -32,6 +33,7 @@ public import LeanMachineLearning.Online.Bandit.BayesRegret
 public import LeanMachineLearning.Online.Bandit.Regret
 public import LeanMachineLearning.Online.Bandit.RewardByCountMeasure
 public import LeanMachineLearning.Online.Bandit.SumRewards
+public import LeanMachineLearning.ReinforcementLearning.MDP.Basic
 public import LeanMachineLearning.SequentialLearning.ActionIndicator
 public import LeanMachineLearning.SequentialLearning.Algorithm
 public import LeanMachineLearning.SequentialLearning.AlgorithmDensity
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/ReinforcementLearning/MDP/Basic.lean b/LeanMachineLearning/ReinforcementLearning/MDP/Basic.lean
new file mode 100644
index 00000000..fe0aa460
--- /dev/null
+++ b/LeanMachineLearning/ReinforcementLearning/MDP/Basic.lean
@@ -0,0 +1,92 @@
+/-
+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 LeanMachineLearning.SequentialLearning.IonescuTulceaSpace
+public import LeanMachineLearning.SequentialLearning.Deterministic
+
+/-!
+# Markov decision processes
+
+-/
+
+@[expose] public section
+
+open MeasureTheory ProbabilityTheory Finset Learning
+
+/-- Markov decision process with state space `𝓢`, action space `𝓐`, and reward space `𝓡`, described
+by a transition kernel `P : Kernel (𝓢 × 𝓐) 𝓢` and a reward kernel `R : Kernel (𝓢 × 𝓐) 𝓡`.
+See `MDP.env` for the environment associated with the MDP. -/
+structure MDP (𝓢 𝓐 𝓡 : Type*) [MeasurableSpace 𝓢] [MeasurableSpace 𝓐] [MeasurableSpace 𝓡] where
+  P : Kernel (𝓢 × 𝓐) 𝓢
+  [hP : IsMarkovKernel P]
+  R : Kernel (𝓢 × 𝓐) 𝓡
+  [hR : IsMarkovKernel R]
+
+namespace Learning.MDP
+
+variable {𝓢 𝓐 𝓡 : Type*} {m𝓢 : MeasurableSpace 𝓢} {m𝓐 : MeasurableSpace 𝓐} {m𝓡 : MeasurableSpace 𝓡}
+
+instance (M : MDP 𝓢 𝓐 𝓡) : IsMarkovKernel M.P := M.hP
+instance (M : MDP 𝓢 𝓐 𝓡) : IsMarkovKernel M.R := M.hR
+
+/-! ### The environment -/
+
+open Classical in
+protected noncomputable def env [h𝓢 : Nonempty 𝓢] (M : MDP 𝓢 𝓐 𝓡) (μ₀ : Measure 𝓢) :
+    Environment 𝓢 𝓐 𝓡 where
+      obs
+      | 0 => Kernel.const _ (if IsProbabilityMeasure μ₀ then μ₀ else Measure.dirac h𝓢.some)
+      | n + 1 => M.P.comap (fun h ↦ ((h (Fin.last n)).obs, (h (Fin.last n)).action)) (by fun_prop)
+      feedback := fun _ ↦ M.R.comap (fun p ↦ (p.1.2, p.2)) (by fun_prop)
+      isMarkovKernel_obs n := by
+        cases n
+        · split_ifs <;> infer_instance
+        · infer_instance
+
+lemma measurable_env [Nonempty 𝓢] (M : MDP 𝓢 𝓐 𝓡) : Measurable M.env := by
+  rw [measurable_environment_iff]
+  refine fun n ↦ ⟨?_, by fun_prop⟩
+  cases n with
+  | zero =>
+    refine Kernel.measurable_const.comp ?_
+    exact Measurable.ite ProbabilityMeasure.measurableSet_isProbabilityMeasure
+      (by fun_prop) (by fun_prop)
+  | succ n => fun_prop
+
+/-! ### Stationary policies and their trajectory laws -/
+
+-- todo: this is a generic definition of a stationary policy, not specific to MDPs.
+/-- The stationary deterministic policy `π` as an algorithm: it plays `π s` in the current
+state `s`. -/
+noncomputable def policyAlg (π : 𝓢 → 𝓐) (hπ : Measurable π) : Algorithm 𝓢 𝓐 𝓡 :=
+  detAlgorithm (fun _ p ↦ π p.2) fun _ ↦ hπ.comp measurable_snd
+
+variable [Nonempty 𝓢]
+
+/-- The law of the trajectory of the policy `π` started at the state `s`: `E^π_s`. -/
+noncomputable def policyMeasure (M : MDP 𝓢 𝓐 𝓡) (π : 𝓢 → 𝓐) (hπ : Measurable π) (s : 𝓢) :
+    Measure (ℕ → Round 𝓢 𝓐 𝓡) :=
+  trajMeasure (policyAlg π hπ) (M.env (Measure.dirac s))
+deriving IsProbabilityMeasure
+
+/-- The law of the state-action trajectory of the policy `π` from the state `s` (no rewards). -/
+noncomputable def stateLaw (M : MDP 𝓢 𝓐 𝓡) (π : 𝓢 → 𝓐) (hπ : Measurable π) (s : 𝓢) :
+    Measure (ℕ → 𝓢 × 𝓐) :=
+  (policyMeasure M π hπ s).map (fun h n ↦ ((h n).obs, (h n).action))
+
+instance (M : MDP 𝓢 𝓐 𝓡) (π : 𝓢 → 𝓐) (hπ : Measurable π) (s : 𝓢) :
+    IsProbabilityMeasure (stateLaw M π hπ s) := Measure.isProbabilityMeasure_map (by fun_prop)
+
+/-! ### Mean rewards -/
+
+
+variable [NormedAddCommGroup 𝓡] [NormedSpace ℝ 𝓡]
+
+/-- The mean reward `r(s, a) = 𝔼[R (s, a)]` of a state-action pair. -/
+noncomputable def meanReward (R : Kernel (𝓢 × 𝓐) 𝓡) [IsMarkovKernel R] (p : 𝓢 × 𝓐) : 𝓡 := (R p)[id]
+
+end Learning.MDP
diff --git a/LeanMachineLearning/SequentialLearning/Algorithm.lean b/LeanMachineLearning/SequentialLearning/Algorithm.lean
index 4003b19d..cf8c80b9 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 𝓞 :=

From 6a548243955bc295d0afc1314906138a84c24eac Mon Sep 17 00:00:00 2001
From: Remy Degenne <remydegenne@gmail.com>
Date: Thu, 10 Sep 2026 17:01:52 +0200
Subject: [PATCH 2/3] fix

---
 LeanMachineLearning/ReinforcementLearning/MDP/Basic.lean | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/LeanMachineLearning/ReinforcementLearning/MDP/Basic.lean b/LeanMachineLearning/ReinforcementLearning/MDP/Basic.lean
index fe0aa460..d2024093 100644
--- a/LeanMachineLearning/ReinforcementLearning/MDP/Basic.lean
+++ b/LeanMachineLearning/ReinforcementLearning/MDP/Basic.lean
@@ -79,7 +79,7 @@ noncomputable def stateLaw (M : MDP 𝓢 𝓐 𝓡) (π : 𝓢 → 𝓐) (hπ :
   (policyMeasure M π hπ s).map (fun h n ↦ ((h n).obs, (h n).action))
 
 instance (M : MDP 𝓢 𝓐 𝓡) (π : 𝓢 → 𝓐) (hπ : Measurable π) (s : 𝓢) :
-    IsProbabilityMeasure (stateLaw M π hπ s) := Measure.isProbabilityMeasure_map (by fun_prop)
+    IsProbabilityMeasure (stateLaw M π hπ s) := by unfold stateLaw; infer_instance
 
 /-! ### Mean rewards -/
 

From 8976e0867abff48e64ede41972e103e4d4034c6b Mon Sep 17 00:00:00 2001
From: Remy Degenne <remydegenne@gmail.com>
Date: Thu, 10 Sep 2026 19:10:33 +0200
Subject: [PATCH 3/3] minor

---
 .../ReinforcementLearning/MDP/Basic.lean       | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/LeanMachineLearning/ReinforcementLearning/MDP/Basic.lean b/LeanMachineLearning/ReinforcementLearning/MDP/Basic.lean
index d2024093..faebfff4 100644
--- a/LeanMachineLearning/ReinforcementLearning/MDP/Basic.lean
+++ b/LeanMachineLearning/ReinforcementLearning/MDP/Basic.lean
@@ -15,7 +15,7 @@ public import LeanMachineLearning.SequentialLearning.Deterministic
 
 @[expose] public section
 
-open MeasureTheory ProbabilityTheory Finset Learning
+open MeasureTheory ProbabilityTheory Learning
 
 /-- Markov decision process with state space `𝓢`, action space `𝓐`, and reward space `𝓡`, described
 by a transition kernel `P : Kernel (𝓢 × 𝓐) 𝓢` and a reward kernel `R : Kernel (𝓢 × 𝓐) 𝓡`.
@@ -38,14 +38,14 @@ instance (M : MDP 𝓢 𝓐 𝓡) : IsMarkovKernel M.R := M.hR
 open Classical in
 protected noncomputable def env [h𝓢 : Nonempty 𝓢] (M : MDP 𝓢 𝓐 𝓡) (μ₀ : Measure 𝓢) :
     Environment 𝓢 𝓐 𝓡 where
-      obs
-      | 0 => Kernel.const _ (if IsProbabilityMeasure μ₀ then μ₀ else Measure.dirac h𝓢.some)
-      | n + 1 => M.P.comap (fun h ↦ ((h (Fin.last n)).obs, (h (Fin.last n)).action)) (by fun_prop)
-      feedback := fun _ ↦ M.R.comap (fun p ↦ (p.1.2, p.2)) (by fun_prop)
-      isMarkovKernel_obs n := by
-        cases n
-        · split_ifs <;> infer_instance
-        · infer_instance
+  obs
+  | 0 => Kernel.const _ (if IsProbabilityMeasure μ₀ then μ₀ else Measure.dirac h𝓢.some)
+  | n + 1 => M.P.comap (fun h ↦ ((h (Fin.last n)).obs, (h (Fin.last n)).action)) (by fun_prop)
+  feedback := fun _ ↦ M.R.comap (fun p ↦ (p.1.2, p.2)) (by fun_prop)
+  isMarkovKernel_obs n := by
+    cases n
+    · split_ifs <;> infer_instance
+    · infer_instance
 
 lemma measurable_env [Nonempty 𝓢] (M : MDP 𝓢 𝓐 𝓡) : Measurable M.env := by
   rw [measurable_environment_iff]