SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Guanren Qiao $^{1}$ , Guiliang Liu $^{*1}$ , Pascal Poupart $^{2,3}$ , Zhiqiang Xu $^{4}$
$^{1}$ School of Data Science, The Chinese University of Hong Kong, Shenzhen,
$^{2}$ University of Waterloo,  $^{3}$ Vector Institute,  $^{4}$ Mohamed bin Zayed University of Artificial Intelligence
223040241@link.cuhk.edu.cn, liuguiliang@cuhk.edu.cn,
ppoupart@uwaterloo.ca,zhiqiang.xu@mbzuai.ac.ae
Inverse Constraint Reinforcement Learning (ICRL) aims to recover the underlying constraints respected by expert agents in a data-driven manner. Existing ICRL algorithms typically assume that the demonstration data is generated by a single type of expert. However, in practice, demonstrations often comprise a mixture of trajectories collected from various expert agents respecting different constraints, making it challenging to explain expert behaviors with a unified constraint function. To tackle this issue, we propose a Multi-Modal Inverse Constrained Reinforcement Learning (MMICRL) algorithm for simultaneously estimating multiple constraints corresponding to different types of experts. MMICRL constructs a flow-based density estimator that enables unsupervised expert identification from demonstrations, so as to infer the agent-specific constraints. Following these constraints, MMICRL imitates expert policies with a novel multi-modal constrained policy optimization objective that minimizes the agent-conditioned policy entropy and maximizes the unconditioned one. To enhance robustness, we incorporate this objective into the contrastive learning framework. This approach enables imitation policies to capture the diversity of behaviors among expert agents. Extensive experiments in both discrete and continuous environments show that MMICRL outperforms other baselines in terms of constraint recovery and control performance. Our implementation is available at: https://github.com/qiaoguanren/Multi-Modal-Inverse-Constrained-Reinforcement-Learning.
A fundamental prerequisite for achieving safe Reinforcement Learning (RL) is that the agents' policies must adhere to the underlying constraints in the environment [1, 2]. However, in many real-world applications (e.g., robot control and autonomous driving), the ideal constraints are time-varying, context-dependent, and inherent to experts' own experience. These constraints are hard to specify mathematically and may not be readily available to RL agents in policy updates.
Figure 1: The flowchart of MMICRL.
A promising approach for learning latent constraints is Inverse Constraint Reinforcement
Learning (ICRL) [3, 4, 5]. As a data-driven technique, ICRL recovers the underlying constraints
*Corresponding author: Guiliang Liu, email: liuguiliang@cuhk.edu.cn
respected by expert agents from their demonstrations and utilizes these constraints to support downstream applications. Existing ICRL methods [3, 4, 6, 7, 8] commonly assumed that all expert demonstrations follow the same constraints, and they approximated these constraints with a unified constraint model, whereas, in practice, the demonstration data may be collected from various agents, and these agents might follow different or even conflicting constraints. It is problematic to leverage a single constraint model to explain the behaviors of diverse agents. For example, in the context of autonomous driving [9], the vehicle-distance constraints followed by trucks and cars should differ, and misapplying these constraints could lead to serious traffic accidents.
To differentiate expert demonstrations, previous RL approaches [10, 11, 12, 13] inferred the latent structure of expert demonstrations and identified expert agents by examining their behavioral preferences. However, these methods were specifically designed for imitation learning rather than constraint inference. Moreover, as their identification process primarily relies on an agent classifier that evaluates state-action pairs, there is no theoretical guarantee that the optimal model is identifiable.
In this work, we propose the Multi-Modal Inverse Constrained Reinforcement Learning (MMICRL) algorithm for estimating agent-specific constraints from a mixture of expert demonstrations (Figure 1). MMICRL extends the traditional maximum entropy framework [14] with agent-conditioned entropy minimization and unconditioned entropy maximization subject to a permissibility constraint. The resulting policy representation facilitates inferring agent-specific constraints by alternating between the following steps: 1) Unsupervised Agent Identification. MMICRL conducts trajectory-based agent identification using a flow-based density estimator. The optimal results are identifiable since each agent's policy must correspond to a unique occupancy measure [15]. 2) Agent-Specific Constraint Inference. Leveraging the identified demonstrations, MMICRL estimates a permissibility function for distinguishing expert demonstrations from sub-optimal ones, based on which we construct constraints for each type of agent. 3) Multi-Modal Policy Optimization. MMICRL measures the accuracy of inferred constraint models by comparing the similarity between expert trajectories and those generated by imitation policies under the inferred constraints. To capture the diversity of behaviors exhibited by multiple agents, we incorporate policy optimization within the contrastive learning framework. We treat the generated trajectories as noisy embeddings of agents, which serve as compact representations of their behaviors. Utilizing the contrastive estimation methods [16], we can enhance the similarity between embeddings for agents of the same type, while simultaneously maintaining the distinctiveness of diverse agent types.
We empirically demonstrate the performance of our method by conducting experiments in both discrete (e.g., Gridworld) [7] and continuous environments (e.g., MuJoCo) [17]. MMICRL significantly surpasses other baselines in terms of distinguishing different agents, adhering to the true constraints, and optimizing control performance. To examine the robustness of MMICRL, we investigate its ability to recover from incorrect agent identification and subsequently infer the correct constraints.
Inverse Constrained Reinforcement Learning. Prior ICRL methods typically learned constraints from demonstrations under the maximum entropy framework [14]. Some research [18, 19] employed constraints to differentiate between feasible and infeasible state-action pairs in Constrained MDP, but these studies were restricted to inferring discrete constraints in environments with known dynamics. A subsequent work [4] extended this approach to continuous state-action spaces with unknown transition models by utilizing neural networks to approximate constraints. Inspired by Bayesian Inverse Reinforcement Learning [20, 21, 22], [8] inferred probability distributions over constraints. To better model demonstrations, [6] extended ICRL to infer soft constraints rather than hard ones, and [23] explored ICRL under the multi-agent setting. Striving for efficient comparisons, [7] established an ICRL benchmark across various RL domains, such as Gridworld, robot control, and autonomous driving. However, these algorithms primarily target inferring constraints for a single agent type, without considering the distinct constraints associated with multiple agent types.
Learning from a Mixture of Expert Demonstrations. Multi-task Inverse Reinforcement Learning (IRL) [24, 25] aims to learn from a mixture of expert demonstrations. Some previous studies [10, 26, 12, 27] utilized the Generative Adversarial Imitation Learning algorithm [28] to model the behaviors of multiple agents. [10, 11, 13, 29] learned interpretable representations of behavioral policies and inferred the latent representation of expert demonstrations in an unsupervised way. Specifically, [30] learned a Variational Auto-Encoder (VAE) [31] where the encoder infers the latent
factors of variation from mixed demonstrations and the decoder models different types of expert behaviors. To strengthen the consistency between the learned policy and the types of agents, [12, 32] included "burn-in demonstrations" for updating the imitation policies. These methods, however, were proposed for imitation learning or rewards recovery (i.e., for IRL) [33] instead of constraint inference.
Constrained Mixture-Agent Markov Decision Process. To support constraint inference for different agents, we formulate the environment as a Constrained Mixture-Agent Markov Decision Process (CMA-MDP) [34]  $\mathcal{M}^{\phi}$ , which can be defined by a tuple  $(S, \mathcal{A}, \mathcal{Z}, \mathrm{R}, p_{\mathcal{T}}, \{(p_{\mathcal{C}_i}, \epsilon_i)\}_{\forall i, \gamma, \mu_0})$  where: 1)  $S$ ,  $\mathcal{A}$  and  $\mathcal{Z}$  denote the space of states, actions, and latent code (for specifying expert agents). 2)  $p_{\mathcal{T}}(s'|s, a)$  and  $\mathrm{R}(s, a)$  define the transition and reward functions. 3)  $p_{\mathcal{C}_j}(c|s, a, z)$  refers to an agent-specific cost model with an associated bound  $\epsilon_j(z)$ , where  $j$  indicates the index of a constraint. 4)  $\gamma$  is the discount factor and  $\mu_0(s)$  defines the initial state distribution. CMA-MDP assumes the agents are differentiable by examining their policies, implying that different agent types cannot share identical policies. We aim to elucidate these differences with the constraints associated with each agent type.
In contrast to Multi-Agent Reinforcement Learning (MARL) [35], where multiple agents can act concurrently, CMA-MDP allows only one agent to operate at a given time. Nevertheless, since the agents in CMA-MDP might adhere to distinct constraints, they can develop different optimal control policies or strategies. However, standard MARL frameworks do not explicitly differentiate between agent types, nor do they distinguish their respective policies and constraints.
Policy Update under Conditional Constraints. We introduce Constrained Reinforcement Learning (CRL) based on CMA-MDP. For an agent identified by  $z$ , the goal of CRL is to find a policy  $\pi^{*}(a|s,z)$  that maximizes expected discounted rewards under the conditional constraints:
$$
\mathcal {J} (\pi | z) = \max  _ {\pi} \mathbb {E} _ {\mu_ {0}, p _ {\mathcal {T}}, \pi} \left[ \sum_ {t = 0} ^ {T} \gamma^ {t} r _ {t} \right] + \beta \mathcal {H} (\pi) \text {s . t .} \mathbb {E} _ {\tau \sim \left(\mu_ {0}, p _ {\mathcal {T}}, \pi\right), p _ {\mathcal {C} _ {j}}} \left[ c _ {j} (\tau | z) \right] \leq \epsilon_ {j} (z) \forall j \in [ 0, J ] \tag {1}
$$
where  $\mathcal{H}(\pi)$  denotes the policy entropy weighted by  $\beta$ , and we follow [4] to define the trajectory cost  $c(\tau |z) = 1 - \prod_{(s,a)\in \tau}\phi (s,a|z)$  where the permissibility function  $\phi (s,a|z)$  indicates the probability that performing action  $a$  under a state  $s$  is safe for agent  $z$ . CRL commonly assumes that the constraints are known, whereas in practice, instead of directly observing the constraint signals, we often have access to expert demonstrations that follow the underlying constraints, and thus the agent must recover the constraints from the demonstration dataset.
Constraint Inference from a Mixture of Expert Dataset. Based on the CMA-MDP, the goal of ICRL is to discover the underlying constraints respected by different types of expert agents from their demonstrations. To achieve it, we consider the Maximum Entropy framework [4, 14] and represent the likelihood function as follows:
$$
p \left(\mathcal {D} _ {E} | \phi\right) = \frac {1}{\left(Z _ {\mathcal {M} ^ {\hat {\epsilon} _ {\phi}}}\right) ^ {N}} \prod_ {i = 1} ^ {N} \sum_ {z} \mathbb {1} _ {\tau^ {i} \in \mathcal {D} _ {z}} \exp \left[ r \left(\tau^ {(i)}\right) \right] \mathbb {1} ^ {\mathcal {M} ^ {\hat {\epsilon} _ {\phi}}} (\tau^ {(i)} | z) \tag {2}
$$
where 1)  $N$  denotes the number of trajectories in the demonstration dataset, 2)  $Z_{\mathcal{M}^{\hat{c}_{\phi}}}$  is a normalizing term, 3) the permissibility indicator  $\mathbb{1}^{\hat{\mathcal{M}}_{\phi}}(\tau^{(i)})$  can be defined by  $\phi (\tau^{(i)}|z) = \prod_{t = 1}^{T}\phi_{t}(s_{t}^{i},a_{t}^{i}|z)$ , and 4) the agent identifier  $\mathbb{1}_{\tau^{(i)}\in \mathcal{D}_z}$  determines whether the trajectory  $\tau^{(i)}$  is generated by the agent  $z$ .
By following [4], constraint inference can be formulated as inferring  $\phi_t$  by maximizing this likelihood function. A common approach is parameterizing  $\phi_t$  with neural network models and updating the model parameters with the agent-specific expert data [4, 6, 7]. However, the major challenge of this task lies in our lack of knowledge about the agent's identity that generates the expert trajectories (i.e.,  $\mathbb{1}_{\tau^{(i)} \in \mathcal{D}_z}$  is unknown). For example, a vehicle trajectory dataset does not label the specific type of vehicles, or these labels are coarse-grained, incomplete, and noisy. On the other hand, to explain the optimal policy with the learned constraints, the types of agents and the constraints must be consistent. If we evaluate the policy of an agent  $z$  under the constraint for another type of agent  $z'$  (e.g., evaluate the driving strategy of a truck under the constraint for a car), the performance will be substantially compromised. As a result, for a specific agent  $z$ , the expert trajectories  $\{\tau_z^*\}$  are optimal while the trajectories generated by other experts  $\{\tau_{z'}^*, \tau_{z''}^*, \tau_{z'''', \ldots}\}$  are supposed to be sub-optimal
or even infeasible. To capture the correct constraints, the ICLR algorithm must identify the agent corresponding to each trajectory in an unsupervised manner.
4 Inverse Constrained Reinforcement Learning for a Mixture of Experts
In order to infer constraints for various expert agent types, we introduce the Multi-Modal Inverse Constrained Reinforcement Learning (MMICRL) algorithm (Figure 1). MMICRL employs a conditional imitation policy to capture the diverse agent behaviors, facilitating unsupervised agent identification. Technically, we require the agents' policies to 1) exhibit high entropy when the agent type (specified by  $z$ ) is unknown, and 2) collapse to a specific behavioral mode when the type is determined. The objective for MMICRL can be expressed as:
$$
\text {M i n i m i z e} - \alpha_ {1} \mathcal {H} [ \pi (\tau) ] + \alpha_ {2} \mathcal {H} [ \pi (\tau | z) ] \tag {3}
$$
$$
\text {S u b j e c t} \int \pi (\tau | z) f _ {z} (\tau) \mathrm {d} \tau = \frac {1}{N} \sum_ {\tau \in \mathcal {D} _ {z}} f (\tau), \int \pi (\tau | z) \mathrm {d} \tau = 1, \text {a n d} \int \pi (\tau | z) \log \phi (\tau | z) \mathrm {d} \tau \geq \epsilon
$$
(Proof is provided in Appendix B.) where  $f(\cdot)$  represents a latent feature extractor,  $\mathcal{H}[\pi(\tau|z)]$  denotes the agent-specific policy entropy, and  $\mathcal{H}[\pi(\tau)]$  signifies the entropy without knowledge of the agent type. The weighting parameters  $\alpha_{1}$  and  $\alpha_{2}$  (both  $\geq 0$ ) determine the balance between conditional entropy minimization and general entropy maximization. This objective differs from the traditional Maximum Entropy Inverse Reinforcement Learning (MEntIRL) [14] in two ways: 1) the objective also minimizes an entropy conditioned on agent types, and 2) it incorporates an additional constraint related to the policy's permissibility (the last constraint). Given this objective, the optimal representation for the trajectory likelihood (i.e., trajectory policy) is:
Proposition 4.1. Let  $p(z|\tau)$  denote the trajectory-level agent identifier, let  $r(\tau) = \frac{\lambda_0}{\alpha_2 - \alpha_1} f(\tau)$  denote the trajectory rewards, let  $Z_{\mathcal{M}^{\hat{\epsilon}\phi}}$  denote a normalizing term. The optimal policy of the above optimization problem can be represented as:
$$
\pi (\tau | z) = \frac {1}{Z _ {\mathcal {M} ^ {\hat {c}} \phi}} \exp \left[ \frac {\alpha_ {1} \mathbb {E} _ {z \sim p (z)} [ \log (p (z | \tau)) ]}{\alpha_ {2} - \alpha_ {1}} + r (\tau) \right] \phi (\tau | z) ^ {\frac {\lambda_ {2}}{\alpha_ {1} - \alpha_ {2}}} \tag {4}
$$
Building upon this policy representation (4), we present the key steps of MMICRL, which involves iteratively performing: 1) unsupervised agent identification for calculating  $p(z|\tau)$  (Section 4.1), 2) conditional inverse constraint inference (Section 4.2) for deducing  $\phi(\tau|z)$ , and 3) multi-modal policy update (Section 4.3) for approximating  $\pi(\tau|z)$ . MMICRL alternates between these steps until the imitation policies reproduce expert trajectories, signifying that the inferred constraints align with the ground-truth constraints.
4.1 Unsupervised Agent Identification
MMICRL identifies expert trajectories (i.e., learning  $p(z|\tau)$ ) in an unsupervised manner. Previous works commonly determined the agents' identities by examining the state-action features [10, 11, 12] with the classifier  $p(z|s_z, a_z)$ . Nevertheless, different agents may exhibit similar behaviors under certain contexts or at specific time steps within a trajectory, which makes this point-wise identification problematic (e.g., see our experimental results of InfoGAIL [10] in Section 5).
To derive a reliable agent identifier, MMICRL performs trajectory-level identification by estimating an agent-specific density. Specifically, we define a state-action density (i.e., normalized occupancy measure)  $\rho_{\pi}(s,a) = (1 - \gamma)\pi (a|s)\sum_{t = 0}^{\infty}\gamma^{t}p(s_{t} = s|\pi)$  where  $p(s_{t} = s|\pi)$  is the probability density of state  $s$  at time step  $t$  following policy  $\pi$ . Based on this density, we consider the theorem:
Proposition 4.2 (Theorem 2 of [15]). Suppose  $\rho$  is the occupancy measure that satisfies the Bellman flow constraints:  $\sum_{a}\rho (s,a) = \mu_{0}(s) + \gamma \sum_{s^{\prime},a}\rho (s^{\prime},a)P\tau (s^{\prime}|s,a)$  and  $\rho (s,a)\geq 0$ . Let the policy defined by:  $\pi_{\rho}(a|s) = \frac{\rho(s,a)}{\int\rho(s,a')\mathrm{d}a'}$ , then  $\pi_{\rho}$  is the only policy whose occupancy measure is  $\rho$ .
Density Estimation. The aforementioned theorem provides a crucial insight for agent identification: "one can identify an expert agent by examining the occupancy measures in the expert trajectories". Leveraging this insight, we design a state-action density estimator to compute a density using a Conditional Flow-based Density Estimator (CFDE). CFDE estimates the density of input variables in the training data distribution under an auto-regressive constraint. Moreover, to enhance our
density estimator's sensitivity to the behavior of different agents, the estimator also conditions on the agent type. so  $p(\pmb {x}|z) = \prod_{i}p(\pmb {x}_i|\pmb{x}_{1:i - 1},z)$  where  $\pmb {x}\coloneqq (s,a)$  defines an event. We implement  $p(\pmb {x}_i|\pmb{x}_{1:i - 1},z) = \mathcal{N}(\pmb {x}_i|\mu_i,(\exp (\alpha_i))^2)$  where  $\mu_{i} = \psi_{\mu_{i}}(\pmb{x}_{1:i - 1},z)$  and  $\alpha_{i} = \psi_{\alpha_{i}}(\pmb{x}_{1:i - 1},z)$ . The neural function  $\psi$  is implemented by stacking multiple MADE layers [36]. The corresponding agent identifier can be represented by Bayesian rule, so  $p_{\psi}(z|\tau) = \frac{p(z)\cdot p_{\psi}(\tau|z)}{\sum_zp(z)\cdot p_{\psi}(\tau|z)}$ . We assume a uniform prior  $p(z)$ , thereby deriving the following form:  $p_{\psi}(z|\tau) = \frac{\prod_{(s,a)\in\tau}p_{\psi}(s,a|z)}{\sum_{z'}\prod_{(s,a)\in\tau}p_{\psi}(s,a|z')}$ .
Agent Identification. After learning the density model  $p_{\psi}(s,a|z)$  with CFDE, we divide  $\mathcal{D}_E$  into sub-datasets  $\{\mathcal{D}_z\}_{z=1}^{|\mathcal{Z}|}$  by: 1) initializing the dataset  $\mathcal{D}_z = \emptyset$  and 2)  $\forall \tau_i \in \mathcal{D}_E$ , adding  $\tau^i$  into  $\mathcal{D}_z$  if  $z = \arg \max_z \sum_{(s,a) \in \tau_i} \log[p_{\psi}(s,a|z)]$ . We repeat the above steps for all  $z \in \mathcal{Z}$ .
4.2 Agent-Specific Constraint Inference
Based on the identified expert dataset, we have  $p_{\omega}(z|\tau_z) = 1, \forall \tau_z \in \mathcal{D}_z$ , and thus  $\log p_{\omega}(z|\tau_z) = 0$  so the likelihood function (4) can be simplified to:
$$
p \left(\mathcal {D} _ {z} \mid \phi , z\right) = \prod_ {i = 1} ^ {N} \frac {1}{Z _ {\mathcal {M} ^ {\hat {\epsilon}} \phi}} \exp \left[ r \left(\tau_ {z} ^ {(i)}\right) \right] \phi \left(\tau_ {z} ^ {(i)}\right) ^ {\frac {\lambda_ {2}}{\alpha_ {1} - \alpha_ {2}}} \tag {5}
$$
where  $\phi (\tau_z^{(i)}) = \prod_{(s,a)\in \tau_z^{(i)}}\phi_t(s,a|z)$  and the normalizer  $Z_{\mathcal{M}^c\phi} = \int \exp [r(\tau)]\phi (\tau)^{\frac{\lambda_2}{\alpha_1 - \alpha_2}}\mathrm{d}\tau$  . By defining  $\eta = \frac{\lambda_2}{\alpha_1 - \alpha_2}$  , we can then define the log-likelihood  $\log [p(\mathcal{D}_z|\phi ,z)]$  as:
$$
\sum_ {i = 1} ^ {N} \left[ r \left(\tau_ {z} ^ {(i)}\right) + \eta \log \prod_ {t = 0} ^ {T} \phi \left(s _ {t} ^ {(i)}, a _ {t} ^ {(i)} | z\right) \right] - N \log \int \exp [ r (\hat {\tau}) ] \left[ \prod_ {t = 0} ^ {T} \phi \left(\hat {s} _ {t}, \hat {a} _ {t} | z\right) \right] ^ {\eta} \mathrm {d} \hat {\tau} \tag {6}
$$
We parameterize the instantaneous permissibility function with  $\omega$ , i.e., construct  $\phi_{\omega}(s_t, a_t | z)$  and update the parameters by computing the gradient of the above likelihood function (the derivation resembles that of [4]), so  $\nabla_{\omega} \log [p(\mathcal{D}_z | \phi, z)]$  can be defined as:
$$
\sum_ {i = 1} ^ {N} \left[ \nabla_ {\phi} \sum_ {t = 0} ^ {T} \eta \log \left[ \phi_ {\omega} \left(s _ {t} ^ {(i)}, a _ {t} ^ {(i)} \mid z\right) \right] \right] - N \mathbb {E} _ {\hat {\tau} \sim \pi_ {\mathcal {M} ^ {\phi} (\cdot | z)}} \left[ \nabla_ {\phi} \sum_ {t = 0} ^ {T} \eta \log \left[ \phi_ {\omega} \left(\hat {s} _ {t}, \hat {a} _ {t} \mid z\right) \right] \right] \tag {7}
$$
This inverse constraint objective relies on the nominal trajectories  $\hat{\tau}$  sampled with the conditional policy  $\pi_{\mathcal{M}\hat{\phi}}(\tau |z)$  (also see Figure 1). For simplicity, we denote it as  $\pi (\tau |z)$ . In the following, we will introduce our approach to learning  $\pi (\tau |z)$ .
4.3 Multi-Modal Policy Optimization
By definition, the policy  $\hat{\pi} (\tau |z)$  is trained to maximize cumulative rewards subject to constraint  $\mathbb{E}_{\tau \sim \pi (\cdot |z)}[\log \phi (\tau |z)]\geq \epsilon$  . To be consistent with our MMICRL objective in formula (3), we design the multi-modal policy optimization objective in the following:
$$
\min  _ {\pi} - \mathbb {E} _ {\pi (\cdot | z)} \left[ \sum_ {t = 0} ^ {T} \gamma^ {t} r \left(s _ {t}, a _ {t}\right) \right] - \alpha_ {1} \mathcal {H} [ \pi (\tau) ] + \alpha_ {2} \mathcal {H} [ \pi (\tau | z) ] \text {s . t .} \mathbb {E} _ {\pi (\cdot | z)} \left(\sum_ {t = 0} ^ {h} \gamma^ {t} \log \phi_ {\omega} (s, a, z)\right) \geq \epsilon
$$
This objective extends maximum entropy policy optimization by minimizing an additional agent-conditioned entropy  $\mathcal{H}[\pi (\tau |z)]$ , which limits the variance of policy distribution for a specific type of agent. The balance between these entropy terms is controlled by  $\alpha_{1}$  and  $\alpha_{2}$ . Since  $\mathcal{H}[\pi (\tau)] = \mathcal{H}[\pi (\tau |z)] + \mathbb{E}_{z\sim p(z),\tau \sim \pi (\tau |z)}\Bigl [\log [p_{\psi}(z|\tau)]\Bigr ] + \mathcal{H}(z)$ , and by removing the uniform prior  $p(z)$  (since it is independent of the policy update), the objective (8) can be simplified to:
$$
\min  _ {\pi} - \mathbb {E} _ {\pi (\cdot | z)} \left[ r (\tau) + \alpha_ {1} \log [ p _ {\psi} (z | \tau) ] \right] + (\alpha_ {2} - \alpha_ {1}) \mathcal {H} [ \pi (\tau | z) ] \text {s . t .} \mathbb {E} _ {\pi (\cdot | z)} \left(\sum_ {t = 0} ^ {h} \gamma^ {t} \log \phi_ {\omega} (s, a, z)\right) \geq \epsilon
$$
Intuitively, this objective expands the reward signals with a log-probability term  $\log[p_{\psi}(z|\tau)]$ , which encourages the policy to generate trajectories from high-density regions for a specific agent type. This approach ensures that the learned policies  $\pi(\cdot | z)_{z=1}^{Z}$  are differentiable.
Learning Diverse Policies via Contrastive Estimation. In practice, directly augmenting the reward with a log-probability term (as in objective 9) may lead to a sub-optimal policy [7]. This issue