feat: the policy-mixture lemma for occupancy measures [AJKS, Lemma 14.1]

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`.
@@ -42,20 +50,25 @@ from its marginals `M.stateMarginal π μ n ⊗ₘ π n`.
   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
   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
@@ -204,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`. -/
@@ -255,4 +276,99 @@ 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 [occupancy_fst, occupancy_fst, tvDist_add_smul_add_smul]
+  gcongr
+  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
+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 π
+
+/-- **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}
+
+/-- **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γ 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
+`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γ fun s ↦ hπ' s ▸ tvDist_one_sub_smul_add_smul_le hα
+
+end Perturbation
+
 end Learning.MDP

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`. 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
@@ -34,16 +41,29 @@ 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_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 μ ν`.
+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`. 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
 
 open Set
+open scoped NNReal
 
 namespace MeasureTheory
 
@@ -109,4 +129,68 @@ 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)
+  rw [← ENNReal.ofReal_le_ofReal_iff (by positivity), ofReal_tvDist,
+    ENNReal.ofReal_mul zero_le_two, ENNReal.ofReal_ofNat]
+  exact he
+
+/-! ### Sums and mixtures -/
+
+/-- 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]
+  simp [← mul_sub, abs_mul]
+
+/-- **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_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