Adversarial Distributed Multi-Task Meta-Inverse Reinforcement Learning with Theory of Mind and Mean-Field Method

Abstract

Maximum entropy adversarial inverse reinforcement learning (ME-AIRL) has garnered widespread attention for its ability to learn rewards and optimize policies from expert demonstrations. In complex multi-task environments, applying meta-learning ME-AIRL to acquire rewards requires a substantial volume of homogeneous expert demonstrations across all tasks, which is often impractical in real-world scenarios. Moreover, interference between tasks further escalates computational complexity. To solve these challenges, this paper proposes a distributed multi-task meta ME-AIRL framework based on theory of mind and mean field, referred to as TMMF-MTAIRL. In TMMF-MTAIRL, the theory of mind is used to capture the relationships and representational information among multiple tasks. Furthermore, TMMF-MTAIRL integrates mean-field theory to transform interactions between complex tasks into interactions between the main task and the average of the remaining tasks. Furthermore, additional latent variables are introduced to enhance adaptation to novel tasks. We evaluate the proposed TMMF-MTAIRL on point-maze benchmarks and a real-world rolling bearing fault diagnosis dataset using metrics such as classification accuracy, mean rewards or cumulative rewards. TMMF-MTAIRL achieves the best performance across all tasks, with an average improvement of 0.16 in accuracy of fault classification over the strongest baseline.

1. Introduction

In the field of artificial intelligence, reinforcement learning (RL) has become one of the core technologies for solving complex decision-making problems by enabling interactive exploration between agents and their environments [1,2,3,4,5]. Traditional RL largely depends on pre-designed reward functions. However, crafting a reward function that both accurately captures task objectives and maintains generalization capability often proves highly challenging [6,7,8]. Therefore, inverse reinforcement learning (IRL) has been proposed to infer implicit reward functions and optimize policies by observing expert demonstrations, thereby avoiding the subjectivity and high costs associated with manual reward design [9,10,11,12]. As real-world scenarios become increasingly complex, single-task IRL has begun to reveal significant limitations. In domains such as autonomous driving, medical robotics, and industrial automation, agents must simultaneously handle multiple interrelated tasks. However, conventional single-task IRL models struggle to capture shared knowledge among tasks and to resolve conflicting objectives [13,14]. These challenges have spurred the emergence of multi-task inverse reinforcement learning (MT-IRL), which integrates IRL’s reward inference capability with knowledge-sharing mechanisms derived from multi-task learning, making it a prominent research focus in recent years [15,16,17].

The core objective of MT-IRLs is to infer either shared or task-specific latent reward functions across multiple tasks. Researchers have enabled cross-task knowledge transfer by introducing hierarchical Bayesian models and meta-learning framework [18,19,20]. For instance, IRL approaches based on probabilistic graphical models can map expert demonstrations from different tasks into a unified latent space, thereby extracting shared reward features [21,22]. By integrating deep neural networks (e.g., Transformer), researchers have proposed end-to-end multi-task IRL models that can automatically extract inter-task correlations from high-dimensional, homogeneous expert data [23,24,25]. For example, the dynamic properties of different tasks can be captured through multi-head attention mechanisms, and adaptive reward functions can be generated [26]. To solve the black box problem of traditional IRL in complex tasks, the causal graph models have been integrated into the multi-task framework. By analyzing causal relationships in expert demonstrations, the models can not only infer reward functions but also identify causal dependencies between tasks, thereby enhancing both policy interpretability and generalization performance [27]. Despite these advances, a significant challenge remains: limited or suboptimal expert demonstrations can adversely affect both reward learning and policy optimization.

To address this challenge in MT-IRLs, generative adversarial networks (GANs) have been integrated into MT-IRLs, aiming to resolve large-scale task problems in complex environments [28,29,30]. To tackle the challenge of autonomous reward acquisition in large-scale, high-dimensional environments with unknown dynamics, Finn et al. incorporated GANs into guided cost learning to learn reward functions that are robust to dynamic changes [31]. To address the poor transferability and robustness of recoverable rewards in varying environments, an adversarial IRL algorithm was proposed to learn hierarchically disentangled rewards through policy learning [32]. Additionally, semantic rewards were incorporated into the adversarial IRL to further enhance and stabilize its performance [33]. To make the graph computation end-to-end microscopic, the model-based adversarial IRL employed a dynamic self-attention model to optimize low-variance policies [34]. Moreover, to improve the adaptability of both adversarial IRL and multi-task adversarial IRL, meta-adversarial inverse reinforcement learning (Meta-AIRL) integrated meta-learning with adversarial IRL to learn adaptive policies [21,35].

However, existing adversarial IRL methods are often designed under the assumption of homogeneous expert presentations [35]. In real-world reward learning scenarios, however, updated mixed expert demonstrations for different tasks may include both human-generated experience data and algorithm-generated data, where state-action visitation frequencies can vary significantly across tasks. Such variations lead to heterogeneous modalities and distinct distribution frequencies. Moreover, these partially heterogeneous multi-task settings require a large number of expert demonstrations to learn rewards, which increases the computational complexity. To address these challenges, researchers have developed a mean-field RL model and a multi-agent IRL framework based on the mean-field game, which effectively mitigate high-dimensional state spaces and heavy computational burdens through simplified modeling approaches [36,37,38,39,40]. Moreover, prior studies in RL and IRL have incorporated the theory of mind (ToM) to enhance information interaction, thus improving the understanding and prediction of other other agents’ behaviors and the intentions of others [41,42,43,44]. Inspired by the application of mean-field theory and ToM, and aiming to improve joint reward learning and knowledge transfer in multi-task IRL under complex partially heterogeneous environments, this paper proposes a meta-adversarial distributed multi-task inverse reinforcement learning based on the theory of mind (TMMF-MTAIRL). Given updated and mixed expert demonstrations, the proposed algorithm constructs the adversarial mean-field inverse reinforcement learning based on ToM and deep latent variables to enable explicit modeling of relationships between multiple tasks, enhancing information utilization among tasks, and facilitating rapid adaptation to novel tasks. Therefore, by simultaneously improving task interpretability and adaptability, TMMF-MTAIRL can more effectively infer potential reward functions, and optimize agents’ policies.

In summary, the main contributions of this paper are as follows:

(1)

To enhance knowledge transfer between tasks, this paper proposes a meta-adversarial IRL framework based on theory of mind, which uses the theory of mind to make better use of the relationships between partially heterogeneous tasks and capture the representational information between tasks.

(2)

To reduce the computational complexity of reward and strategy optimization, the mean-field theory is introduced into TMMF-MTAIRL, which can transform interactions between complex tasks into the interaction between the main task and the rest of the average tasks.

(3)

By introducing additional latent variables to characterize modal information within mixed expert demonstrations, the TMMF-MTAIRL enables rapid adaptation to novel tasks during meta-adversarial training of the discriminator and generator, even with limited expert demonstration data.

(4)

The effectiveness of TMMF-MTAIRL is evaluated on both the maze-point benchmark and rolling bearing fault diagnosis tasks. Experimental results indicate that TMMF-MTAIRL outperforms conventional multi-task IRL methods, achieving state-of-the-art performance in both reward learning and policy optimization.

2. Related Works

Adversarial IRL and multi-task adversarial IRL: Adversarial IRL represents a state-of-the-art research direction that combines GANs with IRL, aiming to efficiently infer latent reward functions of expert behaviors through adversarial training mechanisms. Initially, Finn et al. introduced GANs into guided cost learning [31], and used the distribution of state-action pairs to optimize the overall goal of adversarial IRL: ${\mathbb{E}}_{\pi }\left[{\sum }_{t=0}^{T-1}{\gamma }^{t}r\left(s_t,a_t\right)\right]={\mathbb{E}}_{\pi }\left[{\sum }_{t=0}^{T-1}{f}_{\theta }\left(s_t,a_t\right)-\text{log }\pi \left(a_t|s_t\right)\right]$. In the adversarial IRL algorithm, sampling-based methods are used to calculate the gradient of the loss function with respect to $\theta$, and the training is performed according to the trajectory distribution [32]. In practice, however, such estimates based on trajectory distributions exhibit high variance. Consequently, to cope with complex dynamic environments, researchers have proposed a series of extensions, including behavioral fusion IRL [45], model-based IRL [34], and semantic-based augmented adversarial IRL [33], etc. To address the challenge of inferring reward functions from expert trajectories in multi-agent environments, multi-agent generative adversarial IRL leverages diverse agents’ behaviors to learn shared reward functions [46]. Furthermore, for multi-task scenarios, several approaches have been developed, such as hierarchical multi-task adversarial IRL [19], meta-adversarial IRL [20,21,22], variational multi-task IRL [23], and multi-task IRL based on self-attention dynamic model [26]. In addition, by incorporating causal concepts into reward learning and policy optimization for multi-task IRL, researchers have proposed causal entropy-based multi-task IRL [27]. To further enhance the adaptability of adversarial multi-task IRL, an IRL based on meta-learning is developed to learn adaptive policies through different update frequencies and meta-learning rate [35].

IRL based on mean field and ToM: In complex multi-task environments, agents must interact with highly intricate surroundings, resulting in substantial computational complexity that makes reward learning and policy optimization particularly challenging [20,23]. As the number of agents grows significantly, researchers have proposed mean-field-based RL algorithms to reduce computational complexity, which introduces mean-field theory into multi-agent RL. Typical examples include mean-field Q-learning, mean-field actor–critic algorithms [36], GAT-MF [39], and HMF-based DRL [46], which are representative RL methods based on mean field. Building on these studies, researchers have gradually introduced mean-field methods into IRL to better address the challenges of learning reward functions and optimizing policies [37,47]. Specifically, based on the adversarial IRL framework combined with meta-learning, researchers proposed MFG-AIRL based on the mean field [40]. Furthermore, to enhance belief modeling and intention inference in learning rewards, researchers have proposed ToM-based IRL [40], Bayesian MoT-based IRL [41], intrinsically motivated ToM-based IRL [42], and goal-oriented IRL based on ToM [43]. These approaches aim to predict an agent’s future actions by analyzing its past trajectories. Overall, mean-field theory and ToM have achieved preliminary research progress in RL and IRL [44], effectively facilitating reward learning and policy optimization. These developments provide valuable inspiration for the theoretical framework proposed in this paper.

Table 1. Comparisons of related adversarial IRL methods based on mean field and ToM.

Methods	Faced Problems	Advantages	References
Behavioral fusion IRL	Difficult to adapt to dynamic environments	Estimations based on trajectory distributions have small differences	[45]
Model-based IRL			[34]
Semantic-based augmented adversarial IRL			[33]
Hierarchical multi-task adversarial IRL	Difficult to adapt to multi-task scenarios	Capturing information among multiple tasks in real time	[19]
Meta-adversarial IRL			[20,21,22]
Variational multi-task IRL			[23]
Causal entropy-based multi-task IRL			[27]
Mean-field actor–critic algorithms	Complexity of large-scale interactive computing and modeling	Reducing computational complexity, improving generalization ability	[36]
GAT-MF			[39]
HMF-based DRL			[46]
ToM-based IRL	Modeling of multi-agent interactions is difficult; knowledge transfer for heterogeneous tasks is challenging	Enhanced cognitive depth and generalization ability of agents
Bayesian MoT-based IRL			[41]
Intrinsically motivated ToM-based IRL			[42]
Goal-oriented IRL based on ToM			[43]

shows the faced problems, advantages, and representative references of the related adversarial IRL methods based on mean field and ToM.

3. Preliminaries

3.1. Markov Games

A Markov decision process (MDP) is formulated as $\left(\mathcal{S},\mathcal{A},R,\gamma ,P\right)$, where $\mathcal{S}$ represents the set of states, $\mathcal{A}$ represents the set of actions, $R_s^a=\mathbb{E}\left[r_{t+1}|s_t=s,a_t=a\right]$ indicates the set of reward function, and $\gamma \in \left[0,1\right]$ represents the discount factor. The function $P_{s{s}^{\prime }}^a=\mathbb{P}\left(s_{t+1}|s_t,a_t\right)$ is the state transition distribution. The agent takes actions $a_t$ from the state $s_t$ to the next state $s_{t+1}$ with probability $p\left(s_{t+1}|s_t,a_t\right)$. The agent aims to maximize the the expected cumulative discounted rewards ${\mathbb{E}}_{\pi }\left[{\sum }_{t=1}^T{\gamma }_t{r}_t\left(s_t,a_t\right)|s_0=s,a_0=a\right]$. Based on the MDP in RL, if we cannot obtain the reward function of the IRL, the MDP in IRL is formulated as $\left(\mathcal{S},\mathcal{A},\gamma ,P\right)$. Given the expert demonstrations $\mathcal{D}$, the main objectives of the IRL are as follows: (1) to recover the reward function, and (2) to realize policy optimization by making the learner’s policy approach the expert’s policy.

3.2. Adversarial Inverse Reinforcement Learning

To solve the reward functions under expert demonstrations $\mathcal{D}$, Ziebart et al. [42] proposed maximum entropy inverse reinforcement learning (MaxEnt IRL), which formulates constrained optimization problems as maximum entropy models [17,35]:

$$\begin{array}{cc}\hfill & p\left({\xi }_j|\theta \right)={\displaystyle \frac{1}{z\left(\theta \right)}}{e}^{{\sum }_{s_t\in {\xi }_j}{\theta }^{\mathrm{T}}{\varphi }_{s_t}}{\prod }_{\left(s_t,a_t,s_{t+1}\right)\in {\xi }_j}p\left(s_{t+1}|s_t,a_t\right)\hfill \end{array}$$

(1)

where ${\varphi }_{s_t}$ represents the feature function of each reward, the partition function $z\left(\theta \right)={\sum }_{{\xi }_j}{e}^{{\sum }_{s_t\in {\xi }_j}{\theta }^{\mathrm{T}}{\varphi }_{s_t}}{\prod }_{\left(s_t,a_t,s_{t+1}\right)\in {\xi }_j}p\left(s_{t+1}|s_t,a_t\right)$ is a normalized constant, and $p\left(s_{t+1}|s_t,a_t\right)$ represents the state transition probability. In Equation (1), the trajectory ${\xi }_j$ that the expert demonstrates is exponentially proportional to the cumulative rewards along that trajectory, implying that trajectories with higher rewards are more likely to be preferred by the agent. Furthermore, the MaxEnt IRL algorithm is extended to the multi-task environment. Given the expert demonstrations $\mathcal{D}={\left\{{\xi }_j\right\}}_{j=1}^M$, the agent can obtain the reward function and optimize the policy by imitating the expert strategy ${\pi }_E$, aiming to maximize the objective function:

$$\begin{array}{cc}\hfill & \underset{\theta }{max}J\left(\theta \right)=\underset{\theta }{max}{E}_{{\xi }_j\sim \mathcal{D}}J\left[\text{log }{p}_{\theta }\left({\xi }_j\right)\right]\hfill \end{array}$$

(2)

where the likelihood probability distribution $p_{\theta }\left({\xi }_j\right)$ of the expert trajectory ${\xi }_j$ is parametrized as $p_{\theta }\left({\xi }_j\right)\propto p\left(s_0\right){\prod }_{t=0}^{T-1}p\left(s_{t+1}|s_t,a_t\right)e^{r_{\theta }\left(s_t,a_t\right)}$, which explains the likelihood probability of the observed data in terms of the reward function; that is, the probability distribution of the expert trajectory ${\xi }_j$ is induced by r.

The reward function is obtained by minimizing $D_{KL}\left(p_{{\pi }_E}\left(\xi \right)\left|\right|p_{\theta }\left(\xi \right)\right)$:

$$\begin{array}{cc}\hfill & \underset{\theta }{\text{arg min }}{D}_{KL}\left[p_{{\pi }_E}\left({\xi }_j\right)\left|\right|p_{\theta }\left({\xi }_j\right)\right]=\underset{\theta }{\text{arg min }}{\mathbb{E}}_{{\xi }_j\sim p_{{\pi }_E}\left({\xi }_j\right)}\left[\text{log }{p}_{\theta }\left({\xi }_j\right)\right]\hfill \\ \hfill & =\underset{\theta }{\text{arg min }}{\mathbb{E}}_{{\xi }_j\sim p_{{\pi }_E}\left({\xi }_j\right)}\left[\sum _{t=1}^T{r}_{\theta }\left(s_t,a_t\right)\right]-\text{log }{Z}_{\theta }\hfill \end{array}$$

(3)

However, the partition function $Z_{\theta }$ is difficult to calculate. To solve this problem, adversarial IRL has been proposed as an efficient framework for MaxEnt IRL. In adversarial IRL, a discriminator $D_{\theta }\left(s,a\right)$ is trained:

$$\begin{array}{cc}\hfill & D_{\theta }\left(s,a\right)={\displaystyle \frac{e^{r_{\theta }\left(s,a\right)}}{e^{r_{\theta }\left(s,a\right)}+\pi \left(a|s\right)}}\hfill \end{array}$$

(4)

To make policy ${\pi }_{\theta }$ more closely approximate the expert’s policy, the discriminator minimizes the cross-entropy loss between expert demonstration and generated samples:

$$\begin{array}{cc}\hfill & L\left(\theta \right)={\sum }_{t=0}^T-E_D\left[\text{log }{D}_{\theta }\left(s_t,a_t\right)\right]-E_{{\pi }_t}\left[\text{log}\left(1-D_{\theta }\left(s_t,a_t\right)\right)\right]\hfill \end{array}$$

(5)

The optimal objective of the discriminator is achieved when $\pi ={\pi }_E$, where $\pi$ denotes the policy learned by the generator, and ${\pi }_E$ is the expert’s policy. At this point, the optimal output of the discriminator is 0.5, indicating that $e^{r_{\theta }\left(s,a\right)}={\pi }_E\left(a|s\right)$.

To learn reward functions that can adapt to different tasks, meta-learning is introduced within the adversarial maximum entropy IRL (AMEIRL) to infer task characteristics from expert policy samples and learn rewards [29,31]. The likelihood function distribution of samples with different trajectories is formulated as

$$\begin{array}{cc}\hfill & p\left(\xi =\left\{s_{0:T},a_{0:T}\right\}|\upsilon \right)={\displaystyle \frac{1}{Z_{\theta }}}\eta \left(s_0\right)\underset{t=0}{\prod ^T}p\left(s_{t+1}|s_t,a_t\right)e^{{\displaystyle \sum _{t=0}^T}{r}_{\theta }\left(s_t,a_t,\upsilon \right)}\hfill \end{array}$$

(6)

where $\upsilon$ represents the parameter of meta learning, and $r_{\theta }\left(s_t,a_t,\upsilon \right)$ denotes the reward function.

However, in the complex multi-task learning environment, tasks may interfere with each other, which brings challenges to learning rewards and optimizing policies for the agents. Effectively leveraging the relationships between tasks, capturing the representational information between tasks, and reducing the complexity of optimizing rewards and policies are thus important problems to address. Motivated by mean-field theory in IRL, which primarily addresses interactions and computational challenges among large-scale agents, we incorporate mean-field approximation into multi-task IRL to reduce the complexity of learning rewards and optimizing policies.

4. Multi-Task Meta-Adversarial Mean-Field IRL Based on Theory of Mind

In complex, partially heterogeneous multi-task environments, using meta-learning adversarial IRL to learn reward functions traditionally requires a large set of homogeneous expert demonstrations, and the presence of multiple tasks may lead to mutual interference, which significantly increase the algorithm’s computational complexity. However, in practical task environments, ensuring that the acquired expert demonstrations across different tasks are homogeneous remains a significant challenge. To address these challenges, this paper presents a meta-adversarial mean-field IRL framework based on the context-conditional policy to learn rewards and optimize policy from the constantly updated mixed expert demonstrations. The algorithm incorporates theory of mind to enhance knowledge transfer between multiple tasks, thereby capturing inter-task relationships and fully leveraging environmental information. Furthermore, the mean field is used to model the interactions between the main task and the average of the remaining tasks, which reduces computational complexity. Additionally, latent task variables are accurately inferred from sampled trajectories for the new task, ultimately deriving task-specific adaptive reward functions and optimal policies under the maximum entropy framework. Moreover, the overall generative adversarial framework can continuously be updated and optimized to improve the mixed expert demonstrations.

Different from existing mean-field IRL and ToM-based IRL, we introduce a ToM-conditioned mean field that aggregates cross-task interactions into an embedding of the averaged state of mind, and derive a mutual-information regularized adversarial learning objective with tractable gradient estimators, which can quickly adapt to new tasks with latent task variables and continuously improve mixed expert demonstrations.

4.1. Adversarial Mean-Field IRL Based on Probabilistic Context Variable

In complex, partially heterogeneous multi-task environments, the context-conditional meta-adversarial mean-field IRL framework first samples M expert demonstration trajectories ${\xi }_j$ from the update mixed expert demonstrations $\mathcal{D}$, then learns reward functions and optimizes policies under each trajectory ${\xi }_j$. Each reward and policy learning process with an expert demonstration trajectory is treated as a task. These mixed demonstration trajectories exhibit diverse behavioral patterns and varying distribution frequencies. Therefore, the reward learning and policy optimization task can be viewed as M partially heterogeneous distributed learning tasks, aiming to infer latent reward functions across multiple tasks.

Using the update mixed expert demonstrations, the adversarial mean-field multi-task IRL (MMF-MTAIRL) introduces the probabilistic contextual variable $\upsilon \in \mathcal{U}$ to describe the distribution of task characteristics, where $\mathcal{U}$ is the value space that can enhance the learning precision of rewards for partially heterogeneous multiple tasks. In this framework, both rewards and policies $\pi$ depend on the probabilistic implicit variable $\upsilon$. We make the following definitions: the expert policy ${\pi }_E\left(a|s,\upsilon \right)$ and corresponding reward functions $R\left(s,a,\eta ,\upsilon \right)$ that depend on mean field and state of mind for different tasks, $\eta$-conditional mean-field flow ${\eta }_E\left(·|\upsilon \right)$ and policy flow ${\pi }_E\left(·|\upsilon \right)$, and marginal distribution of expert policy $p_{{\eta }_E,{\pi }_E}\left(\tau \right)=\int p_{{\eta }_E,{\pi }_E}\left(\tau |\upsilon \right)p\left(\upsilon \right)d\upsilon$, where $p\left(\upsilon \right)$ indicates a prior distribution. MMF-MTAIRL with probabilistic contextual variable aims to use meta-learning to learn an inference model $q\left(\upsilon |\tau \right)$ and rewards with samples of the $p_{{\pi }_E}\left(\tau \right)$.

4.2. The Embedding of Theory of Mind for Improved MMF-MTAIRL

In complex distributed and partially heterogeneous multi-task environments, the interconnections among multiple partially heterogeneous tasks significantly increase computational complexity. To facilitate knowledge transfer across multiple tasks, we incorporate theory of mind into MMF-MTAIRL (TMMF-MTAIRL), which explicitly models the latent relationships among tasks.

First, following the task partitioning in Section 4.1 under partially heterogeneous expert demonstrations, we assume there are M tasks $\left\{{\mathcal{T}}_1,{\mathcal{T}}_2,\dots ,{\mathcal{T}}_M\right\}$. The assumption that there exist M tasks reflects a common practical setting in which expert demonstrations are collected from multiple sources or under different operating conditions, leading to partially heterogeneous data. In such cases, demonstrations can be naturally partitioned into task-specific subsets, either explicitly (e.g., by data source) or implicitly through latent context inference. Modeling these subsets as distinct tasks enables a principled multi-task IRL formulation and provides a foundation for scalable task interaction modeling using the mean-field approximation. Importantly, this assumption does not require precise task boundaries or manual task labeling, and it mainly serves as a modeling abstraction to facilitate tractable learning and analysis.

For the improved MMF-MTAIRL with theory of mind, each task corresponds to rewards $r_i\left(s,a,s_i^{ToM}\right)$ and policies ${\pi }_i\left(a|s,s_i^{ToM}\right)$ that depend on the state of mind. Then the state of mind is encoded using an encoder network. Before defining the state of mind for ToM, we provide an intuitive interpretation of its components. In our multi-task setting, each task not only has its own learning status but also maintains an internal belief about other tasks, which determines how knowledge should be transferred across tasks. Therefore, we define the state of mind $s_i^{ToM}$ as a compact representation consisting of three parts: (i) a loss function $g_i$, (ii) belief variables $b_i^j$ over other tasks, and (iii) a confidence term ${\sigma }_i$. Specifically, the loss function $g_i$ reflects the learning status of task ${\mathcal{T}}_i$ and helps the model capture task-specific dynamics. The belief term $b_i^j$ represents the task i’s belief over other tasks j to model inter-task relations. And the confidence term ${\sigma }_i$ measures the confidence of such inter-task influences and controls the strength of cross-task influence, thereby reducing negative transfer when task relationships are weak or noisy. Based on these considerations, the state of mind for task $T_i$ is defined as $s_i^{ToM}=\left[g_i,b_i^1,\dots ,b_i^{M-1},{\sigma }_i\right]$. In this encoder, an LSTM is used to extract state of mind $s_i^{ToM}$ from historical data ${\xi }_i$ of the task.

The architecture of the ToM encoder is as follows. For each task ${\mathcal{T}}_i$, we construct a historical observation sequence from the sampled trajectory ${\xi }_i={\left\{(s_t,a_t)\right\}}_{t=1}^T$. At each time step t, the state-action pair $(s_t,a_t)$ is first mapped into a feature function ${\varphi }_t$ through a learnable embedding (a fully connected layer). The feature sequence ${\left\{{\varphi }_t\right\}}_{t=1}^T$ is then fed into an LSTM encoder to capture temporal dependencies and behavioral patterns in ${\xi }_i$:

$$h_t,c_t=\mathrm{LSTM}({\varphi }_t,h_{t-1},c_{t-1})$$

(7)

where $h_t$ and $c_t$ denote the hidden state and cell state, respectively. The final hidden representation $h_T$ is used as the summary of the task’s historical behavior and is projected via a multilayer perceptron to obtain the embedding for the state of mind:

$$s_i^{ToM}=f_{ToM}\left(h_T\right)$$

(8)

The output $s_i^{ToM}$ is designed to encode the loss function $g_i$, belief $b_i^j$, and the confidence ${\sigma }_i$. During training, the ToM state is dynamically updated, and the average state of mind $s_{avg}^{ToM}$ is recomputed by the LSTM every K steps based on the most recent mixed expert demonstration trajectories. We use a 2-layer LSTM with hidden size H, and the output dimension of $s_i^{ToM}$ is $d_{ToM}$.

TMMF-MTAIRL, based on theory of mind, simplifies multi-task interactions by modeling each task in interaction with an average task $\tilde{\mathcal{T}}$. The average state of mind is calculated as:

$$\begin{array}{cc}\hfill & s_{avg}^{ToM}={\mathbb{E}}_{i\sim u\left(\right\{1,\dots ,M\}\setminus \{j\left\}\right)}\left[s_i^{ToM}\right]\approx {\displaystyle \frac{1}{M-1}}\sum _{i\ne j}{s}_i^{ToM}\hfill \end{array}$$

(9)

where $s_i^{ToM}$ represents the state of mind of the expert demonstration trajectory ${\xi }_i$, and u indicates the uniform distribution. Each task ${\mathcal{T}}_i$ interacts only with the average task $\tilde{\mathcal{T}}$, rather than with all other tasks, thereby reducing computational complexity. Note that $s_i^{ToM}$ represents continuous embedding produced by the ToM encoder; thus, the averaging in Equation (9) is implemented as an element-wise mean over the embedding dimensions, which can be interpreted as an empirical expectation under the mean-field approximation.

In MMF-MTAIRL with ToM, according to the Bayes’ theorem, these probabilistic contextual variables follow a prior distribution $p\left(\upsilon \right)$. In this complex partially heterogeneous multi-task environment, the mean-field context-conditional policy based on theory of mind with different $\upsilon$ represents the task preference with shared structures. The policy, mean field, mind, and reward function all depend on $\upsilon$, which are defined as: $\pi :\mathcal{S}\times \mathcal{U}\to \Delta \left(\mathcal{A}\right)$, $\eta :\mathcal{U}\to \Delta \left(\mathcal{S}\right)$, $R:\mathcal{S}\times \mathcal{A}\times \Delta \left(\mathcal{S}\right)\times \mathcal{U}\to \mathbb{R}$. For simplicity, we assume that the state space, action space, state transition function, and initial state distribution associated with mean field and mind are independent of $\upsilon$. Based on the above definition, the mind-based context-conditional trajectory distribution induced by a policy $\pi$ is formulated as:

$$\begin{array}{cc}\hfill & p_{\eta ,\pi }=\left(\xi |\upsilon ,\theta \right)={\eta }_1\left(s_1\right)\underset{t=1}{\prod ^T}\pi \left(a_t|\left(s_t,s_i^{ToM}\right),\upsilon \right)P\left(s_{t+1}|\left(s_t,s_i^{ToM}\right),a_t,{\eta }_t\left(·|\upsilon \right)\right)\hfill \end{array}$$

(10)

In Equation (10), $\theta$ denotes the parameters of the reward $r_{\theta }(·)$ obtained by Equation (6). Therefore, the trajectory distribution depends on $\theta$ through the reward function $r_{\theta }$ and the corresponding MaxEnt-induced policy.

With the prior distribution $p\left(\upsilon \right)$ of the latent mind-based context variable and the conditional distribution in Equation (10), the optimal mind-based mean-field context-conditional policy aims to maximize the entropy-regularized expected rewards:

$$\begin{array}{cc}\hfill & {\pi }^{*}=\underset{\eta ,\pi }{\text{arg max }}{\mathbb{E}}_{\upsilon \sim p\left(\upsilon \right),\tau \sim p_{\eta ,\pi }\left(\xi |\upsilon \right)}\left[\sum _{t=1}^{T}r\left(\left(s_t,s_i^{ToM}\right),a_t,{\eta }_t,\upsilon \right)-\text{log }{\pi }_t\left(a_t|\left(s_t,s_i^{ToM}\right),\upsilon \right)\right]\hfill \end{array}$$

(11)

where ${\eta }_t\left(s_t,s_i^{ToM}|\upsilon \right)$ follows the consistency between the mean-field flow and policy flow.

Let $R\left(s,a,\eta ,\upsilon \right)$ indicate the reward function of expert, ${\pi }_E\left(a_t|s_t,\upsilon \right)$ represent the expert policy corresponding to the optimal solution of Equation (11), ${\eta }_E\left(·|\upsilon \right)$ indicate $\upsilon$-conditional mind flow, and $\xi$ indicate the expert demonstration sampled from the induced marginal distribution $p_{{\eta }_E,{\pi }_E}\left(\xi \right)=\int p_{{\eta }_E,{\pi }_E}\left(\xi |\upsilon ,\theta \right)p\left(\upsilon ,\theta \right)d\upsilon$. Within the framework of TMMF-MTAIRL, given the updated mixed expert demonstrations, the meta-learned inference model $q\left(\upsilon |{\xi }_E\right)$ first infers the task-specific latent probability variable $\tilde{\upsilon }\sim q\left(\upsilon |{\xi }_E\right)$. This enables the corresponding reward function $r\left(s,a,\eta ,\tilde{\upsilon }\right)$ for tasks to guide the learning of an optimal policy that approximates the policies of expert’s rewards $R\left(s,a,\eta ,\tilde{\upsilon }\right)$.

TMMF-MTAIRL explicitly models latent relationships between multiple tasks by embedding the ToM, which enhances both reward learning and policy optimization through cross-task observation and utilization, while also providing more comprehensive environmental information. This improves interpretability and adaptability in multi-task learning.

4.3. TMMF-MTAIRL with Mutual Information Regularization

Under TMMF-MTAIRL, the context variable inference model is parametrized as $q_{\psi }\left(\upsilon |\xi \right)$, and the reward function is parametrized as $r_{\theta }\left(\left(s_t,s_i^{ToM}\right),a,\eta ,\upsilon \right)$, where $\upsilon$ is inferred from $q_{\psi }\left(\upsilon |\xi \right)$. Let $\left({\eta }_{\theta }\left(·|\upsilon \right),{\pi }_{\theta }\left(·|\left(s_t,s_i^{ToM}\right),\upsilon \right)\right)$ represent the mind-based mean field induced by $r_{\theta }$. Therefore, under the $\left({\eta }_{\theta }\left(·|\upsilon \right),{\pi }_{\theta }\left(·|\left(s_t,s_i^{ToM}\right),\upsilon \right)\right)$, the trajectory distribution $p_{\theta }\left(\xi |\upsilon \right)$ induced by the reward function under different potential variables can be expressed as:

$$\begin{array}{cc}\hfill & p_{\theta }\left(\xi |\upsilon \right)={\displaystyle \frac{1}{Z\left(\theta \right)}}\left[\underset{t=1}{\prod ^T}{\eta }_{\theta }^t\left(s_t,s_i^{ToM}\right)\right]e^{{\displaystyle \sum _{t=1}^T}{r}_{\theta }\left(\left(s_t,s_i^{ToM}\right),a_t,{\eta }_{\theta }^t,\upsilon \right)}\hfill \end{array}$$

(12)

where $Z\left(\theta \right)$ represents the partition function, and ${\eta }_{\theta }$ denotes the induced mean-field flow consistent with ${\pi }_{\theta }$.

The maximum likelihood function of the conditional distribution $p_{\theta }\left(\xi |\upsilon \right)$ in Equation (12) is taken to learn rewards corresponding to multiple tasks using the proposed TMMF-MTAIRL. Because $\upsilon$ is inferred from $q_{\psi }\left(\upsilon |\xi \right)$, if the probabilistic context variable $\upsilon$ is not further restricted, $\upsilon$ will be inferred in advance using $q_{\psi }\left(\upsilon |\xi \right)$ in the process of solving the maximum likelihood function, resulting in rewards that are independent of $\xi$. Therefore, it is necessary to introduce an implicit latent context variable and its relationship with the rewards. We first consider the mutual information between the probabilistic contextual variable and the trajectory $\xi$ under the joint distribution $p_{\theta }\left(\upsilon ,\xi \right)=p\left(\upsilon \right)p_{\theta }\left(\xi |\upsilon \right)$, which is defined as

$$\begin{array}{cc}\hfill & I_{p_{\theta }}\left(\upsilon ;\xi \right)={\mathbb{E}}_{\upsilon \sim p\left(\upsilon \right),\tau \sim p_{\theta }\left(\xi |\upsilon \right)}\left[\text{log }{p}_{\theta }\left(\upsilon |\xi \right)-\text{log }p\left(\upsilon \right)\right]\hfill \end{array}$$

(13)

where $p_{\theta }\left(\upsilon |\xi \right)$ is the posterior distribution corresponding to the conditional distribution $p_{\theta }\left(\xi |\upsilon \right)$ calculated by in Equation (12); the trajectory $\xi$ is sampled from a trajectory distribution induced by the parameterized reward function.

Since we cannot directly obtain a prior distribution $p\left(\upsilon \right)$, a posterior distribution $p_{\theta }\left(\upsilon |\xi \right)$, $\upsilon$-induced conditional mean field ${\eta }_E\left(·|\upsilon \right)$, or the conditional mind ${\mu }_E\left(·|\upsilon \right)$, optimizing the solution of interactive information in Equation (13) becomes challenging. To address this challenge, following the InfoGAIL’s formula reasoning in [48], $q_{\psi }\left(\upsilon |\xi \right)$ is first used to estimate experience value of ${\eta }_{\theta }\left(·|\upsilon \right)$ from the updated mixed expert demonstration, thereby converting the $\upsilon$-induced conditional mean dependence on $\upsilon$ into dependence on $\psi$. Subsequently, the inference network $q_{\psi }\left(\upsilon |\xi \right)$ can serve as a variational approximation of the posterior distribution $p_{\theta }\left(\upsilon |\xi \right)$, allowing both the inference of task uncertainty and the approximation of sampling from $p\left(\upsilon \right)$. The detailed estimation process will be elaborated later. Let $p_{{\pi }_E}\left(\xi \right)$ represent the expert trajectory distribution.

In addition to the above mutual information constraints, the learning of rewards and inference network must satisfy two objectives: Goal 1) $\text{arg min }{\mathbb{E}}_{\upsilon \sim p_{\upsilon }}\left[D_{KL}\left(p_{{\pi }_E}\left(\xi |\upsilon \right)\right)\left|\right|p_{\theta }\left(\xi |\upsilon \right)\right]$; Goal 2) $\text{arg min }{\mathbb{E}}_{\xi \sim p_{\theta }\left(\xi \right)}\left[D_{KL}\left(p_{\theta }\left(\upsilon |\xi \right)\right)\left|\right|q_{\psi }\left(\upsilon |\xi \right)\right]$. Goal 1 aims to learn a reward function $r_{\theta }$ that induces data optimally explaining the expert policy. This involves taking mean over the prior for a potential variable $\upsilon$, ensuring both that the trajectory likelihood distributions are as similar as possible, and that $p_{{\pi }_E}\left(\xi \right)$ closely matches $p_{\theta }\left(\xi \right)$. When the first objective is satisfied, $p_{\theta }\left(\xi |\upsilon \right)$ can be considered to represent $p_{{\pi }_E}\left(\xi |\upsilon \right)$. Goal 2 requires that the inference network $p_{\theta }\left(\upsilon |\xi \right)$ approximates the true posterior distribution $q_{\psi }\left(\upsilon |\xi \right)$ of latent variables, enabling $q_{\psi }\left(\upsilon |\xi \right)$ to infer latent context variables from new expert presentations. With the objective of the mutual information in Equation (13) and the constraints of Goal 1 and Goal 2, the mind-based meta IRL with mutual information regularization can be formulated as a constrained optimization problem. We consider using the Lagrangian function to solve this optimal problem,

$$\begin{array}{cc}\hfill & \underset{\theta ,\psi }{min}-I_{p_{\theta }}\left(\upsilon ;\xi \right)+\alpha {\mathbb{E}}_{p\left(\upsilon \right)}\left[D_{KL}\left(p_{{\pi }_E}\left(\xi |\upsilon \right)\left|\right|p_{\theta }\left(\xi |\upsilon \right)\right)\right]+\beta ·{\mathbb{E}}_{p_{\theta }\left(\xi \right)}\left[D_{KL}\left(p_{\theta }\left(\upsilon |\xi \right)\left|\right|q_{\psi }\left(\upsilon |\xi \right)\right)\right]\hfill \end{array}$$

(14)

In Equation (14), let $\alpha =1$ and $\beta =1$; the Lagrangian function can be rewritten as

$$\begin{array}{c}\underset{\theta ,\psi }{min}{\mathbb{E}}_{p\left(\upsilon \right)}\left[D_{KL}\left(p_{{\pi }_E}\left(\xi |\upsilon \right)\left|\right|p_{\theta }\left(\xi |\upsilon \right)\right)\right]+{\mathbb{E}}_{p_{\theta }\left(\upsilon ,\xi \right)}\left[\text{log }{\displaystyle \frac{p\left(\upsilon \right)}{p_{\theta }\left(\upsilon |\xi \right)}}+\text{log }{\displaystyle \frac{p_{\theta }\left(\upsilon |\xi \right)}{p_{\psi }\left(\upsilon |\xi \right)}}\right]\hfill \\ =\underset{\theta ,\psi }{min}-{\mathbb{E}}_{p\left(\upsilon \right)}\left[D_{KL}\left(p_{{\pi }_E}\left(\xi |\upsilon \right)\left|\right|p_{\theta }\left(\xi |\upsilon \right)\right)\right]+{\mathbb{E}}_{\upsilon \sim p\left(\upsilon \right),\xi \sim p_{\theta }\left(\xi |\upsilon \right)}\left[\text{log }{q}_{\psi }\left(\upsilon |\xi \right)\right]\hfill \\ =\underset{\theta ,\psi }{min}-{\mathbb{E}}_{p\left(\upsilon \right)}\left[D_{KL}\left(p_{{\pi }_E}\left(\xi |\upsilon \right)\left|\right|p_{\theta }\left(\xi |\upsilon \right)\right)\right]+{\mathcal{K}}_{mut}\left(\theta ,\psi \right)\hfill \end{array}$$

(15)

Because the negative term $-H_p\left(\upsilon \right)={\mathbb{E}}_{p_{\theta }\left(\upsilon ,\xi \right)}\left[\text{log }p\left(\upsilon \right)\right]={\mathbb{E}}_{p\left(\upsilon \right)}\left[\text{log }p\left(\upsilon \right)\right]$ can be regarded as a constant during the optimization of the parameters $\theta$ and $\psi$, the $-H_p\left(\upsilon \right)$ is therefore omitted.

For calculating Equation (15), the unknown ${\eta }_{\theta ,t}\left(·|\upsilon \right)$ makes it difficult to estimate the KL divergence between the empirical expert distribution and the energy-based trajectory distribution $p_{\theta }\left(\xi |\upsilon \right)$. To calculate the ${\eta }_{\theta ,t}\left(·|\upsilon \right)$, we first consider using the estimate ${\widehat{\eta }}_{E,t}\left(·|\upsilon \right)$ to implement the calculation of ${\eta }_{\theta ,t}\left(·|\upsilon \right)$. However, the dependence of ${\widehat{\eta }}_{E,t}\left(·|\upsilon \right)$ on $\upsilon$ makes it difficult to obtain the posterior distribution $p\left(\upsilon |\xi \right)$. Fortunately, the adversarial maximum entropy-based IRL provides an efficient sampling method for estimating the maximum likelihood function. Consequently, ${\eta }_{\theta ,t}\left(·|\upsilon \right)$ can be used to approximate the $p\left(\upsilon |\xi \right)$, and ${\widehat{\eta }}_{E,t}\left(·|\upsilon \right)$ can be estimated:

$$\begin{array}{cc}\hfill & {\widehat{\eta }}_{\psi ,t}\left(\left(s,s_i^{ToM}\right)|\upsilon \right)={\mathbb{E}}_{{\xi }_E\sim p_{{\eta }_E,{\pi }_E}\left(\zeta \right)}\left[q_{\psi }\left(\upsilon |{\xi }_E\right)·{\mathbb{I}}_{s_t,s_i^{ToM}=s,s_i^{ToM}}\right]\hfill \end{array}$$

(16)

When ${\widehat{\eta }}_{\psi ,t}\left(s,s_i^{ToM}|\upsilon \right)$ replaces ${\eta }_{\theta ,t}\left(·|\upsilon \right)$, $p_{\theta }\left(\xi |\upsilon \right)$ additionally depends on $\psi$. Therefore, based on Equation (16), the conditional distribution can be rewritten as

$$\begin{array}{cc}\hfill & p_{\theta ,\psi }\left(\xi |\upsilon \right)={\displaystyle \frac{1}{Z\left(\theta ,\psi \right)}}\left[\underset{t=1}{\prod ^T}{\widehat{\eta }}_{\psi ,t}\left(s_t,s_i^{ToM}\right)\right]·e^{{\displaystyle \sum _{t=1}^T}{r}_{\theta }\left(\left(s_t,,s_i^{ToM}\right),a_t,{\widehat{\eta }}_{\psi ,t},\upsilon \right)}\hfill \end{array}$$

(17)

Lemma 1.

In the adversarial framework of the mind-based meta IRL with mutual information regularization, the adaptive samplers are defined as ${\pi }_{\theta ,t}\left(a|\left(s,s_i^{ToM}\right),\upsilon \right)$ and the discriminator is written as $D_{\theta }\left(\left(s,s_i^{ToM}\right),a,{\widehat{\eta }}_{\psi ,t},\upsilon \right)={\displaystyle \frac{e^{r_{\theta }\left(\left(s_t,s_i^{ToM}\right),a_t,{\widehat{\eta }}_{\psi ,t}\right)}}{e^{r_{\theta }\left(\left(s_t,s_i^{ToM}\right),a_t,{\widehat{\eta }}_{\psi ,t}\right)}+{\pi }_{\theta }\left(a|\left(s_t,s_i^{ToM}\right),\upsilon \right)}}$. During adversarial training process with deterministic dynamics, the adaptive samplers ${\pi }_{\theta }$ eventually converge to the optimal value with the reward signal $\left(\text{log }{D}_{\theta }-\text{log }\left(1-D_{\theta }\right)\right)$. At this point, the trajectory distribution induced by $\left({\widehat{\eta }}_{\psi },{\pi }_{\theta }^{*}\right)$ corresponds to the maximum entropy conditional trajectory distribution $p_{\theta ,\psi }\left(\xi |\upsilon \right)$:

$$p_{\left({\widehat{\eta }}_{\psi },{\pi }_{\theta }^{*}\right)}\left(\xi |\upsilon \right)={\displaystyle \frac{1}{Z\left(\theta ,\psi \right)}}\left[\kappa \left(s_1\right)\underset{t=1}{\prod ^T}{\widehat{\eta }}_{\psi ,t}\left(s_t,s_i^{ToM}\right)\right]·e^{{\displaystyle \sum _{t=1}^T}{r}_{\theta }\left(\left(s_t,s_i^{ToM}\right),a_t,{\widehat{\eta }}_{\psi ,t},\upsilon \right)}=p_{\theta ,\psi }\left(\xi |\upsilon \right)$$

(18)

Proof. 

See Appendix A.    □

Lemma 2.

The gradient of ${\mathcal{K}}_{mut}\left(\theta ,\psi \right)$ with respect to θ can be estimated by Equation (19):

$$\begin{array}{cc}\hfill & {\mathbb{E}}_{\upsilon \sim p\left(\upsilon \right),\xi \sim p_{{\pi }_{\omega }^{*}}\left(\xi |\upsilon \right)}\left[\text{log }{q}_{\psi }\left(\upsilon |\xi \right)\left[\sum _{t=1}^T{\displaystyle \frac{\partial }{\partial \theta }}{r}_{\theta }\left(\left(s_t,s_i^{ToM}\right),a_t,{\widehat{\eta }}_{\psi ,t},\upsilon \right)\right.\right.\hfill \\ \hfill & \left.\left.-{\mathbb{E}}_{{\xi }^{\prime }\sim p_{{\pi }_{\omega }^{*}}\left(\xi |\upsilon \right)}{\displaystyle \frac{\partial }{\partial \theta }}{r}_{\theta }\left(\left(s_{t+1},s_i^{ToM}\right),a_{t+1},{\widehat{\eta }}_{\psi ,t+1},\upsilon \right)\right]\right].\hfill \end{array}$$

(19)

Proof. 

See Appendix B.    □

Using Lemma 2, $\upsilon$ can be sampled through the generative process and the trajectories are sampled using the conditional trajectory distribution $p_{\theta }^{*}\left(\xi |\upsilon \right)$ to estimate ${\displaystyle \frac{\partial }{\partial \theta }}{\mathcal{K}}_{mut}\left(\theta ,\psi \right)$. Therefore, the whole training goal is:

$$\begin{array}{cc}\hfill & \underset{\omega }{\text{min }}\underset{\theta ,\psi }{\text{max }}{\mathbb{E}}_{{\xi }_E\sim p_{{\pi }_E}\left(\xi \right),\upsilon \sim q_{\psi }\left(\upsilon |{\tau }_E\right),\left(s,a\right)\sim {\rho }_{{\pi }_{\omega }}\left(s,a|\upsilon \right)}\text{log }\left(1-D_{\theta }\left(\left(s_t,s_i^{ToM}\right),a_t,{\widehat{\eta }}_{\psi ,t},{\upsilon }_t\right)\right)\hfill \\ \hfill & +{\mathbb{E}}_{{\xi }_E\sim p_{{\pi }_E}\left(\xi \right),\upsilon \sim q_{\psi }\left(\upsilon |{\tau }_E\right),\left(s,a\right)\sim {\rho }_{{\pi }_{\omega }}\left(s,a|\upsilon \right)}\text{log }\left(D_{\theta }\left(\left(s_t,s_i^{ToM}\right),a_t,{\widehat{\eta }}_{\psi ,t},{\upsilon }_t\right)\right)+{\mathcal{K}}_{mut}\left(\theta ,\psi \right)\hfill \end{array}$$

(20)

where $D_{\theta }\left(\left(s_t,s_i^{ToM}\right),a_t,{\widehat{\eta }}_{\psi ,t},{\upsilon }_t\right)={\displaystyle \frac{e^{r\left(\left(s_t,s_i^{ToM}\right),a_t,{\upsilon }_t\right)}}{e^{r\left(\left(s_t,s_i^{ToM}\right),a_t,{\upsilon }_t\right)}+\pi \left(a_t|\left(s_t,s_i^{ToM}\right),{\upsilon }_t\right)}}$.

In the discriminator, the inputs consist of both states $s_t$ and the average state of mind $s_{avg}^{ToM}$, while the outputs are task-specific rewards ${\widehat{r}}_i\left(s_t,a_t,s_i^{ToM}\right)$. Incorporating of average state of mind introduces additional dependency conditions for task-specific rewards, ensuring that multiple learned rewards do not correspond to identical policies, thereby resolving the problem of reward ambiguity.

The pseudo-code of the multi-task meta-adversarial mean-field IRL based on theory of mind in the training process is as shown in Algorithm 1. In the proposed TMMF-MTAIRL, all neural network parameters are initialized using standard He initialization. Training is terminated when the maximum number of training iterations is reached. The same stopping criterion is applied across all methods to ensure fair comparison.

As well, the overall structural framework of the multi-task meta-adversarial mean-field IRL based on theory of mind is shown in Figure 1. As depicted in Figure 1, the original expert demonstration data is first combined with the data generated by the generator to create mixed expert demonstration data, from which M expert demonstration trajectories are sampled. The learning of rewards and policies under these M mixed expert demonstration trajectories is treated as M correlated tasks. Next, the states of mind are initialized, and an LSTM is employed to encode them, obtaining the task-specific state of mind. The averaged state of mind across other tasks is computed. Furthermore, a multi-task adversarial IRL framework is constructed. This framework extends adversarial IRL to average theory of mind and introduces probabilistic context variables to characterize task feature distributions. This enables simultaneous learning of multi-task reward functions and average state of mind. During training, the discriminator takes states $s_t$, actions $a_t$, and average state of mind $s_{avg}^{ToM}$ as inputs, and outputs both task-specific rewards ${\widehat{r}}_i\left(s_t,a_t,s_i^{ToM}\right)$ and predicted average state of mind ${\widehat{s}}_{avg}^{ToM}$. To achieve dynamic updates of state of mind, the average state of mind ${\widehat{s}}_{avg}^{ToM}$ is recomputed using the LSTM every K steps based on the most recent expert demonstration trajectory data. Simultaneously, the framework minimizes the discriminator’s adversarial loss function ${\mathcal{L}}_{AIAL}={\mathbb{E}}_{\xi \sim \pi }\left[\text{log }{D}_{\theta }\left(s_t,a_t,{\upsilon }_t\right)\right]+{\mathbb{E}}_{\xi \sim {\pi }^{*}}\left[\text{log }\left(1-D_{\theta }\left(s_t,a_t,{\upsilon }_t\right)\right)\right]+\lambda {∥{\widehat{s}}_{svg}^{ToM}-s_{svg}^{ToM}∥}^2$. During the training of the generator, the generator takes states $s_t$ and states of mind $s_i^{ToM}$ as inputs, and outputs actions $a_t^i\sim {\pi }_i\left(a_t^i|s_t,s_i^{ToM}\right)$, aiming to maximize the cumulative rewards ${\mathbb{E}}_{{\pi }_i}\left[{\sum }_t{r}_i\left(s_t,a_t,s_i^{ToM}\right)\right]$. Finally, the generator’s mixed data is utilized to update the mixed expert demonstration dataset, enhancing the quality of expert demonstrations and yielding superior demonstration data. Subsequently, reward learning is performed on the newly acquired expert demonstrations in preparation for the next training iteration, and this process continues until convergence is achieved.

Algorithm 1 Multi-task meta-adversarial mean-field IRL based on theory of mind

Input: Expert trajectories $D_E=\left\{{\xi }_j^E\right\}$, initial parameters of $r_{\theta }$, ${\pi }_{\theta }$, $q_{\psi }$   for iteration = $1,\dots ,N$ do do      Sample training tasks ${\mathcal{T}}_i\sim p\left(\mathcal{T}\right),i\in \left\{1,2,\dots ,M\right\}$;      for all tasks ${\mathcal{T}}_i$ do do        Obtain expert demos ${\mathcal{D}}_{E,{\mathcal{T}}_i}$;        for For update = $1,\dots ,K$ do do           Obtain generated demos ${\mathcal{D}}_{G,{\mathcal{T}}_i}$ and mixed demons ${\mathcal{D}}_{{\mathcal{T}}_i}$;           Infer state of mind and a batch of latent context variables from the sampled demonstrations: $s_{avg}^{ToM}$ and $\upsilon \sim q_{\psi }\left(\upsilon |\xi \right),\xi \in {\mathcal{D}}_{{\mathcal{T}}_i}$;           Update $\psi$ to increase ${\mathcal{K}}_{mut}\left(\theta ,\psi \right)$ with gradients with samples from $D_{{\mathcal{T}}_i}$;           Update $\theta$ to increase ${\mathcal{K}}_{mut}\left(\theta ,\psi \right)$ with gradients with samples from $D_{{\mathcal{T}}_i}$;           Update $\theta$ to minimize the loss function of the discriminator: ${\mathcal{L}}_{AIAL}={\mathbb{E}}_{\xi \sim \pi }\left[\text{log }{D}_{\theta }\left(s_t,a_t,{\widehat{\eta }}_{\psi ,t},{\upsilon }_t\right)\right]+{\mathbb{E}}_{\xi \sim {\pi }^{*}}\left[\text{log }\left(1-D_{\theta }\left(s_t,a_t,{\widehat{\eta }}_{\psi ,t},{\upsilon }_t\right)\right)\right]+\lambda {∥{\widehat{s}}_{avg}^{ToM}-s_{avg}^{ToM}∥}^2$;           Update $\omega$ with TRPO to increase the following objective:           ${\mathbb{E}}_{\left(s_t,a_t,{\widehat{\eta }}_{\psi ,t},{\upsilon }_t\right)\sim {\mathcal{D}}_{{\mathcal{T}}_i}}\left[\text{log }{D}_{\theta }\left(s_t,a_t,{\widehat{\eta }}_{\psi ,t},{\upsilon }_t\right)\right]$;        end for      end for   end for   Output: Learned inference model $q_{\psi }\left(\upsilon |\xi \right)$, reward function $r_{\theta }\left(s,a,\upsilon \right)$ and policies ${\pi }_{\theta }\left(a|s,\upsilon \right)$;

We analyze the computational complexity of the proposed TMMF-MTAIRL with respect to the number of tasks M, the trajectory length T (for $\xi ={\left\{(s_t,a_t)\right\}}_{t=1}^T$), and the number of inner updates K in Algorithm 1. In each update, for each task ${\mathcal{T}}_i$, we sample trajectories $\xi \in \mathcal{D}$ and perform: (i) context inference of the latent task via $q_{\psi }\left(\upsilon \right|\xi )$ (and the mutual-information-related updates), (ii) reward update of the discriminator $D_{\theta }$, and (iii) policy update for ${\pi }_{\omega }$. Let $c_{\psi }$, $c_{\theta }$, and $c_{\omega }$ denote the per-step (forward/backward) computational costs of $q_{\psi }$, $D_{\theta }$, and ${\pi }_{\omega }$, respectively. Then the per-task cost of one inner update scales as $\mathcal{O}\left(T(c_{\psi }+c_{\theta }+c_{\omega })\right)$, and the overall cost per outer iteration (over all M tasks and K inner updates) is $\mathcal{O}\left(MK·T(c_{\psi }+c_{\theta }+c_{\omega })\right)$. The key to reducing complexity through mean-field lies in the modeling approach for cross-task interactions. If traditional pair-wise task interaction is adopted, the interaction-related computations usually increase quadratically with the number of tasks, and the complexity is $\mathcal{O}\left(M^2\right)$. In contrast, our method aggregates the influence of other tasks through the averaged state of mind, so that the main task interacts only with this mean-field representation instead of all $M-1$ tasks individually. As a result, the interaction-related complexity is reduced from quadratic to linear, that is, $\mathcal{O}\left(M^{2}KT\right)\to \mathcal{O}\left(MKT\right)$. Moreover, since $s_{avg}^{ToM}$ is recomputed by the ToM encoder every K steps, the ToM recomputation cost is amortized over K inner updates, further improving practical efficiency. Therefore, the proposed algorithm reduces the computational complexity.

5. Experiments and Simulations

In this section, we use the point-maze environment and the rolling bearing fault diagnosis experiment to represent the performance of the proposed TMMF-MTAIRL algorithm for learning rewards and optimizing policies.

5.1. Point-Maze Environments

This environment is refactored from the D4RL repository, as introduced in [49]. The task involves a 2-DoF ball that is actuated by forces in the cartesian directions x and y, and must navigate toward a target goal within a closed maze. The environment can be initialized with different maze configurations of varying complexity. In the MuJoCo simulation, the target goal can be visualized as a red static ball, while the actuated ball is green. The control frequency of the ball is f = 10 Hz.

Table 2. Values of control parameters in data collection.

Control Parameters	Values
Number of iterations	2000
Batch size	20,000
Maximum path length	500
Discount factor	0.99
Step size	0.01

shows the values of control parameters used during the collection of the expert demonstrations. The values of control parameters used in the training process of TMMF-MTAIRL are listed in

Table 3. Values of control parameters in TMMF-MTAIRL.

Control Parameters	Values
Number of iterations	3000
Batch size	16
Meta batch size	50
Maximum path length	100
Discount factor	0.99
Step size	0.01
Number of pre-training epochs	1000
Entropy weight	1.0
Information coefficient	0.1

. Most hyperparameters listed in Table 2 and Table 3 were selected based on the successful configuration of similar environments and algorithms. Unless otherwise stated, we use the same set of hyperparameters across all tasks and environments. Empirically, we observe that the performance of TMMF-MTAIRL is robust to moderate variations of these parameters, which facilitates practical deployment.

For the convenience of experimental comparison, we make the following definitions. TaskRmean denotes the mean episodic rewards averaged over all tasks and runs; Mulinfo denotes the average mutual information between multiple tasks in Equation (13); MeanKL denotes the average KL divergence used for evaluating policy distribution shift across tasks; and dLoss denotes the discriminator’s loss function ${\mathcal{L}}_{AIAL}$ averaged over training iterations; Rmean denotes the mean episodic rewards averaged over all runs; times indicates the running times.

To verify the performance of the proposed algorithm, we compared TMMF-MTAIRL with PEMIRL, and Meta-IRLSOT. All experimental results were obtained from five independent runs. Meta-inverse reinforcement learning for mean-field games via probabilistic context variables (PEMIRL) [20]: This method provides a deep latent variable MFG-IRL model for IRL in mean-field games, which can infer reward functions from different but structurally similar expert demonstration tasks. Importantly, it does not require prior knowledge of the underlying context or modifications to the MFG model itself. A meta-inverse reinforcement learning method for fast adaptation of trajectory prediction networks (Meta-IRLSOT) [18]: This Meta-learning-based RL method incorporates an IRL framework to explore trajectory-scene relationships, achieving task-level scene understanding. By introducing meta-learning for co-training, the model enhances its generalization to novel scenarios.

Figure 2a–f show the visualization of the learned reward function using the proposed TMMF-MTAIR at iteration 0, 20, 40, 60, 80, 99, respectively. The horizontal and vertical axes now denote the spatial coordinates of the point-maze environment, and the colorbar indicates the magnitude of the learned reward values. From Figure 2, it can be seen that the positional intuitive schematic diagrams of the target goal and the agent point are obtained during the learning process of rewards. The target goal’s position changes dynamically, and the agent aims to reach it.

Figure 3a–f show the results of loss 1, loss 2, loss 3, policy likelihood, KL, total loss for PEMIRL, Meta-IRLSOT, and TMMF-MTAIRL in the maze left-v0 environment. From these results, we can obtain that our proposed TMMF-MTAIRL has lower loss 1, loss 2, loss 3, policy likelihood, KL, and total loss than those of PEMIRL and Meta-IRLSOT during the pre-training process in the point-maze left-v0 environment. These results demonstrate that TMMF-MTAIRL outperforms both PEMIRL and Meta-IRLSOT.

Table 4. The training data obtained in point-maze environments.

Environments	Methods	TaskRmean	Mulinfo	MeanKL	dLoss
Point-maze	PEMIRL	$-0.143\pm 0.022$	$-0.275\pm 0.026$	$0.106\pm 0.003$	$0.002\pm 0.001$
left-v0	Meta-IRLSOT	$-0.128\pm 0.010$	$-0.248\pm 0.016$	$0.317\pm 0.007$	$1.3\times {10}^{-5}\pm 0.000$
	TMMF-MTAIRL	$-0.125\pm 0.009$	$-0.112\pm 0.013$	$0.079\pm 0.002$	$0.0004\pm 0.000$
Point-maze	PEMIRL	$-0.171\pm 0.018$	$-0.420\pm 0.028$	$0.068\pm 0.004$	$0.0007\pm 0.000$
right-v0	Meta-IRLSOT	$-0.187\pm 0.013$	$-0.432\pm 0.020$	$0.037\pm 0.006$	$0.0009\pm 0.000$
	TMMF-MTAIRL	$-0.171\pm 0.008$	$-0.069\pm 0.019$	$0.029\pm 0.001$	$0.0002\pm 0.000$
Point-maze	PEMIRL	$-0.179\pm 0.021$	$-0.384\pm 0.024$	$0.105\pm 0.003$	$0.005\pm 0.001$
cont-v0	Meta-IRLSOT	$-0.195\pm 0.015$	$-0.367\pm 0.020$	$0.075\pm 0.005$	$0.007\pm 0.001$
	TMMF-MTAIRL	$-0.115\pm 0.004$	$-0.283\pm 0.017$	$0.083\pm 0.002$	$0.003\pm 0.000$

shows the training data obtained in point-maze left-v0, point-maze right-v0, and point-maze cont-0 for TMMF-MTAIRL, PEMIRL, and Meta-IRLSOT. From the results of point-maze left-v0 and point-maze right-v0, it can be seen that our proposed TMMF-MTAIRL has higher TaskRmean and Mulinfo, while exhibiting lower MeanKL, and dLoss, than PEMIRL and Meta-IRLSOT. Although some metrics obtained by our proposed TMMF-MTAIRL algorithm are not optimal, overall performance analysis indicates that it outperforms the other algorithms. To further illustrate the effectiveness of the proposed algorithm, we also collected the testing data in point-maze cont-v0 using TMMF-MTAIRL, PEMIRL, and Meta-IRLSOT to test invisible target points. The results of Rmean, TaskRmean, Mulinfo, MeanKL, and dLoss indicate that our proposed TMMF-MTAIRL consistently performs better than the other algorithms during testing.

Figure 4a,b illustrate that TMMF-MTAIRL achieves higher Rmean and lower times than PEMIRL, Meta-IRLSOT, and TMMF-MTAIRL without mean field in point-maze environments. Therefore, TMMF-MTAIRL demonstrates superior performance relative to the other algorithms used in this paper.

5.2. Rolling Bearing Fault Diagnosis Experiments

For comprehensive performance validation of the proposed method, the TMMF-MTAIRL algorithm is first used to learn the reward function, while simultaneously optimizing expert demonstrations throughout the training phase. Subsequently, RL is applied to classify bearing fault diagnosis data based on the learned reward function. All experimental results were obtained by averaging over 10 independent runs.

Figure 5 shows the rolling bearing fault test rig from Case Western Reserve University, which consists of a fan-end bearing, a motor, a drive-end bearing, a torque transducer/encoder, and a power meter [50]. The experiment collected normal bearing signals and (single-point) drive-end fault signals at sampling rates of 12,000 samples/second and 48,000 samples/second, respectively, whereas all fan-end fault signals were acquired at 12,000 samples/second. To validate the proposed method, the acquired vibration data were divided into four datasets (A, B, C, and D). Each dataset contains 10 distinct classes, with no unseen classes among these four datasets. On average, each class within a given dataset contains approximately 480,000 samples, indicating that the laboratory equipment was operated under each class condition to capture an average of 480,000 observations from vibration sensors.

In the fault classification process, the design of the reward function plays a crucial role in determining the accuracy of fault diagnosis. Therefore, we first employ TMMF-MTAIRL to learn the reward function from expert demonstrations, and then utilize the deep Q-Learning (DQN) algorithm to perform fault diagnosis and classification. The main framework for fault diagnosis based on the TMMF-MTAIRL and DQN algorithms is illustrated in Figure 6. First, the dataset is preprocessed using cross-validation to obtain the training set and validation set. Subsequently, a stacked autoencoder (SAE) is employed to extract the latent features of the data while simultaneously updating and preserving optimal parameters. We then select a portion of the training data as expert demonstration data and sample expert demonstration trajectories from it. Based on the varying preferences among these expert demonstration trajectories, the proposed TMMF-MTAIRL algorithm is employed to learn rewards and optimize the expert demonstration data, treating each trajectory’s reward learning as an independent task. Finally, a DQN algorithm that incorporates Bayesian search methodology and learned reward functions is implemented to perform fault classification.

Expert demonstrations $\mathcal{D}$ are constructed by sampling labeled trajectories $\xi$ from the training set, which refer to the high-quality labeled vibration samples collected under controlled operating conditions from the Case Western Reserve University bearing test rig. A trajectory corresponds to a sequential decision process formed by sliding-window segmentation of a vibration time series. Therefore, a trajectory $\xi ={\left\{(s_t,a_t)\right\}}_{t=1}^T$ represents a sequence of vibration windows and predicted diagnostic decisions. Additionally, each fault type can be viewed as a task because it defines a distinct diagnostic objective and induces a different state distribution.

In the bearing fault diagnosis experiment, to evaluate the performance of the TMMF-MTAIRL algorithm, the proposed TMMF-MTAIRL integrated with cross-validated Bayesian DQN (DQNB-TMMF-MTAIRL-5) was compared with the following four algorithms: (1) traditional deep Q-Learning algorithm (DQN) [51]; (2) Bayesian-based deep Q-Learning algorithm (DQN-Bayesian) [52]; (3) the algorithm combining TMMF-MTAIRL with Bayesian DQN (DQNB-TMMF-MTAIRL); and (4) a deep Q-learning algorithm incorporating Bayesian inference and cross-validation (DQN-Bayesian-5).

To further demonstrate the performance of the proposed algorithm, DQN, DQN-Bayesian, DQNB-TMMF-MTAIRL, DQN-Bayesian-5, and DQNB-TMMF-MTAIRL-5 are applied to fault diagnosis using the bearing fault diagnosis dataset B.

During the training process, the complexity of classification data and the design of reward function constrain the accuracy of fault classification using RL. Therefore, we analyze the cumulative rewards obtained during training for fault classification using DQN, DQN-Bayesian, DQNB-TMMF-MTAIRL, DQN-Bayesian-5, and DQNB-TMMF-MTAIRL-5, as shown in Figure 7. During the training process for fault diagnosis, the cumulative rewards learned by DQNB-TMMF-MTAIRL and DQNB-TMMF-MTAIRL-5 converge to 200 more rapidly than those achieved by DQN, DQN-Bayesian, and DQN-Bayesian-5 with manually designed reward functions. To further enhance model performance, the bearing fault diagnosis data was split into five folds via five-fold cross-validation, where four folds were used for training and one fold for validation in each iteration. As a result, DQNB-TMMF-MTAIRL-5 demonstrates the fastest convergence of cumulative rewards, indicating its superior capability in learning reward functions.

Figure 8 shows the accuracy of fault classification during the training process for the DQN, DQN-Bayesian, DQNB-TMMF-MTAIRL, DQN-Bayesian-5, and DQNB-TMMF-MTAIRL-5. Figure 9 shows the validation accuracy of the fault classification obtained using DQN-Beyasian, DQNB-TMMF-MTAIRL, DQN-Beyasian-5, and DQNB-TMMF-MTAIRL-5 algorithms. The proposed TMMF-MTAIRL, which explicitly models inter-task relationships, learns reward functions that mitigate the difficulty of manual reward design in RL. In Figure 8 and Figure 9, compared with DQN, DQN-Bayesian, DQN-Bayesian-5, and DQNB-AMEIRL-5, DQNB-TMMF-MTAIRL and DQNB-TMMF-MTAIRL-5 achieve better fault classification, resulting in higher accuracy in both fault classification and verification. The comparison with DQN and DQN-Bayesian is intended to demonstrate the advantage of learned rewards over manually engineered rewards in RL-based fault diagnosis. Furthermore, to improve the utility of bearing fault data, we adopted five-fold cross-validation for data preprocessing. This approach maximizes data utilization, reduces the risks of both overlearning and underlearning, and enhances the credibility of experimental findings. DQNB-TMMF-MTAIRL-5 without mean-field algorithm suffers from the lack of mean field, resulting in low exploration efficiency. Therefore, DQNB-TMMF-MTAIRL-5 algorithm achieves the highest fault classification accuracy on both the test set and the verification set.

To verify the robustness of the proposed TMMF-MTAIRL, we obtained the accuracy of fault classification as 0.97, 0.98, and 0.98, respectively when the multi-task scale M was 10, 15, and 20. The accuracy of this algorithm remains basically stable at around 0.97∼0.98 as the scale of multiple tasks changes, which indicates that TMMF-MTAIRL has good robustness.

The proposed TMMF-MTAIRL is optimized via an adversarial minimax objective. Under standard assumptions commonly adopted in adversarial IRL—such as sufficient capacity of discriminator network, accurate inner-loop optimization, and appropriate learning rates—the training dynamics admit stationary points where the induced policy distribution matches the expert trajectory distribution. Furthermore, the mutual-information regularization and mean-field approximation improve stability by preventing degenerate latent inference and reducing high-variance cross-task interactions. In addition, the results in Figure 7, Figure 8 and Figure 9 show that the proposed DQNB-TMMF-MTAIRL can converge. Therefore, the proposed framework enjoys conditional convergence properties and demonstrates stable empirical behavior.

From an application perspective, it is also important to distinguish the proposed approach from alternative fault diagnosis paradigms. Deep transfer learning methods for fault diagnosis typically rely on labeled data and supervised objectives, achieving strong performance when sufficient annotations are available. In contrast, the proposed TMMF-MTAIRL focuses on learning reward functions and policies from heterogeneous demonstrations without requiring explicit fault labels. These two paradigms address complementary problem settings, and integrating transfer learning representations into the proposed IRL framework is a promising direction for future research.

The proposed TMMF-MTAIRL can explicitly model the relationship between tasks and thus improve the explainability and adaptability across tasks, enabling the algorithm to learn more effective rewards and optimization policies compared to the other algorithms discussed in this paper.

6. Conclusions and Future Works

To overcome the challenges of learning rewards in the complex multi-task environment, a multi-task meta-adversarial mean-field IRL framework based on the theory of mind is proposed to enhance information interaction between partially heterogeneous tasks. In TMMF-MTAIRL, the mean-field theory reduces the computational complexity of learning rewards from fixed expert demonstrations in multi-task settings; the theory of mind improves the information sharing and transferability between multiple tasks, and the deep potential variables enhance the adaptability to novel tasks. These improvements jointly increase learning accuracy and reduce computational complexity in learning reward functions. Experiments in a maze-point environment and a rolling bearing fault diagnosis task demonstrate that TMMF-MTAIRL can learn higher-quality reward functions and policies through continuous iterative updates, thereby outperforming the original expert demonstrations.

While the experimental evaluation focuses on representative benchmark environments and a real-world fault diagnosis dataset, the proposed framework is not limited to these settings. The modular design of the ToM encoder, mean-field approximation, and adversarial IRL objective allows straightforward extension to other domains. Evaluating the method on a wider range of environments and large-scale benchmarks will be explored in future work. In the future work, we will focus on the following three aspects: (1) Extending the proposed TMMF-MTAIRL algorithm to real-world environments for policy optimization, beyond the point-maze and rolling bearing fault experiments used for initial validation; (2) Further enhancing the reward learning performance of TMMF-MTAIRL in ultra-large-scale multi-agent and highly complex multi-task environments; and (3) Deep transfer learning methods and the proposed TMMF-MTAIRL address complementary problem settings, and integrating transfer learning representations into the proposed IRL framework is a promising direction for future research.