SITCOM: Scaling Inference-Time COMpute for VLAs

TL;DR

SITCOM targets the fragility of vision-language-action robotics in long-horizon tasks by introducing inference-time MPC-style planning that augments pretrained VLAs with multi-step rollouts and reward-based trajectory selection. A transformer-based dynamics model, pretrained on BridgeData V2 and fine-tuned on SIMPLER, enables efficient $l$-step rollouts from $n$ candidate actions, while a learned reward model ranks trajectories to select the best plan for real execution; a DAgger-inspired adaptation reduces distribution drift during planning. Key contributions include the two-stage dynamics model training, VLA finetuning with 100 expert trajectories, and comprehensive evaluation showing dramatic improvements in task success (from around $48\%$ baseline to $72\%$ with SITCOM) in SIMPLER tasks. This approach bridges the Real2Sim gap and demonstrates scalable, inference-time planning that can transform one-shot VLAs into robust long-horizon robotic planners with practical compute considerations.

Abstract

Learning robust robotic control policies remains a major challenge due to the high cost of collecting labeled data, limited generalization to unseen environments, and difficulties in planning over long horizons. While Vision-Language-Action (VLA) models offer a promising solution by grounding natural language instructions into single-step control commands, they often lack mechanisms for lookahead and struggle with compounding errors in dynamic tasks. In this project, we introduce Scaling Inference-Time COMpute for VLAs (SITCOM), a framework that augments any pretrained VLA with model-based rollouts and reward-based trajectory selection, inspired by Model Predictive Control algorithm. SITCOM leverages a learned dynamics model to simulate multi-step action rollouts to select the best candidate plan for real-world execution, transforming one-shot VLAs into robust long-horizon planners. We develop an efficient transformer-based dynamics model trained on large-scale BridgeV2 data and fine-tuned on SIMPLER environments to bridge the Real2Sim gap, and score candidate rollouts using rewards from simulator. Through comprehensive evaluation across multiple tasks and settings in the SIMPLER environment, we demonstrate that SITCOM when combined with a good reward function can significantly improve task completion rate from 48% to 72% using trained dynamics model.

Figures (8)

• Figure 1: Method Diagram explaining the overall SITCOM architecture. Given an initial frame and goal in text description, a VLA predicts actions which is fed into the dynamics model to get the next frame. This process is then iteratively repeated with the output of dynamics model to generate action sequence rollouts. Finally, a reward model is used to rank the action sequences and pick the best one among them to execute in the real-world.
• Figure 2: Dynamics Model Architecture
• Figure 3: SIMPLER environment and four different tasks with WidowX arm
• Figure 4: Scaling with number of rollouts for SITCOM-EnvSim and SITCOM-Dynamics models.
• Figure 5: Scaling with varying rollout lengths. The performance of the dynamics model decreases with the increase in rollout length, explaining the drop in overall performance for some of the tasks.