strands-vjepa

V-JEPA 2 as a world model, planner, and robot policy for Strands Agents.

Install

pip install strands-vjepa

Extras

Install only what you need:

# Video decoding (decord + av)
pip install strands-vjepa[video]

# HuggingFace pretrained models (transformers + hub)
pip install strands-vjepa[hf]

# Development (pytest, ruff)
pip install strands-vjepa[dev]

# Everything
pip install strands-vjepa[all]

Note: PyTorch is a required dependency. Install it first following pytorch.org for your platform/CUDA version.

What you get

Quick start

Encode + plan (zero-shot, no training)

from strands_vjepa import WorldModel, CEMPlanner
from strands_vjepa.planner import CEMConfig
import torch

wm = WorldModel.from_pretrained("facebook/vjepa2-ac-vitl", device="cuda")
obs = wm.encode(torch.rand(1, 3, 16, 256, 256))
goal = wm.encode(torch.rand(1, 3, 16, 256, 256))

planner = CEMPlanner(wm, CEMConfig(horizon=8, n_samples=256))
result = planner.plan(obs, goal)
print(result.actions.shape)  # [8, 7]

HuggingFace pretrained (real weights)

from strands_vjepa import HFWorldModel

wm = HFWorldModel.from_hf("facebook/vjepa2-vitl-fpc16-256-ssv2")
latents = wm.encode(video)          # [1, 2048, 1024]
logits = wm.classify(video)         # [1, 174] (SSv2 classes)

As Strands Agent tools

from strands import Agent
from strands_vjepa.tools import vjepa_encode, vjepa_plan, vjepa_rollout

agent = Agent(tools=[vjepa_encode, vjepa_plan, vjepa_rollout])
agent("Plan 8 actions from obs.jpg to goal.jpg for a 7-DoF arm")

As a robot policy

import strands_vjepa  # auto-registers with strands-robots
from strands_robots.policies import create_policy

policy = create_policy("vjepa", world_model=wm, action_dim=7, horizon=8)
actions = await policy.get_actions(obs_dict, "goal.jpg")

Fine-tune on your data

from strands_vjepa.training import ACFineTuner, ACTrainingConfig

config = ACTrainingConfig(data_path="/data/trajectories", epochs=100)
ACFineTuner(config).train()

Architecture

strands_vjepa/
├── world_model.py      # Encoder + AC predictor
├── planner.py          # CEM planner
├── policy.py           # strands-robots Policy
├── tools/              # 7 @tool-decorated agent tools
├── training/           # AC fine-tune + JEPA pretrain + DROID dataset
├── eval/               # Frozen-backbone classification probes
└── vendored/           # V-JEPA 2 model code (MIT, from Meta)

Upstream

Built on V-JEPA 2 (Meta AI, Apache 2.0).
This package depends on cagataycali/vjepa2 — a fork with 10 community patches (MPS, ST-A², bugfixes).

License

Apache 2.0