Trajectory-Level Data Augmentation for Offline Reinforcement Learning
Abstract
We propose a data augmentation method for offline reinforcement learning, motivated by active positioning problems. Particularly, our approach enables the training of off-policy models from a limited number of suboptimal trajectories. We introduce a trajectory-based augmentation technique that exploits task structure and the geometric relationship between rewards, value functions, and mathematical properties of logging policies. During data collection, our augmentation supports suboptimal logging policies, leading to higher data quality and improved offline reinforcement learning performance. We provide theoretical justification for these strategies and validate them empirically across positioning tasks of varying dimensionality and under partial observability.
Lay Summary
Teaching robots and machines through trial-and-error in the real world is often expensive or unsafe. To avoid these risks, AI systems can be trained using pre-recorded data, a method known as offline reinforcement learning. However, in many industrial settings, this data is collected by rigid, pre-programmed routines that take slow, zig-zagging, and suboptimal paths to reach their goal. However, the performance of the AI system is fundamentally limited by the quality of the demonstration data. To tackle this, we developed a method called LIFT that automatically identifies "shortcuts" within the recorded data. By mathematically analyzing the geometry of the machine's movements, our method acts like a smart editor, cutting out unnecessary detours to create a more direct, synthetic path. We then train on these refined paths, effectively teaching AI systems to perform better than the original routine that collected the data. Our approach allows AI systems to learn highly efficient, direct movements from flawed, real-world data without needing costly additional trials. This makes it easier, faster, and safer to automate high-precision tasks, such as aligning camera lenses or controlling robotic arms in manufacturing.