VBA: Vision-Brain-Action Model

Replace language with brain activity as the instruction signal for robot manipulation.

---

The Idea

Standard VLAs (Vision-Language-Action models) take a text instruction + camera image and output robot actions. We replace the text encoder with TRIBE v2's brain encoding. The robot is conditioned on a spatial brain activity pattern instead of a flat text string.

VLA (existing):    text string  +  camera image  →  robot action
VBA (this project): brain pattern +  camera image  →  robot action

The brain pattern is richer than text — when TRIBE v2 processes "pick up the red cup," the output encodes motor planning, visual expectations, spatial attention, and semantic understanding simultaneously across 20,484 cortical vertices.

---

Why This Might Actually Be Better Than Text

Text is a bottleneck. When you say "pick up the red cup," you compress an enormous amount of information into 5 words. The VLA's text encoder has to reconstruct all the implicit meaning:

• What does a red cup look like? (text doesn't say)
• Where is it likely to be? (text doesn't say)
• How should the hand approach? (text doesn't say)
• What grip to use? (text doesn't say)

A brain pattern from TRIBE v2 encodes ALL of that implicitly. The predicted brain response to "pick up the red cup" includes:

A single 20k-dimensional brain pattern packages all of these into one conditioning signal.

---

Architecture

Standard VLA (OpenVLA for reference)

"pick up the red cup" → text tokenizer → LLM backbone →
                                                          } → action tokens → 7-DoF
                          camera image → ViT encoder    →

VBA (our model)

"pick up the red cup" → TRIBE v2 → brain pattern (20484,) → BrainEncoder →
                                                                             } → action tokens → 7-DoF
                                           camera image → ViT encoder      →

And the video conditioning mode (no text at all)

video of human doing task → TRIBE v2 → brain pattern (20484,) → BrainEncoder →
                                                                                } → action tokens → 7-DoF
                                                camera image → ViT encoder    →

---

Detailed Architecture

BrainEncoder: (20484,) → (n_tokens, d_model)

See src/vba/model/brain_encoder.py for the full implementation. It:

1. Reshapes the 20,484-vertex brain pattern into 128 spatial patches of ~160 vertices.
2. Linearly embeds each patch and adds a CLS token + positional embeddings.
3. Runs a 6-layer Transformer encoder over the patch tokens.
4. Uses a Q-Former style cross-attention with 32 learned queries to compress into a fixed token count.
5. Projects to the VLA backbone's hidden dim (4096 for OpenVLA / LLaMA 7B).

Output shape: (batch, 32, 4096) — a drop-in replacement for text instruction tokens.

BrainEncoder params: ~45M (trainable). The VLA backbone (OpenVLA 7B) stays frozen.

---

Training Strategy

The key insight: you have paired (text, brain) data for free

TRIBE v2 gives you a deterministic mapping from text → brain pattern. So for every (text, image, action) triple in an existing VLA dataset, you can add the brain pattern. See scripts/02_enrich_dataset.py.

Three training phases

Phase 1 — Brain-text alignment pretraining (5 epochs). See scripts/03_train_alignment.py. Align BrainEncoder tokens with the VLA's own text encoder tokens using MSE. After this, brain tokens live in the same space as text tokens.

Phase 2 — End-to-end action fine-tuning (20 epochs). See scripts/04_train_vba.py. Train with action MSE loss. VLA backbone + ViT stay frozen; only BrainEncoder trains.

Phase 3 — Video-conditioned training (10 epochs). See scripts/05_train_video_cond.py. Replace text-derived brain patterns with video-derived brain patterns. The model learns to reproduce actions from watching a human do them, with the brain as the intermediate representation.

---

Dataset

Option A: Use an existing VLA dataset + enrich with TRIBE v2

For each episode, you already have (instruction_text, camera_images, actions). Run each instruction through TRIBE v2 once to get the brain pattern. Embarrassingly parallel.

Option B: Generate in Isaac Sim

See src/vba/data/isaac_collector.py for a skeleton Franka Panda episode collector.

Enrichment time estimate

TRIBE v2 forward pass on a short text: ~1–2 seconds on A100. 60k episodes (Bridge V2): ~17–33 GPU-hours. One overnight job.

---

Evaluation

All four experiments from the paper are run from scripts/06_evaluate.py and scripts/07_ablation.py:

1. Brain tokens vs text tokens — does the brain signal actually help?
2. Generalization to novel instructions — does brain embedding transfer better than text?
3. Video conditioning (no text) — video → brain → action, vs direct video → action.
4. Brain region ablation — which cortical areas matter for robot control?

---

Full Pipeline Demo

See scripts/08_demo.py for the end-to-end demo (Mode A: text instruction; Mode B: video demonstration, no text).

---

Project Structure

vba/
├── README.md
├── pyproject.toml
├── configs/
│   ├── brain_encoder.yaml       # BrainEncoder architecture
│   ├── training.yaml            # training schedule, losses
│   └── eval.yaml                # evaluation settings
├── scripts/
│   ├── 01_install_tribev2.sh    # setup script
│   ├── 02_enrich_dataset.py     # add brain patterns to VLA dataset
│   ├── 03_train_alignment.py    # Phase 1: brain-text alignment
│   ├── 04_train_vba.py          # Phase 2: end-to-end action training
│   ├── 05_train_video_cond.py   # Phase 3: video conditioning
│   ├── 06_evaluate.py           # all experiments
│   ├── 07_ablation.py           # brain region ablation study
│   └── 08_demo.py               # full pipeline demo
├── src/vba/
│   ├── model/
│   │   ├── brain_encoder.py     # BrainPatchEmbedding + BrainEncoder
│   │   ├── vba_model.py         # BrainEncoder + frozen OpenVLA backbone
│   │   └── losses.py            # alignment loss + action loss
│   ├── data/
│   │   ├── enriched_dataset.py  # VLA dataset + brain patterns
│   │   ├── brain_cache.py       # cache TRIBE v2 outputs to disk
│   │   └── isaac_collector.py   # collect episodes in Isaac Sim
│   └── viz/
│       ├── brain_render.py      # TRIBE v2 PlotBrain wrapper
│       └── demo_dashboard.py    # split-screen brain + robot viz
├── tests/
│   ├── test_brain_encoder.py    # shape tests
│   ├── test_enrichment.py       # verify brain patterns cached correctly
│   └── test_vba_forward.py      # end-to-end forward pass
└── notebooks/
    ├── 01_explore_brain_patterns.ipynb
    ├── 02_brain_vs_text_similarity.ipynb
    └── 03_ablation_analysis.ipynb

---

Hardware Requirements

Total VRAM at inference: TRIBE v2 (~12–16 GB) + OpenVLA 7B frozen (~14 GB) + BrainEncoder (~0.1 GB) + Isaac Sim (~2–4 GB) → fits on A100 40 GB / GH200.

---

Quickstart

# 1. Install
pip install -e ".[tribev2,vla,dev]"

# 2. Sanity check the BrainEncoder
python -c "import torch; from vba.model.brain_encoder import BrainEncoder; \
  m = BrainEncoder(); print(m(torch.randn(2, 20484)).shape)"
# -> torch.Size([2, 32, 4096])

# 3. Run tests
pytest tests/ -v

# 4. See each training script's CLI
python scripts/03_train_alignment.py --help
python scripts/04_train_vba.py --help