From 8cb6848e285bff0d2a5043f74f6fdee654e04f75 Mon Sep 17 00:00:00 2001
From: Remy Degenne <remydegenne@gmail.com>
Date: Sun, 21 Jun 2026 06:14:33 +0200
Subject: [PATCH] use trajectory

---
 .../Online/Bandit/SumRewards.lean             | 22 ++++++++++---------
 .../IonescuTulceaSpace.lean                   | 10 ++++-----
 2 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/LeanMachineLearning/Online/Bandit/SumRewards.lean b/LeanMachineLearning/Online/Bandit/SumRewards.lean
index 77b73365..85de702e 100644
--- a/LeanMachineLearning/Online/Bandit/SumRewards.lean
+++ b/LeanMachineLearning/Online/Bandit/SumRewards.lean
@@ -140,15 +140,15 @@ variable {𝓐 Ω Ω' : Type*} [DecidableEq 𝓐] {m𝓐 : MeasurableSpace 𝓐}
 
 lemma sumRewards_eq_comp :
     sumRewards A R a n =
-     (fun p ↦ ∑ i ∈ range n, if (p i).1 = a then (p i).2 else 0) ∘ (fun ω n ↦ (A n ω, R n ω)) := by
+     (fun p ↦ ∑ i ∈ range n, if (p i).1 = a then (p i).2 else 0) ∘ (trajectory A R) := by
   ext
-  simp [sumRewards]
+  simp [sumRewards, trajectory]
 
 lemma pullCount_eq_comp :
     pullCount A a n =
-      (fun p ↦ ∑ i ∈ range n, if (p i).1 = a then 1 else 0) ∘ (fun ω n ↦ (A n ω, R n ω)) := by
+      (fun p ↦ ∑ i ∈ range n, if (p i).1 = a then 1 else 0) ∘ (trajectory A R) := by
   ext
-  simp [pullCount]
+  simp [pullCount, trajectory]
 
 -- todo: write those lemmas with IdentDistrib instead of equality of maps
 lemma _root_.Learning.IsAlgEnvSeq.law_sumRewards_unique [MeasurableSingletonClass 𝓐]
@@ -199,12 +199,12 @@ lemma _root_.Learning.IsAlgEnvSeq.law_pullCount_sumRewards_unique' [MeasurableSi
       refine measurable_sum _ fun i hi ↦ Measurable.ite ?_ (by fun_prop) (by fun_prop)
       exact (measurableSet_singleton _).preimage (by fun_prop)
   have h_eq_comp : (fun ω a ↦ (pullCount A a n ω, sumRewards A R a n ω))
-      = f ∘ (fun ω n ↦ (A n ω, R n ω)) := by
+      = f ∘ (trajectory A R) := by
     ext ω a : 2
     rw [pullCount_eq_comp (R := R), sumRewards_eq_comp]
     grind
   have h_eq_comp2 : (fun ω a ↦ (pullCount A₂ a n ω, sumRewards A₂ R₂ a n ω))
-      = f ∘ (fun ω n ↦ (A₂ n ω, R₂ n ω)) := by
+      = f ∘ (trajectory A₂ R₂) := by
     ext ω a : 2
     rw [pullCount_eq_comp (R := R₂), sumRewards_eq_comp]
     grind
@@ -231,13 +231,15 @@ lemma _root_.Learning.IsAlgEnvSeq.identDistrib_pullCount_sumRewards [MeasurableS
     (∑ i ∈ range n, if (τ i).1 = a then 1 else 0,
      ∑ i ∈ range n, if (τ i).1 = a then (τ i).2 else 0)
   have hc1 : (fun ω n a ↦ (pullCount A a n ω, sumRewards A R a n ω)) =
-      f ∘ (fun ω n ↦ (A n ω, R n ω)) := by
+      f ∘ (trajectory A R) := by
     ext ω n a : 3
-    simp_rw [Function.comp, f, pullCount, card_filter, sumRewards]
+    simp_rw [Function.comp, f, pullCount, card_filter, sumRewards, trajectory]
+    rfl
   have hc2 : (fun ω' n a ↦ (pullCount A₂ a n ω', sumRewards A₂ R₂ a n ω')) =
-      f ∘ (fun ω' n ↦ (A₂ n ω', R₂ n ω')) := by
+      f ∘ (trajectory A₂ R₂) := by
     ext ω' n a : 3
-    simp_rw [Function.comp, f, pullCount, card_filter, sumRewards]
+    simp_rw [Function.comp, f, pullCount, card_filter, sumRewards, trajectory]
+    rfl
   have hf : Measurable f := by
     simp_rw [f, measurable_pi_iff]
     intro n a
diff --git a/LeanMachineLearning/SequentialLearning/IonescuTulceaSpace.lean b/LeanMachineLearning/SequentialLearning/IonescuTulceaSpace.lean
index 4112eb5f..733e951c 100644
--- a/LeanMachineLearning/SequentialLearning/IonescuTulceaSpace.lean
+++ b/LeanMachineLearning/SequentialLearning/IonescuTulceaSpace.lean
@@ -49,8 +49,8 @@ variable {Ω Ω' : Type*} {mΩ : MeasurableSpace Ω} {mΩ' : MeasurableSpace Ω'
   {P : Measure Ω} [IsProbabilityMeasure P] {P' : Measure Ω'} [IsProbabilityMeasure P']
   {A₁ : ℕ → Ω → 𝓐} {R₁ : ℕ → Ω → 𝓨} {A₂ : ℕ → Ω' → 𝓐} {R₂ : ℕ → Ω' → 𝓨} {N : ℕ}
 
-lemma eq_trajMeasure_of_isAlgEnvSeq (h : IsAlgEnvSeq A₁ R₁ alg env P) :
-    P.map (fun ω n ↦ (A₁ n ω, R₁ n ω)) = trajMeasure alg env := by
+lemma IsAlgEnvSeq.map_trajectory (h : IsAlgEnvSeq A₁ R₁ alg env P) :
+    P.map (trajectory A₁ R₁) = trajMeasure alg env := by
   rw [trajMeasure]
   have h := (Kernel.hasLaw_trajMeasure (Y := fun n ω ↦ (A₁ n ω, R₁ n ω)) (P := P)
     (μ₀ := alg.p0 ⊗ₘ env.ν0) (κ := stepKernel alg env) (fun n ↦ ?_) ?_ (fun n ↦ ?_)).map_eq
@@ -76,14 +76,14 @@ lemma eq_trajMeasure_map_frestrictLe_of_isAlgEnvSeqUntil
 is unique: it does not depend on the probability space used. -/
 lemma isAlgEnvSeq_unique (h1 : IsAlgEnvSeq A₁ R₁ alg env P)
     (h2 : IsAlgEnvSeq A₂ R₂ alg env P') :
-    P.map (fun ω n ↦ (A₁ n ω, R₁ n ω)) = P'.map (fun ω n ↦ (A₂ n ω, R₂ n ω)) := by
-  rw [eq_trajMeasure_of_isAlgEnvSeq h1, eq_trajMeasure_of_isAlgEnvSeq h2]
+    P.map (trajectory A₁ R₁) = P'.map (trajectory A₂ R₂) := by
+  rw [h1.map_trajectory, h2.map_trajectory]
 
 /-- The law of the sequence of actions and observations generated by an algorithm-environment pair
 is unique: it does not depend on the probability space used. -/
 theorem IsAlgEnvSeq.identDistrib_trajectory (h1 : IsAlgEnvSeq A₁ R₁ alg env P)
     (h2 : IsAlgEnvSeq A₂ R₂ alg env P') :
-    IdentDistrib (fun ω n ↦ (A₁ n ω, R₁ n ω)) (fun ω' n ↦ (A₂ n ω', R₂ n ω')) P P' where
+    IdentDistrib (trajectory A₁ R₁) (trajectory A₂ R₂) P P' where
   aemeasurable_fst := (measurable_pi_iff.2 fun n ↦ (h1.measurable_action n).prodMk
     (h1.measurable_feedback n)).aemeasurable
   aemeasurable_snd := (measurable_pi_iff.2 fun n ↦ (h2.measurable_action n).prodMk