remove Unit

Author: RemyDegenne

LeanMachineLearning/SequentialLearning/BayesStationaryEnv.lean

@@ -16,25 +16,25 @@ A Bayesian stationary environment is an environment that draws a parameter `e :
 `stationaryEnv (κ.sectR e)`. Following the "announced variables" mechanism of
 `LeanMachineLearning/SequentialLearning/Announce.lean`, the parameter is not hidden: it is part of
 the environment's move, and the algorithm is the one that ignores it. Concretely,
-`bayesStationaryEnv Q κ : Environment (𝓔 × Unit) 𝓐 𝓨` announces `e` in every observation and runs
-against `alg.comapObs Prod.snd`, for an `alg : Algorithm Unit 𝓐 𝓨`.
+`bayesStationaryEnv Q κ : Environment 𝓔 𝓐 𝓨` announces `e` as the observation of every round and
+runs against `alg.comapObs (fun _ ↦ ())`, for an `alg : Algorithm Unit 𝓐 𝓨`.
 
 The predicate `IsBayesAlgEnvSeq` is not a new notion of run: it is `IsAlgEnvSeq` for that pair,
 for the observation process that announces the parameter `E` at every round.
 
 ## Main definitions
 
 * `bayesStationaryEnv Q κ`: the environment that draws a parameter from `Q` before the first
-  round, announces it in the first component of every observation, and returns feedback `κ (e, a)`
-  when the parameter is `e` and the action is `a`.
+  round, announces it as the observation of every round, and returns feedback `κ (e, a)` when the
+  parameter is `e` and the action is `a`.
 * `IsBayesAlgEnvSeq Q κ alg E A Y P`: states that the parameter `E : Ω → 𝓔` has law `Q` and that
   the sequences of actions `A : ℕ → Ω → 𝓐` and feedbacks `Y : ℕ → Ω → 𝓨` are generated by the
   algorithm `alg : Algorithm Unit 𝓐 𝓨` interacting with `bayesStationaryEnv Q κ`, which it sees
-  through `Algorithm.comapObs Prod.snd`. Equivalently, `A` and `Y` are generated by `alg`
+  through `Algorithm.comapObs (fun _ ↦ ())`. Equivalently, `A` and `Y` are generated by `alg`
   interacting with the stationary environment `stationaryEnv (κ.sectR (E ω))`.
 * `bayesTrajMeasure Q κ alg`: for any choice of probability measure `Q : Measure 𝓔`, Markov kernel
   `κ : Kernel (𝓔 × 𝓐) 𝓨`, and algorithm `alg : Algorithm Unit 𝓐 𝓨`, provides a probability measure
-  `P : Measure (ℕ → Round (𝓔 × Unit) 𝓐 𝓨)` on a space that carries `E`, `A`, and `Y` such that
+  `P : Measure (ℕ → Round 𝓔 𝓐 𝓨)` on a space that carries `E`, `A`, and `Y` such that
   `IsBayesAlgEnvSeq Q κ alg E A Y P`.
 * `bayesTrajMeasurePosterior Q κ alg n`: a `Kernel (Hist Unit 𝓐 𝓨 n) 𝓔` that represents the
   posterior over `E` given the history before time `n` (the `n` first rounds) under
@@ -44,8 +44,8 @@ for the observation process that announces the parameter `E` at every round.
 
 ## Main results
 
-* `IsAlgEnvSeq.isBayesAlgEnvSeq`: a run of `alg.comapObs Prod.snd` against `bayesStationaryEnv Q κ`
-  is a Bayesian algorithm-environment sequence for the announced parameter.
+* `IsAlgEnvSeq.isBayesAlgEnvSeq`: a run of `alg.comapObs (fun _ ↦ ())` against
+  `bayesStationaryEnv Q κ` is a Bayesian algorithm-environment sequence for the announced parameter.
 * `ae_IsAlgEnvSeq h`: if `h : IsBayesAlgEnvSeq Q κ alg E A Y P`, for `Q`-almost every `e : 𝓔`,
   `IsAlgEnvSeq O' A' Y' alg (stationaryEnv (κ.sectR e)) (condDistrib (trajectory _ A Y) E P e)` for
   some sequence of actions `A' : ℕ → (ℕ → Round Unit 𝓐 𝓨) → 𝓐` and sequence of feedbacks
@@ -69,46 +69,44 @@ variable [MeasurableSpace 𝓔] [MeasurableSpace 𝓐] [MeasurableSpace 𝓨] [M
 
 section BayesEnv
 
-/-- The environment that draws a parameter `e : 𝓔` from `Q` before the first round, announces it in
-the first component of every observation, and returns a feedback drawn from `κ (e, a)` when the
-action is `a`. The algorithm is meant to ignore the announced parameter, that is, to run through
-`Algorithm.comapObs Prod.snd`. -/
+/-- The environment that draws a parameter `e : 𝓔` from `Q` before the first round, announces it as
+the observation of every round, and returns a feedback drawn from `κ (e, a)` when the action is
+`a`. The algorithm is meant to ignore the announced parameter, that is, to run through
+`Algorithm.comapObs (fun _ ↦ ())`. -/
 noncomputable
 def bayesStationaryEnv (Q : Measure 𝓔) [IsProbabilityMeasure Q] (κ : Kernel (𝓔 × 𝓐) 𝓨)
-    [IsMarkovKernel κ] : Environment (𝓔 × Unit) 𝓐 𝓨 where
+    [IsMarkovKernel κ] : Environment 𝓔 𝓐 𝓨 where
   obs
-    | 0 => Kernel.const _ (Q.prod (Measure.dirac ()))
+    | 0 => Kernel.const _ Q
     | _ + 1 => Kernel.deterministic (fun h ↦ (h 0).obs) (by fun_prop)
-  feedback _ := κ.comap (fun p ↦ (p.1.2.1, p.2)) (by fun_prop)
+  feedback _ := κ.comap (fun p ↦ (p.1.2, p.2)) (by fun_prop)
   isMarkovKernel_obs n := by cases n <;> infer_instance
 
 variable {Q : Measure 𝓔} [IsProbabilityMeasure Q] {κ : Kernel (𝓔 × 𝓐) 𝓨} [IsMarkovKernel κ]
 
 @[simp]
-lemma obs_bayesStationaryEnv_zero :
-    (bayesStationaryEnv Q κ).obs 0 = Kernel.const _ (Q.prod (Measure.dirac ())) := rfl
+lemma obs_bayesStationaryEnv_zero : (bayesStationaryEnv Q κ).obs 0 = Kernel.const _ Q := rfl
 
 @[simp]
 lemma obs_bayesStationaryEnv_succ (n : ℕ) :
     (bayesStationaryEnv Q κ).obs (n + 1)
-      = Kernel.deterministic (fun h : Hist (𝓔 × Unit) 𝓐 𝓨 (n + 1) ↦ (h 0).obs) (by fun_prop) := rfl
+      = Kernel.deterministic (fun h : Hist 𝓔 𝓐 𝓨 (n + 1) ↦ (h 0).obs) (by fun_prop) := rfl
 
 @[simp]
 lemma feedback_bayesStationaryEnv (n : ℕ) :
-    (bayesStationaryEnv Q κ).feedback n = κ.comap (fun p ↦ (p.1.2.1, p.2)) (by fun_prop) := rfl
+    (bayesStationaryEnv Q κ).feedback n = κ.comap (fun p ↦ (p.1.2, p.2)) (by fun_prop) := rfl
 
 @[simp]
-lemma obs0_bayesStationaryEnv : (bayesStationaryEnv Q κ).obs0 = Q.prod (Measure.dirac ()) := rfl
+lemma obs0_bayesStationaryEnv : (bayesStationaryEnv Q κ).obs0 = Q := rfl
 
 @[simp]
-lemma ν0_bayesStationaryEnv :
-    (bayesStationaryEnv Q κ).ν0 = κ.comap (fun p ↦ (p.1.1, p.2)) (by fun_prop) := rfl
+lemma ν0_bayesStationaryEnv : (bayesStationaryEnv Q κ).ν0 = κ := rfl
 
 end BayesEnv
 
-/-- Insert an announced parameter `e` into every round of an observable history. -/
-def announceHist (e : 𝓔) {n : ℕ} (h : Hist Unit 𝓐 𝓨 n) : Hist (𝓔 × Unit) 𝓐 𝓨 n :=
-  fun i ↦ ((e, ()), (h i).action, (h i).feedback)
+/-- Insert an announced parameter `e` as the observation of every round of an observable history. -/
+def announceHist (e : 𝓔) {n : ℕ} (h : Hist Unit 𝓐 𝓨 n) : Hist 𝓔 𝓐 𝓨 n :=
+  fun i ↦ (e, (h i).action, (h i).feedback)
 
 @[fun_prop]
 lemma measurable_announceHist (n : ℕ) :
@@ -123,13 +121,13 @@ interacting with an underlying environment that depends on `E` and `κ`
 (`stationaryEnv (κ.sectR (E ω))`).
 
 This is `IsAlgEnvSeq` for the announcing environment `bayesStationaryEnv Q κ` and the algorithm
-`alg.comapObs Prod.snd` that ignores the announced parameter: the observation at every round is
-`(E ω, ())`. -/
+`alg.comapObs (fun _ ↦ ())` that ignores the announced parameter: the observation at every round
+is `E ω`. -/
 def IsBayesAlgEnvSeq (Q : Measure 𝓔) [IsProbabilityMeasure Q] (κ : Kernel (𝓔 × 𝓐) 𝓨)
     [IsMarkovKernel κ] (alg : Algorithm Unit 𝓐 𝓨)
     (E : Ω → 𝓔) (A : ℕ → Ω → 𝓐) (Y : ℕ → Ω → 𝓨)
     (P : Measure Ω) [IsProbabilityMeasure P] : Prop :=
-  IsAlgEnvSeq (fun _ ω ↦ (E ω, ())) A Y (alg.comapObs Prod.snd) (bayesStationaryEnv Q κ) P
+  IsAlgEnvSeq (fun _ ↦ E) A Y (alg.comapObs (fun _ : 𝓔 ↦ ())) (bayesStationaryEnv Q κ) P
 
 namespace IsBayesAlgEnvSeq
 
@@ -154,43 +152,38 @@ lemma mk (hasLaw_env : HasLaw E Q P)
     (measurable_action : ∀ n, Measurable (A n) := by fun_prop)
     (measurable_feedback : ∀ n, Measurable (Y n) := by fun_prop) :
     IsBayesAlgEnvSeq Q κ alg E A Y P := by
-  have hO : ∀ _ : ℕ, Measurable (fun ω ↦ ((E ω, ()) : 𝓔 × Unit)) :=
-    fun _ ↦ measurable_param.prodMk measurable_const
-  refine IsAlgEnvSeq.mk hO measurable_action measurable_feedback ?_ ?_ ?_
+  refine IsAlgEnvSeq.mk (fun _ ↦ measurable_param) measurable_action measurable_feedback ?_ ?_ ?_
   · intro n
     cases n with
     | zero =>
       rw [history_zero, obs_bayesStationaryEnv_zero]
-      refine HasLaw.hasCondDistrib_const ⟨(hO 0).aemeasurable, ?_⟩
-      rw [Kernel.const_apply, Measure.prod_dirac, ← hasLaw_env.map_eq,
-        AEMeasurable.map_map_of_aemeasurable (by fun_prop) hasLaw_env.aemeasurable]
-      rfl
+      exact hasLaw_env.hasCondDistrib_const
     | succ n =>
       rw [obs_bayesStationaryEnv_succ]
       exact hasCondDistrib_deterministic _
-        (measurable_history hO measurable_action measurable_feedback (n + 1)).aemeasurable
-        (ae_of_all _ fun _ ↦ rfl)
+        (measurable_history (fun _ ↦ measurable_param) measurable_action measurable_feedback
+          (n + 1)).aemeasurable (ae_of_all _ fun _ ↦ rfl)
   · intro n
     exact HasCondDistrib.comp_right
-      (f := fun q : 𝓔 × (Hist Unit 𝓐 𝓨 n × Unit) ↦ (announceHist q.1 q.2.1, (q.1, q.2.2)))
+      (f := fun q : 𝓔 × (Hist Unit 𝓐 𝓨 n × Unit) ↦ (announceHist q.1 q.2.1, q.1))
       (hf := by fun_prop)
       (Z := fun ω ↦ (E ω, (history (noObs Ω) A Y n ω, noObs Ω n ω)))
       (hasCondDistrib_action n)
   · intro n
     exact HasCondDistrib.comp_right
       (f := fun q : (Hist Unit 𝓐 𝓨 n × Unit) × (𝓔 × 𝓐) ↦
-        ((announceHist q.2.1 q.1.1, (q.2.1, q.1.2)), q.2.2))
+        ((announceHist q.2.1 q.1.1, q.2.1), q.2.2))
       (hf := by fun_prop)
       (Z := fun ω ↦ ((history (noObs Ω) A Y n ω, noObs Ω n ω), (E ω, A n ω)))
       (hasCondDistrib_feedback n)
 
 /-- A Bayesian algorithm-environment sequence is an algorithm-environment sequence for the
 announcing environment `bayesStationaryEnv Q κ`. -/
 lemma isAlgEnvSeq (h : IsBayesAlgEnvSeq Q κ alg E A Y P) :
-    IsAlgEnvSeq (fun _ ω ↦ (E ω, ())) A Y (alg.comapObs Prod.snd) (bayesStationaryEnv Q κ) P := h
+    IsAlgEnvSeq (fun _ ↦ E) A Y (alg.comapObs (fun _ : 𝓔 ↦ ())) (bayesStationaryEnv Q κ) P := h
 
 lemma measurable_param (h : IsBayesAlgEnvSeq Q κ alg E A Y P) : Measurable E :=
-  (h.isAlgEnvSeq.measurable_obs 0).fst
+  h.isAlgEnvSeq.measurable_obs 0
 
 lemma measurable_action (h : IsBayesAlgEnvSeq Q κ alg E A Y P) (n : ℕ) : Measurable (A n) :=
   h.isAlgEnvSeq.measurable_action n
@@ -202,22 +195,17 @@ lemma measurable_feedback (h : IsBayesAlgEnvSeq Q κ alg E A Y P) (n : ℕ) : Me
 lemma hasLaw_env (h : IsBayesAlgEnvSeq Q κ alg E A Y P) : HasLaw E Q P := by
   have h0 := h.isAlgEnvSeq.hasCondDistrib_obs 0
   rw [history_zero] at h0
-  have h1 : HasLaw (fun ω ↦ (E ω, ())) (Q.prod (Measure.dirac ())) P := h0.hasLaw_of_const'
-  refine ⟨h.measurable_param.aemeasurable, ?_⟩
-  have h2 : E = Prod.fst ∘ (fun ω ↦ (E ω, ())) := rfl
-  rw [h2, ← Measure.map_map measurable_fst
-    (h.measurable_param.prodMk (measurable_const : Measurable fun _ : Ω ↦ ())), h1.map_eq]
-  exact Measure.fst_prod
+  exact h0.hasLaw_of_const'
 
 /-- The action at time `n` has the correct conditional distribution given the parameter and the
 history: it depends only on the history. -/
 lemma hasCondDistrib_action (h : IsBayesAlgEnvSeq Q κ alg E A Y P) (n : ℕ) :
     HasCondDistrib (A n) (fun ω ↦ (E ω, (history (noObs Ω) A Y n ω, noObs Ω n ω)))
       ((alg.policy n).prodMkLeft _) P :=
   HasCondDistrib.comp_right
-    (f := fun p : Hist (𝓔 × Unit) 𝓐 𝓨 n × (𝓔 × Unit) ↦ (p.2.1, (Hist.mapObs Prod.snd p.1, p.2.2)))
+    (f := fun p : Hist 𝓔 𝓐 𝓨 n × 𝓔 ↦ (p.2, (Hist.mapObs (fun _ ↦ ()) p.1, ())))
     (hf := by fun_prop)
-    (Z := fun ω ↦ (history (fun _ ω ↦ (E ω, ())) A Y n ω, (E ω, ())))
+    (Z := fun ω ↦ (history (fun _ ↦ E) A Y n ω, E ω))
     (h.isAlgEnvSeq.hasCondDistrib_action n)
 
 /-- The feedback at time `n` has the correct conditional distribution given the history, the
@@ -226,10 +214,10 @@ lemma hasCondDistrib_feedback (h : IsBayesAlgEnvSeq Q κ alg E A Y P) (n : ℕ)
     HasCondDistrib (Y n) (fun ω ↦ ((history (noObs Ω) A Y n ω, noObs Ω n ω), (E ω, A n ω)))
       (κ.prodMkLeft _) P :=
   HasCondDistrib.comp_right
-    (f := fun p : (Hist (𝓔 × Unit) 𝓐 𝓨 n × (𝓔 × Unit)) × 𝓐 ↦
-      ((Hist.mapObs Prod.snd p.1.1, p.1.2.2), (p.1.2.1, p.2)))
+    (f := fun p : (Hist 𝓔 𝓐 𝓨 n × 𝓔) × 𝓐 ↦
+      ((Hist.mapObs (fun _ ↦ ()) p.1.1, ()), (p.1.2, p.2)))
     (hf := by fun_prop)
-    (Z := fun ω ↦ ((history (fun _ ω ↦ (E ω, ())) A Y n ω, (E ω, ())), A n ω))
+    (Z := fun ω ↦ ((history (fun _ ↦ E) A Y n ω, E ω), A n ω))
     (h.isAlgEnvSeq.hasCondDistrib_feedback n)
 
 lemma hasCondDistrib_action' (h : IsBayesAlgEnvSeq Q κ alg E A Y P) (n : ℕ) :
@@ -325,16 +313,16 @@ end IsBayesAlgEnvSeq
 section IsAlgEnvSeq
 
 variable {Q : Measure 𝓔} [IsProbabilityMeasure Q] {κ : Kernel (𝓔 × 𝓐) 𝓨} [IsMarkovKernel κ]
-variable {alg : Algorithm Unit 𝓐 𝓨} {O : ℕ → Ω → 𝓔 × Unit} {A : ℕ → Ω → 𝓐} {Y : ℕ → Ω → 𝓨}
+variable {alg : Algorithm Unit 𝓐 𝓨} {O : ℕ → Ω → 𝓔} {A : ℕ → Ω → 𝓐} {Y : ℕ → Ω → 𝓨}
 variable {P : Measure Ω} [IsProbabilityMeasure P]
 
 /-- Under `bayesStationaryEnv Q κ`, the announced parameter is almost surely the same at every
 round. -/
 lemma IsAlgEnvSeq.ae_obs_eq_obs_zero [StandardBorelSpace 𝓔]
-    (h : IsAlgEnvSeq O A Y (alg.comapObs Prod.snd) (bayesStationaryEnv Q κ) P) (n : ℕ) :
-    (fun ω ↦ ((O 0 ω).1, ())) =ᵐ[P] O n := by
+    (h : IsAlgEnvSeq O A Y (alg.comapObs (fun _ : 𝓔 ↦ ())) (bayesStationaryEnv Q κ) P) (n : ℕ) :
+    O 0 =ᵐ[P] O n := by
   cases n with
-  | zero => exact ae_of_all _ fun _ ↦ rfl
+  | zero => rfl
   | succ n =>
     have h1 := h.hasCondDistrib_obs (n + 1)
     rw [obs_bayesStationaryEnv_succ] at h1
@@ -344,46 +332,40 @@ lemma IsAlgEnvSeq.ae_obs_eq_obs_zero [StandardBorelSpace 𝓔]
     rw [hω]
     simp [history_apply]
 
-/-- A run of `alg.comapObs Prod.snd` against the announcing environment `bayesStationaryEnv Q κ`
-is a Bayesian algorithm-environment sequence for the announced parameter. -/
+/-- A run of `alg.comapObs (fun _ ↦ ())` against the announcing environment
+`bayesStationaryEnv Q κ` is a Bayesian algorithm-environment sequence for the announced
+parameter. -/
 lemma IsAlgEnvSeq.isBayesAlgEnvSeq [StandardBorelSpace 𝓔]
-    (h : IsAlgEnvSeq O A Y (alg.comapObs Prod.snd) (bayesStationaryEnv Q κ) P) :
-    IsBayesAlgEnvSeq Q κ alg (fun ω ↦ (O 0 ω).1) A Y P :=
-  h.congr (fun _ ↦ (h.measurable_obs 0).fst.prodMk measurable_const)
-    h.measurable_action h.measurable_feedback h.ae_obs_eq_obs_zero
-    (fun _ ↦ .rfl) (fun _ ↦ .rfl)
+    (h : IsAlgEnvSeq O A Y (alg.comapObs (fun _ : 𝓔 ↦ ())) (bayesStationaryEnv Q κ) P) :
+    IsBayesAlgEnvSeq Q κ alg (O 0) A Y P :=
+  h.congr (fun _ ↦ h.measurable_obs 0) h.measurable_action h.measurable_feedback
+    h.ae_obs_eq_obs_zero (fun _ ↦ .rfl) (fun _ ↦ .rfl)
 
 end IsAlgEnvSeq
 
 namespace IT
 
 /-- A measure `P` on a measurable space that carries random variables `E`, `A`, and `Y` such that
-`IsBayesAlgEnvSeq Q κ alg E A Y P`. -/
+`IsBayesAlgEnvSeq Q κ alg E A Y P`. The parameter is the observation of the first round,
+`IT.obs 0`. -/
 noncomputable
 def bayesTrajMeasure (Q : Measure 𝓔) [IsProbabilityMeasure Q] (κ : Kernel (𝓔 × 𝓐) 𝓨)
-    [IsMarkovKernel κ] (alg : Algorithm Unit 𝓐 𝓨) : Measure (ℕ → Round (𝓔 × Unit) 𝓐 𝓨) :=
-  trajMeasure (alg.comapObs Prod.snd) (bayesStationaryEnv Q κ)
+    [IsMarkovKernel κ] (alg : Algorithm Unit 𝓐 𝓨) : Measure (ℕ → Round 𝓔 𝓐 𝓨) :=
+  trajMeasure (alg.comapObs (fun _ : 𝓔 ↦ ())) (bayesStationaryEnv Q κ)
 deriving IsProbabilityMeasure
 
-/-- The parameter announced by `bayesStationaryEnv Q κ`, read on the trajectory space. -/
-def param (τ : ℕ → Round (𝓔 × Unit) 𝓐 𝓨) : 𝓔 := (IT.obs 0 τ).1
-
-@[fun_prop]
-lemma measurable_param : Measurable (param (𝓔 := 𝓔) (𝓐 := 𝓐) (𝓨 := 𝓨)) := by
-  unfold param; fun_prop
-
 lemma isBayesAlgEnvSeq_bayesTrajMeasure [StandardBorelSpace 𝓔]
     (Q : Measure 𝓔) [IsProbabilityMeasure Q] (κ : Kernel (𝓔 × 𝓐) 𝓨) [IsMarkovKernel κ]
     (alg : Algorithm Unit 𝓐 𝓨) :
-    IsBayesAlgEnvSeq Q κ alg param action feedback (bayesTrajMeasure Q κ alg) :=
+    IsBayesAlgEnvSeq Q κ alg (obs 0) action feedback (bayesTrajMeasure Q κ alg) :=
   (isAlgEnvSeq_trajMeasure _ _).isBayesAlgEnvSeq
 
 /-- A kernel that represents the posterior over `E` given the history before time `n`. -/
 noncomputable
 def bayesTrajMeasurePosterior [StandardBorelSpace 𝓔] [Nonempty 𝓔]
     (Q : Measure 𝓔) [IsProbabilityMeasure Q] (κ : Kernel (𝓔 × 𝓐) 𝓨) [IsMarkovKernel κ]
     (alg : Algorithm Unit 𝓐 𝓨) (n : ℕ) : Kernel (Hist Unit 𝓐 𝓨 n) 𝓔 :=
-  condDistrib param (history (noObs _) action feedback n) (bayesTrajMeasure Q κ alg)
+  condDistrib (obs 0) (history (noObs _) action feedback n) (bayesTrajMeasure Q κ alg)
 deriving IsMarkovKernel
 
 /-- The posterior given the empty history is the prior. -/