Improving Sample Efficiency of Reinforcement Learning with Background Knowledge from Large Language Models

We consider sequential decision-making problems in an open environment $\mathcal{E}$, in which a specific task $\mathcal{T}$ corresponds to a reinforcement learning (RL) problem. An agent interacts with the environment in discrete time to realize specific requirements defined by the task. Typically, an RL task can be modeled as a Markov decision process (MDP)  [1] $\mathcal{T}=\left(\mathcal{S},\mathcal{A},r,P,\rho,\gamma\right)$. An agent starting at initial state $s_{0}\sim\rho(\cdot)$ selects action $a_{t}\in\mathcal{A}$ and leads to a transition to $s_{t+1}\sim P(\cdot\mid s_{t},a_{t})$ with acquiring reward $r_{t}=r(s_{t},a_{t})$. An RL algorithm usually learns a policy $\pi(a\mid s)$ by maximizing the discounted cumulative reward, a.k.a. return, $R=\sum_{t=0}^{\infty}\gamma^{t}r_{t}$. Since an agent often perceives a view of the open environment, the received input in our experiments may be a partial observation. With a little misuse of notation, we do not explicitly distinguish states and observations in our methodology for simplicity since we do not delve into the theoretical properties of partial observability. A practical approach is to use the trajectory of agent history $\tau_{t}=(s_{0},a_{0},\dots,s_{t-1},a_{t-1},s_{t})$ to represent agent state. Our work focuses on training policies $\pi(a\mid s)$ in any task $\mathcal{T}$ from an environment, where the task variance commonly comes from different targets and maps in our experiments.

To tackle the problem of poor sample efficiency in RL, potential-based reward shaping [26, 45] is a useful technique for altering RL processes. Typically, this technique adds an additional reward term $F(s,s^{\prime})$ to the original environmental reward $r$, whose value is computed with a defined potential function $\phi(s)$:

where $s$ and $s^{\prime}$ are the current state and the next state respectively and $\gamma$ is the discount factor of the MDP. An advantage of applying potential-based reward shaping is that it can maintain the policy optimality of original problems [26]. Consider an MDP $\mathcal{T}$ where $r_{t}$ is the extrinsic reward received from the environment at the time step $t$. We can denote $G=\sum_{t=0}^{\infty}r_{t}$ as the return of an episode in the original MDP. The concept of reward shaping refers to adding additional shaped reward $F(s,s^{\prime})$ to the extrinsic reward $r$:

Note that reward shaping is not to modify the reward function of the environment but to supplement additional rewards for computation. The purpose of the shaping function $F$ is often to provide heuristic domain knowledge to the problem when the agent transitions from state $s$ to $s^{\prime}$. We thus define the return of an episode by

Potential-based reward shaping is a particular type of reward shaping. For potential-based reward shaping, the shaping function is in the form of

where $\gamma$ is the exact discount factor from the MDP. We call $\phi$ the potential function and $\phi(s)$ the potential of state $s$. Therefore, we can define the potential function $\phi:\mathcal{S}\rightarrow\mathbb{R}$ instead of defining $F:\mathcal{S}\times\mathcal{S}\rightarrow\mathbb{R}$. The reason for choosing this form of reward shaping is that it can converge to the optimal policy. Consider the return for one episode:

where we can decompose $G^{\prime}$ into the cumulative reward on the original MDP and the potential of the initial state $\phi(s_{0})$. As the initial state $s_{0}$ can be any arbitrary state, we can easily extend this equality to a shaped Q-function, for example, $Q^{\prime}(s,a)=Q(s,a)-\phi(s)$. Therefore, any RL algorithm that maximizes the cumulative reward or the Q-values will derive the same optimal policies after reward shaping, since the term $\phi(s)$ is not related to action selection.

The optimality invariance of reward shaping is found in early research [26, 45] and widely applied in recent RL works [14, 16]. We note that although potential-based reward shaping can provide guarantees of the final results, it may alter the optimization dynamics in a way that either accelerates or slows down policy learning. A well-designed potential function $\phi(s)$ can decrease the time to convergence while bad ones may increase the time. In our work, we propose three approaches to design potential functions that can reflect the background knowledge of the environment from LLM feedback. We prove that our designs of the potential function can significantly improve sample efficiency, leading to a shorter convergence time.

To ground LLMs for decision-making tasks, it is conventional to transform trajectories into text captions. Though this limitation may be addressed by further improvement on multimodal LLMs, we posit the existence of text captions in our pre-collected data for querying LLMs. Typically, the description of a trajectory can be derived from state captions and action names.

In this section, we describe our framework that harnesses LLMs to obtain background knowledge of an environment, and further use it for accelerating downstream RL tasks. As there exist various ways to guide RL with LLM feedback, we summarize the desiderata of our framework below to better illustrate our requirements:

• •

  Interaction-free. The process of acquiring background knowledge from LLMs should not be attached to an online RL process, which will involve a considerable amount of interactions with the environment.

• •

  Task-agnostic. The background knowledge acquired in our framework should represent a general understanding of the environment, rather than involving behaviors about specific tasks.

• •

  Optimality-invariant. We expect to not change policy optimality after injecting background knowledge into RL. We provide background knowledge for improving sample efficiency.

Based on these three principles, we design a framework to extract background knowledge for reward shaping as previously shown in Figure 1. To satisfy the desiderata, we highlight three key procedures in the framework:

Data collection. To follow the interaction-free property, we pre-collect a dataset $\mathcal{D}=\{(s,a)\}$ by directly deploying an RL algorithm in an open environment without specifying its task goal. Therefore, the dataset only contains states and actions to delineate possible interactions in the environment. Specifically, we adopt RND [12] and save its interactive data periodically, following the standard in offline RL literature [46]. As an intrinsically motivated RL approach, RND tends to explore novel states in the open environment. Therefore, the collected data can contain diverse and useful behaviors for concluding background knowledge.

Background knowledge representation. To ground LLMs for the environment, we iteratively sample trajectories from the pre-collected data and ask LLMs for their feedback. During this procedure, the LLM makes assessments based on captions of agent trajectories but keeps unknown about the task goal, producing task-agnostic background knowledge. The knowledge is then represented as a potential function $\phi(s)$ by gathering LLM feedback over data samples, where we defer different variants for background knowledge representation to the next section.

RL with reward shaping. Using the potential function $\phi(s)$, we can adopt potential-based reward shaping for downstream RL tasks according to Equation (1), which shapes RL processes with preserving policy optimality. We adopt the PPO algorithm [47] based on the implementation from CleanRL [48] to train policies in different tasks.

The above three phases achieve our proposed desiderata, composing a common framework to extract and reuse background knowledge for downstream RL tasks. We note that the implementation of data collection and the RL process can be realized by different approaches. Specifically, we use the PPO algorithm [47] with the CleanRL [48] implementation, a stable and succinct PPO version for adopting our reward-shaping technique. The only modification to the training process is to add the auxiliary reward to the original environmental reward. To adapt to our selected environments, we use convolutional neural networks to extract features from the input observation. Then the features are fed into different multi-layer perceptrons (MLPs) to compute the action logits or value functions. Due to the different observation spaces of our evaluated environments, the feature processing networks are slightly different. We show the network structure of RL agents and also report the hyperparameters of the algorithms in our appendix, where most of the hyperparameters are directly taken from the CleanRL implementation despite some coefficients are related to environment sampling and auxiliary rewards.

With pretraining in code data, LLMs exhibit powerful abilities in programming and reasoning [49, 50, 51], propelling successful applications to write code for guiding RL [24, 25, 23]. Most of these works propose to prompt an LLM to write code of reward functions to depict desired targets. As coding is a direct approach to conveying LLM knowledge, we thus introduce an iterative prompting procedure to harness LLMs to program and improve their written code. We design straightforward prompts to ask LLMs to provide a Python function whose arguments come from the current observation and agent histories. By providing agent trajectories, LLMs can analyze agent behaviors based on more thorough considerations. Our prompt contains three portions:

• •

  Environmental information. We attach the environment description and necessary constants and variables of the environment, serving as a header of the code.

• •

  Trajectory samples. We sample trajectory data from the pre-collected data and caption them to better ground LLMs for the environment.

• •

  Feedback from samples. We use the current code to compute results from sampled data and append the results as an evaluation.

To be more specific, we prompt LLMs to write code for a function def compute_reward(obs, past_obs, past_actions). Although we use the computation results as the potential values $\phi(s)$, we name the function name compute_reward with a little misuse, since it is better for LLMs to understand. The parameters of the function include the current observation, historical observations, and historical actions, where we truncate the length of the historical sequence to $5$ for efficiency consideration. The LLM is supposed to write a function that returns the expected value of the potential function and also a dictionary that describes different portions of the result. This additional dictionary gives the LLM diverse feedback and thus helps improve the code. This helpful information will also be attached to the next query. There may be situations where the code provided by the LLM cannot pass the Python interpreter, especially for weak LLMs. In this case, we attach the error log reported by the interpreter and ask the LLM to fix potential issues. The LLM will have multiple retry chances to write a runnable code. Otherwise, we will stop the iteration and use the latest successful code as the final result. For each iteration, we sample trajectories from the dataset and feed them into the written code for the output values and the information dictionary. Afterward, we attach the text caption of the trajectories along with the dictionary information and ask the LLM to improve the code. We repeat this process for $20$ rounds, which is a moderate number since we find that involving additional iterations may result in too much unnecessary or inaccurate code.

As LLMs show the potential to be a helpful evaluation tool [52, 53], many prior works utilize the feedback of LLMs to align learned models [54, 55, 56]. We note that the preference over trajectories can also be a useful metric to delineate background knowledge and thus develop an approach that learns the potential function from LLM feedback. We sample two trajectories $\tau^{0},\tau^{1}$ from pre-collected data and ask LLMs to provide preferences over the given pair. The sampled trajectory sequences are truncated to $H$ steps $\tau=(s_{1},a_{1},\dots,s_{H},a_{H})$ for fair comparisons. The preference label can be $y=0$ for $\tau^{0}>\tau^{1}$ or $y=1$ for $\tau^{1}>\tau^{0}$, resulting in a preference dataset $\mathcal{D}_{\text{pref}}=\{(\tau^{0},\tau^{1},y)\}$. Following previous preference-based RL studies [57, 58], we can define a preference predictor following the Bradley-Terry model [59]:

where $\phi(s)$ is the potential function. Therefore, we can learn $\phi(s)$ by maximizing the likelihood of Equation (6) in our annotated dataset $\mathcal{D}_{\text{pref}}$ in the following objective:

where $\theta$ is the parameters of $\phi$. Following [60], we adopt the Transformer structure to compute $\phi(s)$ with historical information. For each preference label, we sample two trajectories $\tau^{0},\tau^{1}$ from pre-collected data and ask LLMs to provide preferences over the given pair. The trajectory sequence is truncated to a length of $H=5$ and is then captioned to compose the input prompt. We annotate labels including $y=0$ for $\tau^{0}>\tau^{1}$ and $y=1$ for $\tau^{1}>\tau^{0}$, forming a dataset $\mathcal{D}$ of labeled data. In addition, we also find that LLMs may refuse to provide a rank in the process because they equally preferred the provided two trajectories. To utilize this portion of data, we add an additional log-likelihood term to the original loss on the unlabeled data $\mathcal{D}^{\prime}$:

We mix $\mathcal{L}(\theta)$ and $\mathcal{L}^{\prime}(\theta)$ when optimizing the preference predictor. To parameterize the potential function $\phi(s)$, it is conventional to use a neural network with the input of agent state $s$. However, we find that it can be more helpful to make $\phi$ a non-Markovian function according to the previous work Preference Transformer [60] since the previous observations and actions contain useful information to judge the potential of the current state. Therefore, we adopt the Transformer structure to capture sequential information. We feed the input sequence $(s_{1},a_{1},\dots,s_{H},a_{H})$ into the Transformer and aggregate the output representations on each state $s_{t}$ to compute the logits with an additional preference attention layer proposed by [60]. The embedding dimension of the Transformer is $128$ with $1$ layer and $1$ header. We adopt the same feature extractor of the RL policy to process input observation and a simple embedding layer to process the action.

Another way to involve background knowledge is to guide the agent with potential goals of the environment. Recent works [21, 41, 61] propose several techniques to query LLMs for helpful goals and thus guide the agent to visit states of interest. In our approach, we try to discover such potential goals in an offline manner. We first caption sampled trajectories $\tau$ from the dataset and ask LLMs for potential goals the agent can reach. Then we store the output goals $g$ along with the trajectory caption, forming a text-based goal library for further utilization. During the RL process, we introduce a retrieval procedure to select appropriate goals from the goal library according to the agent history. When the agent is in state $s_{t}$, we use the text form of agent history $\tau_{t_{1}}$ to retrieve an exhaustive goal list $g_{1},\dots,g_{k}$ from the top-$K$ most similar trajectories. The similarity metrics are calculated using embeddings provided by a pretrained language model as the sentence encoder [62]. Then we compute the cosine similarity $\sigma(s_{t},g_{i})$ between the current trajectory $\tau_{t}$ and the goal $g_{i}$ as the achievement of each goal. The value of the potential function in $s$ is the maximal similarity of potential goals $\phi(s)=\max_{i}\sigma(s,g_{i})$ since it provides the degree of how well the agent completes the most potential goal.

To harness LLMs for goal suggestion, we also caption the trajectory sequence and add the caption to the prompt. We then ask the LLM to provide potential goals based on the history of the agent. We ask the LLMs to list all possible goals in the form of an unordered list and stone all the goals along with the trajectory caption. The pair of trajectory captions and corresponding goals compose our desired goal library. For downstream RL, we retrieve the top-$K$ similar pairs according to the similarity between the query text and the text caption from the library, where $K$ is set to $3$. For our approach, we deploy a naive method by computing the cosine similarity for each text caption, but we note that approaches to accelerating this retrieval process are well-studied, which goes beyond our scope. We adopt a small-scale pretrained BERT model [63] from HuggingFace¹¹1https://huggingface.co/prajjwal1/bert-small as our sentence encoder.

We summarize the properties of our proposed three variants on background knowledge representation in Table I. Although these methods all leverage LLMs to express background knowledge, the differences in representing potential functions result in different behaviors during downstream RL tasks. BK-Code does not require an additional parameterized model but calculates auxiliary rewards solely using written code. In contrast, BK-Pref and BK-Goal require pretrained models to express potential functions. The similarity computation in BK-Goal requires text-form environmental descriptions, limiting its applications to environments that can provide a text captioner. The other two methods only require text captions in collected data samples, making it practical to only caption data samples rather than implementing a general captioner. Besides, BK-Code is infeasible for RL tasks with unstructured input such as images and texts since we cannot directly parse such input using code.

Principles on prompt design. The prompts in background knowledge representation typically contain two portions, the environmental information and the mission description. An exception comes from BK-Code, which requires additional code snippets for programming. We make the designed prompts modularized and succinct, applicable to different environments with moderate efforts in replacing the environment-related prompt. In our supplementary materials, we list all used prompts, discuss how to adapt to other domains, and provide some example responses from different environments and LLMs.

Minigrid [64] contains a spectrum of configurable tasks in a grid environment. We adopt the Minigrid environment [64] (Figure 5(a)) from its official GitHub repository³³3https://github.com/Farama-Foundation/Minigrid. For this environment, we mainly use its goto and pickup tasks types from the BabyAI domain [67] to build tasks. To specify the domain knowledge for a comprehensive analysis, we mainly focus on the goto tasks from the BabyAI series [67], which is a subset of original Minigrid tasks. A goto task typically requires the agent to navigate to a specific object in a map with multiple distractors. The target object type can be a ball, a box, or a key with a specific color from red, blue, green, or purple. Similarly, a pickup task requires the agent navigate to an object and perform a pickup action. The original environment registration in the code repository is text-conditioned, where the agent may receive different text instructions for separate episodes. To unify the task goal, we make specifications for the environment, making it generate consistent goals during RL processes. We create a series of goto and pickup tasks for our experiments. A goto task typically requires the agent to navigate to a specific kind of object, e.g., a red ball in the grid world. As for its extension, the goto-seq task requires the agent to sequentially navigate to two different objects in one episode, bringing more difficulties for agent exploration. A pickup task additionally asks the agent to perform a pickup action after navigating to the object. The original BabyAI environments usually contain map sizes ranging from $5$ to $8$, which can be simple for exploration methods. To create more challenging benchmarks, we scale the map size to $20-30$ to conduct our experiments. To introduce diverse downstream tasks for evaluation, we design different kinds of unseen tasks on purpose. The data for background knowledge representation is collected from tasks without targets of the object type key and the color purple. However, during the downstream RL training, we aim to train policies in these tasks within one-time knowledge representation. Therefore, we can evaluate the effectiveness of acquired background knowledge in a spectrum of seen and unseen tasks.

In Figure 3, we show the average episodic returns of different methods in four downstream tasks with different targets. We find that our proposed three variants generally outperform the compared baselines in all tasks. Notably, when conducting experiments in tasks with unseen object types and colors, the improvement of sample efficiency of our methods is still obvious. The results indicate that our way of representing background knowledge does acquire task-agnostic background knowledge of the environment, thus accelerating policy learning in unseen tasks. We also find that Lang-ND and L-NovelD, which both adopt state captions to provide intrinsic motivation, perform better than classic exploration baselines, demonstrating that leveraging text features can be useful to improve sample efficiency.

Crafter [65] is a 2D survival game (Figure 5(b)) drawing inspiration from the popular Minecraft game. The environment is procedurally generated and partially observable, where the agent needs to complete achievements such as collecting and crafting. The original Crafter game rewards agents based on the accomplishment of achievements and health changes. Based on the achievement list, we create downstream tasks corresponding to each achievement where the agent can only get a reward when completing it. We also adopt the environment from its official repository⁴⁴4https://github.com/danijar/crafter. However, the original action space in Crafter merges different operations like eating, collecting, and attacking into a single ‘do’ action. Though successfully reducing the action space for exploration, this implementation may hinder understanding of environmental logic. Following the solution in [21], we split this action into several actions including ‘eat’, ‘drink’, ‘attack’, and ‘collect’, making it more aligned to actual agent behaviors. The transformation on the action space also helps explain actions during text captioning. We create different downstream tasks in the Crafter environment by splitting the achievement list of the environment into different tasks. We evaluate our proposed methods in these tasks except for BK-Code since the observation space in Crafter is image-based. As shown in Figure 4, we find that the two applicable approaches, BK-Pref and BK-Goal, still exhibit superior performance in most downstream tasks. For crafting tasks that require specific action sequences like making a wood pick axe, introducing background knowledge can significantly improve sample efficiency.

Some recent works also consider learning RL policies with LLM knowledge besides our compared baselines, such as Motif [56] and ELLM [21]. However, these methods are not directly comparable to our method due to different training paradigms and the inability of prompt design in each task. Motif also takes the idea of learning from LLM-labeled preference, also known as RL from artificial intelligence feedback (RLAIF) [54, 55, 68], which directly uses the preference data to augment rewards. Unlike Motif, our framework focuses on more general background knowledge of a domain to avoid complex prompt design and costly LLM queries. In addition, ELLM adopts a different paradigm to generate possible goals during policy pretraining. In contrast, our framework does not require the engagement of LLMs within the interactions with the LLMs but extracts background knowledge from LLMs in an interaction-free manner. Although these methods cannot realize an efficient paradigm like ours, we can compare them in each single task. In Figure 6, we show the results in a map with size 30 to enlarge its difficulty, where these baselines extract knowledge from each single task data while our methods still use previous acquired knowledge. We find that our three methods can still perfrom well compared to these baselines, indicating the effectiveness of integrating LLM knowledge via reward shaping.

In this section, we further test the generalizability of acquired background knowledge, aiming to find whether our framework of using background knowledge can accelerate policy learning for more distinct tasks. We conduct experiments in tasks with different task types and map sizes from the highly configurable Minigrid environment to evaluate the performance of our methods.

Effectiveness on emerging task types. The data used for acquiring background knowledge contains experiences from the goto tasks, which only require the agent to navigate to one object. Here we further introduce two additional task types to examine whether the background knowledge can provide more general training signals. The extra task types include (1) the goto-seq task, where the agent should navigate to two distinct objects sequentially, and (2) the pickup task, where the agent needs to execute a “pickup” action after successfully navigating to the target. As shown in Figure 7, our methods still exhibit superior performance compared to other baselines. Specifically, the goto-seq task type is relatively difficult as it requires exhaustive exploration to find two desired objects. Our compared baselines can hardly solve this kind of task while our methods can solve the problem without previously encountering these tasks. We find that BK-Code is less stable and shows less promising performance than the other two variants in some tasks, indicating that this form of knowledge representation may not be as generalizable as other variants.

Scaling to larger map sizes. When the map becomes larger, the agent needs to execute more actions precisely to acquire the reward, resulting in increasing difficulty in exploring and exploiting such reward signals. To this end, we configure a series of goto tasks with different map sizes to examine whether reward shaping with background knowledge can scale well to larger maps. We evaluate our methods in this task series and plot the average episodic returns at the $2$M-th time step in Figure 8. Notably, our proposed methods can all maintain high sample efficiency with increasing task difficulty. In contrast, the compared baselines mostly exhibit significant performance drops when the map becomes larger. The results indicate that our three variants can scale to larger maps well.

In this section, we conduct experiments to find out the sensitivity of learned background knowledge under the conditions of different language models and data. In Figure 9, we evaluate our three variants with two LLMs, gpt-3.5-turbo and gpt-4, separately. The results show that BK-Code is the most sensitive approach when we turn to a weaker LLM. Since the quality of code strongly correlates to the capability of LLMs, we also find that weaker LLMs have a higher probability of generating code with runtime errors. The performance drop when using gpt-3.5-turbo in BK-Pref and BK-Goal is hardly observed in Minigrid but is more evident in the crafting task from Crafter, where an LLM with lower capability may have a more superficial environment understanding. To further investigate the effectiveness of our framework under much weaker open-source language models, we also test the performance of BK-Pref and BK-Goal when using a 7B version of Llama-2 chat model in Figure 10. We omit the BK-Code method with the Llama version since we find that the Llama-2 chat model is not capable to generate valid code in our case. We find that BK-Pref and BK-Goal using Llama models perform slightly worse than the methods with GPT models, indicating that the capability of LLMs may affect the quality of extracted knowledge. However, thanks to our effective knowledge representations from preferences or goals, the performance loss is not significant and these algorithms can still solve the task with knowledge from a weak Llama-2 7B model.

Besides, we try to discover the connection between used data and RL performance, as LLMs may not be able to derive sufficient knowledge from low-quality data. we present the policy performance when using data collected by random policies for background knowledge representation, abbreviated as the random data approach. As shown in Figure 11, we still observe significant sample efficiency improvements for these variants. For Minigrid, we find that background knowledge derived from random data is sufficient to accelerate policy learning in goto tasks. For the more difficult task of making a wood pickaxe in Crafter, random policies may not be able to provide useful trajectories, limiting the knowledge LLMs can provide.

As stated in our main paper, we deploy an RND algorithm in the environment and collect data periodically. To be specific, we train RND for $5$M steps and evaluate the learned policy for each $500000$ step. During the evaluation process, we run the policy in $50$ episodes and save the data as our dataset. For Minigrid, we run the RND algorithm simultaneously in multiple tasks without fixing the task types, which means that the collected unlabeled trajectories may come from a distribution of tasks. We specify the task range by restricting the target types as mentioned in our experiments section, where the object type key and the color purple do not appear in the pre-collected data. For the Crafter environment, we directly deploy an RND algorithm in the original environment. The agent will try to maximize the raw Crafter reward based on completing achievements and keeping healthy. We present the properties of collected datasets in Figure 13.

We use the PPO algorithm [47] with the CleanRL [48] implementation, a stable and succinct PPO version for adopting our reward-shaping technique. The only modification to the training process is to add the auxiliary reward to the original environmental reward. To adapt to our selected environments, we use convolutional neural networks to extract features from the input observation. Then the features are fed into different multi-layer perceptrons (MLPs) to compute the action logits or value functions. We show the network structure of RL agents in Figure 12. Due to the different observation spaces of the two environments, the feature processing networks are slightly different. We also report the hyperparameters of the RL algorithms in Table III, where most of the hyperparameters are directly taken from the CleanRL implementation despite some coefficients related to environment sampling and auxiliary rewards.

Our method is less time-consuming since it only requires one-time knowledge extraction and does not call LLM during RL. We show the time cost of all methods in the BabyAI-GoToRedBall-S30 task as an example. In Table II, we find that our methods are time-efficient, except for BK-Goal which needs to compute text embeddings during training. Generally, the time cost of all our methods is comparable with these baselines. In addition, we also present typical values of LLM calls and costs for our three approaches in Table IV. Our framework is cost-efficient since it does not involve LLM calls during RL. In contrast, some prior works like ELLM typically requires millions of LLM calls since it needs interactions with LLMs in the RL pretraining process.