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

🕒 Published (v1): 2025-11-11 11:44 UTC · Source: Arxiv · Venue: AAAI 2026 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

TL;DR¶

SafeMIL addresses offline safe imitation learning without per-timestep reward or cost annotations by framing cost function learning as a Multiple Instance Learning (MIL) problem. It uses a small set of non-preferred (constraint-violating) trajectories and a large unlabeled dataset to learn a parameterized cost function, then applies weighted behavior cloning over trajectories the learned cost deems safe. On MuJoCo and navigation benchmarks, it achieves 3.7× better median safety performance than the best baseline.

Problem¶

Offline safe RL methods require per-timestep cost annotations, which are often impractical (e.g., toxicity scoring in conversational agents, surgical robotics). Standard imitation learning assumes demonstrations are safe. There is no method that learns a safe offline policy using only trajectory-level labels—specifically a small set of trajectories known to violate constraints and a larger unlabeled pool with unknown constraint status, and no reward or cost signal at any timestep.

Method¶

SafeMIL reframes cost-function learning as a MIL bag-classification problem. Bags are constructed by sampling K trajectories: negative bags draw exclusively from the non-preferred dataset D_N; unlabeled bags draw from D_U (which contains a mix). By Lemma 1, for sufficiently large K the unlabeled bag contains at least one preferred trajectory with probability 1−(1−α)^K, ensuring Score(B_n) > Score(B_u) with high probability.

The bag score uses the Deep Sets parameterization:

Score(B) = (1/K) Σ_{τ∈B} Σ_{t=0}^{T-1} γ^t ĉ_θ(s_t, a_t)

where ĉ_θ is the learned cost network. The cost network is trained via a Bradley-Terry pairwise loss pushing negative bag scores above unlabeled bag scores (Eq. 9). Partial trajectories of length H are supported to reduce compute, robust to mislabeling when K is large.

Policy learning uses soft trajectory-weighted behavior cloning over D_U: w(τ) = exp(−Σ_t γ^t ĉ_θ(s_t, a_t)/β), upweighting low-cost trajectories.

Key Contributions¶

First application of MIL to offline safe imitation learning; no per-timestep cost or reward labels required.
Theoretical guarantees: Lemma 1 and Theorem 1 quantify the probability that unlabeled bags contain preferred trajectories as a function of K and α (preferred fraction).
Partial-trajectory extension enabling computationally efficient training without degrading safety, stable across H ∈ {1, 5, 10} at large K.
Soft trajectory-weighted BC objective that leverages the full unlabeled dataset rather than hard threshold filtering.

Results¶

Velocity-constrained tasks (Walker2d, Swimmer, Ant): SafeMIL outperforms all baselines on normalized cost and CVaR@20% cost while matching or exceeding normalized return.
Navigation tasks (Point-Circle2, Point-Goal1, Point-Button1): SafeMIL achieves better safety on Point-Goal1; competitive with the best baseline (T-REX-WBC) on Point-Circle2 and Point-Button1.
Aggregate: Median safety performance is 3.7× better than the best baseline across all six environments (Appendix Tables 3, 4).
Baseline comparison: SafeDICE, DWBC-NU, T-REX-WBC, BC-Unlabeled; oracle upper-bound is COptiDICE with ground-truth reward+cost.
Sensitivity: Safety improves monotonically with bag size K on Swimmer-Velocity; safety is stable across trajectory lengths H ∈ {1, 5, 10} at K=128.

Limitations¶

Assumes non-preferred trajectories have similar cost magnitude; heterogeneous violation severities may break the score ordering guarantee.
Requires knowing/tuning the constraint threshold b̂ (hard filtering) or β (soft weighting) without access to ground-truth costs.
Lemma 1 requires knowing or estimating α (fraction of preferred trajectories in D_U); the paper does not discuss sensitivity to misspecification of α.
Evaluated only on MuJoCo locomotion and 2D navigation; generalization to high-dimensional state spaces (e.g., image-based robotics) is untested.
Bag size K and partial trajectory length H introduce additional hyperparameters not present in simpler baselines.

Relevance to Agentic AI / LLM Agents¶

SafeMIL directly addresses a core challenge for deploying LLM-based agents in safety-critical settings: acquiring safety constraints from naturally occurring failure examples (toxic outputs, policy violations) without requiring dense per-token cost annotation. The MIL framing—treating a conversation or action sequence as a "bag" and labeling only the full trajectory as non-preferred—maps naturally onto how human feedback is actually collected for conversational agents. This connects to the broader RLHF/RLAIF line where learning from negative demonstrations or refusals is preferable to engineering explicit scalar penalties. The offline, no-online-interaction requirement is also highly relevant for agents where environment rollouts are expensive or risky.