V-JEPA 2 Drake Simulation

Understand and use V-JEPA 2 using Drake simulation library for RL and ML tasks

Python,
Drake,
uv,
Pytorch,
RL

V-JEPA 2 Integration with Drake Simulation

V-JEPA 2 pose estimator running in Drake simulation

Overview

This project looks into finding if a pretrained JEPA video embeddings can capture meaningful differences in robot motion and provide a useful signal for control. I use a Drake simulation with a KUKA iiwa arm to generate trajectories, extract V-JEPA2 embeddings from the rendered motion, and evaluate whether the latent space separates different movement behaviors. I then use embedding similarity as a dense progress signal for a reaching task, with the goal of connecting learned video representations to practical robot control.

In Progress

V-JEPA 2 also introduced a Action-Conditioned world model V-JEPA 2-AC, the goal is to use the visual latent space from V-JEPA 2 and the action of the arm to predict the next latent space. My understanding of metas work flow:

Freeze the pretrained V-JEPA 2
Train the model using arm trajectories
Sample actions, predict the next timestep, compare to the goal embedding and run the best action
Repeat in a MPC loop.

Goal is to fine tune their model for the KUKA iiwa, Meta used the Frank Panda which is included with Drake but for learning purposes using KUKA iiwa seems more interesting.

RL Reward

We can utilize the latent space from V-JEPA2 to create a reward function based only on the image. For this our reward function will take the cosine similarity between our embedding within a dt window (in this case 1 second) and a goal embedding, which comes from the last second of a single demonstration video.

Reward Function

$$r(t) = \cos\big(\text{embed}(\text{last 1 s}),\ \text{embed}(\text{goal})\big)$$

We have 4 states to show this somewhat works: a nominal reach that finishes the task, one that stalls halfway through, one that undershoots the target, and one doing the wrong motion entirely. The reward ranks all 4 correctly without knowing anything about the task itself. So we get a signal that says "warmer" or "colder" at every timestep, from one demo video and zero reward engineering.

V-JEPA2 reward curves for the four rollouts

Two caveats. The raw cosine lives in a narrow band (~0.85 to 1.0) because the embedding space is anisotropic, so a real RL loop would likely likely normalize per task, e.g. pin the start frame at 0 and the goal at 1.
Second, the margins are small enough that compression matters decoding H.264 encoded frames shifts the embeddings, so everything here is scored on raw rendered frames, not decoded mp4s.

Still Ongoing Project

I'm still working on this to finetune V-JEPA 2-AC to work in both drakes simulation software and the KUKA iiwa, where as Meta used Frank Panda in thier model.

GitHub