Improving the Data-efficiency of Reinforcement Learning by Warm-starting with LLM

TL;DR

The algorithm, LORO, can both converge to an optimal policy and have a high sample efficiency thanks to the LLM's good starting policy, and empirically demonstrate that LORO outperforms baseline algorithms such as pure LLM-based policies, pure RL, and a naive combination of the two.

Abstract

We investigate the usage of Large Language Model (LLM) in collecting high-quality data to warm-start Reinforcement Learning (RL) algorithms for learning in some classical Markov Decision Process (MDP) environments. In this work, we focus on using LLM to generate an off-policy dataset that sufficiently covers state-actions visited by optimal policies, then later using an RL algorithm to explore the environment and improve the policy suggested by the LLM. Our algorithm, LORO, can both converge to an optimal policy and have a high sample efficiency thanks to the LLM's good starting policy. On multiple OpenAI Gym environments, such as CartPole and Pendulum, we empirically demonstrate that LORO outperforms baseline algorithms such as pure LLM-based policies, pure RL, and a naive combination of the two, achieving up to $4 \times$ the cumulative rewards of the pure RL baseline.

Figures (27)

• Figure 1: The LLM Off-policy pre-train, RL On-policy (LORO) algorithm. Image inspired by levine2020offline.
• Figure 2: Our algorithm, LORO, outperforms the LLM policies (Qwen 7B, Qwen 32B) and the vanilla On-Policy RL baselines. We demonstrate the effectiveness of LORO in six environments, reducing the data required to learn the optimal policy from two to ten times the vanilla On-Policy baseline. LORO and the On-policy baseline learn the optimal policy in the first four environments. Even when not converged to the optimal solution, LORO outperforms other baselines in the last two environments. LLM's base policies can perform very well in the Pendulum experiment, maybe because of the extra information that we provide for this setting, described in Appendix \ref{['sec:appendix_env_impl']}. All results are shown with standard error over five random seeds. In the CliffWalking experiment, some baselines are not shown in the figure since their episode rewards are too small. Similarly, multiple baselines overlap at -200 on the MountainCar experiment.
• Figure 3: Comparing pre-training (then removing the collected data) versus mixing the LLM's collected data with on-policy data without pre-training. It's clear that pre-training is necessary for LORO to achieve superior performance compared to naively mixing the data.
• Figure 4: Comparing pre-training with LLM's data versus random and on-policy data. The main finding here is that pre-training is only useful with LLM's data.
• Figure 5: Comparing the effect of different amounts of pre-training data for the CartPole environment. This is a rare edge case where using LORO is worse than the On-policy baseline.