From 01a1320e0bade6a9d123df45b349a5cb193cab4c Mon Sep 17 00:00:00 2001 From: Remy Degenne Date: Thu, 5 Nov 2026 08:03:00 +0100 Subject: [PATCH 1/4] feat: the policy-mixture lemma for occupancy measures [AJKS, Lemma 14.1] MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stationary policies within ε of each other in total variation at every state have occupancy measures within ε / (1 - γ), and state marginals within γ ε / (1 - γ). For a mixture (1 - α) π + α π'' this gives the policy-mixture lemma, with bounds 2 α γ / (1 - γ) and 2 α / (1 - γ). General tools: eventwise characterizations of tvDist for measures of the same mass, tvDist of sums and mixtures, and the data processing inequality for tvDist under Markov kernels and composition-products. Co-Authored-By: Claude Code --- .../MarkovDecisionProcess/Occupancy.lean | 124 ++++++++++++++++- .../Measure/TotalVariation/Basic.lean | 75 +++++++++- .../Probability/Kernel/TotalVariation.lean | 130 ++++++++++++++++++ 3 files changed, 326 insertions(+), 3 deletions(-) create mode 100644 LMLExplore/Probability/Kernel/TotalVariation.lean diff --git a/LMLExplore/MarkovDecisionProcess/Occupancy.lean b/LMLExplore/MarkovDecisionProcess/Occupancy.lean index 814a0bdb..ad64fe65 100644 --- a/LMLExplore/MarkovDecisionProcess/Occupancy.lean +++ b/LMLExplore/MarkovDecisionProcess/Occupancy.lean @@ -7,6 +7,7 @@ module public import LMLExplore.MarkovDecisionProcess.RunValue public import LMLExplore.Probability.Kernel.Resolvent +public import LMLExplore.Probability.Kernel.TotalVariation /-! # Discounted occupancy measures of a Markov decision process @@ -29,6 +30,13 @@ The occupancy measure depends on the run only through the laws of the pairs `(S can be computed along every run `IsAlgEnvSeq S A Y alg (M.env μ) P`, and for a Markov policy `π` from its marginals `M.stateMarginal π μ n ⊗ₘ π n`. +Occupancy measures depend continuously on stationary policies: if `tvDist (π' s) (π s) ≤ ε` at +every state, then the state marginals of the occupancy measures of `π'` and `π` are within +`γ * ε / (1 - γ)` in total variation, and the occupancy measures themselves within `ε / (1 - γ)`. +For the mixture `π' = (1 - α) π + α π''` this is the policy-mixture lemma, with the bounds +`2 * α * γ / (1 - γ)` and `2 * α / (1 - γ)`: a step of size `α` towards another policy moves +the occupancy measure by `O(α)`, the estimate behind the analysis of conservative policy updates. + ## Main definitions * `Learning.MDP.occupancy M γ alg μ`: the discounted state-action occupancy measure of `alg`. @@ -49,13 +57,18 @@ from its marginals `M.stateMarginal π μ n ⊗ₘ π n`. policy, it is the image of the initial state-action law by the resolvent of the state-action kernel, scaled by `1 - γ`; `Learning.MDP.isCondKernel_occupancy_markov_const`: the occupancy measure of a stationary policy `π` disintegrates as its state marginal composed with `π`. +* `Learning.MDP.tvDist_fst_occupancy_le`, `Learning.MDP.tvDist_occupancy_le`: stationary + policies within `ε` of each other in total variation at every state have occupancy measures + within `ε / (1 - γ)`, and state marginals within `γ * ε / (1 - γ)`; +* `Learning.MDP.tvDist_fst_occupancy_le_of_mixture`, `Learning.MDP.tvDist_occupancy_le_of_mixture`: + the policy-mixture lemma [AJKS, Lemma 14.1]. ## References * [P] M. L. Puterman, *Markov Decision Processes: Discrete Stochastic Dynamic Programming*, Wiley, 1994, §6.9. * [AJKS] A. Agarwal, N. Jiang, S. M. Kakade, W. Sun, *Reinforcement Learning: Theory and - Algorithms*, 2022, §1.5. + Algorithms*, 2022, §1.5 and Lemma 14.1. -/ public section @@ -255,4 +268,113 @@ instance isCondKernel_occupancy_markov_const {π : Kernel 𝓢 𝓐} [IsMarkovKe end Markov +/-! ### Occupancy measures of close stationary policies -/ + +section Perturbation + +variable {μ : Measure 𝓢} [IsProbabilityMeasure μ] {π π' : Kernel 𝓢 𝓐} [IsMarkovKernel π] + [IsMarkovKernel π'] {ε : ℝ} + +/-- The flow equation contracts total variation: the state marginals of the occupancy measures +of two stationary policies are `γ` times closer than the occupancy measures themselves. -/ +private lemma tvDist_fst_occupancy_le_mul (hγ : γ < 1) : + tvDist (M.occupancy γ (Algorithm.markov fun _ ↦ π') μ).fst + (M.occupancy γ (Algorithm.markov fun _ ↦ π) μ).fst ≤ + γ * tvDist (M.occupancy γ (Algorithm.markov fun _ ↦ π') μ) + (M.occupancy γ (Algorithm.markov fun _ ↦ π) μ) := by + have := isProbabilityMeasure_occupancy (M := M) (alg := Algorithm.markov fun _ ↦ π') (μ := μ) hγ + have := isProbabilityMeasure_occupancy (M := M) (alg := Algorithm.markov fun _ ↦ π) (μ := μ) hγ + -- both marginals are `(1 - γ) • μ + γ • (M.P ∘ₘ x)` for their occupancy measure `x` + rw [Measure.fst, Measure.fst, occupancy_map_fst, occupancy_map_fst, ← ENNReal.coe_one, + ← ENNReal.coe_sub, Measure.coe_nnreal_smul, Measure.coe_nnreal_smul, Measure.coe_nnreal_smul] + refine (tvDist_add_smul_le (by rw [Measure.comp_apply_univ, Measure.comp_apply_univ, + measure_univ, measure_univ]) γ).trans ?_ + gcongr + exact tvDist_comp_right_le (by simp) M.P + +/-- The occupancy measures of two stationary policies `π'` and `π` with +`tvDist (π' s) (π s) ≤ ε` for every state `s` are within `ε` plus the distance of their state +marginals. -/ +private lemma tvDist_occupancy_le_add (hγ : γ < 1) (hε : ∀ s, tvDist (π' s) (π s) ≤ ε) : + tvDist (M.occupancy γ (Algorithm.markov fun _ ↦ π') μ) + (M.occupancy γ (Algorithm.markov fun _ ↦ π) μ) ≤ + ε + tvDist (M.occupancy γ (Algorithm.markov fun _ ↦ π') μ).fst + (M.occupancy γ (Algorithm.markov fun _ ↦ π) μ).fst := by + have := isProbabilityMeasure_occupancy (M := M) (alg := Algorithm.markov fun _ ↦ π') (μ := μ) hγ + have := isProbabilityMeasure_occupancy (M := M) (alg := Algorithm.markov fun _ ↦ π) (μ := μ) hγ + -- each occupancy measure disintegrates along its policy + conv_lhs => rw [← Measure.IsCondKernel.disintegrate + (ρ := M.occupancy γ (Algorithm.markov fun _ ↦ π') μ) (ρCond := π'), + ← Measure.IsCondKernel.disintegrate + (ρ := M.occupancy γ (Algorithm.markov fun _ ↦ π) μ) (ρCond := π)] + calc _ ≤ tvDist ((M.occupancy γ (Algorithm.markov fun _ ↦ π') μ).fst ⊗ₘ π') + ((M.occupancy γ (Algorithm.markov fun _ ↦ π') μ).fst ⊗ₘ π) + + tvDist ((M.occupancy γ (Algorithm.markov fun _ ↦ π') μ).fst ⊗ₘ π) + ((M.occupancy γ (Algorithm.markov fun _ ↦ π) μ).fst ⊗ₘ π) := tvDist_triangle _ _ _ + _ ≤ _ := by + gcongr + · simpa using tvDist_compProd_right_le + (μ := (M.occupancy γ (Algorithm.markov fun _ ↦ π') μ).fst) (ae_of_all _ hε) + · exact tvDist_compProd_left_le (by simp) π + +/-- **Occupancy measures of close policies**: if the stationary policies `π'` and `π` satisfy +`tvDist (π' s) (π s) ≤ ε` at every state `s`, then for `γ < 1` the state marginals of their +occupancy measures from `μ` are within `γ * ε / (1 - γ)` in total variation. -/ +theorem tvDist_fst_occupancy_le (hγ : γ < 1) (hε : ∀ s, tvDist (π' s) (π s) ≤ ε) : + tvDist (M.occupancy γ (Algorithm.markov fun _ ↦ π') μ).fst + (M.occupancy γ (Algorithm.markov fun _ ↦ π) μ).fst ≤ γ * ε / (1 - γ) := by + have h1 := tvDist_fst_occupancy_le_mul (M := M) (π := π) (π' := π') (μ := μ) hγ + have h2 := tvDist_occupancy_le_add (M := M) (μ := μ) hγ hε + have hγ' : (γ : ℝ) < 1 := by exact_mod_cast hγ + rw [le_div_iff₀ (sub_pos.2 hγ')] + nlinarith [γ.coe_nonneg] + +/-- **Occupancy measures of close policies**: if the stationary policies `π'` and `π` satisfy +`tvDist (π' s) (π s) ≤ ε` at every state `s`, then for `γ < 1` their state-action occupancy +measures from `μ` are within `ε / (1 - γ)` in total variation. -/ +theorem tvDist_occupancy_le (hγ : γ < 1) (hε : ∀ s, tvDist (π' s) (π s) ≤ ε) : + tvDist (M.occupancy γ (Algorithm.markov fun _ ↦ π') μ) + (M.occupancy γ (Algorithm.markov fun _ ↦ π) μ) ≤ ε / (1 - γ) := by + have h1 := tvDist_fst_occupancy_le_mul (M := M) (π := π) (π' := π') (μ := μ) hγ + have h2 := tvDist_occupancy_le_add (M := M) (μ := μ) hγ hε + have hγ' : (γ : ℝ) < 1 := by exact_mod_cast hγ + rw [le_div_iff₀ (sub_pos.2 hγ')] + nlinarith [γ.coe_nonneg] + +variable {π'' : Kernel 𝓢 𝓐} [IsMarkovKernel π''] {α : ℝ≥0} + +omit [Nonempty 𝓢] [IsMarkovKernel π'] in +/-- A mixture `(1 - α) • π s + α • π'' s` of stationary policies is within `2 * α` of `π s` in +total variation. -/ +private lemma tvDist_mixture_le (hα : α ≤ 1) (hπ' : ∀ s, π' s = (1 - α) • π s + α • π'' s) + (s : 𝓢) : tvDist (π' s) (π s) ≤ 2 * α := by + rw [hπ'] + refine (tvDist_smul_add_smul_le (by simp) hα).trans ?_ + rw [mul_comm] + gcongr + simpa [one_add_one_eq_two] using tvDist_le_add (μ := π'' s) (ν := π s) + +/-- **The policy-mixture lemma** [AJKS, Lemma 14.1], for state marginals: for `γ < 1` and +`α ≤ 1`, the mixture `π'` of the stationary policies `π` and `π''` with weights `1 - α` and `α` +has an occupancy measure from `μ` whose state marginal is within `2 * α * γ / (1 - γ)` of that +of `π` in total variation. -/ +theorem tvDist_fst_occupancy_le_of_mixture (hγ : γ < 1) (hα : α ≤ 1) + (hπ' : ∀ s, π' s = (1 - α) • π s + α • π'' s) : + tvDist (M.occupancy γ (Algorithm.markov fun _ ↦ π') μ).fst + (M.occupancy γ (Algorithm.markov fun _ ↦ π) μ).fst ≤ 2 * α * γ / (1 - γ) := by + rw [show (2 * α * γ : ℝ) = γ * (2 * α) by ring] + exact tvDist_fst_occupancy_le hγ (tvDist_mixture_le hα hπ') + +/-- **The policy-mixture lemma** [AJKS, Lemma 14.1], for state-action occupancy measures: for +`γ < 1` and `α ≤ 1`, the mixture `π'` of the stationary policies `π` and `π''` with weights +`1 - α` and `α` has an occupancy measure from `μ` within `2 * α / (1 - γ)` of that of `π` in +total variation. -/ +theorem tvDist_occupancy_le_of_mixture (hγ : γ < 1) (hα : α ≤ 1) + (hπ' : ∀ s, π' s = (1 - α) • π s + α • π'' s) : + tvDist (M.occupancy γ (Algorithm.markov fun _ ↦ π') μ) + (M.occupancy γ (Algorithm.markov fun _ ↦ π) μ) ≤ 2 * α / (1 - γ) := + tvDist_occupancy_le hγ (tvDist_mixture_le hα hπ') + +end Perturbation + end Learning.MDP diff --git a/LMLExplore/MeasureTheory/Measure/TotalVariation/Basic.lean b/LMLExplore/MeasureTheory/Measure/TotalVariation/Basic.lean index e6f29d6c..a3ebf5f6 100644 --- a/LMLExplore/MeasureTheory/Measure/TotalVariation/Basic.lean +++ b/LMLExplore/MeasureTheory/Measure/TotalVariation/Basic.lean @@ -8,8 +8,10 @@ module public import Mathlib.MeasureTheory.Integral.Bochner.Basic public import Mathlib.MeasureTheory.Measure.TotalVariation.Defs +import Mathlib.MeasureTheory.Integral.Bochner.Set import Mathlib.MeasureTheory.Integral.IntegrableOn import Mathlib.MeasureTheory.VectorMeasure.Integral +import TauCeti.MeasureTheory.Measure.TotalVariation /-! # Integrals against measures close in total variation @@ -18,7 +20,12 @@ The total variation distance `tvDist μ ν` of two finite measures (with values probability measures) controls the difference of the integrals of bounded functions: if `‖f‖ ≤ C`, then `‖∫ f dμ - ∫ f dν‖ ≤ C * tvDist μ ν`. When `μ` and `ν` have the same mass, the bound depends only on the oscillation of `f`: if `f` takes values in `[a, b]`, then -`|∫ f dμ - ∫ f dν| ≤ (b - a) / 2 * tvDist μ ν`. +`|∫ f dμ - ∫ f dν| ≤ (b - a) / 2 * tvDist μ ν`. For measures of the same mass, the total +variation distance is also controlled by events: `|μ s - ν s| ≤ tvDist μ ν / 2` for every +measurable `s`, and conversely `tvDist μ ν ≤ 2 * c` as soon as `μ s ≤ ν s + c` for every +measurable `s`. From these, adding a common finite measure and scaling both measures by `t` +multiplies the distance by at most `t`, which bounds the distance from a mixture +`(1 - α) • μ + α • ν` to `μ` by `α * tvDist ν μ`. These are the bounds through which errors in transition kernels, measured in total variation, propagate to expected values, as in the simulation lemma for Markov decision processes and in the @@ -34,16 +41,27 @@ contraction of the oscillation of functions by a Markov kernel (Dobrushin's coef * `MeasureTheory.abs_integral_sub_integral_le_mul_tvDist_of_mem_Icc`: for measures of the same mass and `f` with values in `[a, b]` almost everywhere, `|∫ f dμ - ∫ f dν| ≤ (b - a) / 2 * tvDist μ ν`. +* `MeasureTheory.abs_measureReal_sub_measureReal_le_tvDist_div_two`, + `MeasureTheory.tvDist_le_two_mul_of_forall_measureReal_le_add`: for measures of the same mass, + the total variation distance bounds the difference of the masses of every event, and eventwise + bounds bound the total variation distance. +* `MeasureTheory.tvDist_add_smul_le`: `tvDist (ξ + t • μ) (ξ + t • ν) ≤ t * tvDist μ ν` for + measures `μ` and `ν` of the same mass. +* `MeasureTheory.tvDist_smul_add_smul_le`: the mixture `(1 - α) • μ + α • ν` is within + `α * tvDist ν μ` of `μ`. ## Implementation notes The first bound is the bound `MeasureTheory.VectorMeasure.norm_integral_le_of_norm_le_const` on -integrals against the signed measure `μ - ν`, whose variation has mass `tvDist μ ν`. +integrals against the signed measure `μ - ν`, whose variation has mass `tvDist μ ν`. The eventwise +characterization is Tau Ceti's `MeasureTheory.etvDist_le_two_mul_iff`, transferred to the real +distance `tvDist`. -/ public section open Set +open scoped NNReal namespace MeasureTheory @@ -109,4 +127,57 @@ theorem abs_integral_sub_integral_le_mul_tvDist_of_mem_Icc [IsFiniteMeasure μ] integral_sub (Integrable.of_bound hfν _ (hfν'.mono fun _ ↦ hB)) (integrable_const _)] at h simpa [measureReal_def, hμν] using h +/-! ### Eventwise bounds -/ + +/-- **Events of measures of the same mass close in total variation**: for finite measures `μ` +and `ν` of the same mass, every measurable event `s` satisfies +`|μ s - ν s| ≤ tvDist μ ν / 2`. -/ +theorem abs_measureReal_sub_measureReal_le_tvDist_div_two [IsFiniteMeasure μ] + [IsFiniteMeasure ν] (hμν : μ univ = ν univ) {s : Set 𝓧} (hs : MeasurableSet s) : + |μ.real s - ν.real s| ≤ tvDist μ ν / 2 := by + have h := abs_integral_sub_integral_le_mul_tvDist_of_mem_Icc hμν (f := s.indicator 1) + (a := 0) (b := 1) (by fun_prop (disch := exact hs)) (by fun_prop (disch := exact hs)) + (ae_of_all _ fun x ↦ by by_cases hx : x ∈ s <;> simp [hx]) + (ae_of_all _ fun x ↦ by by_cases hx : x ∈ s <;> simp [hx]) + rw [integral_indicator_one hs, integral_indicator_one hs] at h + exact h.trans_eq (by ring) + +/-- **Total variation from eventwise bounds**: for finite measures `μ` and `ν` of the same mass, +if every measurable event `s` satisfies `μ s ≤ ν s + c`, then `tvDist μ ν ≤ 2 * c`. -/ +theorem tvDist_le_two_mul_of_forall_measureReal_le_add [IsFiniteMeasure μ] [IsFiniteMeasure ν] + (hμν : μ univ = ν univ) {c : ℝ} (h : ∀ s, MeasurableSet s → μ.real s ≤ ν.real s + c) : + tvDist μ ν ≤ 2 * c := by + have hc : 0 ≤ c := by simpa using h ∅ MeasurableSet.empty + have he : etvDist μ ν ≤ 2 * ENNReal.ofReal c := by + refine (etvDist_le_two_mul_iff hμν).2 fun s hs ↦ ?_ + rw [← ofReal_measureReal (μ := μ), ← ofReal_measureReal (μ := ν), + ← ENNReal.ofReal_add measureReal_nonneg hc] + exact ENNReal.ofReal_le_ofReal (h s hs) + calc tvDist μ ν = (etvDist μ ν).toReal := (rfl) + _ ≤ (2 * ENNReal.ofReal c).toReal := ENNReal.toReal_mono (by finiteness) he + _ = 2 * c := by simp [ENNReal.toReal_ofReal hc] + +/-! ### Sums and mixtures -/ + +/-- Adding a common finite measure to two finite measures `μ` and `ν` of the same mass, after +scaling both by `t`, multiplies their total variation distance by at most `t`: +`tvDist (ξ + t • μ) (ξ + t • ν) ≤ t * tvDist μ ν`. -/ +theorem tvDist_add_smul_le {ξ : Measure 𝓧} [IsFiniteMeasure ξ] [IsFiniteMeasure μ] + [IsFiniteMeasure ν] (hμν : μ univ = ν univ) (t : ℝ≥0) : + tvDist (ξ + t • μ) (ξ + t • ν) ≤ t * tvDist μ ν := by + rw [← mul_div_cancel_left₀ (t * tvDist μ ν) (two_ne_zero' ℝ), mul_div_assoc] + refine tvDist_le_two_mul_of_forall_measureReal_le_add (by simp [hμν]) fun s hs ↦ ?_ + have h := (abs_le.1 (abs_measureReal_sub_measureReal_le_tvDist_div_two hμν hs)).2 + rw [measureReal_add_apply, measureReal_add_apply, measureReal_nnreal_smul_apply, + measureReal_nnreal_smul_apply] + linear_combination mul_le_mul_of_nonneg_left h t.coe_nonneg + +/-- **Total variation of a mixture**: for finite measures `μ` and `ν` of the same mass and +`α ≤ 1`, the mixture `(1 - α) • μ + α • ν` is within `α * tvDist ν μ` of `μ`. -/ +theorem tvDist_smul_add_smul_le [IsFiniteMeasure μ] [IsFiniteMeasure ν] + (hμν : ν univ = μ univ) {α : ℝ≥0} (hα : α ≤ 1) : + tvDist ((1 - α) • μ + α • ν) μ ≤ α * tvDist ν μ := by + conv_lhs => enter [2]; rw [← one_smul ℝ≥0 μ, ← tsub_add_cancel_of_le hα, add_smul] + exact tvDist_add_smul_le hμν α + end MeasureTheory diff --git a/LMLExplore/Probability/Kernel/TotalVariation.lean b/LMLExplore/Probability/Kernel/TotalVariation.lean new file mode 100644 index 00000000..a014d3bb --- /dev/null +++ b/LMLExplore/Probability/Kernel/TotalVariation.lean @@ -0,0 +1,130 @@ +/- +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.Composition.MeasureCompProd +public import LMLExplore.MeasureTheory.Measure.TotalVariation.Basic + +import Mathlib.MeasureTheory.Integral.IntegrableOn +import Mathlib.Probability.Kernel.Composition.MeasureComp + +/-! +# Total variation distance under Markov kernels + +Composition with a Markov kernel does not increase the total variation distance between two +measures of the same mass (the data processing inequality for the total variation distance), and +replacing the kernel of a composition-product `μ ⊗ₘ κ` by a kernel `η` with +`tvDist (κ x) (η x) ≤ ε` for almost every `x` moves the measure by at most `μ univ * ε`. + +These are the two ways a perturbation propagates along a Markov chain: through the law of the +current state, and through the transition kernel. Together they bound the distance between the +laws of Markov chains with close kernels, such as the occupancy measures of two close policies of +a Markov decision process. + +## Main statements + +* `MeasureTheory.tvDist_comp_right_le`: `tvDist (κ ∘ₘ μ) (κ ∘ₘ ν) ≤ tvDist μ ν` for a Markov + kernel `κ` and finite measures `μ` and `ν` of the same mass; +* `MeasureTheory.tvDist_compProd_left_le`: `tvDist (μ ⊗ₘ κ) (ν ⊗ₘ κ) ≤ tvDist μ ν`, under the + same hypotheses; +* `MeasureTheory.tvDist_compProd_right_le`: `tvDist (μ ⊗ₘ κ) (μ ⊗ₘ η) ≤ μ.real univ * ε` for + Markov kernels `κ` and `η` with `tvDist (κ x) (η x) ≤ ε` for `μ`-almost every `x`. + +## Implementation notes + +All three bounds are proved eventwise, through +`MeasureTheory.tvDist_le_two_mul_of_forall_measureReal_le_add`: the mass of an event under a +composition is the integral of a function with values in `[0, 1]`, to which +`MeasureTheory.abs_integral_sub_integral_le_mul_tvDist_of_mem_Icc` applies. +-/ + +public section + +open Set ProbabilityTheory +open scoped ENNReal + +namespace MeasureTheory + +variable {α β : Type*} {mα : MeasurableSpace α} {mβ : MeasurableSpace β} {μ ν : Measure α} + +/-- The mass of an event `s` under the composition of a measure `μ` with a Markov kernel `κ`, +as the integral of `x ↦ κ x s`. -/ +private lemma measureReal_comp_eq_integral (κ : Kernel α β) [IsMarkovKernel κ] {s : Set β} + (hs : MeasurableSet s) : + (κ ∘ₘ μ).real s = ∫ x, (κ x).real s ∂μ := by + simp only [measureReal_def] + rw [Measure.bind_apply hs κ.aemeasurable, + integral_toReal (κ.measurable_coe hs).aemeasurable (ae_of_all _ fun _ ↦ measure_lt_top _ _)] + +/-- The mass of an event `s` under the composition-product of a finite measure `μ` with a Markov +kernel `κ`, as the integral of `x ↦ κ x (Prod.mk x ⁻¹' s)`. -/ +private lemma measureReal_compProd_eq_integral [IsFiniteMeasure μ] (κ : Kernel α β) + [IsMarkovKernel κ] {s : Set (α × β)} (hs : MeasurableSet s) : + (μ ⊗ₘ κ).real s = ∫ x, (κ x).real (Prod.mk x ⁻¹' s) ∂μ := by + simp only [measureReal_def] + rw [Measure.compProd_apply hs, + integral_toReal (Kernel.measurable_kernel_prodMk_left hs).aemeasurable + (ae_of_all _ fun _ ↦ measure_lt_top _ _)] + +/-- **Data processing inequality** for the total variation distance: composing two finite +measures `μ` and `ν` of the same mass with a Markov kernel `κ` does not increase their total +variation distance. -/ +theorem tvDist_comp_right_le [IsFiniteMeasure μ] [IsFiniteMeasure ν] (hμν : μ univ = ν univ) + (κ : Kernel α β) [IsMarkovKernel κ] : + tvDist (κ ∘ₘ μ) (κ ∘ₘ ν) ≤ tvDist μ ν := by + rw [← mul_div_cancel₀ (tvDist μ ν) (two_ne_zero' ℝ)] + refine tvDist_le_two_mul_of_forall_measureReal_le_add + (by rw [Measure.comp_apply_univ, Measure.comp_apply_univ, hμν]) fun s hs ↦ ?_ + rw [measureReal_comp_eq_integral κ hs, measureReal_comp_eq_integral κ hs] + -- `x ↦ κ x s` takes values in `[0, 1]` + have hf : Measurable fun x ↦ (κ x).real s := (κ.measurable_coe hs).ennreal_toReal + have h01 : ∀ x, (κ x).real s ∈ Icc 0 1 := fun _ ↦ ⟨measureReal_nonneg, measureReal_le_one⟩ + have h := abs_integral_sub_integral_le_mul_tvDist_of_mem_Icc hμν (a := 0) (b := 1) + (f := fun x ↦ (κ x).real s) hf.aestronglyMeasurable hf.aestronglyMeasurable + (ae_of_all _ h01) (ae_of_all _ h01) + linarith [(abs_le.1 h).2] + +/-- Composition-products with a common Markov kernel `κ` do not increase the total variation +distance between finite measures `μ` and `ν` of the same mass. -/ +theorem tvDist_compProd_left_le [IsFiniteMeasure μ] [IsFiniteMeasure ν] (hμν : μ univ = ν univ) + (κ : Kernel α β) [IsMarkovKernel κ] : + tvDist (μ ⊗ₘ κ) (ν ⊗ₘ κ) ≤ tvDist μ ν := by + rw [Measure.compProd_eq_comp_prod, Measure.compProd_eq_comp_prod] + exact tvDist_comp_right_le hμν _ + +/-- Replacing the kernel `κ` of the composition-product `μ ⊗ₘ κ` by a Markov kernel `η` with +`tvDist (κ x) (η x) ≤ ε` for `μ`-almost every `x` moves it by at most `μ.real univ * ε` in total +variation. -/ +theorem tvDist_compProd_right_le [IsFiniteMeasure μ] {κ η : Kernel α β} [IsMarkovKernel κ] + [IsMarkovKernel η] {ε : ℝ} (hε : ∀ᵐ x ∂μ, tvDist (κ x) (η x) ≤ ε) : + tvDist (μ ⊗ₘ κ) (μ ⊗ₘ η) ≤ μ.real univ * ε := by + rw [← mul_div_cancel₀ (μ.real univ * ε) (two_ne_zero' ℝ)] + refine tvDist_le_two_mul_of_forall_measureReal_le_add + (by rw [Measure.compProd_apply_univ, Measure.compProd_apply_univ]) fun s hs ↦ ?_ + rw [measureReal_compProd_eq_integral κ hs, measureReal_compProd_eq_integral η hs] + have hκ : Measurable fun x ↦ (κ x).real (Prod.mk x ⁻¹' s) := + (Kernel.measurable_kernel_prodMk_left hs).ennreal_toReal + have hη : Measurable fun x ↦ (η x).real (Prod.mk x ⁻¹' s) := + (Kernel.measurable_kernel_prodMk_left hs).ennreal_toReal + have hint {f : α → ℝ} (hf : Measurable f) (hf1 : ∀ x, f x ∈ Icc 0 1) : Integrable f μ := + Integrable.of_bound hf.aestronglyMeasurable 1 (ae_of_all _ fun x ↦ by + rw [Real.norm_of_nonneg (hf1 x).1]; exact (hf1 x).2) + have hκ1 (x : α) : (κ x).real (Prod.mk x ⁻¹' s) ∈ Icc 0 1 := + ⟨measureReal_nonneg, measureReal_le_one⟩ + have hη1 (x : α) : (η x).real (Prod.mk x ⁻¹' s) ∈ Icc 0 1 := + ⟨measureReal_nonneg, measureReal_le_one⟩ + -- the integrands differ by at most `ε / 2` almost everywhere + have hle : ∫ x, (κ x).real (Prod.mk x ⁻¹' s) ∂μ ≤ + ∫ x, ((η x).real (Prod.mk x ⁻¹' s) + ε / 2) ∂μ := by + refine integral_mono_ae (hint hκ hκ1) ((hint hη hη1).add (integrable_const _)) ?_ + filter_upwards [hε] with x hx + have h := (abs_le.1 (abs_measureReal_sub_measureReal_le_tvDist_div_two (μ := κ x) (ν := η x) + (by simp) (measurable_prodMk_left (x := x) hs))).2 + linarith + rw [integral_add (hint hη hη1) (integrable_const _), integral_const, smul_eq_mul] at hle + linarith + +end MeasureTheory From f7802f6a09260e6120adbbd9cdcce795284ff9f5 Mon Sep 17 00:00:00 2001 From: Remy Degenne Date: Thu, 5 Nov 2026 08:57:52 +0100 Subject: [PATCH 2/4] refactor: drop equal-mass hypotheses from the tvDist sum and kernel bounds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - `tvDist_add_smul_add_smul` (equality, via finite partitions) and `tvDist_one_sub_smul_add_smul` replace the `_le` versions and need no equal-mass hypothesis; add `tvDist_one_sub_smul_add_smul_le` (≤ 2α for probability measures), replacing the private MDP helper. - `tvDist_comp_right_le` / `tvDist_compProd_left_le` hold for all finite measures, proved on a Hahn decomposition. - Make `Measure.measureReal_bind_eq_integral` and `Measure.measureReal_compProd_eq_integral` public, for finite kernels. - Add the flow equation `occupancy_fst` with `ℝ≥0` weights and use it. Co-Authored-By: Claude Code --- .../MarkovDecisionProcess/Occupancy.lean | 36 ++++----- .../Measure/TotalVariation/Basic.lean | 65 +++++++++------ .../Probability/Kernel/TotalVariation.lean | 81 ++++++++++++------- 3 files changed, 104 insertions(+), 78 deletions(-) diff --git a/LMLExplore/MarkovDecisionProcess/Occupancy.lean b/LMLExplore/MarkovDecisionProcess/Occupancy.lean index ad64fe65..86845f1b 100644 --- a/LMLExplore/MarkovDecisionProcess/Occupancy.lean +++ b/LMLExplore/MarkovDecisionProcess/Occupancy.lean @@ -50,8 +50,8 @@ the occupancy measure by `O(α)`, the estimate behind the analysis of conservati probability measure; * `Learning.MDP.runValue_eq_integral_occupancy`: for `γ < 1` and bounded real rewards, the discounted value of `alg` is `(1 - γ)⁻¹ * ∫ r d(M.occupancy γ alg μ)`; -* `Learning.MDP.occupancy_map_fst`, `Learning.MDP.occupancy_prod_univ`: the flow equation, for a - probability measure `μ`; +* `Learning.MDP.occupancy_map_fst`, `Learning.MDP.occupancy_fst`, + `Learning.MDP.occupancy_prod_univ`: the flow equation, for a probability measure `μ`; * `Learning.MDP.occupancy_markov`, `Learning.MDP.occupancy_markov_const`: the occupancy measures of Markov and stationary policies, in terms of their state-action marginals; for a stationary policy, it is the image of the initial state-action law by the resolvent of the state-action @@ -217,6 +217,14 @@ theorem occupancy_map_fst : simp_rw [pow_succ', mul_assoc, ENNReal.tsum_mul_left] ring +/-- **The flow equation** of occupancy measures, for the state marginal `Measure.fst` with +`ℝ≥0` weights: the state marginal of the occupancy measure `x` of every algorithm from the +probability measure `μ` is `(1 - γ) • μ + γ • (M.P ∘ₘ x)`. -/ +theorem occupancy_fst : + (M.occupancy γ alg μ).fst = (1 - γ) • μ + γ • (M.P ∘ₘ M.occupancy γ alg μ) := by + ext B hB + simp [Measure.fst, occupancy_map_fst, ENNReal.coe_sub] + /-- **The flow equation** of occupancy measures, setwise: the occupancy measure `x` of every algorithm from the probability measure `μ` satisfies, for every measurable set `B` of states, `x (B ×ˢ univ) = (1 - γ) * μ B + γ * ∫⁻ p, M.P p B ∂x`. -/ @@ -285,12 +293,9 @@ private lemma tvDist_fst_occupancy_le_mul (hγ : γ < 1) : have := isProbabilityMeasure_occupancy (M := M) (alg := Algorithm.markov fun _ ↦ π') (μ := μ) hγ have := isProbabilityMeasure_occupancy (M := M) (alg := Algorithm.markov fun _ ↦ π) (μ := μ) hγ -- both marginals are `(1 - γ) • μ + γ • (M.P ∘ₘ x)` for their occupancy measure `x` - rw [Measure.fst, Measure.fst, occupancy_map_fst, occupancy_map_fst, ← ENNReal.coe_one, - ← ENNReal.coe_sub, Measure.coe_nnreal_smul, Measure.coe_nnreal_smul, Measure.coe_nnreal_smul] - refine (tvDist_add_smul_le (by rw [Measure.comp_apply_univ, Measure.comp_apply_univ, - measure_univ, measure_univ]) γ).trans ?_ + rw [occupancy_fst, occupancy_fst, tvDist_add_smul_add_smul] gcongr - exact tvDist_comp_right_le (by simp) M.P + exact tvDist_comp_right_le M.P /-- The occupancy measures of two stationary policies `π'` and `π` with `tvDist (π' s) (π s) ≤ ε` for every state `s` are within `ε` plus the distance of their state @@ -315,7 +320,7 @@ private lemma tvDist_occupancy_le_add (hγ : γ < 1) (hε : ∀ s, tvDist (π' s gcongr · simpa using tvDist_compProd_right_le (μ := (M.occupancy γ (Algorithm.markov fun _ ↦ π') μ).fst) (ae_of_all _ hε) - · exact tvDist_compProd_left_le (by simp) π + · exact tvDist_compProd_left_le π /-- **Occupancy measures of close policies**: if the stationary policies `π'` and `π` satisfy `tvDist (π' s) (π s) ≤ ε` at every state `s`, then for `γ < 1` the state marginals of their @@ -343,17 +348,6 @@ theorem tvDist_occupancy_le (hγ : γ < 1) (hε : ∀ s, tvDist (π' s) (π s) variable {π'' : Kernel 𝓢 𝓐} [IsMarkovKernel π''] {α : ℝ≥0} -omit [Nonempty 𝓢] [IsMarkovKernel π'] in -/-- A mixture `(1 - α) • π s + α • π'' s` of stationary policies is within `2 * α` of `π s` in -total variation. -/ -private lemma tvDist_mixture_le (hα : α ≤ 1) (hπ' : ∀ s, π' s = (1 - α) • π s + α • π'' s) - (s : 𝓢) : tvDist (π' s) (π s) ≤ 2 * α := by - rw [hπ'] - refine (tvDist_smul_add_smul_le (by simp) hα).trans ?_ - rw [mul_comm] - gcongr - simpa [one_add_one_eq_two] using tvDist_le_add (μ := π'' s) (ν := π s) - /-- **The policy-mixture lemma** [AJKS, Lemma 14.1], for state marginals: for `γ < 1` and `α ≤ 1`, the mixture `π'` of the stationary policies `π` and `π''` with weights `1 - α` and `α` has an occupancy measure from `μ` whose state marginal is within `2 * α * γ / (1 - γ)` of that @@ -363,7 +357,7 @@ theorem tvDist_fst_occupancy_le_of_mixture (hγ : γ < 1) (hα : α ≤ 1) tvDist (M.occupancy γ (Algorithm.markov fun _ ↦ π') μ).fst (M.occupancy γ (Algorithm.markov fun _ ↦ π) μ).fst ≤ 2 * α * γ / (1 - γ) := by rw [show (2 * α * γ : ℝ) = γ * (2 * α) by ring] - exact tvDist_fst_occupancy_le hγ (tvDist_mixture_le hα hπ') + exact tvDist_fst_occupancy_le hγ fun s ↦ hπ' s ▸ tvDist_one_sub_smul_add_smul_le hα /-- **The policy-mixture lemma** [AJKS, Lemma 14.1], for state-action occupancy measures: for `γ < 1` and `α ≤ 1`, the mixture `π'` of the stationary policies `π` and `π''` with weights @@ -373,7 +367,7 @@ theorem tvDist_occupancy_le_of_mixture (hγ : γ < 1) (hα : α ≤ 1) (hπ' : ∀ s, π' s = (1 - α) • π s + α • π'' s) : tvDist (M.occupancy γ (Algorithm.markov fun _ ↦ π') μ) (M.occupancy γ (Algorithm.markov fun _ ↦ π) μ) ≤ 2 * α / (1 - γ) := - tvDist_occupancy_le hγ (tvDist_mixture_le hα hπ') + tvDist_occupancy_le hγ fun s ↦ hπ' s ▸ tvDist_one_sub_smul_add_smul_le hα end Perturbation diff --git a/LMLExplore/MeasureTheory/Measure/TotalVariation/Basic.lean b/LMLExplore/MeasureTheory/Measure/TotalVariation/Basic.lean index a3ebf5f6..a1c07024 100644 --- a/LMLExplore/MeasureTheory/Measure/TotalVariation/Basic.lean +++ b/LMLExplore/MeasureTheory/Measure/TotalVariation/Basic.lean @@ -23,9 +23,9 @@ bound depends only on the oscillation of `f`: if `f` takes values in `[a, b]`, t `|∫ f dμ - ∫ f dν| ≤ (b - a) / 2 * tvDist μ ν`. For measures of the same mass, the total variation distance is also controlled by events: `|μ s - ν s| ≤ tvDist μ ν / 2` for every measurable `s`, and conversely `tvDist μ ν ≤ 2 * c` as soon as `μ s ≤ ν s + c` for every -measurable `s`. From these, adding a common finite measure and scaling both measures by `t` -multiplies the distance by at most `t`, which bounds the distance from a mixture -`(1 - α) • μ + α • ν` to `μ` by `α * tvDist ν μ`. +measurable `s`. Adding a common finite measure and scaling both measures by `t` multiplies the +distance by `t`, so the distance from a mixture `(1 - α) • μ + α • ν` to `μ` is +`α * tvDist ν μ`, at most `2 * α` for probability measures. These are the bounds through which errors in transition kernels, measured in total variation, propagate to expected values, as in the simulation lemma for Markov decision processes and in the @@ -45,17 +45,19 @@ contraction of the oscillation of functions by a Markov kernel (Dobrushin's coef `MeasureTheory.tvDist_le_two_mul_of_forall_measureReal_le_add`: for measures of the same mass, the total variation distance bounds the difference of the masses of every event, and eventwise bounds bound the total variation distance. -* `MeasureTheory.tvDist_add_smul_le`: `tvDist (ξ + t • μ) (ξ + t • ν) ≤ t * tvDist μ ν` for - measures `μ` and `ν` of the same mass. -* `MeasureTheory.tvDist_smul_add_smul_le`: the mixture `(1 - α) • μ + α • ν` is within - `α * tvDist ν μ` of `μ`. +* `MeasureTheory.tvDist_add_smul_add_smul`: `tvDist (ξ + t • μ) (ξ + t • ν) = t * tvDist μ ν`. +* `MeasureTheory.tvDist_one_sub_smul_add_smul`, `MeasureTheory.tvDist_one_sub_smul_add_smul_le`: + the mixture `(1 - α) • μ + α • ν` is at distance `α * tvDist ν μ` from `μ`, at most `2 * α` + for probability measures. ## Implementation notes The first bound is the bound `MeasureTheory.VectorMeasure.norm_integral_le_of_norm_le_const` on integrals against the signed measure `μ - ν`, whose variation has mass `tvDist μ ν`. The eventwise characterization is Tau Ceti's `MeasureTheory.etvDist_le_two_mul_iff`, transferred to the real -distance `tvDist`. +distance `tvDist`. The behaviour under sums and scalings is read off the supremum over finite +measurable partitions, `MeasureTheory.tvDist_eq_iSup_finPartition_abs`, so it needs no assumption +on the masses. -/ public section @@ -159,25 +161,36 @@ theorem tvDist_le_two_mul_of_forall_measureReal_le_add [IsFiniteMeasure μ] [IsF /-! ### Sums and mixtures -/ -/-- Adding a common finite measure to two finite measures `μ` and `ν` of the same mass, after -scaling both by `t`, multiplies their total variation distance by at most `t`: -`tvDist (ξ + t • μ) (ξ + t • ν) ≤ t * tvDist μ ν`. -/ -theorem tvDist_add_smul_le {ξ : Measure 𝓧} [IsFiniteMeasure ξ] [IsFiniteMeasure μ] - [IsFiniteMeasure ν] (hμν : μ univ = ν univ) (t : ℝ≥0) : - tvDist (ξ + t • μ) (ξ + t • ν) ≤ t * tvDist μ ν := by - rw [← mul_div_cancel_left₀ (t * tvDist μ ν) (two_ne_zero' ℝ), mul_div_assoc] - refine tvDist_le_two_mul_of_forall_measureReal_le_add (by simp [hμν]) fun s hs ↦ ?_ - have h := (abs_le.1 (abs_measureReal_sub_measureReal_le_tvDist_div_two hμν hs)).2 +/-- Adding a common finite measure to two finite measures `μ` and `ν`, after scaling both by +`t`, multiplies their total variation distance by `t`: +`tvDist (ξ + t • μ) (ξ + t • ν) = t * tvDist μ ν`. -/ +theorem tvDist_add_smul_add_smul {ξ : Measure 𝓧} [IsFiniteMeasure ξ] [IsFiniteMeasure μ] + [IsFiniteMeasure ν] (t : ℝ≥0) : + tvDist (ξ + t • μ) (ξ + t • ν) = t * tvDist μ ν := by + rw [tvDist_eq_iSup_finPartition_abs, tvDist_eq_iSup_finPartition_abs, + Real.mul_iSup_of_nonneg t.coe_nonneg] + -- on every event, the differences of masses are scaled by `t` + refine iSup_congr fun P ↦ ?_ + rw [Finset.mul_sum] + refine Finset.sum_congr rfl fun p _ ↦ ?_ rw [measureReal_add_apply, measureReal_add_apply, measureReal_nnreal_smul_apply, - measureReal_nnreal_smul_apply] - linear_combination mul_le_mul_of_nonneg_left h t.coe_nonneg - -/-- **Total variation of a mixture**: for finite measures `μ` and `ν` of the same mass and -`α ≤ 1`, the mixture `(1 - α) • μ + α • ν` is within `α * tvDist ν μ` of `μ`. -/ -theorem tvDist_smul_add_smul_le [IsFiniteMeasure μ] [IsFiniteMeasure ν] - (hμν : ν univ = μ univ) {α : ℝ≥0} (hα : α ≤ 1) : - tvDist ((1 - α) • μ + α • ν) μ ≤ α * tvDist ν μ := by + measureReal_nnreal_smul_apply, add_sub_add_left_eq_sub, ← mul_sub, abs_mul, NNReal.abs_eq] + +/-- **Total variation of a mixture**: for finite measures `μ` and `ν` and `α ≤ 1`, the mixture +`(1 - α) • μ + α • ν` is at distance `α * tvDist ν μ` from `μ`. -/ +theorem tvDist_one_sub_smul_add_smul [IsFiniteMeasure μ] [IsFiniteMeasure ν] {α : ℝ≥0} + (hα : α ≤ 1) : + tvDist ((1 - α) • μ + α • ν) μ = α * tvDist ν μ := by conv_lhs => enter [2]; rw [← one_smul ℝ≥0 μ, ← tsub_add_cancel_of_le hα, add_smul] - exact tvDist_add_smul_le hμν α + exact tvDist_add_smul_add_smul α + +/-- **Total variation of a mixture of probability measures**: for probability measures `μ` and +`ν` and `α ≤ 1`, the mixture `(1 - α) • μ + α • ν` is within `2 * α` of `μ`. -/ +theorem tvDist_one_sub_smul_add_smul_le [IsProbabilityMeasure μ] [IsProbabilityMeasure ν] + {α : ℝ≥0} (hα : α ≤ 1) : + tvDist ((1 - α) • μ + α • ν) μ ≤ 2 * α := by + rw [tvDist_one_sub_smul_add_smul hα, mul_comm] + gcongr + simpa [one_add_one_eq_two] using tvDist_le_add (μ := ν) (ν := μ) end MeasureTheory diff --git a/LMLExplore/Probability/Kernel/TotalVariation.lean b/LMLExplore/Probability/Kernel/TotalVariation.lean index a014d3bb..585927b2 100644 --- a/LMLExplore/Probability/Kernel/TotalVariation.lean +++ b/LMLExplore/Probability/Kernel/TotalVariation.lean @@ -9,13 +9,14 @@ public import Mathlib.Probability.Kernel.Composition.MeasureCompProd public import LMLExplore.MeasureTheory.Measure.TotalVariation.Basic import Mathlib.MeasureTheory.Integral.IntegrableOn +import Mathlib.MeasureTheory.Measure.TotalVariation.Basic import Mathlib.Probability.Kernel.Composition.MeasureComp /-! # Total variation distance under Markov kernels Composition with a Markov kernel does not increase the total variation distance between two -measures of the same mass (the data processing inequality for the total variation distance), and +finite measures (the data processing inequality for the total variation distance), and replacing the kernel of a composition-product `μ ⊗ₘ κ` by a kernel `η` with `tvDist (κ x) (η x) ≤ ε` for almost every `x` moves the measure by at most `μ univ * ε`. @@ -26,8 +27,11 @@ a Markov decision process. ## Main statements +* `MeasureTheory.Measure.measureReal_bind_eq_integral`, + `MeasureTheory.Measure.measureReal_compProd_eq_integral`: the mass of an event under `κ ∘ₘ μ` + or `μ ⊗ₘ κ`, as an integral against `μ`; * `MeasureTheory.tvDist_comp_right_le`: `tvDist (κ ∘ₘ μ) (κ ∘ₘ ν) ≤ tvDist μ ν` for a Markov - kernel `κ` and finite measures `μ` and `ν` of the same mass; + kernel `κ` and finite measures `μ` and `ν`; * `MeasureTheory.tvDist_compProd_left_le`: `tvDist (μ ⊗ₘ κ) (ν ⊗ₘ κ) ≤ tvDist μ ν`, under the same hypotheses; * `MeasureTheory.tvDist_compProd_right_le`: `tvDist (μ ⊗ₘ κ) (μ ⊗ₘ η) ≤ μ.real univ * ε` for @@ -35,10 +39,11 @@ a Markov decision process. ## Implementation notes -All three bounds are proved eventwise, through -`MeasureTheory.tvDist_le_two_mul_of_forall_measureReal_le_add`: the mass of an event under a -composition is the integral of a function with values in `[0, 1]`, to which -`MeasureTheory.abs_integral_sub_integral_le_mul_tvDist_of_mem_Icc` applies. +The mass of an event under a composition is the integral of a function with values in +`[0, 1]`. The data processing inequality evaluates the distance on a Hahn decomposition `s` of +the two compositions, where it is the gap between the integrals of `x ↦ κ x sᶜ - κ x s`, bounded +by `MeasureTheory.norm_integral_sub_integral_le_mul_tvDist`. The bound for a change of kernel is +proved eventwise, through `MeasureTheory.tvDist_le_two_mul_of_forall_measureReal_le_add`. -/ public section @@ -50,19 +55,19 @@ namespace MeasureTheory variable {α β : Type*} {mα : MeasurableSpace α} {mβ : MeasurableSpace β} {μ ν : Measure α} -/-- The mass of an event `s` under the composition of a measure `μ` with a Markov kernel `κ`, +/-- The mass of an event `s` under the composition of a measure `μ` with a finite kernel `κ`, as the integral of `x ↦ κ x s`. -/ -private lemma measureReal_comp_eq_integral (κ : Kernel α β) [IsMarkovKernel κ] {s : Set β} +theorem Measure.measureReal_bind_eq_integral (κ : Kernel α β) [IsFiniteKernel κ] {s : Set β} (hs : MeasurableSet s) : (κ ∘ₘ μ).real s = ∫ x, (κ x).real s ∂μ := by simp only [measureReal_def] rw [Measure.bind_apply hs κ.aemeasurable, integral_toReal (κ.measurable_coe hs).aemeasurable (ae_of_all _ fun _ ↦ measure_lt_top _ _)] -/-- The mass of an event `s` under the composition-product of a finite measure `μ` with a Markov -kernel `κ`, as the integral of `x ↦ κ x (Prod.mk x ⁻¹' s)`. -/ -private lemma measureReal_compProd_eq_integral [IsFiniteMeasure μ] (κ : Kernel α β) - [IsMarkovKernel κ] {s : Set (α × β)} (hs : MeasurableSet s) : +/-- The mass of an event `s` under the composition-product of a measure `μ` with a finite kernel +`κ`, as the integral of `x ↦ κ x (Prod.mk x ⁻¹' s)`. -/ +theorem Measure.measureReal_compProd_eq_integral [SFinite μ] (κ : Kernel α β) [IsFiniteKernel κ] + {s : Set (α × β)} (hs : MeasurableSet s) : (μ ⊗ₘ κ).real s = ∫ x, (κ x).real (Prod.mk x ⁻¹' s) ∂μ := by simp only [measureReal_def] rw [Measure.compProd_apply hs, @@ -70,30 +75,43 @@ private lemma measureReal_compProd_eq_integral [IsFiniteMeasure μ] (κ : Kernel (ae_of_all _ fun _ ↦ measure_lt_top _ _)] /-- **Data processing inequality** for the total variation distance: composing two finite -measures `μ` and `ν` of the same mass with a Markov kernel `κ` does not increase their total -variation distance. -/ -theorem tvDist_comp_right_le [IsFiniteMeasure μ] [IsFiniteMeasure ν] (hμν : μ univ = ν univ) - (κ : Kernel α β) [IsMarkovKernel κ] : +measures `μ` and `ν` with a Markov kernel `κ` does not increase their total variation +distance. -/ +theorem tvDist_comp_right_le [IsFiniteMeasure μ] [IsFiniteMeasure ν] (κ : Kernel α β) + [IsMarkovKernel κ] : tvDist (κ ∘ₘ μ) (κ ∘ₘ ν) ≤ tvDist μ ν := by - rw [← mul_div_cancel₀ (tvDist μ ν) (two_ne_zero' ℝ)] - refine tvDist_le_two_mul_of_forall_measureReal_le_add - (by rw [Measure.comp_apply_univ, Measure.comp_apply_univ, hμν]) fun s hs ↦ ?_ - rw [measureReal_comp_eq_integral κ hs, measureReal_comp_eq_integral κ hs] - -- `x ↦ κ x s` takes values in `[0, 1]` - have hf : Measurable fun x ↦ (κ x).real s := (κ.measurable_coe hs).ennreal_toReal - have h01 : ∀ x, (κ x).real s ∈ Icc 0 1 := fun _ ↦ ⟨measureReal_nonneg, measureReal_le_one⟩ - have h := abs_integral_sub_integral_le_mul_tvDist_of_mem_Icc hμν (a := 0) (b := 1) - (f := fun x ↦ (κ x).real s) hf.aestronglyMeasurable hf.aestronglyMeasurable - (ae_of_all _ h01) (ae_of_all _ h01) + -- on a Hahn set `s` of the two compositions, the distance is the gap between the integrals of + -- `f x = κ x sᶜ - κ x s`, which has absolute value at most `1` + obtain ⟨s, hs⟩ := exists_isHahnDecomposition (κ ∘ₘ μ) (κ ∘ₘ ν) + have hm := hs.measurableSet + have hf : Measurable fun x ↦ (κ x).real sᶜ - (κ x).real s := + (κ.measurable_coe hm.compl).ennreal_toReal.sub (κ.measurable_coe hm).ennreal_toReal + have hf1 (x : α) : ‖(κ x).real sᶜ - (κ x).real s‖ ≤ 1 := by + rw [Real.norm_eq_abs, abs_le] + constructor <;> linarith [measureReal_nonneg (μ := κ x) (s := s), + measureReal_nonneg (μ := κ x) (s := sᶜ), measureReal_le_one (μ := κ x) (s := s), + measureReal_le_one (μ := κ x) (s := sᶜ)] + have hint (ρ : Measure α) [IsFiniteMeasure ρ] {t : Set β} (ht : MeasurableSet t) : + Integrable (fun x ↦ (κ x).real t) ρ := + Integrable.of_bound (κ.measurable_coe ht).ennreal_toReal.aestronglyMeasurable 1 + (ae_of_all _ fun x ↦ by + rw [Real.norm_of_nonneg measureReal_nonneg]; exact measureReal_le_one) + have h := norm_integral_sub_integral_le_mul_tvDist hf.aestronglyMeasurable + hf.aestronglyMeasurable (ae_of_all μ hf1) (ae_of_all ν hf1) + rw [integral_sub (hint μ hm.compl) (hint μ hm), integral_sub (hint ν hm.compl) (hint ν hm), + one_mul, Real.norm_eq_abs] at h + rw [hs.tvDist_eq, Measure.measureReal_bind_eq_integral κ hm, + Measure.measureReal_bind_eq_integral κ hm, Measure.measureReal_bind_eq_integral κ hm.compl, + Measure.measureReal_bind_eq_integral κ hm.compl] linarith [(abs_le.1 h).2] /-- Composition-products with a common Markov kernel `κ` do not increase the total variation -distance between finite measures `μ` and `ν` of the same mass. -/ -theorem tvDist_compProd_left_le [IsFiniteMeasure μ] [IsFiniteMeasure ν] (hμν : μ univ = ν univ) - (κ : Kernel α β) [IsMarkovKernel κ] : +distance between finite measures `μ` and `ν`. -/ +theorem tvDist_compProd_left_le [IsFiniteMeasure μ] [IsFiniteMeasure ν] (κ : Kernel α β) + [IsMarkovKernel κ] : tvDist (μ ⊗ₘ κ) (ν ⊗ₘ κ) ≤ tvDist μ ν := by rw [Measure.compProd_eq_comp_prod, Measure.compProd_eq_comp_prod] - exact tvDist_comp_right_le hμν _ + exact tvDist_comp_right_le _ /-- Replacing the kernel `κ` of the composition-product `μ ⊗ₘ κ` by a Markov kernel `η` with `tvDist (κ x) (η x) ≤ ε` for `μ`-almost every `x` moves it by at most `μ.real univ * ε` in total @@ -104,7 +122,8 @@ theorem tvDist_compProd_right_le [IsFiniteMeasure μ] {κ η : Kernel α β} [Is rw [← mul_div_cancel₀ (μ.real univ * ε) (two_ne_zero' ℝ)] refine tvDist_le_two_mul_of_forall_measureReal_le_add (by rw [Measure.compProd_apply_univ, Measure.compProd_apply_univ]) fun s hs ↦ ?_ - rw [measureReal_compProd_eq_integral κ hs, measureReal_compProd_eq_integral η hs] + rw [Measure.measureReal_compProd_eq_integral κ hs, + Measure.measureReal_compProd_eq_integral η hs] have hκ : Measurable fun x ↦ (κ x).real (Prod.mk x ⁻¹' s) := (Kernel.measurable_kernel_prodMk_left hs).ennreal_toReal have hη : Measurable fun x ↦ (η x).real (Prod.mk x ⁻¹' s) := From a55c7239544a6dc4a9c0c0a2b80381186e440dd8 Mon Sep 17 00:00:00 2001 From: Remy Degenne Date: Thu, 5 Nov 2026 09:13:46 +0100 Subject: [PATCH 3/4] refactor: shorten two tvDist proofs per proof-quality review Use Mathlib's ofReal_tvDist instead of unfolding tvDist by rfl, and close the sum-and-mixture identity with simp instead of a long named rewrite chain. Co-Authored-By: Claude Code --- .../MeasureTheory/Measure/TotalVariation/Basic.lean | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/LMLExplore/MeasureTheory/Measure/TotalVariation/Basic.lean b/LMLExplore/MeasureTheory/Measure/TotalVariation/Basic.lean index a1c07024..7dfdb7e3 100644 --- a/LMLExplore/MeasureTheory/Measure/TotalVariation/Basic.lean +++ b/LMLExplore/MeasureTheory/Measure/TotalVariation/Basic.lean @@ -155,9 +155,9 @@ theorem tvDist_le_two_mul_of_forall_measureReal_le_add [IsFiniteMeasure μ] [IsF rw [← ofReal_measureReal (μ := μ), ← ofReal_measureReal (μ := ν), ← ENNReal.ofReal_add measureReal_nonneg hc] exact ENNReal.ofReal_le_ofReal (h s hs) - calc tvDist μ ν = (etvDist μ ν).toReal := (rfl) - _ ≤ (2 * ENNReal.ofReal c).toReal := ENNReal.toReal_mono (by finiteness) he - _ = 2 * c := by simp [ENNReal.toReal_ofReal hc] + rw [← ENNReal.ofReal_le_ofReal_iff (by positivity), ofReal_tvDist, + ENNReal.ofReal_mul zero_le_two, ENNReal.ofReal_ofNat] + exact he /-! ### Sums and mixtures -/ @@ -173,8 +173,8 @@ theorem tvDist_add_smul_add_smul {ξ : Measure 𝓧} [IsFiniteMeasure ξ] [IsFin refine iSup_congr fun P ↦ ?_ rw [Finset.mul_sum] refine Finset.sum_congr rfl fun p _ ↦ ?_ - rw [measureReal_add_apply, measureReal_add_apply, measureReal_nnreal_smul_apply, - measureReal_nnreal_smul_apply, add_sub_add_left_eq_sub, ← mul_sub, abs_mul, NNReal.abs_eq] + rw [measureReal_add_apply, measureReal_add_apply] + simp [← mul_sub, abs_mul] /-- **Total variation of a mixture**: for finite measures `μ` and `ν` and `α ≤ 1`, the mixture `(1 - α) • μ + α • ν` is at distance `α * tvDist ν μ` from `μ`. -/ From f37bdd14ea2d279c21a376e6789d89846648033a Mon Sep 17 00:00:00 2001 From: Remy Degenne Date: Thu, 5 Nov 2026 09:26:32 +0100 Subject: [PATCH 4/4] refactor: use Mathlib's IsMarkovKernel.integrable in tvDist_comp_right_le Co-Authored-By: Claude Code --- LMLExplore/Probability/Kernel/TotalVariation.lean | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/LMLExplore/Probability/Kernel/TotalVariation.lean b/LMLExplore/Probability/Kernel/TotalVariation.lean index 585927b2..f1aa1a03 100644 --- a/LMLExplore/Probability/Kernel/TotalVariation.lean +++ b/LMLExplore/Probability/Kernel/TotalVariation.lean @@ -11,6 +11,7 @@ public import LMLExplore.MeasureTheory.Measure.TotalVariation.Basic import Mathlib.MeasureTheory.Integral.IntegrableOn import Mathlib.MeasureTheory.Measure.TotalVariation.Basic import Mathlib.Probability.Kernel.Composition.MeasureComp +import Mathlib.Probability.Kernel.Integral /-! # Total variation distance under Markov kernels @@ -91,15 +92,12 @@ theorem tvDist_comp_right_le [IsFiniteMeasure μ] [IsFiniteMeasure ν] (κ : Ker constructor <;> linarith [measureReal_nonneg (μ := κ x) (s := s), measureReal_nonneg (μ := κ x) (s := sᶜ), measureReal_le_one (μ := κ x) (s := s), measureReal_le_one (μ := κ x) (s := sᶜ)] - have hint (ρ : Measure α) [IsFiniteMeasure ρ] {t : Set β} (ht : MeasurableSet t) : - Integrable (fun x ↦ (κ x).real t) ρ := - Integrable.of_bound (κ.measurable_coe ht).ennreal_toReal.aestronglyMeasurable 1 - (ae_of_all _ fun x ↦ by - rw [Real.norm_of_nonneg measureReal_nonneg]; exact measureReal_le_one) have h := norm_integral_sub_integral_le_mul_tvDist hf.aestronglyMeasurable hf.aestronglyMeasurable (ae_of_all μ hf1) (ae_of_all ν hf1) - rw [integral_sub (hint μ hm.compl) (hint μ hm), integral_sub (hint ν hm.compl) (hint ν hm), - one_mul, Real.norm_eq_abs] at h + rw [integral_sub (Kernel.IsMarkovKernel.integrable μ κ hm.compl) + (Kernel.IsMarkovKernel.integrable μ κ hm), + integral_sub (Kernel.IsMarkovKernel.integrable ν κ hm.compl) + (Kernel.IsMarkovKernel.integrable ν κ hm), one_mul, Real.norm_eq_abs] at h rw [hs.tvDist_eq, Measure.measureReal_bind_eq_integral κ hm, Measure.measureReal_bind_eq_integral κ hm, Measure.measureReal_bind_eq_integral κ hm.compl, Measure.measureReal_bind_eq_integral κ hm.compl]