Added belief space binning (#18)

* Add SparseArrays entry

* Added binning option and fixed max_steps stopping criteria

* Added binning

* Stopped updating the bin manager if use_binning was false

* added separate test sets for use with and without binning

* Removed unnecessary code

* added version of `entropy` for `::SparseVector`

* Updated binning method + format updates

* Update tests with new BinManager

Author: dylan-asmar

Project.toml

@@ -15,4 +15,5 @@ LinearAlgebra = "1"
 POMDPTools = "0.1, 1"
 POMDPs = "0.9, 1"
 Printf = "1"
+SparseArrays = "1"
 julia = "1.7"

src/sample.jl

@@ -18,8 +18,13 @@ function sample_points(sol::SARSOPSolver, tree::SARSOPTree, b_idx::Int, L, U, t,
     V̲, V̄ = tree.V_lower[b_idx], tree.V_upper[b_idx]
     γ = discount(tree)
 
-    V̂ = V̄ #TODO: BAD, binning method
-    if V̂ ≤ V̲ + sol.kappa*ϵ*γ^(-t) || (V̂ ≤ L && V̄ ≤ max(U, V̲ + ϵ*γ^(-t)))
+    if sol.use_binning
+        V̂ = get_bin_value(tree, b_idx)
+    else
+        V̂ = V̄
+    end
+    
+    if V̄ ≤ V̲ + sol.kappa*ϵ*γ^(-t) || (V̂ ≤ L && V̄ ≤ max(U, V̲ + ϵ*γ^(-t)))
         return
     else
         Q̲, Q̄, a′ = max_r_and_q(tree, b_idx)

src/solver.jl

@@ -9,6 +9,7 @@ Base.@kwdef struct SARSOPSolver{LOW,UP} <: Solver
     init_lower::LOW     = BlindLowerBound(bel_res = 1e-2)
     init_upper::UP      = FastInformedBound(bel_res=1e-2)
     prunethresh::Float64= 0.10
+    use_binning::Bool   = true
 end
 
 function POMDPTools.solve_info(solver::SARSOPSolver, pomdp::POMDP)
@@ -20,7 +21,7 @@ function POMDPTools.solve_info(solver::SARSOPSolver, pomdp::POMDP)
 
     t0 = time()
     iter = 0
-    while time()-t0 < solver.max_time && root_diff(tree) > solver.precision
+    while iter <= solver.max_steps && time()-t0 < solver.max_time && root_diff(tree) > solver.precision
         sample!(solver, tree)
         backup!(tree)
         prune!(solver, tree)

src/tree.jl

@@ -4,6 +4,72 @@ mutable struct PruneData
     prune_threshold::Float64
 end
 
+struct BinData
+    bin_value::Matrix{Float64}
+    bin_count::Matrix{Int}
+    bin_error::Matrix{Float64}
+end
+
+struct BinNode
+    key::Tuple{Int,Int}
+    prev_error::Float64
+end
+
+struct BinManager
+    lowest_ub::Float64
+    num_levels::Int
+    num_bins_per_level::Vector{Int}
+    bin_levels_intervals::Vector{NamedTuple{(:ub, :entropy),Tuple{Float64,Float64}}}
+    bin_levels_nodes::Vector{Dict{Int,BinNode}}
+    bin_levels::Vector{BinData}
+    previous_lowerbound::Dict{Int,Float64}
+end
+
+function BinManager(Vs_upper::Vector{Float64}, num_bins_per_level=[5, 10])
+    num_levels = length(num_bins_per_level)
+    lowest_ub = minimum(Vs_upper)
+    highest_ub = maximum(Vs_upper)
+
+    # [level][:ub|:entropy] => value
+    bin_levels_intervals = Vector{NamedTuple{(:ub, :entropy),Tuple{Float64,Float64}}}(undef, num_levels)
+
+    # [level][b_idx][:key|:prev_error] => (ub_interval_idx, entropy_interval_idx)|previous_error
+    bin_levels_nodes = Vector{Dict{Int,BinNode}}(undef, num_levels)
+
+    # [level][:bin_value|:bin_count|:bin_error][(ub_interval_idx, entropy_interval_idx)] => Float64|Int|Float64
+    bin_levels = Vector{BinData}(undef, num_levels)
+
+    num_states = length(Vs_upper)
+    max_e = max_entropy(num_states)
+    for level_i in 1:num_levels
+        num_bins = num_bins_per_level[level_i]
+
+        ub = (highest_ub - lowest_ub) / num_bins
+        ent = max_e / num_bins
+        bin_levels_intervals[level_i] = (ub=ub, entropy=ent)
+
+        bin_levels_nodes[level_i] = Dict{Int,BinNode}()
+
+        bin_levels[level_i] = BinData(
+            zeros(Float64, num_bins, num_bins), # bin_value
+            zeros(Int, num_bins, num_bins), # bin_count
+            zeros(Float64, num_bins, num_bins) # bin_error
+        )
+    end
+
+    previous_lowerbound = Dict{Int,Float64}() # b_idx => lowerbound
+
+    return BinManager(
+        lowest_ub,
+        num_levels,
+        num_bins_per_level,
+        bin_levels_intervals,
+        bin_levels_nodes,
+        bin_levels,
+        previous_lowerbound
+    )
+end
+
 struct SARSOPTree
     pomdp::ModifiedSparseTabular
 
@@ -20,7 +86,7 @@ struct SARSOPTree
 
     _discount::Float64
     is_terminal::BitVector
-    is_terminal_s::SparseVector{Bool, Int}
+    is_terminal_s::SparseVector{Bool,Int}
 
     #do we need both b_pruned and ba_pruned? b_pruned might be enough
     sampled::Vector{Int} # b_idx
@@ -32,20 +98,23 @@ struct SARSOPTree
     prune_data::PruneData
 
     Γ::Vector{AlphaVec{Int}}
+
+    use_binning::Bool
+    bm::BinManager
 end
 
 
-function SARSOPTree(solver, pomdp::POMDP)
+function SARSOPTree(solver, pomdp::POMDP; num_bins_per_level=[5, 10])
     sparse_pomdp = ModifiedSparseTabular(pomdp)
     cache = TreeCache(sparse_pomdp)
 
     upper_policy = solve(solver.init_upper, sparse_pomdp)
     corner_values = map(maximum, zip(upper_policy.alphas...))
 
-    tree = SARSOPTree(
-        sparse_pomdp,
+    bin_manager = BinManager(corner_values, num_bins_per_level)
 
-        Vector{Float64}[],
+    tree = SARSOPTree(
+        sparse_pomdp, Vector{Float64}[],
         Vector{Int}[],
         corner_values, #upper_policy.util,
         Float64[],
@@ -63,8 +132,10 @@ function SARSOPTree(solver, pomdp::POMDP)
         Vector{Int}(),
         BitVector(),
         cache,
-        PruneData(0,0,solver.prunethresh),
-        AlphaVec{Int}[]
+        PruneData(0, 0, solver.prunethresh),
+        AlphaVec{Int}[],
+        solver.use_binning,
+        bin_manager
     )
     return insert_root!(solver, tree, _initialize_belief(pomdp, initialstate(pomdp)))
 end
@@ -93,7 +164,7 @@ function insert_root!(solver, tree::SARSOPTree, b)
     pomdp = tree.pomdp
 
     Γ_lower = solve(solver.init_lower, pomdp)
-    for (α,a) ∈ alphapairs(Γ_lower)
+    for (α, a) ∈ alphapairs(Γ_lower)
         new_val = dot(α, b)
         push!(tree.Γ, AlphaVec(α, a))
     end
@@ -118,7 +189,7 @@ function update(tree::SARSOPTree, b_idx::Int, a, o)
     ba_idx = tree.b_children[b_idx][a]
     bp_idx = tree.ba_children[ba_idx][o]
     V̲, V̄ = if tree.is_terminal[bp_idx]
-        0.,0.
+        0.0, 0.0
     else
         lower_value(tree, tree.b[bp_idx]), upper_value(tree, tree.b[bp_idx])
     end
@@ -139,7 +210,7 @@ function add_belief!(tree::SARSOPTree, b, ba_idx::Int, o)
     push!(tree.is_terminal, terminal)
 
     V̲, V̄ = if terminal
-        0., 0.
+        0.0, 0.0
     else
         lower_value(tree, b), upper_value(tree, b)
     end
@@ -163,6 +234,9 @@ function fill_belief!(tree::SARSOPTree, b_idx::Int)
     else
         fill_populated!(tree, b_idx)
     end
+    if tree.use_binning
+        update_bin_node!(tree, b_idx)
+    end
 end
 
 """
@@ -186,8 +260,8 @@ function fill_populated!(tree::SARSOPTree, b_idx::Int)
             bp_idx, V̲, V̄ = update(tree, b_idx, a, o)
             b′ = tree.b[bp_idx]
             po = tree.poba[ba_idx][o]
-            Q̄ += γ*po*V̄
-            Q̲ += γ*po*V̲
+            Q̄ += γ * po * V̄
+            Q̲ += γ * po * V̲
         end
 
         Qa_upper[a] = Q̄
@@ -219,7 +293,7 @@ function fill_unpopulated!(tree::SARSOPTree, b_idx::Int)
         tree.ba_children[ba_idx] = ba_children
 
         n_b += N_OBS
-        pred = dropzeros!(mul!(tree.cache.pred, pomdp.T[a],b))
+        pred = dropzeros!(mul!(tree.cache.pred, pomdp.T[a], b))
         poba = zeros(Float64, N_OBS)
         Rba = belief_reward(tree, b, a)
 
@@ -230,15 +304,15 @@ function fill_unpopulated!(tree::SARSOPTree, b_idx::Int)
             # belief update
             bp = corrector(pomdp, pred, a, o)
             po = sum(bp)
-            if po > 0.
+            if po > 0.0
                 bp.nzval ./= po
                 poba[o] = po
             end
 
             bp_idx, V̲, V̄ = add_belief!(tree, bp, ba_idx, o)
 
-            Q̄ += γ*po*V̄
-            Q̲ += γ*po*V̲
+            Q̄ += γ * po * V̄
+            Q̲ += γ * po * V̲
         end
         Qa_upper[a] = Q̄
         Qa_lower[a] = Q̲
@@ -249,3 +323,140 @@ function fill_unpopulated!(tree::SARSOPTree, b_idx::Int)
     tree.V_lower[b_idx] = lower_value(tree, tree.b[b_idx])
     tree.V_upper[b_idx] = maximum(tree.Qa_upper[b_idx])
 end
+
+function initialize_bin_node!(tree::SARSOPTree, b_idx::Int)
+    lb_val = tree.V_lower[b_idx]
+    ub_val = tree.V_upper[b_idx]
+    node_entropy = entropy(tree.b[b_idx])
+
+    for level_i in 1:tree.bm.num_levels
+        ub_interval_idx = get_interval_idx(
+            ub_val, tree.bm.lowest_ub, tree.bm.bin_levels_intervals[level_i][:ub],
+            tree.bm.num_bins_per_level[level_i]
+        )
+
+        entropy_interval_idx = get_interval_idx(
+            node_entropy, 0.0, tree.bm.bin_levels_intervals[level_i][:entropy],
+            tree.bm.num_bins_per_level[level_i]
+        )
+
+        key = (ub_interval_idx, entropy_interval_idx)
+        prev_error = 0.0
+
+        bin_count = tree.bm.bin_levels[level_i].bin_count[ub_interval_idx, entropy_interval_idx]
+        if bin_count > 0
+            err = tree.bm.bin_levels[level_i].bin_value[ub_interval_idx, entropy_interval_idx] - lb_val
+            prev_error = err * err
+            tree.bm.bin_levels[level_i].bin_error[ub_interval_idx, entropy_interval_idx] += prev_error
+            value = (tree.bm.bin_levels[level_i].bin_value[ub_interval_idx, entropy_interval_idx] * bin_count + lb_val) / (bin_count + 1)
+            tree.bm.bin_levels[level_i].bin_count[ub_interval_idx, entropy_interval_idx] += 1
+        else
+            err = ub_val - lb_val
+            prev_error = err * err
+            tree.bm.bin_levels[level_i].bin_error[ub_interval_idx, entropy_interval_idx] = prev_error
+            value = lb_val
+            tree.bm.bin_levels[level_i].bin_count[ub_interval_idx, entropy_interval_idx] = 1
+        end
+        tree.bm.bin_levels[level_i].bin_value[ub_interval_idx, entropy_interval_idx] = value
+        tree.bm.bin_levels_nodes[level_i][b_idx] = BinNode(key, prev_error)
+
+    end
+    tree.bm.previous_lowerbound[b_idx] = lb_val
+end
+
+function update_bin_node!(tree::SARSOPTree, b_idx::Int)
+    lb_val = tree.V_lower[b_idx]
+    up_val = tree.V_upper[b_idx]
+
+    if !haskey(tree.bm.bin_levels_nodes[1], b_idx)
+        return initialize_bin_node!(tree, b_idx)
+    end
+
+    for level_i in 1:tree.bm.num_levels
+        node = tree.bm.bin_levels_nodes[level_i][b_idx]
+        key = node.key
+        ub_interval_idx, entropy_interval_idx = key
+        prev_error = 0.0
+
+        bin_count = tree.bm.bin_levels[level_i].bin_count[ub_interval_idx, entropy_interval_idx]
+        if bin_count == 1
+            err = up_val - lb_val
+            prev_error = err * err
+            tree.bm.bin_levels[level_i].bin_error[ub_interval_idx, entropy_interval_idx] = prev_error
+        else
+            err = tree.bm.bin_levels[level_i].bin_value[ub_interval_idx, entropy_interval_idx] - lb_val
+            tree.bm.bin_levels[level_i].bin_error[ub_interval_idx, entropy_interval_idx] -= node.prev_error
+            prev_error = err * err
+            tree.bm.bin_levels[level_i].bin_error[ub_interval_idx, entropy_interval_idx] += prev_error
+        end
+
+        tree.bm.bin_levels_nodes[level_i][b_idx] = BinNode(key, prev_error)
+        tree.bm.bin_levels[level_i].bin_value[ub_interval_idx, entropy_interval_idx] = (tree.bm.bin_levels[level_i].bin_value[ub_interval_idx, entropy_interval_idx] * bin_count + lb_val - tree.bm.previous_lowerbound[b_idx]) / bin_count
+    end
+    tree.bm.previous_lowerbound[b_idx] = lb_val
+end
+
+function get_bin_value(tree::SARSOPTree, b_idx::Int)
+
+    lb_val = tree.V_lower[b_idx]
+    ub_val = tree.V_upper[b_idx]
+
+    node = tree.bm.bin_levels_nodes[1][b_idx]
+    key = node.key
+    ub_interval_idx, entropy_interval_idx = key
+    if tree.bm.bin_levels[1].bin_count[ub_interval_idx, entropy_interval_idx] == 1
+        return ub_val
+    else
+        smallest_error = Inf
+        best_level = 0
+        best_key = key
+        for level_i in 1:tree.bm.num_levels
+            node = tree.bm.bin_levels_nodes[level_i][b_idx]
+            key = node.key
+            ub_interval_idx, entropy_interval_idx = key
+            if tree.bm.bin_levels[level_i].bin_error[ub_interval_idx, entropy_interval_idx] + 1e-10 < smallest_error
+                best_level = level_i
+                smallest_error = tree.bm.bin_levels[level_i].bin_error[ub_interval_idx, entropy_interval_idx]
+                best_key = key
+            end
+        end
+
+        best_ub_interval_idx, best_entropy_interval_idx = best_key
+        best_value = tree.bm.bin_levels[best_level].bin_value[best_ub_interval_idx, best_entropy_interval_idx]
+        if best_value > ub_val + 1e-10
+            return ub_val
+        elseif best_value + 1e-10 < lb_val
+            return lb_val
+        else
+            return best_value
+        end
+    end
+end
+
+function max_entropy(n::Int)
+    return -1 * ((1.0 / n) * log(1.0 / n)) * n
+end
+
+function entropy(b::AbstractVector)
+    ent = 0.0
+    for b_i in b
+        b_i > 0 && (ent -= b_i * log(b_i))
+    end
+    return ent
+end
+
+function entropy(b::SparseVector)
+    ent = 0.0
+    for b_i in b.nzval
+        ent -= b_i * log(b_i)
+    end
+    return ent
+end
+
+function get_interval_idx(value::Float64, lower::Float64, interval::Float64, num_intervals::Int)
+    if interval == 0.0
+        return 1
+    end
+    idx = Int(floor((value - lower) / interval) + 1)
+    return clamp(idx, 1, num_intervals)
+end