Skip to content

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

Merged
lmlexplore-automation[bot] merged 4 commits into
mainfrom
roadmap/policy-mixture-lemma-worker1
Nov 5, 2026
Merged

lmlexplore-automation[bot] merged 4 commits into
mainfrom
roadmap/policy-mixture-lemma-worker1

Conversation

@RemyDegenne

Copy link
Copy Markdown
Contributor

This PR proves the policy-mixture lemma [AJKS, Lemma 14.1], the last open item of Layer 12.2 (Corollaries) of the ReinforcementLearning roadmap: for γ < 1, α ≤ 1 and stationary policies with π' s = (1 - α) • π s + α • π'' s at every state, the state marginals of the occupancy measures of π' and π from μ are within 2 * α * γ / (1 - γ) in total variation (Learning.MDP.tvDist_fst_occupancy_le_of_mixture), and the state-action occupancy measures within 2 * α / (1 - γ) (Learning.MDP.tvDist_occupancy_le_of_mixture). Both follow from a general continuity statement, proved on measurable spaces: stationary policies with tvDist (π' s) (π s) ≤ ε at every state have occupancy measures within ε / (1 - γ), and state marginals within γ * ε / (1 - γ) (Learning.MDP.tvDist_occupancy_le, Learning.MDP.tvDist_fst_occupancy_le). The proof runs through the flow equation occupancy_map_fst and the disintegration isCondKernel_occupancy_markov_const that are already on main: the marginal distance is at most γ times the occupancy distance, which is at most ε plus the marginal distance.

Roadmap: ReinforcementLearning

Milestone: Layer 12.2, "The policy-mixture lemma [AJKS, Lemma 14.1]". After this PR every listed item of Layer 12.2 is on main or in an open PR (the pessimistic decomposition is #152). The consumer is Layer 17.8 (conservative policy iteration): its improvement bound V^{π_{t+1}}_μ - V^{π_t}_μ ≥ α/(1-γ) (𝔸_t - 2αγ/(1-γ)²) combines the performance difference lemma with tvDist_fst_occupancy_le_of_mixture. What remains for 17.8 is the CPI algorithm and its oracle; it can build the mixture kernel itself and pass it through the hπ' hypothesis.

The general tools go where their subject lives:

  • LMLExplore/MeasureTheory/Measure/TotalVariation/Basic.lean: for measures of the same mass, |μ s - ν s| ≤ tvDist μ ν / 2 (abs_measureReal_sub_measureReal_le_tvDist_div_two) and its converse (tvDist_le_two_mul_of_forall_measureReal_le_add); tvDist (ξ + t • μ) (ξ + t • ν) ≤ t * tvDist μ ν (tvDist_add_smul_le); and the mixture bound tvDist ((1 - α) • μ + α • ν) μ ≤ α * tvDist ν μ (tvDist_smul_add_smul_le).
  • new LMLExplore/Probability/Kernel/TotalVariation.lean: the data processing inequality tvDist (κ ∘ₘ μ) (κ ∘ₘ ν) ≤ tvDist μ ν for a Markov kernel (tvDist_comp_right_le, named after Mathlib's klDiv_comp_right_le), its composition-product form tvDist_compProd_left_le, and tvDist_compProd_right_le: tvDist (μ ⊗ₘ κ) (μ ⊗ₘ η) ≤ μ.real univ * ε when tvDist (κ x) (η x) ≤ ε almost everywhere.

Reuse: the eventwise characterization is Tau Ceti's MeasureTheory.etvDist_le_two_mul_iff (already compiled at the pinned version), moved to the real tvDist. The integral bound is LMLExplore's abs_integral_sub_integral_le_mul_tvDist_of_mem_Icc. Neither Mathlib, Tau Ceti nor LML has total-variation bounds for kernel compositions at the pinned versions (searched for tvDist/etvDist together with comp, bind, compProd and map; Mathlib's VectorMeasure.variation_map_le covers only deterministic maps). The data-processing bounds assume equal masses because their proofs are eventwise; every use here is between probability measures. The Dobrushin contraction of Layer 1.6 will sharpen tvDist_comp_right_le by the factor δ(κ).

🤖 Prepared with Claude Code

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 <noreply@github.com>
@lmlexplore-automation lmlexplore-automation Bot added awaiting-CI CI has not yet reported on the latest commit awaiting-review CI is green; waiting for review verdicts and removed awaiting-CI CI has not yet reported on the latest commit labels Nov 5, 2026
@RemyDegenne

RemyDegenne commented Nov 5, 2026 •

Copy link
Copy Markdown
Contributor Author

AI review — approved

Each rubric is judged independently by multiple review agents; the PR merges only once every rubric is green — any rubric that is not green (changes requested, blocked, errored, stale, or not yet run) blocks the merge. See the rubrics.

rubric state judge summary
✅ correctness approved claude/claude-opus-5-5 Mathlib's tvDist is the full variation norm of μ - ν (it is etvDist, taken as a real number), so the bounds 2αγ/(1-γ) and 2α/(1-γ) match AJKS Lemma 14.1 as stated in the L1 norm. The general lemmas carry the right hypotheses (equal masses where the proof works event by event, Markov kernels, α ≤ 1, γ < 1), and none is vacuous or hides its content in an assumption.
✅ reuse approved claude/claude-opus-5-5 The new tvDist bounds (eventwise, scaling/mixture, Markov-kernel composition and compProd) have no counterpart in Mathlib's TotalVariation lemmas or in LMLExplore, and they reuse the existing integral bound, IsHahnDecomposition.tvDist_eq, tvDist_eq_iSup_finPartition_abs, and the occupancy flow and disintegration results. occupancy_fst is a ℝ≥0-weighted, Measure.fst restatement of occupancy_map_fst, and it is the form tvDist_add_smul_add_smul needs.
✅ scope approved claude/claude-opus-5-5 The PR proves the ReinforcementLearning Layer 12.2 item 'The policy-mixture lemma [AJKS, Lemma 14.1]' (roadmap README.md:1008) with both stated bounds. Every new total-variation and kernel lemma is used in that proof, so the PR is one topic.
✅ attribution approved claude/claude-opus-5-5 The policy-mixture lemma credits [AJKS, Lemma 14.1] in its docstrings, main-statements list and References, and the reuse of Tau Ceti's etvDist_le_two_mul_iff is credited in the implementation notes. The kernel data-processing bounds are a routine Hahn-decomposition argument that follows no particular source, and the roadmap is credited only in the PR description.
✅ api-design approved claude/claude-opus-5-5 The new general total-variation and kernel bounds are public and in their canonical homes; Mathlib has none of them, and the only private lemmas are steps specific to the MDP proof. The PR adds no definitions, so no characterizing API is missing, and there is no compatibility surface.
✅ generality approved claude/claude-opus-5-5 Equal-mass hypotheses appear only in the eventwise lemmas that need them; the sum, mixture and data-processing bounds hold for any finite measures with Markov kernels, and nothing in ./mathlib already covers them. The occupancy continuity bounds are proved once for policies within ε of each other, and the two mixture lemmas follow from them in one line each.
✅ placement approved claude/claude-opus-5-5 The general total-variation bounds are in TotalVariation/Basic.lean, the kernel data-processing bounds are in the new Probability/Kernel/TotalVariation.lean, and the occupancy results are in Occupancy.lean. Every added import is used (Bochner.Set for integral_indicator_one, IntegrableOn for Integrable.of_bound, TauCeti for etvDist_le_two_mul_iff, the kernel file in Occupancy.lean) and none is too broad.
✅ naming approved claude/claude-opus-5-5 New names describe their conclusions in standard Mathlib terminology, consistent with klDiv_comp_right_le and the left/right convention for the changing argument; _le/equality suffixes match statement strength and no notation is introduced.
✅ documentation approved claude/claude-opus-5-5 The new kernel file and both updated modules have module docstrings that state what each file contains and what it is for, and every new public declaration has a docstring that matches its statement and hypotheses without overclaiming. The [AJKS, Lemma 14.1] citation matches the ReinforcementLearning roadmap entry at README.md:1008.
✅ proof-quality approved claude/claude-opus-5-5 The earlier finding is fixed: tvDist_comp_right_le now uses Mathlib's Kernel.IsMarkovKernel.integrable, and the remaining local hint in tvDist_compProd_right_le covers a set that depends on x, which no Mathlib lemma handles. The other proofs are short and rely on automation, with no undocumented change/show or definitional-equality tricks.

♻️ = approved on an earlier commit, re-run before merge.

Reviewing this diff at head f37bdd1. rubrics @ 32ce206. CLI @ 32ce206. Review spend: $6.89.

@lmlexplore-automation lmlexplore-automation Bot added awaiting-author A review requested changes; author action needed and removed awaiting-review CI is green; waiting for review verdicts labels Nov 5, 2026
…ounds

- `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 <noreply@github.com>
@lmlexplore-automation lmlexplore-automation Bot added awaiting-CI CI has not yet reported on the latest commit awaiting-review CI is green; waiting for review verdicts awaiting-author A review requested changes; author action needed and removed awaiting-author A review requested changes; author action needed awaiting-CI CI has not yet reported on the latest commit awaiting-review CI is green; waiting for review verdicts labels Nov 5, 2026
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 <noreply@github.com>
@lmlexplore-automation lmlexplore-automation Bot added awaiting-CI CI has not yet reported on the latest commit awaiting-review CI is green; waiting for review verdicts awaiting-author A review requested changes; author action needed and removed awaiting-author A review requested changes; author action needed awaiting-CI CI has not yet reported on the latest commit awaiting-review CI is green; waiting for review verdicts labels Nov 5, 2026
…t_le

Co-Authored-By: Claude Code <noreply@github.com>
@lmlexplore-automation lmlexplore-automation Bot added awaiting-CI CI has not yet reported on the latest commit awaiting-review CI is green; waiting for review verdicts and removed awaiting-author A review requested changes; author action needed awaiting-CI CI has not yet reported on the latest commit awaiting-review CI is green; waiting for review verdicts labels Nov 5, 2026
@lmlexplore-automation lmlexplore-automation Bot added the ready-to-merge Current head passes all automated merge prerequisites; queue reservations may delay entry label Nov 5, 2026
@lmlexplore-automation
lmlexplore-automation Bot added this pull request to the merge queue Nov 5, 2026
Merged via the queue into main with commit 92f20d2 Nov 5, 2026
22 of 25 checks passed
@lmlexplore-automation lmlexplore-automation Bot removed the ready-to-merge Current head passes all automated merge prerequisites; queue reservations may delay entry label Nov 5, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant