proofs/lean/Agent/Gym.lean

/-
  Authors: Antje Worring, Zach Kelling

  Agent Training Environment (Gym) Safety Properties

  Formal proofs about the Gym training platform's environment safety.
  Gym is an open-source AI model training platform supporting 100+
  model architectures with training-free GRPO integration.

  Maps to:
  - zoo/papers/gym-training-platform.tex: Gym specification
  - zoo/papers/hllm-training-free-grpo.tex: HLLM/TF-GRPO foundation
  - github.com/zooai/gym: implementation

  Key results:
  - Environment state stays within defined bounds
  - Reward function is consistent across episodes
  - Exploration strategy covers all reachable states

  Author: Zach Kelling
-/

import Mathlib.Data.Nat.Defs
import Mathlib.Data.List.Basic
import Mathlib.Data.Finset.Basic
import Mathlib.Tactic

namespace Agent.Gym

/-! ## Core Structures -/

/-- Environment state is a bounded integer vector (abstracted as Nat with bounds). -/
structure EnvState where
  /-- State value (abstracted to single dimension for formal treatment) -/
  value : Nat
  /-- Lower bound on state -/
  lo : Nat
  /-- Upper bound on state -/
  hi : Nat
  /-- Invariant: state is within bounds -/
  hBounded : lo ≤ value ∧ value ≤ hi
  deriving DecidableEq

/-- An action modifies the state within bounds. -/
structure Action where
  /-- Delta applied to state (can be negative, represented as Int) -/
  delta : Int

/-- An episode is a sequence of (state, action, reward) tuples. -/
structure Step where
  state : EnvState
  action : Action
  reward : Nat

abbrev Episode := List Step

/-- A reward function maps (state, action) pairs to rewards. -/
abbrev RewardFn := EnvState → Action → Nat

/-- The set of reachable states from a given state. -/
abbrev ReachableSet := Finset Nat

/-! ## Theorem 1: Environment State Bounded -/

/-- Clamping an Int between lo and hi then converting to Nat preserves bounds,
    provided the window is well-formed (`lo ≤ hi`, which every `EnvState`
    guarantees via `hBounded`). The lower bound holds because `max lo _ ≥ lo ≥ 0`
    survives `Int.toNat`; the upper bound because `min raw hi ≤ hi` and `lo ≤ hi`
    give `max lo (min raw hi) ≤ hi`. -/
theorem clamp_bounded (lo hi : Nat) (raw : Int) (hle : lo ≤ hi) :
    lo ≤ (max (lo : Int) (min raw (hi : Int))).toNat ∧
    (max (lo : Int) (min raw (hi : Int))).toNat ≤ hi := by
  constructor
  · have h1 : (lo : Int) ≤ max (lo : Int) (min raw (hi : Int)) := le_max_left _ _
    have h0 : (0 : Int) ≤ (lo : Int) := Int.ofNat_nonneg lo
    rw [Int.le_toNat (le_trans h0 h1)]
    exact h1
  · have hub : max (lo : Int) (min raw (hi : Int)) ≤ (hi : Int) := by
      apply max_le
      · exact_mod_cast hle
      · exact min_le_right _ _
    have : (max (lo : Int) (min raw (hi : Int))).toNat ≤ ((hi : Int)).toNat :=
      Int.toNat_le_toNat hub
    simpa using this

/-- Apply an action to a state, clamping to bounds. -/
def applyAction (s : EnvState) (a : Action) : EnvState :=
  let raw := (s.value : Int) + a.delta
  let clamped := max (s.lo : Int) (min raw (s.hi : Int))
  { value := clamped.toNat
    lo := s.lo
    hi := s.hi
    hBounded := clamp_bounded s.lo s.hi ((s.value : Int) + a.delta)
      (le_trans s.hBounded.1 s.hBounded.2) }

/-- Environment state remains bounded after any action.
    The Gym platform enforces this via clamping: any action that would
    push the state out of bounds is truncated to the boundary.
    This prevents undefined behavior in the training loop. -/
theorem env_state_bounded (s : EnvState) (a : Action) :
    let s' := applyAction s a
    s'.lo ≤ s'.value ∧ s'.value ≤ s'.hi := by
  exact (applyAction s a).hBounded

/-- Bounds are preserved: the action doesn't change the bound definition. -/
theorem bounds_preserved (s : EnvState) (a : Action) :
    (applyAction s a).lo = s.lo ∧ (applyAction s a).hi = s.hi := by
  simp [applyAction]

/-- After a full episode, all intermediate states are bounded. -/
theorem episode_states_bounded (ep : Episode)
    (h : ∀ step ∈ ep, step.state.lo ≤ step.state.value ∧ step.state.value ≤ step.state.hi) :
    ∀ step ∈ ep, step.state.lo ≤ step.state.value := by
  intro step hmem
  exact (h step hmem).1

/-! ## Theorem 2: Reward Consistency -/

/-- A reward function is consistent if it is deterministic: the same
    (state, action) pair always produces the same reward.
    This is a fundamental requirement for GRPO's group-relative advantage
    computation to be meaningful across episodes. -/
def rewardConsistent (rf : RewardFn) : Prop :=
  ∀ s₁ s₂ : EnvState, ∀ a : Action,
    s₁.value = s₂.value → s₁.lo = s₂.lo → s₁.hi = s₂.hi →
    rf s₁ a = rf s₂ a

/-- If the reward function is deterministic in state (value, lo, hi) and action,
    then two trajectories that agree pointwise on those state components and on the
    chosen actions induce the same sequence of rewards *as computed by `rf`*.

    Note: we phrase the conclusion over `rf` applied at each step
    (`fun st => rf st.state st.action`), since that is the reward the deterministic
    function assigns. The hypotheses give full state agreement (value, lo, hi) so that
    `rewardConsistent` applies. Proof: `List.ext_getElem` reduces to pointwise equality,
    then `rewardConsistent` closes each index. -/
theorem reward_consistent (rf : RewardFn) (ep₁ ep₂ : Episode)
    (hConsistent : rewardConsistent rf)
    (hSameLen : ep₁.length = ep₂.length)
    (hSameStates : ∀ i : Fin (min ep₁.length ep₂.length),
      (ep₁.get ⟨i.val, by omega⟩).state.value = (ep₂.get ⟨i.val, by omega⟩).state.value ∧
      (ep₁.get ⟨i.val, by omega⟩).state.lo = (ep₂.get ⟨i.val, by omega⟩).state.lo ∧
      (ep₁.get ⟨i.val, by omega⟩).state.hi = (ep₂.get ⟨i.val, by omega⟩).state.hi)
    (hSameActions : ∀ i : Fin (min ep₁.length ep₂.length),
      (ep₁.get ⟨i.val, by omega⟩).action = (ep₂.get ⟨i.val, by omega⟩).action) :
    ep₁.map (fun st => rf st.state st.action) = ep₂.map (fun st => rf st.state st.action) := by
  apply List.ext_getElem
  · simp [hSameLen]
  · intro i h1 h2
    simp only [List.getElem_map]
    have hi_lt : i < min ep₁.length ep₂.length := by
      simp only [List.length_map] at h1 h2; omega
    have hs := hSameStates ⟨i, hi_lt⟩
    have ha := hSameActions ⟨i, hi_lt⟩
    have e1 : ep₁[i] = ep₁.get ⟨i, by simp only [List.length_map] at h1; omega⟩ := by rfl
    have e2 : ep₂[i] = ep₂.get ⟨i, by simp only [List.length_map] at h2; omega⟩ := by rfl
    rw [e1, e2, ha]
    exact hConsistent _ _ _ hs.1 hs.2.1 hs.2.2

/-- A bounded reward function produces rewards in [0, rmax]. -/
def rewardBounded (rf : RewardFn) (rmax : Nat) : Prop :=
  ∀ s : EnvState, ∀ a : Action, rf s a ≤ rmax

/-! ## Theorem 3: Exploration Coverage -/

/-- An exploration strategy is a function from state to action distribution.
    Here we model it as producing a list of actions to try. -/
abbrev ExplorationStrategy := EnvState → List Action

/-- The set of states reachable from a starting state via a strategy.
    Computed by applying all actions and collecting resulting state values. -/
def reachableFrom (s : EnvState) (strategy : ExplorationStrategy) : List Nat :=
  (strategy s).map (fun a => (applyAction s a).value)

/-- Coverage property: given enough episodes, the exploration strategy
    visits every reachable state at least once.
    This is critical for GRPO's group-relative advantage computation:
    if some states are never visited, the advantage estimates are biased.

    We state this as: for any reachable state value v, there exists an
    action in the strategy that reaches v from some starting state. -/
theorem exploration_covers (s : EnvState) (strategy : ExplorationStrategy)
    (v : Nat) (hReachable : v ∈ reachableFrom s strategy) :
    ∃ a ∈ strategy s, (applyAction s a).value = v := by
  simp [reachableFrom] at hReachable
  obtain ⟨a, ha_mem, ha_val⟩ := hReachable
  exact ⟨a, ha_mem, ha_val⟩

/-- If the strategy produces at least one action for every state,
    the reachable set is nonempty. -/
theorem exploration_nonempty (s : EnvState) (strategy : ExplorationStrategy)
    (h : (strategy s).length > 0) :
    (reachableFrom s strategy).length > 0 := by
  simp only [reachableFrom, List.length_map]
  exact h

end Agent.Gym