Learning Quality-Controllable Robot Policies with Iterative Hindsight Relabeling

AI summary

Overview

Research area: Robot learning — imitation learning / behavior cloning, quality- and style-conditioned robot policies, and online data collection for policy improvement (cs.RO).

Technical level: Advanced. The paper assumes familiarity with behavior cloning, flow matching policies, vision-language-action (VLA) models, hindsight relabeling, and offline/online RL baselines.

Scope: The paper introduces Quality-Controllable Policy (QCP), a robot policy that is conditioned on quantized execution-quality labels so a user can steer speed, smoothness, or stability at inference time, and pairs it with an iterative hindsight relabeling loop that pushes the controllable quality range beyond the initial demonstrations, validated in Robomimic, RoboEval, and two real-world Franka tasks.

What This Paper Is About

Robot policies are typically trained with behavior cloning on successful demonstrations, which treats every success as equally desirable and therefore reproduces whatever mixture of fast, slow, smooth, or jerky behaviors happens to be in the data. In deployment, however, success alone is not enough: a conveyor robot should prioritize speed, while a robot carrying wine glasses must prioritize stability. The paper's goal is to learn a single policy whose execution quality can be dialed up or down at inference time along one or several quality axes, and then to expand the achievable quality range beyond what the initial demonstrations contain by repeatedly collecting rollouts, relabeling them by their measured quality, and retraining.

Key Contributions

1. Quality-Controllable Policy (QCP). A policy conditioned on discrete quality labels obtained by measuring each demonstration with a rule-based score function and quantizing scores into equal-frequency bins. Conditioning is done with a learned embedding for from-scratch flow matching policies, or by inserting the label as text (e.g., "quality level: 3") into the language prompt for VLA models such as π0.5.
2. Iterative hindsight relabeling. A loop that collects rollouts conditioned on the highest quality bin, keeps only successful ones, relabels them by the quality they actually achieved (not the quality requested), recomputes bin boundaries over the accumulated dataset including the original demonstrations, and retrains the policy from scratch. No manual quality annotation is required because scores are computed automatically.
3. Empirical demonstration of simultaneous quality and success gains. On Robomimic and RoboEval, QCP improves both target quality and task success across iterations, whereas quality-filtered BC, BC on accumulated rollouts, SORL, and SCIQL tend to improve one at the expense of the other.
4. Multi-attribute control and real-world validation. A single policy conditioned on up to six quality labels jointly improves six attributes on RoboEval, and bidirectional control of speed plus lift amplitude is shown on Square. Real-robot experiments on conveyor throughput and wine-glass transport show practical gains.

Main Findings

• Quality and success improve together. Across tasks and quality attributes, QCP improves the target quality metric and the task success rate over iterations, while baselines trade one for the other: BC_acc has no quality signal and improves success alone, BC_filter and SORL discard most successful data and lose success, and SCIQL's advantage weighting creates a trade-off.
• QCP reaches higher final quality than flow-matching RL baselines. At iteration 5 (after five rounds of rollout collection, averaged over three seeds), QCP attains better quality than DSRL and EXPO even when their quality rewards are weighted highly. On Square+Speed, QCP reports 82.2% success at 120.3 steps, versus EXPO's 96.0% success at 162.5 steps and DSRL's 37.0% at 208.8 steps — EXPO's success rate is higher, but its execution is slower. The authors attribute the RL baselines' weakness to those methods being tuned to consume a success signal rather than a new quality axis.
• The controllable range extends past the demonstrations. The Wasserstein-1 distance between initial demonstrations and final successful rollouts, normalized by the interquartile range of initial demonstration scores, exceeds 1.5 for smoothness and joint path length. For speed, the rollout distribution shifts toward shorter completion times, beyond the demonstrated range; the fastest rollouts appear close to the maximum physically achievable speed in this setup.
• The quality knob works even before any relabeling. At iteration 0 on Square, requesting the highest bin yields 193.7 mean steps to success versus 308.1 for the lowest bin, with intermediate bins in between. This ordering holds for every attribute at every iteration, even though new data is only ever collected at the highest bin, and the whole range shifts upward over iterations.
• The gain comes from policy updates, not extra sampling. Spending the same rollout budget as 1,800 rollouts from the policy frozen at iteration 0 does not reproduce the shift: the fixed policy reproduces its iteration-0 distribution, and none of its successful rollouts reaches the median smoothness of the final iteration.
• Quantized labels beat continuous scores. With four quantile bins, roughly the best quarter of trajectories share the best label, giving the policy multiple examples of the requested behavior. On Square, quantized labels execute faster than a fixed continuous target at every iteration, and four bins reach higher task success from iteration 1 onward; more bins buy a small speed gain at a growing cost in success.
• Multi-attribute control works jointly. Four combinations of faster/slower execution and smaller/larger lift amplitude on Square move both attributes in the requested directions, with final quality close to policies trained for each attribute alone. On RoboEval, a single policy conditioned on six labels improves most attributes over initial BC and BC_acc.
• Failed rollouts hurt. The authors tried including failed rollouts as additional low-quality data and found it degraded performance.
• Real-world wine serving. Only 50 of 382 teleoperated attempts (13.1%) succeeded during data collection. Conditioning on the highest stability bin reduces mean tilt angle from 0.0660 rad for BC to 0.0482 rad (a 27.0% reduction) and improves success from 30.7% to 42.0% over iterations; BC_filter improves through iteration 1 but degrades at iteration 2.
• Real-world conveyor work. With 20 cans on a belt moving one can every three seconds, QCP raises mean end-effector speed from 188.6 to 257.6 mm/s (36.6%) and throughput from 8.2 to 13.0 cans per minute (59.3%) relative to BC. BC_filter reaches 217.7 mm/s at iteration 2 and its throughput drops from 8.3 to 6.3 cans per minute, below BC. The final QCP policy exceeds the 173.5 mm/s of human teleoperation.

Methodology in Plain English

Scoring quality. Each quality attribute is measured on a complete trajectory with a simple rule-based function. The simplest example is 1/T for execution speed; others include normalized action differences for smoothness, total joint path length for efficiency, end-effector tilt for stability, and lift amplitude as the difference between maximum and minimum end-effector height. Where lower values are better, the score is inverted so that the highest bin always means the preferred behavior.

Turning scores into labels. Scores are quantized into B equal-frequency bins using the empirical cumulative distribution over the current dataset, so label j means the score lies between the j/B and (j+1)/B quantiles. The experiments use B = 4 bins and request the highest bin, q* = 3.

Training. The policy sees a history of L = 2 observations plus the quality label and predicts a chunk of H future actions. It is a flow matching policy: a velocity field is trained by regression to transport Gaussian noise into the action chunk, and actions are generated at inference by integrating that field from u = 0 to u = 1 (nine Euler steps). For VLA models, the quality label is simply written into the language prompt and the model is fine-tuned with its original recipe, so no architecture change is needed.

The loop. Train on demonstrations; run the policy asking for the highest bin; keep successful rollouts; measure their actual quality and relabel them accordingly; recompute bin boundaries over the enlarged dataset (which still includes the original demonstrations); retrain from scratch; repeat. Two mechanisms are offered for why rollouts can exceed the best demonstrations: the label is attached to every action chunk, so the policy can recombine well-executed segments from different demonstrations, and sampling is stochastic, so conditioning on the top bin raises the chance of landing above the demonstrated quality. Folding those rollouts back in moves the top bin upward.

Why retrain instead of fine-tune. Relabeling changes what each label refers to, so the previous label-to-behavior mapping would have to be unlearned. The authors also note that reinitialization may help avoid the loss of plasticity and primacy bias seen in deep RL, though they do not study this directly.

Setup. Simulation policies are MLPs (three layers of width 512, 128-dimensional flow-time embedding) trained with AdamW for 100k steps on Robomimic and 30k on RoboEval at batch size 1024; real-world experiments fine-tune π0.5 for 10 epochs at batch size 32. Each simulation iteration collects 300 rollouts per task and seed (1,800 total over six iterations at a 400-step limit); real-world iterations collect as many rollouts as there are initial demonstrations (400 for conveyor, 50 for wine glasses). Results are averaged over three seeds. RL baselines were trained for 720k environment steps to match the rollout budget.

Why This Matters

Impact on research. The paper reframes imitation learning's binary success signal as a controllable, multi-dimensional property of execution, and shows that a mechanism analogous to reward-quantile relabeling in language models (Quark) transfers to robot execution quality. It offers a supervised alternative to reward-weighted online RL for flow matching policies, and it shows that the commonly assumed trade-off between behavioral style and task performance can be avoided rather than merely balanced.

Real-world applications:

• Logistics and warehouse automation: conveyor picking, where throughput depends jointly on grasp success and cycle speed, as in the paper's 20-can belt task.
• Service and food robotics: carrying trays, glasses, or liquids, where transport stability (tilt angle) directly determines whether the task succeeds.
• Manufacturing and assembly: tasks where motion smoothness and jerk matter for part handling, wear, or precision, and where a single policy must serve multiple stations with different quality requirements.
• Human-robot coexistence and hardware longevity: smoother, requestable execution profiles for robots operating near people or expensive mechanisms.

Industry relevance. A single deployed policy can be retuned per task by changing a label rather than retraining or maintaining separate checkpoints, and the labels are computed automatically, so the improvement loop requires no human quality annotation. The reported throughput and success gains on real hardware, including surpassing human teleoperation speed on the conveyor task, are the kind of metric that maps directly onto operational cost.

Future Directions

• Plasticity and reinitialization. The authors explicitly leave unstudied whether retraining from scratch each iteration mitigates the loss of plasticity and primacy bias documented in deep RL; this is an open empirical question for longer training loops.
• Richer quality definitions. All quality axes here come from hand-designed rule-based score functions; learning scores or rewards for qualities that are hard to measure (e.g., perceived naturalness or care) is a natural extension, and the paper contrasts its automatic labels with the human comparisons required by Freeform Preference Learning.
• What to do with failures. Including failed rollouts as low-quality data degraded performance, so how to exploit failure data in this relabeling framework remains unresolved.
• Label granularity and scale. The four-bin quantile scheme was chosen empirically; how bin count, number of jointly conditioned attributes, and iteration length interact at larger scale and over longer deployments is not settled.

Target Audience

Researchers and engineers working on robot imitation learning, flow matching and VLA policies, and style- or quality-conditioned control, who want a method that improves execution quality without sacrificing task success. It is also relevant to practitioners deploying manipulation policies in logistics, service, or manufacturing settings who need runtime control over speed, smoothness, or stability, and to readers interested in the transfer of reward-quantile relabeling ideas from language models to robotics. The paper is written at an advanced level and assumes comfort with behavior cloning, flow matching, and standard offline/online RL baselines.

Authors’ abstract

Robot policy evaluation often relies on a binary success signal for task completion. In real-world deployment, however, desired execution behaviors are also defined by quality attributes such as execution speed and motion smoothness, which need to be dynamically controlled according to task requirements. To address this, we introduce Quality-Controllable Policy (QCP), a robot policy that enables inference-time control over both individual quality attributes and their combinations. To extend the controllable quality range beyond initial demonstrations, we further propose an iterative hindsight relabeling loop: QCP collects online rollouts under target quality labels, adds successful trajectories to the dataset, relabels the accumulated trajectories in hindsight according to their measured quality, and fine-tunes the policy. Experiments on Robomimic and RoboEval demonstrate that QCP consistently improves both task execution quality and success rates across iterations, producing successful rollouts that exceed the quality of the initial demonstrations. Finally, we validate QCP on real-world manipulation tasks, where controlling speed increases throughput in conveyor operations and controlling transport stability improves success when carrying wine glasses, demonstrating that quality-controllable policies yield practical gains in physical deployment.