arxiv.org web signal

TD-JEPA mines a temporal cost, lifts OGB-Cube by 14.2 points

robotics generative ai ai-business

TL;DR

  • TD-JEPA keeps the LeWM encoder-predictor backbone and mines a directed temporal cost from reward-free trajectories using same-trajectory step order.
  • Under locked evaluation, the mined cost raises Two-Room success to 100.0% versus LeWM's 97.4% and lifts OGB-Cube by 14.2 points.
  • Ablations show the directed head, cross-trajectory negatives, and rollout-consistency term each contribute; TD-JEPA matches or exceeds the concurrent RC-aux baseline on every environment.

Planning in a learned latent world model has a specific problem: the objective you train the encoder with is not the objective you actually use when you plan. Joint-Embedding Predictive Architectures optimize short-horizon prediction in embedding space, then a planner ranks imagined futures by how close their embeddings are, usually by Euclidean distance. That ranking is a byproduct of the training objective, not something the log data was ever asked to teach.

TD-JEPA, a new method from Jiaxin Bai and Jiaxuan Xiong posted to arXiv, tries to close that gap without replacing the backbone. It retains the LeWM encoder-predictor as-is, then trains a directed head that mines a temporal cost from reward-free trajectories. Same-trajectory step order supplies positive targets, cross-trajectory pairs act as heuristic negatives, and a rollout-consistency term is matched to the planner horizon. The mined cost can be deployed directly at plan time, or it can act as a representation signal that improves the old Euclidean ranking when contact geometry dominates.

Under a locked evaluation protocol, the numbers are concrete. Two-Room success moves to 100.0% from LeWM's 97.4%. On OGB-Cube, shared Euclidean planning on the temporally trained checkpoint sits 14.2 points above LeWM. Push-T improves too, though the abstract does not quantify by how much. Against LeWM and the concurrent RC-aux baseline, TD-JEPA matches or exceeds on every environment tested, and ablations report that the directed head, cross-trajectory negatives, and rollout consistency each carry real weight.

The honest caveat is that the abstract does not report dataset scale, author affiliation, or how the heuristic cross-trajectory negatives are actually sampled, so anyone reproducing this has to read the paper for the recipe rather than the abstract. The Two-Room result also starts from a near-ceiling baseline, so OGB-Cube is the number that carries the story.

For teams already running a JEPA-style planner on offline demonstration logs, the practical read is that the ranking metric is a lever they may not have pulled. Mining a directed temporal cost from the trajectories they already have, then reusing the same encoder-predictor checkpoint, is a cheaper move than retraining a new backbone.