# Neural-to-Tree Policy Distillation with Policy Improvement Criterion

arXiv: https://arxiv.org/abs/2108.06898 · HTML: https://mlsanity.com/paper/2108.06898/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 31 references that are in the map.
- First posted on arXiv 16 Aug 2021.

## Builds on (most cited first, 21 of 21)

- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Towards A Rigorous Science of Interpretable Machine Learning](https://mlsanity.com/paper/1702.08608/) (2017)
- [Reinforcement Learning with Deep Energy-Based Policies](https://mlsanity.com/paper/1702.08165/) (2017)
- [Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review](https://mlsanity.com/paper/1805.00909/) (2018)
- [Policy Distillation](https://mlsanity.com/paper/1511.06295/) (2015)
- [Distributed Prioritized Experience Replay](https://mlsanity.com/paper/1803.00933/) (2018)
- [Distilling a Neural Network Into a Soft Decision Tree](https://mlsanity.com/paper/1711.09784/) (2017)
- [Verifiable Reinforcement Learning via Policy Extraction](https://mlsanity.com/paper/1805.08328/) (2018)
- [Deep Reinforcement Learning for De-Novo Drug Design](https://mlsanity.com/paper/1711.10907/) (2017)
- [Explaining Image Classifiers by Counterfactual Generation](https://mlsanity.com/paper/1807.08024/) (2018)
- [Interpretable Machine Learning -- A Brief History, State-of-the-Art and Challenges](https://mlsanity.com/paper/2010.09337/) (2020)
- [CXPlain: Causal Explanations for Model Interpretation under Uncertainty](https://mlsanity.com/paper/1910.12336/) (2019)
- [Measuring Neural Net Robustness with Constraints](https://mlsanity.com/paper/1605.07262/) (2016)
- [Interpreting CNNs via Decision Trees](https://mlsanity.com/paper/1802.00121/) (2018)
- [Distilling Policy Distillation](https://mlsanity.com/paper/1902.02186/) (2019)
- [Scalable End-to-End Autonomous Vehicle Testing via Rare-event Simulation](https://mlsanity.com/paper/1811.00145/) (2018)
- [Explanation by Progressive Exaggeration](https://mlsanity.com/paper/1911.00483/) (2019)
- [Toward Interpretable Deep Reinforcement Learning with Linear Model U-Trees](https://mlsanity.com/paper/1807.05887/) (2018)
- [A Deep Value-network Based Approach for Multi-Driver Order Dispatching](https://mlsanity.com/paper/2106.04493/) (2021)
- [Robustness Verification of Tree-based Models](https://mlsanity.com/paper/1906.03849/) (2019)
- [Distilling Knowledge from Deep Networks with Applications to Healthcare Domain](https://mlsanity.com/paper/1512.03542/) (2015)

## Built on by (most cited first, 5 of 5)

- [Reinforcement Learning Fine-Tuning Enhances Activation Intensity and Diversity in the Internal Circuitry of LLMs](https://mlsanity.com/paper/2509.21044/) (2025)
- [Mitigating Information Loss in Tree-Based Reinforcement Learning via Direct Optimization](https://mlsanity.com/paper/2408.08761/) (2024)
- [Practical Knowledge Distillation: Using DNNs to Beat DNNs](https://mlsanity.com/paper/2302.12360/) (2023)
- [MSVIPER: Improved Policy Distillation for Reinforcement-Learning-Based Robot Navigation](https://mlsanity.com/paper/2209.09079/) (2022)
- [Fidelity-Induced Interpretable Policy Extraction for Reinforcement Learning](https://mlsanity.com/paper/2309.06097/) (2023)

## Abstract (arXiv)

While deep reinforcement learning has achieved promising results in challenging decision-making tasks, the main bones of its success --- deep neural networks are mostly black-boxes. A feasible way to gain insight into a black-box model is to distill it into an interpretable model such as a decision tree, which consists of if-then rules and is easy to grasp and be verified. However, the traditional model distillation is usually a supervised learning task under a stationary data distribution assumption, which is violated in reinforcement learning. Therefore, a typical policy distillation that clones model behaviors with even a small error could bring a data distribution shift, resulting in an unsatisfied distilled policy model with low fidelity or low performance. In this paper, we propose to address this issue by changing the distillation objective from behavior cloning to maximizing an advantage evaluation. The novel distillation objective maximizes an approximated cumulative reward and focuses more on disastrous behaviors in critical states, which controls the data shift effect. We evaluate our method on several Gym tasks, a commercial fight game, and a self-driving car simulator. The empirical results show that the proposed method can preserve a higher cumulative reward than behavior cloning and learn a more consistent policy to the original one. Moreover, by examining the extracted rules from the distilled decision trees, we demonstrate that the proposed method delivers reasonable and robust decisions.