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

AI summary

Overview

Research area: Offline safe imitation learning (a fusion of imitation learning, safe reinforcement learning, and constrained Markov decision processes).

Technical level: Advanced. The paper assumes familiarity with CMDPs, behavior cloning, distribution correction methods, and multiple instance learning.

Scope: The paper proposes SafeMIL, a method that learns a parameterized safety cost function via a Multiple Instance Learning (MIL) formulation, then uses that cost to filter or reweight an unlabeled trajectory dataset so that a behavior-cloning policy imitates only the behaviors likely to satisfy safety constraints.

What This Paper Is About

In many safety-critical domains, a learning agent cannot safely explore online, and it is hard to specify a reward function and per-timestep safety costs. However, it is often easy to collect a small set of trajectories showing behavior to avoid — red-light running, vehicle near-misses, toxic chatbot outputs — which the authors call "non-preferred trajectories."

The goal is to learn a safe policy from a small number of non-preferred trajectories plus a large pool of unlabeled trajectories containing a mix of safe (preferred) and unsafe (non-preferred) behavior, with no reward or cost labels anywhere. The agent must both imitate good behavior and avoid the risky behavior demonstrated in the non-preferred set.

Key Contributions

1. A MIL formulation for cost learning in offline safe IL. The paper formulates learning a parameterized cost function that predicts whether a state-action pair is risky as a Multiple Instance Learning problem, where bags are sets of trajectories and only bag-level labels (non-preferred vs. unlabeled) are available. The authors state this is the first introduction of an MIL formulation for the offline safe IL setting.

2. An intuitive bag score function. Using a permutation-invariant score function based on the Fundamental Theorem of Symmetric Functions with monomials, with the identity function chosen for the outer transform, the bag score becomes the average discounted cumulative predicted cost over the bag's trajectories. This score is trained with a Bradley-Terry-style loss to rank non-preferred bags above unlabeled bags.

3. Preferred-behavior identification and policy learning. The learned cost is used either to hard-filter unlabeled trajectories whose estimated cumulative cost is below a threshold, or to softly weight each trajectory in a weighted behavior-cloning loss. Both schemes let the agent imitate behavior likely to satisfy CMDP constraints.

4. Strong empirical safety results. Across six DSRL benchmark environments, SafeMIL achieves a final median performance reported as 3.7 times better than the best baseline algorithm.

Main Findings

• Safer policies without reward loss: SafeMIL consistently outperforms or matches state-of-the-art offline safe IL methods on safety, reporting that it can recover low-cost safe policies without compromising reward performance.

• Velocity-constrained tasks: SafeMIL outperforms all baselines on Walker2d-Velocity, Swimmer-Velocity, and Ant-Velocity, with safety performance closely approximating the Constrained-RL reference policy (COptiDICE).

• Navigation tasks: SafeMIL learns a safer policy on Point-Goal1 and is competitive with other baselines on Point-Circle2 and Point-Button1. On Point-Circle2, T-REX-WBC learns a policy with better constraint satisfaction.

• Standard BC is unsafe on mixed data: When the unlabeled dataset contains constraint-violating trajectories, the plain BC-Unlabeled policy has higher cost and can exhibit non-preferred behavior.

• Larger bags improve safety: On Swimmer-Velocity, increasing bag size across K = {1, 8, 16, 64, 128} improves safety (lower cost) while maintaining reasonable episode return, stabilizing after a certain bag size. This matches the theoretical result that larger bags raise the probability of containing at least one preferred trajectory.

• Trajectory length is not critical: With a sufficiently large bag size (K = 128), safety performance on Swimmer-Velocity is stable across partial trajectory lengths H = {1, 5, 10}, supporting the use of partial trajectories as a computationally efficient alternative to full trajectories.

• Weighting scheme choice is minor: Trajectory-level weighting and state-action (transition)-level weighting produce similar safety and return performance on Swimmer-Velocity and Point-Goal1.

Methodology in Plain English

The researchers assume two offline datasets: a small set of non-preferred trajectories (50 trajectories in their experiments) and a much larger unlabeled set (200 trajectories) containing both good and bad behavior. No reward or cost labels exist for either set.

The core idea is to learn a neural cost function that scores how risky a state-action pair is. Because per-timestep risk labels are unavailable, the authors borrow the Multiple Instance Learning trick: instead of labeling individual state-action pairs, they label groups (bags) of trajectories.

• Negative bags are formed by sampling trajectories from the non-preferred dataset — every trajectory in such a bag is assumed risky.
• Unlabeled bags are formed by sampling from the unlabeled dataset. Because the unlabeled set contains some preferred trajectories, a large enough bag will almost certainly contain at least one of them; the probability is 1 − (1 − α)^K, where α is the proportion of preferred trajectories.

Each bag is scored by summing the discounted predicted costs of its trajectories and averaging over the bag. The score function is made permutation-invariant so the ordering of trajectories inside the bag does not matter. The cost network is then trained with a Bradley-Terry ranking loss that pushes the score of negative bags above the score of unlabeled bags. The authors prove that this ordering holds with probability 1 − (1 − α)^K.

With a cost function that flags risky state-action pairs, the unlabeled dataset can be cleaned up. Two options are offered: a hard threshold that keeps only trajectories whose estimated cumulative cost falls below a limit, or a soft exponential weight that down-weights high-cost trajectories (and, in a variant, high-cost individual transitions). A behavior-cloning policy is then trained on this filtered or weighted data.

The authors also extend the method to training with partial trajectory segments of fixed length, arguing that mislabeled fragments matter less as bag size grows because the score reflects an average cost.

Why This Matters

Impact on research: The paper reframes safety cost learning as a ranking problem over groups of trajectories rather than a per-timestep labeling problem. This connects offline safe imitation learning to the well-developed MIL literature and shows that a small amount of "what not to do" data can substitute for costly per-step cost annotation. It also gives a theoretical characterization (via the bag-size probability bound) of when unlabeled bags can serve as positive examples.

Real-world applications:

• Autonomous driving: learn from a small set of crash, red-light-running, or near-miss trajectories plus a large unlabeled human-driving dataset, without hand-labeling safety cost per timestep.
• Conversational agents: treat reported toxic outputs as non-preferred trajectories and use the rest of the interaction logs as unlabeled data to steer the model away from toxic behavior.
• Robotics and surgical robotics: avoid costly or dangerous online exploration when defining safety constraints explicitly is difficult.
• Recommender and content systems: use flagged harmful sessions as non-preferred data while learning from general user logs.

Industry relevance: The setting matches industrial practice closely — companies typically have abundant unlabeled logs and a small, curated set of known-bad examples, but rarely have dense per-step reward and cost annotations. SafeMIL requires no online interaction and no reward/cost labels, which lowers the barrier to deploying imitation-based policies in regulated or high-risk products.

Future Directions

• Where the method currently trails: On Point-Circle2, T-REX-WBC learns a better constraint-satisfying policy, so understanding when the cost-ranking assumption breaks down is an open question.

• Assumption of uniform non-preferred cost: Both the score analysis and Theorem 1 assume non-preferred trajectories have similar cost. Relaxing this to heterogeneous risk levels is a natural extension.

• Scaling to richer data: The experiments use 50 non-preferred and 200 unlabeled trajectories; behavior on far larger, noisier datasets — and on higher-dimensional or image-based tasks — is not reported.

• Practical tuning questions: How to set the constraint threshold, the weighting temperature, and the mixing proportion in practice, and how sensitive results are to the number of non-preferred trajectories, are left for future work.

• Integration with constrained policy optimization: Combining the learned cost with offline safe RL objectives rather than pure behavior cloning could yield further safety gains.

Target Audience

This paper is best suited for graduate students and researchers working on safe reinforcement learning, offline RL, and imitation learning, as well as practitioners who need to train policies from logged data with limited safety annotation. Readers should already be comfortable with CMDPs, behavior cloning, and offline distribution-correction methods; the MIL machinery is explained from first principles but the paper moves quickly into formal statements and loss derivations.

Authors’ 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.