SpeechVLA-WM

Speech-Conditioned Vision-Language-Action models with Future-State World Models — built on top of VLANeXt. Course project for CS 444 (UIUC, Spring 2026) by Aniket Tathe and Amitabh Swain.

We extend VLANeXt in two orthogonal directions:

1. Speech instead of text as the instruction modality (Matcha-TTS → frozen Whisper-small.en → learned projector → Qwen3-VL-2B).
2. World-model auxiliary losses that force the VLM's hidden states to encode future world state — a FutureProprioHead predicting future joint+gripper state by MSE, plus a FutureFeaturesHead predicting future image latents.

All experiments are on LIBERO-spatial.

---

Headline results

Stage-3 with naive training destroys the policy (60 % → 35 %). Our diagnosis (per-term loss logging) traced this to an aggressive learning-rate restart on top of a fresh AdamW state, not the auxiliary losses themselves. With (a) selective per-parameter optimizer-state load on resume and (b) a 10× lower restart LR, Stage-3 (Exp-11) recovers to 56 % — statistically tied with Stage-2.

Same task ✕ same trial seed across three configurations. The matching 5-second video grid is in figures/rollout_compare_grid.mp4.

---

What didn't work — the world-model image generator

VLANeXt's image-generator head, trained jointly with the rest of Stage-3, learned the marginal VQ-token distribution of LIBERO-spatial scenes but not actually how to predict future frames. The 9 → 3 nat training-CE drop is entropy reduction over the unigram distribution, not predictive learning. We verified by directly rendering the head's autoregressive predictions:

All 6 ground-truth-vs-prediction pairs are textured noise on the prediction side. This drove our Plan B FutureFeaturesHead, which predicts the mean Emu3.5 codebook embedding by MSE instead of the next image's discrete tokens by CE — a much smoother target. It converged trivially (loss 0.17 → 0.0002 in 50 steps) and didn't disrupt the policy, but the target turned out to be too easy: mean-pooled codebook embeddings are nearly invariant across LIBERO-spatial frames, so the head learned a global constant. Real future work is somewhere between the two extremes — e.g. per-quadrant spatial pooling, or PCA components of frame-level latents.

---

Training curves

Cosine similarity to Qwen text embeddings hits 0.998 by step 150 in Stage 1; Stage-2 mini-batch loss reaches 0.062 at its best checkpoint; Stage-3 Exp-7's total loss looks like a healthy 12.57 → 2.98 descent but is misleading — the action loss component spikes from 0.13 to 1.3 within 200 steps of warmup. Without per-term logging this is invisible.

---

Repository layout

SpeechVLA-WM/
├── README.md                ← this file
├── figures/                 ← PNG + MP4 result figures
└── work/                    ← helper scripts (data prep, training launches, plotting, eval)
│   ├── generate_speech.py
│   ├── cache_whisper_embeddings.py
│   ├── exp{5,7,9,10,11}_*.sbatch
│   ├── rollout_eval_*.sbatch
│   ├── world_model_sample_frames.py
│   ├── plot_*.py, build_rollout_compare_video.py
│   └── install_env.sh, finish_env.sh, download_libero.sh
└── vlanext/                 ← our edits / new modules on top of VLANeXt
    ├── scripts/
    │   ├── train.py                ← per-term loss CSV + selective AdamW state load
    │   └── train_speech_align.py   ← Stage-1 alignment trainer (ours)
    ├── src/
    │   ├── models/
    │   │   ├── VLANeXt.py              ← speech threading, world-model heads wiring
    │   │   ├── speech_encoder.py       ← SpeechProjector (ours, new)
    │   │   ├── proprio_generator.py    ← FutureProprioHead (ours, new)
    │   │   └── future_features_head.py ← FutureFeaturesHead, Plan B (ours, new)
    │   ├── datasets/libero_act.py      ← emits future_proprio + future_image
    │   └── evaluation/libero_bench/
    │       └── VLANeXt_utils.py        ← speech cache, Stage-3 model build at eval
    └── config/                          ← all YAML configs (Stage-1 / 2 / 3 + eval)

This repo contains only the code we wrote or modified. The upstream VLANeXt codebase (Emu3.5 tokenizer, action head, third-party LIBERO benchmark, etc.) is not redistributed here — to actually run this you would git clone https://github.com/DravenALG/VLANeXt and drop our files in over it.

---

Reproducing a result

Before running anything, set these environment variables (the scripts read them):

export PROJECT_ROOT=/path/to/your/clone     # where this repo + the VLANeXt clone live
export CONDA_ENV_PATH=/path/to/your/cs444   # full path or name of the conda env
export HF_HOME=$HOME/.cache/huggingface     # HuggingFace cache (Qwen, Whisper, Matcha, Emu3.5)

And edit each sbatch's #SBATCH --account=YOUR_SLURM_ACCOUNT to your SLURM allocation.

Then any single sbatch reproduces one stage/experiment:

# Stage 1 — projector alignment
sbatch work/exp5_stage2_long.sbatch                       # (also pulls Stage 1 weights from disk)

# Stage 2 — speech-conditioned policy, 20 k steps  →  60.0 % SR
sbatch work/exp5_stage2_long.sbatch

# Stage 3 — naive both-heads training (the disruption demo)  →  35.0 % SR
sbatch work/exp7_stage3_wm.sbatch

# Stage 3 — Plan B with selective optimizer state load and lower LR  →  56.0 % SR
sbatch work/exp11_stage3_features_lowlr_h200.sbatch

# Closed-loop rollouts
sbatch work/rollout_eval_speech.sbatch                    # against Stage-2 checkpoint
sbatch work/rollout_eval_stage3_h200.sbatch               # against Exp-7 checkpoint
sbatch work/rollout_eval_exp11_h200.sbatch                # against Exp-11 checkpoint

# Rendering world-model predicted future frames (the negative-result figure)
sbatch work/world_model_frames_h200.sbatch

# Figures (run locally, after rollouts finish)
python work/plot_training_curves.py
python work/plot_rollout_compare.py
python work/plot_rollout_compare_grid.py
python work/build_rollout_compare_video.py

---

Two infrastructure changes worth pulling out

These are inside vlanext/scripts/train.py:

1. Per-term loss logging. The training CSV now writes one column per loss component (action_loss, dct_loss, image_loss, proprio_loss, features_loss) instead of just the summed loss. The Stage-3 policy disruption is invisible at the total-loss level (12.57 → 2.98 looks healthy); it's a single-glance read in the per-term CSV (action_loss jumps from 0.13 to 1.3 the instant LR ramps up).

2. Selective per-parameter optimizer-state load on resume. When you add new heads on top of a checkpoint, optimizer.load_state_dict(...) raises because param_groups no longer match. The original train.py fell back to a fresh optimizer state for all parameters, which destroys the accumulated AdamW momentum/variance the prior stage spent thousands of steps building. Our fix walks model.parameters() in order, looks up the corresponding saved state by index, and assigns it to optimizer.state[p] per tensor; only the new heads get fresh state. Combined with a 10× lower restart LR, this is what made Stage-3 train without disrupting the policy.

Anyone training multi-stage VLAs with progressively added auxiliary heads should consider both.

---

Citation / further reading

• VLANeXt — the codebase we built on top of.
• LIBERO — manipulation benchmark.
• Matcha-TTS, Whisper, Qwen3-VL, Emu3.5 — see the upstream model cards.

For the full writeup (8 500 words, 7 figures, per-task tables, discussion), see the course final report (uploaded separately).