PhD Dissertation Defense by LI Wenjun | Diversity-Augmented Training for Generalizable AI Agents

Diversity-Augmented Training for Generalizable AI Agents

ABOUT THE TALK

Reinforcement Learning (RL) has achieved remarkable success in both simulated and real-world tasks. However, a key challenge remains: RL agents often struggle to generalize to novel or unseen scenarios. For instance, a robot trained on flat terrain may fail on uneven ground. This dissertation addresses this limitation by exploring how to improve generalization through diversity-augmented training. 

The work investigates this problem across three major settings: infinite horizon RL, finite horizon RL, and large language model (LLM) agents. Infinite horizon training emphasizes long-term improvement through open-ended learning across a wide range of environments. Finite horizon training, on the other hand, prioritizes learning efficiency by enabling generalization from a limited set of training scenarios. The final setting extends the diversity-based approach to LLMs, aiming to improve their performance in structured tasks such as planning. 

To support generalization in infinite horizon RL, this dissertation introduces novel diversity metrics that guide environment generation within the Unsupervised Environment Design (UED) framework. These metrics promote the creation of training environments that are both challenging and behaviorally diverse. For finite horizon RL, a meta-learning framework is proposed that enables a teacher agent to efficiently design concise, tailored training sequences for each student agent. In the LLM context, a sampling method is developed to select diverse training examples based on data structure rather than language similarity, improving both generalization and sample efficiency in planning tasks.

Overall, this dissertation offers practical algorithms and insights for training generalizable AI agents by explicitly incorporating diversity into the learning process—bridging the gap between training performance and real-world adaptability.

SPEAKER BIOGRAPHY

Wenjun LI is a PhD candidate in Computer Science at the School of Computing and Information Systems (SCIS), supervised by Prof. Pradeep VARAKANTHAM. His research focuses on training generalizable AI agents, large language models, and AI for education.