Dynamic Test-Time Compute Scaling in Control Policy: Difficulty-Aware Stochastic Interpolant Policy

2025-11-25

AI Summary

Existing diffusion- and flow-based robotic control policies employ fixed inference step counts, ignoring task-difficulty heterogeneity—leading to computational redundancy for simple tasks and performance degradation for complex ones. To address this, we propose Difficulty-Aware Stochastic Interpolation (DA-SIP), the first framework enabling runtime dynamic allocation of computational resources for both diffusion and flow controllers, while supporting flexible training and inference. DA-SIP integrates a task difficulty classifier, ODE/SDE-adaptive numerical integration selection, and a fine-grained solver scheduling mechanism. Evaluated across diverse manipulation tasks, DA-SIP reduces total inference computation time by 2.6–4.4× on average, while maintaining success rates comparable to full-step baselines. This yields substantial improvements in computational efficiency and task adaptability without compromising control performance.

Technology Category

• Intelligent Robots: Behavior Learning & Control
• Humans and AI: Human-Aware Planning and Behavior Prediction
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Responsible Web: Machine-in-the-loop, human agency and autonomy
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs

Abstract

Diffusion- and flow-based policies deliver state-of-the-art performance on long-horizon robotic manipulation and imitation learning tasks. However, these controllers employ a fixed inference budget at every control step, regardless of task complexity, leading to computational inefficiency for simple subtasks while potentially underperforming on challenging ones. To address these issues, we introduce Difficulty-Aware Stochastic Interpolant Policy (DA-SIP), a framework that enables robotic controllers to adaptively adjust their integration horizon in real time based on task difficulty. Our approach employs a difficulty classifier that analyzes observations to dynamically select the step budget, the optimal solver variant, and ODE/SDE integration at each control cycle. DA-SIP builds upon the stochastic interpolant formulation to provide a unified framework that unlocks diverse training and inference configurations for diffusion- and flow-based policies. Through comprehensive benchmarks across diverse manipulation tasks, DA-SIP achieves 2.6-4.4x reduction in total computation time while maintaining task success rates comparable to fixed maximum-computation baselines. By implementing adaptive computation within this framework, DA-SIP transforms generative robot controllers into efficient, task-aware systems that intelligently allocate inference resources where they provide the greatest benefit.

Problem

• Adaptive computation scaling for robotic controllers based on task difficulty
• Dynamic adjustment of integration horizon and solver selection in real-time
• Reducing computational costs while maintaining performance across manipulation tasks

Innovation

• Adaptive integration horizon based on task difficulty
• Difficulty classifier dynamically selects computation parameters
• Unified stochastic interpolant framework for generative policies