Fine-tune VLA-JEPA on LIBERO (single H100, ~$10)

cd examples/vla-jepa-libero-finetune
super volume create vla-jepa-cache          # once
super create --sku gpu_1x_h100              # → s-xxxx
super upload s-xxxx .
super exec --on s-xxxx --wait 8m -- bash /workspace/vla-jepa-libero-finetune/run.sh setup

# Warm the cache volume, then fine-tune (spawn = survives spot churn; ~2.5h)
super spawn s-xxxx -- sh -c 'bash /workspace/vla-jepa-libero-finetune/run.sh prefetch && \
    STEPS=10000 BATCH=16 SAVE_FREQ=2500 bash /workspace/vla-jepa-libero-finetune/run.sh train'
super logs s-xxxx <pid> --follow

# Optional: train only on one suite (10x more epochs on the demo tasks) —
# write a JSON list of episode indices (libero_object = 807..1260) and pass it:
#   EPISODES_FILE=/workspace/phase2_episodes.json

# Roll out base vs fine-tuned on the same tasks/seed (each saves videos)
super spawn s-xxxx -- sh -c 'GRIPPER_FIX= POLICY=lerobot/VLA-JEPA-Pretrain TAG=base \
    SUITE=libero_object EPISODES=2 bash /workspace/vla-jepa-libero-finetune/run.sh eval'
super spawn s-xxxx -- sh -c 'POLICY=/workspace/train/vla_jepa_libero/checkpoints/last/pretrained_model \
    TAG=finetuned SUITE=libero_object EPISODES=2 bash /workspace/vla-jepa-libero-finetune/run.sh eval'
super exec --on s-xxxx -- bash /workspace/vla-jepa-libero-finetune/run.sh video
super download s-xxxx /workspace/out/side_by_side/reel.mp4 ./reel.mp4
super stop s-xxxx     # clean stop also commits the cache volume for warm restarts

"""lerobot-eval wrapper that inverts the policy's gripper channel.

Fine-tuning VLA-JEPA-Pretrain on LIBERO with lerobot-train produces a
checkpoint whose unnormalizer inverts the gripper dimension at eval time
(the official VLA-JEPA-LIBERO conversion sidesteps this with an
action-normalization mask that lerobot-train does not reproduce). The model's
raw gripper predictions are correct — verified open-loop against ground-truth
actions — so we flip the sign back at the policy boundary and defer a proper
fix to upstream lerobot.
"""

import sys

import torch

from lerobot.policies.vla_jepa.modeling_vla_jepa import VLAJEPAPolicy
from lerobot.scripts.lerobot_eval import main

_orig = VLAJEPAPolicy.predict_action_chunk


def _fixed(self, batch, *args, **kwargs):
    chunk = _orig(self, batch, *args, **kwargs)
    g = self.config.gripper_dim
    chunk = torch.cat(
        [chunk[..., :g], -chunk[..., g : g + 1], chunk[..., g + 1 :]], dim=-1
    )
    return chunk


VLAJEPAPolicy.predict_action_chunk = _fixed

if __name__ == "__main__":
    sys.exit(main())

#!/usr/bin/env bash
# VLA-JEPA → LIBERO fine-tune on a single GPU, with side-by-side eval videos.
#
# Modes (first argument):
#   setup  — create the venv on the cache volume + warm HF caches (idempotent)
#   train  — fine-tune from lerobot/VLA-JEPA-Pretrain (resumes if checkpoints exist)
#   eval   — roll out a checkpoint in the LIBERO sim, saving videos + metrics
#   video  — stitch base-vs-finetuned rollouts into side-by-side MP4s
set -euo pipefail

MODE="${1:-train}"; shift || true

CACHE=/workspace/cache            # venv + HF hub cache (import into a volume post-run)
VENV="$CACHE/venv"
OUT=/workspace/out                # plain workspace dir: durable during the run
TRAIN_DIR="/workspace/train/${OUT_NAME:-vla_jepa_libero}"
export HF_HOME="$CACHE/hf"
# The datasets arrow decode-cache is large (tens of GB) and fully derived —
# keep it out of the synced /workspace so it doesn't swamp the uploader.
export HF_DATASETS_CACHE=/tmp/hf_datasets
export MUJOCO_GL="${MUJOCO_GL:-egl}"
export PYOPENGL_PLATFORM="${PYOPENGL_PLATFORM:-egl}"
export TOKENIZERS_PARALLELISM=false

# Pinned: last vla_jepa fix on lerobot main (2026-07-31).
LEROBOT_REF="${LEROBOT_REF:-4808d8457e55}"

# From some egress ranges github.com serves the ref advertisement but answers
# git's POST /git-upload-pack with a spurious 401 — only over HTTP/2. Git reads
# that as "authenticate", asks for a username, and with no writer on `super run`'s
# PTY the fetch hangs until the run's whole budget expires. The same fetch works
# over HTTP/1.1, and the refusal is deterministic, so pin it rather than retry.
# Global, because the `git+https://` install below shells out to its own git.
git config --global http.version HTTP/1.1
# If a credential prompt ever does get raised, fail loudly instead of hanging.
export GIT_TERMINAL_PROMPT=0

# The DROID-pretrained base checkpoint names its cameras exterior_1_left/
# exterior_2_left; the LIBERO dataset and env expose image (agentview) and
# image2 (wrist). This map bridges them for train and eval alike.
DEFAULT_RENAME_MAP='{"observation.images.image": "observation.images.exterior_1_left", "observation.images.image2": "observation.images.exterior_2_left"}'

rename_args() {
    # ${VAR-default}: default only when unset — RENAME_MAP= explicitly disables.
    local map="${RENAME_MAP-$DEFAULT_RENAME_MAP}"
    if [ -n "$map" ]; then
        printf -- '--rename_map=%s' "$map"
    fi
}

mkdir -p "$OUT"

setup() {
    if [ ! -x "$VENV/bin/python" ]; then
        python3 -m venv "$VENV"
    fi
    "$VENV/bin/pip" install --upgrade pip
    "$VENV/bin/pip" install \
        "lerobot[vla_jepa,libero,eval,dataset,training,transformers-dep] @ git+https://github.com/huggingface/lerobot@${LEROBOT_REF}"
    "$VENV/bin/python" - <<'PY'
import torch
assert torch.cuda.is_available(), "CUDA not visible"
print("torch", torch.__version__, torch.cuda.get_device_name(0))
PY
}

prefetch() {
    "$VENV/bin/python" - <<'PY'
from huggingface_hub import snapshot_download
for repo in ("lerobot/VLA-JEPA-Pretrain",):
    print("hub:", repo, "->", snapshot_download(repo))
from lerobot.datasets.lerobot_dataset import LeRobotDataset
ds = LeRobotDataset("HuggingFaceVLA/libero")
print("dataset episodes:", ds.num_episodes, "frames:", ds.num_frames)
PY
}

prune_checkpoints() {
    # lerobot keeps every numbered checkpoint (~25GB each with optimizer state);
    # keep only the two newest so the disk survives a long run. "last" is a
    # symlink to the newest, which is all --resume needs.
    while true; do
        { ls -d "$TRAIN_DIR"/checkpoints/[0-9]* 2>/dev/null || true; } \
            | sort | head -n -2 | xargs -r rm -rf
        sleep 120
    done
}

train() {
    local steps="${STEPS:-10000}" batch="${BATCH:-16}"
    prune_checkpoints &
    PRUNER_PID=$!
    trap 'kill "${PRUNER_PID:-}" 2>/dev/null || true' EXIT
    local resume_args=()
    if [ -f "$TRAIN_DIR/checkpoints/last/pretrained_model/config.json" ]; then
        resume_args=(--resume=true --config_path="$TRAIN_DIR/checkpoints/last/pretrained_model/train_config.json")
    fi
    local ra
    ra="$(rename_args)"
    # EPISODES_FILE: JSON list of episode indices to train on (e.g. one suite).
    local ep_args=()
    if [ -n "${EPISODES_FILE:-}" ]; then
        ep_args=(--dataset.episodes="$(cat "$EPISODES_FILE")")
    fi
    "$VENV/bin/lerobot-train" \
        ${ra:+"$ra"} \
        "${ep_args[@]}" \
        --policy.path="${POLICY_PATH:-lerobot/VLA-JEPA-Pretrain}" \
        --policy.push_to_hub=false \
        --dataset.repo_id=HuggingFaceVLA/libero \
        --batch_size="$batch" \
        --steps="$steps" \
        --save_freq="${SAVE_FREQ:-1000}" \
        --log_freq=50 \
        --output_dir="$TRAIN_DIR" \
        "${resume_args[@]}"
}

evaluate() {
    # POLICY: HF repo id or a local checkpoint dir; TAG names the output folder.
    local policy="${POLICY:-lerobot/VLA-JEPA-Pretrain}" tag="${TAG:-base}"
    local suite="${SUITE:-libero_object}" episodes="${EPISODES:-10}"
    local ra
    ra="$(rename_args)"
    # GRIPPER_FIX=1 (default for our fine-tuned checkpoints) routes through
    # eval_gripper_fix.py, which un-inverts the gripper channel that
    # lerobot-train's rebuilt unnormalizer flips at eval time. Set GRIPPER_FIX=
    # (empty) for checkpoints with a proper action mask (e.g. the official
    # lerobot/VLA-JEPA-LIBERO).
    local eval_cmd=("$VENV/bin/lerobot-eval")
    if [ -n "${GRIPPER_FIX-1}" ]; then
        eval_cmd=("$VENV/bin/python" "$(dirname "${BASH_SOURCE[0]}")/eval_gripper_fix.py")
    fi
    # libero's first-run init asks whether to customize its dataset folder;
    # answer N (use defaults) since spawned processes have no stdin.
    printf 'N\n' | "${eval_cmd[@]}" \
        ${ra:+"$ra"} \
        --policy.path="$policy" \
        --env.type=libero \
        --env.task="$suite" \
        --eval.n_episodes="$episodes" \
        --eval.batch_size="${EVAL_BATCH:-1}" \
        --output_dir="$OUT/eval_$tag"
}

video() {
    "$VENV/bin/python" /workspace/vla-jepa-libero-finetune/side_by_side.py \
        --base "$OUT/eval_base" --finetuned "$OUT/eval_finetuned" --out "$OUT/side_by_side"
}

case "$MODE" in
    setup) setup ;;
    prefetch) prefetch ;;
    train) train ;;
    eval) evaluate ;;
    video) video ;;
    *) echo "unknown mode: $MODE" >&2; exit 2 ;;
esac

"""Stitch matching base/finetuned LIBERO rollout videos into labeled side-by-side MP4s.

Pairs videos by relative path under the two eval output dirs and runs ffmpeg
hstack with a "before / after" banner. Produces one MP4 per pair plus a
combined reel of every pair concatenated.
"""

import argparse
import pathlib
import subprocess


def find_videos(root: pathlib.Path) -> dict[str, pathlib.Path]:
    return {str(p.relative_to(root)): p for p in sorted(root.rglob("*.mp4"))}


def stitch(base: pathlib.Path, tuned: pathlib.Path, dest: pathlib.Path) -> None:
    dest.parent.mkdir(parents=True, exist_ok=True)
    label = (
        "drawtext=text='%s':fontcolor=white:fontsize=28:box=1:boxcolor=black@0.55:"
        "boxborderw=8:x=16:y=16"
    )
    filt = (
        f"[0:v]fps=30,scale=512:512,{label % 'base (no fine-tune)'}[l];"
        f"[1:v]fps=30,scale=512:512,{label % 'fine-tuned on supercomputer'}[r];"
        "[l][r]hstack=inputs=2[v]"
    )
    subprocess.run(
        ["ffmpeg", "-y", "-i", str(base), "-i", str(tuned), "-filter_complex", filt,
         "-map", "[v]", "-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", "20", str(dest)],
        check=True, capture_output=True,
    )


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", type=pathlib.Path, required=True)
    ap.add_argument("--finetuned", type=pathlib.Path, required=True)
    ap.add_argument("--out", type=pathlib.Path, required=True)
    args = ap.parse_args()

    base_vids = find_videos(args.base)
    tuned_vids = find_videos(args.finetuned)
    pairs = sorted(set(base_vids) & set(tuned_vids))
    if not pairs:
        raise SystemExit(f"no matching videos between {args.base} and {args.finetuned}")

    outputs = []
    for rel in pairs:
        dest = args.out / rel.replace("/", "_")
        stitch(base_vids[rel], tuned_vids[rel], dest)
        outputs.append(dest)
        print(f"stitched {rel} -> {dest}")

    concat_list = args.out / "concat.txt"
    concat_list.write_text("".join(f"file '{p.resolve()}'\n" for p in outputs))
    reel = args.out / "reel.mp4"
    subprocess.run(
        ["ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", str(concat_list),
         "-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", "20", str(reel)],
        check=True, capture_output=True,
    )
    print(f"reel -> {reel}")


if __name__ == "__main__":
    main()

{
    "sku": "gpu_1x_h100",
    "min_disk_gb": 200,
    "image": {
        "base": "nvidia/cuda:12.8.1-devel-ubuntu24.04",
        "apt": [
            "git",
            "cmake",
            "build-essential",
            "ffmpeg",
            "python3-pip",
            "python3-venv",
            "libegl1",
            "libopengl0",
            "libglib2.0-0",
            "fonts-dejavu-core"
        ],
        "env": {
            "PYTHONUNBUFFERED": "1",
            "HF_HOME": "/workspace/cache/hf",
            "MUJOCO_GL": "egl",
            "PYOPENGL_PLATFORM": "egl"
        }
    },
    "volumes": [
        {"name": "vla-jepa-cache", "mount": "/workspace/cache"}
    ]
}