AssistanceZero: Scalably Solving Assistance Games

We thank the reviewer for their thorough review and helpful feedback. Below are responses to individual points and questions:

• Complexity of environment: While the reviewer refers to MBAG as a toy setting, we argue it is substantially more complex than prior environments in assistance games, which used 5x5 or 9x10 grid worlds with less than a dozen possible goals like collecting lemons or gemstones [1, 2]. MBAG instead involves building houses in Minecraft, a more realistic task with over $10^{400}$ possible goals. That PPO struggles even with reward engineering and auxiliary losses underscores the environment's difficulty.

  Overcooked [3] and Hanabi [4], despite being simpler than MBAG, have enabled significant research on collaborative AI—further suggesting MBAG is a meaningful benchmark.

• “It is hard to articulate the second challenge around lines 104 that human model is needed...” This is a challenge articulated in past work, e.g., Fisac et al. [5] write that “more realistic human models [are] therefore crucial in using the CIRL formulation to solve real-world value-alignment problems” (CIRL = assistance games). While past papers have used human models based on RL or planning [1, 2], we are the first to leverage an imitation learning-based human model for assistance games. We will update the introduction to clarify this.

• “AssistantZero predicts the parameters of the [reward] network. How will this scale?” The reward function $R(s\_t, a\_t^\**H**, a\_t^\**R**; \\theta)$ is parameterized by $\\theta$, but is not generally a neural network. The parameters $\\theta$ can be any information that encodes the task/goal/preferences of the human. In our case, $\\theta$ consists of the goal structure, i.e., the configuration of blocks that the human is attempting to assemble. In the case of an LLM assistant for coding, $\\theta$ might consist of a set of test cases that the human would like to satisfy by implementing some function. We will modify Section 2 to make this more clear.

• "Some details are hard to follow. For example, the details of PPO training…” We used the standard PPO loss, which consists of the clipped surrogate objective, the value function loss, and the entropy bonus. We also experimented with modifications to try to boost PPO's performance, as described in Section 4.1 and Appendix D.2. Equation (1), the AssistanceZero loss function, has the policy loss, value loss, reward loss, previous reward KL loss, and the human action prediction loss. Only the value loss is identical between PPO and AssistanceZero. We did try adding the goal prediction loss to PPO but found it did not improve performance (see Appendix D.2).

• “Why does the assistant need a non-Markovian policy?” Past work on assistance games [6] has shown that finding a best response to a fixed human policy is equivalent to solving a single-agent POMDP. Solving a POMDP requires a non-Markovian policy and thus we use a recurrent policy for the assistant.

• “The policy of human is kept constant which is often not true...” We agree that this is not a completely realistic assumption. If the human is also learning about their own reward/goal, then their behavior may be non-Markovian; this setting of “partially observable assistance games” is studied in a recent paper [7]. Humans could also adapt to a particular assistant policy, making their policy non-constant. However, we find that our assistant still works well with real humans, indicating that our assumption of a fixed human policy is reasonable in practice.

• “How would exploration methods help in setting where the reward function is sparse?” MCTS already inherently provides exploration as part of AssistanceZero. Future work could investigate reward shaping, pretraining, and intrinsic motivation to further improve performance in sparse-reward settings.

• “How expensive is MCTS compared to PPO at test-time?” We use 20 simulations of MCTS at test time, requiring 20 forward passes. Despite this, it runs in real time on a 1080 Ti GPU. Notably, AssistanceZero without test-time MCTS matches performance, using the same compute as PPO:

[1] Woodward et al. Learning to Interactively Learn and Assist.

[2] Zhi-Xuan et al. Pragmatic Instruction Following and Goal Assistance.

[3] Carroll et al. On the Utility of Learning about Humans for Human-AI Coordination.

[4] Bard et al. The Hanabi challenge: A new frontier for AI research.

[5] Fisac et al. Pragmatic-Pedagogic Value Alignment.

[6] Shah et al. Benefits of Assistance Over Reward Learning.

[7] Emmons et al. Observation Interference in Partially Observable Assistance Games.

We thank the reviewer for their thorough review and helpful suggestions. Below are responses to individual questions.

Responses to questions

1. The paper assumes access to $\hat{p}_t(\theta)$ in Line 294. How can this be available for general human tasks?
We believe the reviewer is asking about how $\theta$ (the true reward parameter) is available during training. $\theta$ is not visible to the assistant policy when it is actually acting in the environment; the assistant infers a belief distribution over $\theta$ using the reward parameter prediction head. However, once an episode completes, the $\theta$ value that was used by the human model during that episode is available for updating the network to enable better assistance in the future, and thus we can use it as part of the loss function for AssistanceZero.

2. More generally, is assuming the reward structure as known a restricting assumption that blocks application to general AI training? … What is the reward structure used in MBAG and how large is the $\theta$ space in the experiment?

In our MBAG environment, the goal parameters $\theta$ consist of the specific arrangement of blocks in the goal house that the human is trying to build. This space contains over $10^{400}$ possible goal structures, since each block in the 3D grid that makes up the Minecraft world can be 10 different types. In a more complex setting, like a coding assistant, the goal parameters could instead consist of a set of unit tests that the user wants to satisfy by writing a function. One could collect a large dataset of function specifications along with unit tests and define the reward function based on how many unit tests pass at the end of the human-assistant interaction. In an even more general setting, an LLM judge could be used to give rewards based on detailed task descriptions that are used as the reward parameters. This would be distinct from RLAIF because a separate human model is used to interact with the assistant while a powerful LLM judge evaluates whether the latent task is actually accomplished. While these generalizations of our work still require overcoming additional challenges, we believe that our approach can scale to complex assistance settings.

3. What if the reward structure prior is incomplete or even wrong?

We note that our reward structure in MBAG is likely incomplete. It only describes the goal house and doesn’t represent other preferences the human may have for the assistant’s behavior, such as not getting in their way or obstructing their view. Despite this, our assistant performs well with real humans in the user study.

Nevertheless, there are a few ways in which we can handle an incomplete/incorrect reward structure prior. One way is to detect if the human’s behavior is too unlikely under any reward function we can represent, and perform safe fallback actions. Another approach is to combine assistance games with RLHF by including an RLHF loss term that represents human preferences not described in the structured reward model. This is a good direction for future work, which we will mention in the conclusion.

Key benefits of assistance games

“the key is that the agent is also getting a ground-truth reward signal instead of just a human's thumbs-up/down.”

We agree; receiving reward from the latent “true” reward function rather than (potentially biased) feedback from the human is a key ingredient in assistance games. We appreciate the reviewer’s point that the assistance game approach would encourage the AI to be more proactive in order to maximize the shared reward rather than simply react to and please the human. We'll emphasize this important distinction more clearly in the paper.

Human heterogeneity

“The training of both the human model and the assistive model seem to assume that humans are homogeneous… Is this a limiting assumption that future work should seek to relax? Does the authors observe heterogeneity in human data and preferences during experiment?”

While we use a single human model, it can still represent heterogeneous behavior. Theoretically, since we use an LSTM-based human model, it can represent a mixture of human policies. In fact, we expect this to be the case since we observe significant heterogeneity in the human training data ( and the user study). Participants had very different levels of experience with playing Minecraft and computer/video games, so their building speed and play styles varied greatly.

We thank the reviewer for their thorough review and are glad they found the paper proposes a “novel approach” and that our MBAG environment is “also a strong contribution.” We have responded to individual points and questions below.

Small user study

“The user study is fairly small (only 16 participants), which limits the representativeness of the findings.”

Our user study was in-person (1-1.5 hours per participant) because we couldn’t run it on a crowdsourcing platform due to latency constraints. Our sample size is comparable to other in-person HCI studies: the median number of participants in CHI 2019 mixed-method studies was 18 [3]. Moreover, our participants had different amounts of experience with Minecraft and computer games, making our sample diverse.

Finally, our results are statistically significant based on appropriate statistical tests. Our assistant enables humans to build the goal with fewer place/break actions (one-sided t test $p \< 0.05$).

Markovian human model

“Your model assumes that human behavior is Markovian… How might your approach perform if humans exhibit non-Markovian behavior”

We actually do use a non-Markovian human model by employing an LSTM to condition on past observations. Recurrence improved the model’s accuracy (Appendix D.1.2), suggesting humans exhibit non-Markovian behavior. Prior assistance game work also uses recurrent human models [1]. We kept the Markov notation for consistency with the original assistance game formalism.

RLHF comparison

“Have you considered evaluating AssistanceZero against a full RLHF system?”

Comparing with RLHF would be interesting but challenging because RLHF is not easily applicable to the MBAG environment. First, RLHF is usually formulated as a single-agent problem [2, 3], so the additional human agent in our setting would make it difficult to apply standard techniques. Furthermore, in LLMs, RLHF is applied to only a single step of interaction between the assistant and the user, i.e., the comparison data used by RLHF uses conversations which only differ in the last assistant message. In MBAG, the equivalent would be to compare single assistant actions taken in response to a given history of human and assistant actions. However, it may be quite difficult to judge assistant actions in isolation; for instance, more than half of assistant actions are usually movement, and it is unclear how to judge the relative usefulness of say, moving left versus up. For these reasons, we decided to only compare to an SFT baseline, especially since SFT alone for LLMs can often achieve performance close to that of RLHF [7].

Future work could compare assistance games with RLHF in domains where it's more feasible, like LLM post-training.

LLM post-training

“In the real world, humans do not have full observation of the state…”

Our approach could be applied even if the human is learning about their own reward. The assistant would still be rewarded based on the true goal, and we could train a recurrent human policy on data involving human learning. We could follow approaches like [4], which uses RL in a multi-armed bandit setting where the human learns about their reward parameters through noisy observations. Additionally, the assistant could predict the human’s internal state and learning dynamics to help them learn faster [5]. We also note that such “partially observable assistance games” are studied in [1]. We chose to study the simpler case where the human observes their own goal since it is still challenging.

“the most efficient approach would be to communicate their intentions to the AI at the beginning (via language) so that the AI can have full observation of the goal.”

Often, the human cannot fully communicate their goal to the AI. For example, for complex coding tasks, it would be difficult and time consuming to write long, detailed specs. Additionally, if the human wants the code to produce visual output (e.g., a diagram), natural language is insufficient to provide the AI full goal observability. Thus, an assistance game approach would remain beneficial.

Even with perfect communication, an assistant trained with the assistant game framework would likely be more aligned with the human’s true intentions than RLHF, which could deceive humans since it’s trained to seek positive feedback from users [6].

[1] Emmons et al. Observation Interference in Partially Observable Assistance Games
[2] Christiano et al. Deep reinforcement learning from human preferences
[3] Reinhard. Participants, Incentives, and User Studies: A Survey of CHI 2019
[4] Chan et al. The assistive multi-armed bandit
[5] Tian et al. Towards modeling and influencing the dynamics of human learning
[6] Williams et al. On targeted manipulation and deception when optimizing LLMs for user feedback
[7] Zhou et al. LIMA: Less Is More for Alignment.