From a7178b7083d0bae6747a86c1cd587130a75e4df1 Mon Sep 17 00:00:00 2001
From: Ben Kraske <71412733+bkraske@users.noreply.github.com>
Date: Sat, 22 Jun 2024 14:28:38 -0600
Subject: [PATCH] Fix belief pruning (#17)

* Change all_ba_pruned flag

* Actually change all_ba_pruned flag
---
 src/prune.jl | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/src/prune.jl b/src/prune.jl
index ad33a52..7f19834 100644
--- a/src/prune.jl
+++ b/src/prune.jl
@@ -33,18 +33,17 @@ function prune!(tree::SARSOPTree)
             Qa_upper = tree.Qa_upper[b_idx]
             Qa_lower = tree.Qa_lower[b_idx]
             b_children = tree.b_children[b_idx]
-            ba = tree.b_children[b_idx]
             max_lower_bound = maximum(Qa_lower)
+            all_ba_pruned = true
             for (idx, Qba) ∈ enumerate(Qa_upper)
                 ba_idx = b_children[idx]
-                all_ba_pruned = true
-                if !tree.ba_pruned[ba_idx] && Qa_upper[idx] < max_lower_bound
+                if !tree.ba_pruned[ba_idx] && Qba < max_lower_bound
                     pruneSubTreeBa!(tree, ba_idx)
                 else
                     all_ba_pruned = false
                 end
-                all_ba_pruned && (tree.b_pruned[b_idx] = true)
             end
+            all_ba_pruned && (tree.b_pruned[b_idx] = true)
         end
     end
 end