When Does LeJEPA Learn a World Model? — part 5/9

Cleaned-text code points [62462, 78324). Bibliography omitted here; source.md includes it.

To extract a bound on W nl W_{\mathrm{nl}} , we need an upper bound on ℒ ⁡ ( h ) \mathcal{L}(h) . An encoder achieving exact alignment would satisfy ℒ ⁡ ( h ) = 2 ​ ( 1 − ρ ) ​ tr ​ ( Cov ⁡ ( h ⁡ ( z ) ) ) \mathcal{L}(h)=2(1-\rho)\,\mathrm{tr}(\mathrm{Cov}(h(z))) (attained when h h is linear). We assume the encoder is δ \delta -approximately aligned, meaning ℒ ⁡ ( h ) ≤ 2 ​ ( 1 − ρ ) ​ tr ​ ( Cov ⁡ ( h ⁡ ( z ) ) ) + δ . \mathcal{L}(h)\;\leq\;2(1-\rho)\,\mathrm{tr}(\mathrm{Cov}(h(z)))+\delta. (44) Combining the upper and lower bounds and rearranging: W nl ≤ δ 2 ​ ρ ​ ( 1 − ρ ) . \boxed{W_{\mathrm{nl}}\;\leq\;\frac{\delta}{2\rho(1-\rho)}.} (45) The alignment gap δ \delta directly controls how much variance the encoder can place on nonlinear Hermite components. The denominator 2 ​ ρ ​ ( 1 − ρ ) 2\rho(1-\rho) is the spectral gap between degree-1 and degree-2 eigenvalues, and governs the sensitivity of this bound.

### C.2 Step 2: Whitening constraint forces M M near orthogonal

The covariance of h ⁡ ( z ) h(z) decomposes by Hermite orthogonality as Cov ⁡ ( h ⁡ ( z ) ) = M ​ M ⊤ + N , \mathrm{Cov}(h(z))=MM^{\top}+N, (46) where N i ​ j = ∑ | α | ≥ 2 c α ( i ) ​ c α ( j ) N_{ij}=\sum_{|\alpha|\geq 2}c_{\alpha}^{(i)}\,c_{\alpha}^{(j)} is the contribution from nonlinear terms. Since N N is a Gram matrix (positive semidefinite) with tr ⁡ ( N ) = W nl \mathrm{tr}(N)=W_{\mathrm{nl}} , and for any PSD matrix ‖ N ‖ F ≤ tr ⁡ ( N ) \|N\|_{F}\leq\mathrm{tr}(N) (all eigenvalues are non-negative, so ∑ λ i 2 ≤ ( ∑ λ i ) 2 \sum\lambda_{i}^{2}\leq(\sum\lambda_{i})^{2} ), we have ‖ N ‖ F ≤ W nl \|N\|_{F}\leq W_{\mathrm{nl}} .

The approximate whitening assumption gives: ‖ M ​ M ⊤ − I n ‖ F = ‖ Cov ⁡ ( h ⁡ ( z ) ) − I n − N ‖ F ≤ ε + W nl . \|MM^{\top}-I_{n}\|_{F}=\|\mathrm{Cov}(h(z))-I_{n}-N\|_{F}\leq\varepsilon+W_{\mathrm{nl}}. (47) Let M = Q ​ P M=QP be the polar decomposition with Q ∈ O ⁡ ( n ) Q\in O(n) and P P symmetric positive semidefinite with singular values σ 1 , … , σ n \sigma_{1},\ldots,\sigma_{n} . Then M ​ M ⊤ = Q ​ P 2 ​ Q ⊤ MM^{\top}=QP^{2}Q^{\top} , so ‖ P 2 − I n ‖ F = ‖ M ​ M ⊤ − I n ‖ F \|P^{2}-I_{n}\|_{F}=\|MM^{\top}-I_{n}\|_{F} . Since P P is PSD, each σ i ≥ 0 \sigma_{i}\geq 0 , and | σ i − 1 | ≤ | σ i + 1 | ​ | σ i − 1 | = | σ i 2 − 1 | |\sigma_{i}-1|\leq|\sigma_{i}+1||\sigma_{i}-1|=|\sigma_{i}^{2}-1| . Therefore: ‖ M − Q ‖ F = ‖ P − I n ‖ F ≤ ‖ P 2 − I n ‖ F ≤ ε + W nl . \|M-Q\|_{F}=\|P-I_{n}\|_{F}\leq\|P^{2}-I_{n}\|_{F}\leq\varepsilon+W_{\mathrm{nl}}. (48)

### C.3 Step 3: Pythagorean decomposition

Write h ⁡ ( z ) − Q ​ z = ( M − Q ) ​ z ⏟ degree 1 + ν ⁡ ( z ) ⏟ degree ≥ 2 . h(z)-Qz=\underbrace{(M-Q)z}_{\text{degree 1}}+\underbrace{\nu(z)}_{\text{degree}\geq 2}. (49) Since ( M − Q ) ​ z (M-Q)z lives entirely in Hermite degree 1 and ν ⁡ ( z ) \nu(z) lives entirely in degree ≥ 2 \geq 2 , orthogonality of Hermite degrees under γ n \gamma_{n} gives: 𝔼 ⁡ [ ‖ h ⁡ ( z ) − Q ​ z ‖ 2 ] = 𝔼 ⁡ [ ‖ ( M − Q ) ​ z ‖ 2 ] + 𝔼 ⁡ [ ‖ ν ⁡ ( z ) ‖ 2 ] . \mathbb{E}[\|h(z)-Qz\|^{2}]=\mathbb{E}[\|(M-Q)z\|^{2}]+\mathbb{E}[\|\nu(z)\|^{2}]. (50) The cross term 𝔼 ⁡ [ ν ​ ( z ) ⊤ ​ ( M − Q ) ​ z ] = 0 \mathbb{E}[\nu(z)^{\top}(M-Q)z]=0 vanishes exactly. For the first term, since z ∼ 𝒩 ⁡ ( 0 , I n ) z\sim\mathcal{N}(0,I_{n}) : 𝔼 ⁡ [ ‖ ( M − Q ) ​ z ‖ 2 ] = tr ⁡ ( ( M − Q ) ​ 𝔼 ​ [ z ​ z ⊤ ] ​ ( M − Q ) ⊤ ) = ‖ M − Q ‖ F 2 . \mathbb{E}[\|(M-Q)z\|^{2}]=\mathrm{tr}\!\left((M-Q)\,\mathbb{E}[zz^{\top}]\,(M-Q)^{\top}\right)=\|M-Q\|_{F}^{2}. (51) Combining with ( 45 ) and ( 48 ): 𝔼 ⁡ [ ‖ h ⁡ ( z ) − Q ​ z ‖ 2 ] = ‖ M − Q ‖ F 2 + W nl ≤ ( ε + W nl ) 2 + W nl ≤ ( ε + δ 2 ​ ρ ​ ( 1 − ρ ) ) 2 + δ 2 ​ ρ ​ ( 1 − ρ ) . \mathbb{E}[\|h(z)-Qz\|^{2}]=\|M-Q\|_{F}^{2}+W_{\mathrm{nl}}\leq(\varepsilon+W_{\mathrm{nl}})^{2}+W_{\mathrm{nl}}\leq\left(\varepsilon+\frac{\delta}{2\rho(1-\rho)}\right)^{\!2}+\frac{\delta}{2\rho(1-\rho)}. (52) This completes the proof. ∎

##### Role of the Gaussian on the latent side vs. the output side.

The proof uses the Gaussian assumption on z z in two essential ways: the Hermite basis and Mehler’s formula (which provide the spectral decomposition ( 41 )), and the geometric eigenvalue decay ρ d \rho^{d} (which creates the spectral gap exploited in Step 1). By contrast, the proof uses no distributional assumption on h ⁡ ( z ) h(z) beyond zero mean and approximate whitening. In particular, the full distributional Gaussianity condition h ⁡ ( z ) ∼ 𝒩 ⁡ ( 0 , I n ) h(z)\sim\mathcal{N}(0,I_{n}) of Thm. 5.1 can be relaxed to 𝔼 ⁡ [ h ⁡ ( z ) ] = 0 \mathbb{E}[h(z)]=0 and ‖ Cov ⁡ ( h ⁡ ( z ) ) − I n ‖ F ≤ ε \|\mathrm{Cov}(h(z))-I_{n}\|_{F}\leq\varepsilon without any change to the bound.

## Appendix D Proof and Discussion of Theorem 5.4

###### Proof.

Fix any action sequence a 1 : T ∈ 𝒜 T a_{1:T}\in\mathcal{A}^{T} and initial state z 0 z_{0} , and let z ^ 0 = h ⁡ ( z 0 ) = Q ​ z 0 \hat{z}_{0}=h(z_{0})=Qz_{0} . By construction of the pushforward, the joint law of ( z ^ 0 , z ^ 1 , … , z ^ T ) (\hat{z}_{0},\hat{z}_{1},\ldots,\hat{z}_{T}) under p ^ \hat{p} equals the joint law of ( Q ​ z 0 , Q ​ z 1 , … , Q ​ z T ) (Qz_{0},Qz_{1},\ldots,Qz_{T}) under p p . Applying ( 53 ) with R = Q R=Q , ℓ ⁡ ( z ^ t , a t ) = ℓ ⁡ ( Q ​ z t , a t ) = ℓ ⁡ ( z t , a t ) a.s. , ℓ T ​ ( z ^ T ) = ℓ T ​ ( z T ) a.s. \ell(\hat{z}_{t},a_{t})\;=\;\ell(Qz_{t},a_{t})\;=\;\ell(z_{t},a_{t})\quad\text{a.s.},\qquad\ell_{T}(\hat{z}_{T})\;=\;\ell_{T}(z_{T})\quad\text{a.s.} (54) Summing and taking expectations, 𝔼 p ^ ​ [ ∑ t = 0 T − 1 ℓ ⁡ ( z ^ t , a t ) + ℓ T ​ ( z ^ T ) | z ^ 0 ] = 𝔼 p ​ [ ∑ t = 0 T − 1 ℓ ⁡ ( z t , a t ) + ℓ T ​ ( z T ) | z 0 ] . \mathbb{E}_{\hat{p}}\!\left[\sum_{t=0}^{T-1}\ell(\hat{z}_{t},a_{t})+\ell_{T}(\hat{z}_{T})\,\Big|\,\hat{z}_{0}\right]\;=\;\mathbb{E}_{p}\!\left[\sum_{t=0}^{T-1}\ell(z_{t},a_{t})+\ell_{T}(z_{T})\,\Big|\,z_{0}\right]. (55) Since this holds for every a 1 : T a_{1:T} , the minima coincide, giving V ^ ∗ ​ ( z ^ 0 ) = V ∗ ​ ( z 0 ) \hat{V}^{*}(\hat{z}_{0})=V^{*}(z_{0}) , and any minimizer of one side is a minimizer of the other. ∎

### D.1 Worked Applications

The invariance condition ( 53 ) covers a substantial part of control applications. We spell out three.

##### Goal-reaching.

The terminal cost ℓ T ​ ( z ) = ‖ z − z ∗ ‖ 2 \ell_{T}(z)=\|z-z^{*}\|^{2} is O ⁡ ( n ) O(n) -invariant when the state and goal rotate together: ‖ R ​ z − R ​ z ∗ ‖ 2 = ‖ z − z ∗ ‖ 2 \|Rz-Rz^{*}\|^{2}=\|z-z^{*}\|^{2} . A planner choosing a goal z ^ ∗ ∈ ℝ n \hat{z}^{*}\in\mathbb{R}^{n} in the learned latent and minimizing ‖ z ^ T − z ^ ∗ ‖ 2 \|\hat{z}_{T}-\hat{z}^{*}\|^{2} solves the true goal-reaching problem for the goal z ∗ = Q ⊤ ​ z ^ ∗ z^{*}=Q^{\top}\hat{z}^{*} . The learned goal need never be mapped back to the true latent for this to work; the planner’s actions are correct as a sequence, regardless of the coordinate system they are computed in.

##### Linear-quadratic regulation.

Take linear dynamics z t + 1 = A ​ z t + B ​ a t + ε t z_{t+1}=Az_{t}+Ba_{t}+\varepsilon_{t} with ε t ∼ 𝒩 ⁡ ( 0 , Σ ) \varepsilon_{t}\sim\mathcal{N}(0,\Sigma) , and quadratic costs ℓ ⁡ ( z , a ) = z ⊤ ​ W ​ z + a ⊤ ​ R ​ a \ell(z,a)=z^{\top}Wz+a^{\top}Ra , ℓ T ​ ( z ) = z ⊤ ​ W T ​ z \ell_{T}(z)=z^{\top}W_{T}z , with W , W T ⪰ 0 W,W_{T}\succeq 0 and R ≻ 0 R\succ 0 . The pushforward dynamics in z ^ \hat{z} -coordinates are z ^ t + 1 = ( Q ​ A ​ Q ⊤ ) ​ z ^ t + ( Q ​ B ) ​ a t + Q ​ ε t , Q ​ ε t ∼ 𝒩 ⁡ ( 0 , Q ​ Σ ​ Q ⊤ ) , \hat{z}_{t+1}=(QAQ^{\top})\,\hat{z}_{t}+(QB)\,a_{t}+Q\varepsilon_{t},\qquad Q\varepsilon_{t}\sim\mathcal{N}(0,Q\Sigma Q^{\top}), (56) so the LQR problem in z ^ \hat{z} -space has system matrices ( A ^ , B ^ , Σ ^ ) = ( Q ​ A ​ Q ⊤ , Q ​ B , Q ​ Σ ​ Q ⊤ ) (\hat{A},\hat{B},\hat{\Sigma})=(QAQ^{\top},QB,Q\Sigma Q^{\top}) and cost matrices ( W ^ , W ^ T , R ) = ( Q ​ W ​ Q ⊤ , Q ​ W T ​ Q ⊤ , R ) (\hat{W},\hat{W}_{T},R)=(QWQ^{\top},QW_{T}Q^{\top},R) . The discrete algebraic Riccati equation P = A ⊤ ​ P ​ A − A ⊤ ​ P ​ B ​ ( R + B ⊤ ​ P ​ B ) − 1 ​ B ⊤ ​ P ​ A + W P=A^{\top}PA-A^{\top}PB(R+B^{\top}PB)^{-1}B^{\top}PA+W (57) is covariant under ( A , B , W , P ) ↦ ( Q ​ A ​ Q ⊤ , Q ​ B , Q ​ W ​ Q ⊤ , Q ​ P ​ Q ⊤ ) (A,B,W,P)\mapsto(QAQ^{\top},QB,QWQ^{\top},QPQ^{\top}) : if P P solves the true-latent equation, then P ^ = Q ​ P ​ Q ⊤ \hat{P}=QPQ^{\top} solves the z ^ \hat{z} -space equation. The optimal feedback gain K = ( R + B ⊤ ​ P ​ B ) − 1 ​ B ⊤ ​ P ​ A K=(R+B^{\top}PB)^{-1}B^{\top}PA transforms as K ^ = K ​ Q ⊤ \hat{K}=KQ^{\top} , and the feedback action is a ^ t = − K ^ ​ z ^ t = − ( K ​ Q ⊤ ) ​ ( Q ​ z t ) = − K ​ z t , \hat{a}_{t}=-\hat{K}\hat{z}_{t}=-(KQ^{\top})(Qz_{t})=-Kz_{t}, (58) which is the true LQR action at every state. The value function is likewise V ^ ∗ ​ ( z ^ ) = z ^ ⊤ ​ P ^ ​ z ^ = z ⊤ ​ P ​ z = V ∗ ​ ( z ) \hat{V}^{*}(\hat{z})=\hat{z}^{\top}\hat{P}\hat{z}=z^{\top}Pz=V^{*}(z) . LQR is therefore a special case of Thm. 5.4 in which the invariance is inherited from the quadratic structure and the covariance of all system matrices under the rotation.

##### General rotation-invariant objectives.

Any cost depending on the state only through quantities preserved by rotations satisfies ( 53 ). Examples: norms ‖ z ‖ \|z\| and radial functions f ⁡ ( ‖ z ‖ ) f(\|z\|) ; inner products ⟨ z , z ′ ⟩ \langle z,z^{\prime}\rangle between pairs of states; quadratic forms z ⊤ ​ W ​ z z^{\top}Wz with W W specified in the learned coordinates (which silently corresponds to Q ⊤ ​ W ​ Q Q^{\top}WQ in the true latent); Gaussian belief updates, whose covariance updates are covariant under rotation; and linear value functions w ⊤ ​ z w^{\top}z with w w specified in the learned coordinates. In each case, an engineer who writes down the cost using only these operations has unwittingly written down a cost in the equivalence class of ( 53 ), and Thm. 5.4 certifies that planning is unaffected by the ambiguity.

### D.2 Scope: encoder versus action-conditioned dynamics

Thm. 5.4 concerns the encoder alone. The action-conditioned transition p ^ ​ ( z ^ ′ ∣ z ^ , a ) \hat{p}(\hat{z}^{\prime}\mid\hat{z},a) must still be learned from data [ 9 ] ; our theorems do not prove that it is. What they do guarantee is that the encoder does not corrupt the learning problem for the transition: for h ⁡ ( z ) = Q ​ z h(z)=Qz , the pushforward of the true transition lies in the same function class as the original (linear transitions remain linear, Gaussian transitions remain Gaussian, and any class closed under orthogonal reparametrization of the state is equally expressive in both coordinate systems). The residue of the orthogonal ambiguity is a consistent reparametrization across the encoder, the transition model, and the cost: so long as these three components are trained or specified to be mutually consistent in the learned coordinates, planning in those coordinates is planning in the true ones.

The natural next theoretical step is to prove that the action-conditioned transition itself is identifiable. This is a genuinely different problem: the encoder identifiability proved here relies on stationarity plus Gaussianity of the latent distribution, whereas transition identifiability would trade these for an excitation condition on the action sequence (the actions must explore all latent directions). Persistent excitation is the classical condition in subspace identification [ 66 , 67 ] , and extending LeJEPA-style identifiability to the controlled setting is the subject of ongoing work.

### D.3 Connection to classical optimal control

The covariance of LQR under orthogonal reparametrization has been exploited for decades in control theory: the Kalman filter, the Riccati equation, and Lyapunov stability analysis all transform covariantly under orthogonal changes of basis, and this is why principal-component and whitening transforms are used as preprocessing steps in classical state estimation. What is new here is not the covariance itself but its role as a consequence of an identifiability guarantee . In classical control, the state is assumed given; the engineer chooses a coordinate system. In the LeJEPA setting, the coordinate system is chosen by the learning procedure, and the theorem is that the choice is made consistently enough, orthogonal rather than arbitrary linear (let alone arbitrary nonlinear), that classical control tools continue to apply without modification. Latent-space planning with a learned world model has a long heritage [ 88 , 44 , 50 ] ; Thm. 5.4 provides a representation-learning foundation for the classical pipeline.

## Appendix E Alternative Proof via Dirichlet Energy

We provide a second, independent proof of identifiability using a different noise model and different mathematical machinery. This proof requires stronger regularity ( C 1 C^{1} diffeomorphism and infinitesimal noise limit) but offers complementary geometric insight and makes explicit contact with the local isometry approach used in prior nonlinear ICA [ 74 ] .

### E.1 Setup

Consider the alternative noise model z ′ = z + ϵ ​ η z^{\prime}=z+\sqrt{\epsilon}\,\eta with η ∼ 𝒩 ⁡ ( 0 , I n ) \eta\sim\mathcal{N}(0,I_{n}) , define the normalized loss: ℒ ⁡ ( h ) = lim ϵ → 0 1 ϵ ​ 𝔼 ​ [ ‖ h ⁡ ( z + ϵ ​ η ) − h ⁡ ( z ) ‖ 2 ] . \mathcal{L}(h)=\lim_{\epsilon\to 0}\frac{1}{\epsilon}\,\mathbb{E}\!\left[\|h(z+\sqrt{\epsilon}\eta)-h(z)\|^{2}\right]. (59) We require h : ℝ n → ℝ n h:\mathbb{R}^{n}\to\mathbb{R}^{n} to be a C 1 C^{1} diffeomorphism with h ⁡ ( z ) ∼ 𝒩 ⁡ ( 0 , I n ) h(z)\sim\mathcal{N}(0,I_{n}) when z ∼ 𝒩 ⁡ ( 0 , I n ) z\sim\mathcal{N}(0,I_{n}) .

### E.2 Proof

##### Step 1: Reduction to Dirichlet energy.

By the first-order Taylor expansion h ⁡ ( z + ϵ ​ η ) − h ⁡ ( z ) = ϵ ​ J h ​ ( z ) ​ η + o ⁡ ( ϵ ) h(z+\sqrt{\epsilon}\eta)-h(z)=\sqrt{\epsilon}\,J_{h}(z)\,\eta+o(\sqrt{\epsilon}) , where J h ​ ( z ) J_{h}(z) is the Jacobian of h h at z z . Squaring, dividing by ϵ \epsilon , taking expectations, and using 𝔼 η ​ [ ‖ A ​ η ‖ 2 ] = ‖ A ‖ F 2 \mathbb{E}_{\eta}[\|A\eta\|^{2}]=\|A\|_{F}^{2} for any matrix A A : ℒ ⁡ ( h ) = 𝔼 z ∼ γ n ​ [ ‖ J h ​ ( z ) ‖ F 2 ] , \mathcal{L}(h)=\mathbb{E}_{z\sim\gamma_{n}}\!\left[\|J_{h}(z)\|_{F}^{2}\right], (60) which is the Dirichlet energy of h h with respect to the Gaussian measure γ n \gamma_{n} .

The exchange of limit and expectation is justified by dominated convergence: the C 1 C^{1} regularity of h h combined with the Gaussian decay of γ n \gamma_{n} provides the necessary integrability. The full computation, including careful treatment of the remainder term, is given below.

Starting from the Taylor expansion: h ⁡ ( z + ϵ ​ η ) − h ⁡ ( z ) = ϵ ​ J h ​ ( z ) ​ η + R ⁡ ( z , ϵ ​ η ) , h(z+\sqrt{\epsilon}\eta)-h(z)=\sqrt{\epsilon}\,J_{h}(z)\,\eta+R(z,\sqrt{\epsilon}\eta), (61) where ‖ R ⁡ ( z , ϵ ​ η ) ‖ = o ⁡ ( ϵ ) \|R(z,\sqrt{\epsilon}\eta)\|=o(\sqrt{\epsilon}) . Expanding the squared norm: ‖ h ⁡ ( z + ϵ ​ η ) − h ⁡ ( z ) ‖ 2 = ϵ ​ ‖ J h ​ ( z ) ​ η ‖ 2 + 2 ​ ϵ ​ ( J h ​ ( z ) ​ η ) ⊤ ​ R + ‖ R ‖ 2 . \|h(z{+}\sqrt{\epsilon}\eta)-h(z)\|^{2}=\epsilon\,\|J_{h}(z)\eta\|^{2}+2\sqrt{\epsilon}\,(J_{h}(z)\eta)^{\top}R+\|R\|^{2}. (62) Dividing by ϵ \epsilon and taking ϵ → 0 \epsilon\to 0 , the last two terms vanish, yielding: 1 ϵ ​ 𝔼 z , η ​ [ ‖ h ⁡ ( z + ϵ ​ η ) − h ⁡ ( z ) ‖ 2 ] → ϵ → 0 𝔼 z , η ​ [ ‖ J h ​ ( z ) ​ η ‖ 2 ] . \frac{1}{\epsilon}\,\mathbb{E}_{z,\eta}\!\left[\|h(z{+}\sqrt{\epsilon}\eta)-h(z)\|^{2}\right]\;\xrightarrow{\epsilon\to 0}\;\mathbb{E}_{z,\eta}\!\left[\|J_{h}(z)\eta\|^{2}\right]. (63) Using 𝔼 η ​ [ η ⊤ ​ M ​ η ] = tr ⁡ ( M ) \mathbb{E}_{\eta}[\eta^{\top}M\eta]=\mathrm{tr}(M) for M = J h ​ ( z ) ⊤ ​ J h ​ ( z ) M=J_{h}(z)^{\top}J_{h}(z) : 𝔼 η ​ [ ‖ J h ​ ( z ) ​ η ‖ 2 ] = tr ⁡ ( J h ​ ( z ) ⊤ ​ J h ​ ( z ) ) = ‖ J h ​ ( z ) ‖ F 2 . \mathbb{E}_{\eta}\!\left[\|J_{h}(z)\eta\|^{2}\right]=\mathrm{tr}(J_{h}(z)^{\top}J_{h}(z))=\|J_{h}(z)\|_{F}^{2}. (64)

##### Step 2: The expected log-determinant is zero.