Large Language Model as a Policy Teacher for Training Reinforcement Learning Agents

TL;DR

This paper presents a framework where a large language model acts as a teacher to train a smaller, specialized reinforcement learning agent, resulting in more efficient learning and improved task performance in complex environments.

Contribution

The paper introduces a novel LLM-based teacher-student framework for training RL agents, reducing data requirements and surpassing teacher capabilities through environment feedback.

Findings

Student agents trained with LLM guidance outperform baseline methods.

The approach achieves higher efficiency in data usage for training RL agents.

Student agents surpass their teacher in task performance after further training.

Abstract

Recent studies have uncovered the potential of Large Language Models (LLMs) in addressing complex sequential decision-making tasks through the provision of high-level instructions. However, LLM-based agents lack specialization in tackling specific target problems, particularly in real-time dynamic environments. Additionally, deploying an LLM-based agent in practical scenarios can be both costly and time-consuming. On the other hand, reinforcement learning (RL) approaches train agents that specialize in the target task but often suffer from low sampling efficiency and high exploration costs. In this paper, we introduce a novel framework that addresses these challenges by training a smaller, specialized student RL agent using instructions from an LLM-based teacher agent. By incorporating the guidance from the teacher agent, the student agent can distill the prior knowledge of the LLM into its own model. Consequently, the student agent can be trained with significantly less data. Moreover, through further training with environment feedback, the student agent surpasses the capabilities of its teacher for completing the target task. We conducted experiments on challenging MiniGrid and Habitat environments, specifically designed for embodied AI research, to evaluate the effectiveness of our framework. The results clearly demonstrate that our approach achieves superior performance compared to strong baseline methods. Our code is available at https://github.com/ZJLAB-AMMI/LLM4Teach.