ELASTIC: Efficiently Learning to Adaptively Scale Test-Time Compute for Generative Control Policies

2026-06-30

AI Summary

This work addresses the challenge of dynamically allocating test-time computation for generative policies in robotic control, where the optimal trade-off between sequential refinement and parallel exploration is highly dependent on the current state, task, and policy. To this end, we propose ELASTIC, an algorithm that—without access to the generative policy’s training data—learns a state-dependent computation scheduling strategy via a meta-Markov decision process, adaptively adjusting the number of sequential denoising steps and parallel samples at each inference step. ELASTIC unifies the joint optimization of sequential and parallel computation through a reinforcement learning–based meta-policy operating over a frozen, pre-trained generative model (e.g., a diffusion model). Experiments demonstrate that ELASTIC achieves Pareto superiority over fixed and single-axis scaling baselines in simulation, and on a real robot using the π₀.₅ model, it attains success rates comparable to best-of-10 sampling with 34% lower latency.

Abstract

Generative control policies (GCPs), such as diffusion policies and flow-based vision-language-action models, enable test-time scaling in robot control. Test-time compute can be allocated along two axes: sequential scaling, which increases denoising steps to refine actions, and parallel scaling, which samples multiple candidate actions to search across modes of the policy distribution. However, the optimal allocation of sequential and parallel compute is hard to know a priori as it is state-, task-, and policy-dependent. For example, early stages of a grasp may benefit from broader parallel exploration, while near-contact phases may require more sequential refinement for precision. We present ELASTIC, an algorithm that learns state-dependent test-time compute schedules for GCPs. We formulate compute allocation as a meta-Markov Decision Process in which a meta-policy interacts with a frozen pretrained robot policy and selects sequential steps and parallel samples at each denoising iteration to maximize task success while minimizing compute. Using reinforcement learning, this meta-policy also learns adaptive compute schedules without access to the GCP's training data. Across simulated manipulation benchmarks with diffusion policies, ELASTIC Pareto-dominates fixed and single-axis scaling baselines at matched compute budgets. On real-world robot manipulation with the $π_{0.5}$ vision-language-action model, ELASTIC matches best-of-$10$ success while reducing wall-clock latency by 34%.

Problem

• generative control policies
• test-time compute scaling
• sequential vs parallel scaling
• state-dependent allocation
• robot control

Innovation

• test-time compute scaling
• generative control policies
• adaptive resource allocation
• meta-reinforcement learning
• robot manipulation