VQ-WM

In submission

VQ-WM

Vector-Quantized World Model:
Latent Feedback Motion Planning for Robot Manipulation

Xuezhou Xu*1, Chung-Pang Wang*1, Zhenzhen Li2, Benjamin Riviere1

1New York University  ·  2NVIDIA

*Equal contribution

Paper Code — coming soon

Visual world models have shown strong potential for robot planning, yet long-horizon manipulation remains challenging: prediction errors compound over autoregressive rollouts, search complexity grows with the horizon, and planners propose latent targets that the low-level controller cannot reliably reach. We introduce VQ-WM, a latent feedback motion planning framework that vector-quantizes the latent space by controller reachability, measured as first-passage time, so a discrete planner can search far ahead over reachable states while a continuous controller tracks each planned state with visual feedback.

Overview of VQ-WM: initial and goal images are encoded into latent states and vector-quantized by controller reachability; the planner searches over vector-quantized states q_init, q_1, q_goal, and a feedback controller drives the robot through each region of attraction R(q) toward z_goal.
Fig. 1: Overview of VQ-WM. The initial (green) and goal (red) images are encoded into latent states and vector-quantized according to controller reachability, so that each vector-quantized state serves as a target the controller is expected to reach. The planner performs discrete search over these vector-quantized states using a discrete latent predictor. A continuous controller then drives the robot toward each planned target in turn, using visual feedback to correct deviations along the way.

Real-robot Tower of Hanoi

100%task progress, 15 moves, play data only

0 50 100 100% VQ-WM 27% V-JEPA2-AC* 5% CosmosPolicy 3% π0.5 0% V-JEPA2-AC

Simulated manipulation

75.3%success on Scene, the longest task

0 50 100 75.3% VQ-WM 47.3% TS 44.7% PLDM 42.7% LeWM 40.7% H-JEPA

Walker locomotion

97.3%success rate, images + proprioception

0 50 100 97.3% VQ-WM 22% VQ-WMimages only 12.7% DINO-WMimages only 12% DINO-WM

Method

Discrete search over controller-reachable states, controlled with visual feedback

VQ-WM components and pipeline. (a) Inference: planning over vector-quantized states with the discrete predictor P_Q, a fast controller pi acting every step, and vector quantization phi_Z of the current and goal latents. (b) Training: controller, continuous latent predictor, state codebook construction via greedy set cover, and vector-quantized latent dynamics.
Fig. 2: VQ-WM components and pipeline. (a) Inference encodes the initial and goal images, searches over state codes, and executes the planned targets with visual feedback. (b) Offline training learns the controller and continuous latent predictor, performs state vector quantization using first-passage time, and then learns vector-quantized latent dynamics.

Planning slow

Breadth-first search over vector-quantized states with the discrete predictor PQ returns the shortest path over vector-quantized states from qinit to qgoal. One call plans about five minutes of robot execution in under a second.

Control fast

A goal-conditioned controller π(· | zt, qi) tracks each planned vector-quantized state every step from visual feedback, for the number of calls the first-passage-time estimate assigns to that transition.

First-passage time d(z, z′): the number of controller calls until a closed-loop rollout of π in the learned latent model first reaches z′. It replaces Euclidean distance for building the codebook and assigning states, so every vector-quantized state is a target the controller can actually reach.

Training

  1. Controller π

    Goal-conditioned, trained on state–goal pairs up to R steps apart with an L1 action loss.

  2. Latent predictor PZ

    Supervised on closed-loop rollouts with π, the same procedure used to estimate first-passage time.

  3. State codebook Q

    Greedy set cover over estimated regions of attraction; each latent state is assigned to the vector-quantized state it reaches in the fewest controller calls.

  4. Vector-quantized latent dynamics PQ

    Latent action encoder, action codebook and discrete latent predictor learned jointly from vector-quantized transitions with cross-entropy and VQ losses.

Simulation

Long-horizon manipulation and locomotion in Stable-WorldModel

The margin over the strongest baseline grows with the task horizon.

TABLE I: Long-horizon Manipulation success rates.
MethodCube-EasyCube-HardScene
LeWM74.0 ± 2.832.0 ± 6.142.7 ± 1.9
TS68.0 ± 2.810.7 ± 0.947.3 ± 2.5
PLDM62.0 ± 4.311.3 ± 3.844.7 ± 4.1
H-JEPA76.7 ± 2.560.0 ± 5.040.7 ± 2.5
VQ-WM (ours)84.7 ± 2.573.3 ± 0.975.3 ± 5.0

Proprioception improves the controller; the codebook adapts; the planner is untouched.

TABLE II: Success rate (%) on Walker with different observation inputs.
MethodImagesImages + proprio.
DINO-WM12.7 ± 3.412.0 ± 4.9
VQ-WM (ours)22.0 ± 5.997.3 ± 2.5

Real World

Four-disk Tower of Hanoi on a WidowX AI arm, from short play data

Four photographs of a WidowX AI arm solving the four-disk Tower of Hanoi, from the initial stack on the left rod to the goal stack on the right rod.
Fig. 3: Real-world control of VQ-WM on the four-disk Tower of Hanoi puzzle.
Bar chart of average task progress for Cosmos Policy, pi 0.5, V-JEPA 2-AC, V-JEPA 2-AC with manual subgoals, and VQ-WM across expert data H1, expert data H15, and non-expert data. VQ-WM reaches 100, 99.6 and 100 percent.
Fig. 4: Real-world results. VLA and WAM baselines without runtime search fail to extrapolate from short play subtasks to full-puzzle solutions, whereas VQ-WM composes local behaviors into complete plans and executes them reliably. V-JEPA 2-AC struggles even on H1, suggesting a limited effective planning horizon. V-JEPA 2-AC* receives manually specified subgoals that supply the solution sequence, leaving the model to perform local planning and control.

Ablation

First-passage time vs. Euclidean distance for state quantization

Simulated Tower of Hanoi, same encoder, controller and training; only the quantization metric changes.

  • First-passage time peaks near 90% success; Euclidean distance peaks at 40%.
  • High success holds across codebook sizes from about 250 to 2000, an eightfold range.
  • Euclidean quantization ignores reachability and switches targets on latent proximity, so plans stall or skip ahead.
Success rate versus codebook size First-passage-time quantization reaches 85 to 89 percent success between 427 and 1288 vector-quantized states and drops at both extremes; Euclidean quantization peaks at 40 percent around 2037 states. Three seeds, 50 trajectories each. 0 25 50 75 100 200 500 1000 2000 5000 codebook size |Q| success rate (%) |Q| = 468: 4.7% ± 1.9 (3 seeds) |Q| = 778: 7.3% ± 0.9 (3 seeds) |Q| = 922: 6.0% ± 0.0 (3 seeds) |Q| = 1348: 16.7% ± 3.8 (3 seeds) |Q| = 2037: 40.0% ± 3.3 (3 seeds) |Q| = 3129: 34.7% ± 9.0 (3 seeds) |Q| = 189: 48.0% ± 4.9 (3 seeds) |Q| = 427: 85.3% ± 0.9 (3 seeds) |Q| = 673: 89.3% ± 2.5 (3 seeds) |Q| = 1288: 82.0% ± 2.8 (3 seeds) |Q| = 3132: 68.7% ± 3.4 (3 seeds) |Q| = 4467: 30.0% ± 5.7 (3 seeds) First-passage timeEuclidean
Fig. 5: State vector quantization with different codebook sizes.

Citation

BibTeX

@misc{xu2026vectorquantizedworldmodellatent,
      title={Vector-Quantized World Model: Latent Feedback Motion Planning for Robot Manipulation}, 
      author={Xuezhou Xu and Chung-Pang Wang and Zhenzhen Li and Benjamin Riviere},
      year={2026},
      eprint={2610.25600},
      archivePrefix={arXiv},
      primaryClass={cs.RO},
      url={https://arxiv.org/abs/2610.25600}, 
}