Prepare before You Act: Learning from Humans to Rearrange Initial States
Yinlong Dai, Andre Keyser, Dylan Losey
AI summary
Problem
Standard imitation learning policies fail on out-of-distribution initial states like occluded or misplaced objects, typically requiring massive datasets to generalize. Humans naturally overcome this by rearranging the environment to simplify the task before execution.
Approach
ReSET uses a scoring network to decide if rearrangement is needed, a flow generation network to predict human-like object movements from videos, and a reduction policy to execute these moves before the main task policy runs.
Key results
- Theoretical proof that restructuring states reduces the generalization gap upper bound
- Flow-based reduction policy trained on action-agnostic human videos and teleoperation data
- Outperforms diffusion policies and VLAs in few-shot out-of-distribution manipulation tasks
- Achieves robust task execution with equal or less total training data
Why it matters
Enables robots to learn complex manipulation tasks efficiently from limited demonstrations, bridging the gap between controlled training and real-world cluttered environments.
Abstract
Imitation learning (IL) has proven effective across a wide range of manipulation tasks. However, IL policies often struggle when faced with out-of-distribution observations; for instance, when the target object is in a previously unseen position or occluded by other objects. In these cases, extensive demonstrations are needed for current IL methods to reach robust and generalizable behaviors. But when humans are faced with these sorts of atypical initial states, we often rearrange the environment for more favorable task execution. For example, a person might rotate a coffee cup so that it is easier to grasp the handle, or push a box out of the way so they can directly grasp their target object. In this work we seek to equip robot learners with the same capability: enabling robots to prepare the environment before executing their given policy. We propose ReSET, an algorithm that takes initial states — which are outside the policy’s distribution — and autonomously modifies object poses so that the restructured scene is similar to training data. Theoretically, we show that this two step process (rearranging the environment before rolling out the given policy) reduces the generalization gap. Practically, our ReSET algorithm combines action-agnostic human videos with task-agnostic teleoperation data to i) decide when to modify the scene, ii) predict what simplifying actions a human would take, and iii) map those predictions into robot action primitives. Comparisons with diffusion policies, VLAs, and other baselines show that using ReSET to prepare the environment enables more robust task execution with equal amounts of total training data. See videos at our anonymous website: https://reset2025paper.github.io