PaperlayerView exact revision on arXiv

Source-linked AI summary

Refine Actions If Needed: Adaptive Action Chunk Refinement for Robot Manipulation

Hanlin Liu, Zichen Zhou, Jiacheng Zhang, Liuyuan Na, Jin Jin, Guangyan Chen, Yufeng Yue, Gen Li, Fangqiang Ding

arXiv:2610.19765v1cs.RO

TL;DR

Executable proposals offer low-latency actions, while iterative generative policies require repeated evaluations and refinement needs vary by state. RAIN predicts candidate errors from offline demonstrations to route proposals before refinement, with reported simulation and real-robot gains and lower latency from dynamic routing.

  • Problem

    Lightweight policies produce executable actions quickly but can err, while iterative generation can improve success at the cost of repeated network evaluations and refinement needs vary across states.

    Source: §I · p. 1

  • Method

    RAIN trains an ensemble on expert-action deviations of unrefined and refined proposals, then uses calibrated pre-refinement risk predictions to choose whether and how much to refine.

    Source: §I · p. 2

  • Results

    Across eight ManiSkill tasks, RAIN achieved 66.6% success at 0.322 ms/action, and on real-robot tasks it improved success over ACT by 17.5 percentage points with comparable average successful completion time.

    Source: §I · p. 2

  • Takeaways & Limitations

    Matched controls show dynamic routing maintains the strongest observed fixed-refinement success level with lower latency, and its value lies in selective allocation rather than balanced refinement use or frequent switching.

    Source: F. How the Router Allocates Computation · p. 8, §V · p. 8

    Action-error targets may not fully reflect refinement’s effect on task success, especially outside the demonstration distribution.

    Source: §V · p. 8

Abstract

from arXiv · show

Computationally lightweight manipulation policies predict executable actions with low inference latency, but prediction errors can limit task success. Iterative generative policies can achieve higher task success rates, but repeated network evaluations increase inference latency. Refining action proposals from computationally lightweight policies can combine these strengths, but the amount of refinement needed varies across states. We introduce RAIN (Refine Actions If Needed), a framework that predicts action errors at different refinement levels to decide whether and how much to refine each executable proposal, before performing any refinement. We train an ensemble to predict these errors by comparing action proposals and their refined versions with expert actions from offline demonstrations, without requiring explicit labels of refinement needs. Calibration balances uncertainty, computation, and temporal consistency. We evaluate RAIN on eight ManiSkill tasks and four real-robot tasks. In simulation, our policy achieves 66.6% success, approaching the best evaluated diffusion-policy success rate with 5.71 times lower policy latency. On the real robot, success increases from 55.0% with Action Chunking with Transformers (ACT) to 72.5%, with comparable average successful completion time. Matched controls show that dynamic routing maintains the strongest observed fixed-refinement success level with 26.6% fewer refiner evaluations. Project page: https://rain-policy.github.io/.

I. INTRODUCTION

RAIN addresses the tension between fast executable action proposals and slower iterative generation by deciding, before refinement, whether a proposal merits correction. Its evaluation reports strong simulation and real-robot results alongside computation savings from dynamic routing.

  • The underlying trade-off is that iterative diffusion and flow-matching generation can improve task success, but repeated network evaluations increase inference latency.Source: §I · p. 1
  • The router must predict refinement outcomes before generating them, because online refinement probes consume the computation it aims to save and demonstrations lack direct labels for refinement needs.Source: §I · p. 2
  • RAIN uses offline comparative supervision and pre-refinement risk prediction to decide whether and how much to refine an executable proposal.Its risk ensemble learns from candidate deviations from expert actions, and offline calibration balances predicted risk, uncertainty, computation, and temporal consistency.Source: §I · p. 2
  • 66.6% success at 0.322 ms/action of amortized policy latency approaches the best evaluated diffusion-policy success rate, while real-robot success improves over ACT by 17.5 percentage points with comparable average successful completion time.Matched controls show dynamic routing maintains the strongest observed fixed-refinement success level while reducing policy latency by 15.7%.Source: §I · p. 2

II. RELATED WORK

Prior work reduces generative-policy cost through shorter generation, informative starting points, and adaptive computation. RAIN instead compares optional correction against direct execution of the same executable proposal, then routes before refinement.

II. RELATED WORK

  • Efficiency approaches shorten or reuse iterative generation, but distillation may trade action accuracy or diversity for speed, and temporal reuse may not match the current state.Source: §II · p. 2
  • Informative sources and adaptive-generation methods reduce the need to start from noise, but do not directly determine whether correction of an executable proposal justifies its cost.RAIN makes that comparison before the first refiner evaluation and retains direct execution as an option.Source: §II · pp. 2–3, Figure 2 · p. 3
  • Unlike correction-learning approaches that use environment interaction or human interventions, RAIN learns refinement allocation from offline demonstrations without online rewards or human corrections.Source: §II · p. 2
  • Other adaptive methods allocate generation steps or schedules, whereas RAIN compares correction amounts for the same proposal, including no correction, and predicts continuous candidate risks.Source: §II · p. 3

III. METHOD

  • RAIN starts from an executable proposal, compares candidate refinement levels with expert actions, trains an error-prediction ensemble, and executes only the selected refinement amount.Source: §III · p. 3

A. Executable Proposals and Optional Refinement

RAIN produces an executable action chunk and an uncertainty signal, then treats flow-matching refinement as optional. Joint proposal–refiner training retains direct action supervision while also encouraging proposals that the refiner can improve.

A. Executable Proposals and Optional Refinement

B. Joint Proposal–Refiner Training

  • Joint training retains the proposal’s supervised action-prediction objective while flow-matching updates both refiner and proposal, encouraging chunks that are directly executable and easier to refine.Both models are frozen afterward for router training.Source: B. Joint Proposal–Refiner Training · p. 4

C. Offline Candidate-Risk Supervision

Candidate-risk supervision compares alternative refinement amounts with the same aligned expert chunk, providing continuous error targets for learning whether correction helps.

C. Offline Candidate-Risk Supervision

  • Each refinement candidate, including direct execution, is compared with the same aligned expert action chunk to obtain a continuous target.Source: C. Offline Candidate-Risk Supervision · p. 4
  • The candidate error is a masked, dimension-normalized arm-action error with temporal weights, and a logarithmic transformation limits the influence of large errors while preserving candidate ordering.Source: C. Offline Candidate-Risk Supervision · p. 4
  • Expert-action deviations provide dense comparative supervision without branching the environment or collecting rollout rewards; smaller error means closer agreement with the demonstration.Source: C. Offline Candidate-Risk Supervision · p. 4

D. Pre-Refinement Risk Prediction

RAIN uses an ensemble to predict candidate action errors and their uncertainty before refinement. Its continuous supervision preserves differences among candidate refinements.

D. Pre-Refinement Risk Prediction

  • Each ensemble member predicts a mean and log variance for every candidate’s transformed error from observation features, the executable proposal, and log proposal uncertainty.Training combines heteroscedastic regression with a Huber term and auxiliary pairwise-ranking and best-candidate objectives.Source: D. Pre-Refinement Risk Prediction · p. 4
  • Continuous error targets retain differences among candidates instead of reducing supervision to a single preferred refinement label.Source: D. Pre-Refinement Risk Prediction · p. 4
  • The predicted mean estimates transformed candidate error, while the uncertainty combines predicted variance and ensemble disagreement.Proposal uncertainty is an input, not a separate refinement threshold.Source: D. Pre-Refinement Risk Prediction · p. 4

E. Offline Calibration and Online Routing

RAIN calibrates a candidate-scoring rule offline, then predicts scores before refinement at deployment. The router balances error, uncertainty, computation, and temporal consistency when choosing direct execution or refinement.

E. Offline Calibration and Online Routing

  • At deployment, candidate scores balance predicted error, uncertainty, refinement computation, and temporal consistency.A refined candidate is selected only when its combined score is preferable to direct execution; more refinement is not assumed to be better.Source: E. Offline Calibration and Online… · p. 4
  • A smoothed transition model reflects how the lowest-error refinement amount changes along demonstration trajectories, favoring plausible sequences without excluding candidates.Source: E. Offline Calibration and Online… · p. 4
  • Offline calibration selects coefficients by trajectory-averaged disagreement with the minimum-error candidate, then favors lower mean refinement among settings within one paired standard error of the best.The coefficients are frozen before closed-loop evaluation.Source: E. Offline Calibration and Online… · p. 4
  • Online, the ensemble scores candidates without running the refiner, and the policy executes the proposal directly or constructs only the selected refinement.Alternative candidates are generated for offline supervision, not online search.Source: E. Offline Calibration and Online… · p. 5

IV. EXPERIMENTS

Across eight ManiSkill tasks, RAIN achieves high success at low policy latency, while on four real-robot tasks it improves success over ACT with comparable successful completion time. Experiments compare multiple policy families and evaluate shared task data under specified protocols.

A. Experimental Setup

  • The study evaluates eight ManiSkill3 tasks and four real-robot tasks, using 500 successful demonstrations per simulated task and 50 training demonstrations per robot task.Simulation uses 100 shared environment seeds per task and real-robot evaluation uses 20 trials per task.Source: A. Experimental Setup · p. 5
  • Simulation comparisons include ACT, A2A, DP, FM, and ProbeFlow; matched Fixed 0–3, Direct-C, and RAIN conditions share checkpoints, preprocessing, seeds, and an eight-action execution block.Fixed k applies k refinements at each policy call, while RAIN selects the amount dynamically.Source: A. Experimental Setup · p. 5
  • Simulation policy latency measures isolated FP32 CUDA-Graph inference through action-chunk readiness, whereas real-robot completion time includes observation processing, inference, and physical execution.Source: A. Experimental Setup · p. 5

B. Overall Success Rate and Efficiency

C. Matched Routing Controls

Matched controls show that dynamic refinement reaches a slightly higher success rate with lower latency than the strongest fixed-refinement setting, while outperforming Fixed 2 and standalone ACT on success.

C. Matched Routing Controls

  • The matched protocol holds the proposal, refiner, preprocessing, seeds, execution horizon, and inference implementation constant, varying only FM refinement allocation.Source: C. Matched Routing Controls · p. 6
  • Fixed 0 outperforms standalone ACT, possibly because variance prediction helps the proposal policy learn task-relevant action variability.Source: C. Matched Routing Controls · pp. 6–7
  • 66.6% success at 0.322 ms/action for RAIN compares with 66.4% at 0.382 ms/action for Fixed 3, with 15.7% lower latency and 26.6% fewer FM evaluations per action.The matched-controls table covers success rate and inference cost across routing variants.Source: C. Matched Routing Controls · p. 7, Table III · p. 7
  • RAIN raises success from 65.8% with Fixed 2 to 66.6% while reducing latency from 0.330 to 0.322 ms/action, despite slightly more FM evaluations.Source: C. Matched Routing Controls · p. 7

D. Continuous Risk versus Direct Refinement Prediction

RAIN predicts action-error estimates for each refinement level rather than directly classifying the best level, and achieves a stronger success-oriented operating point than direct prediction.

D. Continuous Risk versus Direct Refinement Prediction

  • RAIN predicts a numerical action-error estimate for each refinement level, while Direct-C learns a hard label for the level with the lowest expert-action deviation.Source: D. Continuous Risk versus Direct… · p. 7
  • 66.6% success at 0.322 ms/action for RAIN improves success while reducing latency and refinement evaluations relative to Direct-C Argmax.Compared with calibrated Direct-C, RAIN gains 1.5 percentage points of success for 0.012 ms/action additional latency.Source: D. Continuous Risk versus Direct… · p. 7
  • Per-level error estimates preserve whether candidate refinement levels have similar or substantially different errors, letting the router weigh expected error reduction against computation.Source: D. Continuous Risk versus Direct… · p. 7

E. Effect of the Transition Prior

The transition prior modestly reduces refinement cost while maintaining comparable success, complementing current-state risk estimates with transition patterns learned from demonstrations.

E. Effect of the Transition Prior

  • 0.276 FM evaluations per action and 0.322 ms/action with the transition prior compare with 0.288 and 0.325 ms/action without it, while success remains comparable at 66.6% versus 66.4%.Source: E. Effect of the Transition Prior · p. 7
  • The prior favors plausible refinement sequences without excluding candidates, complementing current-state risk estimates with transition patterns learned from offline demonstrations.Source: E. Effect of the Transition Prior · p. 7

F. How the Router Allocates Computation

RAIN allocates refinement selectively across tasks, with different tasks using markedly different refinement amounts. The paper also reports auxiliary diffusion-refinement results and identifies a limitation in using expert-action error as a proxy for task success.

F. How the Router Allocates Computation

  • PushT and StackPyramid use the highest refinement amount in 99.6% and 97.4% of decisions, while PickCube relies almost entirely on zero or one refinement.PushCube, StackCube, and PegInsertionSide use broader mixtures; RAIN does not impose uniformly lower refinement across tasks.Source: F. How the Router Allocates Computation · p. 7

G. Refinement Beyond Flow Matching

  • In the separate ACT–diffusion setting, refinement raises success from 52.9% for the unrefined proposal to 63.6% at the best evaluated fixed-refinement setting.The paper's dynamic-routing results remain evaluated with the flow-matching refiner.Source: G. Refinement Beyond Flow Matching · p. 8

V. CONCLUSION AND LIMITATIONS

  • Across eight ManiSkill tasks and four real-robot tasks, dynamic routing achieves high success at low simulation latency and improves real-robot success with successful completion times comparable to ACT.Matched simulation controls also show slightly higher success and lower latency than the strongest fixed-refinement setting.Source: §V · p. 8
  • The router’s expert-action-error targets may not fully reflect refinement’s effect on task success, especially outside the demonstration distribution.Learning refinement decisions that better predict task outcomes remains an important direction.Source: §V · p. 8

Related papers

Loading 2610.19765v1…