VJEPA2ACModel
PretrainedModelBackboneMixinVJEPA2ACModel(config: VJEPA2ACConfig)V-JEPA 2-AC encoder plus action-conditioned latent predictor.
Parameters
configVJEPA2ACConfigNotes
The encoder is applied independently to each frame, repeated to fill
the patch embedding's tubelet. A video with T frames therefore
produces T action-model steps, each containing
(image_size // patch_size) ** 2 spatial tokens — one step per
frame, as the released training loop and energy-landscape example
both do. forward takes one conditioning row per transition,
so T - 1 of them. Planning and controller search stay outside the
model.
Used by 2
Constructors
1Properties
1Instance methods
6Encode every frame to its own latent grid, (B, T * P, D).
One action-model step is one frame, not one two-frame tubelet.
The encoder's patch embedding is three-dimensional and wants
tubelet_size frames, so each frame is repeated to fill one —
which is what the released code does in both its training loop
and its energy-landscape example:
c.permute(0, 2, 1, 3, 4).flatten(0, 1).unsqueeze(2).repeat(1, 1, 2, 1, 1)
Pairing two distinct frames instead would halve the temporal resolution and take one action per two frames, which is not the rate the released checkpoint was post-trained at.
forward(video: Tensor, actions: Tensor, states: Tensor, extrinsics: Tensor | None = None)One teacher-forced transition step over a clip.
The clip's first T - 1 frames are the context and its last
T - 1 are the target: the predictor is asked what the next
frame looks like, which is the post-training objective. A
conditioning row is therefore one per transition, so actions,
states and extrinsics carry T - 1 of them.
Parameters
Returns
VJEPA2ACOutputThe predicted next latents, the context they came from, the latents actually observed at those frames, and the L1 between them.
Notes
Both sides are layer-normalised before the loss when
normalize_targets is set, as the released loop does. Comparing
a prediction against the frame it was given — which is what
scoring the unshifted latents would do — is satisfied by copying
the input, since the predictor's attention is frame-causal.
Encode a video into flattened per-step spatial tokens.
Layer-normalise latents, as the released loop does to both sides.
predict(context: Tensor, actions: Tensor, states: Tensor, extrinsics: Tensor | None = None)Predict the next latent grids from encoded context and controls.
trainable_parameters()Return encoder and predictor parameters for post-training.