gaojie
c66855adfc
Sync to site1 / sync (push) Has been cancelled
- 移除 JEPA/lejepa-identifiability 子模块 gitlink - 移除 research/multiply/MultiPLY 子模块 gitlink - 删除 .gitmodules（不再有外部 URL 依赖） - 两个目录内容作为普通文件纳入主仓库追踪 - 删除各自内部 .git 目录，消除嵌套 git 仓库
246 lines
12 KiB
Lean4
246 lines
12 KiB
Lean4
import Mathlib
|
||
|
||
/-!
|
||
# Part D — Planning Equivalence (Corollary)
|
||
|
||
Let h(z) = Qz with Q ∈ O(n) be the encoder at the optimum of Theorem 4.1.
|
||
For any finite-horizon optimal control problem whose stage and terminal
|
||
costs are O(n)-invariant in the state argument, the optimal value function
|
||
and the set of optimal action sequences agree between the learned latent
|
||
and the true latent.
|
||
|
||
The proof reduces — via the rotation-invariance hypothesis and the
|
||
pushforward property of expected costs — to the trivial fact that pointwise
|
||
equal real-valued functions share minimizers.
|
||
|
||
## Verification status
|
||
|
||
| Component | Status |
|
||
|---------------------------------------|-------------|
|
||
| ControlProblem structure | structural |
|
||
| Orthogonal invariance definition | structural |
|
||
| ExpectedCosts abstraction | structural |
|
||
| Total-cost definition | structural |
|
||
| Trajectory pushforward (stage) | axiomatized |
|
||
| Trajectory pushforward (terminal) | axiomatized |
|
||
| Per-step stage cost equivalence | VERIFIED |
|
||
| Terminal cost equivalence | VERIFIED |
|
||
| Total cost equivalence (main step) | VERIFIED |
|
||
| Minimizer equivalence (plan agreement)| VERIFIED |
|
||
| Value equivalence | VERIFIED |
|
||
-/
|
||
|
||
set_option maxHeartbeats 400000
|
||
|
||
open scoped BigOperators
|
||
|
||
noncomputable section
|
||
|
||
abbrev Latent (n : ℕ) := Fin n → ℝ
|
||
abbrev Plan (Action : Type*) (T : ℕ) := Fin T → Action
|
||
|
||
|
||
-- ═══════════════════════════════════════════════════════════════
|
||
-- STRUCTURE: CONTROL PROBLEM AND ROTATION INVARIANCE
|
||
-- ═══════════════════════════════════════════════════════════════
|
||
|
||
/-- A finite-horizon optimal control problem with stage cost ℓ(z,a) and
|
||
terminal cost ℓ_T(z). -/
|
||
structure ControlProblem (n : ℕ) (Action : Type*) where
|
||
stage_cost : Latent n → Action → ℝ
|
||
terminal_cost : Latent n → ℝ
|
||
|
||
/-- The costs of the control problem are O(n)-invariant in the state argument
|
||
under a map Q: ℓ(Q z, a) = ℓ(z, a) for all z, a, and ℓ_T(Q z) = ℓ_T(z)
|
||
for all z. In the corollary, Q is the orthogonal recovery matrix from
|
||
Theorem 4.1; the definition does not itself require Q to be linear or
|
||
orthogonal — only the invariance property is used. -/
|
||
def IsOrthogonalInvariant {n : ℕ} {Action : Type*}
|
||
(cp : ControlProblem n Action) (Q : Latent n → Latent n) : Prop :=
|
||
(∀ z a, cp.stage_cost (Q z) a = cp.stage_cost z a) ∧
|
||
(∀ z, cp.terminal_cost (Q z) = cp.terminal_cost z)
|
||
|
||
|
||
-- ═══════════════════════════════════════════════════════════════
|
||
-- STRUCTURE: EXPECTED COSTS UNDER SOME DYNAMICS
|
||
-- ═══════════════════════════════════════════════════════════════
|
||
|
||
/-- Expected costs along a trajectory under a specific (stochastic) dynamics.
|
||
|
||
`stage_exp a z₀ t c` is the expected value of `c(z_t, a_t)` at time `t`
|
||
along the trajectory starting from `z₀` and following the action sequence
|
||
`a`. `term_exp a z₀ c` is the expected value of `c(z_T)` at the final
|
||
time. Parameterizing over the cost function `c` lets the same dynamics
|
||
object be reused for different costs, and makes the pushforward relation
|
||
(below) statable without explicit measure theory. -/
|
||
structure ExpectedCosts (n : ℕ) (Action : Type*) (T : ℕ) where
|
||
stage_exp :
|
||
Plan Action T → Latent n → Fin T → (Latent n → Action → ℝ) → ℝ
|
||
term_exp :
|
||
Plan Action T → Latent n → (Latent n → ℝ) → ℝ
|
||
|
||
|
||
-- ═══════════════════════════════════════════════════════════════
|
||
-- TOTAL EXPECTED COST
|
||
-- ═══════════════════════════════════════════════════════════════
|
||
|
||
/-- Total expected cost for a plan `a` from initial state `z₀`: the sum of
|
||
per-step stage costs plus the terminal cost. -/
|
||
def totalCost {n : ℕ} {Action : Type*} {T : ℕ}
|
||
(cp : ControlProblem n Action) (E : ExpectedCosts n Action T)
|
||
(a : Plan Action T) (z₀ : Latent n) : ℝ :=
|
||
(∑ t : Fin T, E.stage_exp a z₀ t cp.stage_cost)
|
||
+ E.term_exp a z₀ cp.terminal_cost
|
||
|
||
|
||
-- ═══════════════════════════════════════════════════════════════
|
||
-- AXIOMATIZED: TRAJECTORY PUSHFORWARD
|
||
-- ═══════════════════════════════════════════════════════════════
|
||
|
||
/-- **Stage pushforward** (axiomatized): under the pushforward dynamics
|
||
`E_hat`, the expected value of any cost `c` at time `t` starting from
|
||
`Q z` equals the expected value under the original dynamics `E` starting
|
||
from `z` of the pre-composed cost `c ∘ (Q × id)`.
|
||
|
||
Mathematically this is the content of "the joint law of (ẑ_0, …, ẑ_T)
|
||
under the pushforward dynamics starting from ẑ_0 = Q z equals the joint
|
||
law of (Q z_0, …, Q z_T) under the original dynamics starting from
|
||
z_0 = z", restricted to per-time-step marginals and evaluated against
|
||
arbitrary test functions. -/
|
||
axiom stage_pushforward
|
||
{n : ℕ} {Action : Type*} {T : ℕ}
|
||
(E_hat E : ExpectedCosts n Action T) (Q : Latent n → Latent n)
|
||
(a : Plan Action T) (z : Latent n) (t : Fin T)
|
||
(c : Latent n → Action → ℝ) :
|
||
E_hat.stage_exp a (Q z) t c
|
||
= E.stage_exp a z t (fun z' act => c (Q z') act)
|
||
|
||
/-- **Terminal pushforward** (axiomatized): the same relation at the
|
||
terminal time. -/
|
||
axiom terminal_pushforward
|
||
{n : ℕ} {Action : Type*} {T : ℕ}
|
||
(E_hat E : ExpectedCosts n Action T) (Q : Latent n → Latent n)
|
||
(a : Plan Action T) (z : Latent n) (c : Latent n → ℝ) :
|
||
E_hat.term_exp a (Q z) c = E.term_exp a z (fun z' => c (Q z'))
|
||
|
||
|
||
-- ═══════════════════════════════════════════════════════════════
|
||
-- VERIFIED: PER-STEP COST EQUIVALENCE
|
||
-- ═══════════════════════════════════════════════════════════════
|
||
|
||
/-- **Stage-cost equivalence** (VERIFIED): the per-step expected stage cost
|
||
at `Q z` under the pushforward dynamics equals the per-step expected
|
||
stage cost at `z` under the original dynamics, when the stage cost is
|
||
O(n)-invariant. This is the point where orthogonal invariance of the
|
||
cost (hypothesis) meets trajectory pushforward (axiom). -/
|
||
theorem stage_cost_equiv
|
||
{n : ℕ} {Action : Type*} {T : ℕ}
|
||
(cp : ControlProblem n Action) (Q : Latent n → Latent n)
|
||
(E_hat E : ExpectedCosts n Action T)
|
||
(hinv : IsOrthogonalInvariant cp Q)
|
||
(a : Plan Action T) (z : Latent n) (t : Fin T) :
|
||
E_hat.stage_exp a (Q z) t cp.stage_cost
|
||
= E.stage_exp a z t cp.stage_cost := by
|
||
rw [stage_pushforward E_hat E Q a z t cp.stage_cost]
|
||
have hfun : (fun z' act => cp.stage_cost (Q z') act) = cp.stage_cost := by
|
||
funext z'
|
||
funext act
|
||
exact hinv.1 z' act
|
||
rw [hfun]
|
||
|
||
/-- **Terminal-cost equivalence** (VERIFIED). -/
|
||
theorem terminal_cost_equiv
|
||
{n : ℕ} {Action : Type*} {T : ℕ}
|
||
(cp : ControlProblem n Action) (Q : Latent n → Latent n)
|
||
(E_hat E : ExpectedCosts n Action T)
|
||
(hinv : IsOrthogonalInvariant cp Q)
|
||
(a : Plan Action T) (z : Latent n) :
|
||
E_hat.term_exp a (Q z) cp.terminal_cost
|
||
= E.term_exp a z cp.terminal_cost := by
|
||
rw [terminal_pushforward E_hat E Q a z cp.terminal_cost]
|
||
have hfun : (fun z' => cp.terminal_cost (Q z')) = cp.terminal_cost := by
|
||
funext z'
|
||
exact hinv.2 z'
|
||
rw [hfun]
|
||
|
||
|
||
-- ═══════════════════════════════════════════════════════════════
|
||
-- VERIFIED: TOTAL COST EQUIVALENCE (PLANNING EQUIVALENCE)
|
||
-- ═══════════════════════════════════════════════════════════════
|
||
|
||
/-- **Planning equivalence** (VERIFIED, main step): for any action sequence,
|
||
the total expected cost under the pushforward dynamics at `Q z₀` equals
|
||
the total expected cost under the original dynamics at `z₀`.
|
||
|
||
This is the central computational content of the corollary; everything
|
||
that follows (value and minimizer equivalence) is a consequence. -/
|
||
theorem planning_equivalence
|
||
{n : ℕ} {Action : Type*} {T : ℕ}
|
||
(cp : ControlProblem n Action) (Q : Latent n → Latent n)
|
||
(E_hat E : ExpectedCosts n Action T)
|
||
(hinv : IsOrthogonalInvariant cp Q)
|
||
(a : Plan Action T) (z : Latent n) :
|
||
totalCost cp E_hat a (Q z) = totalCost cp E a z := by
|
||
unfold totalCost
|
||
have hstage :
|
||
(∑ t : Fin T, E_hat.stage_exp a (Q z) t cp.stage_cost)
|
||
= ∑ t : Fin T, E.stage_exp a z t cp.stage_cost := by
|
||
apply Finset.sum_congr rfl
|
||
intro t _
|
||
exact stage_cost_equiv cp Q E_hat E hinv a z t
|
||
have hterm :
|
||
E_hat.term_exp a (Q z) cp.terminal_cost
|
||
= E.term_exp a z cp.terminal_cost :=
|
||
terminal_cost_equiv cp Q E_hat E hinv a z
|
||
rw [hstage, hterm]
|
||
|
||
|
||
-- ═══════════════════════════════════════════════════════════════
|
||
-- VERIFIED: MINIMIZER AND VALUE EQUIVALENCE
|
||
-- ═══════════════════════════════════════════════════════════════
|
||
|
||
/-- **Minimizer equivalence** (VERIFIED): an action sequence minimizes the
|
||
expected cost under the pushforward dynamics at `Q z` iff it minimizes
|
||
the expected cost under the original dynamics at `z`.
|
||
|
||
Consequence: the optimal plan is the same whether it is computed in the
|
||
learned latent or the true latent. -/
|
||
theorem minimizer_equivalence
|
||
{n : ℕ} {Action : Type*} {T : ℕ}
|
||
(cp : ControlProblem n Action) (Q : Latent n → Latent n)
|
||
(E_hat E : ExpectedCosts n Action T)
|
||
(hinv : IsOrthogonalInvariant cp Q)
|
||
(a : Plan Action T) (z : Latent n) :
|
||
(∀ a', totalCost cp E_hat a (Q z) ≤ totalCost cp E_hat a' (Q z)) ↔
|
||
(∀ a', totalCost cp E a z ≤ totalCost cp E a' z) := by
|
||
have h : ∀ a', totalCost cp E_hat a' (Q z) = totalCost cp E a' z :=
|
||
fun a' => planning_equivalence cp Q E_hat E hinv a' z
|
||
constructor
|
||
· intro hmin a'
|
||
have ha := h a
|
||
have ha' := h a'
|
||
have := hmin a'
|
||
linarith
|
||
· intro hmin a'
|
||
have ha := h a
|
||
have ha' := h a'
|
||
have := hmin a'
|
||
linarith
|
||
|
||
/-- **Value equivalence** (VERIFIED): if `a` achieves total cost `V` under
|
||
the original dynamics at `z`, it achieves the same `V` under the
|
||
pushforward dynamics at `Q z`. Combined with `minimizer_equivalence`,
|
||
this gives the corollary's `V̂*(Q z) = V*(z)` statement. -/
|
||
theorem value_equivalence
|
||
{n : ℕ} {Action : Type*} {T : ℕ}
|
||
(cp : ControlProblem n Action) (Q : Latent n → Latent n)
|
||
(E_hat E : ExpectedCosts n Action T)
|
||
(hinv : IsOrthogonalInvariant cp Q)
|
||
(a : Plan Action T) (z : Latent n) (V : ℝ)
|
||
(hV : totalCost cp E a z = V) :
|
||
totalCost cp E_hat a (Q z) = V := by
|
||
rw [planning_equivalence cp Q E_hat E hinv a z, hV]
|
||
|
||
|
||
end
|