🤖 AI Summary
This work addresses the high cost and suboptimality of acquiring high-quality demonstration data in imitation learning, particularly the lack of scalable data sources for goal-conditioned control tasks. To overcome this challenge, the authors propose an efficient data generation and augmentation framework that leverages trajectory optimization to automatically produce thousands of near-optimal trajectories within minutes on a standard laptop. By relabeling intermediate states along these trajectories as new goals, the training dataset is expanded by an order of magnitude. A lightweight goal-conditioned policy trained on this augmented dataset—containing fewer than 80,000 parameters—achieves near-optimal performance and high success rates across multiple tasks. Moreover, its inference speed exceeds that of the trajectory optimization solver by over 6,000×, substantially improving generalization and enabling practical deployment on embedded systems.
📝 Abstract
Imitation learning is a well-established approach for machine-learning-based control. However, its applicability depends on having access to demonstrations, which are often expensive to collect and/or suboptimal for solving the task. In this work, we present GCImOpt, an approach to learn efficient goal-conditioned policies by training on datasets generated by trajectory optimization. Our approach for dataset generation is computationally efficient, can generate thousands of optimal trajectories in minutes on a laptop computer, and produces high-quality demonstrations. Further, by means of a data augmentation scheme that treats intermediate states as goals, we are able to increase the training dataset size by an order of magnitude. Using our generated datasets, we train goal-conditioned neural network policies that can control the system towards arbitrary goals. To demonstrate the generality of our approach, we generate datasets and then train policies for various control tasks, namely cart-pole stabilization, planar and three-dimensional quadcopter stabilization, and point reaching using a 6-DoF robot arm. We show that our trained policies can achieve high success rates and near-optimal control profiles, all while being small (less than 80,000 neural network parameters) and fast enough (up to more than 6,000 times faster than a trajectory optimization solver) that they could be deployed onboard resource-constrained controllers. We provide videos, code, datasets and pre-trained policies under a free software license; see our project website https://jongoiko.github.io/gcimopt/.