Improving Pre-Trained Vision-Language-Action Policies with Model-Based Search

TL;DR

VLAPS enhances pre-trained vision-language-action policies for robotics by integrating model-based search, significantly improving task success rates and safety in out-of-distribution scenarios through a novel planning framework.

Contribution

Introduces VLAPS, a framework that embeds model-based search into VLA policy inference, enabling more robust and efficient robotic task execution.

Findings

VLAPS outperforms VLA-only baselines by up to 67 percentage points.

The method enables better exploration of large, language-conditioned search spaces.

VLAPS improves safety and performance in zero-shot robotic tasks.

Abstract

Pre-trained vision-language-action (VLA) models offer a promising foundation for generalist robot policies, but often produce brittle behaviors or unsafe failures when deployed zero-shot in out-of-distribution scenarios. We present Vision-Language-Action Planning & Search (VLAPS) -- a novel framework and accompanying algorithms that embed model-based search into the inference procedure of pre-trained VLA policies to improve their performance on robotic tasks. Specifically, our method biases a modified Monte Carlo Tree Search (MCTS) algorithm -- run using a model of the target environment -- using action priors defined by the VLA policy. By using VLA-derived abstractions and priors in model-based search, VLAPS efficiently explores language-conditioned robotics tasks whose search spaces would otherwise be intractably large. Conversely, by integrating model-based search with the VLA policy's inference procedure, VLAPS yields behaviors that are more performant than those obtained by directly following the VLA policy's action predictions. VLAPS offers a principled framework to: i) control test-time compute in VLA models, ii) leverage a priori knowledge of the robotic environment, and iii) integrate established planning and reinforcement learning techniques into the VLA inference process. Across all experiments, VLAPS significantly outperforms VLA-only baselines on language-specified tasks that would otherwise be intractable for uninformed search algorithms, increasing success rates by as much as 67 percentage points.