# Visual Imitation Learning of Non-Prehensile Manipulation Tasks with Dynamics-Supervised Models

arXiv: https://arxiv.org/abs/2410.19379 · HTML: https://mlsanity.com/paper/2410.19379/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Embodied and robot world models, one of 11 directions in World models.
- Builds on 17 of its 21 references that are in the map.
- First posted on arXiv 25 Oct 2024.

## Placed in

- World models / [Embodied and robot world models](https://mlsanity.com/wm/embodied-and-robot-world-models/)
- Offline RL and imitation learning from data / [Robot manipulation imitation learning](https://mlsanity.com/offlinerl/robot-manipulation-imitation-learning/)

## Builds on (most cited first, 17 of 17)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [RT-1: Robotics Transformer for Real-World Control at Scale](https://mlsanity.com/paper/2212.06817/) (2022)
- [Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware](https://mlsanity.com/paper/2304.13705/) (2023)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [End-to-End Training of Deep Visuomotor Policies](https://mlsanity.com/paper/1504.00702/) (2015)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [GAIA-1: A Generative World Model for Autonomous Driving](https://mlsanity.com/paper/2309.17080/) (2023)
- [Universal Manipulation Interface: In-The-Wild Robot Teaching Without In-The-Wild Robots](https://mlsanity.com/paper/2402.10329/) (2024)
- [FlowFormer: A Transformer Architecture for Optical Flow](https://mlsanity.com/paper/2203.16194/) (2022)
- [Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images](https://mlsanity.com/paper/1506.07365/) (2015)
- [Visual Imitation Made Easy](https://mlsanity.com/paper/2008.04899/) (2020)
- [DITTO: Offline Imitation Learning with World Models](https://mlsanity.com/paper/2302.03086/) (2023)
- [Deep Predictive Learning: Motion Learning Concept inspired by Cognitive Robotics](https://mlsanity.com/paper/2306.14714/) (2023)
- [Keep it Upright: Model Predictive Control for Nonprehensile Object Transportation with Obstacle Avoidance on a Mobile Manipulator](https://mlsanity.com/paper/2305.17484/) (2023)
- [Relational Object-Centric Actor-Critic](https://mlsanity.com/paper/2310.17178/) (2023)
- [Force Map: Learning to Predict Contact Force Distribution from Vision](https://mlsanity.com/paper/2304.05803/) (2023)

## Built on by (most cited first, 1 of 1)

- [A Flexible Field-Based Policy Learning Framework for Diverse Robotic Systems and Sensors](https://mlsanity.com/paper/2512.19148/) (2025)

## Abstract (arXiv)

Unlike quasi-static robotic manipulation tasks like pick-and-place, dynamic tasks such as non-prehensile manipulation pose greater challenges, especially for vision-based control. Successful control requires the extraction of features relevant to the target task. In visual imitation learning settings, these features can be learnt by backpropagating the policy loss through the vision backbone. Yet, this approach tends to learn task-specific features with limited generalizability. Alternatively, learning world models can realize more generalizable vision backbones. Utilizing the learnt features, task-specific policies are subsequently trained. Commonly, these models are trained solely to predict the next RGB state from the current state and action taken. But only-RGB prediction might not fully-capture the task-relevant dynamics. In this work, we hypothesize that direct supervision of target dynamic states (Dynamics Mapping) can learn better dynamics-informed world models. Beside the next RGB reconstruction, the world model is also trained to directly predict position, velocity, and acceleration of environment rigid bodies. To verify our hypothesis, we designed a non-prehensile 2D environment tailored to two tasks: "Balance-Reaching" and "Bin-Dropping". When trained on the first task, dynamics mapping enhanced the task performance under different training configurations (Decoupled, Joint, End-to-End) and policy architectures (Feedforward, Recurrent). Notably, its most significant impact was for world model pretraining boosting the success rate from 21% to 85%. Although frozen dynamics-informed world models could generalize well to a task with in-domain dynamics, but poorly to a one with out-of-domain dynamics.