Imitation learning is great, but needs us to have (near) optimal data. We throw away most other data (failures, evaluation data, suboptimal data, undirected play data), even though this data can be really useful and way cheaper! In our new work - RISE, we show a simple way to... *use all of this non-optimal data to robustify imitation learning* with minimal requirements beyond BC. Key idea: use non-expert data to learn how to *recover* back to expert data with a minimal frills offline RL that works under sparse data coverage. Allows usage of *all* available data, not just expert data - never throw your data away! Paper: Website: A 🧵(1/10)show more

To start - we know that BC can be fragile. Particularly, we know that BC struggles when placed in OOD configurations with sparse data coverage. Now we can of course collect more data - but this requires expensive high-quality demonstrations. But non-expert data can be plentiful - play data, failed attempts during teleoperation, or evaluation rollouts of existing policies. But it’s not clear how to actually use this data with BC. (2/10)

Our idea in Robust Imitation by Stitching from Experts (RISE) is simple - while you can’t just imitate non-optimal data, you can use it to learn how to *get back to expert states* using offline RL. It’s a particularly simple offline RL problem - no reward models needed; label expert data with a reward of 1, and non-expert data with 0, allowing dynamic programming to stitch together non-expert data with expert data. Importantly, it’s dead simple - throw in expert data with reward 1, non-expert data with reward 0 into the buffer and do offline RL with your favorite offline RL method. (3/10)

Simple enough right? - not so different from what @sidgreddy and @svlevine proposed in SQIL. Now, the issue is that plain old offline RL approaches often fail to stitch together disjoint trajectories under relatively sparse data coverage (especially from vision). So while there is a promise of using suboptimal data for recovery, it doesn’t work quite as well we’d want it to! (4/10)

Looking into it a bit deeper, Q functions are able to stitch reasonably well, but we find that the marginal action distribution captured by the policy is overly conservative, with little action coverage. This lack of coverage in the policy distribution makes it hard to actually make use of the stitching captured by the Q function. Simply put - the algorithm can't find the right actions to take because the policy distribution doesn't cover it! (5/10)

The fix is easy! - allow the policy to “borrow” actions from nearby states. In doing so, policies can widen their action distributions in controlled ways, significantly improving their ability to stitch suboptimal data! This can be implemented in simple ways - e.g enforcing Lipschitz smoothness on the policy, or through simple data augmentation. Think of it like “fuzzing” nearby states to make stitching easier even in low data coverage settings. But the impact on policy recovery and success is huge! (6/10)

Let’s see how well this works! First, using RISE allows your imitation learning policies to be a *lot* more robust, using non-optimal data to recover from OOD scenarios. On the flipside, just imitation non-optimal data or standard offline RL is far less effective (as we see quantitatively) (7/10)

Secondly, we find that RISE can make use of failed or suboptimal demonstrations, significantly improving learned policies over standard imitation learning. Moreover this works across different tasks, including with deformables! (8/10)

Interestingly, this also extends to policy evaluations. We can iteratively keep reincorporating policy evalution rollouts into RISE to keep improving. Evaluation data is not wasted! (9/10)

What should you take away: 1) Use all your data, just throw it in the buffer with a 0. Don’t waste any of your failed data! 2) even suboptimal data can teach you how to recover back to experts 3) stitching can be hard in low coverage settings, some smoothness assumptions can really help! Fun diving deep into the guts of offline RL and learning a bunch about stitching behavior and recovery. Led by @kevin_huang8, Rosario Scalise, Cleah Winston along with Ayush Agrawal, @yunchuzh, @RohanBaijal, Markus Grotz, Byron Boots, @Ben_Burchfiel, @MashaItkina, Paarth Shah Curious to hear thoughts/feedback! Paper: Website: (10/10)

really like using "auxiliary" data in general, since BC is quite restrictive on data needs. does the nearest-neighbor trick work for images? i wonder if u can generalize this to images by training a BC policy on some coarsened version of state (like a diffusion policy ;) )

Yes works on images! We didn’t actually do nearest neighbors explicitly, but instead we did data augmentation with DINO features and we enforced Lipschitz smoothness with spectral norm. So all results are from images :)

Abhishek, you're spot on! It's a waste to discard all that extra data; maybe we're missing out on some real gems, no?

Imitation learning is a dead end. World models are AGI. Learn everything you can about them

🤩

RISE's approach to leveraging non-expert data is brilliant! Solving sparse coverage while maintaining task robustness - this could transform robotics pipelines 🛠️

Interesting; RISE is a good mnemonic for imitation learning work.