PiJEPA: Language-Conditioned Visual Nav

• PiJEPA is a two-stage framework that combines a reactive, language-conditioned policy with a JEPA-based world model for latent space planning.
• It leverages a shared frozen pretrained visual encoder, enabling consistent latent representations for both reactive policy and trajectory optimization.
• The method uses a policy-guided MPPI warm-start to overcome poor initialization in high-dimensional continuous action spaces, resulting in improved positional accuracy.

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 (Chahe et al., 26 Mar 2026).

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 (Chahe et al., 26 Mar 2026).

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 $o_g$0 and on the instruction $o_g$1, where language is encoded using a pretrained T5 model with 16 tokens. It represents the conditional action distribution

$o_g$2

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 $o_g$3, a learnable predictor $o_g$4, and a learnable action encoder $o_g$5. Given a context of past encoded observations and actions, it predicts

$o_g$6

The predictor uses causal attention so it can operate on varying context lengths up to $o_g$7, 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 (Chahe et al., 26 Mar 2026).

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: $o_g$8 Here $o_g$9 denotes autoregressive rollout of the world model for $\ell$0 steps, with each predicted $\ell$1 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

$\ell$2

where

$\ell$3

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 (Chahe et al., 26 Mar 2026).

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 $\ell$4 action chunks from the policy,

$\ell$5

transforms them into the world-model action frame,

$\ell$6

and initializes the MPPI Gaussian with empirical moments: $\ell$7

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

$\ell$8

the heading is accumulated as

$\ell$9

and the displacement is rotated into the local frame: $(o_t,\ell)$0 The transformed action is then normalized to $(o_t,\ell)$1 using dataset bounds.

MPPI refinement then proceeds by sampling candidate trajectories

$(o_t,\ell)$2

scoring them with the latent planning objective, retaining the top-$(o_t,\ell)$3 elites, and updating the Gaussian by weighted moment matching. Elite weights are

$(o_t,\ell)$4

followed by

$(o_t,\ell)$5

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 (Chahe et al., 26 Mar 2026).

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 $(o_t,\ell)$6 iterations, $(o_t,\ell)$7 candidate samples per iteration, $(o_t,\ell)$8 elites, inverse temperature $(o_t,\ell)$9, $E_\phi$0 policy diffusion samples for the prior, and variance clamp $E_\phi$1. 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 (Chahe et al., 26 Mar 2026).

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 (Yee et al., 1 Apr 2026). JetParticle-JEPA is a particle-level JEPA method for jet tagging in high-energy physics (Letellier et al., 12 Jun 2026). UR-JEPA is presented as a projector-space regularization alternative relevant to PiJEPA-style methods, but it does not describe PiJEPA architecture or experiments (Le, 31 May 2026). 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 (Chahe et al., 26 Mar 2026).