feat: leave-one-out evaluation of the separated plug-in policy [LWCGC24, Lemma 5]

LMLExplore/Analysis/SpecialFunctions/Log/Basic.lean

@@ -50,6 +50,10 @@ import Mathlib.Analysis.SpecialFunctions.Log.Deriv
 * `Real.add_one_mul_log_sub_le_one`: `(m + 1) log (m + 1) - m log m - log (m + 1 / 2) ≤ 1` for
   `m : ℕ`.
 * `Real.log_div_le_log_one_div_add_log`: `log (x / δ) ≤ log (1 / δ) + log x` for `δ ∈ (0, 1]`.
+* `Real.natCeil_log_inv_one_sub_add_one_le`: `⌈log (1 - γ)⁻¹⌉₊ + 1 ≤ 2 log (e / (1 - γ))` for
+  `0 ≤ γ < 1`.
+* `Real.mul_exp_neg_log_div_le`: `c exp (-log (x / δ)) ≤ δ / n` for `c n ≤ x`, the count of
+  confidence levels in union bounds over `n` events.
 * `Real.eventually_log_le_mul_of_isBigO`: if `r n = O(n ^ p)`, then `log (r n) ≤ ε n` for all
   large `n`, for every `ε > 0`.
 * `Real.eventually_exists_add_le_mul_and_forall_ofReal_add_lt_mul`: for `1 < κ c` and
@@ -420,6 +424,27 @@ theorem log_div_le_log_one_div_add_log (x : ℝ) {δ : ℝ} (hδ : 0 < δ) (hδ1
   · rw [log_div hx hδ.ne', one_div, log_inv]
     linarith
 
+/-- For `0 ≤ γ < 1`, `⌈log (1 - γ)⁻¹⌉₊ + 1 ≤ 2 log (e / (1 - γ))`. The left-hand side counts the
+levels `0, …, ⌈log (1 - γ)⁻¹⌉₊` of a recursion run for `log (1 - γ)⁻¹` steps, as in the union
+bounds of model-based policy evaluation with a generative model. -/
+theorem natCeil_log_inv_one_sub_add_one_le {γ : ℝ} (hγ₀ : 0 ≤ γ) (hγ : γ < 1) :
+    (⌈log (1 - γ)⁻¹⌉₊ : ℝ) + 1 ≤ 2 * log (exp 1 / (1 - γ)) := by
+  have hu : 0 < 1 - γ := sub_pos.2 hγ
+  have hy : 0 ≤ log (1 - γ)⁻¹ := log_nonneg ((one_le_inv₀ hu).2 (by linarith))
+  have h := Nat.ceil_lt_add_one hy
+  rw [log_div (exp_pos 1).ne' hu.ne', log_exp]
+  rw [log_inv] at h hy ⊢
+  linarith
+
+/-- For `c, n, δ > 0` with `c n ≤ x`, `c exp (-log (x / δ)) ≤ δ / n`: if each of `n` bad events
+has probability at most `c exp (-L)` at the confidence level `L = log (x / δ)`, their union has
+probability at most `δ`. -/
+theorem mul_exp_neg_log_div_le {c x n δ : ℝ} (hc : 0 < c) (hn : 0 < n) (hδ : 0 < δ)
+    (h : c * n ≤ x) : c * exp (-log (x / δ)) ≤ δ / n := by
+  have hx : 0 < x := (mul_pos hc hn).trans_le h
+  rw [exp_neg, exp_log (div_pos hx hδ), inv_div, ← mul_div_assoc, div_le_div_iff₀ hx hn]
+  nlinarith
+
 /-- If `r n = O(n ^ p)`, then `log (r n) ≤ ε n` for all large `n`, for every `ε > 0`. -/
 theorem eventually_log_le_mul_of_isBigO {r : ℕ → ℝ} {p : ℝ}
     (hr : r =O[atTop] fun n ↦ (n : ℝ) ^ p) {ε : ℝ} (hε : 0 < ε) :

LMLExplore/MarkovDecisionProcess/MarkovPolicy.lean

@@ -34,7 +34,8 @@ policy `π`, these marginals are given by powers of two kernels:
 Integrating a function against `M.stateKernel π s` averages over the actions `a` of `π s` its
 integral against `M.P (s, a)` (`integral_stateKernel`, and `integral_stateKernel_eq_sum` for
 finitely many actions); integrating a function against `M.saKernel π p` integrates over `M.P p`
-its average over the actions of `π` (`integral_saKernel`).
+its average over the actions of `π` (`integral_saKernel`). The state kernel of a deterministic
+stationary policy `d` is `s ↦ M.P (s, d s)` (`stateKernel_deterministic_apply`).
 
 The deterministic stationary policies `policyAlg π hπ` are Markov policies
 (`policyAlg_eq_markov`).
@@ -123,6 +124,14 @@ lemma stateKernel_comp_eq_comp_compProd (μ : Measure 𝓢) [SFinite μ] :
 lemma stateKernel_apply (s : 𝓢) : M.stateKernel π s = M.P ∘ₘ (Measure.dirac s ⊗ₘ π) := by
   rw [← stateKernel_comp_eq_comp_compProd, Measure.dirac_bind (Kernel.measurable _)]
 
+/-- The state kernel of the deterministic stationary policy `d` draws the next state from
+`M.P (s, d s)`. -/
+@[simp]
+lemma stateKernel_deterministic_apply {d : 𝓢 → 𝓐} (hd : Measurable d) (s : 𝓢) :
+    M.stateKernel (Kernel.deterministic d hd) s = M.P (s, d s) := by
+  rw [stateKernel_apply, Measure.compProd_deterministic, Measure.map_dirac' (by fun_prop),
+    Measure.dirac_bind M.P.measurable]
+
 /-- Composing a measure `ν` on state-action pairs with the state-action kernel of `π` draws a
 next state from `M.P ∘ₘ ν`, then an action from `π`. -/
 lemma saKernel_comp_eq_compProd (ν : Measure (𝓢 × 𝓐)) [SFinite ν] :

LMLExplore/MarkovDecisionProcess/ReturnVariance.lean

@@ -33,19 +33,24 @@ model.
 
 For a nonnegative reward `r_π`, the bound can be stated in terms of the value rather than of the
 rewards [LWCGC24, Lemma 11]: if `v^π ≤ B`, the resolvent of the state kernel `P_π` applied to the
-standard deviations `√(Var_{P_π}(v^π))` is at most `B / (γ √(1 - γ))`. The theorems of this
-file apply only to the value `v^π`. The values `V^(l) = R_γ(P_π) √(Var_{P_π}(V^(l-1)))` of the
-successive standard deviations, with `V^(0) = v^π`, are not policy values; the generic form for a
-Markov kernel, `ProbabilityTheory.Kernel.integral_sqrt_variance_integral_resolvent_le`, applied
-with `f = √(Var_{P_π}(V^(l-2)))` (and `f = r_π` for `l = 1`), so that its resolvent is
-`V = V^(l-1)`, bounds each `V^(l)` by `(γ √(1 - γ))⁻¹` times a bound on `V^(l-1)`.
+standard deviations `√(Var_{P_π}(v^π))` is at most `B / (γ √(1 - γ))`. The values
+`V^(l) = R_γ(P_π) √(Var_{P_π}(V^(l-1)))` of the successive standard deviations, with
+`V^(0) = v^π`, are the levels `M.valueLevel γ π l` of [LWCGC24, Lemma 2]; they are not policy
+values. The generic form for a Markov kernel,
+`ProbabilityTheory.Kernel.integral_sqrt_variance_integral_resolvent_le`, applied with
+`f = √(Var_{P_π}(V^(l-2)))` (and `f = r_π` for `l = 1`), so that its resolvent is `V = V^(l-1)`,
+bounds each `V^(l)` by `(γ √(1 - γ))⁻¹` times a bound on `V^(l-1)` [LWCGC24, (73)].
 This recursion is how [LWCGC24] evaluates a fixed policy with a generative model from sample
-sizes of order `(1 - γ)⁻¹`.
+sizes of order `(1 - γ)⁻¹`: if the state kernel `P'_π` of a second MDP `M'` with the same rewards
+deviates from `P_π` on the first `m ≥ log (1 - γ)⁻¹` levels by a Bernstein-type bound, the values
+of `π` in `M` and `M'` are close [LWCGC24, Lemmas 1 and 2].
 
 ## Main definitions
 
 * `Learning.MDP.returnVariance M γ π`: the variance `Σ^π (s, a)` of the discounted return of the
   mean rewards along the state-action chain of `π` started at `(s, a)`.
+* `Learning.MDP.valueLevel M γ π l`: the levels `V 0 = v^π` and
+  `V (l + 1) = R_γ(P_π) √Var_{P_π}(V l)` of [LWCGC24, Lemma 2] for a stationary policy `π`.
 
 ## Main statements
 
@@ -64,6 +69,13 @@ sizes of order `(1 - γ)⁻¹`.
   `Learning.MDP.integral_sqrt_variance_value_stateKernel_resolvent_le`: for `r_π ≥ 0` and
   `v^π ≤ B`, `γ ^ 2 R_γ(P_π) Var_{P_π}(v^π) ≤ 2 B v^π - (v^π) ^ 2` and
   `R_γ(P_π) √(Var_{P_π}(v^π)) ≤ B / (γ √(1 - γ))` [LWCGC24, Lemma 11].
+* `Learning.MDP.valueLevel_mem_Icc`: with rewards in `[0, C]`, the level `l` is in
+  `[0, C / ((1 - γ) (γ √(1 - γ)) ^ l)]` [LWCGC24, (73)].
+* `Learning.MDP.abs_value_sub_value_le_of_abs_integral_stateKernel_sub_le`: the deterministic part
+  of model-based policy evaluation [LWCGC24, Lemmas 1 and 2], for two MDPs with the same rewards,
+  with the deviations in the form given by Bernstein's inequality at a level `L` and `ε = 2 L / N`.
+  It is shared by `ReinforcementLearning/Generative/PolicyEvaluation.lean` [LWCGC24, Lemma 1] and
+  `ReinforcementLearning/Generative/LeaveOneOut.lean` [LWCGC24, Lemma 5].
 
 ## References
 
@@ -72,7 +84,8 @@ sizes of order `(1 - γ)⁻¹`.
 * [M. G. Azar, R. Munos and H. J. Kappen, *Minimax PAC bounds on the sample complexity of
   reinforcement learning with a generative model*, 2013][AMK13]
 * [G. Li, Y. Wei, Y. Chi, Y. Gu and Y. Chen, *Breaking the sample size barrier in model-based
-  reinforcement learning with a generative model*, Operations Research, 2024][LWCGC24], Lemma 11
+  reinforcement learning with a generative model*, Operations Research, 2024][LWCGC24], Lemmas 1,
+  2 and 11 (arXiv:2005.12900)
 -/
 
 public section
@@ -259,3 +272,92 @@ theorem integral_sqrt_variance_value_stateKernel_resolvent_le (hγ₀ : 0 < γ)
     (abs_policyReward_le hR) hr hvB s
 
 end Learning.MDP
+
+/-! ### The levels of [LWCGC24, Lemma 2] -/
+
+namespace Learning.MDP
+
+open Real (log)
+
+variable {𝓢 𝓐 : Type*} {m𝓢 : MeasurableSpace 𝓢} {m𝓐 : MeasurableSpace 𝓐}
+
+/-- The levels of [LWCGC24, Lemma 2] for a stationary policy `π`, with state kernel `P_π`:
+`V 0 = R_γ(P_π) r_π`, which is the value `v^π` for bounded rewards and `γ < 1`, and
+`V (l + 1) = R_γ(P_π) √Var_{P_π}(V l)`. -/
+noncomputable def valueLevel (M : MDP 𝓢 𝓐 ℝ) (γ : ℝ≥0) (π : Kernel 𝓢 𝓐) : ℕ → 𝓢 → ℝ
+  | 0, x => ∫ y, M.policyReward π y ∂(Kernel.resolvent γ (M.stateKernel π) x)
+  | l + 1, x =>
+    ∫ y, √Var[valueLevel M γ π l; M.stateKernel π y] ∂(Kernel.resolvent γ (M.stateKernel π) x)
+
+@[simp]
+lemma valueLevel_zero (M : MDP 𝓢 𝓐 ℝ) (γ : ℝ≥0) (π : Kernel 𝓢 𝓐) (x : 𝓢) :
+    M.valueLevel γ π 0 x = ∫ y, M.policyReward π y ∂(Kernel.resolvent γ (M.stateKernel π) x) :=
+  (rfl)
+
+@[simp]
+lemma valueLevel_succ (M : MDP 𝓢 𝓐 ℝ) (γ : ℝ≥0) (π : Kernel 𝓢 𝓐) (l : ℕ) (x : 𝓢) :
+    M.valueLevel γ π (l + 1) x =
+      ∫ y, √Var[M.valueLevel γ π l; M.stateKernel π y] ∂(Kernel.resolvent γ (M.stateKernel π) x) :=
+  (rfl)
+
+/-- With rewards in `[0, C]` and `0 < γ < 1`, the level `l` of [LWCGC24, Lemma 2] for a stationary
+policy is in `[0, C / ((1 - γ) (γ √(1 - γ)) ^ l)]` [LWCGC24, (73)]. -/
+lemma valueLevel_mem_Icc {M : MDP 𝓢 𝓐 ℝ} {γ : ℝ≥0} (hγ₀ : 0 < γ) (hγ : γ < 1) {C : ℝ}
+    (hR : ∀ p, ∀ᵐ x ∂(M.R p), x ∈ Set.Icc 0 C) (π : Kernel 𝓢 𝓐) [IsMarkovKernel π] (l : ℕ)
+    (x : 𝓢) : M.valueLevel γ π l x ∈ Set.Icc 0 (C / ((1 - γ) * (γ * √(1 - γ)) ^ l)) :=
+  Kernel.mem_Icc_of_eq_integral_resolvent_sqrt_variance hγ₀ hγ measurable_policyReward
+    (fun y ↦ (policyReward_mem_Icc hR y).1) (fun y ↦ (policyReward_mem_Icc hR y).2)
+    (valueLevel_zero M γ π) (valueLevel_succ M γ π) l x
+
+/-- **The deterministic part of model-based policy evaluation** [LWCGC24, Lemmas 1 and 2]. Let
+`M` and `M'` be MDPs with the same rewards, in `[0, C]`, `0 < γ < 1`, and `π` a stationary policy,
+with state kernels `P_π` in `M` and `P'_π` in `M'`. Let `V = M.valueLevel γ π` be the levels
+`V 0 = v^π` and `V (l + 1) = R_γ(P_π) √Var_{P_π}(V l)` of [LWCGC24, Lemma 2], bounded by
+`b l = C / ((1 - γ) (γ √(1 - γ)) ^ l)`, and `log (1 - γ)⁻¹ ≤ m`. Let `L ≥ 0`, `N > 0` and
+`18 L ≤ (1 - γ) N`. If `P'_π` deviates from `P_π` on the levels `l ≤ m` by the bound
+`√(2 Var_{P_π}(V l) L / N) + b l L / (3 N)` of Bernstein's inequality at the level `L`, then the
+value of `π` in `M'` satisfies `|v'^π - v^π| ≤ 3 C √(2 L / (N (1 - γ) ^ 3))`.
+
+This is `ProbabilityTheory.Kernel.abs_integral_resolvent_sub_integral_resolvent_le_of_abs_sub_le`
+with `ε = 2 L / N`, for the resolvent forms `v^π = R_γ(P_π) r_π` and `v'^π = R_γ(P'_π) r_π` of the
+two values. -/
+theorem abs_value_sub_value_le_of_abs_integral_stateKernel_sub_le {M M' : MDP 𝓢 𝓐 ℝ}
+    (hM' : M'.R = M.R) {γ : ℝ≥0} (hγ₀ : 0 < γ) (hγ : γ < 1) {C : ℝ}
+    (hR : ∀ p, ∀ᵐ x ∂(M.R p), x ∈ Set.Icc 0 C) (π : Kernel 𝓢 𝓐) [IsMarkovKernel π] {m : ℕ}
+    (hm : log (1 - (γ : ℝ))⁻¹ ≤ m) {N L : ℝ} (hN : 0 < N) (hL : 0 ≤ L) (hNL : 18 * L ≤ (1 - γ) * N)
+    (h : ∀ l ≤ m, ∀ x,
+      |∫ y, M.valueLevel γ π l y ∂(M'.stateKernel π x) -
+          ∫ y, M.valueLevel γ π l y ∂(M.stateKernel π x)| ≤
+        √(2 * Var[M.valueLevel γ π l; M.stateKernel π x] * L / N) +
+          C / ((1 - γ) * (γ * √(1 - γ)) ^ l) * L / (3 * N)) (s : 𝓢) :
+    |M'.value γ π s - M.value γ π s| ≤ 3 * C * √(2 * L / (N * (1 - γ) ^ 3)) := by
+  have hRabs (p : 𝓢 × 𝓐) : ∀ᵐ x ∂(M.R p), |x| ≤ C :=
+    (hR p).mono fun x hx ↦ (abs_of_nonneg hx.1).trans_le hx.2
+  have hf (y : 𝓢) := policyReward_mem_Icc (π := π) hR y
+  -- both values are resolvents of the same reward `r_π`
+  have hfe : M'.policyReward π = M.policyReward π :=
+    funext fun y ↦ by rw [policyReward_def, hM', policyReward_def]
+  rw [value_eq_integral_resolvent hγ (hM' ▸ hRabs), value_eq_integral_resolvent hγ hRabs, hfe]
+  -- Lemma 2 with `ε = 2 L / N`, for which `9 ε ≤ 1 - γ`
+  have hε9 : 9 * (2 * L / N) ≤ 1 - γ := by
+    rw [← mul_div_assoc, div_le_iff₀ hN]
+    linarith
+  refine (Kernel.abs_integral_resolvent_sub_integral_resolvent_le_of_abs_sub_le
+    (M'.stateKernel π) hγ₀ hγ measurable_policyReward (fun y ↦ (hf y).1) (fun y ↦ (hf y).2)
+    (valueLevel_zero M γ π) (valueLevel_succ M γ π) hm (by positivity) hε9
+    (fun l hl x ↦ (h l hl x).trans (add_le_add (le_of_eq ?_) ?_)) s).trans_eq ?_
+  · congr 1
+    ring
+  · -- the bound `b l` of the level `l` is nonnegative
+    have hb := valueLevel_mem_Icc hγ₀ hγ hR π l x
+    have hbL : 0 ≤ C / ((1 - γ) * (γ * √(1 - γ)) ^ l) * L / N := by
+      have := hb.1.trans hb.2
+      positivity
+    calc C / ((1 - γ) * (γ * √(1 - γ)) ^ l) * L / (3 * N)
+        = C / ((1 - γ) * (γ * √(1 - γ)) ^ l) * L / N / 3 := by
+          rw [div_div, mul_comm (3 : ℝ)]
+      _ ≤ C / ((1 - γ) * (γ * √(1 - γ)) ^ l) * L / N * 2 := by linarith
+      _ = 2 * L / N * (C / ((1 - γ) * (γ * √(1 - γ)) ^ l)) := by ring
+  · rw [div_div]
+
+end Learning.MDP

LMLExplore/MarkovDecisionProcess/Value.lean

@@ -92,6 +92,15 @@ lemma abs_policyReward_le [IsMarkovKernel π] (hR : ∀ p, ∀ᵐ x ∂(M.R p),
     |M.policyReward π s| ≤ C :=
   abs_integral_le_of_abs_le fun a ↦ abs_meanReward_le hR (s, a)
 
+/-- With rewards in `[0, C]`, the reward of a stationary policy is in `[0, C]`. -/
+lemma policyReward_mem_Icc [IsMarkovKernel π] (hR : ∀ p, ∀ᵐ x ∂(M.R p), x ∈ Set.Icc 0 C)
+    (s : 𝓢) : M.policyReward π s ∈ Set.Icc 0 C :=
+  ⟨integral_nonneg fun a ↦ by
+      rw [meanReward_def]
+      exact integral_nonneg_of_ae ((hR _).mono fun _ hx ↦ hx.1),
+    (le_abs_self _).trans
+      (abs_policyReward_le (fun p ↦ (hR p).mono fun _ hx ↦ (abs_of_nonneg hx.1).trans_le hx.2) s)⟩
+
 /-- The reward of a deterministic stationary policy `d` is the mean reward of the action `d s`. -/
 @[simp]
 lemma policyReward_deterministic {d : 𝓢 → 𝓐} (hd : Measurable d) (s : 𝓢) :
@@ -131,6 +140,14 @@ lemma value_def (s : 𝓢) :
     M.value γ π s = ∑' n, (γ : ℝ) ^ n * ∫ s', M.policyReward π s' ∂((M.stateKernel π ^ n) s) :=
   (rfl)
 
+variable (M π) in
+/-- For `γ = 0`, the value of a stationary policy is its reward: `v^π = r_π`. -/
+@[simp]
+lemma value_zero [IsSFiniteKernel π] (s : 𝓢) : M.value 0 π s = M.policyReward π s := by
+  rw [value_def, tsum_eq_single 0 fun n hn ↦ by simp [hn]]
+  simp [Kernel.one_def, Kernel.id_apply,
+    integral_dirac' _ _ measurable_policyReward.stronglyMeasurable]
+
 section Value
 
 variable [IsMarkovKernel π]

LMLExplore/ReinforcementLearning/Generative/Basic.lean

@@ -109,6 +109,12 @@ lemma empirical_congr (M : MDP 𝓢 𝓐 𝓡) (N : ℕ) [NeZero N] {ω ω' : 
   MDP.ext (by rw [empirical_P, empirical_P, empiricalKernel_congr h])
     (by rw [empirical_R, empirical_R])
 
+/-- For `γ = 0`, the value of a stationary policy in the empirical MDP is its value in `M`: both
+are its reward. -/
+lemma value_empirical_zero (M : MDP 𝓢 𝓐 ℝ) (N : ℕ) [NeZero N] (ω : ℕ → 𝓢 × 𝓐 → 𝓢)
+    (π : Kernel 𝓢 𝓐) [IsSFiniteKernel π] : (M.empirical N ω).value 0 π = M.value 0 π :=
+  funext fun s ↦ by rw [value_zero, value_zero, policyReward_def, empirical_R, policyReward_def]
+
 /-- For finite state and action spaces, every function of the empirical MDP is a measurable
 function of the array of samples: the empirical MDP only depends on the finitely many samples
 `ω i p`, `i < N`, which take finitely many values. -/