Source-linked AI summary
Verifiable Reinforcement Learning via Policy Extraction
Osbert Bastani, Yewen Pu, Armando Solar-Lezama
TL;DR
Deep reinforcement learning can solve difficult control problems, but the safety of learned DNN policies is difficult to verify efficiently. The paper introduces VIPER, which uses Q-function-guided imitation learning to extract decision trees from DNN oracles. These trees retain perfect task performance while supporting formal verification across several control settings.
Problem
The safety of deep reinforcement learning policies is difficult to ensure because verifying complex DNN policies is often inefficient or infeasible.
Method
VIPER uses imitation learning to extract decision tree policies from pretrained DNN oracles while leveraging their Q-functions.
Results
VIPER learns small decision trees that play perfectly on symbolic Atari Pong, toy Pong, and cart-pole, while enabling verification of robustness, correctness, and stability.
Takeaways & Limitations
Structured decision tree policies can preserve the performance of learned DNN controllers while making important behavioral properties efficiently verifiable.
Takeaways & Limitations
The approach requires polynomial dynamics and focuses on verifying stability locally around the origin, often using Taylor approximations.
Abstract
from arXiv · showhide
While deep reinforcement learning has successfully solved many challenging control tasks, its real-world applicability has been limited by the inability to ensure the safety of learned policies. We propose an approach to verifiable reinforcement learning by training decision tree policies, which can represent complex policies (since they are nonparametric), yet can be efficiently verified using existing techniques (since they are highly structured). The challenge is that decision tree policies are difficult to train. We propose VIPER, an algorithm that combines ideas from model compression and imitation learning to learn decision tree policies guided by a DNN policy (called the oracle) and its Q-function, and show that it substantially outperforms two baselines. We use VIPER to (i) learn a provably robust decision tree policy for a variant of Atari Pong with a symbolic state space, (ii) learn a decision tree policy for a toy game based on Pong that provably never loses, and (iii) learn a provably stable decision tree policy for cart-pole. In each case, the decision tree policy achieves performance equal to that of the original DNN policy.
1 Introduction
VIPER addresses the difficulty of safely verifying deep reinforcement learning policies by extracting structured decision trees from DNN oracles. It uses Q-function-guided imitation learning to produce smaller trees that retain performance while enabling verification of robustness, correctness, and stability.
- Motivation: Decision tree policies combine nonparametric representational capacity with structure that makes safety, stability, and robustness easier to verify.Existing verification approaches struggle to scale to high-dimensional state spaces and are often inefficient for DNN policies.
- Method: VIPER distills a pretrained DNN policy, called the oracle, into a decision tree using imitation learning and the oracle’s Q-function.The approach builds on DAGGER: the oracle generates labeled data, while Q-information guides extraction beyond the selected action alone.
- Method: VIPER learns decision trees an order of magnitude smaller than DAGGER’s, making them easier to verify.The oracle provides cumulative rewards for state-action pairs, which Q-DAGGER and VIPER exploit during policy extraction.
- Results: The extracted policies support verification of robustness for symbolic Atari Pong, non-loss for toy Pong, and a stability region for cart-pole.For Atari Pong, verification uses a symbolic abstraction; for cart-pole, stability is evaluated around the goal state.
- Results: VIPER learns relatively small decision trees with fewer than 1000 nodes that play perfectly on symbolic Atari Pong, toy Pong, and cart-pole.The reported policies match perfect reward across all three tasks.
2 Policy Extraction
The paper formulates policy extraction as learning a policy within a constrained class using an oracle DNN and its Q-function, then develops Q-DAGGER and VIPER for structured decision trees. Its loss prioritizes errors according to their impact on future performance, yielding bounds that can improve over prior DAGGER guarantees.
- Policy extraction framework: Q-DAGGER is a policy-extraction algorithm with theoretical guarantees, while VIPER adapts it to extract decision tree policies.VIPER uses decision-tree training within the Q-DAGGER policy-extraction framework.
- Policy extraction framework: The objective is to learn the best policy in a class Π by leveraging an oracle π∗ and its Q-function Q(π∗).The formulation uses a finite-horizon MDP with states, actions, transitions, and rewards.
- Theoretical guarantees: The proposed bound is O(T) when the extracted policy is accurate on critical states, compared with a prior O(T^2) bound regardless of that accuracy.The prior bound also contains u, which may be O(T) when an incorrect action forfeits all subsequent rewards.
- Policy extraction framework: The weighted loss can rank policies differently from 0-1 loss: πright has ℓ(πright) = O(T −1) while πleft has ℓ(πleft) = O(1).The corresponding 0-1 losses reverse this ordering: g(πleft) = O(T −1) and g(πright) = O(1).
- Decision tree extraction: VIPER trains decision trees heuristically with CART on a resampled dataset whose sampling weights encode the weighted imitation loss.The algorithm aggregates oracle-labeled trajectories, resamples state-action pairs proportionally to the loss, and selects the best policy by cross-validation.
3 Verification
The paper describes verification of correctness, stability, and robustness properties for decision-tree control policies using structured mathematical formulations and existing verification tools. These properties remain difficult to verify for DNN policies, while the presented approaches face scaling and modeling limitations.
- Correctness for toy Pong: Correctness for toy Pong is formulated as avoiding an unsafe region after starting from a safe state, under bounded vertical velocity assumptions.The proof uses an inductive invariant requiring the ball to re-enter the safe set within tmax = ⌈2ymax/vmin⌉ steps.
- Correctness for toy Pong: Piecewise-linear dynamics and decision-tree controllers partition the state space into regions where linear constraints encode system evolution and correctness.The resulting formulas combine conjunctions and disjunctions of linear inequalities, enabling satisfiability-based verification.
- Verification limitations: Verification has concrete scope boundaries: continuous-state correctness is not amenable to enumeration, and stability checks can scale exponentially with state dimension.The stability approach also requires polynomial dynamics and focuses on local verification around the origin, often using approximations.
- Correctness for cart-pole: For cart-pole, correctness requires the pole to remain within a specified angular deviation from upright over the first Tmax = 10 time steps.The nonlinear transition function is approximated by f(s, a) ≈ As + Ba, which is considered accurate when the pole-angle deviation is small.
- Stability: Stability is verified through a Lyapunov function whose positivity and decreasing derivative establish a region of attraction around the target.For polynomial dynamics, sum-of-squares optimization can construct the function, check the Lyapunov properties, and compute the region of attraction.
- Robustness: Decision-tree robustness at a state is computed by finding the smallest L∞ distance to a region routed to a different-action leaf.A linear program computes the distance from the reference state to each such leaf region, yielding the largest certified robustness radius.
4 Evaluation
VIPER extracts compact decision-tree policies that match strong DNN performance while enabling verification of correctness, robustness, and stability properties. Evaluation spans Atari Pong, toy Pong, cart-pole, and half-cheetah, with comparisons against fitted Q-iteration and DAGGER.
- Atari Pong: 21.0 reward is achieved by a 769-node VIPER tree on symbolic Atari Pong, matching the DQN oracle.Robustness is computed at five random states, with ε values from 0.5 to 2.8.
- Toy Pong: VIPER’s 31-node toy-Pong tree initially exposes a counterexample in which paddle oscillation can miss a ball near the screen edge.Adding a top-level safer-action node or extending the paddle length to L = 9/2 removes all additional counterexamples.
- Cart-pole: 200.0 reward is achieved by a 3-node VIPER tree on cart-pole, and Z3 verifies the desired safety property in 1.5 seconds.The extracted tree matches the oracle’s perfect reward of 200.0.
- Cart-pole stability: Neither the extracted nor original cart-pole policy is stable when trained only to keep the pole above a reward threshold.A separate iLQR-guided 3-node tree achieves reward 200.0, while SOS verification certifies stability for ||s||∞≤0.03 in 3.9 seconds.
- Baseline comparisons: After more than 200K rollouts, fitted Q-iteration achieves only 104.3 reward on cart-pole.The comparison uses reward as a function of training rollouts, including oracle-training rollouts for VIPER.
- Baseline comparisons: VIPER produces substantially smaller Atari Pong trees than DAGGER, including 769 versus 7967 nodes at reward 21.Other reported comparisons are 31 versus 127 nodes at reward 0 and 127 versus 3459 nodes at reward 20.
- Half-cheetah: A 9757-node half-cheetah regression tree achieves cumulative reward 4014, compared with 4189 for the neural-network oracle.This result demonstrates high-quality extraction but not parity with the oracle on that instance.
5 Conclusion
The paper presents efficiently verifiable decision-tree policies and identifies extensions needed to broaden their verification and repair capabilities.
- The approach learns decision-tree policies that can be verified efficiently.The authors identify wider verification tasks, reduced cart-pole approximations, automatic repair, and applications in safe reinforcement learning as future directions.
- Cart-pole correctness verification currently relies on several approximations, motivating invariant-set and piecewise-linear-bound alternatives.
- The paper considers a limited variety of verification tasks and leaves automatic repair of discovered policy errors for future work.
A.1 Proof of Lemma 2.1
The proof rewrites the relevant summand and applies an equality justified by a stated condition, yielding the desired form of the third term.
- The proof rewrites the summand of the third term into an equivalent expression.
- The final algebraic step follows from a separately stated equality.
- The proof invokes T(s) = 0 for all s ∈ S together with the definition of d(π).
A.2 Proof of Theorem 2.2
The theorem proof applies a high-probability bound and the assumption on N to obtain the target loss bound, then invokes Lemma 2.1.
- With probability at least 1−δ, the proof uses γ_N = Õ(1/N) to control the bound.
- The assumption on N yields ℓ(π) ≤ Tε_N + Õ(1).
- The result follows by applying Lemma 2.1.
A.3 Proof of Theorem 2.3
The proof characterizes the optimal policy by showing how it moves between states toward the higher-reward state ˜s before transitioning to send. It then compares πleft and πright through their disagreements with the optimal policy, yielding different g-values.
- Both πleft and πright reach send in k + 1 steps, after which the proof analyzes the optimal actions during those steps.
- The optimal policy moves down at s−(k−1), right at sk, and left otherwise during the relevant time steps.
- State ˜s is reachable from every si with −(k −1) ≤ i ≤ k −1 in fewer than 2(k + 1) steps, and only ˜s and sk provide nonzero rewards.
- Because ˜s yields higher reward than sk, π∗ generally moves toward ˜s, with specified exceptions at s−k, s−(k−1), sk, and send.
- πleft differs from π∗ only at s−(k−1), giving g(πleft) = T −1, whereas πright differs except at sk, giving g(πright) = 1 −T −1.