2608.07151-Interpretable reinforcement learning with decision-tree pruning

page_by_page

Video file (mp4)

In short

The hosts discuss a paper by Ringer and Tokic on making reinforcement learning policies interpretable by converting neural networks into decision trees and pruning them. They highlight auditable pruning, DACP's reward-aware method, and results showing pruned trees sometimes outperform teachers. They conclude interpretability is an incremental, transparent process.

Key concepts

Decision-tree pruning
A method to simplify a decision tree by removing branches or leaves, making it easier for humans to read. The paper uses three strategies: max-depth, max-impurity, and DACP, each with different criteria for what to cut, all guided by a reward guard to maintain performance.
DACP (Adaptive pruning with reward guard)
A pruning algorithm that tracks how often each node is visited during real rollouts and removes the least-used nodes, but only if the reward stays above a threshold. If a batch of cuts drops reward too much, it splits the batch and retries, like binary search.
Reward guard
A mechanism that re-executes the policy after each pruning edit and accepts the edit only if the reward doesn't drop below a dynamic threshold, controlled by tolerance and stability factors. This ensures each simplification step is safe and auditable.
Interpretability proxy
A measurable stand-in for how understandable a policy is. The paper uses leaf node count, assuming fewer leaves mean easier reading. However, they admit it's a proxy—a smaller tree isn't always clearer if attributes are confusing.

This episode discusses

Transcript

Introduction to the show: ident: Paper Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Next we'll be talking about the paper "Interpretable reinforcement learning with decision-tree pruning".

Jane: The paper was written by Mark Ringer and Michel Tokic from Ludwig-Maximilians-University Munich and Siemens AG.

Tom: Stay tuned as we take you through the paper and discuss its implications.

Paper summary: Tom: So we've seen the basics — the authors take a trained neural network policy, convert it into a decision tree, and then prune that tree down to something a human can actually read. What strikes me first is that they treat pruning as a series of auditable edits, not just one blunt operation.

Jane: Right, and that's a big shift from the usual "train a smaller model" approach. Here every cut is recorded, like a changelog for the policy, so you can see exactly which branch was removed and what it did to the reward.

Lu: I'd add that they don't stop at structural pruning. Their third algorithm, DACP, actually tracks how often each node is visited during real rollouts, then tries removing the least-used nodes while checking the reward stays above a threshold. That's much closer to how a human would simplify a rule set.

Meng: And they handle failed removals cleverly — if cutting a batch of nodes drops the reward too much, they split the batch in half and try again. That divide-and-conquer bit reminded me of binary search, but for pruning decisions.

Tom: Exactly. And the reward guard uses two knobs: a tolerance factor for how much reward you're willing to lose per step, and a stability factor for the absolute floor you'll accept. So it's not just "keep performance above X" — it's a dynamic target that adapts as you prune.

Jane: Their results show something I found genuinely surprising: in Acrobot, pruning actually improved the reward above the original teacher's level. They argue that's overfitting — the tree had memorized noise, and cutting those branches made it generalize better.

Lu: That happens in supervised learning all the time, but seeing it in a distilled RL policy is nice confirmation. Their table also shows the learner sometimes beats the teacher, like in LunarLander, where the teacher scored lower than its own decision tree.

Meng: The trade-off curves are the heart of the paper though. For every environment, there's a clear drop-off point where the reward suddenly collapses. Before that, you can shrink the tree substantially with almost no loss. That's exactly the sweet spot you'd want in practice.

Tom: And across most environments, DACP held up best. The authors make a strong claim from that: structural pruning alone isn't enough, you need those mid-pruning reward evaluations to guide your cuts.

Jane: But they're honest about the limits. They only distilled trees up to 1024 leaves, and they admit that leaf count is just a proxy — a smaller tree isn't always easier for a human to understand if the attributes are confusing.

Lalam: That's the part that connects to a bigger conversation. In safety-critical settings, you don't just need a policy that works — you need to explain to a regulator why it works. A pruning trace gives you a narrative: "we removed these branches, the reward stayed flat, then it dropped, so we stopped." That's auditability in a form that a human can actually hold onto.

Tom: And that's the real contribution, I think — not just smaller trees, but making the simplification process itself transparent. The paper calls it a transformation trajectory, which is a nice way to frame interpretability as something you can inspect, not just a final artifact.

Jane: Right, you can watch the CartPole tree shrink step by step, see which gray branches are about to disappear, and check the reward table alongside it. That's a lot more satisfying than staring at a black box.

Lu: One thing I'd push back on: the paper's future work section basically admits the leaf-count proxy hasn't been validated with actual human users. So we don't yet know if their "simplified" trees genuinely read better to people. That's the missing experiment.

Meng: Still, as a practical toolkit, this is ready to use. The pruning is post-hoc, it reuses standard libraries, and you get a clear record of every change. For someone shipping a real RL system, that's a huge step up from "here's a neural net, good luck."

Tom: And their open door for combining pruning strategies or testing other operators leaves a lot of room for follow-up work. I'd love to see a user study that pits a 64-leaf pruned tree against a 200-leaf one and asks which people find easier to verify.

Jane: Let's hope that's already in the works. For now, the takeaway from this paper is that interpretability isn't a one-shot deal — you can earn it incrementally, with a paper trail.

Page 1 of the paper: Tom: So the first page already sets up the whole problem. They say you can turn a neural network policy into a decision tree, but those trees often get so huge that nobody can actually read them. That's the gap they're attacking.

Jane: That's a big gap, yeah.

Tom: And there's a line that captures it perfectly: "interpretable policies must be compact enough to read and simulate, not only explicit in form." So you can have something explicit but still completely overwhelming for a human.

Jane: Exactly. So their answer is to prune the tree after it's extracted, but with guardrails. Every time you cut a branch, you re-run the policy, check the reward hasn't dropped too much, and only keep the edit if it passes.

Tom: Wait, so they literally re-run the whole policy for every candidate cut? That sounds expensive.

Jane: It is, and that's why they keep a record of the process. They say each accepted edit is recorded, yielding an auditable trail, so you can trace how the big messy tree became the small clean one.

Tom: That's the part I find genuinely new. It's a controlled edit process, not just a big chop. They even spell it out: "apply a candidate operator, re-execute the policy to measure task return and interpretability proxies, and accept the edit only if it passes a non-inferiority test."

Jane: So interpretability becomes a property of the transformation itself, not just the final artifact. That's a real shift from earlier work where you'd just train a smaller model directly.

Tom: And they're building on Kohler's distillation method, but deliberately restricting to tree policies to push interpretability further. On this page they mention using the actor network from stablebaselines3 and then benchmarking the same way Kohler does.

Jane: Right. The motivation is also spelled out clearly. They say this matters for settings where verification and accountability are required, so it's about trust, not just making things look tidy.

Tom: And they're honest about the trade-off. They claim pruning traces reveal consistent interpretability improvements while maintaining high performance, but that's a claim we'll need to see tested on the benchmarks.

Jane: One thing I appreciate is that they don't pretend pruning is free. The re-execution cost is right there in the method, and they accept it as the price of knowing each cut actually works.

Tom: Yes, and that careful, evidence-based approach is what might make these policies usable in real industrial settings. That's why Siemens is involved, after all.

Page 2 of the paper: Tom: So on page two we get into the actual machinery of how they simplify the trees. They take the trained neural network, run it to collect state-action pairs, and then fit a scikit-learn decision tree classifier on those pairs. That gives them a starting tree that's already a readable policy, though often a big one.

Jane: So the distillation comes first, and then the pruning actually shrinks it down.

Tom: Exactly. And they use leaf node count as their interpretability proxy. The fewer leaves, the easier it is to read and simulate. They chose leaf count because it's invariant to just renaming thresholds or reordering branches, so it captures complexity without getting tricked by syntactic changes.

Jane: But leaf count alone doesn't tell you which branches are actually important.

Tom: That's where the three pruning strategies come in. Max-depth just caps the tree depth and replaces everything beyond that with the majority class. Max-impurity uses the Gini index to stop splitting once a node is mostly one action. And then there's DACP, which counts how often each node gets visited during real runs and removes the least visited ones, but with a reward guard so you don't destroy performance.

Jane: So DACP is usage-aware, while the other two are purely structural. That's a key distinction.

Tom: Right. And every pruning step is followed by a subtree collapsing pass that merges sibling leaves with the same action. That's a pure cleanup that doesn't change the policy at all, just removes redundancy. They even mention that the distilling algorithm doesn't punish unnecessary splits, so you end up with uniform subtrees that can be collapsed for free.

Jane: And they record every accepted edit, so you can trace the whole simplification history.

Tom: Yes, that's the auditable trail. They re-execute the policy after each candidate edit, measure the reward, and only accept it if it doesn't drop below a threshold. That's what makes the transformation process itself transparent, not just the final tree. Earlier we talked about Kohler training compact policies directly, but here you can watch the tree shrink step by step and see exactly which cut caused what.

Jane: So page two is really laying out the toolkit. Three pruning strategies, a collapse operation, and a reward-based acceptance rule. The rest of the paper then compares those strategies on benchmarks.

Page 3 of the paper: Tom: So page three introduces the actual pruning strategies, and I have to say the max-impurity one is beautifully simple. It just checks how mixed the actions are in each node using that Gini formula, and if a node is mostly one action, it turns it into a leaf.

Jane: Right, so it's basically saying, "this branch isn't adding much information, let's stop splitting here." But then they go further with the adaptive pruning method, which is the clever one for me.

Tom: Oh, the DACP one? That's where the tree keeps track of how often each node is actually visited during a rollout.

Jane: Exactly. It prunes the least visited nodes first, but it's careful not to wreck the reward. They define a minimum acceptable reward for each step, and if pruning a batch drops below that, the batch gets split in half and they try again recursively.

Tom: So it's like a binary search for which nodes you can safely remove. That's a smart way to limit the number of expensive benchmark calls, because running the full policy evaluation is costly.

Jane: And they point out a real subtlety: you'd think a node visited rarely must be unimportant, but sometimes those rare nodes are the ones that handle critical edge cases. That's why they have those tolerance and stability factors guarding the reward.

Tom: Right, the tolerance controls how much you can lose per iteration, and the stability factor sets a hard floor so the policy never sinks too low. That gives you a clear audit trail of every edit.

Jane: Then at the start of the results, they mention capping the tree at 1024 leaf nodes for comparability. And I love that detail about MountainCar—it never even grew past 340 leaves no matter how large they allowed the tree to be.

Tom: That's a nice hint that the environment's intrinsic complexity matters more than the cap. So even before any pruning, the distilled tree sizes are telling you something about the task itself.

Page 4 of the paper: Jane: So this page is where the actual numbers come in, and the first thing that jumps out is Table 1, comparing the original neural network teachers against the distilled decision trees.

Tom: And the results aren't uniform at all — for Acrobot and CartPole the learner basically matches the teacher, but for HalfCheetah and Walker2D the distilled policy drops a lot.

Jane: Right, they say that's probably due to the limited number of leaf nodes in the transformation, which makes sense because those environments need finer control.

Tom: But then there's LunarLander, where the teacher actually scored 149 and the learner scored 233 — the tree policy beat its own teacher.

Jane: That's wild. They explain it as the teacher being overfitted, and pruning away some of that complexity made the policy generalize better.

Tom: Now they also define their own solved thresholds for Pendulum and Walker2D, since gymnasium doesn't provide them, and they're pretty careful about it.

Jane: For Pendulum they set it at -200, which is basically a near-upright pendulum accumulating only -1 per timestep, and for Walker2D they use 1500, meaning constant forward locomotion.

Tom: And then the observations from Figure 1, the reward-size trade-off curves — they see a broadly monotonic decrease in reward as pruning progresses.

Jane: But that's not the whole story, because they also find a distinct drop-off point where further pruning suddenly crashes performance.

Tom: Exactly, and in a few environments like Acrobot, pruning temporarily improves reward above the original policy, which ties back to that overfitting idea.

Jane: They also note only small deviations between the three pruning algorithms in the middle stages, but then conclude DACP performs superior over most environments overall.

Tom: That's a strong claim — structural pruning alone isn't enough, and the reward-aware backtracking in DACP is what makes the difference.

Jane: The limitations are honest too: they only went up to 1024 leaf nodes, so larger or smaller trees might behave differently.

Tom: And they admit the leaf-node count is just a proxy for interpretability — a bigger tree with clearer attributes might actually be easier for a human to read.

Jane: So this page really sets up the central tension of the whole paper: pruning helps a lot, but you can't prune too far, and the metric itself is still up for debate.

Page 5 of the paper: Tom: So once they actually run these three pruning strategies, the first surprise is how bloated the starting trees really are. They let every tree grow to a maximum of 1024 leaves, and even CartPole fills that completely, even though a tiny tree could solve the task. MountainCar, on the other hand, never goes beyond 340 leaves, no matter how much room it's given. So the distillation step itself is creating a lot of unnecessary structure.

Jane: That's a rough starting point for the pruning. And when they compare the distilled tree to the original neural network in Table 1, how much performance does the conversion cost?

Tom: For most environments the learner basically matches the teacher — CartPole gets 488 against the teacher's 500, and Pendulum and Swimmer are identical. But HalfCheetah drops from 8898 down to 5023, and Walker2d roughly halves from 3917 to 1815. The authors blame the leaf cap during transformation for those complex continuous tasks. And then there's LunarLander, where the tree actually beats its teacher, 233 versus 149, which they explain by overfitting in the original network.

Jane: That's the same pattern they see later in pruning, with Acrobot reward temporarily going above the original. So removing overfitted branches helps in both directions.

Tom: Exactly. In Figure 1 every environment shows a broadly monotonic decrease in reward as pruning progresses, but there's always a distinct drop-off point, and before that point the simplification is nearly free. On Acrobot, pruning even improves reward above the original policy. Between the three strategies, the deviations are small in the intermediate stages, but DACP wins on most environments — the paper's conclusion is that structural pruning alone isn't enough, and you need the reward-guided backtracking. And they had to invent solved thresholds for Pendulum and Walker2d since gymnasium doesn't provide them, so some of those comparisons rest on their own definitions.

Jane: And they're careful about what the leaf count actually measures. They flag that a larger tree with clearly understandable attributes could be more readable than a smaller one, so the proxy has real limits. That's why they point to user studies as the necessary next step, to check whether these pruned trees are actually easier for humans to follow. Their trace for LunarLanderContinuous in Table 2 does show the pruning trajectory step by step, with reward fluctuating around the teacher's level, so you can audit each accepted edit — but whether a person can read those trees and predict the agent's next action is still an open question.

Tom: Right, and that's the honest limitation on this page. The visual trail, like the CartPole tree going from eight leaves to six while holding performance, is genuinely useful for inspection. But interpretability measured by leaf count is not the same as interpretability measured by a human sitting down and reading the rules.

Conclusion: Tom: So the big takeaway here is that they’ve turned the pruning process itself into something you can inspect, not just the final policy.

Jane: Exactly. Every time they cut a branch, they re-run the policy, measure the reward, and only keep the cut if performance holds up. That gives you a trail of changes.

Tom: And that trail matters, because if you’re putting these policies in a real system, you need to know why it got simpler and whether any decision logic actually changed.

Jane: Right. They found that in most environments you can shrink the tree quite a bit before you hit a cliff where reward drops sharply. That suggests there’s a comfortable middle ground.

Tom: What I liked was that pruning even improved things in a couple of cases, like Acrobot. Overfitting branches got removed and the policy generalized better.

Jane: Yeah, that’s a nice counterintuitive result. Simpler didn’t just mean more readable, it sometimes meant more robust.

Tom: Though they’re careful to say that leaf count is just a proxy. A smaller tree isn’t automatically easier for a human to follow.

Jane: Right, so they’re calling for user studies to actually test how people read these trees. That feels like the natural next step.

Tom: For now, the work gives engineers a practical way to take a black-box neural policy and turn it into something a human can audit and even edit.

Jane: And the auditable edit trail is the part that could make a difference in regulated settings, where you need to document every change you make to a model.

Tom: Realistically, this won’t replace neural policies in every application. But for tasks where accountability matters, this is a meaningful step.

Jane: I’m glad we covered it. Let’s move on to the next paper.

More episodes

← Home