Data Agent: Learning to Select Data via End-to-End Dynamic Optimization

Dynamic Data selection aims to accelerate training by prioritizing informative samples during online training. However, existing methods typically rely on task-specific handcrafted metrics or static/snapshot-based criteria to estimate sample importance, limiting scalability across learning paradigms and making it difficult to capture the evolving utility of data throughout training. To address this challenge, we propose Data Agent, an end-to-end dynamic data selection framework that formulates data selection as a training-aware sequential decision-making problem. The agent learns a sample-wise selection policy that co-evolves with model optimization, guided by a composite reward that integrates loss-based difficulty and confidence-based uncertainty signals. The reward signals capture complementary objectives of optimization impact and information gain, together with a tuning-free adaptive weighting mechanism that balances these signals over training. Extensive experiments across a wide range of datasets and architectures demonstrate that Data Agent consistently accelerates training while preserving or improving performance, e.g., reducing costs by over 50\% on ImageNet-1k and MMLU with lossless performance. Moreover, its dataset-agnostic formulation and modular reward make it plug-and-play across tasks and scenarios, e.g., robustness to noisy datasets, highlighting its potential in real-world scenarios. Code is available at https://github.com/Jackbrocp/Data-Agent.

Training modern AI models often requires processing very large amounts of data, which can be expensive and time-consuming. However, not every training example is equally useful throughout the whole training process. Some examples may be more helpful at early stages, while others become more important later. In this paper, we introduce Data Agent, a method that learns how to dynamically choose useful training data while a model is being trained. Instead of relying on a fixed human-designed rule, Data Agent observes the training process and adapts its data selection strategy over time. We show that Data Agent can reduce training cost while maintaining, and sometimes improving, model performance. It works across different types of AI tasks, including image recognition, object detection, semantic segmentation, and language model fine-tuning. By making training more efficient, Data Agent can help researchers and practitioners build strong AI models with less computation, lower cost, and better scalability.