Paper: Aligning Agents like Large Language Models (ICLR-2024)

Reviewer comment (Novelty):
Although the idea is interesting, the technique is too incremental, integrating IL for pretraining, IL for SFT, and REINFORCE for alignment.

Author response:
Regarding novelty, we agree that the individual components of our work, and even the combination of imitation learning and reinforcement learning from preferences, is not novel, as discussed extensively in our related work in Section 2. However, there have since been many improvements to the basic pipeline of imitation learning followed by RLHF (introduced by [1]) in the context of large language models which have not been investigated back in the context of agents. Our main contribution is to apply this modern pipeline to a visually rich and complex domain with real human data, with the application of recent advances to improve the efficiency of the costly online alignment. For example: 

 
**1. Separation of imitation learning into pre-training on general data and fine-tuning on limited task-specific data**

Pre-training on general text followed by instruction fine-tuning is essential for training helpful LLMs, but in the context of agents the usual procedure is to only imitate task-specific data. We demonstrate in Sections 4.1 and 4.2 that similar pre-training followed by task-specific fine-tuning is beneficial for task performance in our context for agents. We’ve also now added an additional ablation in Appendix D that demonstrates that the pre-training stage is beneficial over imitating the task-specific behaviour alone. 

**2. Use of Agent Model / Encoder for Initialising Reward Model**

We demonstrate in Section 4.4 that utilising the agent model trained with an imitation learning loss as the initialisation for the reward model significantly improves the performance of the reward model and improves the feasibility of training a reward model with limited costly human feedback. While this is now typical for LLM reward models, the demonstration that this also applies in our case of reward modelling from pixels for decision-making agents across data sizes is a valuable contribution. 

**3. Introduction of Preference Fine-Tuning on Highest Reward Trajectories**

We demonstrate in Section 4.5 that additional fine-tuning on a subset of the trajectories used for training the reward model with greatest reward significantly improves the efficiency of the subsequent alignment of the model with online reinforcement learning. This approach was concurrently introduced for language models [2], but is novel in the context of agents as far as we are aware. This is a big quality of life improvement for an end-user or game-designer that would like to align a (large) generalist agent, since fine-tuning on existing data is much less expensive than online reinforcement learning in a complex domain.

Reviewer comment (Experiments):
With the incremental technique, I would like to see more sufficient and interesting empirical findings. However, this paper does not contain the discussion including but not limited to scaling law — whether model size scales, data scales, then the performance improves still holds for this setting.

Author response:
Regarding the scaling of the model, we’ve also added additional results on how model size affects performance in our added ablation of the pre-training stage in Appendix D. As well as demonstrating the benefit of the unsupervised pre-training stage with supervised fine-tuning rather than training from scratch on the fine-tuning dataset, we found that larger model sizes (up to the 100M parameters we considered) are beneficial for our task even for training only on the more limited fine-tuning data.

Reviewer comment (Experiments):
With the incremental technique, I would like to see more sufficient and interesting empirical findings. However, this paper does not contain the discussion including but not limited to in-context learning — whether preferences with some demonstrations in the model context can be captured.

Author response:
Investigating whether demonstrations could be incorporated in the agent’s context for in-context learning is an interesting idea, but out of scope of this work on efficiently aligning agents with preferences.

Reviewer comment (Experiments):
With the incremental technique, I would like to see more sufficient and interesting empirical findings. However, this paper does not contain the discussion including but not limited to longer model context — currently, the context length is 32 timesteps. Compared to typical LLMs with at least 2k steps, isn’t it too small to memorize useful information in the context?

Author response:
Our context corresponds to 32 previous timesteps as you mention. Since this is sufficient to determine the agent’s current location and motion, it is possible for the agent to learn our desired navigation behaviour with this context. While longer context could be interesting in more general settings given the partial observability of the environment, we note that this would also be significantly more expensive in terms of compute and memory for little gain, given that there is a lot more information contained within each image of context than there is in each word of context.

Reviewer comment (Experiments):
With the incremental technique, I would like to see more sufficient and interesting empirical findings. However, this paper does not contain the discussion including but not limited to alignment methods — during alignment, try PPO algorithm, which is verified as effective by many papers. Furthermore, why this paper utilizes an **undiscounted** version of REINFORCE? How about direct policy optimization (DPO) that utilizes supervised learning style alignment? As this paper only finetunes the last layer, how about full-parameter finetuning? How about LoRA?

Author response:
We utilised REINFORCE for simplicity and found it to be sufficient given we are only aiming to refine the policy to learn specific behaviours that are not too far from the distribution of the base model. We found discounting makes little difference as we want to reinforce entire preferred trajectories. In general however, using discounting and PPO to better control instability from large updates to the policy may be beneficial as we note in Section 4.5. 
 
Given our results on the effectiveness of preference fine-tuning, we hypothesise that Direct Policy Optimisation (DPO) could be very effective in our setting, but is out of scope of what we could properly analyse in a single conference paper and left for further work as discussed in Appendix J. Due to the offline to online distribution shift in our setting we expect that some online component is likely to be beneficial however, as discussed in the new Appendix E. Similarly we also believe LoRA could be effective in our setting as we mention in Section 4.5, but is also out of scope and left for further work as discussed in Appendix J. Full-finetuning was not possible due to the memory requirements of the dedicated hardware required to run the game online.

Reviewer comment (Experiments):
With the incremental technique, I would like to see more sufficient and interesting empirical findings. However, this paper does not contain the discussion including but not limited to comparison and discussion of offline-to-online reinforcement learning.

Author response:
We agree this is a good point for analysis that we had missed and we thank you for highlighting this. We have added a discussion of this in the new Appendix E (referenced from Section 4.1). We hope this explains the challenges of the significant offline to online distribution shift we have in our environment, and how our approach mitigates many issues associated with these challenges.

Reviewer comment (Experiments):
With the incremental technique, I would like to see more sufficient and interesting empirical findings. However, this paper does not contain the discussion including but not limited to generalization — how about more environments or multi-task setting?

Author response:
Transfer to other environments is out of scope of this work, as discussed in Section 4.1 and Appendix J. Our work is multi-task in the sense that we aim for the base agent to capture diverse behaviour corresponding to many different tasks (such as navigating to different jumppads) and then aligning the agent to perform a particular task.

Reviewer comment (Experiments):
Width: The authors claim they "investigate how the procedure for aligning LLMs can be applied to aligning agents from pixels in a complex 3D environment". Although the environment can be visually perplexing, the task itself might not be challenging enough. Questions remain such as: Is the goal hard to recognize or identify due to the complexity of the 3D environment? Are there many distractions or interferences? Is the goal semantically rich enough to make understanding it a challenge? Does human preference over such a goal go beyond simple multiple-choice selection? Many of these aspects of interest when exploring pretraining–instruction-tuning–RLHF in game AI are unchecked in this paper. Since this scheme has been verified in the language domain with many such challenges and has worked well, investigating it in a different domain (game AI) without comparable challenges could be less illuminating to the community.

Author response:
We thank you for your very constructive and helpful review. We’re pleased that you appreciate the importance of the topic, the valuable contribution we have made to exploring this idea, and our thorough investigation.  
 
To address your concerns regarding the scope of our analysis: 

> -Width: the authors have claimed they "investigate how the procedure for aligning LLMs can be applied to aligning agents from pixels in a complex 3D environment". Although I do agree that the environment can be visually perplexed, the task, however, might not be challenging enough. Is the goal hard to be recognized/identified due to the complextity of the 3D environment? Are there many distractions or interferences? Is the goal semantically rich therefore understanding it could become a challenge? Does the human preference over such goal go beyond simple multi-choice selection? Unfortunately, I find many of these aspects of interest when exploring pretraining - instruction-tuning - RLHF in game AI unchecked in this paper. Point is, this scheme has been verified in the language domain teaming with many of these aforementioned challenges and it actually work pretty well. Therefore, investigating this on a different domain (game AI), but without on par challenges, could be less illuminating to the community. 

We believe that our task is sufficiently challenging, and more importantly, illustrative for our analysis. Since we are using a real AAA video game with real human data from visual-only states, obtaining any behaviour reliably is challenging. The navigation task we considered is complex enough to necessitate our procedure, since there is no reward function to use for RL (and the preference reward model is too sparse for tabula rasa RL given the 3D environment), meaning RL alone is not possible. Therefore we must use limited task-relevant / demonstration data for imitation learning. However, imitation learning may not provide exactly the preferred behaviour and is not very robust to online distribution shift, particularly when learning from pixels, necessitating final alignment with online preference-based reinforcement learning to reliably achieve the preferred behaviour, as demonstrated by our results in Section 4. The task also still involves all the complexities of real human data, with the full visual input and gamepad action output spaces. 
 
As you highlight, this approach has demonstrated success in the language domain, but it is not trivial to apply it to agents in the game domain. In particular, there is significant offline to online distribution shift in our task due to character selection and visual modifications as we discuss in our new Appendix E, as well as the context (equivalent to prompting) that the agent is provided with at initialization as we discuss in Section 4.1. Additionally, learning from pixels and in an environment in which online reinforcement learning can only be performed in real time (due to game engine constraints) provides a much more challenging domain than language benchmarks. Our work therefore demonstrates what can be learned and applied from progress in the language domain, while highlighting additional challenges in the game domain, as in Section 4.5.

Reviewer comment (Reproducibility):
Depth: Although the paper offers rich variants of reward models in the RLHF experiments section, some aspects are missing, such as details on the scale of the model (both the base model and the reward model) and the error range. Including these would help provide a better understanding of applying RLHF in a new domain.

Author response:
> -Depth: although the paper offers rich variants of reward models in the RLHF experiments section, some aspects can still be missing. To name a few: the scale of the model(both the base model and the reward model), error range, etc. These will help with a better understanding on applying RLHF in a new domain. 
 

Our main contribution is to apply the modern LLM training pipeline to a visually rich and complex domain with real human gameplay data, with the application of multiple recent advances this pipeline for language models back to agents to improve the efficiency of the costly online alignment. Some of the contributions of our analysis of applying RLHF in this domain included:  

**1. Separation of imitation learning into pre-training on general data and fine-tuning on limited task-specific data**

Pre-training on general text followed by instruction fine-tuning is essential for training helpful LLMs, but in the context of agents the usual procedure is to only imitate task-specific data. We demonstrate in Sections 4.1 and 4.2 that similar pre-training followed by task-specific fine-tuning is beneficial for task performance in our context for agents. We’ve also now added an additional ablation in Appendix D that demonstrates that the pre-training stage is beneficial over imitating the task-specific behaviour alone. 

**2. Use of Agent Model / Encoder for Initialising Reward Model**

We demonstrate in Section 4.4 that utilising the agent model trained with an imitation learning loss as the initialisation for the reward model significantly improves the performance of the reward model and improves the feasibility of training a reward model with limited costly human feedback. While this is now typical for LLM reward models, the demonstration that this also applies in our case of reward modelling from pixels for decision-making agents across data sizes is a valuable contribution. 

**3. Introduction of Preference Fine-Tuning on Highest Reward Trajectories**

We demonstrate in Section 4.5 that additional fine-tuning on a subset of the trajectories used for training the reward model with greatest reward significantly improves the efficiency of the subsequent alignment of the model with online reinforcement learning. This approach was concurrently introduced for language models [1], but is novel in the context of agents as far as we are aware. This is a big quality of life improvement for an end-user or game-designer that would like to align a (large) generalist agent, since fine-tuning on existing data is much less expensive than online reinforcement learning in a complex domain. 
 
The combination of these additions in this domain makes it possible to align agents with only limited preference data and online RL updates. 
 
Regarding the scaling of the model, we’ve also added additional results on how model size affects performance in our added ablation of the pre-training stage in Appendix D. As well as demonstrating the benefit of the unsupervised pre-training stage over supervised fine-tuning only, we found that larger model sizes (up to the $\sim100M$ parameters we considered) are beneficial for our task even for training only on the more limited fine-tuning data.  
 
We hope that we have addressed your concerns and clarified the value of our contributions towards bridging the gap between the LLM and gaming agent communities. 
 

**References:**

[1] Reinforced Self-Training (ReST) for Language Modeling, https://arxiv.org/abs/2308.08998

Reviewer comment (Novelty):
The paper shows that the same overall process for finetuning LLMs can also be applied to agents within this game. However, the overall method has limited novelty – as noted in the related works section, the general process of bootstrapping RL with imitation learning and further finetuning agent behaviors with reward models learned from human preferences have both been studied in prior work training RL agents outside of LLMs.

Author response:
Regarding novelty, we agree that the individual components of our work, and even the combination of imitation learning and reinforcement learning from preferences, is not novel, as discussed extensively in our related work in Section 2. However, there have since been many improvements to the basic pipeline of imitation learning followed by RLHF (introduced by [1]) in the context of large language models which have not been investigated back in the context of agents. Our main contribution is to apply this modern pipeline to a visually rich and complex domain with real human data, with the application of recent advances to improve the efficiency of the costly online alignment. For example: 

**1. Separation of imitation learning into pre-training on general data and fine-tuning on limited task-specific data**

Pre-training on general text followed by instruction fine-tuning is essential for training helpful LLMs, but in the context of agents the usual procedure is to only imitate task-specific data. We demonstrate in Sections 4.1 and 4.2 that similar pre-training followed by task-specific fine-tuning is beneficial for task performance in our context for agents. We’ve also now added an additional ablation in Appendix D that demonstrates that the pre-training stage is beneficial over imitating the task-specific behaviour alone. 

**2. Use of Agent Model / Encoder for Initialising Reward Model**

We demonstrate in Section 4.4 that utilising the agent model trained with an imitation learning loss as the initialisation for the reward model significantly improves the performance of the reward model and improves the feasibility of training a reward model from pixels with limited costly human feedback. While this is now typical for LLM reward models, the demonstration that this also applies in our case of reward modelling from pixels for decision-making agents across data sizes is a valuable contribution. 

**3. Introduction of Preference Fine-Tuning on Highest Reward Trajectories**

We demonstrate in Section 4.5 that additional fine-tuning on a subset of the trajectories used for training the reward model with greatest reward significantly improves the efficiency of the subsequent alignment of the model with online reinforcement learning. This approach was concurrently introduced for language models [2], but is novel in the context of agents as far as we are aware. This is a big quality of life improvement for an end-user or game-designer that would like to align a (large) generalist agent, since fine-tuning on existing data is much less expensive than online reinforcement learning in a complex domain.  

We hope these three main examples highlight the novelty of our contribution in the context of aligning agents. 

**References:**

[1] Reward learning from human preferences and demonstrations in Atari, https://arxiv.org/abs/1811.06521 

[2] Reinforced Self-Training (ReST) for Language Modeling, https://arxiv.org/abs/2308.08998

Reviewer comment (Experiments):
The authors note that unlike prior work, their goal is to train "a generally capable base agent that can be aligned by game designers or end-users to perform different tasks or behave with different styles". This is a meaningful contribution even if the overall process itself is not novel, however, I believe that the current set of behaviors studied in this work (navigating to 3 different locations) is too limited. The paper would be greatly strengthened if the sets of behaviors or tasks demonstrated by the base and fine-tuned agents were more complex. While I understand the merit of first studying the approach in a simpler setting to disentangle what the agent is learning at each stage, the paper would be more complete with additional experiments showing the methodology also works on learning more complex behaviors, as this is the underlying motivation of the work.

Author response:
In order to provide a meaningful analysis of what is beneficial to transfer from current LLM training pipelines to align agents with user preferences, we intentionally utilise a minimal environment with a quantifiable metric to disentangle the effects of each stage of the pipeline. Our current task provides such an environment which maintains the complexities of the full game, with clear interpretation of how each stage of the training procedure affects the jumppad distribution. The task still *necessitates our procedure* however, since there is no reward function to use for RL (and the preference reward model is too sparse for tabula rasa RL given the 3D environment), meaning RL alone is not possible. Therefore we must use limited task-relevant / demonstration data for imitation learning. However, imitation learning may not provide exactly the preferred behaviour and is not very robust to online distribution shift, particularly when learning from pixels, necessitating final alignment with online preference-based reinforcement learning to reliably achieve the preferred behaviour, as demonstrated by our results in Section 4. The task also still involves all the complexities of real human data involving many characters and visual modifications and associated online distribution shift, with the full visual input and gamepad action output spaces.

Reviewer comment (Presentation):
Rather than just having the bar plots in the first two stages and then reward curves for the final stage, the paper presentation would be strengthened with an overarching connecting figure showing a heat-map of how the agent behaviours across the map evolve at each stage (first showing diverse locations visited, then focusing on the jumppads evenly, then focusing on just the left/right one).

Author response:
We agree that a summary heatmap demonstrating how agent behaviours evolve during alignment is useful addition to demonstrate our core contribution, so thank you for this helpful suggestion. We have added this figure to Appendix H in our paper and referenced it from Section 4.5, with colours corresponding to our original Figure 1.

Reviewer comment (Experiments):
While the more poor performance from the agent finetuning to the right jumppad is likely due to the lack of data and diversity as noted by the authors, this is also exactly the problem that RL should be able to address. Have the authors tried more sophisticated RL algorithms for encouraging diverse exploration (adding in an entropy bonus, injecting randomness, more complex intrinsic rewards on top of the human preferences)?

Author response:
Since our environment is a real AAA video game involving an open 3D world, which must be run in real time for online RL, significant exploration would be infeasibly expensive (particularly for end users). Our motivation is instead to capture a diverse multi-modal distribution of human behaviour, before refining this behaviour through fine-tuning and aligning with preferences to select specific behaviours from this initial base distribution. This is what enables us to align our agent in only 600 RL updates. We believe that a more effective approach would be to have less aggressive fine-tuning, or to fine-tune on a larger task-relevant dataset to prevent this premature loss of diversity in our pipeline, as we discuss in Section 4.5.2. Further improving the sample efficiency of online alignment with additional offline training would be a good direction for further work. 

We hope this helps to explain the challenges of our environment and motivation for our approach, as well as the value of our contribution towards bridging the gap between the LLM and gaming agent communities.