Technical Summary: Learning When to Stop: Selective Imitation Learning Under Arbitrary Dynamics Shift

1. Problem Formulation

The paper addresses the challenge of Imitation Learning (IL) in settings where the test environment's transition dynamics differ arbitrarily from the training environment. In classical offline IL, a learner trained on expert demonstrations from a source environment $M$ is deployed in a target environment $N$. If the state supports of $M$ and $N$ are disjoint or significantly shifted, a standard policy may perform arbitrarily poorly because it cannot recognize when its training data no longer provides reliable guidance.

The authors propose a framework for Selective Imitation Learning, where the learner is granted the option to abstain (stop execution) mid-trajectory when it detects that the current state is outside the reliable support of its training data. The goal is to learn a selective policy $(\pi, \tau)$, consisting of a base policy $\pi$ and a stopping time $\tau$, that satisfies two competing objectives:

1. Completeness: The policy should rarely abstain in the training environment $M$ (i.e., $\alpha_M(\pi, \tau) \le \epsilon$).
2. Soundness: The policy should incur low regret in the test environment $N$ before it stops (i.e., $Regret_N(\pi, \tau) \le \epsilon$).

The setting assumes access to:

• Labeled training rollouts: State-action pairs $(s, a)$ from the expert $\pi^\star$ in environment $M$.
• Unlabeled test rollouts: State-only trajectories from the expert $\pi^\star$ in environment $N$. This asymmetry models real-world scenarios where action telemetry is costly or proprietary (e.g., autonomous driving dashcam footage without steering commands).

2. Methodology: SeqRejectron

The core algorithm, SeqRejectron, generalizes the Rejectron algorithm from selective classification (Goldwasser et al., 2020) to sequential decision-making. The key insight is that instead of making independent decisions at each time step (which leads to a horizon-dependent sample complexity), the algorithm certifies entire trajectory prefixes using a sparse set of validator policies.

Core Mechanism

1. Base Policy: The learner selects a base policy $\pi_0$ consistent with the training data (e.g., Maximum Likelihood Estimation).
2. Validator Sets: The algorithm constructs a small set of "validator" policies $\Phi$ drawn from the version space (policies consistent with training data).
3. Stopping Rule: The learner executes $\pi_0$ and continues only as long as every validator in $\Phi$ agrees with $\pi_0$. The first point of disagreement triggers a stop (abstention).
4. Sparse Validator Distribution: A critical theoretical contribution is the proof that a sparse validator set (size independent of the horizon $H$) is sufficient to guarantee safety. This is achieved via a minimax argument over a zero-sum game between a "challenger" (picking a competing policy) and a "validator" (picking a set of validators).

Variants for Different Settings

The paper adapts this framework to three increasingly general settings:

• Deterministic Policies:

  • Mechanism: Uses a binary disagreement test. The stopping time $\tau$ is the first step where any validator $\pi \in \Phi$ disagrees with $\pi_0$.
  • Implementation: The validator distribution is constructed using no-regret dynamics (specifically, an oracle-efficient reduction to Weighted Realizable Multiple-Instance Learning).
  • Result: Achieves horizon-free sample complexity $\tilde{O}(\sqrt{\log|\Pi|/n})$ under sparse costs.

• Stochastic Policies:

  • Challenge: Binary disagreement is too coarse; distributions may differ "softly" at every step without a single alarming discrepancy.
  • Mechanism: Replaces the binary test with a cumulative Hellinger distance stopping time. The learner stops when the accumulated squared Hellinger distance between $\pi_0$ and any validator exceeds a threshold $\theta$.
  • Result: Provides finite-sample bounds on the Hellinger distance of stopped trajectories. The paper establishes an $\Omega(1/\epsilon^3)$ lower bound for labeled sample complexity in the single-step stochastic setting, matching the cost-driven part of the upper bound.

• Misspecified Policy Classes:

  • Challenge: The expert $\pi^\star$ may not lie within the learner's policy class $\Pi$.
  • Mechanism: Replaces the hard version space constraint with a symmetrically regularized game. Validators are allowed to disagree with the base policy on source data, but this disagreement is penalized.
  • Result: Guarantees degrade gracefully with the amount of misspecification (measured by the best-in-class misspecification error $\Delta$).

3. Key Contributions and Results

Theoretical Guarantees

• Horizon-Free Complexity: For deterministic policies, the sample complexity is independent of the horizon $H$ (except through the total cost scale $C_{max}$), scaling as $\tilde{O}(\sqrt{\log|\Pi|/n})$. This improves upon naive step-wise reductions which suffer an $O(H^2)$ penalty.
• Sparse Validation: The paper proves the existence of a validator distribution supported on sets of size $O(1/\epsilon)$, independent of the state space, action space, or horizon.
• Lower Bounds: The authors prove an $\Omega(1/\epsilon^3)$ lower bound for labeled sample complexity in the stochastic single-step setting, showing that the cubic dependence on $1/\epsilon$ is unavoidable for cost-driven regimes.
• Misspecification Robustness: The framework extends to agnostic settings where the expert is not in the class, with regret bounds that scale with the square root of the misspecification error.

Empirical Evaluation

The authors evaluate SeqRejectron in the LunarLander-v3 environment.

• Setup: Training occurs in a windless environment; testing occurs in a windy environment with different initial impulses.
• Findings: The algorithm successfully detects out-of-distribution dynamics. By tuning a safety threshold $\theta$, the authors demonstrate a tradeoff where the policy abstains minimally on in-distribution (training) data while reliably stopping on out-of-distribution (test) data, preventing catastrophic failures. The number of validators was found to be robust, with performance saturating quickly.

4. Significance and Claims

The paper claims to establish abstention as a principled, theoretically grounded response to unanticipated dynamics shifts in imitation learning.

• Beyond Heuristics: Unlike robust RL or domain randomization approaches that require assumptions about the nature of the shift or online access to the test distribution, SeqRejectron provides distribution-free PAC guarantees using only offline data.
• Sequential Analogue of PQ Learning: The work positions itself as a sequential analogue to the PQ learning framework (Goldwasser et al., 2020), extending the ability to abstain from single predictions to entire trajectory prefixes.
• Safety Mechanism: By decoupling the decision of what to do from when to stop, the framework offers a safety mechanism that prevents the learner from driving the system into unrecoverable states when its training data is uninformative.

The authors acknowledge limitations, noting that the stochastic setting currently lacks an oracle-efficient computational reduction (unlike the deterministic case) and that the variance-driven sample complexity rates for stochastic policies remain an open gap between upper and lower bounds. They also suggest future work could incorporate partial structural knowledge (e.g., bounded TV distance) to tighten test-side stopping rates.