KALM: Knowledgeable Agents by Offline Reinforcement Learning from Large Language Model Rollouts

Thank you for your time and valuable comments. Below, we address each of your questions and provide detailed responses to your concerns.

Q1: Constrained by the limited context length, generating a complete trajectory seems unreliable for more complex tasks with longer trajectories. Moreover, there is a lack of details regarding the trajectory length of the experimental tasks.

Comment 1: Constrained by the limited context length, generating a complete trajectory seems unreliable for more complex tasks with longer trajectories.

Response 1: This is a good point. We understand your concern about the trajectory generation on long-horizon tasks. We clarity that current experiments primarily demonstrate the efficacy of our proposed methodology, which extracts knowledge from LLM to facilitate low-level control. For tasks with longer trajectories, there have been established techniques verified to be effective, e.g., uncertainty estimation and branch rollout in MOPO [1], and long-range coherence in Sora [2], which is capable of generating video exceeding one minute. While these techniques are compatible with KALM, they were not employed in this study since the average trajectory length was approximately 70 timesteps (as shown below), which did not necessitate their integration. It would be interesting for future research to investigate how these techniques could be incorporated to handle longer trajectory tasks within the KALM framework.

Comment 2: there is a lack of details regarding the trajectory length of the experimental tasks.

Response 2: Thanks for your suggestion. We have included the details regarding the trajectory length of offline data, as shown in the following table.

Q2: The proposed method requires the environment-built-in reward function to obtain and incorporate rewards, which may not be accessible for some other tasks. What if training the LLM to generate the rewards?

We acknowledge that the availability of a built-in reward function can be a limitation in environments where such functions are not readily accessible or well-defined. However, we would like to emphasize that the main focus of this work is to explore the effective utilization of knowledge in LLM for low-level control, and the proposed method is not inherently dependent on the existence of such a reward function. Considering the potential absence of such functions, training LLM to generate rewards is a potential solution, while it may add burden to the LLM training. As an alternative, there have been other well-validated methods utilizing LLM to generate reward function [3,4] or dealing with sparse reward problems [5,6].

Q3: Questions about experiment setting.

Comment 1: How is the dataset collected? What is the quality of the policy used to collect data?

Response 1: We collect the offline dataset by utilizing expert policies specifically trained to complete natural language goal for each task. To simulate the noise in real-world data, we introduce a small portation (10%) of rollouts generated by the random policy. In this way, the dataset reflects a degree of the variability of the environment. We would incorporate these details in the experiment setting in the revised version.

Comment 2: Why use only 6400 rollout-goal pairs to train offline RL, given there are 100,000 rollout-goal pairs available?

Response 2: We apologize for the confusion brought by the dataset utilization. This is attributed to an issue of computation resources. The generation of imaginary rollouts is computationally expensive and time-consuming, with the capacity to generate approximately 6,400 pairs in a 24-hour period. Therefore, to prevent potential sample bias and maintain a balanced dataset, we choose an equivalent number of offline rollouts from the available 100,000 pairs.

Nevertheless, we acknowledge it worth exploring to utilize all the 100,000 offline rollouts, by balancing the imaginary / offline rollout in each training batch. The experiment results on CLEVR-Robot are shown in the table:

The result suggests that an increase in the amount of offline rollouts can lead to improvements in method performance. Overall, the KALM still outperforms the baseline under this setting.

Reference

[1] MOPO: Model-based Offline Policy Optimization. Tianhe Yu, et al. NeurIPS 2020.

[2] Video generation models as world simulators. OpenAI team. 2024.

[3] Text2Reward: Reward Shaping with Language Models for Reinforcement Learning. Tianbao Xie, et al. ICLR 2024.

[4] Eureka: Human-Level Reward Design via Coding Large Language Models. Yecheng Jason Ma, et al. ICLR 2024.

[5] Learning by playing solving sparse reward tasks from scratch. Riedmiller, et al. ICML 2018.

[6] Overcoming exploration in reinforcement learning with demonstrations. Ashvin Nair. ICRA 2018.

Thanks for your response and positive assessment. While current experiments have verified the feasibility of our motivation, we believe it would be valuable supplement by evaluating KALM on more complex tasks. During the rebuttal phase, we have investigated KALM's adaptability on visual tasks. In future work, we would further explore KALM's performance on tasks with longer horizon, employing the established techniques discussed in our response to Q1.

We appreciate your insightful comments and finding our idea interesting. To address your concern, we have devoted effort to include discussion and experiment on model-based RL. Please find our response below.

Q1: Model-based offline RL methods like MOPO, MOReL and COMBO demonstrate strong performances and are also capable of generating imaginary rollouts just like KALM. They should be discussed and included in the baselines.

Comment 1: Discussion about model-based RL methods.

Response 1: We agree with you at this point. In our initial submission, we touch upon model-based RL (MBRL) methods (MAPLE and RedM) in Section 2.2, referring to them as environment models. However, we acknowledge the need for a more comprehensive discussion about prevalent MBRL methods:

MBRL algorithms learn a dynamic model from offline data, which can then be used by any policy learning algorithm to recover the policy. MOPO [1] and MOReL [2] use uncertainty quantification to construct a lower bound, aiming to avoid issues like model bias and distribution shift. COMBO [3] employs both the offline dataset and model-generated rollouts to train a value function and regularize it on out-of-support state-action tuples. Despite both MBRL and KALM methods utilized generated rollouts for policy training, they are different from motivation. KALM, in particular, extracts knowledge from pre-trained LLM to build knwoledgeable agent. Leveraging the LLM's general world knowledge, we demonstrate that LLM has the potential to generate rollouts for unseen goals, extending beyond the scope of offline data and enabling the acquisition of novel skills.

We would like to incorporate the discussion into the revised version.

Comment 2: I will increase my score to 6 upon seeing comparisons between KALM and model-based offline RL. This comparison directly demonstrates the effectiveness of using LLMs instead of traditional neural networks to generate rollouts.

Response 2: Thanks for your suggestions. We have conducted experiments to compare KALM and two representative baselines: COMBO and MOPO, as the results in the table:

KALM surpasses two compared model-based offline RL methods, demonstrating the effectiveness imaginary of rollouts generated from LLM.

Q2: The policy in KALM also conditions on G, a hidden vector generated by BERT. I think this introduces additional information to the policy. Does the other baseline also use this information?

In our experiments, all baseline methods utilize the BERT to convert textual goals into vector representations, ensuring a fair comparison. We appologize for the omission of this detail in the description of experiment setting. To further improve the comprehensiveness of our experiment, we have now included an addition baseline, DT [4], that utilizes both LLM and offline dataset. DT utilizes Llama-2-7b-chat-hf as the backbone policy model and offline data to training the policy. It treats decision-making as a sequence modeling problem, using a transformer architecture to predict actions based on the desired future return. In this experiment, DT trains on the offline data same as other methods, and the results are as follow:

The results show that KALM outperforms DT on all four types of tasks.

Q3: This is an open question. What do you think are the major differences between KALM and world models? Is there any possibilities to extend KALM to more complicated environments?

Comment 1: What do you think are the major differences between KALM and world models?

Response 1: Both world models and KALM train model to generate environmental rollouts, but their motivations are different: KALM leverages the embedded knowledge within pre-trained LLM, utilizing this pre-existing information to generate rollouts. Conversely, world models begin with no prior knowledge and learn from extensive datasets, constructing knowledge from foundational principles.

Besides, world models and KALM utilize the offline data for different purposes. World models particularly excel in visual tasks, processing observations such as high-dimensional images. They utilize offline datasets to abstract observation features and learn environmental dynamics. On the other hand, KALM utilizes offline data to train policy and applies the offline data to bridge the gap between LLM and the environmental.

Comment 2: Is there any possibilities to extend KALM to more complicated environments (even real-world cases)?

Response 2: As we discuss in Section 5, we have considered extending KALM to more complex tasks, e.g., tasks with visual input. We make attempt towards this direction on Meta-world environment with visual observation (visual input is closer to the real world). We utilize a randomly initialized ViT as the vision encoder/decoder to process visual input. Figure 3 in the attached PDF file shows the example of the generated rollout, given a language goal unseen during the training. The generated rollout successfully captures the core information of the environment, reflecting the reasonable robotic movements tendency. This result provides evidence of the potential for extending KALM to more complicated environments.

Reference:

[1] MOPO: Model-based Offline Policy Optimization. Tianhe Yu, et al. NeurIPS 2020.

[2] MOReL: Model-Based Offline Reinforcement Learning. Rahul Kidambi, et al. NeurIPS 2020.

[3] COMBO: Conservative Offline Model-Based Policy Optimization. Tianhe Yu, et al. NeurIPS 2021.

[4] Decision Transformer: Reinforcement Learning via Sequence Modeling. Lili Chen, et al. NeurIPS 2021.

Thanks for the reply. Your experiment results demonstrate LLM-based environment models are better than previous models in MBRL, even though the generated rollouts are numeric (which, to my best of knowledge, is a weakness of current LLMs). Am I getting your conclusions right?

Thank you for raising your score. We are glad the additional comparisons and discussions address your concerns. We ensure these content would be incorporated in the future version.

In addition to the above rebuttal, we would like to clarify some misunderstandings.

It seems that the reviewer thought that KALM uses the LLM to learn from interacting with the environment by employing reinforcement learning (or DPO as the reviewer asked) to maximize the reward from the environment. This way could have been an ONLINE RL setting, where interactions with the environment are allowed.

However, our whole setting is very close to the OFFLINE RL, where no online interactions are allowed during training, but only a fixed set of offline trajectories are available, plus a general-purpose LLM model. The LLM is fine-tuned using the offline trajectories, and is then used to generate trajectories in a purely imaginary way, just the same as language generation. The real and imaginary trajectories are then fed to offline RL methods. Therefore, comparing with previous offline RL methods that can only learn from offline trajectories is not unfair.

We will revise to make our setting easier to understand, and will add more LLM-related baselines as in the responses to other reviewers.

Thank you for your valuable comments and time in reviewing this paper. It appears there may be some confusion regarding the fairness of the experimental comparison. Below, we provide a comprehensive response to address these concerns.

Q1: Concerns about the fairness of comparison.

Comment 1: I am not sure that this comparison is fair and reasonable. Since KALM uses imaginary rollout data in addition to an offline dataset, it may not be considered as an offline method.

Response 1: We first clarify the fairness of the comparison: all methods in our experiments are offline methods, which does not require online interaction with the environment (we would discuss the environment interaction in detail in response 2 below). KALM fine-tunes a general purpose LLM with the offline data, and generates imaginary rollouts without online interaction, thus adhering to the offline setting and ensuring the fair comparison. In the initial submission, we have made a comparison between KALM and a LLM-related baseline (LLM as policy), as shown in Figure 5. To further justify the proposed method and address your concern, we add an additional baseline for comparison, Decision-transformer (DT) [1], which utilizes Llama-2-7b-chat-hf as the backbone policy model and offline data to training the policy. DT treats decision-making as a sequence modeling problem, using a transformer architecture to predict actions based on the desired future return. In this experiment, DT trains on the offline data same as other methods, and the results are as follow (as well shown in Figure 1 in the attached PDF file):

The results show that KALM outperforms DT on all four types of tasks. We would like to update the paper to incorporate the additional result.

Comment 2: When generating imaginary rollouts, does KALM interact with the environment?

Response 2: KAML does not require interaction with the environment, and the imaginary rollouts can be generated in pure offline manner. We appologize for any confusion brought by Figure 2, which involves an optional online process. For clarity, it should be noted that in our comparative experiments, KALM is implemented without the online procedure. The inclusion of the optional online process in Figure 2 is intended to illustrate the method's extendibility and capability for future integration with real-time environment interactions.

Q2: Does KALM use rewards from the environment to filter out low-quality rollouts?

KALM does not employs rollout filter. We use all imaginary rollouts generated by LLM. We suggest that offline RL algorithm can also learn useful knowledge from failure experience. Besides, these failure rollouts may contribute to the robustness of the resulting policy by providing a diverse range of scenarios for the model to learn from, although they are not successful rollouts. This insight is aligned with the comment proposed by Reviewer rv4h.

Q3: Can we use DPO (Direct Preference Optimization) instead of offline RL algorithms? Are there any pros and cons?

DPO can not be directly applied to our problem setting, as it requires preference data annotated by human or AI and is initally designed for aligning LLMs with human value. A potential adaptation of DPO for integration with KALM is treating the offline rollout as the preferred data and the generated rollouts as the less-preferred data. However, this is still unfeasible as we do not have offline rollouts (i.e., preferred data for DPO) that correspond to unseen language goals. Note that the primary motivation of this research is to leverage the knowledge of LLMs to develop agents with enhanced knowledge, this goal is not dependent on the use of any specific offline RL optimization technique.

We hope that our response has addressed your concern and questions satisfactorily. If you had any further concerns, we are glad for discussion.

Reference:

[1] Decision Transformer: Reinforcement Learning via Sequence Modeling. Lili Chen, et al. NeurIPS 2021.

Thank you for providing thoughtful responses to my questions. I could understand KALM more. Especially, I appreciate the effort that the authors have made in providing additional experiments with Decision Transformers (DT). Accordingly, I raise my rating.

Hi reviewer 1PB3. As the discussion period comes to an end, we want to follow up to see if the response addresses your concerns. If you have any further questions, please let us know. Thank you again!

We appreciate your highlight of the method novelty and the significance of the problem setting. Please find our response to each comment as follow.

Q1: While the method seems to work in these simple simulated settings, the method itself seems susceptible to language model hallucinations. It would be interesting to see statistics on how often the model imagines objects / obstacles that don’t exist, or implausible rollouts in the simulator. Alternatively, it would be interesting to see that having hallucinations can actually lead to increased robustness.

Hallucination is indeed a central issue to consider when we use LLM. The results in Table 1 provide a general result in term of matching rate between generated rollouts and the given goal. In response to your comment, we have conducted more detailed statistics on the quality of the generated rollouts. Note that current experiments focus on the environment state in numerical vector, where each dimension in the vector has specific semantic. Thus the LLM would not generate rollouts containing impossible objects. Here we focus on investigating the implausible rollouts as follow topics: 1. the object is out of the table; 2. the object floats in the air; 3. implausible robotics joint pose (e.g., out of the joint rotation angle bound) and 4. exceeding dynamics limits between two steps. The statistics results are as follow:

The results indicate while there is a certain level of hallucination present in the outputs of the LLM, the majority of the imaginary rollouts remain within reasonable scope. Besides, the results on CLEVR-Robot demonstrate lower anomaly ratio. These hallucinations may serve as a form of domain randomization that improves the policy's robustness.

Q2: Given that the motivation is to improve generalizability via LLMs, it would be helpful to also include a comparison against RL methods that use prior knowledge from LLMS (e.g. the cited LLaRP). While the LLM as a policy method touches on this, it would support the motivations of the paper more if there were direct comparisons against methods that use LLMs for world knowledge & priors.

Good point. We agree with you that an additional baseline utilizing both LLM and offline data more would be a strong supplement to our experiments. In response to your concern, we conduct comparison with an additional baseline, Decision-transformer (DT) [1], with Llama-2-7b-chat-hf as the backbone policy model. DT treats decision-making as a sequence modeling problem, using a transformer architecture to predict actions based on the desired future return. In this experiment, DT trains on the offline data same as other methods, and the results are as follow:

The results show that KALM outperforms DT on all four types of tasks. We would like to update the paper to incorporate the additional result.

Q3: Why do the authors think that the ablation study for Table 1 shows that adding components of the method actually leads to a degradation of performance on rollout accuracy on the unseen (easy) set? The paper presents an explanation that the current objectives don’t address learning action semantics, but this also applies to the unseen (hard) tasks and does not explain a degradation in performance.

Answer: We suggest that the performance degradation is attributed to the specific nature of the unseen (easy) task, whose objective is to predict one-step transitions given unseen language goals. To be more specifically, we would discuss the two components of LLM fine-tuning in KALM (i.e., dynamics prediction and rollout explanation) respectively. For dynamics prediction objective, unseen (easy) task objective (predicting $a_t$ and $s_{t+1}$ given $G$ and $s_t$) shares similarity, yet diverges from the dynamics prediction objective (predict $s_{t+1}$ given $s_t$ and $a_t$). This difference introduces a potential conflict in the LLM SFT. This is evidenced by the empirical results that KALM w/o dynamics prediction achieves highest rollout accuracy for unseen (easy) task. For rollout explanation objective, it focuses on giving explanation on a long rollout sequences. While this objective enriches the model's capability to provide coherence over temporal sequence, it may inadvertently detract from the model's ability to capture the immediate logic of transitions between adjacent two steps.

We would like to update the paper to reflect the above discussion about the experiment results.

Q4: There are a few grammatical and writing errors in the paper.

We have fixed your mentioned typos and would make a more thorough check to fix the grammatical and writing errors.

We believe the experiment and the result analysis have been clearly enhanced based on your c comments. If you had any further concerns, please let us know.

Reference:

[1] Decision Transformer: Reinforcement Learning via Sequence Modeling. Lili Chen, et al. NeurIPS 2021.