AdaptiveExploration.lean

import Std

/- Finite adaptive experiment design, Lean 4.24.0.
   The proof is fully
   constructive over finite choices. No placeholders/sorry/axioms are used.
   Bayesian likelihoods are instantiated in state transitions, with weighted
   charges and rewards; this generic optimizer does not establish calibration. -/
namespace AdaptiveExploration

/-- Observable finite-state decision problem. Actions/terminal choices depend
    on state, so remaining budget, rights and deadlines can affect both. -/
structure Model (S A O D : Type) where
  outcomes : List O
  actions : S → List A
  terminals : S → List D
  outside : D
  outside_mem : ∀ s, outside ∈ terminals s
  payoff : S → D → Rat
  outside_zero : ∀ s, payoff s outside = 0
  next : S → A → O → S
  charge : S → A → Rat
  budget : S → Nat
  debit : S → A → Nat
  affordable : ∀ s a, a ∈ actions s → debit s a ≤ budget s
  child_budget : ∀ s a o, a ∈ actions s → o ∈ outcomes →
    budget (next s a o) = budget s - debit s a

/-- Policies are defined independently of the optimizer. Terminal choice is
    part of the policy, rather than an oracle maximization hidden in eval. -/
inductive Policy (A O D : Type) : Nat → Type where
  | stop {n : Nat} (d : D) : Policy A O D n
  | step {n : Nat} (a : A) (next : O → Policy A O D n) : Policy A O D (n+1)

/-- List rather than type enumeration: instantiate outcomes without duplicates
    and normalized likelihoods for the probability interpretation. -/
def sumOver {α : Type} : List α → (α → Rat) → Rat
  | [], _ => 0
  | x :: xs, f => f x + sumOver xs f

def evaluate (m : Model S A O D) :
    (n : Nat) → S → Policy A O D n → Rat
  | 0, s, .stop d => m.payoff s d
  | _n+1, s, .stop d => m.payoff s d
  | n+1, s, .step a k =>
      -m.charge s a + sumOver m.outcomes (fun o => evaluate m n (m.next s a o) (k o))

def Feasible (m : Model S A O D) :
    (n : Nat) → S → Policy A O D n → Prop
  | 0, s, .stop d => d ∈ m.terminals s
  | _n+1, s, .stop d => d ∈ m.terminals s
  | n+1, s, .step a k =>
      a ∈ m.actions s ∧ ∀ o, o ∈ m.outcomes → Feasible m n (m.next s a o) (k o)

/-- A nonempty finite argmax, with deterministic tie preference. -/
def choose {α : Type} (score : α → Rat) (base : α) : List α → α
  | [] => base
  | x :: xs =>
      let tail := choose score base xs
      if score tail < score x then x else tail

theorem choose_base_le {α : Type} (score : α → Rat) (base : α) (xs : List α) :
    score base ≤ score (choose score base xs) := by
  induction xs with
  | nil => exact Rat.le_refl
  | cons x xs ih =>
    simp only [choose]
    split <;> grind

theorem choose_member_le {α : Type} (score : α → Rat) (base : α)
    (xs : List α) (x : α) (hx : x ∈ xs) :
    score x ≤ score (choose score base xs) := by
  induction xs with
  | nil => simp at hx
  | cons y ys ih =>
    simp only [List.mem_cons] at hx
    simp only [choose]
    rcases hx with hxy | hxy
    · subst x
      split <;> grind
    · have h := ih hxy
      split <;> grind

theorem choose_preserves {α : Type} (score : α → Rat) (base : α)
    (xs : List α) (P : α → Prop) (hb : P base)
    (hx : ∀ x, x ∈ xs → P x) : P (choose score base xs) := by
  induction xs with
  | nil => exact hb
  | cons x xs ih =>
    have hhead : P x := hx x (by simp)
    have htail : P (choose score base xs) :=
      ih (fun y hy => hx y (by simp [hy]))
    simp only [choose]
    split <;> assumption

theorem sumOver_mono {α : Type} (xs : List α) (f g : α → Rat)
    (h : ∀ x, x ∈ xs → f x ≤ g x) : sumOver xs f ≤ sumOver xs g := by
  induction xs with
  | nil => exact Rat.le_refl
  | cons x xs ih =>
    have hh := h x (by simp)
    have ht := ih (fun y hy => h y (by simp [hy]))
    simp only [sumOver]
    grind

/-- Backward induction produces an actual policy, not only a value. -/
def optimize (m : Model S A O D) : (n : Nat) → S → Policy A O D n
  | 0, s =>
      .stop (choose (m.payoff s) m.outside (m.terminals s))
  | n+1, s =>
      choose (evaluate m (n+1) s) (.stop m.outside)
        ((m.terminals s).map (fun d => Policy.stop d) ++
         (m.actions s).map (fun a => Policy.step a (fun o => optimize m n (m.next s a o))))

def value (m : Model S A O D) (n : Nat) (s : S) : Rat :=
  evaluate m n s (optimize m n s)

theorem optimize_feasible (m : Model S A O D) (n : Nat) (s : S) :
    Feasible m n s (optimize m n s) := by
  induction n generalizing s with
  | zero =>
    change (choose (m.payoff s) m.outside (m.terminals s)) ∈ m.terminals s
    exact choose_preserves _ _ _ (fun d => d ∈ m.terminals s)
      (m.outside_mem s) (fun _ h => h)
  | succ n ih =>
    unfold optimize
    apply choose_preserves
    · exact m.outside_mem s
    · intro p hp
      simp only [List.mem_append, List.mem_map] at hp
      rcases hp with hp | hp
      · rcases hp with ⟨d, hd, rfl⟩
        exact hd
      · rcases hp with ⟨a, ha, rfl⟩
        exact ⟨ha, fun o _ => ih (m.next s a o)⟩

/-- GLOBAL finite-policy optimality, not only a declared comparison among a
    few handpicked policies. It covers every recursively feasible policy. -/
theorem bellman_dominates_policy (m : Model S A O D) (n : Nat) (s : S)
    (p : Policy A O D n) (hp : Feasible m n s p) :
    evaluate m n s p ≤ value m n s := by
  induction n generalizing s with
  | zero =>
    cases p with
    | stop d =>
      exact choose_member_le (m.payoff s) m.outside (m.terminals s) d hp
  | succ n ih =>
    cases p with
    | stop d =>
      have hd : d ∈ m.terminals s := hp
      apply choose_member_le (evaluate m (n+1) s) (.stop m.outside)
        ((m.terminals s).map (fun d => Policy.stop d) ++
         (m.actions s).map (fun a => Policy.step a (fun o => optimize m n (m.next s a o))))
        (.stop d)
      simp only [List.mem_append, List.mem_map]
      exact Or.inl ⟨d, hd, rfl⟩
    | step a k =>
      obtain ⟨ha, hk⟩ := hp
      let q : Policy A O D (n+1) :=
        .step a (fun o => optimize m n (m.next s a o))
      have hsum := sumOver_mono m.outcomes
        (fun o => evaluate m n (m.next s a o) (k o))
        (fun o => evaluate m n (m.next s a o) (optimize m n (m.next s a o)))
        (fun o ho => ih (m.next s a o) (k o) (hk o ho))
      have hpq : evaluate m (n+1) s (.step a k) ≤ evaluate m (n+1) s q := by
        simp only [evaluate, q]
        grind
      have hq : evaluate m (n+1) s q ≤ value m (n+1) s := by
        apply choose_member_le (evaluate m (n+1) s) (.stop m.outside)
          ((m.terminals s).map (fun d => Policy.stop d) ++
           (m.actions s).map (fun a => Policy.step a (fun o => optimize m n (m.next s a o)))) q
        simp only [List.mem_append, List.mem_map]
        exact Or.inr ⟨a, ha, rfl⟩
      exact Rat.le_trans hpq hq

/-- The upper bound is attainable by a feasible policy. -/
theorem bellman_attained (m : Model S A O D) (n : Nat) (s : S) :
    ∃ p : Policy A O D n, Feasible m n s p ∧ evaluate m n s p = value m n s := by
  exact ⟨optimize m n s, optimize_feasible m n s, rfl⟩

theorem stop_lower_bound (m : Model S A O D) (n : Nat) (s : S)
    (d : D) (hd : d ∈ m.terminals s) : m.payoff s d ≤ value m n s := by
  have hp : Feasible m n s (.stop d) := by cases n <;> exact hd
  have hv := bellman_dominates_policy m n s (.stop d) hp
  cases n <;> exact hv

theorem optional_stop_nonnegative (m : Model S A O D) (n : Nat) (s : S) :
    0 ≤ value m n s := by
  have h := stop_lower_bound m n s m.outside (m.outside_mem s)
  rw [m.outside_zero] at h
  exact h

/-- One actual child transition preserves the budget ledger. Induct this
    identity along a path for cumulative debit+remaining=initial budget. -/
theorem budget_step_conservation (m : Model S A O D) (s : S) (a : A) (o : O)
    (ha : a ∈ m.actions s) (ho : o ∈ m.outcomes) :
    m.debit s a + m.budget (m.next s a o) = m.budget s := by
  rw [m.child_budget s a o ha ho]
  have h := m.affordable s a ha
  omega

/-- A menu expansion is valuable only if all old policies really remain
    feasible with equal utilities. This is an abstract embedding theorem;
    concrete inclusion of action menus establishes its hypothesis. -/
theorem feasible_emulation_monotone (m₁ m₂ : Model S A O D)
    (n : Nat) (s : S)
    (hemulate : ∀ p : Policy A O D n, Feasible m₁ n s p →
      Feasible m₂ n s p ∧ evaluate m₁ n s p = evaluate m₂ n s p) :
    value m₁ n s ≤ value m₂ n s := by
  have h := hemulate (optimize m₁ n s) (optimize_feasible m₁ n s)
  have hb := bellman_dominates_policy m₂ n s (optimize m₁ n s) h.1
  change evaluate m₁ n s (optimize m₁ n s) ≤ value m₂ n s
  rw [h.2]
  exact hb


inductive Experiment where
  | first | second
  deriving DecidableEq, Repr
inductive Observation where
  | plus | minus | missing
  deriving DecidableEq, Repr
inductive Deployment where
  | none | first | second
  deriving DecidableEq, Repr

/-- Joint weights: first coordinate project A, second project B. -/
structure Weights where
  hh : Rat
  hl : Rat
  lh : Rat
  ll : Rat
  deriving DecidableEq, Repr

structure State where
  weights : Weights
  aAvailable : Bool
  bAvailable : Bool
  remaining : Nat
  deriving DecidableEq, Repr

def mass (w : Weights) : Rat := w.hh+w.hl+w.lh+w.ll

def nextWeights (w : Weights) (a : Experiment) (o : Observation) : Weights :=
  match a, o with
  | _, .missing => ⟨0,0,0,0⟩
  | .first, .plus => ⟨w.hh*(3/4),w.hl*(3/4),w.lh*(1/4),w.ll*(1/4)⟩
  | .first, .minus => ⟨w.hh*(1/4),w.hl*(1/4),w.lh*(3/4),w.ll*(3/4)⟩
  | .second, .plus => ⟨w.hh*(3/4),w.hl*(1/4),w.lh*(3/4),w.ll*(1/4)⟩
  | .second, .minus => ⟨w.hh*(1/4),w.hl*(3/4),w.lh*(1/4),w.ll*(3/4)⟩

def next (s : State) (a : Experiment) (o : Observation) : State :=
  { weights := nextWeights s.weights a o
    aAvailable := if a = .first then false else s.aAvailable
    bAvailable := if a = .second then false else s.bAvailable
    remaining := s.remaining-1 }

def actions (s : State) : List Experiment :=
  if 0 < s.remaining then
    (if s.aAvailable then [.first] else []) ++ (if s.bAvailable then [.second] else [])
  else []

def payoff (s : State) : Deployment → Rat
  | .none => 0
  | .first => 4*(s.weights.hh+s.weights.hl-s.weights.lh-s.weights.ll)
  | .second => 4*(s.weights.hh+s.weights.lh-s.weights.hl-s.weights.ll)

def clean : Model State Experiment Observation Deployment where
  outcomes := [.plus,.minus,.missing]
  actions := actions
  terminals := fun _ => [.none,.first,.second]
  outside := .none
  outside_mem := by intro s; simp
  payoff := payoff
  outside_zero := by intro s; rfl
  next := next
  charge := fun s _ => mass s.weights / 2
  budget := State.remaining
  debit := fun _ _ => 1
  affordable := by
    intro s a ha
    by_cases h : 0 < s.remaining
    · omega
    · simp [actions, h] at ha
  child_budget := by intro s a o ha ho; rfl

def initial : State := ⟨⟨1/4,1/4,1/4,1/4⟩,true,true,2⟩

/-- Uses only observed feedback, including a harmless feasible continuation
    at the zero-mass missing branch. -/
def adaptive : Policy Experiment Observation Deployment 2 :=
  .step .first (fun o =>
    match o with
    | .plus => .stop .first
    | .minus | .missing => .step .second (fun p =>
      match p with
      | .plus => .stop .second
      | .minus | .missing => .stop .none))

/-- Terminal deployment is posterior-optimal even for fixed experiment sets. -/
def terminal (s : State) : Policy Experiment Observation Deployment 0 :=
  .stop (choose (payoff s) .none [.none,.first,.second])

def onlyA : Policy Experiment Observation Deployment 1 :=
  .step .first (fun o => terminal (next initial .first o))
def onlyB : Policy Experiment Observation Deployment 1 :=
  .step .second (fun o => terminal (next initial .second o))
def both : Policy Experiment Observation Deployment 2 :=
  .step .first (fun o => .step .second (fun p => terminal (next (next initial .first o) .second p)))

def bothReverse : Policy Experiment Observation Deployment 2 :=
  .step .second (fun o => .step .first (fun p => terminal (next (next initial .second o) .first p)))

inductive FixedSet where
  | none | onlyA | onlyB | both | bothReverse
  deriving DecidableEq, Repr

def fixedValue : FixedSet → Rat
  | .none => evaluate clean 0 initial (.stop .none)
  | .onlyA => evaluate clean 1 initial onlyA
  | .onlyB => evaluate clean 1 initial onlyB
  | .both => evaluate clean 2 initial both
  | .bothReverse => evaluate clean 2 initial bothReverse

theorem clean_optimum : value clean 2 initial = 3/4 := by decide +kernel
theorem adaptive_value : evaluate clean 2 initial adaptive = 3/4 := by decide +kernel
theorem onlyA_value : fixedValue .onlyA = 1/2 := by decide +kernel
theorem onlyB_value : fixedValue .onlyB = 1/2 := by decide +kernel
theorem both_value : fixedValue .both = 1/2 := by decide +kernel
theorem both_reverse_value : fixedValue .bothReverse = 1/2 := by decide +kernel

theorem every_fixed_set_bounded (f : FixedSet) : fixedValue f ≤ 1/2 := by
  cases f <;> decide +kernel

theorem adaptive_strictly_beats_every_fixed_set (f : FixedSet) :
    fixedValue f < value clean 2 initial := by
  rw [clean_optimum]
  have h := every_fixed_set_bounded f
  grind

theorem child_mass_conservation (w : Weights) (a : Experiment) :
    mass (nextWeights w a .plus) + mass (nextWeights w a .minus) +
    mass (nextWeights w a .missing) = mass w := by
  cases a <;> unfold nextWeights mass <;> grind

/-- Local is optimized adaptively under its explicitly restricted rights. -/
def restricted : Model State Experiment Observation Deployment := clean

def localInitial : State := { initial with bAvailable := false }

theorem restricted_local_optimum : value restricted 2 localInitial = 1/2 := by
  decide +kernel

/-- Mandatory operating-mode activation fee. Stop retains zero value. These
    scalar identities are interpretations of the clean optimum, not a new
    claimed model optimization with a hidden local bypass. -/
def mandatoryMode (fee : Rat) : Rat :=
  let active := value clean 2 initial-fee
  if 0 ≤ active then active else 0

def optionalService (fee : Rat) : Rat :=
  let active := value clean 2 initial-fee
  let localValue := value restricted 2 localInitial
  if localValue ≤ active then active else localValue

theorem low_fee_shared_wins : mandatoryMode (1/8) - value restricted 2 localInitial = 1/8 := by
  decide +kernel

theorem high_fee_mandatory_loses : mandatoryMode (1/2) - value restricted 2 localInitial = -1/4 := by
  decide +kernel

theorem high_fee_optional_bypasses : optionalService (1/2) = value restricted 2 localInitial := by
  decide +kernel

/-- Censored-feedback fixture: priorH=3/4, project value +4 versus −8. -/
inductive BinaryDeployment where
  | none | deploy
  deriving DecidableEq, Repr

structure BinaryState where
  high : Rat
  low : Rat
  remaining : Nat
  deriving DecidableEq, Repr

def trueKernel (high : Bool) : Observation → Rat
  | .plus => if high then 3/4 else 1/4
  | .minus => 0
  | .missing => if high then 1/4 else 3/4

def assumedKernel (high : Bool) : Observation → Rat
  | .plus => if high then 3/8 else 1/8
  | .minus => if high then 1/8 else 3/8
  | .missing => 1/2

def binaryActions (s : BinaryState) : List Unit :=
  if 0 < s.remaining then [()] else []

def binaryPayoff (s : BinaryState) : BinaryDeployment → Rat
  | .none => 0
  | .deploy => 4*s.high-8*s.low

def binaryModel (kernel : Bool → Observation → Rat) :
    Model BinaryState Unit Observation BinaryDeployment where
  outcomes := [.plus,.minus,.missing]
  actions := binaryActions
  terminals := fun _ => [.none,.deploy]
  outside := .none
  outside_mem := by intro s; simp
  payoff := binaryPayoff
  outside_zero := by intro s; rfl
  next := fun s _ o => ⟨s.high*kernel true o,s.low*kernel false o,s.remaining-1⟩
  charge := fun s _ => (s.high+s.low)/4
  budget := BinaryState.remaining
  debit := fun _ _ => 1
  affordable := by
    intro s a ha
    by_cases h : 0 < s.remaining
    · omega
    · simp [binaryActions, h] at ha
  child_budget := by intro s a o ha ho; rfl

def binaryInitial : BinaryState := ⟨3/4,1/4,1⟩

def calibratedPolicy : Policy Unit Observation BinaryDeployment 1 :=
  .step () (fun o => match o with
    | .plus => .stop .deploy
    | .minus | .missing => .stop .none)

def misspecifiedPolicy : Policy Unit Observation BinaryDeployment 1 :=
  .step () (fun o => match o with
    | .plus | .missing => .stop .deploy
    | .minus => .stop .none)

theorem calibrated_optimum : value (binaryModel trueKernel) 1 binaryInitial = 3/2 := by
  decide +kernel

theorem calibrated_policy_value :
    evaluate (binaryModel trueKernel) 1 binaryInitial calibratedPolicy = 3/2 := by
  decide +kernel

theorem assumed_optimum : value (binaryModel assumedKernel) 1 binaryInitial = 9/8 := by
  decide +kernel

theorem misspecified_policy_planned_value :
    evaluate (binaryModel assumedKernel) 1 binaryInitial misspecifiedPolicy = 9/8 := by
  decide +kernel

theorem misspecified_policy_true_value :
    evaluate (binaryModel trueKernel) 1 binaryInitial misspecifiedPolicy = 3/4 := by
  decide +kernel

theorem baseline_deploy_value : binaryPayoff binaryInitial .deploy = 1 := by
  decide +kernel

theorem misspecified_learning_harms :
    evaluate (binaryModel trueKernel) 1 binaryInitial misspecifiedPolicy -
    binaryPayoff binaryInitial .deploy = -1/4 := by
  decide +kernel

theorem true_kernel_normalized (h : Bool) :
    trueKernel h .plus + trueKernel h .minus + trueKernel h .missing = 1 := by
  cases h <;> decide +kernel

theorem assumed_kernel_normalized (h : Bool) :
    assumedKernel h .plus + assumedKernel h .minus + assumedKernel h .missing = 1 := by
  cases h <;> decide +kernel

end AdaptiveExploration