SafeMIL: Learning Offline Safe Imitation Policy from Non-Preferred Trajectories

TL;DR

SafeMIL tackles offline safe imitation learning when only limited non-preferred trajectories and unlabeled data are available, lacking explicit per-timestep rewards or costs. It learns a parameterized cost via a Multiple Instance Learning formulation, using bag-level risk signals to distinguish high-cost (non-preferred) from potentially safe behavior within unlabeled data. The learned cost guides a behavior- cloning policy to imitate trajectories likely to satisfy CMDP constraints, with optional soft weighting to emphasize safer samples. Empirical results across velocity-constrained and navigation tasks show SafeMIL achieves substantially safer policies without sacrificing reward performance, outperforming or matching state-of-the-art baselines and delivering a median improvement of about 3.7× in safety metrics across environments.

Abstract

In this work, we study the problem of offline safe imitation learning (IL). In many real-world settings, online interactions can be risky, and accurately specifying the reward and the safety cost information at each timestep can be difficult. However, it is often feasible to collect trajectories reflecting undesirable or risky behavior, implicitly conveying the behavior the agent should avoid. We refer to these trajectories as non-preferred trajectories. Unlike standard IL, which aims to mimic demonstrations, our agent must also learn to avoid risky behavior using non-preferred trajectories. In this paper, we propose a novel approach, SafeMIL, to learn a parameterized cost that predicts if the state-action pair is risky via Multiple Instance Learning. The learned cost is then used to avoid non-preferred behaviors, resulting in a policy that prioritizes safety. We empirically demonstrate that our approach can learn a safer policy that satisfies cost constraints without degrading the reward performance, thereby outperforming several baselines.

Figures (15)

• Figure 1: Performance Comparison. We report the final bootstrapped mean performance of the algorithm on Walker2d-Velocity, Swimmer-Velocity, Ant-Velocity task after 1 million training steps. Mean and 95% CIs over 5 seeds. We observe that our method outperforms all the baselines and can recover low cost safe policies without compromising reward performance. We also report the learning curves for all the algorithms in Fig. 8 in Appendix.
• Figure 2: Performance Comparison. We report the final bootstrapped mean performance of the algorithm on Point-Circle2, Point-Goal1, Point-Button1 tasks after 1 million training steps. Mean and 95% CIs over 5 seeds. We observe that the Point-Goal1 environment, our algorithm performs better, while maintaining competitive performance in the remaining environments. We also report the learning curves for all the algorithms in Fig. 9 in Appendix.
• Figure 3: Sensitivity to Bag Size. We report the final mean performance of the algorithm on the Swimmer-Velocity environment for different bag sizes $K=\{1, 8, 16, 64, 128\}$, after training for 1 million steps. Mean and 95% CIs over 5 seeds. We observe that increasing the bag size $(K)$ lead to a higher probability of finding preferred trajectories within the unlabeled bag. Thereby, improving the algorithm's safety performance while maintaining reasonable episode return.
• Figure 4: Sensitivity to Trajectory Length. We report the final mean performance of the algorithm on the Swimmer-Velocity environment for different partial trajectory lengths $H=\{1, 5, 10\}$, after training for 1 million steps. Mean and 95% CIs over 5 seeds. For a sufficient bag size of $128$, we observe that safety performance is stable across different trajectory lengths.
• Figure 5: MuJoCo-based velocity-constrained tasks of DSRL environment. For each task, the agent needs to move as fast as possible while adhering to the velocity limits.

Theorems & Definitions (7)

• Definition 1: Preferred Trajectory
• Definition 2: Non-preferred Trajectory
• Lemma 1
• proof
• Theorem 1
• proof
• proof