diff --git a/LMLExplore/MarkovDecisionProcess/Episodic/EpisodeEnv.lean b/LMLExplore/MarkovDecisionProcess/Episodic/EpisodeEnv.lean new file mode 100644 index 00000000..fb3c1e2d --- /dev/null +++ b/LMLExplore/MarkovDecisionProcess/Episodic/EpisodeEnv.lean @@ -0,0 +1,427 @@ +/- +Copyright (c) 2026 LMLExplore contributors. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: LMLExplore contributors +-/ +module + +public import LMLExplore.MarkovDecisionProcess.Episodic.Policy +public import LeanMachineLearning.SequentialLearning.IdentificationAlg +public import LeanMachineLearning.SequentialLearning.StationaryEnv + +/-! +# Episodes as rounds + +In online episodic learning, each round of the interaction is a whole episode of an episodic MDP +`M : EpisodicMDP 𝓒 𝓐` with horizon `H`: at the start of every episode, the learner chooses an +`H`-step policy from the previous episodes, plays it from an initial state drawn from the initial +law `ΞΌ`, and observes the episode. This file builds the corresponding environments of LML's +interaction protocol. + +An episode played with the Markov policy `Ο€ : β„• β†’ Kernel 𝓒 𝓐` is the run of `augMarkovPolicy Ο€` +in the time-augmented MDP `M.augmented`, started from `ΞΌ` at the step `0`. Its law is a measure on +trajectories of `M.augmented`; the states `s_0, …, s_H` of an episode are `episodeStates H` of the +trajectory and its rewards `r_0, …, r_{H-1}` are `episodeRewards H`. The law of the state sequence +of an episode is `M.statesLaw H ΞΌ Ο€`, and the joint law of the state sequence and the reward +sequence is `M.episodeLaw H ΞΌ Ο€`. + +Under `M.statesLaw H ΞΌ Ο€`, the first state has law `ΞΌ`, and the state of the step `h + 1` has +conditional law `∫ P_h(Β· | s_h, a) dΟ€_h(a | s_h)` given the states `s_0, …, s_h` +(`hasCondDistrib_statesLaw`). This comes from the same property along every run of `Ο€` +(`hasCondDistrib_obs_fst_add_one_states`). Under `M.episodeLaw H ΞΌ Ο€`, the expected sum of the +rewards is the value `∫ V^Ο€_0 dΞΌ` of `Ο€` (`integral_sum_episodeLaw`). + +For a finite state space and a countable action space, the `H`-step policies form a countable +discrete type. The episode-level environments are then the stationary environments, without +observations, of the kernels from `H`-step policies to the laws of an episode: +* `M.statesEnv H ΞΌ`: the feedback of a round is the state sequence of the episode. It models + learners that know the rewards, as when the rewards are deterministic; +* `M.episodeEnv H ΞΌ`: the feedback of a round is the pair of the state sequence and the reward + sequence. A learner in it also learns unknown stochastic rewards. + +Forgetting the rewards maps the episode laws of `M.episodeEnv H ΞΌ` to those of `M.statesEnv H ΞΌ` +(`map_fst_episodeKernel`). A best policy identification algorithm is an LML identification +algorithm in `M.statesEnv H ΞΌ` with an `H`-step policy as output (`BPIAlg`). + +The definitions `statesLaw`, `statesKernel`, `statesEnv` and `BPIAlg` are adapted, with their +names and namespace, from the formalization entropy-mdp (`github.com/RemyDegenne/entropy-mdp`, +file `Essakine2026Tight/ForLMLExplore/MarkovDecisionProcess/Episodic/Basic.lean`, by RΓ©my +Degenne, Apache 2.0) of [EV26], from a fixed initial state to an initial law, and from +deterministic to Markov policies for `statesLaw`. + +## Main definitions + +* `Learning.MDP.Episodic.episodeStates H`, `Learning.MDP.Episodic.episodeRewards H`: the state + and reward sequences of an episode of horizon `H`. +* `Learning.MDP.Episodic.EpisodicMDP.statesLaw M H ΞΌ Ο€`: the law of the state sequence of an + episode of the Markov policy `Ο€`. +* `Learning.MDP.Episodic.EpisodicMDP.episodeLaw M H ΞΌ Ο€`: the joint law of the state and reward + sequences of an episode of `Ο€`. +* `Learning.MDP.Episodic.EpisodicMDP.statesKernel M H ΞΌ`, + `Learning.MDP.Episodic.EpisodicMDP.episodeKernel M H ΞΌ`: the same laws as kernels from + `H`-step policies. +* `Learning.MDP.Episodic.EpisodicMDP.statesEnv M H ΞΌ`, + `Learning.MDP.Episodic.EpisodicMDP.episodeEnv M H ΞΌ`: the episode-level environments. +* `Learning.MDP.Episodic.BPIAlg 𝓒 𝓐 H`: the best policy identification algorithms. + +## Main statements + +* `Learning.MDP.Episodic.EpisodicMDP.hasCondDistrib_obs_fst_add_one_states`: along a run of the + Markov policy `Ο€` in `M.augmented` started at the step `hβ‚€`, the state of round `n + 1` has + conditional law `∫ P_{hβ‚€+n}(Β· | s_n, a) dΟ€_{hβ‚€+n}(a | s_n)` given the states of the rounds + `0, …, n`. +* `Learning.MDP.Episodic.EpisodicMDP.hasLaw_statesLaw`, + `Learning.MDP.Episodic.EpisodicMDP.hasLaw_episodeLaw`: the episodes of every run have these + laws. +* `Learning.MDP.Episodic.EpisodicMDP.hasLaw_statesLaw_zero`, + `Learning.MDP.Episodic.EpisodicMDP.hasCondDistrib_statesLaw`: the law of the first state and + the conditional law of each next state under `M.statesLaw H ΞΌ Ο€`. +* `Learning.MDP.Episodic.EpisodicMDP.integral_sum_episodeLaw`: the expected return of an episode + is the value `∫ V^Ο€_0 dΞΌ`. +* `Learning.MDP.Episodic.EpisodicMDP.map_fst_episodeKernel`: forgetting the rewards maps the + episode laws of `M.episodeEnv H ΞΌ` to those of `M.statesEnv H ΞΌ`. +* `Learning.MDP.Episodic.EpisodicMDP.hasCondDistrib_feedback_statesEnv`, + `Learning.MDP.Episodic.EpisodicMDP.hasCondDistrib_feedback_episodeEnv`: in a run of these + environments, the episode of round `t` has the law of an episode of the policy of round `t`. + +## References + +* [EV26] A. Essakine and C. Vernade, *Tight sample complexity bounds for entropic best policy + identification*, COLT 2026, arXiv:2605.13717. +-/ + +public section + +open MeasureTheory ProbabilityTheory Learning Finset + +namespace Learning.MDP.Episodic + +variable {𝓒 𝓐 : Type*} {m𝓒 : MeasurableSpace 𝓒} {m𝓐 : MeasurableSpace 𝓐} + +/-! ### The states and rewards of an episode -/ + +/-- The states `s_0, …, s_H` of the first `H + 1` rounds of a trajectory of the time-augmented +MDP. -/ +def episodeStates (H : β„•) (x : β„• β†’ Round (𝓒 Γ— β„•) 𝓐 ℝ) : Fin (H + 1) β†’ 𝓒 := + fun h ↦ (IT.obs h x).1 + +@[simp] +lemma episodeStates_apply (H : β„•) (x : β„• β†’ Round (𝓒 Γ— β„•) 𝓐 ℝ) (h : Fin (H + 1)) : + episodeStates H x h = (IT.obs h x).1 := + (rfl) + +@[fun_prop] +lemma measurable_episodeStates (H : β„•) : Measurable (episodeStates (𝓒 := 𝓒) (𝓐 := 𝓐) H) := + Measurable.of_eval fun h ↦ (IT.measurable_obs h).fst + +/-- The rewards `r_0, …, r_{H-1}` of the first `H` rounds of a trajectory of the time-augmented +MDP. -/ +def episodeRewards (H : β„•) (x : β„• β†’ Round (𝓒 Γ— β„•) 𝓐 ℝ) : Fin H β†’ ℝ := + fun h ↦ IT.feedback h x + +@[simp] +lemma episodeRewards_apply (H : β„•) (x : β„• β†’ Round (𝓒 Γ— β„•) 𝓐 ℝ) (h : Fin H) : + episodeRewards H x h = IT.feedback h x := + (rfl) + +@[fun_prop] +lemma measurable_episodeRewards (H : β„•) : Measurable (episodeRewards (𝓒 := 𝓒) (𝓐 := 𝓐) H) := + Measurable.of_eval fun h ↦ IT.measurable_feedback h + +namespace EpisodicMDP + +/-! ### The next state given the past states, along a run -/ + +/-- The states of a history of `n` rounds followed by the state `s` of round `n`. -/ +private def snocStates {n : β„•} (p : Hist (𝓒 Γ— β„•) 𝓐 ℝ n Γ— 𝓒) : Fin (n + 1) β†’ 𝓒 := + Fin.snoc (Ξ± := fun _ ↦ 𝓒) (fun i ↦ (p.1 i).obs.1) p.2 + +private lemma measurable_snocStates (n : β„•) : + Measurable (snocStates (𝓒 := 𝓒) (𝓐 := 𝓐) (n := n)) := by + refine Measurable.of_eval fun i ↦ ?_ + induction i using Fin.lastCases with + | last => simpa [snocStates] using measurable_snd + | cast i => simp only [snocStates, Fin.snoc_castSucc]; fun_prop + +/-- The kernel of the next augmented state given the past and the action of round `n`, composed +with the policy and restricted to the states at the step `k`, is the kernel of the next state +given the past states. -/ +private lemma comap_map_fst_snd_compProd_augmented (M : EpisodicMDP 𝓒 𝓐) (Ο€ : β„• β†’ Kernel 𝓒 𝓐) + [βˆ€ h, IsMarkovKernel (Ο€ h)] (n k : β„•) : + ((Kernel.prodMkLeft (Hist (𝓒 Γ— β„•) 𝓐 ℝ n) (augMarkovPolicy Ο€) βŠ—β‚– + M.augmented.P.comap (fun p ↦ (p.1.2, p.2)) (by fun_prop)).snd.map Prod.fst).comap + (fun p : Hist (𝓒 Γ— β„•) 𝓐 ℝ n Γ— 𝓒 ↦ (p.1, (p.2, k))) (by fun_prop) = + ((M.trans k βˆ˜β‚– (Kernel.id Γ—β‚– Ο€ k)).comap (fun Ο„ ↦ Ο„ (Fin.last n)) + (measurable_pi_apply _)).comap (snocStates (𝓐 := 𝓐)) (measurable_snocStates n) := by + ext p t ht + -- unfold the composite kernels at `p`, then integrate against the Dirac mass at `p.2` + have := (M.trans k).measurable_coe ht + simp (disch := measurability) [Kernel.map_apply', Kernel.snd_apply', Kernel.compProd_apply, + Kernel.comp_apply', Kernel.prod_apply, Kernel.id_apply, Measure.dirac_prod, lintegral_map, + snocStates, Set.preimage_preimage] + +section Runs + +variable [Nonempty 𝓒] {M : EpisodicMDP 𝓒 𝓐} {Ο€ : β„• β†’ Kernel 𝓒 𝓐} [βˆ€ h, IsMarkovKernel (Ο€ h)] + {Ξ© : Type*} {mΞ© : MeasurableSpace Ξ©} {P : Measure Ξ©} [IsProbabilityMeasure P] + {O : β„• β†’ Ξ© β†’ 𝓒 Γ— β„•} {A : β„• β†’ Ξ© β†’ 𝓐} {Y : β„• β†’ Ξ© β†’ ℝ} + +/-- **Markov property of the states**: along a run of the Markov policy `Ο€` in the augmented MDP +started at the step `hβ‚€`, the state of round `n + 1` has conditional law +`∫ P_{hβ‚€+n}(Β· | s_n, a) dΟ€_{hβ‚€+n}(a | s_n)` given the states `s_0, …, s_n` of the rounds +`0, …, n`. -/ +theorem hasCondDistrib_obs_fst_add_one_states {ΞΌ : Measure (𝓒 Γ— β„•)} [IsProbabilityMeasure ΞΌ] + {hβ‚€ : β„•} (hΞΌ : βˆ€α΅ p βˆ‚ΞΌ, p.2 = hβ‚€) + (h : IsAlgEnvSeq O A Y (Algorithm.markov fun _ ↦ augMarkovPolicy Ο€) (M.augmented.env ΞΌ) P) + (n : β„•) : + HasCondDistrib (fun Ο‰ ↦ (O (n + 1) Ο‰).1) (fun Ο‰ (i : Fin (n + 1)) ↦ (O i Ο‰).1) + ((M.trans (hβ‚€ + n) βˆ˜β‚– (Kernel.id Γ—β‚– Ο€ (hβ‚€ + n))).comap (fun Ο„ ↦ Ο„ (Fin.last n)) + (measurable_pi_apply _)) P := by + -- the next augmented state given the history, the state and the action of round `n` + have hf : Measurable fun x : Hist (𝓒 Γ— β„•) 𝓐 ℝ (n + 1) ↦ + ((fun i : Fin n ↦ x i.castSucc, (x (Fin.last n)).obs), (x (Fin.last n)).action) := by + fun_prop + have hO : HasCondDistrib (O (n + 1)) (fun Ο‰ ↦ ((history O A Y n Ο‰, O n Ο‰), A n Ο‰)) + (M.augmented.P.comap (fun p ↦ (p.1.2, p.2)) (by fun_prop)) P := + HasCondDistrib.comp_right (Z := history O A Y (n + 1)) + (ΞΊ := M.augmented.P.comap (fun p ↦ (p.1.2, p.2)) (by fun_prop)) (hf := hf) + (MDP.hasCondDistrib_obs_add_one_history h n) + -- integrate out the action, drawn from `Ο€` given the history and the state of round `n` + have hS := ((h.hasCondDistrib_action_history_obs_markov n).compProd hO).snd.comp_left + (f := Prod.fst) measurable_fst + -- the step of round `n` is almost surely `hβ‚€ + n` + have hae : (fun Ο‰ ↦ (history O A Y n Ο‰, O n Ο‰)) =ᡐ[P] + (fun p : Hist (𝓒 Γ— β„•) 𝓐 ℝ n Γ— 𝓒 ↦ (p.1, (p.2, hβ‚€ + n))) ∘ + fun Ο‰ ↦ (history O A Y n Ο‰, (O n Ο‰).1) := by + filter_upwards [ae_obs_snd_eq hΞΌ h n] with Ο‰ hΟ‰ + simp [← hΟ‰] + have he : MeasurableEmbedding fun p : Hist (𝓒 Γ— β„•) 𝓐 ℝ n Γ— 𝓒 ↦ (p.1, (p.2, hβ‚€ + n)) := + MeasurableEmbedding.id.prodMap (measurableEmbedding_prod_mk_right (hβ‚€ + n)) + have hS' := (hS.congr hae.symm .rfl).of_measurableEmbedding_comp_right he + rw [comap_map_fst_snd_compProd_augmented] at hS' + -- the resulting kernel only depends on the states of the rounds `0, …, n` + refine (HasCondDistrib.comp_right hS').congr (.of_forall fun Ο‰ ↦ ?_) .rfl + funext i + induction i using Fin.lastCases with + | last => simp [snocStates] + | cast i => simp [snocStates, history_apply] + +end Runs + +/-! ### The laws of an episode -/ + +variable (M : EpisodicMDP 𝓒 𝓐) (H : β„•) (ΞΌ : Measure 𝓒) + +/-- Under the initial law `ΞΌ` placed at the step `0`, the step is almost surely `0`. -/ +private lemma ae_snd_map_prodMk_zero : βˆ€α΅ p βˆ‚(ΞΌ.map fun s ↦ (s, 0)), p.2 = 0 := + (ae_map_iff measurable_prodMk_right.aemeasurable + (measurable_snd (measurableSet_singleton 0))).2 (.of_forall fun _ ↦ rfl) + +section Laws + +variable [Nonempty 𝓒] (Ο€ : β„• β†’ Kernel 𝓒 𝓐) [βˆ€ h, IsMarkovKernel (Ο€ h)] + +/-- The trajectories of the episodes of `Ο€` from `ΞΌ` form a run of `Ο€` in the augmented MDP. -/ +private lemma isAlgEnvSeq_episode : + IsAlgEnvSeq IT.obs IT.action IT.feedback (Algorithm.markov fun _ ↦ augMarkovPolicy Ο€) + (M.augmented.env (ΞΌ.map fun s ↦ (s, 0))) + (trajMeasure (Algorithm.markov fun _ ↦ augMarkovPolicy Ο€) + (M.augmented.env (ΞΌ.map fun s ↦ (s, 0)))) := + IT.isAlgEnvSeq_trajMeasure _ _ + +/-- The law of the state sequence `s_0, …, s_H` of an episode of horizon `H` played with the +Markov policy `Ο€` from the initial law `ΞΌ`. -/ +noncomputable def statesLaw : Measure (Fin (H + 1) β†’ 𝓒) := + (trajMeasure (Algorithm.markov fun _ ↦ augMarkovPolicy Ο€) + (M.augmented.env (ΞΌ.map fun s ↦ (s, 0)))).map (episodeStates H) + +/-- The joint law of the state sequence `s_0, …, s_H` and the reward sequence +`r_0, …, r_{H-1}` of an episode of horizon `H` played with the Markov policy `Ο€` from the initial +law `ΞΌ`. -/ +noncomputable def episodeLaw : Measure ((Fin (H + 1) β†’ 𝓒) Γ— (Fin H β†’ ℝ)) := + (trajMeasure (Algorithm.markov fun _ ↦ augMarkovPolicy Ο€) + (M.augmented.env (ΞΌ.map fun s ↦ (s, 0)))).map fun x ↦ (episodeStates H x, episodeRewards H x) + +instance : IsProbabilityMeasure (M.statesLaw H ΞΌ Ο€) := + (Measure.isProbabilityMeasure_map_iff (measurable_episodeStates H).aemeasurable).2 inferInstance + +instance : IsProbabilityMeasure (M.episodeLaw H ΞΌ Ο€) := + (Measure.isProbabilityMeasure_map_iff + ((measurable_episodeStates H).prodMk (measurable_episodeRewards H)).aemeasurable).2 + inferInstance + +/-- Forgetting the rewards maps the law of an episode to the law of its state sequence. -/ +@[simp] +lemma map_fst_episodeLaw : (M.episodeLaw H ΞΌ Ο€).map Prod.fst = M.statesLaw H ΞΌ Ο€ := by + rw [episodeLaw, Measure.map_map measurable_fst (by fun_prop), statesLaw] + simp [Function.comp_def] + +variable {M H ΞΌ Ο€} + +/-- Along every run of the Markov policy `Ο€` in the augmented MDP started from `ΞΌ` at the step +`0`, the states of the rounds `0, …, H` have law `M.statesLaw H ΞΌ Ο€`. -/ +lemma hasLaw_statesLaw {Ξ© : Type*} {mΞ© : MeasurableSpace Ξ©} {P : Measure Ξ©} + [IsProbabilityMeasure P] {O : β„• β†’ Ξ© β†’ 𝓒 Γ— β„•} {A : β„• β†’ Ξ© β†’ 𝓐} {Y : β„• β†’ Ξ© β†’ ℝ} + (h : IsAlgEnvSeq O A Y (Algorithm.markov fun _ ↦ augMarkovPolicy Ο€) + (M.augmented.env (ΞΌ.map fun s ↦ (s, 0))) P) : + HasLaw (fun Ο‰ (i : Fin (H + 1)) ↦ (O i Ο‰).1) (M.statesLaw H ΞΌ Ο€) P := + -- `episodeStates H ∘ trajectory O A Y` is the function of the statement by definition + HasLaw.comp ⟨(measurable_episodeStates H).aemeasurable, rfl⟩ h.hasLaw_trajectory + +/-- Along every run of the Markov policy `Ο€` in the augmented MDP started from `ΞΌ` at the step +`0`, the states of the rounds `0, …, H` and the rewards of the rounds `0, …, H - 1` have law +`M.episodeLaw H ΞΌ Ο€`. -/ +lemma hasLaw_episodeLaw {Ξ© : Type*} {mΞ© : MeasurableSpace Ξ©} {P : Measure Ξ©} + [IsProbabilityMeasure P] {O : β„• β†’ Ξ© β†’ 𝓒 Γ— β„•} {A : β„• β†’ Ξ© β†’ 𝓐} {Y : β„• β†’ Ξ© β†’ ℝ} + (h : IsAlgEnvSeq O A Y (Algorithm.markov fun _ ↦ augMarkovPolicy Ο€) + (M.augmented.env (ΞΌ.map fun s ↦ (s, 0))) P) : + HasLaw (fun Ο‰ ↦ (fun i : Fin (H + 1) ↦ (O i Ο‰).1, fun i : Fin H ↦ Y i Ο‰)) + (M.episodeLaw H ΞΌ Ο€) P := + -- the composite with `trajectory O A Y` is the function of the statement by definition + HasLaw.comp ⟨((measurable_episodeStates H).prodMk (measurable_episodeRewards H)).aemeasurable, + rfl⟩ h.hasLaw_trajectory + +variable [IsProbabilityMeasure ΞΌ] + +/-- Under the law of the state sequence of an episode, the first state has law `ΞΌ`. -/ +lemma hasLaw_statesLaw_zero : HasLaw (fun Ο„ ↦ Ο„ 0) ΞΌ (M.statesLaw H ΞΌ Ο€) := by + have h0 := MDP.hasLaw_obs_zero (isAlgEnvSeq_episode M ΞΌ Ο€) + refine ⟨(measurable_pi_apply 0).aemeasurable, ?_⟩ + have he : (fun Ο„ : Fin (H + 1) β†’ 𝓒 ↦ Ο„ 0) ∘ episodeStates (𝓐 := 𝓐) H = Prod.fst ∘ IT.obs 0 := by + funext x + simp + rw [statesLaw, Measure.map_map (measurable_pi_apply 0) (measurable_episodeStates H), he, + ← Measure.map_map measurable_fst (IT.measurable_obs 0), h0.map_eq, + Measure.map_map measurable_fst measurable_prodMk_right] + exact Measure.map_id + +/-- **Markov property of the states of an episode**: under the law of the state sequence of an +episode of the Markov policy `Ο€`, the state of the step `h + 1` has conditional law +`∫ P_h(Β· | s_h, a) dΟ€_h(a | s_h)` given the states `s_0, …, s_h`. -/ +theorem hasCondDistrib_statesLaw (h : Fin H) : + HasCondDistrib (fun Ο„ : Fin (H + 1) β†’ 𝓒 ↦ Ο„ h.succ) + (fun Ο„ (i : Fin (h + 1)) ↦ Ο„ (Fin.castLE (Nat.succ_le_succ h.isLt.le) i)) + ((M.trans h βˆ˜β‚– (Kernel.id Γ—β‚– Ο€ h)).comap (fun Ο„ ↦ Ο„ (Fin.last h)) + (measurable_pi_apply _)) (M.statesLaw H ΞΌ Ο€) := by + have hc := hasCondDistrib_obs_fst_add_one_states (ae_snd_map_prodMk_zero ΞΌ) + (isAlgEnvSeq_episode M ΞΌ Ο€) h + rw [zero_add] at hc + have hX : Measurable fun (Ο„ : Fin (H + 1) β†’ 𝓒) (i : Fin (h + 1)) ↦ + Ο„ (Fin.castLE (Nat.succ_le_succ h.isLt.le) i) := + Measurable.of_eval fun _ ↦ measurable_pi_apply _ + -- transfer the conditional law from the trajectories to their state sequences + refine ⟨(hX.prodMk (measurable_pi_apply _)).aemeasurable, ?_⟩ + rw [statesLaw, Measure.map_map hX (measurable_episodeStates H), + Measure.map_map (hX.prodMk (measurable_pi_apply _)) (measurable_episodeStates H)] + convert hc.map_eq using 3 <;> simp [Function.comp_def] + +/-- **The expected return of an episode** is the value `∫ V^Ο€_0 dΞΌ` of the Markov policy `Ο€`, +for rewards bounded by `C`. -/ +theorem integral_sum_episodeLaw {C : ℝ} (hM : M.RewardsIn (Set.Icc (-C) C)) : + ∫ e, βˆ‘ h, e.2 h βˆ‚(M.episodeLaw H ΞΌ Ο€) = ∫ s, M.value H Ο€ 0 s βˆ‚ΞΌ := by + have hv := integral_sum_feedback_eq_integral_value (H := H) hM (ae_snd_map_prodMk_zero ΞΌ) + (isAlgEnvSeq_episode M ΞΌ Ο€) + have hmap : ∫ p : 𝓒 Γ— β„•, M.value H Ο€ 0 p.1 βˆ‚(ΞΌ.map fun s ↦ (s, 0)) = + ∫ s, M.value H Ο€ 0 s βˆ‚ΞΌ := + integral_map measurable_prodMk_right.aemeasurable + ((measurable_value 0).comp measurable_fst).aestronglyMeasurable + rw [Nat.sub_zero, hmap] at hv + rw [episodeLaw, integral_map (by fun_prop) (by fun_prop), ← hv] + congr 1 with x + exact Fin.sum_univ_eq_sum_range (fun n ↦ IT.feedback n x) H + +end Laws + +/-! ### The episode-level environments -/ + +section EpisodeEnv + +variable [Nonempty 𝓒] [Finite 𝓒] [Countable 𝓐] [MeasurableSingletonClass 𝓒] + [MeasurableSingletonClass 𝓐] [Nonempty 𝓐] + +/-- The kernel from `H`-step policies to the law of the state sequence of an episode played with +the policy from the initial law `ΞΌ`. -/ +noncomputable def statesKernel : Kernel (Policy 𝓒 𝓐 H) (Fin (H + 1) β†’ 𝓒) where + toFun Ο€ := M.statesLaw H ΞΌ (Ο€.toKernel fun _ ↦ measurable_of_countable _) + measurable' := measurable_of_countable _ + +@[simp] +lemma statesKernel_apply (Ο€ : Policy 𝓒 𝓐 H) : + M.statesKernel H ΞΌ Ο€ = M.statesLaw H ΞΌ (Ο€.toKernel fun _ ↦ measurable_of_countable _) := + (rfl) + +instance : IsMarkovKernel (M.statesKernel H ΞΌ) := + ⟨fun Ο€ ↦ by rw [statesKernel_apply]; infer_instance⟩ + +/-- The kernel from `H`-step policies to the joint law of the state sequence and the reward +sequence of an episode played with the policy from the initial law `ΞΌ`. -/ +noncomputable def episodeKernel : + Kernel (Policy 𝓒 𝓐 H) ((Fin (H + 1) β†’ 𝓒) Γ— (Fin H β†’ ℝ)) where + toFun Ο€ := M.episodeLaw H ΞΌ (Ο€.toKernel fun _ ↦ measurable_of_countable _) + measurable' := measurable_of_countable _ + +@[simp] +lemma episodeKernel_apply (Ο€ : Policy 𝓒 𝓐 H) : + M.episodeKernel H ΞΌ Ο€ = M.episodeLaw H ΞΌ (Ο€.toKernel fun _ ↦ measurable_of_countable _) := + (rfl) + +instance : IsMarkovKernel (M.episodeKernel H ΞΌ) := + ⟨fun Ο€ ↦ by rw [episodeKernel_apply]; infer_instance⟩ + +/-- Forgetting the rewards maps the episode laws of `M.episodeEnv H ΞΌ` to those of +`M.statesEnv H ΞΌ`. -/ +@[simp] +lemma map_fst_episodeKernel : (M.episodeKernel H ΞΌ).map Prod.fst = M.statesKernel H ΞΌ := by + ext Ο€ : 1 + rw [Kernel.map_apply _ measurable_fst, episodeKernel_apply, map_fst_episodeLaw, + statesKernel_apply] + +/-- The episode-level environment of the episodic MDP `M` with horizon `H` and initial law `ΞΌ`, +in which the learner observes the state sequences only, the rewards being known: at each round, +the learner plays an `H`-step policy and observes the state sequence of the episode. -/ +noncomputable def statesEnv : Environment Unit (Policy 𝓒 𝓐 H) (Fin (H + 1) β†’ 𝓒) := + Environment.bandit (M.statesKernel H ΞΌ) + +/-- The episode-level environment of the episodic MDP `M` with horizon `H` and initial law `ΞΌ`, +in which the learner observes the state sequences and the reward sequences: at each round, the +learner plays an `H`-step policy and observes the states and the rewards of the episode. -/ +noncomputable def episodeEnv : + Environment Unit (Policy 𝓒 𝓐 H) ((Fin (H + 1) β†’ 𝓒) Γ— (Fin H β†’ ℝ)) := + Environment.bandit (M.episodeKernel H ΞΌ) + +lemma statesEnv_def : M.statesEnv H ΞΌ = Environment.bandit (M.statesKernel H ΞΌ) := + (rfl) + +lemma episodeEnv_def : M.episodeEnv H ΞΌ = Environment.bandit (M.episodeKernel H ΞΌ) := + (rfl) + +variable {M H ΞΌ} {Ξ© : Type*} {mΞ© : MeasurableSpace Ξ©} {P : Measure Ξ©} [IsFiniteMeasure P] + {O : β„• β†’ Ξ© β†’ Unit} {X : β„• β†’ Ξ© β†’ Policy 𝓒 𝓐 H} + +/-- In a run of the episode-level environment `M.statesEnv H ΞΌ`, the state sequence of the +episode `t` has the law of an episode of the policy `X t` of the episode `t`. -/ +lemma hasCondDistrib_feedback_statesEnv {Y : β„• β†’ Ξ© β†’ Fin (H + 1) β†’ 𝓒} + {alg : Algorithm Unit (Policy 𝓒 𝓐 H) (Fin (H + 1) β†’ 𝓒)} + (h : IsAlgEnvSeq O X Y alg (M.statesEnv H ΞΌ) P) (t : β„•) : + HasCondDistrib (Y t) (X t) (M.statesKernel H ΞΌ) P := + h.hasCondDistrib_feedback_bandit t + +/-- In a run of the episode-level environment `M.episodeEnv H ΞΌ`, the state and reward sequences +of the episode `t` have the law of an episode of the policy `X t` of the episode `t`. -/ +lemma hasCondDistrib_feedback_episodeEnv {Y : β„• β†’ Ξ© β†’ (Fin (H + 1) β†’ 𝓒) Γ— (Fin H β†’ ℝ)} + {alg : Algorithm Unit (Policy 𝓒 𝓐 H) ((Fin (H + 1) β†’ 𝓒) Γ— (Fin H β†’ ℝ))} + (h : IsAlgEnvSeq O X Y alg (M.episodeEnv H ΞΌ) P) (t : β„•) : + HasCondDistrib (Y t) (X t) (M.episodeKernel H ΞΌ) P := + h.hasCondDistrib_feedback_bandit t + +end EpisodeEnv + +end EpisodicMDP + +/-- A best policy identification algorithm with horizon `H`: an identification algorithm that +plays an `H`-step policy at each episode, observes the state sequence of the episode, and +outputs an `H`-step policy. -/ +abbrev BPIAlg (𝓒 𝓐 : Type*) [MeasurableSpace 𝓒] [MeasurableSpace 𝓐] (H : β„•) := + IdentAlg Unit (Policy 𝓒 𝓐 H) (Fin (H + 1) β†’ 𝓒) (Policy 𝓒 𝓐 H) + +end Learning.MDP.Episodic diff --git a/LMLExplore/MarkovDecisionProcess/Episodic/Policy.lean b/LMLExplore/MarkovDecisionProcess/Episodic/Policy.lean new file mode 100644 index 00000000..17d6d810 --- /dev/null +++ b/LMLExplore/MarkovDecisionProcess/Episodic/Policy.lean @@ -0,0 +1,128 @@ +/- +Copyright (c) 2026 LMLExplore contributors. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: LMLExplore contributors +-/ +module + +public import LMLExplore.MarkovDecisionProcess.Episodic.Value + +/-! +# Deterministic `H`-step policies of an episodic MDP + +A deterministic `H`-step policy `Ο€ : Policy 𝓒 𝓐 H = Fin H β†’ 𝓒 β†’ 𝓐` plays the action `Ο€ h s` at +the step `h < H` in the state `s`. As for stochastic processes, it is not bundled with its +measurability: the statements that need it take the hypothesis `hΟ€ : βˆ€ h, Measurable (Ο€ h)`, +which holds for every policy when the state space is countable and discrete. + +When the action space is nonempty, `Policy.extend Ο€ : β„• β†’ 𝓒 β†’ 𝓐` extends `Ο€` to all steps by a +fixed action beyond the horizon, and `Policy.toKernel Ο€ hΟ€ : β„• β†’ Kernel 𝓒 𝓐` is the +corresponding deterministic Markov policy, to which the values of +`LMLExplore/MarkovDecisionProcess/Episodic/Value.lean` apply. The steps `h β‰₯ H` never matter for +an episode of horizon `H`. + +The definitions `Policy` and `Policy.extend` are adapted, with their names and namespace, from +the formalization entropy-mdp (`github.com/RemyDegenne/entropy-mdp`, file +`Essakine2026Tight/ForLMLExplore/MarkovDecisionProcess/Episodic/Basic.lean`, by RΓ©my Degenne, +Apache 2.0) of [EV26]. + +## Main definitions + +* `Learning.MDP.Episodic.Policy 𝓒 𝓐 H`: the deterministic `H`-step policies. +* `Learning.MDP.Episodic.Policy.extend Ο€`: the policy `Ο€` extended beyond the horizon. +* `Learning.MDP.Episodic.Policy.toKernel Ο€ hΟ€`: the deterministic Markov policy of `Ο€`. + +## Main statements + +* `Learning.MDP.Episodic.Policy.extend_injective`: extension beyond the horizon is injective. +* `Learning.MDP.Episodic.Policy.measurable_extend`: the extension of a measurable policy is + measurable at every step. +* `Learning.MDP.Episodic.EpisodicMDP.value_toKernel_fin`: before the horizon, the value of + `Ο€` is its action value at the action it plays. + +## References + +* [EV26] A. Essakine and C. Vernade, *Tight sample complexity bounds for entropic best policy + identification*, COLT 2026, arXiv:2605.13717. +-/ + +public section + +open MeasureTheory ProbabilityTheory + +namespace Learning.MDP.Episodic + +variable {𝓒 𝓐 : Type*} {H : β„•} + +/-- A deterministic `H`-step policy: the action `Ο€ h s` played at the step `h : Fin H` in the +state `s`. -/ +abbrev Policy (𝓒 𝓐 : Type*) (H : β„•) := Fin H β†’ 𝓒 β†’ 𝓐 + +namespace Policy + +variable [Nonempty 𝓐] + +/-- The `H`-step policy `Ο€` extended to all steps by a fixed action beyond the horizon. -/ +noncomputable def extend (Ο€ : Policy 𝓒 𝓐 H) (h : β„•) (s : 𝓒) : 𝓐 := + if hh : h < H then Ο€ ⟨h, hh⟩ s else Classical.arbitrary 𝓐 + +@[simp] +lemma extend_of_lt (Ο€ : Policy 𝓒 𝓐 H) {h : β„•} (hh : h < H) (s : 𝓒) : + Ο€.extend h s = Ο€ ⟨h, hh⟩ s := by + simp [extend, hh] + +lemma extend_fin (Ο€ : Policy 𝓒 𝓐 H) (h : Fin H) (s : 𝓒) : Ο€.extend h s = Ο€ h s := + Ο€.extend_of_lt h.2 s + +lemma extend_of_le (Ο€ : Policy 𝓒 𝓐 H) {h : β„•} (hh : H ≀ h) (s : 𝓒) : + Ο€.extend h s = Classical.arbitrary 𝓐 := by + simp [extend, not_lt.2 hh] + +lemma extend_injective : Function.Injective (extend : Policy 𝓒 𝓐 H β†’ β„• β†’ 𝓒 β†’ 𝓐) := by + intro Ο€ Ο€' h + funext k s + rw [← Ο€.extend_fin k s, ← Ο€'.extend_fin k s, h] + +variable [MeasurableSpace 𝓒] [MeasurableSpace 𝓐] + +/-- The extension of a policy that is measurable at every step is measurable at every step. -/ +lemma measurable_extend {Ο€ : Policy 𝓒 𝓐 H} (hΟ€ : βˆ€ h, Measurable (Ο€ h)) (h : β„•) : + Measurable (Ο€.extend h) := by + by_cases hh : h < H + Β· rw [show Ο€.extend h = Ο€ ⟨h, hh⟩ from funext (Ο€.extend_of_lt hh)] + exact hΟ€ ⟨h, hh⟩ + Β· rw [show Ο€.extend h = fun _ ↦ Classical.arbitrary 𝓐 from funext (Ο€.extend_of_le (not_lt.1 hh))] + exact measurable_const + +/-- The deterministic Markov policy of the `H`-step policy `Ο€`, measurable at every step: at the +step `h`, the Dirac mass at the action `Ο€.extend h s`. -/ +noncomputable def toKernel (Ο€ : Policy 𝓒 𝓐 H) (hΟ€ : βˆ€ h, Measurable (Ο€ h)) (h : β„•) : + Kernel 𝓒 𝓐 := + Kernel.deterministic (Ο€.extend h) (measurable_extend hΟ€ h) + +@[simp] +lemma toKernel_apply (Ο€ : Policy 𝓒 𝓐 H) (hΟ€ : βˆ€ h, Measurable (Ο€ h)) (h : β„•) (s : 𝓒) : + Ο€.toKernel hΟ€ h s = Measure.dirac (Ο€.extend h s) := + Kernel.deterministic_apply _ _ + +instance (Ο€ : Policy 𝓒 𝓐 H) (hΟ€ : βˆ€ h, Measurable (Ο€ h)) (h : β„•) : + IsMarkovKernel (Ο€.toKernel hΟ€ h) := by + unfold toKernel; infer_instance + +end Policy + +namespace EpisodicMDP + +variable [MeasurableSpace 𝓒] [MeasurableSpace 𝓐] [Nonempty 𝓐] (M : EpisodicMDP 𝓒 𝓐) + +/-- Before the horizon, the value of the `H`-step policy `Ο€` is its action value at the action +`Ο€ h s` it plays. -/ +theorem value_toKernel_fin {Ο€ : Policy 𝓒 𝓐 H} (hΟ€ : βˆ€ h, Measurable (Ο€ h)) (h : Fin H) + (s : 𝓒) : + M.value H (Ο€.toKernel hΟ€) h s = M.qValue H (Ο€.toKernel hΟ€) h s (Ο€ h s) := by + have := M.value_deterministic_of_lt (H := H) (Policy.measurable_extend hΟ€) h.2 s + rwa [Policy.extend_fin] at this + +end EpisodicMDP + +end Learning.MDP.Episodic