Learning When to Stop: Selective Imitation Learning Under Arbitrary Dynamics Shift
Pith reviewed 2026-05-20 21:59 UTC · model grok-4.3
The pith
A selective imitation learner stops in uncertain states under dynamics shift and achieves horizon-free sample complexity using a small fixed set of validator policies.
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
Core claim
The paper establishes that the SeqRejectron algorithm constructs a stopping rule from a validator set of size independent of the horizon and policy class, yielding a selective policy that is complete on the training distribution and sound on the test distribution under arbitrary dynamics shift. For deterministic policies and sparse costs this gives horizon-free sample complexity Õ(log|Π|/ε²). Analogous horizon-free guarantees hold for stochastic policies via a cumulative Hellinger stopping time. The same framework extends to misspecified experts and differing expert policies across environments, with performance degrading gracefully in the amount of misspecification.
What carries the argument
The SeqRejectron algorithm that builds a stopping rule from a small set of validator policies whose cardinality is bounded independently of horizon and policy class size.
If this is right
- For deterministic policies the sample complexity remains horizon-free and scales as Õ(log|Π|/ε²) under sparse costs.
- For stochastic policies analogous horizon-free bounds hold through a cumulative Hellinger stopping time.
- The framework extends to misspecified experts and different expert policies across train and test, with graceful degradation.
- The validator set size stays independent of horizon and policy class size.
Where Pith is reading between the lines
- The selective stopping mechanism could be paired with a separate fallback policy or human query to handle the states where the learner elects to stop.
- The validator-based construction may transfer to other sequential tasks that experience gradual distribution shift over long horizons.
- Empirical tests with controlled amounts of dynamics mismatch could quantify how often the learned policy chooses to stop as a function of shift severity.
Load-bearing premise
The costs are sparse and unlabeled state trajectories from the same expert are available in the test environment.
What would settle it
An experiment or counterexample in which costs have dense support or no unlabeled test trajectories are supplied, causing either the sample complexity to grow with the horizon or the low-regret guarantee before stopping to fail.
Figures
read the original abstract
Behavior cloning provides strong imitation learning guarantees when training and test environments share the same dynamics. However, in many deployment settings the test environment's transitions differ from training, and classical offline IL offers no recourse: the learner must commit to an action at every state, even when its demonstrations are uninformative and could lead to arbitrary degradation of performance. This motivates the study of selective imitation, where the learner may choose to stop when it cannot act reliably. We introduce a model for selective imitation under arbitrary dynamics shift: given labeled expert demonstrations from a training environment and unlabeled state trajectories from the same expert in a test environment, the learner outputs a selective policy that is complete (rarely stops in training) and sound (incurs low regret before stopping in test). Our algorithm, SeqRejectron, constructs a stopping rule using a small set of validator policies whose size is independent of the horizon or policy class. For deterministic policies, this yields horizon-free $\tilde{O}(\log|\Pi|/\epsilon^2)$ sample complexity, assuming sparse costs. For stochastic policies, we obtain analogous horizon-free guarantees using a cumulative Hellinger stopping time. We extend the framework to misspecified experts and different expert policies across train and test and obtain results that gracefully degrade with the amount of misspecification.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper introduces a model and algorithm (SeqRejectron) for selective imitation learning under arbitrary train-test dynamics shift. Given expert demonstrations from training and unlabeled state trajectories from the same expert in the test environment, the learner outputs a selective policy that is complete (rarely stops on training data) and sound (low regret before stopping on test data). The algorithm uses a small set of validator policies (size independent of horizon and policy class) to construct a stopping rule. For deterministic policies it claims horizon-free Õ(log|Π|/ε²) sample complexity under sparse costs; for stochastic policies it uses a cumulative Hellinger stopping time. Extensions to misspecified experts and differing train/test expert policies are also provided, with graceful degradation.
Significance. If the derivations and bounds hold, the work offers a principled way to handle arbitrary dynamics shift in imitation learning without committing to actions in uninformative states. The horizon-free sample complexity (conditional on sparse costs) and validator construction independent of horizon are technically notable strengths for long-horizon settings. The framework's extensions to misspecification add practical value, and the emphasis on both completeness and soundness provides clear, falsifiable guarantees.
major comments (2)
- [abstract and §4] The horizon-free Õ(log|Π|/ε²) bound for deterministic policies is stated to hold under the sparse-costs assumption (abstract and §4). The manuscript should explicitly derive or cite the step where sparsity prevents pre-stopping regret from scaling linearly with H under arbitrary shift; without this, the stopping-time regret bound appears to reintroduce horizon dependence, undermining the central claim.
- [§5] The soundness guarantee for the validator-based stopping rule (SeqRejectron) relies on unlabeled test trajectories from the expert. The paper should clarify whether the sample complexity remains horizon-free if these trajectories are unavailable or if the test expert policy differs substantially from training, as the extension in §5 appears to degrade only under the same sparsity premise.
minor comments (2)
- [§3] Notation for the validator set size and its independence from H and |Π| should be introduced earlier (e.g., before the main theorem) to improve readability.
- [§2] The abstract mentions 'sparse costs' without a formal definition; add a precise statement (e.g., per-step cost zero except on a small-measure set) in the model section.
Simulated Author's Rebuttal
We thank the referee for their careful reading and constructive feedback. We address each major comment below and indicate planned revisions to strengthen the clarity of the claims.
read point-by-point responses
-
Referee: [abstract and §4] The horizon-free Õ(log|Π|/ε²) bound for deterministic policies is stated to hold under the sparse-costs assumption (abstract and §4). The manuscript should explicitly derive or cite the step where sparsity prevents pre-stopping regret from scaling linearly with H under arbitrary shift; without this, the stopping-time regret bound appears to reintroduce horizon dependence, undermining the central claim.
Authors: We agree that an explicit derivation is needed for full transparency. Under the sparse-costs assumption, the analysis in the proof of the main theorem for deterministic policies bounds the pre-stopping regret by showing that costs are incurred only on a measure-zero or low-probability set of states; combined with the validator-based stopping rule, this ensures the accumulated regret before stopping remains independent of H even under arbitrary dynamics shifts. We will add a dedicated lemma and step-by-step derivation in the revised §4 (with a forward reference from the abstract) to make this explicit and cite the sparsity lemma directly. revision: yes
-
Referee: [§5] The soundness guarantee for the validator-based stopping rule (SeqRejectron) relies on unlabeled test trajectories from the expert. The paper should clarify whether the sample complexity remains horizon-free if these trajectories are unavailable or if the test expert policy differs substantially from training, as the extension in §5 appears to degrade only under the same sparsity premise.
Authors: The framework is defined to require unlabeled test trajectories from the expert for constructing the validator set and stopping rule; without them the algorithm cannot be executed as stated, so we make no horizon-free claim in that setting. For differing train/test expert policies, §5 already shows graceful degradation of the bounds while retaining horizon-freeness under sparsity. We will add an explicit clarifying paragraph in §5 and the discussion section stating the necessity of test trajectories and confirming that horizon-freeness holds only under the stated assumptions including sparsity. revision: yes
Circularity Check
No circularity: derivation relies on explicit assumptions and standard concentration
full rationale
The paper's central result (horizon-free Õ(log|Π|/ε²) sample complexity for deterministic policies) is stated to hold under the explicit assumption of sparse costs. The SeqRejectron algorithm constructs a stopping rule from a fixed-size set of validator policies, with the bound following from this construction plus standard concentration inequalities. No quoted step reduces a prediction to a fitted input by construction, nor does any load-bearing claim rest on a self-citation chain or self-definitional loop. The sparsity premise is introduced as an assumption that enables the horizon-free property rather than being smuggled in via prior work or renaming. The derivation is therefore self-contained.
Axiom & Free-Parameter Ledger
axioms (2)
- domain assumption Sparse costs assumption for deterministic policies
- domain assumption Availability of unlabeled state trajectories from the expert in the test environment
Reference graph
Works this paper leans on
-
[1]
Advances in Neural Information Processing Systems , volume=
Beyond perturbations: Learning guarantees with arbitrary adversarial test examples , author=. Advances in Neural Information Processing Systems , volume=
-
[2]
Advances in Neural Information Processing Systems , volume=
Is behavior cloning all you need? understanding horizon in imitation learning , author=. Advances in Neural Information Processing Systems , volume=
-
[3]
2017 IEEE/RSJ international conference on intelligent robots and systems (IROS) , pages=
Domain randomization for transferring deep neural networks from simulation to the real world , author=. 2017 IEEE/RSJ international conference on intelligent robots and systems (IROS) , pages=. 2017 , organization=
work page 2017
-
[4]
Proceedings of the IEEE/CVF international conference on computer vision , pages=
Exploring the limitations of behavior cloning for autonomous driving , author=. Proceedings of the IEEE/CVF international conference on computer vision , pages=
-
[5]
Advances in neural information processing systems , volume=
Alvinn: An autonomous land vehicle in a neural network , author=. Advances in neural information processing systems , volume=
-
[6]
Efficient reductions for imitation learning , author=. Proceedings of the thirteenth international conference on artificial intelligence and statistics , pages=. 2010 , organization=
work page 2010
-
[7]
A reduction of imitation learning and structured prediction to no-regret online learning , author=. Proceedings of the fourteenth international conference on artificial intelligence and statistics , pages=. 2011 , organization=
work page 2011
-
[8]
IEEE Transactions on information theory , volume=
On optimum recognition error and reject tradeoff , author=. IEEE Transactions on information theory , volume=. 2003 , publisher=
work page 2003
- [9]
-
[10]
Algorithmic Learning Theory , pages=
Efficient learning with arbitrary covariate shift , author=. Algorithmic Learning Theory , pages=. 2021 , organization=
work page 2021
-
[11]
Advances in Neural Information Processing Systems , volume=
Tolerant algorithms for learning with arbitrary covariate shift , author=. Advances in Neural Information Processing Systems , volume=
-
[12]
Advances in neural information processing systems , volume=
Selective classification for deep neural networks , author=. Advances in neural information processing systems , volume=
-
[13]
Journal of Computer and System Sciences , volume=
Efficient algorithms for online decision problems , author=. Journal of Computer and System Sciences , volume=. 2005 , publisher=
work page 2005
-
[14]
Behavioral Cloning from Observation
Behavioral cloning from observation , author=. arXiv preprint arXiv:1805.01954 , year=
work page internal anchor Pith review Pith/arXiv arXiv
-
[15]
Advances in Neural Information Processing Systems , volume=
Toward the fundamental limits of imitation learning , author=. Advances in Neural Information Processing Systems , volume=
-
[16]
Proceedings of the forty-eighth annual ACM symposium on Theory of Computing , pages=
The computational power of optimization in online learning , author=. Proceedings of the forty-eighth annual ACM symposium on Theory of Computing , pages=
-
[17]
Off-dynamics reinforcement learning: Training for transfer with domain classifiers , author=. arXiv preprint arXiv:2006.13916 , year=
-
[18]
Advances in Neural Information Processing Systems , volume=
Robust inverse reinforcement learning under transition dynamics mismatch , author=. Advances in Neural Information Processing Systems , volume=
-
[19]
International Conference on Machine Learning , pages=
Robust imitation learning against variations in environment dynamics , author=. International Conference on Machine Learning , pages=. 2022 , organization=
work page 2022
-
[20]
arXiv preprint arXiv:2002.11879 , year=
State-only imitation with transition dynamics mismatch , author=. arXiv preprint arXiv:2002.11879 , year=
-
[21]
Advances in Neural Information Processing Systems , volume=
An imitation from observation approach to transfer learning with dynamics mismatch , author=. Advances in Neural Information Processing Systems , volume=
-
[22]
arXiv preprint arXiv:2512.14895 , year=
Imitation Learning for Multi-turn LM Agents via On-policy Expert Corrections , author=. arXiv preprint arXiv:2512.14895 , year=
-
[23]
The International Journal of Robotics Research , volume=
Learning dexterous in-hand manipulation , author=. The International Journal of Robotics Research , volume=. 2020 , publisher=
work page 2020
-
[24]
2019 International Conference on Robotics and Automation (ICRA) , pages=
Safe reinforcement learning with model uncertainty estimates , author=. 2019 International Conference on Robotics and Automation (ICRA) , pages=. 2019 , organization=
work page 2019
-
[25]
Advances in Neural Information Processing Systems , volume=
Bridging offline reinforcement learning and imitation learning: A tale of pessimism , author=. Advances in Neural Information Processing Systems , volume=
-
[26]
Advances in neural information processing systems , volume=
Conservative q-learning for offline reinforcement learning , author=. Advances in neural information processing systems , volume=
-
[27]
Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems
Offline reinforcement learning: Tutorial, review, and perspectives on open problems , author=. arXiv preprint arXiv:2005.01643 , year=
work page internal anchor Pith review Pith/arXiv arXiv 2005
-
[28]
Efficient imitation under misspecification
Efficient imitation under misspecification , author=. arXiv preprint arXiv:2503.13162 , year=
-
[29]
International Conference on Algorithmic Learning Theory , pages=
On the hardness of domain adaptation and the utility of unlabeled target samples , author=. International Conference on Algorithmic Learning Theory , pages=. 2012 , organization=
work page 2012
-
[30]
arXiv preprint arXiv:2110.03239 , year=
Understanding domain randomization for sim-to-real transfer , author=. arXiv preprint arXiv:2110.03239 , year=
-
[31]
CAD2RL: Real Single-Image Flight without a Single Real Image
Cad2rl: Real single-image flight without a single real image , author=. arXiv preprint arXiv:1611.04201 , year=
work page internal anchor Pith review Pith/arXiv arXiv
-
[32]
Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages=
Sim-to-real via sim-to-sim: Data-efficient robotic grasping via randomized-to-canonical adaptation networks , author=. Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages=
-
[33]
2018 IEEE international conference on robotics and automation (ICRA) , pages=
Sim-to-real transfer of robotic control with dynamics randomization , author=. 2018 IEEE international conference on robotics and automation (ICRA) , pages=. 2018 , organization=
work page 2018
-
[34]
Sim-to-Real: Learning Agile Locomotion For Quadruped Robots
Sim-to-real: Learning agile locomotion for quadruped robots , author=. arXiv preprint arXiv:1804.10332 , year=
work page internal anchor Pith review Pith/arXiv arXiv
-
[35]
Solving Rubik's Cube with a Robot Hand
Solving rubik's cube with a robot hand , author=. arXiv preprint arXiv:1910.07113 , year=
work page internal anchor Pith review Pith/arXiv arXiv 1910
-
[36]
arXiv preprint arXiv:2502.12310 , year=
Domain randomization is sample efficient for linear quadratic control , author=. arXiv preprint arXiv:2502.12310 , year=
-
[37]
The Fourteenth International Conference on Learning Representations , year=
Statistical Guarantees for Offline Domain Randomization , author=. The Fourteenth International Conference on Learning Representations , year=
-
[38]
2019 international conference on robotics and automation (ICRA) , pages=
Closing the sim-to-real loop: Adapting simulation randomization with real world experience , author=. 2019 international conference on robotics and automation (ICRA) , pages=. 2019 , organization=
work page 2019
-
[39]
Preparing for the Unknown: Learning a Universal Policy with Online System Identification
Preparing for the unknown: Learning a universal policy with online system identification , author=. arXiv preprint arXiv:1702.02453 , year=
work page internal anchor Pith review Pith/arXiv arXiv
-
[40]
RMA: Rapid Motor Adaptation for Legged Robots
Rma: Rapid motor adaptation for legged robots , author=. arXiv preprint arXiv:2107.04034 , year=
work page internal anchor Pith review Pith/arXiv arXiv
-
[41]
Conference on Robot Learning , pages=
Active domain randomization , author=. Conference on Robot Learning , pages=. 2020 , organization=
work page 2020
-
[42]
Conference on robot learning , pages=
Sim-to-real robot learning from pixels with progressive nets , author=. Conference on robot learning , pages=. 2017 , organization=
work page 2017
-
[43]
Proceedings of the IEEE , volume=
A game theoretic approach to controller design for hybrid systems , author=. Proceedings of the IEEE , volume=. 2000 , publisher=
work page 2000
-
[44]
2019 18th European control conference (ECC) , pages=
Control barrier functions: Theory and applications , author=. 2019 18th European control conference (ECC) , pages=. 2019 , organization=
work page 2019
-
[45]
International workshop on hybrid systems: Computation and control , pages=
Safety verification of hybrid systems using barrier certificates , author=. International workshop on hybrid systems: Computation and control , pages=. 2004 , organization=
work page 2004
-
[46]
2017 IEEE 56th annual conference on decision and control (CDC) , pages=
Hamilton-jacobi reachability: A brief overview and recent advances , author=. 2017 IEEE 56th annual conference on decision and control (CDC) , pages=. 2017 , organization=
work page 2017
-
[47]
Mathematics of Operations Research , volume=
Robust dynamic programming , author=. Mathematics of Operations Research , volume=. 2005 , publisher=
work page 2005
-
[48]
Robust control of Markov decision processes with uncertain transition matrices , author=. Operations Research , volume=. 2005 , publisher=
work page 2005
-
[49]
Mathematics of Operations Research , volume=
Robust Markov decision processes , author=. Mathematics of Operations Research , volume=. 2013 , publisher=
work page 2013
-
[50]
International conference on machine learning , pages=
Robust adversarial reinforcement learning , author=. International conference on machine learning , pages=. 2017 , organization=
work page 2017
-
[51]
EPOpt: Learning Robust Neural Network Policies Using Model Ensembles
Epopt: Learning robust neural network policies using model ensembles , author=. arXiv preprint arXiv:1610.01283 , year=
work page internal anchor Pith review Pith/arXiv arXiv
-
[52]
Advances in neural information processing systems , volume=
Robust deep reinforcement learning against adversarial perturbations on state observations , author=. Advances in neural information processing systems , volume=
-
[53]
International Conference on Machine Learning , pages=
Action robust reinforcement learning and applications in continuous control , author=. International Conference on Machine Learning , pages=. 2019 , organization=
work page 2019
-
[54]
International Conference on Artificial Intelligence and Statistics , pages=
Sample complexity of robust reinforcement learning with a generative model , author=. International Conference on Artificial Intelligence and Statistics , pages=. 2022 , organization=
work page 2022
-
[55]
Journal of Machine Learning Research , volume=
Distributionally robust model-based offline reinforcement learning with near-optimal sample complexity , author=. Journal of Machine Learning Research , volume=
-
[56]
Advances in neural information processing systems , volume=
Generative adversarial imitation learning , author=. Advances in neural information processing systems , volume=
-
[57]
Learning Robust Rewards with Adversarial Inverse Reinforcement Learning
Learning robust rewards with adversarial inverse reinforcement learning , author=. arXiv preprint arXiv:1710.11248 , year=
work page internal anchor Pith review Pith/arXiv arXiv
-
[58]
arXiv preprint arXiv:1912.05032 , year=
Imitation learning via off-policy distribution matching , author=. arXiv preprint arXiv:1912.05032 , year=
- [59]
-
[60]
international conference on machine learning , pages=
Dropout as a bayesian approximation: Representing model uncertainty in deep learning , author=. international conference on machine learning , pages=. 2016 , organization=
work page 2016
-
[61]
Advances in neural information processing systems , volume=
Simple and scalable predictive uncertainty estimation using deep ensembles , author=. Advances in neural information processing systems , volume=
-
[62]
Uncertainty-Aware Reinforcement Learning for Collision Avoidance
Uncertainty-aware reinforcement learning for collision avoidance , author=. arXiv preprint arXiv:1702.01182 , year=
work page internal anchor Pith review Pith/arXiv arXiv
-
[63]
IEEE Robotics and Automation Letters , volume=
Safe planning in dynamic environments using conformal prediction , author=. IEEE Robotics and Automation Letters , volume=. 2023 , publisher=
work page 2023
-
[64]
Learning for Dynamics and Control Conference , pages=
Adaptive conformal prediction for motion planning among dynamic agents , author=. Learning for Dynamics and Control Conference , pages=. 2023 , organization=
work page 2023
-
[65]
Conformal prediction for stl runtime verification , author=. Proceedings of the ACM/IEEE 14th International Conference on Cyber-Physical Systems (with CPS-IoT Week 2023) , pages=
work page 2023
-
[66]
2024 IEEE 63rd Conference on Decision and Control (CDC) , pages=
Single trajectory conformal prediction , author=. 2024 IEEE 63rd Conference on Decision and Control (CDC) , pages=. 2024 , organization=
work page 2024
-
[67]
Foundations and Trends in Machine Learning , volume=
Conformal prediction: A gentle introduction , author=. Foundations and Trends in Machine Learning , volume=. 2023 , publisher=
work page 2023
-
[68]
Advances in neural information processing systems , volume=
Conformal prediction under covariate shift , author=. Advances in neural information processing systems , volume=
-
[69]
Advances in Neural Information Processing Systems , volume=
Adaptive conformal inference under distribution shift , author=. Advances in Neural Information Processing Systems , volume=
-
[70]
Journal of statistical planning and inference , volume=
Improving predictive inference under covariate shift by weighting the log-likelihood function , author=. Journal of statistical planning and inference , volume=. 2000 , publisher=
work page 2000
- [71]
-
[72]
A theory of learning from different domains , author=. Machine learning , volume=. 2010 , publisher=
work page 2010
-
[73]
Domain adaptation: Learning bounds and algorithms
Domain adaptation: Learning bounds and algorithms , author=. arXiv preprint arXiv:0902.3430 , year=
-
[74]
Reinforcement learning: An introduction , author=. 1998 , publisher=
work page 1998
-
[75]
Handbooks in operations research and management science , volume=
Markov decision processes , author=. Handbooks in operations research and management science , volume=. 1990 , publisher=
work page 1990
-
[76]
Efficient Discrepancy Testing for Learning with Distribution Shift , author=. NeurIPS , year=
-
[77]
37th Annual Conference on Learning Theory, COLT 2024 (to appear) , year=
Testable Learning with Distribution Shift , author=. 37th Annual Conference on Learning Theory, COLT 2024 (to appear) , year=
work page 2024
-
[78]
37th Annual Conference on Learning Theory, COLT 2024 (to appear) , year=
Learning Intersections of Halfspaces with Distribution Shift: Improved Algorithms and SQ Lower Bounds , author=. 37th Annual Conference on Learning Theory, COLT 2024 (to appear) , year=
work page 2024
-
[79]
International Conference on Machine Learning , pages=
Of moments and matching: A game-theoretic framework for closing the imitation gap , author=. International Conference on Machine Learning , pages=. 2021 , organization=
work page 2021
-
[80]
Advances in Neural Information Processing Systems , volume=
Minimax optimal online imitation learning via replay estimation , author=. Advances in Neural Information Processing Systems , volume=
discussion (0)
Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.