- 移除 JEPA/lejepa-identifiability 子模块 gitlink - 移除 research/multiply/MultiPLY 子模块 gitlink - 删除 .gitmodules（不再有外部 URL 依赖） - 两个目录内容作为普通文件纳入主仓库追踪 - 删除各自内部 .git 目录，消除嵌套 git 仓库
This commit is contained in:
gaojie
@@ -0,0 +1,245 @@
|
||||
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
|
||||
Reference in New Issue
Block a user