---
title: 'PiJEPA: Language-Conditioned Visual Nav'
url: https://www.emergentmind.com/topics/pijepa
type: topic
---

# PiJEPA: Language-Conditioned Visual Nav

Searching arXiv for the PiJEPA paper and closely related JEPA-family context.
PiJEPA is a two-stage framework for **language-conditioned visual navigation** that combines a reactive instruction-following policy with latent world-model planning. In the setting studied, the agent receives a current egocentric RGB observation \(o_t\), a goal image \(o_g\), and a natural-language instruction \(\ell\), and must output navigation actions that move toward the visually specified goal while respecting the instruction. The method addresses two complementary weaknesses: reactive policies struggle with long-horizon reasoning, while world-model planning in continuous high-dimensional action spaces is hindered by poor initialization. PiJEPA closes this gap by using a finetuned policy to provide an informative action prior and a JEPA world model to refine that prior through latent-space planning [2603.25981].

## 1. Problem setting and conceptual motivation

PiJEPA is formulated for navigation tasks in which local action plausibility is not sufficient for global success. A purely reactive policy can map \((o_t,\ell)\) directly to actions, but it cannot explicitly evaluate future consequences over long horizons. Conversely, a world-model planner can reason over long horizons, yet sampling-based optimization is inefficient when it begins from an uninformed prior over action sequences.

The framework’s central idea is to split responsibilities. The policy contributes **instruction grounding** and short-horizon competence; the world model contributes **look-ahead** and trajectory optimization. This suggests a decomposition in which language semantics enter through the policy, while the planner refines policy-proposed behaviors using latent dynamics rollouts. A key detail is that both stages operate in the embedding space of the same **frozen pretrained visual encoder** \(E_\phi\), so the policy and world model share a consistent latent representation. Given an image \(o_t\), the encoder produces
\[
z_t = E_\phi(o_t) \in \mathbb{R}^{n \times d},
\]
and likewise
\[
z_g = E_\phi(o_g).
\]

The paper studies two encoder choices: **DINOv2 ViT-S** and **V-JEPA-2 ViT-L**. This encoder-sharing design is fundamental to PiJEPA’s architecture, because the policy prior and the planner’s latent evaluation are defined in the same visual feature space [2603.25981].

## 2. Two-stage architecture

The first stage is a finetuned **Octo-based generalist policy**. Octo is a transformer VLA policy with a **diffusion action head** that models a distribution over action chunks. In PiJEPA, the original Octo visual encoder is replaced by the frozen pretrained encoder \(E_\phi\) followed by a learnable projection
\[
W_{\text{proj}} \in \mathbb{R}^{d \times d_{\text{octo}}}.
\]
The policy is conditioned on the encoded current observation \(z_t\) and on the instruction \(\ell\), where language is encoded using a pretrained **T5** model with 16 tokens. It represents the conditional action distribution
\[
\pi_\theta(a \mid o_t,\ell),
\]
and produces action chunk samples through diffusion sampling.

Those samples are not only candidate actions; they also define the initialization for planning. This is the architectural novelty of PiJEPA: the policy is not merely evaluated as a baseline but used to define the **initial search distribution** for the planner.

The second stage is a **Joint-Embedding Predictive Architecture world model (JEPA-WM)**. Rather than reconstructing pixels, it predicts future **latent visual states** in the embedding space of the same frozen encoder. The world model consists of the frozen encoder \(E_\phi\), a learnable predictor \(P_\psi\), and a learnable action encoder \(A_\psi\). Given a context of past encoded observations and actions, it predicts
\[
\hat{z}_{t+1} = P_\psi\!\left(z_{t-w:t},\, A_\psi(a_{t-w:t})\right).
\]

The predictor uses **causal attention** so it can operate on varying context lengths up to \(w\), and action information is injected at every layer using **Adaptive Layer Normalization (AdaLN)** conditioning. The resulting system is therefore asymmetric by design: the policy is language-conditioned and reactive, whereas the world model is a language-agnostic latent dynamics model that supports planning [2603.25981].

## 3. Latent dynamics model and planning objective

The JEPA world model is trained by minimizing mean-squared error between predicted latent states and target latent states over multi-step rollouts:
\[
L_{\text{total}} = \sum_{k=1}^{K_{\text{roll}}} \frac{1}{B}\sum_{b=1}^{B} \left\| F_{\phi,\psi}(z_t^b, a_{t:t+k-1}^b) - z_{t+k}^b \right\|_2^2.
\]
Here \(F_{\phi,\psi}\) denotes autoregressive rollout of the world model for \(k\) steps, with each predicted \(\hat z\) fed back as context. Training uses **truncated backpropagation through time** on CAST trajectory segments, while the encoder remains frozen.

At inference time, the planner evaluates candidate action sequences by rolling them forward in latent space and comparing the final predicted latent state to the encoded goal. The planning cost is
\[
L^p(z_t, a_{1:H}, z_g) = \frac{1}{n}\sum_{i=1}^{n} \left\| \hat{z}_{t+H}^{(i)} - z_g^{(i)} \right\|_2^2,
\]
where
\[
\hat z_{t+H} = F_{\phi,\psi}(z_t, a_{1:H}).
\]

This objective is explicitly **terminal latent-state matching**. There is no learned reward model and no explicit intermediate waypoint reward in the reported method. The world model itself is also explicitly **language-agnostic**: it models only visual-action dynamics. Language affects planning indirectly through the policy prior rather than through the world-model loss. This separation of responsibilities is one of PiJEPA’s defining assumptions. A plausible implication is that the framework is especially well matched to settings where language primarily selects among visually grounded trajectories rather than changing the underlying transition dynamics [2603.25981].

## 4. Policy-guided MPPI warm-start

PiJEPA’s central mechanism is the use of policy samples to warm-start **Model Predictive Path Integral (MPPI)** planning. Standard MPPI begins from an uninformed Gaussian in action space. The paper argues that this is inefficient for long-horizon navigation in continuous spaces because much of the sampling budget is spent on implausible trajectories.

PiJEPA instead samples \(N_\pi\) action chunks from the policy,
\[
a^{(i)} \sim \pi_\theta(\cdot \mid o_t,\ell), \qquad i=1,\dots,N_\pi,
\]
transforms them into the world-model action frame,
\[
\tilde a^{(i)} = T(a^{(i)}),
\]
and initializes the MPPI Gaussian with empirical moments:
\[
\mu^0 = \frac{1}{N_\pi}\sum_{i=1}^{N_\pi} \tilde a^{(i)}, \qquad
\sigma^0 = \operatorname{clamp}\!\left(\operatorname{std}(\{\tilde a^{(i)}\}),\, \sigma_{\min},\, \sigma_{\max}\right).
\]

A systems detail is essential here: the policy and world model use different action coordinate frames. Octo outputs **global-frame displacements**, while the world model expects **local body-frame actions**. If
\[
a_t = (\Delta x_t, \Delta y_t, \sin\Delta\phi_t, \cos\Delta\phi_t),
\]
the heading is accumulated as
\[
\phi_t = \sum_{\tau=0}^{t-1} \operatorname{atan2}(\sin\Delta\phi_\tau, \cos\Delta\phi_\tau),
\]
and the displacement is rotated into the local frame:
\[
\begin{pmatrix}
\Delta x_t^\ell \\
\Delta y_t^\ell
\end{pmatrix}
=
\begin{pmatrix}
\cos\phi_t & \sin\phi_t \\
-\sin\phi_t & \cos\phi_t
\end{pmatrix}
\begin{pmatrix}
\Delta x_t \\
\Delta y_t
\end{pmatrix}.
\]
The transformed action is then normalized to \([-1,1]\) using dataset bounds.

MPPI refinement then proceeds by sampling candidate trajectories
\[
a^{(i)}_{1:H} = \operatorname{clamp}\big(\mu^j + \sigma^j \odot \epsilon^{(i)},\, -1,\, 1\big), \qquad \epsilon^{(i)} \sim \mathcal{N}(0,I),
\]
scoring them with the latent planning objective, retaining the top-\(K\) elites, and updating the Gaussian by weighted moment matching. Elite weights are
\[
w_k = \frac{\exp\!\big(\lambda (c_{\min}-c_k)\big)}
{\sum_{k'=1}^{K}\exp\!\big(\lambda (c_{\min}-c_{k'})\big)},
\]
followed by
\[
\mu^{j+1} = \sum_k w_k\, a^{(k)}, \qquad
(\sigma^{j+1})^2 = \sum_k w_k \big(a^{(k)} - \mu^{j+1}\big)^2.
\]

The paper’s algorithm is therefore: encode current observation and goal, sample policy actions, transform them into world-model coordinates, initialize MPPI from their empirical mean and standard deviation, run iterative latent-space planning, and execute the resulting trajectory. The authors explicitly describe this as using the policy to define the **initial search distribution** rather than treating policy execution and planning as separate alternatives [2603.25981].

## 5. Training setup and empirical performance

Both policy and world model are trained and evaluated on **CAST**, a visual navigation dataset with **counterfactual instruction-action augmentation**. CAST is described as designed to reduce posterior collapse in instruction-conditioned navigation, where observations alone can otherwise dominate the action prediction. This is particularly relevant to PiJEPA because its planner relies on a policy prior that must remain sensitive to language.

The policy network is **Octo-Small** with a **3-layer MLP with residual connections** as the action head, trained with the **DDPM diffusion objective** using a cosine noise schedule. The world model predictor is a **Vision Transformer** with **frame-causal attention** and **AdaLN action conditioning** at every transformer block. For **V-JEPA-2**, each image frame is encoded via a **frame-duplication strategy**, and **layer normalization** is applied to encoder outputs to stabilize both prediction targets and the planning cost landscape. The paper reports that all experiments were run on **four NVIDIA H200 GPUs**.

The reported MPPI hyperparameters are \(J=4\) iterations, \(N=32\) candidate samples per iteration, \(K=4\) elites, inverse temperature \(\lambda=0.8\), \(N_\pi=4\) policy diffusion samples for the prior, and variance clamp \([\sigma_{\min},\sigma_{\max}] = [0.01,0.05]\). Inference latency is reported as about **2.13 s** for Octo diffusion sampling plus **0.35 s** for MPPI, for a total of about **2.48 s** for an 8-step trajectory.

Evaluation on **CAST validation episodes** uses **ATE XY**, **ATE Heading**, **RPE XY**, and **RPE Heading**, with lower values better for all. Four methods are compared: **Octo**, **MPPI**, **Octo-WM**, and **PiJEPA**. Under **DINOv2**, PiJEPA achieves **ATE XY RMSE 1.78, Mean 1.42, Final 3.12**, compared with **Octo** at **1.98, 1.65, 3.19**, **MPPI** at **1.85, 1.48, 3.23**, and **Octo-WM** at **1.80, 1.43, 3.15**. Under **V-JEPA-2**, PiJEPA achieves **ATE XY RMSE 1.65, Mean 1.32, Final 2.88**, while **Octo-WM** slightly exceeds it on final position at **2.87**, with PiJEPA best on most positional metrics. The paper’s stated conclusion is that **PiJEPA achieves the best positional accuracy across both encoders**, while uninformed MPPI often performs best on heading metrics and Octo-WM captures much of the benefit of planning without iterative optimization [2603.25981].

## 6. Limitations, failure modes, and position within the JEPA family

PiJEPA makes several strong modeling assumptions. Most importantly, it uses a **decoupled design** in which language enters only through the policy, while the world model is language-agnostic. This yields a clean modular decomposition, but it also means the planner’s objective does not directly score instruction semantics beyond the structure already encoded by the policy prior. A second limitation is the **terminal-only latent matching objective**: the planner optimizes final latent proximity to the goal image, with no explicit dense reward or waypoint objective.

The main reported failure case is the instruction **“Follow the building.”** The paper identifies two reasons. First, the instruction is ambiguous because multiple buildings are visible, so the policy may ground the language to the wrong referent. Second, the world model can **stagnate**, producing nearly identical latent predictions across the rollout horizon regardless of action, which makes planning unable to distinguish trajectories. In that case, PiJEPA still performs better than either component alone because the policy prior keeps search in a productive region, but it still undershoots ground truth. The paper lists future directions including improved dynamics architectures to prevent rollout stagnation, diversity-promoting rollout mechanisms, intermediate waypoint costs, richer action spaces, and closed-loop replanning.

Within the broader JEPA literature, PiJEPA should be distinguished from similarly named methods. **PI-JEPA** denotes **Physics-Informed Joint Embedding Predictive Architecture** for label-free surrogate pretraining in coupled multiphysics simulation, not navigation [2604.01349]. **JetParticle-JEPA** is a particle-level JEPA method for jet tagging in high-energy physics [2606.14813]. **UR-JEPA** is presented as a projector-space regularization alternative relevant to PiJEPA-style methods, but it does not describe PiJEPA architecture or experiments [2606.01443]. Accordingly, PiJEPA’s specific contribution is not a generic JEPA regularizer or a domain-agnostic latent predictor. Its novelty lies in using a **language-conditioned policy as a prior over action sequences** to warm-start **MPPI planning in the latent space of a JEPA world model**, thereby addressing poor action initialization in long-horizon navigation [2603.25981].

Source: https://www.emergentmind.com/topics/pijepa