Refactor Algorithm and Environment

LMLTutorial/Pages/DefiningAlgorithm.lean

@@ -33,42 +33,52 @@ We then illustrate them with the UCB bandit algorithm.
 # Algorithm and environment
 
 In LML, we prove theorems about the interaction of an algorithm with an environment.
-An algorithm takes actions, to which the environment responds with feedback (e.g., rewards for the bandit case, the gradient of a function in optimization problems).
-In general, both action and feedback can depend on the entire history up to the current time and can be randomized.
+Each round of the interaction consists of three stages: the environment draws an observation (e.g., the context in a contextual bandit), the algorithm takes an action based on that observation, and the environment responds with feedback (e.g., rewards for the bandit case, the gradient of a function in optimization problems).
+In general, observation, action and feedback can depend on the entire history up to the current time and can be randomized.
+One round is recorded by the `Round` abbreviation and a history of `n` rounds by the `Hist` abbreviation.
+
+{docstring Round}
+
+{docstring Hist}
+
 The `Algorithm` structure is defined as follows:
 
 {docstring Algorithm}
 
-This structure refers to two types, the type of actions `𝓐` and the type of feedback `𝓨`.
-Both are measurable spaces, since we consider stochastic algorithms and environments.
-Before time `n`, there is a history of actions and feedbacks `Fin n → 𝓐 × 𝓨` (the `n` pairs of action and feedback at times `0, ..., n - 1`; the processes are 0-indexed).
-The `policy` field contains for each time `n` a kernel from that history to the action space.
-That is, it maps every possible history to a random action at time `n` (and that map is measurable).
-The `h_policy` field records that the measure describing the action is a probability measure (and it is in square brackets to tell Lean to infer it automatically whenever possible).
-At time `0` the history is empty: `Fin 0 → 𝓐 × 𝓨` has a unique element, and the distribution of the first action is `policy 0` applied to that element.
-That distribution is called `Algorithm.p0`.
+This structure refers to three types, the type of observations `𝓞`, the type of actions `𝓐` and the type of feedback `𝓨`.
+All three are measurable spaces, since we consider stochastic algorithms and environments.
+Before time `n`, there is a history of `n` complete rounds `Hist 𝓞 𝓐 𝓨 n = Fin n → 𝓞 × 𝓐 × 𝓨` (the observation-action-feedback triples at times `0, ..., n - 1`; the processes are 0-indexed).
+The `policy` field contains for each time `n` a kernel from that history together with the observation at time `n` to the action space.
+That is, it maps every possible history and current observation to a random action at time `n` (and that map is measurable).
+The `isMarkovKernel_policy` field records that the measure describing the action is a probability measure (and it is in square brackets to tell Lean to infer it automatically whenever possible).
+At time `0` the history is empty: `Hist 𝓞 𝓐 𝓨 0` has a unique element, and the distribution of the first action given the first observation is `policy 0` applied to that element.
+That kernel is called `Algorithm.p0`.
+
+Many settings have no observations at all: the algorithm sees only the past rounds. Those are described by taking `𝓞 = Unit`, and we write `noObs Ω` for the corresponding (constant) observation process.
 
-If the algorithms actions are not random, we can use the `detAlgorithm` definition to build an algorithm from the data of a measurable function for the action at each time, as a function of the history before that time.
+If the algorithms actions are not random, we can use the `detAlgorithm` definition to build an algorithm from the data of a measurable function for the action at each time, as a function of the history before that time and of the current observation.
 The first action is the value of that function at time `0` on the empty history.
 
 {docstring detAlgorithm}
 
 We can see here that we did not need to prove that the kernels are `IsMarkovKernel`.
 Lean knows that deterministic kernels are Markov.
 
-The `Environment` structure is the mirror of the `Algorithm` structure, with a kernel for the feedback instead of the actions.
+The `Environment` structure is the mirror of the `Algorithm` structure, with a kernel for the observation and a kernel for the feedback instead of the actions.
 
 {docstring Environment}
 
-`feedback n` gives the distribution of the feedback at time `n` given the history before `n` and the action at time `n`.
-The distribution of the first feedback given the first action is `feedback 0` applied to the empty history; it is called `Environment.ν0`.
+`obs n` gives the distribution of the observation at time `n` given the history before `n`.
+`feedback n` gives the distribution of the feedback at time `n` given the history before `n`, the observation and the action at time `n`.
+The distribution of the first observation is `obs 0` applied to the empty history; it is called `Environment.obs0`.
+The distribution of the first feedback given the first observation and action is `feedback 0` applied to the empty history; it is called `Environment.ν0`.
 
-In many applications the feedback depends only on the last action and not on the prior history.
+In many applications there is no observation and the feedback depends only on the last action, not on the prior history.
 We provide an `obliviousEnv` definition that builds an environment for those cases.
 
 {docstring obliviousEnv}
 
-`(ν n).prodMkLeft _` is the kernel `ν n` seen as a `Kernel ((Fin n → 𝓐 × 𝓨) × 𝓐) 𝓨` by ignoring the history.
+`(ν n).prodMkLeft _` is the kernel `ν n` seen as a `Kernel ((Hist Unit 𝓐 𝓨 n × Unit) × 𝓐) 𝓨` by ignoring the history and the observation.
 
 If furthermore the feedback kernel does not change with time, we can use the `stationaryEnv` definition to build the environment.
 
@@ -82,7 +92,7 @@ This is done by the `IsAlgEnvSeq` structure.
 
 {docstring IsAlgEnvSeq}
 
-This structure takes as input two sequences of random variables (two stochastic processes), `A` and `Y`, which represent the actions and feedback generated by the interaction of the algorithm with the environment.
+This structure takes as input three sequences of random variables (three stochastic processes), `O`, `A` and `Y`, which represent the observations, actions and feedback generated by the interaction of the algorithm with the environment.
 It states that those sequences are measurable and that they have the correct conditional distributions given by the algorithm and environment.
 The measurable space `Ω` and the measure `P` are not imposed: they can be chosen as we want, as long as the conditions of `IsAlgEnvSeq` are satisfied.
 This definition requires `𝓐` and `𝓨` to be nonempty standard Borel spaces, because Mathlib's theory about conditional distributions requires those assumptions.
@@ -149,7 +159,7 @@ The theorem gives an upper bound on the expected regret of UCB at time `n`.
 
 # Building vs analyzing algorithms
 
-When building an algorithm, we describe it with functions from the history `(Fin n → 𝓐 × R)` to the action space `𝓐`.
+When building an algorithm, we describe it with functions from the history and the current observation `(Hist 𝓞 𝓐 R n × 𝓞)` to the action space `𝓐`.
 Thus, to construct UCB, we used the following empirical mean function.
 
 {docstring empMean'}

LeanMachineLearning.lean

@@ -50,6 +50,7 @@ public import LeanMachineLearning.SequentialLearning.Algorithms.RandomSampling.T
 public import LeanMachineLearning.SequentialLearning.Algorithms.RoundRobin
 public import LeanMachineLearning.SequentialLearning.Algorithms.Uniform
 public import LeanMachineLearning.SequentialLearning.BayesStationaryEnv
+public import LeanMachineLearning.SequentialLearning.Comap
 public import LeanMachineLearning.SequentialLearning.Deterministic
 public import LeanMachineLearning.SequentialLearning.DivergenceDecomposition
 public import LeanMachineLearning.SequentialLearning.EvaluationEnv

LeanMachineLearning/ForMathlib/MeasureTheory/MeasurableSpace/Embedding.lean

@@ -26,6 +26,18 @@ learning algorithms (elements of `Fin n → 𝓐 × 𝓨` or `Iic n → 𝓐 ×
 
 open Finset Preorder
 
+/-- `Prod.mk x` is a measurable embedding as soon as `{x}` is measurable. This generalises
+`measurableEmbedding_prodMk_left`, which assumes `MeasurableSingletonClass`. -/
+lemma measurableEmbedding_prodMk_left_of_measurableSet {α β : Type*} [MeasurableSpace α]
+    [MeasurableSpace β] {x : α} (hx : MeasurableSet {x}) :
+    MeasurableEmbedding (Prod.mk x : β → α × β) where
+  injective _ _ h := (Prod.ext_iff.mp h).2
+  measurable := by fun_prop
+  measurableSet_image' s hs := by
+    convert! hx.prod hs
+    ext p
+    simp [Prod.ext_iff, eq_comm, and_left_comm]
+
 lemma coe_default_Iic_zero : ((default : Iic 0) : ℕ) = 0 := rfl
 
 namespace MeasurableEquiv

LeanMachineLearning/ForMathlib/Probability/HasCondDistrib.lean

@@ -145,6 +145,17 @@ lemma HasLaw.prod_of_hasCondDistrib {P : Measure β}
     HasLaw (fun ω ↦ (X ω, Y ω)) (P ⊗ₘ κ) μ :=
   ⟨by fun_prop, by rw [h2.map_eq, h1.map_eq]⟩
 
+/-- `HasCondDistrib` only depends on the almost everywhere equivalence classes of the two random
+variables. -/
+lemma HasCondDistrib.congr {X' : α → β} {Y' : α → Ω} (h : HasCondDistrib Y X κ μ)
+    (hX : X' =ᵐ[μ] X) (hY : Y' =ᵐ[μ] Y) :
+    HasCondDistrib Y' X' κ μ := by
+  have h_pair : (fun a ↦ (X' a, Y' a)) =ᵐ[μ] fun a ↦ (X a, Y a) := by
+    filter_upwards [hX, hY] with a h1 h2
+    rw [h1, h2]
+  exact ⟨h.aemeasurable.congr h_pair.symm, by rw [Measure.map_congr h_pair,
+    Measure.map_congr hX, h.map_eq]⟩
+
 lemma HasCondDistrib.hasLaw_comp [SFinite μ] [IsSFiniteKernel κ] (h : HasCondDistrib Y X κ μ) :
     HasLaw Y (κ ∘ₘ (μ.map X)) μ := by
   refine ⟨by fun_prop, ?_⟩
@@ -160,6 +171,13 @@ lemma HasCondDistrib.prod {Z : α → Ω'} {η : Kernel (β × Ω) Ω'}
     AEMeasurable.map_map_of_aemeasurable (by fun_prop) (by fun_prop)]
   rfl
 
+lemma hasCondDistrib_comp_self [SFinite μ] {f : β → Ω} (hf : Measurable f)
+    (hX : AEMeasurable X μ) :
+    HasCondDistrib (f ∘ X) X (Kernel.deterministic f hf) μ := by
+  refine ⟨hX.prodMk (hf.comp_aemeasurable hX), ?_⟩
+  rw [Measure.compProd_deterministic, AEMeasurable.map_map_of_aemeasurable (by fun_prop) hX]
+  rfl
+
 lemma ae_eq_of_hasCondDistrib_deterministic [MeasurableEq Ω] [SFinite μ] {f : β → Ω}
     (hf : Measurable f) (hX : AEMeasurable X μ)
     (hY : AEMeasurable Y μ) (h : HasCondDistrib Y X (Kernel.deterministic f hf) μ) :
@@ -169,6 +187,14 @@ lemma ae_eq_of_hasCondDistrib_deterministic [MeasurableEq Ω] [SFinite μ] {f :
     AEMeasurable.map_map_of_aemeasurable (by fun_prop) (by fun_prop)]
   rfl
 
+lemma hasCondDistrib_deterministic_iff [MeasurableEq Ω] [SFinite μ] {f : β → Ω}
+    (hf : Measurable f) (hX : AEMeasurable X μ) (hY : AEMeasurable Y μ) :
+    HasCondDistrib Y X (Kernel.deterministic f hf) μ ↔ Y =ᵐ[μ] f ∘ X := by
+  refine ⟨ae_eq_of_hasCondDistrib_deterministic hf hX hY, fun h ↦ ?_⟩
+  refine HasCondDistrib.congr ?_ ?_ h (X := X)
+  · exact hasCondDistrib_comp_self hf hX
+  · rfl
+
 section Const
 
 section CompRight
@@ -350,5 +376,4 @@ lemma HasCondDistrib.hasCondDistrib_sectR [IsFiniteMeasure μ] [StandardBorelSpa
   rw [Kernel.map_apply _ hf] at ha
   filter_upwards [hc, ha] with b hcb hab using hcb.trans hab
 
-
 end ProbabilityTheory

LeanMachineLearning/ForMathlib/Probability/WithDensity.lean

@@ -120,6 +120,30 @@ lemma compProd_withDensity_left {κ : Kernel α β} {η : Kernel (α × β) γ}
     _ = ((κ ⊗ₖ η).withDensity (fun a bc ↦ f a bc.1)) a := by
         rw [← compProd_apply_eq_compProd_sectR, Kernel.withDensity_apply _ (by fun_prop)]
 
+lemma sectR_withDensity {η : Kernel (α × β) γ} {g : α × β → γ → ℝ≥0∞} [IsSFiniteKernel η]
+    (hg : Measurable (Function.uncurry g)) (a : α) :
+    (η.withDensity g).sectR a = (η.sectR a).withDensity (fun b ↦ g (a, b)) := by
+  ext b s hs
+  rw [Kernel.sectR_apply, Kernel.withDensity_apply' _ hg,
+    Kernel.withDensity_apply' _ (by fun_prop), Kernel.sectR_apply]
+
+lemma compProd_withDensity_right {κ : Kernel α β} {η : Kernel (α × β) γ} {g : α × β → γ → ℝ≥0∞}
+    [IsSFiniteKernel κ] [IsSFiniteKernel η] [IsSFiniteKernel (η.withDensity g)]
+    (hg : Measurable (Function.uncurry g)) :
+    κ ⊗ₖ (η.withDensity g) = (κ ⊗ₖ η).withDensity (fun a bc ↦ g (a, bc.1) bc.2) := by
+  ext a : 1
+  have h_sf : IsSFiniteKernel ((η.sectR a).withDensity (fun b ↦ g (a, b))) := by
+    rw [← sectR_withDensity hg]
+    infer_instance
+  calc (κ ⊗ₖ (η.withDensity g)) a
+      = (κ a) ⊗ₘ ((η.withDensity g).sectR a) := compProd_apply_eq_compProd_sectR ..
+    _ = (κ a) ⊗ₘ ((η.sectR a).withDensity (fun b ↦ g (a, b))) := by rw [sectR_withDensity hg]
+    _ = ((κ a) ⊗ₘ (η.sectR a)).withDensity (fun p ↦ g (a, p.1) p.2) := by
+        refine Measure.compProd_withDensity ?_
+        fun_prop
+    _ = ((κ ⊗ₖ η).withDensity (fun a bc ↦ g (a, bc.1) bc.2)) a := by
+        rw [← compProd_apply_eq_compProd_sectR, Kernel.withDensity_apply _ (by fun_prop)]
+
 lemma withDensity_rnDeriv_eq' {κ η : Kernel α β} [MeasurableSpace.CountableOrCountablyGenerated α β]
     [IsFiniteKernel κ] [IsFiniteKernel η] (h : ∀ a, κ a ≪ η a) :
     η.withDensity (κ.rnDeriv η) = κ :=