CoveredAdaptiveActions.lean

1 import KomorebiLSM.SequentialConfidence
2 import Mathlib.Probability.Kernel.Composition.IntegralCompProd
3 import Mathlib.MeasureTheory.Integral.Bochner.SumMeasure
4 import Mathlib.Tactic
5 
6 /-! Covered action evidence from a literal conditional experiment.
7 The experiment assumption is equality of joint measures, not an assumed
8 conditional expectation of the importance-weighted observation. -/
9 set_option autoImplicit false
10 namespace Komorebi.LSM.CoveredAdaptiveActions
11 open MeasureTheory ProbabilityTheory
12 open scoped ENNReal BigOperators
13 
14 section KernelCalculation
15 variable {H A Y : Type*} [mH : MeasurableSpace H] [MeasurableSpace A] [MeasurableSpace Y]
16 variable [Fintype A] [DecidableEq A] [MeasurableSingletonClass A]
17 variable (π : Kernel H A) (Q : Kernel (H × A) Y)
18 variable [IsMarkovKernel π] [IsMarkovKernel Q]
19 
20 /-- The recorded target-action probability of the actual policy kernel. -/
21 noncomputable def propensity (a : A) (h : H) : ℝ := (π h).real {a}
22 
23 omit [Fintype A] [DecidableEq A] [IsMarkovKernel π] in
24 theorem propensity_measurable (a : A) : Measurable (propensity π a) :=
25   (π.measurable_coe (measurableSet_singleton a)).ennreal_toReal
26 
27 /-- A predictable bounded test evaluated under the target observation kernel. -/
28 noncomputable def targetMean (φ : H × Y → ℝ) (a : A) (h : H) : ℝ :=
29   ∫ y, φ (h,y) ∂Q (h,a)
30 
31 omit [Fintype A] [DecidableEq A] [MeasurableSingletonClass A] in
32 theorem targetMean_measurable (φ : H × Y → ℝ) (hφ : Measurable φ) (a : A) :
33     Measurable (targetMean Q φ a) := by
34   let Qa : Kernel H Y := Q.comap (fun h => (h,a)) (measurable_id.prodMk measurable_const)
35   exact hφ.stronglyMeasurable.integral_kernel_prod_right' (κ := Qa) |>.measurable
36 
37 omit [Fintype A] [DecidableEq A] [MeasurableSingletonClass A] in
38 theorem targetMean_mem_Icc (φ : H × Y → ℝ) (hφ : Measurable φ)
39     (hb : ∀ h y,φ (h,y)∈Set.Icc (0:ℝ) 1) (a : A) (h : H) :
40     targetMean Q φ a h∈Set.Icc (0:ℝ) 1 := by
41   have hi : Integrable (fun y => φ (h,y)) (Q (h,a)) := Integrable.of_mem_Icc (0:ℝ) 1
42     ((hφ.comp (measurable_const.prodMk measurable_id)).aemeasurable)
43     (Filter.Eventually.of_forall (hb h))
44   refine ⟨integral_nonneg (fun y => (hb h y).1),?_⟩
45   have h := integral_mono_ae hi (integrable_const (1:ℝ))
46     (Filter.Eventually.of_forall (fun y => (hb h y).2))
47   simpa only [targetMean,integral_const,probReal_univ,one_smul] using h
48 
49 /-- Importance weighting before sampling the action/observation pair. -/
50 noncomputable def weightedTest (φ : H × Y → ℝ) (a : A) (p : H × (A × Y)) : ℝ :=
51   if p.2.1=a then φ (p.1,p.2.2)/propensity π a p.1 else 0
52 
53 omit [Fintype A] [IsMarkovKernel π] in
54 theorem weightedTest_measurable (φ : H × Y → ℝ) (hφ : Measurable φ) (a : A) :
55     Measurable (weightedTest π φ a) := by
56   apply Measurable.ite
57     ((measurableSet_singleton a).preimage (measurable_fst.comp measurable_snd))
58     ((hφ.comp (measurable_fst.prodMk (measurable_snd.comp measurable_snd))).div
59       ((propensity_measurable π a).comp measurable_fst)) measurable_const
60 
61 omit [MeasurableSpace Y] [Fintype A] [MeasurableSingletonClass A] [IsMarkovKernel π] in
62 theorem weightedTest_mem_Icc (φ : H × Y → ℝ)
63     (hb : ∀ h y,φ (h,y)∈Set.Icc (0:ℝ) 1) (a : A)
64     (c : ℝ) (hc : 0<c) (hcover : ∀ h,c≤propensity π a h) (p : H × (A × Y)) :
65     weightedTest π φ a p∈Set.Icc (0:ℝ) (1/c) := by
66   have hp : 0<propensity π a p.1 := hc.trans_le (hcover p.1)
67   by_cases ha : p.2.1=a
68   · simp only [weightedTest,ha,ite_true,Set.mem_Icc]
69     refine ⟨div_nonneg (hb _ _).1 hp.le,?_⟩
70     apply (div_le_div_iff₀ hp hc).mpr
71     have := (hb p.1 p.2.2).2
72     nlinarith [hcover p.1]
73   · simp only [weightedTest,ha,ite_false,Set.mem_Icc,le_refl,true_and]
74     exact one_div_nonneg.mpr hc.le
75 
76 /-- Sampling the policy and then its action-conditioned observation cancels
77 the strictly positive target propensity exactly. -/
78 theorem integral_weightedTest (φ : H × Y → ℝ) (hφ : Measurable φ)
79     (hb : ∀ h y,φ (h,y)∈Set.Icc (0:ℝ) 1) (a : A)
80     (c : ℝ) (hc : 0<c) (hcover : ∀ h,c≤propensity π a h) (h : H) :
81     (∫ r,weightedTest π φ a (h,r)∂(π ⊗ₖ Q) h)=targetMean Q φ a h := by
82   have hmeas : Measurable (fun r => weightedTest π φ a (h,r)) :=
83     (weightedTest_measurable π φ hφ a).comp (measurable_const.prodMk measurable_id)
84   have hi : Integrable (fun r => weightedTest π φ a (h,r)) ((π ⊗ₖ Q) h) :=
85     Integrable.of_mem_Icc (0:ℝ) (1/c) hmeas.aemeasurable
86     (Filter.Eventually.of_forall (fun r => weightedTest_mem_Icc π φ hb a c hc hcover (h,r)))
87   rw [ProbabilityTheory.integral_compProd hi]
88   have hinner (b : A) : (∫ y,weightedTest π φ a (h,(b,y))∂Q (h,b))=
89       if b=a then targetMean Q φ a h/propensity π a h else 0 := by
90     by_cases ha : b=a
91     · subst b
92       simp only [weightedTest,ite_true,targetMean,integral_div]
93     · simp [weightedTest,ha]
94   rw [integral_fintype hi.integral_compProd]
95   simp_rw [hinner,smul_eq_mul]
96   simp only [mul_ite,mul_zero,Finset.sum_ite_eq',Finset.mem_univ,ite_true]
97   exact mul_div_cancel₀ _ (ne_of_gt (hc.trans_le (hcover h)))
98 
99 end KernelCalculation
100 
101 section ConditionalExperiment
102 variable {Ω A Y : Type*} [mΩ : MeasurableSpace Ω] [mA : MeasurableSpace A]
103 variable [mY : MeasurableSpace Y] [Fintype A] [DecidableEq A] [MeasurableSingletonClass A]
104 variable {μ : Measure Ω} [IsProbabilityMeasure μ]
105 variable {m : MeasurableSpace Ω} (hm : m≤mΩ)
106 
107 /-- A literal history/action/observation joint law. Its factorization says that
108 the recorded action is drawn from π before drawing from Q at that action.
109 The preaction information m and the full sample space are arbitrary. -/
110 structure Experiment (action : Ω → A) (observation : Ω → Y) where
111   policy : @Kernel Ω A m mA
112   observationKernel : @Kernel (Ω × A) Y (m.prod mA) mY
113   policy_markov : IsMarkovKernel policy
114   observation_markov : IsMarkovKernel observationKernel
115   action_measurable : @Measurable Ω A mΩ mA action
116   observation_measurable : @Measurable Ω Y mΩ mY observation
117   joint_law :
118     @Measure.map Ω (Ω × (A × Y)) mΩ (m.prod (mA.prod mY))
119       (fun ω => (ω,(action ω,observation ω))) μ =
120       (μ.trim hm) ⊗ₘ (policy ⊗ₖ observationKernel)
121 
122 attribute [instance] Experiment.policy_markov Experiment.observation_markov
123 
124 variable {action : Ω → A} {observation : Ω → Y}
125 variable (E : Experiment (mΩ := mΩ) (μ := μ) hm action observation)
126 
127 omit [Fintype A] [DecidableEq A] [MeasurableSingletonClass A] [IsProbabilityMeasure μ] in
128 include hm E in
129 theorem recorded_joint_measurable :
130     @Measurable Ω (Ω × (A × Y)) mΩ (m.prod (mA.prod mY))
131       (fun ω => (ω,(action ω,observation ω))) :=
132   (measurable_id'' hm).prodMk ((Experiment.action_measurable (mΩ := mΩ) E).prodMk (Experiment.observation_measurable (mΩ := mΩ) E))
133 
134 /-- The actual importance-weighted random observation on the given experiment. -/
135 noncomputable def evidence (φ : Ω × Y → ℝ) (a : A) (ω : Ω) : ℝ :=
136   weightedTest (mH := m) (Experiment.policy (mΩ := mΩ) E) φ a (ω,(action ω,observation ω))
137 
138 omit [Fintype A] [IsProbabilityMeasure μ] in
139 theorem evidence_measurable (φ : Ω × Y → ℝ)
140     (hφ : @Measurable (Ω × Y) ℝ (m.prod mY) _ φ) (a : A) :
141     @Measurable Ω ℝ mΩ _ (evidence (mΩ := mΩ) hm E φ a) :=
142   (weightedTest_measurable (mH := m) (Experiment.policy (mΩ := mΩ) E) φ hφ a).comp (recorded_joint_measurable (mΩ := mΩ) hm E)
143 
144 omit [Fintype A] [MeasurableSingletonClass A] [IsProbabilityMeasure μ] in
145 theorem evidence_mem_Icc (φ : Ω × Y → ℝ)
146     (hb : ∀ ω y,φ (ω,y)∈Set.Icc (0:ℝ) 1) (a : A)
147     (c : ℝ) (hc : 0<c) (hcover : ∀ ω,c≤propensity (mH := m) (Experiment.policy (mΩ := mΩ) E) a ω) (ω : Ω) :
148     evidence (mΩ := mΩ) hm E φ a ω∈Set.Icc (0:ℝ) (1/c) :=
149   weightedTest_mem_Icc (mH := m) (Experiment.policy (mΩ := mΩ) E) φ hb a c hc hcover _
150 
151 omit [Fintype A] [IsProbabilityMeasure μ] in
152 /-- Recording the action and observation at the next filtration level makes
153 the importance-weighted evidence adapted at that level. -/
154 theorem evidence_measurable_at {m' : MeasurableSpace Ω} (hmm' : m≤m')
155     (ha : @Measurable Ω A m' mA action) (hy : @Measurable Ω Y m' mY observation)
156     (φ : Ω × Y → ℝ) (hφ : @Measurable (Ω × Y) ℝ (m.prod mY) _ φ) (a : A) :
157     @Measurable Ω ℝ m' _ (evidence (mΩ := mΩ) hm E φ a) :=
158   (weightedTest_measurable (mH := m) (Experiment.policy (mΩ := mΩ) E) φ hφ a).comp
159     ((measurable_id'' hmm').prodMk (ha.prodMk hy))
160 
161 /-- Joint-law factorization and exact finite-action cancellation identify the
162 actual conditional expectation. No conditional-mean equation is an input. -/
163 theorem evidence_condExp (φ : Ω × Y → ℝ)
164     (hφ : @Measurable (Ω × Y) ℝ (m.prod mY) _ φ)
165     (hb : ∀ ω y,φ (ω,y)∈Set.Icc (0:ℝ) 1) (a : A)
166     (c : ℝ) (hc : 0<c) (hcover : ∀ ω,c≤propensity (mH := m) (Experiment.policy (mΩ := mΩ) E) a ω) :
167     μ[evidence (mΩ := mΩ) hm E φ a | m] =ᵐ[μ] targetMean (mH := m) (Experiment.observationKernel (mΩ := mΩ) E) φ a := by
168   let : MeasurableSpace Ω := mΩ
169   have hW := weightedTest_measurable (mH := m) (Experiment.policy (mΩ := mΩ) E) φ hφ a
170   have hWi : Integrable (weightedTest (mH := m) (Experiment.policy (mΩ := mΩ) E) φ a)
171       ((μ.trim hm) ⊗ₘ ((Experiment.policy (mΩ := mΩ) E) ⊗ₖ (Experiment.observationKernel (mΩ := mΩ) E))) :=
172     Integrable.of_mem_Icc (0:ℝ) (1/c) hW.aemeasurable
173       (Filter.Eventually.of_forall (weightedTest_mem_Icc (mH := m) (Experiment.policy (mΩ := mΩ) E) φ hb a c hc hcover))
174   have hevi : Integrable (evidence (mΩ := mΩ) hm E φ a) μ :=
175     Integrable.of_mem_Icc (0:ℝ) (1/c) (evidence_measurable (mΩ := mΩ) hm E φ hφ a).aemeasurable
176       (Filter.Eventually.of_forall (evidence_mem_Icc (mΩ := mΩ) hm E φ hb a c hc hcover))
177   have hg := targetMean_measurable (mH := m) (Experiment.observationKernel (mΩ := mΩ) E) φ hφ a
178   have hgi : Integrable (targetMean (mH := m) (Experiment.observationKernel (mΩ := mΩ) E) φ a) μ :=
179     Integrable.of_mem_Icc (0:ℝ) 1 (hg.mono hm le_rfl).aemeasurable
180       (Filter.Eventually.of_forall (targetMean_mem_Icc (mH := m) (Experiment.observationKernel (mΩ := mΩ) E) φ hφ hb a))
181   apply (ae_eq_condExp_of_forall_setIntegral_eq hm hevi
182     (fun _ _ _ => hgi.integrableOn) ?_ hg.stronglyMeasurable.aestronglyMeasurable).symm
183   intro s hs _
184   calc
185     (∫ ω in s,targetMean (mH := m) (Experiment.observationKernel (mΩ := mΩ) E) φ a ω∂μ)=
186         ∫ ω in s,targetMean (mH := m) (Experiment.observationKernel (mΩ := mΩ) E) φ a ω∂μ.trim hm :=
187       setIntegral_trim hm hg.stronglyMeasurable hs
188     _ = ∫ p in s ×ˢ Set.univ,weightedTest (mH := m) (Experiment.policy (mΩ := mΩ) E) φ a p∂
189         (μ.trim hm) ⊗ₘ ((Experiment.policy (mΩ := mΩ) E) ⊗ₖ (Experiment.observationKernel (mΩ := mΩ) E)) := by
190       rw [Measure.setIntegral_compProd hs MeasurableSet.univ hWi.integrableOn]
191       simp only [Measure.restrict_univ]
192       apply setIntegral_congr_fun hs
193       intro ω _
194       exact (integral_weightedTest (mH := m) (Experiment.policy (mΩ := mΩ) E) (Experiment.observationKernel (mΩ := mΩ) E) φ hφ hb a c hc hcover ω).symm
195     _ = ∫ p in s ×ˢ Set.univ,weightedTest (mH := m) (Experiment.policy (mΩ := mΩ) E) φ a p∂
196         @Measure.map Ω (Ω × (A × Y)) mΩ (m.prod (mA.prod mY))
197           (fun ω => (ω,(action ω,observation ω))) μ := by rw [(Experiment.joint_law (mΩ := mΩ) E)]
198     _ = ∫ ω in s,evidence (mΩ := mΩ) hm E φ a ω∂μ := by
199       let : MeasurableSpace (Ω × (A × Y)) := m.prod (mA.prod mY)
200       rw [setIntegral_map (hs.prod MeasurableSet.univ) hW.aestronglyMeasurable
201         (recorded_joint_measurable (mΩ := mΩ) hm E).aemeasurable]
202       have hpre : (fun ω => (ω,(action ω,observation ω))) ⁻¹' (s ×ˢ Set.univ)=s := by
203         ext ω
204         simp
205       rw [hpre]
206       rfl
207 
208 end ConditionalExperiment
209 end Komorebi.LSM.CoveredAdaptiveActions

Back to source controls