Unveiling the latent dynamics in social cognition with multi-agent inverse reinforcement learning

We thank the reviewer for great suggestions. We uploaded a revised pdf with changes highlighted in blue and addressed the reviewer's concerns point by point:

Related work. We did another round of literature review and rewrote the related work session. Out of the references the reviewer mentioned, Rabinowitz et al 2018 seems irrelevant because the proposed ToM-net architecture seeks to maximize predictability of other agent’s actions and doesn’t explicitly infer a reward function. Carrol et al. incorporated behavioral cloning (i.e. a supervised classifier to predict action) to aid in human-computer interaction, which is irrelevant to inverse RL.

POMDP v.s. MDP We apologize for the confusion about references. Baker et al 2017 used POMDP while Baker et al 2009 and Velez-Ginorio 2017 use MDP. We revised them in the new paragraph. The ‘While fruitful…’ sentence was paraphrased from the review paper by Jara-Ettinger in 2019 (Fig. 1 legend). It might be confusing to mention it here without context. We deleted this as well as the comparison of MDP and POMDP. As noted in the revised ‘related work’ part, most cognitive sciences work adopted single agent POMDP and multi-agent MDP work to model social interactions. We chose to use multi-agent MDP in our work and look forward to adopting it into a multi-agent POMDP scenario.

John Maynard Keynes' (1937) "general theory of employment" is cited... The beauty contest in game theory is a behavioral finance concept that describes how people make predictions based on what they think others will do, instead of the intrinsic value of an asset. The game is named after John Maynard Keynes, who used the metaphor in his book to explain price fluctuations in the stock market. Burnham et al. (2009) conducted a Keynes beauty contest in 658 subjects and reported value choice from 17 to 50, indicating that the majority of participants exhibited a cognitive reasoning level of approximately zero or one. We have updated the references accordingly in the manuscript.

We cannot be sure that the inference method works for more complex tasks... We appreciate the reviewer’s insightful comments. First, we would like to clarify that our algorithm successfully distinguishes between level-1 and level-0 agents. Level-1 agents make predictions about the other agent’s behavior, while level-0 agents act non-strategically but still pursue self-serving actions using the basic features of the environment. This distinction demonstrates the algorithm's ability to infer differing reasoning capabilities. However, we acknowledge that the model is currently unable to distinguish between two types of level-1 agents that employ different strategies for predicting the other agent's behavior. Regarding the statement, "this simple task does not require a sophisticated mentalization process," we meant that the cooperative foraging task may not demand nuanced or complex predictions about the other agent. Importantly, our results indeed revealed that mentalizing at level-1 of the cognitive hierarchy is critical for success in this task. To prevent any confusion, we have now removed the aforementioned sentence and revised the section accordingly.

Egocentric agents. Egocentric agents perform so differently from other kinds of agents because it takes them significantly more steps to perform the cooperative task (Fig S6). Therefore, we could distinguish these agents even without modeling fitting. We still updated the fit results of egocentric models to Fig.3A. They are substantially worse.

Eq. (2) It should be $\alpha$. We modified Eq.2

Hyperparameter selection. We first tested the learning rate of 0.001, 0.005, 0.01, 0.05. Based on convergence rate and training stability, we picked 0.01 for weight optimization and 0.005 for value map optimization. Fixing these two learning rates and $\lambda_2=1$, we tested $\lambda_1$ from the pool of 0, 1, 5, 10, 20, 30 and got the model evidence curve with respect to this hyperparameter. The log likelihood decreases when $\lambda_2$ increases and the recovered map looks similar for any $\lambda_2>0$ (Fig.S12). We picked $\lambda_1=5$ but also think this choice will not affect most conclusions in the paper. Note that we only care more about the ratio between $\lambda_1$ and $\lambda_2$ because the value map is robust to scaling.

Column distance and row distance In the hallway task, two agents cannot occupy the same grid. Agent 1 (on the left) tends to move to the right, while Agent 2 tends to move to the left (Fig.S10). A state where the column distance is zero indicates that the two agents are on different rows and successfully passing each other. However, a row distance of one can also occur when the agents are still far apart column-wise, which does not necessarily ensure successful task completion.

Fig. 5 We added the fitting results for ‘success’ trials in Fig.5. Their results look very similar to ‘expert’ pairs.

We thank the reviewer for great suggestions. We uploaded a revised pdf with changes highlighted in blue and addressed the reviewer's concerns point by point:

Choice of cognitive hierarchy. We first define a forward process for different levels of cognition, generate a policy matrix based on the reward value map and chosen reasoning process, and then calculate the likelihood of observed state-actions pairs. The appropriate level of cognitive hierarchy was chosen by comparing the modeling fitness over the observed trajectories based on AIC or other potential model fitness criteria.

Computational cost. Assume we have two agents executing independently in a task with $N$ states and $A$ actions, this expands to $N^2$ states and $A^2$ actions when considering joint state and action spaces. In our algorithm, we leverage PyTorch’s automatic differentiation to perform inference of the value function to maximize the likelihood of the observed state-action pairs. The computational complexity is equivalent to the complexity of the forward policy calculation process multiplied by the number of iterations required for convergence of the inference process, which is dependent on the specifics of the optimizer and the loss landscape. Our approach of value decomposition decreases the number of parameters to infer from $O(N^2)$ to $O(N)$, greatly accelerating the iteration process. But it’s hard to give an accurate estimate of the number of iterations needed due to complexity of the loss function. Therefore, we focus on the computational complexity of the forward policy calculation through value iteration.

A level-0 agent follows a standard single-agent value iteration process, requiring the formation of a policy matrix of size $N \times A$. This process incurs a computational cost of $O(K \times N^2 \times A)$ , where K is the number of iterations needed for the value iteration process. A level-1 agent first predicts that its counterpart is a level-0 agent, and then calculates its own policy based on this assumption, involving the formation of a joint optimal policy matrix of size $N^2 \times A^2$. This increases the computational cost to $O(K \times N^4 \times A^2)$ due to the complexity of value iteration at this level. A level-2 agent first forms a prediction model of a level-1 agent, which incurs a computational cost of $O(K \times N^4 \times A^2)$. Subsequently, it computes its own policy through matrix multiplication, maintaining a computational complexity of $O(2K \times N^4 \times A^2)$ overall. For level-$k$ ($k>2$) agent, the computational complexity is $O(k \times K \times N^4 \times A^2)$. In summary, the entire computational complexity equals the cost for policy calculation times the number of iterations in the inference process. The cost of computing the policy increases exponentially with the number of agents and grows linearly with the cognitive level.

Trajectories availability. Our algorithm is mainly devised for cognitive or neuroscience usage. The expert trajectories could come from recorded animal trajectories such as bat mimicking flying, fish schooling, rodent interactions. An existing open dataset is MaBe challenge https://www.aicrowd.com/challenges/multi-agent-behavior-challenge-2022.

Justify the recovered latent belief. We believe the value maps recovered in Fig.2DE are valid because the actual reward location is lit up in these maps. In Fig.F, the value function of Euclidean distance also reflects the essence of a cooperative foraging task where animals need to be simultaneously at the same location to receive reward. We also plotted the occupancy maps (Fig. S7) of the observed trajectories ${(s_1^t, s_2^t, a_1^t, a_2^t)}$ as $a(s_1) = \sum_t \\delta(s_1^t)$, $b(s_2) = \sum_t \\delta(s_2^t)$ and $c(d) = \sum_t \delta (d(s_1^t,s_2^t))$. In the foraging case, $a,b$ looks like a more diffused version of recovered value map $m,n$ but the empirically observed interaction term $c$ has a peak when distance = 1 and the recovered value map highlights the joint states where distance=0. This is because exploration terminates at the rewarded state, causing a smaller occupancy frequency. But MAIRL still recovers distance=0 as the critical state. In the hallway task case, the occupancy map (Fig. S10) has very high values on the middle row but the recovered value map in Fig.4C only highlights the reward locations at the two ends.

Accommodate for partial observability. We included this as a limitation.

Irrational agents Yes. We assume the agents act in a goal-directed manner governed by underlying value functions. The decisions are always made in an optimal manner given a chosen value function. If agents behave irrationally by doing random explorations, our algorithm will discover a value function contaminated by some random variables. For example, we performed MAIRL on random walk trajectories. The recovered value map and interaction map are totally random (Fig.S5)

We thank the reviewer for great suggestions. We uploaded a revised pdf with changes highlighted in blue and addressed the reviewer's concerns point by point:

The presentation of Figure 1 The blue boxes belong to agent 2's policy generation process. Map means a value function which takes location or Euclidean distance as input and outputs a real number as the value of that location. To ensure precise terminology and better alignment with Eq. 2, we have updated Fig. 1 to more accurately reflect its relationship with the equation.

$\hat{a^2}$ represent an action instead of a policy Yes. it should be action. Thanks for catching the typo.

**This work includes human datasets in the hallway task, yet there is no mention of ethics approval or dataset accessibility ** Sorry for the missing information. This dataset was collected by Ho et al in 2016 and is publicly available. The ethics statement was included in the original publication. We revised the corresponding text to emphasize that we were analyzing an open dataset.

The decomposition approach is basic, and the experiments are limited in complexity We respectfully argue that the decomposition method is far from trivial, as kindly highlighted by reviewers KWEN and 85mY. The decomposition approach not only enables efficient inference but also enhances the interpretability of value function estimation. Additionally, we provided a formal analysis of its validity in Session 3.2 and Supplementary Section A.1, which allows users to assess the approximation error based on the task setup and chosen interaction terms before executing the full inference algorithm. This analysis also demonstrated the effectiveness of the decomposition method in a three-agent cooperative foraging task. We appreciate the reviewer’s thoughtful suggestion regarding task selection and will certainly explore the Overcooked task in future work. Thank you for your valuable feedback.

"theory of mind (ToM)." the second line of Introduction should be "theory of mind (ToM)". We’d appreciate it if the reviewer could clarify on the question. They are exactly the same phrases to us.

We thank the reviewer for great suggestions. We uploaded a revised pdf with changes highlighted in blue and addressed the reviewer's concerns point by point:

Limitations of value approximation. We did a mathematical analysis of the validity of value decomposition and numerically evaluated the approximation errors in the Session 3.2 and Supplementary Session A1

Abstract and introduction. We updated them

The considered hierarchical models of behavior are highly simplified... We provided modeling frameworks for level-0, level-1 agents with chance predictions and level-1 agents with perfect predictions. These forward modeling processes may be mathematically simple but definitely not trivial in terms of modeling human/animal behavior as human exhibited a cognitive reasoning level of approximately zero or one. When applied to multi-state, multi-action and multi-agent cases, even level-1 agents could induce complicated behaviors and equilibriums. Using AIC to compare models is a well-established approach in statistical modeling. We'd be happy to incorporate additional model comparison metrics if the reviewer could kindly suggest alternatives.

Limits of the Euclidian interaction map... A good pick of the interaction term is the key to reveal an interpretable value estimation in social tasks. In most cases, it’s hard to interpret the full joint value function because its function domain is too large. In the cooperative foraging task, it is reasonable to use a Euclidean distance interaction term because the task requires the agents to occupy the same grid simultaneously. And indeed, value estimated expert trajectories highlight a huge value at distance equals to zero. In the hallway task, the goal is to switch locations with fewer steps. It is not immediately clear to us which interaction term could be most essential. Therefore, we tested four different interaction terms and concluded that column distance is the most essential one. This reveals the essence of the task: switching column indices. The monkey experiment is to showcase the applicability of our method to a non-cooperative task.

Limitations. We rewrote the paragraph.

Where do you infer the agent's beliefs? We intend to infer the belief from expert trajectories, usually collected through cognitive or neuroscience experiments. An example dataset: https://www.aicrowd.com/challenges/multi-agent-behavior-challenge-2022.

Where do you infer world models and mental states? We removed this phrase.

Line 67... We revised the sentence as “While MARL with recursive reasoning typically models social behavior in goal-directed tasks, ...” In real-world interactions, explicit value functions are often unavailable or poorly defined, unlike in constrained task or game setups. Naturalistic behaviors typically emerge from complex, latent processes involving subjective preferences, hidden goals, and environmental factors. For instance, animal behaviors such as foraging or group movement frequently involve implicit trade-offs (e.g., balancing safety and food acquisition) that are not explicitly represented as value functions.

Line 94... The sentence in question was from the review paper by Jara-Ettinger (2019, Fig. 1 legend). However, we recognize that mentioning it here without sufficient context could lead to confusion and removed it.

Line 101 We removed it.

Eq. 1 We assumed state-action pairs are independent along the temporal sequence, a common assumption used in most IRL literature.

Figure 1 The blue boxes belong to agent 2's policy generation process. We updated Figure 1.

Line 170 Please refer the Supplementary Session A1 and Figure S1, S2 and S3.

Eq2 We assumed K=1. To avoid confusion, we removed the summation over $k$. (Eq.2) Line 243: We have plugged in the marginalization expressions ($P(\hat{a_2}|s)$) based on random explorations.

Figure 2F We replotted Fig.2F with the horizontal axis as the Euclidean distance between agents.

Column distance in the hallway task In the hallway task, two agents cannot occupy the same grid. Agent 1 (on the left) tends to move to the right, while Agent 2 tends to move to the left (Fig.S10). A state where the column distance is zero indicates that the two agents are on different rows and successfully passing each other. However, a row distance of one can also occur when the agents are still far apart column-wise, which does not necessarily ensure successful task completion. Choosing the interaction map depends on the specific task. However, we could validate whether the decomposition is valid or tolerable using the analysis in Supplementary Session A1.

Fig. 4C The trial types were predetermined by the rule. Failure trial: only one participant arrives at the goal. Success trials: both participants get to the goal simultaneously. Expert trials: both participants get to the goal simultaneously in minimum steps (7 steps).