JEPA Policy

English | 中文

JEPA Policy is a diffusion-free imitation-learning policy that predicts robot actions and a future observation's representation together. Both targets come from the same demonstration. A shared Transformer makes the joint prediction in two forward passes, with the visual encoder trained from scratch.

This repository contains training and evaluation code, matched baseline configurations, and the real-robot deployment stack. It builds on the Minimum Flow Policies codebase.

A 1 min 45 sec overview of the method, Coffee Preparation and Tool Hang comparisons, inference latency, and five real-robot tasks. Play here or watch on the project page. The video is silent; full-screen playback is available for the method diagrams.

Architecture comparison with Diffusion Policy, plain JEPA, and ACT-JEPA. JEPA Policy uses shared attention for action and future-representation prediction. This figure compares designs, not performance; click to view the original.

• Joint action and future prediction, with two-step MIP training and inference
• End-to-end visual encoder training, with stop-gradient on future targets
• Image/state observation support and Hydra training/evaluation configurations
• robomimic, LIBERO, and MimicGen task configurations
• ARX5/X5 deployment for JEPA Policy, MIP, and Diffusion Policy

Success rates on nine simulated tasks, including Hammer Cleanup. Lines show the three-seed mean best-checkpoint success rate; shaded bands show the seed range. The radial axis spans 40–100%; the downward triangle marks a seed below the axis floor. See the project page and paper for the full results and evaluation protocol.

JEPA Policy requires Python 3.12 and PyTorch. The default environment supports robomimic:

git clone https://github.com/jiejie567/JEPA-Policy.git
cd JEPA-Policy
uv sync --extra dev

For headless MuJoCo rendering:

export MUJOCO_GL=egl

LIBERO and MimicGen use different robosuite versions. Install them in separate environments following their official installation instructions, then install this repository in editable mode in each environment.

Train JEPA Policy on robomimic Square:

uv run examples/train_robomimic.py \
  -cn exps/jepa_policy.yaml \
  task=square_ph_image

Train the matched action-only MIP baseline:

uv run examples/train_robomimic.py \
  -cn exps/mip_action_only.yaml \
  task=square_ph_image

Evaluate a checkpoint:

uv run examples/train_robomimic.py \
  -cn exps/jepa_policy.yaml \
  task=square_ph_image \
  mode=eval \
  optimization.model_path=/path/to/checkpoint.pt

Override any Hydra option from the command line, for example:

uv run examples/train_robomimic.py \
  -cn exps/jepa_policy.yaml \
  task=tool_hang_ph_image \
  optimization.seed=41 \
  optimization.gradient_steps=300000

The real_robot/ directory contains the inference stack used for the ARX5/X5 experiments. It evaluates JEPA Policy, the matched MIP baseline, and Diffusion Policy through the same observation, action, safety, recording, and operator-labeling pipeline.

The real-robot release includes:

• the reported deployment source snapshots and pinned ARX5 SDK revision;
• offline and observation-only dry-runs before hardware control is enabled;
• launchers for all five reported tasks and all three methods;
• the external checkpoint/statistics layout and a 75-checkpoint artifact manifest; and
• explicit hardware, camera, control-frequency, calibration, and emergency-stop requirements.

Checkpoints, datasets, robot recordings, camera serial numbers, compiled SDK binaries, and internal machine paths are intentionally excluded from Git. Start with the real-robot installation and safety guide before running any launcher.

The exact action-only MIP preset and the aligned Diffusion Policy recipe used for the main comparison are documented in baselines/. The Diffusion Policy material pins the upstream commit and supplies only the task configurations and small protocol overlay needed to reproduce the paper; it does not duplicate the third-party repository.

The repository includes configurations for standard robomimic image/state tasks and the image tasks used to validate JEPA Policy:

• robomimic: Square, Tool Hang, and four-camera Transport
• LIBERO: MokaMoka and MugMug
• MimicGen: Coffee Preparation, Kitchen, and Three Piece Assembly

robomimic datasets are resolved through the dataset repository configured in examples/configs/task/robomimic_base.yaml.

For LIBERO, set the installation and dataset paths:

export LIBERO_ROOT=/path/to/LIBERO

For MimicGen, download the three task datasets and set their root:

bash tools/download_mimicgen_core.sh
export MIMICGEN_DATA_ROOT=$PWD/datasets/mimicgen/core

The public JEPA Policy preset uses:

• action horizon 10 for robomimic/MimicGen and 16 for LIBERO
• two MIP prediction steps with interpolation time 0.9
• one future-representation token
• future horizon 4
• adaptive future-loss ratio 0.1
• shared temporal crop for current and future observations

Task files define environment-specific horizons, observations, action dimensions, and dataset locations. The method configuration remains shared across tasks.

uv run pytest -q \
  tests/test_joint_action_future_mip.py \
  tests/test_future_loss_ratio.py \
  tests/test_temporal_consistent_crop.py

To produce a history-free review archive from the exact current commit:

bash tools/export_anonymous.sh /path/to/JEPA-Policy-anonymous.tar.gz

Set ANON_REPOSITORY_URL to replace the public clone URL with the anonymous review URL. The exporter refuses to package known author identifiers or machine-specific paths.

This codebase extends Minimum Flow Policies and uses robomimic, robosuite, LIBERO, and MimicGen.

Released under the MIT License. Third-party benchmarks and datasets retain their own licenses.