DIAL: Distribution-Informed Adaptive Learning of Multi-Task Constraints for Safety-Critical Systems

The CRL approach considers the environment as a constrained Markov decision process (CMDP) [49] to learn an optimal policy that maximizes the discounted cumulative rewards while ensuring the agent adheres to a set of safety constraints. To maintain consistency in notation, we redefine the CMDP from a distribution perspective as a following tuple $(\mathcal{S},\mathcal{A},\mathcal{P}_{T},\mathcal{R},\mu,\gamma,\mathcal{C})$. Here $\mathcal{S}\in\mathbb{R}^{|\mathcal{S}|}$ and $\mathcal{A}\in\mathbb{R}^{|\mathcal{A}|}$ are state and action space, respectively, where $s\in\mathcal{S},a\in\mathcal{A}$ are the elements for each space. $\mathcal{P}_{T}(s^{\prime}|s,a)$ indicates the state transition dynamics. $\mathcal{R}$ represents the set of all reward functions, where $r(s,a)\in\mathcal{R}:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$ is a reward function. Let $\mu$ denote the initial state distribution, $\gamma\in(0,1)$ the discount factor, and $\mathcal{C}=\{(\mathcal{P}_{C_{i}},\epsilon_{i})\}_{i=1}^{K}$ the set of $K$ constraints. For the $i$-th constraint, $\mathcal{P}_{C_{i}}(c|s,a)$ gives the probability of incurring cost $c$ in state $s$ and action $a$, and $\epsilon_{i}\geq 0$ is a budget, which is an upper bound on expected cumulative costs. A cost function $c_{i}(s,a)\in C_{i}:\mathcal{S}\times\mathcal{A}\rightarrow\{0,1\}$ is represented as an indicator function for unsafe conditions, where $C_{i}$ is the set of all cost functions. We denote a policy as $\pi(a|s)\in\Pi:\mathcal{S}\times\mathcal{A}\rightarrow[0,1]$, which maps states to probability distributions over actions, where $\Pi$ is the set of all policies. Let us refer to the $\gamma$-discounted cumulative sum over an infinite time horizon as the return. To be specific, we denote reward-return as $r(\tau)=\sum_{t=0}^{\infty}\gamma^{t}r(s_{t},a_{t})$ and cost-return as $c_{i}(\tau)=\sum_{t=0}^{\infty}\gamma^{t}c_{i}(s_{t},a_{t})$, respectively. Here $\tau=(s_{0},a_{0},s_{1},a_{1},\dots)$ is a trajectory of state-action pair and a set of trajectories is $\mathcal{D}=\{\tau\}_{j=1}^{N}$. When $\pi(\tau)=\mu(s_{0})\prod_{t=0}^{\infty}\pi(a_{t}|s_{t})\mathcal{P}_{T}(s_{t+1}|s_{t},a_{t})$ is defined as the probability that policy yield trajectory with $s_{0}\sim\mu$, $a_{t}\sim\pi(\cdot|s_{t})$, and $s_{t+1}\sim\mathcal{P}_{T}(\cdot|s_{t},a_{t})$ for $t\geq 0$, we can represent the expected reward-return as $\mathbb{E}_{\tau\sim\pi(\cdot)}[r(\tau)]$ and a constraint with expected cost-return as $\mathbb{E}_{\tau\sim\pi(\cdot)}[c_{i}(\tau)]\leq\epsilon_{i}$, respectively. Note that this type of constraint is typically referred to as ”soft”, meaning that it allows some trajectories to exceed the cost-return threshold $\epsilon_{i}$, as long as the expected cost-return stays within $\epsilon_{i}$. In contrast, a ”hard” constraint demands that each trajectory individually remains within the cost bound of $\epsilon_{i}$. To enforce a hard constraint, one could set $\epsilon_{i}=0$ to ensure all costs remain non-negative. In this context, the objective of maximum entropy (MaxEnt) CRL is to determine an optimal policy that maximizes the expected entropy-regularized reward-return while satisfying the given constraints. This is represented by the following formulation [15]:

where augmented term $\mathcal{H}(\pi)=-\int_{\tau}\pi(\tau)\log\pi(\tau)d\tau$ with coefficient $\beta\rightarrow\infty$ promotes randomness in action selection, supporting a certain level of exploration to prevent getting stuck in local optima while satisfying constraints.

In practical applications, it is challenging to manually specify all constraints to obtain an optimal policy in CRL. Nevertheless, it may be feasible to obtain expert demonstrations that satisfy safety requirements. ICRL effectively recovers the constraints from expert trajectories, assuming that the reward is available separately. To formalize this, previous studies [42, 20] that build on our starting point extend the MaxEnt model [39]. This model represents the probability of a trajectory under the policy $\pi$ and is adapted to CMDP as follows:

Here $\mathbf{1}(c_{i}(\tau)\leq\epsilon_{i}\,\forall i)$ serves as a feasibility indicator, representing whether a trajectory meets all constraints. Since checking this indicator becomes intractable when the state and action spaces are continuous, it is replaced by a binary classifier $\zeta(\tau)$, which approximates the indicator using a differentiable neural network with sigmoid activation. This leads to the following maximum likelihood problem [20]:

$\zeta(\tau)=\prod_{t}\prod_{i}\zeta_{i}(s_{t},a_{t})$ can be decomposed into a product of feasibility factors $\zeta_{i}(s_{t},a_{t}):\mathcal{S}\times\mathcal{A}\rightarrow(0,1)$ for each state-action pair. The feasibility that each state-action pair for all $K$ constraints is denoted as $\zeta(s_{t},a_{t})=\prod_{i}\zeta_{i}(s_{t},a_{t})$. Note that $\zeta(s,a)\notin\{0,1\}$ due to the properties of the sigmoid output. From here, we set $c(s,a)$ with $\bar{\zeta}(s,a)=1-\zeta(s,a):\mathcal{S}\times\mathcal{A}\rightarrow(0,1)$ and interpret $\zeta(\tau)$ as an estimate of the probability that $\tau$ will be feasible, similar to the case of soft constraints. In this context, the $\epsilon$ value is usually set just above zero, allowing it to operate like a hard constraint. When calculating the gradient of the constraint model to optimize the log-likelihood in Eq. 3, the reward term can be ignored. The partition function $\log Z(\zeta)$ can also be estimated through a sample-based approximation using nominal policy $\pi$ learned in the forward step. Thus, the update of the constraint function is derived by matching the expected gradient of $\log\zeta$ for each expert and nominal trajectories as follows [20]:

where $\omega(\tau)=\frac{\zeta}{\zeta^{\prime}}$ represents importance sampling weights, defined as the ratio relative to $\zeta^{\prime}$ computed in earlier iterations. Additionally, the subscript E indicates elements associated with the expert. Assuming that the expert policy $\pi_{E}$ is known, we have access to expert trajectories $\tau_{E}$ sampled from $\mathcal{D}_{E}=\{\tau_{E}\}_{j=1}^{N_{E}}$. The standard ICRL method uses an iterative approach that alternates between updating the policy and the constraint function based on Eq. 1 and Eq. 4, as illustrated on the right side of Fig. 1a. In practice, the constraint function is approximated using the complement of $\zeta$ obtained from the inverse step, treating it as $c_{i}\approx\bar{\zeta}_{i}=1-\zeta_{i}$ to establish the constraints. The nominal policy is then updated using the PPO Lagrange algorithm [50].

When designing the distribution-aware constraint in Fig. 1b, we aim to capture the diverse risk preferences in expert demonstrations that meet safety requirements, even for randomly assigned tasks. To achieve this, we employ a learning approach that infers distorted constraint distributions with properly selected risk levels. Instead of maximum likelihood estimation for $\zeta(\tau)$, Bayesian methods [52, 45, 46] can estimate a posterior distribution $p(\zeta(\tau)|\mathcal{D})$ by combining a prior $p(\zeta(\tau))$ with the trajectory likelihood $p(\mathcal{D}|\zeta(\tau))\approx\pi(\tau)$, as specified in Eq. 2. Previous methods calculate the constraint $\mathbb{E}_{\tau\sim\pi(\cdot),\zeta\sim p(\cdot|\mathcal{D})}[\bar{\zeta}(\tau)]\leq\epsilon$ to ensure the mean of the estimated distribution remains below the threshold. However, this approach is vulnerable to long or heavy-tailed distributions, rendering it insufficiently stringent for handling rare but extreme situations. The mean constraint ultimately lacks the flexibility to adjust risk levels for changes in the unsafe set, shown as the green polygon in Fig. 1b. This limitation becomes even more pronounced in multi-task settings, where the attraction region, marked by the blue boundary, expands to cover more tasks. Although a previous study [46] addresses this issue by employing a risk-sensitive criterion like CVaR, it focuses solely on optimizing the policy for a single task during the forward step. In contrast, we emphasize using CVaR in the inverse step to learn constraints capable of handling scalable novel tasks.

When estimating statistical metrics like the mean or CVaR of a posterior distribution, it is often challenging because arbitrary distributions lack a precise closed form. Therefore, it is necessary to approximate them with more manageable distributions. In our case, we would like to model the probability that a trajectory is safe as a random variable that falls within a finite interval between 0 and 1. Consequently, we choose an ideal Beta distribution as the posterior distribution and set $p(\zeta(\tau)|\mathcal{D})\approx q(\zeta(\tau)|\alpha)$. The Beta distribution is particularly useful as it can represent asymmetric or heavy-tailed distributions. This distribution reflects the ratio of binary outcomes, such as successes and failures, of safe paths and is characterized by two parameters, $\alpha=[\alpha_{1},\alpha_{2}]$. For simplicity in notation, we omit $i$ and assume a single constraint case. To handle multiple constraints, we can factorize $q(\zeta(\tau)|\alpha)=\prod_{i}q(\zeta_{i}(\tau)|\alpha^{i})$, expressing it as a product of independent components.

We use CVaR to measure distorted risk, taking into account the probability of a path being safe at a specified risk level within the Beta distribution. Fig. 3 illustrates how CVaR is calculated in contrast to the mean. First, we determine the value at risk (VaR), which represents the threshold probability of safety at a given risk level $\lambda$. This value is the $\lambda$-percentile $F_{\zeta(\tau)}^{-1}(\lambda;\alpha)$, with $\text{VaR}_{\lambda}$ defined as the lower $\lambda$-quantile of the distribution:

where $F_{\zeta(\tau)}(x;\alpha)$ represents the cumulative distribution function (CDF) of the $q(\zeta(\tau)|\alpha)$ distribution. Next, CVaR is obtained by calculating the expected value over the region below $\text{VaR}_{\lambda}$ as follows:

This represents the average of extreme values at a given risk level. Note that if $\lambda=1$, CVaR is equal to the mean. We employ a neural network $f_{\phi}(\alpha|\tau)$, which takes a trajectory $\tau$ as input and outputs the parameter $\alpha$ of a Beta distribution. This approach facilitates constraint learning by capturing the risk distribution and incorporating CVaR. Since the two parameters of the Beta distribution are positive, the final layer is implemented with a softplus activation. We interpret the network’s output as variables sampled from two Gamma distributions, treating the approximated network as a prior model for the Beta distribution. This approach allows the shape of the Beta distribution to be determined by random variables generated by the Gamma distributions, rather than fixed parameters, thereby capturing uncertainty in the data. This flexibility makes probabilistic modeling more adaptable and relevant to the context. Therefore, we update the constraint model by applying the risk-sensitive criterion $\Gamma_{\phi}^{\lambda}(\tau)\doteq\mathbb{E}_{\alpha\sim f_{\phi}(\cdot|\tau)}[\text{CVaR}_{\lambda}]$ in place of $\zeta(\tau)$ in Eq. 4. Consequently, the gradient of the loss $\nabla_{\phi}\mathcal{L}_{C}(\phi,\lambda)$ is formulated as:

In addition, considering diverse risk levels is known to be beneficial for acquiring risk-sensitive knowledge [53]. For this reason, we update the network using risk level $\lambda$ sampled from the uniform distribution $\mathcal{U}(0,1)$ in the safe IL stage where multi-task learning is conducted. In the later safe TL stage, where achieving high performance on the target task is crucial, fine-tuning is performed using grid search to determine $\lambda$. Following this, we include a regularization term when optimizing the evidence lower bound (ELBO) for the approximate posterior as follows:

Given that both distributions $q(\cdot)$ and $p(\cdot)$ are Beta distributions, the Kullback-Leibler (KL) divergence can be computed in closed form, as shown in [54].

We introduce a novel policy improvement method within the ICRL framework that ensures safety by allowing limited exploration at an acceptable risk level, enabling effective learning across diverse tasks. Through this proposed approach, we obtain a TASE policy that helps in learning shared knowledge across various tasks. This enables the recovery of a more flexible constraint function that supports robust inference by accounting for varying risk levels and task-specific environmental changes. Our approach builds on the observation that using data spanning a wider range of tasks improves the accuracy of generalizable constraint learning [21]. Unlike prior approaches that either set constraints tailored to specific tasks [47] or apply the same constraints across all tasks [26], our method provides improved generalization and adaptability to new situations.

Learning a policy in IL involves aligning it to frequently visit the state-action spaces observed in the demonstrations. Effective training requires the nominal policy to aim for an even coverage of a wide range of states. This approach enables divergence in the probability density of state-action coverage from the expert policy, thereby facilitating learning. In practice, this approach involves increasing the entropy of the probability density over the state-action pairs that the agent visits. To achieve this, Eq. 1 includes an entropy regularizer term $\mathcal{H}(\pi)$ that directly measures the probability distribution $\pi(\tau)$ of the nominal policy, thereby enhancing exploration. However, this naive method has limitations as it focuses solely on the probability distribution of the policy without capturing the correlations between states. In addition, simply relying on random actions for exploration can be inefficient, particularly in high-dimensional spaces or complex tasks. When the policy repeatedly selects high-reward actions early in training, confidence in these actions increases, resulting in lower entropy and reduced exploration. This can eventually hinder learning or prevent the achievement of goals.

To address these limitations, we promote structured exploration by incorporating correlations between states into our entropy estimation. This approach is implemented under constraints that ensure safety. By focusing on task-relevant states, we interpret the increase in entropy as an exploration bonus in the policy update process, effectively enhancing task entropy. To extend to complex domains, we can indirectly model the state density function $\rho:\mathcal{S}\rightarrow\mathbb{R}_{\geq 0}$ using $M$ particle groups $\mathcal{S}=\{s_{i}\}_{i=1}^{M}$ [55]. Using these particle groups, we represent the average state density visited by the policy $\pi$ as $\rho_{\pi}(s)=\frac{1}{T}\sum_{t=0}^{T}\gamma\rho(s_{t}=s|\pi)$ for time horizon $T$, where $\int_{\mathcal{S}}\rho_{\pi}(s)\,\text{d}s=1$. The entropy of the state density $\rho_{\pi}$ is computed as $\mathcal{H}(\rho_{\pi})=-\int_{s\in\mathcal{S}}\rho_{\pi}(s)\ln\rho_{\pi}(s)\,ds$. In practice, we approximate this entropy using the particle groups, yielding $-\sum_{i=1}^{M}\hat{\rho}_{\pi}(s_{i})\ln\hat{\rho}_{\pi}(s_{i})\Delta s_{i}$, where $\Delta s_{i}$ represents the interval width around each particle $s_{i}$. The approximated density $\hat{\rho}_{\pi}=\frac{k}{M\cdot V_{i}^{k}}$ for these particles is determined by estimating the volume of a hypersphere formed by the radius $R_{i}=|x_{i}-x_{i}^{k\text{-NN}}|$. This calculation uses the $k$-NN method for each particle, where $x_{i}^{k\text{-NN}}$ represents the position of the $k$-th nearest neighbor ($k$-NN) particle to $x_{i}$. This density gives the $k$-NN entropy estimator [56] as follows:

Here, $\Gamma(\cdot)$ represents the gamma function, and $\Psi(\cdot)$ is the digamma function, which is the logarithmic derivative of the gamma function. The last two terms correct the average bias introduced to address entropy underestimation for small $k$ values. To integrate this entropy estimator into RL, we use samples from the old policy $\pi_{\bar{\theta}}$ to approximate the state entropy of the current policy $\pi_{\theta}$. The policy $\pi$ is parameterized by a neural network with parameters $\theta$. For simplicity, we denote $\rho_{\pi_{\theta}}$ as $\rho_{\theta}$, omitting $\pi$ in the notation. Subsequently, to handle the discrepancy between the sampling policy $\pi_{\bar{\theta}}$ and the target policy $\pi_{\theta}$, we apply an importance-weighted (IW) $k$-NN estimator [57] as follows:

Here, $\mathcal{N}_{i}^{k}$ is the set of indices of $k$-NN of $x_{i}$. In this context, when $\theta=\bar{\theta}$ and a uniform weight $w_{j}=\frac{1}{N}$is applied, we obtain $\hat{\mathcal{H}}_{k}(\rho_{\theta}|\rho_{\bar{\theta}})=\hat{\mathcal{H}}_{k}(\rho_{\theta})$. Notably, in the IW $k$-NN estimation approach, trajectories are sampled independently, while the states within each trajectory account for the correlations among neighboring particles. Furthermore, to stabilize the convergence, we utilize a KL estimator $\hat{\mathcal{D}}_{KL}$ [18]. This is computed as the difference between the IW $k$-NN estimator in Eq. 10 and the estimator in Eq. 9, with the bias correction term canceling out. As long as the updated policy satisfies $\hat{\mathcal{D}}_{KL}[\rho_{\theta}\parallel\rho_{\bar{\theta}}]\leq\delta$, we can optimize the policy multiple times, serving as a trust-region constraint.

Eventually, we can replace the constraint with the risk-sensitive criterion and the entropy regularizer term with the IW k-NN estimator in Eq. 1. We then solve an unconstrained min-max optimization problem by applying the Lagrangian method to the objective function [58]. This approach allows us to handle the original constrained problem as an equivalent unconstrained problem, which we define as follows:

where $\mathcal{J}_{R}(\theta)=\mathbb{E}_{\tau\sim\pi_{\theta}(\cdot)}\left[r(\tau)\right]$ encourages the maximization of expected reward-return. Moreover, $\mathcal{J}_{H}(\theta)=\hat{\mathcal{H}}_{k}(\rho_{\theta}|\rho_{\bar{\theta}})$ is an entropy-based term, scaled by $\beta$ to control the degree of entropy regularization, thereby encouraging exploration in the policy. Lastly, $\mathcal{J}_{\kappa}(\phi,\lambda)=\mathbb{E}_{\tau\sim\pi_{\theta}(\cdot)}[\bar{\Gamma}_{\phi}^{\lambda}(\tau)]-\epsilon$ is a constraint term that ensures the policy remains within the constraint threshold with an expected risk $\bar{\Gamma}_{\phi}^{\lambda}=1-\Gamma_{\phi}^{\lambda}$. In this setup, $\kappa$, also referred to as the safety weight, is a Lagrange multiplier associated with the constraint term $\mathcal{J}_{\kappa}$. By optimizing this Lagrangian formulation, we effectively balance maximizing rewards and entropy while minimizing the constraint violation. Algo. 1 and 2 summarize our training procedure for the safe IL and safe TL stages in DIAL, respectively.

Safe IL and TL are evaluated within the navigation and control domain, encompassing scenarios with risky situations involving static or dynamic obstacles across state spaces ranging from low to high dimensions. In safe IL, the aim is to perform exploration that navigates through as many safe states as feasible for arbitrarily given tasks. In safe TL, the objective is to conduct safe exploitation that effectively tackles a given target task while ensuring it remains within the constraints. To demonstrate the advantages of these two components, we have composed urban driving tasks that address real-world navigation and robot control tasks to validate performance across diverse environments, further details are described in the following paragraphs.

Urban Driving: To compare environments characterized by dynamic risks in high-dimensional states, we use a modified version of the intersection environment provided by HighwayEnv [59], referring [26]. Fig. 4 illustrates safe IL, where the agent learns shared constraints (4b) from demonstrations in the multi-task scenario (4a), and safe TL, where the agent applies this learned knowledge to adapt in the meta-task scenario (4c). The observation space includes 7 types of kinematic information, such as position and speed, for the agent and 15 surrounding vehicles. To achieve a permutation-invariant representation of the $\mathbb{R}^{15\times 7}$ input independently of the order of surrounding vehicles, we use an attention-based encoder as the backbone network for the constraint model $f_{\phi}(\alpha|\tau)$. This encoder captures the relative importance of each surrounding vehicle and calculates a weighted sum to represent the input as a latent variable in $\mathbb{R}^{32}$. We adopt the same controllers, reward function, and constraints as specified in [26] to ensure fair comparisons with the baselines. Each vehicle follows a fixed path by adjusting acceleration and steering angle through a linearized controller parameterized in $\mathbb{R}^{5}$ [51]. Surrounding vehicles maintain fixed parameters, while the agent’s parameters are optimized using a constrained cross-entropy method (CEM) [32] rather than the PPO Lagrangian in Eq.11. This approach addresses multiple constraints by first ranking parameters according to the number of constraint violations, then assessing them based on violation magnitude and reward. The reward function is designed as a linear combination of random variables and features to reflect various driving preferences, including reaching the target, driving speed, and lane changes:

where $v_{t}$ represents the vehicle’s speed, $\angle_{t}$ indicates the difference angle between the heading of the vehicle and the target lane, and $\xi_{v}\sim\mathcal{N}(0.1,0.1)$ and $\xi_{\angle}\sim\mathcal{N}(-0.2,0.1)$ are random variables. The target lane is the path extending to the given target location along the road network provided by the environment. The constraint function limits dangerous events by defining features composed of four metrics, $\varphi(s_{t},a_{t}):\mathcal{S}\times\mathcal{A}\rightarrow\{0,1\}^{4}$. These features include cases where the speed exceeds 15, the distance to the vehicle ahead is less than 10, a collision occurs, and the vehicle departs from the road. A safe situation is defined as one where the probability of each event occurring in an episode remains below the ground truth constraint thresholds $\epsilon=[0.2,0.2,0.05,0.1]$. In this problem, we consider $1-\max(0,\frac{1}{T}\sum_{t}^{T}\varphi(s_{t},a_{t})-\hat{\epsilon})$ as the feasibility function and learn to infer thresholds $\hat{\epsilon}$ for the four defined features.

Robot Control: To compare environments with static risks in low-dimensional states, we use a modified version of OpenAI Gym [60], shown in Fig. 5a-c. Each environment has a safety area marked with a red line. For MountainCar, the goal is to reach the point where the right flag is located without the car going to the left of the red line. However, the agent is penalized so that a large action $a_{t}$ is not performed. For CartPole, the goal is to raise the pole angle $\theta_{t}$ vertically while keeping the cart inside both red lines. For BasicNav, the goal is to reduce the goal distance $d_{t}=s_{goal}-s_{t}$ while avoiding the circular hazard region in the center. For evaluating responses to risks involving randomness in high-dimensional states, we use Safety Gym [50], illustrated in Fig. 5d. For PointGoal, we aim to control the point robot to reach a random goal described as a green cylinder. In this environment, the agent is restricted from entering the blue circles depicted on the ground or pushing the vase marked with a cyan block. A vase moving upon contact and stationary blue circles are randomly generated within a certain range around the target point.

The rewards and costs used to train the RL agent are given the prefix ”extrinsic” for those provided by the environment and ”intrinsic” for those recovered through IL. Each environment provides an extrinsic cost whenever unsafe interactions occur, but this is used only for performance evaluation and not for training. However, the extrinsic rewards from the environment and the recovered intrinsic costs are both used for training. Further details on each environment are provided in the Appendix. A.

Our approach, DIAL, is an algorithm designed to learn constraints that ensure safety when addressing new meta-tasks. These tasks are set in similar but slightly different environments from those demonstrated in multi-task settings. Several approaches based on IRL or ICL serve as natural starting points to address our constraint learning problem. This makes them suitable for comparison with our proposed method. We adopt the following three methods as baselines: MERL [39], MECL [20], and COCL [26]. MERL and MECL both belong to the MaxEnt-based IL family. They can derive the optimal policy that imitates the expert from single-task demonstrations. The key difference between these methods lies that MERL recovers rewards, while MECL focuses on recovering constraints under the assumption that rewards are known. In MERL, the average of the inferred individual rewards for each task demonstration, $\sum_{i}^{N}\hat{r}_{i}$, can be added as a penalty term to the reward $r_{\text{eval}}$ for a new meta-task. In MECL, the average of the inferred individual constraints across multi-task demonstrations, $\sum_{i}^{N}\hat{c}_{i}$, is incorporated as a constraint term. This term is considered separately from the reward $r_{\text{eval}}$, under the assumption that the rewards for multi-task demonstrations are already known. COCL does not belong to the IL family that cannot obtain the policy. However, it can recover a shared constraint $\hat{c}$ even without knowing the individual rewards $r_{i}$ for each task. This is achieved by constructing the convex hull of the safety set based on feature expectations from the demonstration data. Subsequently, if $r_{\text{eval}}$ and $\hat{c}$ are known, a policy can be obtained by optimizing Eq. 1. In the robot control environment illustrated in Fig. 5, COCL is excluded as a baseline because it cannot be applied due to the lack of directly defined feature vectors for the state. We highlight that our approach, like COCL, learns constraints across multi-task demonstrations. However, DIAL can derive the TASE policy and considers the distribution of constraints. These aspects distinguish our method and provide notable advantages, particularly in helping to safely adapt to changed environments when addressing new tasks.

DIAL trains two learnable models: the constraint function and the policy. Both models have a two-layer neural network architecture with 256 hidden units and ReLU activation across all methods. The policy handles continuous action spaces by outputting the mean and variance of a Gaussian distribution. The mean is constrained to a specific range using a tanh output, while the variance ensures positive values using a softplus output. In our approach, the constraint function approximates the parameters of a Beta distribution, which is why we use a softplus output instead of softmax. All neural network parameters are updated using the Adam [61] optimizer. As an exception, in the Urban driving environment, we employ an encoder that embeds inputs into permutation-invariant representations as the backbone model for the constraint function in all methods. Additionally, in this environment, the policy is represented by a linearized controller instead of neural network and is optimized using CEM [32]. The hyperparameters used in the experiments were tuned through a coarse grid search. Appendix. B provides detailed descriptions of the expert demonstration collection method, hyperparameter selection, and stabilization techniques for training to facilitate experiment reproduction. We observed experimentally that DIAL achieves asymptotic performance within approximately 150K environmental steps in urban driving of Fig. 4 and within 20K, 300K, 300K, and 1M environmental steps in robot control of Fig. 5. We believe these results are due to the use of a distribution-aware constraint function, which allows for flexible adjustment of risk levels and cautious exploration of new tasks. All experiments were conducted on a PC with an Intel Xeon Gold 6248R CPU (3.00GHz, 48 cores), an NVIDIA GeForce RTX 3090 GPU, and 256GB of memory.

We use the following metrics to evaluate performance: reward-return (RR), cost-return (CR), constraint violation rate (CV), and state entropy (SE). RR and CR represent the average of the total explicit rewards and costs obtained by the agent during an episode, averaged across multiple episodes. RR reflects task performance, while CR evaluates safety performance. CV indicates the proportion of episodes where the value of CR divided by $T$ exceeds the specified constraint budget $\epsilon$, representing the likelihood of violating constraints in a single episode. When handling multiple constraints, the sum of all CVs is used. SE is calculated by discretizing two primary states in each environment and measuring the frequency with which the policy visits each state during all episodes. This metric is used to compare the exploration level of different policies. In particular, the map that assigns visit frequencies or inferred constraints to each discretized state is useful for qualitative comparisons. The two primary states are empirically selected for each environment to clearly highlight differences. For more details, please refer to Tab. LABEL:table:environmental_settings in Appendix. A. All metrics are presented as the average values measured over 20 episodes for 5 seeds. The shaded areas in the plots represent the standard deviation.

This section evaluates the constraint function related to safety requirements and the policy associated with task success through safe imitation learning (IL) from demonstrations collected while safely performing multi-task. Fig. 6a shows the RR measured according to the number of expert demonstrations used for training, as well as the sum of the four CVs presented in Fig. 4b. Our proposed DIAL method achieves both higher task performance and lower CV as more demonstrations are used, demonstrating the best results among the compared baselines. While there is a slight increase in the variance of RR in the last segment of the plot, the values show a gradually stable increase. This can be interpreted as a natural phenomenon due to the differences in scale among the various tasks. COCL maintains relatively stable performance but is not as efficient as DIAL. MERL shows little improvement in RR even when many demonstrations are used, and it decreases at the end due to overfitting. Additionally, CV remains high. Although MECL effectively reduces CV, RR decreases as more demonstrations are added due to excessive conservatism.

To confirm whether DIAL has successfully learned all the designed constraints when trained with 300 expert trajectories, Fig. 7 presents learning curves showing the violation rates for each condition. As training progresses, the violation rates for all constraints gradually decrease and ultimately stabilize at low levels, demonstrating the stability and reliability of the proposed method. Furthermore, to examine the proposed method that approximates each constraint budget $\hat{\epsilon}_{i}$ as variables of a Beta distribution, Fig. 8 shows how the density function of this distribution evolves for training. These results indicate that DIAL can learn to adhere to the constraints more effectively. At the beginning of training, all conditions start with low parameter values, resulting in a distribution that is evenly spread across the entire threshold range. As training progresses and the $\alpha_{2}$ value increases rapidly, the distribution becomes concentrated in the lower threshold region. This rapid increase can cause the agent to become overly conservative, potentially limiting task performance. To prevent this, the proposed method distorts the distribution based on randomly sampled risk levels, encouraging flexible exploration for specific conditions. Additionally, the term in Eq. 8 prevents overfitting of the distribution by incorporating a given prior probability. Due to this design, soft constraints like speeding or maintaining distance from the vehicle ahead can allow some violations during the middle stages of training, enabling finer adjustments. Towards the later stages of training, there is a noticeable tendency for the density to increase at specific thresholds and for the distribution to narrow. This implies that the model has gradually achieved stable performance.

In the robot control environments, we evaluate the degree of safe exploration of policy by comparing the SE and CR based on the number of expert trajectories used, as shown in Tab. I. DIAL demonstrates higher SE and simultaneously lower CR compared to other methods, even when using a relatively small number of trajectories in all environments. MERL exhibits high exploration performance in terms of SE but has a high CR, leading to reduced safety. MECL maintains a relatively low CR but has low SE, indicating insufficient exploration performance. Although DIAL’s SE is slightly lower than MERL’s, the difference is minimal as shown by the underlined numbers and there is a significant difference in CR. These results demonstrate that our method successfully balances diverse state exploration and safety.

Furthermore, Fig. 9 visually demonstrates the effectiveness of DIAL’s design in finely adjusting the deviation in risk levels by carefully selecting $\lambda$ with limited data. Even when using the same expert data shown on the far left in Fig. 9a, the exploratory tendencies of the learned policy vary depending on the choice of $\lambda$. Policy tend to avoid risks when $\lambda$ is low and take risks when $\lambda$ is high. This result suggests that setting a low $\lambda$ is advantageous for balancing when expert states are near the safety boundary. Conversely, when the data is far from the boundary, setting a high $\lambda$ is more appropriate. Fig. 9b shows that we can infer a constraint that most closely resembles the ground truth (GT) located on the far left by finely tuning $\lambda$ in DIAL. In contrast, MERL fails to properly capture the constraint distribution, and MECL can have its distribution’s center and shape distorted differently from the GT due to reliance on the characteristics of the data used for training.

In this section, we address safe TL in environments where the safety requirements are the same as safe IL, but the explicit reward functions for the target task are given. In safe TL, the key concern is whether the agent can maintain compliance with the constraints without forgetting them despite changes in the objective function for solving the target task. To demonstrate that using the function recovered through safe IL reduces the burden of cost design, we use the extrinsic cost only to evaluate safety and do not use it in safe TL. Details of the hyperparameters used for training the agent are provided in Tab. LABEL:table:safeTL_hyperparameters in Appendix. B. Fig. 6b compares the performance of meta-task learning in the urban driving environment illustrated in Fig. 4c. For DIAL, as the number of demonstrations increases, RR steadily improves while CV decreases. In the case of COCL, although it has slightly better performance than DIAL when the number of demonstrations is small, it shows only a marginal reduction in CV and almost no improvement in RR, even with an increased number of demonstrations. MERL has significant limitations in terms of safety, while MECL suffers from degraded performance due to overly conservative behavior. Overall, DIAL effectively balances performance and safety, utilizing the available demonstrations to achieve the best overall results.

Fig. 10 illustrates how the DIAL agent gradually acquires the ability to perform tasks safely and efficiently during the learning process. The average speed fluctuates significantly in the early stages, However, it gradually stabilizes between approximately 11 and 12 as the episodes progress. The vehicle’s deviation from the lane center also reaches a stable level below 0.2. The goal achievement rate and feasible rewards steadily increase and converge to certain values, where feasible rewards are the sum of rewards the agent obtains only when all safety requirements are met. In the $20$% of cases where the goal is not achieved, timeouts occur due to deadlocks caused by congestion or collisions in the surrounding traffic. To verify how the agent trained with DIAL performs the meta-task safely, we visualize the distribution of environmental data obtained by the agent while executing tasks in a single episode, as shown in Fig. 11. The vehicle’s speed remains mostly concentrated near the threshold without exceeding the speed limit of 15. The distance to the vehicle ahead is evenly distributed over values greater than the threshold of 10. The lateral distance from the center of the lane is concentrated within the range of -1 to 1 to prevent lane departure, where the deviations in values are due to reaching the destination. These results show that the agent learns and retains behaviors such as maintaining speed limits, safe distances, and staying centered in the lane, even when performing novel tasks.

Tab. II presents a comparison of RR and CR measured using the fully trained policies on robot control tasks. DIAL consistently exhibits the lowest CR across all environments, satisfying the condition that the average is less than the threshold for each environment, as indicated in parentheses. In terms of RR, DIAL shows the highest performance in MountainCar, and while MERL has the highest RR in the other environments, DIAL achieves performance that is nearly comparable to MERL, as seen in the underlined numbers. For MERL, although the RR is relatively high, the CR is extremely large, which limits its ability to address safety issues. In the case of MECL, the RR is generally low, and the CR is higher than that of DIAL, exceeding the threshold for each environment.

Fig. 12 visually depicts the improvement process of the policy learned through DIAL in safe TL. In the early stages of training, the agent visits as many safe states as possible. During the middle of training, unsafe interactions occur at the boundaries of the safe area. However, The agent gradually converges to maximize the reward for the given target task while satisfying safety constraints. For detailed information on the rewards for each environment, please refer to Figure 13 in Appendix A. These results demonstrate that even when the objective function is altered to maximize the reward of a new target task, the proposed method does not lose the safety information previously learned.