EB-VLA: Vision-Language-Action Model with an Energy-Based Action Head

An implementation of a Vision-Language-Action (VLA) model that replaces the usual diffusion action head with an Energy-Based Transformer (EBT) head. Instead of predicting actions directly, the head learns a scalar energy landscape E(a, C) over candidate action chunks a conditioned on multimodal context C (vision, language, proprioception). Actions are produced by running MCMC gradient descent on this landscape at inference time, and the converged energy value can be used as a built-in confidence signal.

What's in this repo

The core contribution lives in a single file:

• prismatic/models/action_heads.py — the energy-based action head. Key classes:
  • VLAEnergyTransformer — the transformer that maps (action_chunk, VLM_context) to a scalar energy per timestep, with bidirectional self-attention over action tokens, cross-attention into VLM features, RoPE, RMSNorm, and a SwiGLU FFN.
  • ActionHeadEBT — the wrapper that runs MCMC inference (Langevin-perturbed gradient descent on the energy, with energy-scaled per-timestep step size) and exposes the training loss.

Everything else is a fairly standard VLA training/eval stack adapted from the OpenVLA-OFT / Prismatic / VLA-Adapter lineage.

Repository layout

├── prismatic/
│   └── models/
│       └── action_heads.py        # EBT energy head (main contribution)
├── vla-scripts/
│   ├── train.py                   # training loop
│   ├── finetune.py                # LoRA fine-tuning of the VLM backbone
│   └── vla_evaluation.py          # LIBERO evaluation
├── experiments/
│   └── robot/libero/              # LIBERO + LIBERO-Plus eval entrypoint
├── configs/
│   ├── train/                     # training configs (ebt_spatial*.yaml, ...)
│   └── eval/                      # evaluation configs
├── run_scripts/
│   ├── slurm/                     # SLURM launchers (train + perturbation evals)
│   └── train_scripts/             # plain bash training launchers
├── scripts/                       # analysis utilities (energy plots, scoring, ...)
├── run.py                         # unified entry point (train / eval)
└── requirements_vla_ebt.txt

Installation

conda create -n eb-vla python=3.10 -y
conda activate eb-vla

# PyTorch — pick the build that matches your CUDA
pip install torch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0

pip install -e .
pip install -r requirements_vla_ebt.txt

For LIBERO evaluation you additionally need the LIBERO simulator, and for perturbation evaluation the LIBERO-Plus extension.

Training

The unified entry point is run.py. It takes a YAML config and a --mode:

python run.py --config configs/train/ebt_spatial.yaml --mode train

Any config value can be overridden from the command line:

python run.py --config configs/train/ebt_spatial.yaml --mode train \
    --batch_size=8 --max_steps=120000 --data.data_root_dir=/path/to/libero

Relevant training configs in configs/train/:

Ready-to-use launchers (SLURM and plain bash) live in run_scripts/, e.g. run_scripts/slurm/train_ebt_energy_scaled.sh or run_scripts/train_scripts/train_ebt_spatial.sh.

Evaluation

Standard LIBERO

python experiments/robot/libero/run_libero_eval.py \
    --pretrained_checkpoint=/path/to/checkpoint \
    --task_suite_name=libero_10 \
    --num_trials_per_task=10 \
    --ebt_mcmc_step_size=0.22 \
    --num_open_loop_steps=8 \
    --use_proprio=True \
    --num_images_in_input=2

LIBERO-Plus perturbations

Scripts that sweep over perturbation types (camera viewpoint, object layout, lighting, background texture, sensor noise, language, robot init state) and difficulty levels are in run_scripts/slurm/:

• eval_perturbed_single.sh — runs all (task × perturbation × difficulty) combinations sequentially
• eval_perturbed_array.sh — same, as a SLURM array job
• Per-perturbation scripts: eval_perturbed_{light,language,sensor_noise,...}.sh

Inference-time knobs

The final energy value returned by the head can be logged per rollout and used as a confidence score for out-of-distribution detection or failure prediction; see analysis utilities in scripts/:

• scripts/score_training_energy.py — score expert actions to calibrate baseline energy statistics
• scripts/plot_energy_landscape.py — visualize the learned energy surface around an action
• scripts/analyze_uncertainty.py — threshold sweeps for OOD detection

Acknowledgments

The VLM backbone and overall training pipeline build on Prismatic VLMs, OpenVLA-OFT, and VLA-Adapter. The energy-based transformer formulation follows the EBT line of work.

hpc_code