SlowGuess/NeurIPS_2023_OCR huggingface dataset:
arises because the log-probability term assigns a large penalty to trajectories with low  $p_{\psi}(z|\tau)$ . In such cases, the controlling policy becomes more sensitive to density estimation (since  $p_{\psi}(z|\tau) = \mathrm{softmax}[p_{\psi}(\tau|z)]$ , see Section 4.1) rather than the reward signals. Balancing the trade-off between reward and density maximization by identifying an optimal weight  $\alpha_{1}$  is challenging, especially when predictions from  $p_{\psi}(\tau|z)$  are less accurate at the beginning of training (i.e., during the cold start).
To resolve the above issues, we consider replacing the identification probability with a contrastive estimation method by constructing the following objective for policy optimization: (10)
$$
\min  _ {\pi} - \mathbb {E} _ {\pi (\cdot | z)} \left(r (\tau) + \alpha_ {1} L _ {c e} (\tau , \mathcal {V} _ {1, \dots , | Z |})\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
$$
where  $\mathcal{V}$  defines the probing sets (constructed with the density estimator in Algorithm 1). Given a specific agent type  $z$ , these probing vectors are among the most representative data points since they are located in a high-density region conditioning on  $z$  and a low-density region conditioning on other agent types ( $\tilde{z} \neq z$ ). Inspired by the InfoNCE loss [16, 37],  $L_{ce}$  can be calculated as:
$$
L _ {c e} (\tau , \mathcal {V} _ {1, \dots , | Z |}) = \sum_ {t = 0} ^ {T} \log \frac {\exp \left[ \sum_ {(\hat {s} _ {z} , \hat {a} _ {z}) \in \mathcal {V} _ {z}} f _ {s} [ (s _ {t} , a _ {t}) , (\hat {s} _ {z} , \hat {a} _ {z}) ] \right]}{\sum_ {\tilde {z} \in \mathcal {Z}} \exp \left[ \sum_ {(\tilde {s} , \tilde {a}) \in \mathcal {V} _ {\tilde {z}}} f _ {s} [ (s _ {t} , a _ {t}) , (\tilde {s} , \tilde {a}) ] \right]} \tag {11}
$$
where  $f_{s}$  denotes the score function for measuring the similarity between the features from different state-action pairs (we use cosine similarity in the experiment). To interpret  $L_{ce}(\tau, \mathcal{V}_{1,\dots,|Z|})$ , we can treat  $(s,a) \in \{\tau, \mathcal{V}_z\}$  as positive embeddings for  $\pi(\cdot|z)$  since they are generated by this policy. On the other hand,  $(\tilde{s},\tilde{a}) \in \{\mathcal{V}_{\tilde{z}}\}_{\tilde{z} \neq z}$  are negative embeddings for  $\pi(\cdot|z)$  since they are generated by controlling with other policies  $\pi(\cdot|\tilde{z})$  (where  $\tilde{z} \neq z$ ). Considering the generation is influenced by the stochasticity in environment dynamics (e.g., transition functions), we can equivalently view the generation process as injecting the environmental noise into the policy, and thus Noise Contrastive Estimation (NCE) [38] becomes a compatible tool for learning differentiable policies. Specifically, since embeddings in  $\mathcal{V}_z$  belong to a high conditional density region in  $p(\cdot|z)$  (see Algorithm 1), the knowledge from the density estimator has been integrated into policy updates.
In essence, the integration of contrastive learning into policy optimization helps the algorithm to better understand the relationships between agents, their behaviors, and the corresponding expert trajectories, resulting in improved performance for tasks involving diverse agent types. Algorithm 2 introduces the detailed implementation of MMICRL.
Algorithm 1: Probing_Sets
Input: Agent type  $z$ , trajectory dataset  $\mathcal{D}_z$  and conditional density model  $p(\cdot | z)$  Initialize a probing sets  $\mathcal{V}_z = \{\emptyset\}$ ;
for  $\tau_z\in \hat{\mathcal{D}}_z$  do Find  $(\hat{s},\hat{a}) = \arg \max_{(s,a)\in \tau_z}\log p_{\psi_i}(s,a|z) - \frac{1}{|\mathcal{Z}| - 1}\sum_{\tilde{z}}\mathbb{1}_{\tilde{z}\neq z}\log p_{\psi_i}(s,a|\tilde{z})$  Store the probing points  $\mathcal{V}_z = \mathcal{V}_z\cup \{(\hat{s},\hat{a})\}$
Output:  $\nu_{z}$
Running Setting. For a comprehensive comparison, we employ consistent evaluation metrics across all environments. These metrics include 1) Constraint Violation Rate, which assesses the likelihood of a policy violating a constraint in a given trajectory, and 2) Feasible Cumulative Rewards, which calculates the total rewards accumulated by the agent before violating any constraints. The demonstration data are assumed to have a zero violation rate. We run experiments with three different seeds and present the mean  $\pm$  std results for each algorithm. To ensure a fair comparison, we maintain uniform settings for all comparison baselines. Appendix A.2 reports the detailed settings.
Comparison Methods. We employ an ablation strategy to progressively remove components from MMICRL (Algorithm 2) and consider the following variants: 1) MMICRL-LD excludes our contrastive estimation approach, using the objective (9) directly for policy updates. 2) Inspired by [10], InfoGAIL-ICRL replaces the trajectory density model with a discriminator to identify agents based on state-action pairs, extending GAIL [28] to distinguish between distinct agents' trajectories. 3) MEICRL [4] eliminates the agent identifier and expands traditional Maximum
Algorithm 2: Multi-Modal Inverse Constrained Reinforcement Learning (MMICRL)
Input: Expert demonstration  $\mathcal{D}_E$ , upper level iterations  $N$ , lower level iterations  $B$   
Initialize the policy  $\pi_{\theta_0}(\cdot|\cdot,z)$ , constraint function  $\phi_{\omega_0}(\cdot|z)$  and density model  $p_{\psi_0}(\cdot|z)$ ;  
for  $i = 1,2,\ldots,N$  do  
Initialize an ensemble of imitation datasets  $\{\hat{\mathcal{D}}_z\}_{z\in\mathcal{Z}}$  and each dataset  $\hat{\mathcal{D}}_z = \{\emptyset\}$ ;  
for  $j = 1,2,\ldots,B$  do  
Randomly sample a  $z_j$ , and generate imitation trajectories  $\tau_{z_j}$  with  $\pi_{\theta_i}(\cdot|\cdot,z_j)$ ;  
Update the density model by maximizing:  $\log p_{\psi_i}(\tau_{z_j}|z_j)$ ;  
Store the generated trajectories and the code  $\hat{\mathcal{D}}_z = \hat{\mathcal{D}}_z \cup \{(\tau_{z_j},z_j)\}$ ;  
end  
Divide  $\mathcal{D}_E$  into sub-datasets  $\{\mathcal{D}_{E,z}\}_{z=1}^{|\mathcal{Z}|}$  by utilizing the density model (see Section 4.1);  
for  $k = 1,2,\ldots,B$  do  
Sample an agent type from the prior  $\hat{z} \sim p(z)$ ;  
Sample an expert trajectory  $\tau_{E,\hat{z}}$  from the  $\hat{z}^{th}$  expert dataset:  $\tau_{E,\hat{z}} \sim \mathcal{D}_{E,\hat{z}}$ ;  
Sample a nominal trajectory from the  $\hat{z}^{th}$  generated dataset:  $\hat{\tau} \sim \hat{\mathcal{D}}_{\hat{z}}$ ;  
Update the constraint function  $\phi_{\omega_i}$  with the objective (7);  
end  
for  $z \in \mathcal{Z}$  do  
Construct a probing set  $\mathcal{V}_z = \text{ProbingSets}(z,\mathcal{D}_z,p_\psi(\cdot|z))$  (Algorithm 1);  
update the policy with the objective (10);  
end
Entropy (ME) IRL methods to infer Markovian constraints in a model-free environment. 4) Binary Classifier Constraint Learning (B2CL) constructs a deterministic binary classifier directly for constraint inference, bypassing the need for the maximum entropy framework.
5.1 Empirical Evaluations in Discrete Environments
The discrete environments we use are based on Gridworld, a widely studied RL environment that enables us to visualize the recovered constraints and trajectories generated by various agents. We create a 7x7 Gridworld map and design four distinct constraint map settings. In Figure 2, the leftmost column illustrates expert trajectories and constraints for each configuration. Notably, the first three settings (rows 1-3 in Figure 2) incorporate two different types of constraints, while the final setting (the last row in Figure 2) includes three constraint types. The primary goal for each agent is to navigate from the starting point to the endpoint while minimizing the number of steps and avoiding their respective constraints. To facilitate constraint inference, a demonstration dataset containing expert trajectories is provided for each environment [7].
Figure 2 displays the ground-truth trajectory map alongside the learned trajectory maps. Appendix C.1 summarizes the detailed performance in terms of feasible cumulative rewards, constraint violation rate, and destination-reaching rate. Without implementing agent identification, both B2CL and MEICRL can only restore a single constraint type. Although InfoGAIL-ICRL incorporates agent identification, its performance is unsatisfactory, resulting in incomplete destination-reaching. In contrast, MMICRL-LD and MMICRL exhibit significant improvements; they can identify various agent types corresponding to different constraints and generate accurate trajectories that successfully reach the endpoint. Notably, the enhancements provided by MMICRL are more prominent, and the trajectories produced by the MMICRL algorithm closely resemble expert demonstrations.
5.2 Empirical Evaluations in Continuous Environments
In continuous environments, we evaluate ICLR algorithms by whether they can infer location constraints for various types of robots (robots can only operate within designated areas). The task involves directing the robot to navigate out of the danger zone, represented by the constraint.
Table 1: Continuous environment setting.
Setting	Dim.	1stconstraint	2ndconstraint
Half-cheetah	24	x≤ -3	x≥ 3
Antwall	121	x≤ -3	x≥ 3
Swimmer	12	x≤ -0.01	x≥ 0.01
Walker	24	x≤ -0.1	x≥ 0.1
Figure 2: We utilize blue, red, and orange colors to represent the constraints or trajectories corresponding to each agent type. The red circle and blue triangle indicate the starting and ending points of each environment, respectively. Since MECL and B2CL recover only a single, unified constraint, they produce just one type of trajectory (Marked by the red color).
We base our continuous environments on MuJoCo [17]. For each environment, we configure two types of robots—one moving forward and the other moving backward—each with opposing constraints. Table 1 provides an overview of the settings. Our goal is to distinguish their respective constraints and enable each robot type to maximize rewards without violating constraints.
Figure 4 and Table 2 illustrate the constraint violation rate and feasible cumulative rewards for each agent across various environments. To evaluate the accuracy of the discovered constraints, Figure 3 illustrates the distribution of x-coordinate values for all states visited by agents of different types. Compared to other baselines, MMICRL consistently exhibits lower constraint violation rates and higher cumulative rewards, averaged over all agents. This superiority can be attributed to MMICRL's ability to capture the diversity of constraints by inferring agent-specific constraints and learning heterogeneous policies through contrastive estimation. While certain
Figure 3: From left to right, the environments are blocked Antwall, Blocked Half-Cheetah, Blocked Swimmer, and Blocked Walker. The first row shows the results of InfoGAIL-ICRL, followed by MMICRL. The Blue and orange areas represent the trajectories of two types of agents.
bases, but the algorithm is still not efficient.
5.3 Robust Test: Can MMICRL Recover from the Early Mistakes?
To test the robustness of MMICRL, we enforce MMICRL to make incorrect agent identifications and assess whether MMICRL can recover from these errors. Specifically, in the first round of training, rather than using our density estimator, we randomly assign incorrect labels to  $20\%$  of expert trajectories. These incorrect identifications can impact ongoing constraint inference and policy updates. To recover from these mistakes, the MMICRL algorithm must accurately model diverse agent preferences and rectify the incorrect trajectories to match the appropriate agent type.
Table 2: MuJoCo testing performance. We report the average feasible rewards and constraint violation rate in 100 runs. The best average performance is highlighted in bold.
Method	Blocked Half-Cheetah	Blocked Ant	Blocked Swimmer	Blocked Walker
Feasible Cumulative Rewards
B2CL	5.03E+1 / 2.12E+3	1.33E+4 / 3.11E+3	3.28E+2 / 7.40 E-1	2.19E+1 / 7.88E+1
MEICRL	4.85E+1 / 3.45E+3	1.80E+4 / 3.31E+3	2.01E+2 / 8.30E-1	2.61E+1 / 7.33E+1
InfoGAIL-ICRL	2.22E+2 / 1.33E+2	3.73E+2 / 1.11E+2	2.38E+1 / 7.90E-1	2.15E+1 / 1.58E+3
MMICRL-LD	4.32E+3 / 2.56E+3	1.82E+4 / 2.21E+4	2.66E+2 / 6.10E+2	9.02E+2 / 5.12E+2
MMICRL	6.12E+3 / 3.06E+3	2.13E+4 / 2.17E+4	4.07E+2 / 6.48E+2	8.98E+2 / 1.53E+3
Constraint Violation Rate
B2CL	100%±0% / 67%±24%	33%±24% / 67±24%	62%±12% / 100%±0%	100%±0% / 100%±0%
MEICRL	100%±0% / 50%±25%	0%±0% / 67±24%	79%±14% / 100%±0%	100%±0% / 100%±0%
InfoGAIL-ICRL	25%±16% / 93%±5%	34%±18% / 37%±19%	73%±12% / 100%±0%	100%±0% / 0%±0%
MMICRL-LD	33%±24% / 34%±24%	0%±0% / 0%±0%	71%±16% / 34%±23%	52%±25% / 50%±25%
MMICRL	0%±0% / 0%±0%	0%±0% / 0%±0%	55%±23% / 28%±19%	31%±22% / 25%±22%
Figure 4: The feasible cumulative rewards (left two columns) and constraint violation rate (right two columns). We denote the results for different agents with  $z_0$  and  $z_1$ . From top to bottom, the environments are blocked Antwall, Blocked Half-Cheetah, Blocked Walker, and Blocked Swimmer.
Figure 5 displays the performance of MMICRL in both continuous and discrete settings. We observe that it can successfully recover multiple constraints corresponding to different types of agents. The errors made at the beginning of training do not have a significant impact on the overall performance, which demonstrates the robustness of MMICRL.
Figure 5: The constraint violation rate (left two columns), and feasible cumulative rewards (right two columns) in Blocked Ant and Gridworld-Setting1 (Appendix C.3 shows complete results).
5.4 Empirical Evaluations in Realistic Environments
To demonstrate the generalization capability of the model, we conducted some preliminary studies in a realistic autonomous driving environment. This environment is constructed by utilizing the HighD dataset (For more details, check [7] and [9]). For extracting features of cars and roads, we use the features collector from Commonroad RL [39]. The constraint that we are interested in are 1) Car distance  $\geq 20\mathrm{m}$  (agent 0) and 2) Car distance  $\geq 40\mathrm{m}$  (agent 1). Figure 6 shows the distribution of car distance in expert demonstrations for agent 0 and agent 1. We aim to investigate whether the MMICRL algorithm can differentiate between two different types of cars based on distance constraints.
Figure 6: Car distance distribution.
Figure 7: MMICRL performance in a realistic environment. We test four different settings (each column refers to one setting.). We visualize the car distance in the states generated by the imitation policies learned under the inferred constraints for agent 0 (upper row) and 1(lower row).
In the four settings shown in Figure 7, we find that in the first row, the observed distances of the first type of car are all above 20 meters, while in the second row, the distances of the other type of car are mostly above 40 meters. This suggests that the MMICRL algorithm shows promising preliminary results in this context. In the future, we will further optimize the algorithm to achieve a lower constraint violation rate and higher effective rewards in a wider range of real-world environments.
Omitting Agent Interactions: MMICRL does not consider the interactions between agents or how they can potentially collaborate or compete to satisfy a joint constraint. Future research could extend the game theory to ICRL for implementing multi-agent constraint inference.
Experiment in Virtual Environment. We evaluate our algorithm in virtual games rather than real-world applications (e.g., autonomous driving). This is due to the lack of an ICRL benchmark for a mixture of experts. Future work can explore the application of MMICRL in realistic scenarios.
In this work, we introduce the MMICRL algorithm to differentiate multiple constraints corresponding to various types of agents. MMICRL incorporates unsupervised constraint inference, agent-specific constraint inference, and multi-modal policy optimization. To demonstrate the advantages of our method over other baselines, we investigate whether MMICRL can accurately perform multiple constraint inference in both discrete and continuous environments.
The work is supported in part by the National Key R&D Program of China under grant No2022ZD0116004, by Shenzhen Science and Technology Program ZDSYS20211021111415025, and by the Start-up Fund UDF01002911 of the Chinese University of Hong Kong, Shenzhen. We acknowledge the funding from the Canada CIFAR AI Chairs program, and the support of the Natural Sciences and Engineering Research Council of Canada (NSERC).
[35] Kaiqing Zhang, Zhuoran Yang, and Tamer Basar. Multi-agent reinforcement learning: A selective overview of theories and algorithms. Handbook of reinforcement learning and control, pages 321-384, 2021.
[36] Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle. MADE: masked autoencoder for distribution estimation. In International Conference on Machine Learning (ICML), volume 37, pages 881-889, 2015.
[37] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross B. Girshick. Momentum contrast for unsupervised visual representation learning. In International Conference on Computer Vision and Pattern Recognition (CVPR), pages 9726-9735, 2020.
[38] Michael Gutmann and Aapo Hyvarinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In International Conference on Artificial Intelligence and Statistics, (AISTATS), volume 9, pages 297–304, 2010.
[39] Xiao Wang, Hanna Krasowski, and Matthias Althoff. Commonroad-rl: A configurable reinforcement learning environment for motion planning of autonomous vehicles. In International Intelligent Transportation Systems Conference (ITSC), 2021.