Probabilistic Labelling for ProductProcesses

Project.toml

@@ -8,6 +8,7 @@ CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 HiGHS = "87dc4568-4c63-4d18-b0c0-bb2238e4078b"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
+JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
@@ -30,6 +31,7 @@ GPUArrays = "10, 11"
 HiGHS = "1.19.0"
 JSON = "0.21.4"
 JuMP = "1.29.0"
+JuliaFormatter = "2.1.6"
 LLVM = "7, 8, 9"
 NCDatasets = "0.13, 0.14"
 StyledStrings = "1.11.0"

src/bellman.jl

@@ -217,6 +217,30 @@ function bellman!(
     lf = labelling_function(model)
     dfa = automaton(model)
 
+    return _bellman_helper!(
+        workspace,
+        strategy_cache,
+        Vres,
+        V,
+        dfa,
+        mp,
+        lf,
+        upper_bound,
+        maximize,
+    )
+end
+
+function _bellman_helper!(
+    workspace::ProductWorkspace,
+    strategy_cache::AbstractStrategyCache,
+    Vres,
+    V,
+    dfa::DFA,
+    mp::IntervalMarkovProcess,
+    lf::DeterministicLabelling,
+    upper_bound = false,
+    maximize = true,
+)
     W = workspace.intermediate_values
 
     @inbounds for state in dfa
@@ -244,6 +268,57 @@ function bellman!(
     return Vres
 end
 
+function _bellman_helper!(
+    workspace::ProductWorkspace,
+    strategy_cache::AbstractStrategyCache,
+    Vres,
+    V::AbstractArray{R},
+    dfa::DFA,
+    mp::IntervalMarkovProcess,
+    lf::ProbabilisticLabelling,
+    upper_bound = false,
+    maximize = true,
+) where {R}
+    W = workspace.intermediate_values
+
+    @inbounds for state in dfa
+        local_strategy_cache = localize_strategy_cache(strategy_cache, state)
+
+        # Select the value function for the current DFA state
+        # according to the appropriate DFA transition function
+        map!(W, CartesianIndices(state_values(mp))) do idx
+            v = zero(R)
+
+            for (label, prob) in enumerate(lf[idx])
+                new_dfa_state = dfa[state, label]
+                v += prob * V[idx, new_dfa_state]
+            end
+
+            return v
+
+            # A = lf[idx]     #(Lx1) for IMDP state s', get all the p(f | s') for each label 
+            # B = V[idx, :]   #(Mx1) fir IMDP state s', get all V(s', z') for each DFA state
+            # C = onehot(transition(dfa))[state, :, :] #(LxM)
+
+            # return B' * C' * A
+        end
+
+        # For each state in the product process, compute the Bellman operator
+        # for the corresponding Markov process
+        bellman!(
+            workspace.underlying_workspace,
+            local_strategy_cache,
+            selectdim(Vres, ndims(Vres), state),
+            W,
+            mp;
+            upper_bound = upper_bound,
+            maximize = maximize,
+        )
+    end
+
+    return Vres
+end
+
 function localize_strategy_cache(strategy_cache::NoStrategyCache, dfa_state)
     return strategy_cache
 end

src/models/ProductProcess.jl

@@ -55,7 +55,7 @@ function checkproduct(
 )
 
     # check labelling states (input) match MDP states
-    if size(labelling_func) != state_values(mdp)
+    if state_values(labelling_func) != state_values(mdp)
         throw(
             DimensionMismatch(
                 "The mapped states $(size(labelling_func)) in the labelling function is not equal the fRMDP state variables $(state_values(mdp)).",

src/probabilities/DeterministicLabelling.jl

@@ -0,0 +1,84 @@
+"""
+    struct DeterministicLabelling{
+        T  <: Integer, 
+        AT <: AbstractArray{T}
+    }
+
+A type representing the labelling of IMDP states into DFA inputs.
+
+Formally, let ``L : S \\to 2^{AP}`` be a labelling function, where 
+- ``S`` is the set of IMDP states, and
+- ``2^{AP}`` is the power set of atomic propositions
+
+Then the ```DeterministicLabelling``` type is defined as vector which stores the mapping. 
+
+### Fields
+- `map::AT`: mapping function where indices are (factored) IMDP states and stored values are DFA inputs.
+- `num_outputs::Int32`: number of labels accounted for in mapping.
+
+"""
+struct DeterministicLabelling{T <: Integer, AT <: AbstractArray{T}} <: AbstractLabelling
+    map::AT
+    num_outputs::Int32
+
+    function DeterministicLabelling(map::AT) where {T <: Integer, AT <: AbstractArray{T}}
+        num_outputs = checklabelling(map)
+
+        return new{T, AT}(map, Int32(num_outputs))
+    end
+end
+
+function checklabelling(map::AbstractArray{<:Integer})
+    labels = unique(map)
+
+    if any(labels .< 1)
+        throw(ArgumentError("Labelled state index cannot be less than 1"))
+    end
+
+    # Check that labels are consecutive integers
+    sort!(labels)
+    if any(diff(labels) .!= 1)
+        throw(ArgumentError("Labelled state indices must be consecutive integers"))
+    end
+
+    return last(labels)
+end
+
+"""
+    mapping(dl::DeterministicLabelling)
+
+Return the mapping array of the labelling function. 
+"""
+mapping(dl::DeterministicLabelling) = dl.map
+
+"""
+    size(dl::DeterministicLabelling)
+
+Returns the shape of the input range of the labeling function ``L : S \\to 2^{AP}``, which can be multiple dimensions in case of factored IMDPs. 
+"""
+Base.size(dl::DeterministicLabelling) = size(dl.map)
+
+"""
+    num_labels(dl::DeterministicLabelling)
+Return the number of labels (DFA inputs) in the labelling function.
+"""
+num_labels(dl::DeterministicLabelling) = dl.num_outputs
+
+"""
+    state_values(dl::DeterministicLabelling)
+Return a tuple with the number of states for each state variable of the labeling function ``L : S \\to 2^{AP}``, which can be multiple dimensions in case of factored IMDPs. 
+"""
+state_values(dl::DeterministicLabelling) = size(dl.map)
+
+"""
+    num_states(dl::DeterministicLabelling)
+Return the number of states of the labeling function ``L : S \\to 2^{AP}``
+"""
+num_states(dl::DeterministicLabelling) = prod(state_values(dl))
+
+"""
+    getindex(dl::DeterministicLabelling, s...)
+
+Return the label for state s.
+"""
+Base.getindex(dl::DeterministicLabelling, s...) = dl.map[s...]

src/probabilities/Labelling.jl

@@ -1,74 +1,6 @@
-abstract type AbstractLabelling end
-
-"""
-    struct LabellingFunction{
-        T  <: Integer, 
-        AT <: AbstractArray{T}
-    }
-
-A type representing the labelling of IMDP states into DFA inputs.
-
-Formally, let ``L : S \\to 2^{AP}`` be a labelling function, where 
-- ``S`` is the set of IMDP states, and
-- ``2^{AP}`` is the power set of atomic propositions
-
-Then the ```LabellingFunction``` type is defined as vector which stores the mapping. 
-
-### Fields
-- `map::AT`: mapping function where indices are (factored) IMDP states and stored values are DFA inputs.
-- `num_inputs::Int32`: number of IMDP states accounted for in mapping.
-
-"""
-struct LabellingFunction{T <: Integer, AT <: AbstractArray{T}} <: AbstractLabelling
-    map::AT
-    num_outputs::Int32
-end
-
-function LabellingFunction(map::AT) where {T <: Integer, AT <: AbstractArray{T}}
-    num_outputs = checklabelling(map)
-
-    return LabellingFunction(map, Int32(num_outputs))
-end
-
-function checklabelling(map::AbstractArray{<:Integer})
-    labels = unique(map)
-
-    if any(labels .< 1)
-        throw(ArgumentError("Labelled state index cannot be less than 1"))
-    end
-
-    # Check that labels are consecutive integers
-    sort!(labels)
-    if any(diff(labels) .!= 1)
-        throw(ArgumentError("Labelled state indices must be consecutive integers"))
-    end
-
-    return last(labels)
-end
-
-"""
-    mapping(labelling_func::LabellingFunction)
-
-Return the mapping array of the labelling function. 
-"""
-mapping(labelling_func::LabellingFunction) = labelling_func.map
-
-"""
-    size(labelling_func::LabellingFunction)
-
-Returns the shape of the input range of the labeling function ``L : S \\to 2^{AP}``, which can be multiple dimensions in case of factored IMDPs. 
-"""
-Base.size(labelling_func::LabellingFunction) = size(labelling_func.map)
-
-"""
-    num_labels(labelling_func::LabellingFunction)
-Return the number of labels (DFA inputs) in the labelling function.
 """
-num_labels(labelling_func::LabellingFunction) = labelling_func.num_outputs
+    AbstractLabelling
 
+An abstract type for labelling functions.
 """
-    getindex(lf::LabellingFunction, s...)
-
-Return the label for state s.
-"""
-Base.getindex(lf::LabellingFunction, s...) = lf.map[s...]
+abstract type AbstractLabelling end

src/probabilities/ProbabilisticLabelling.jl

@@ -0,0 +1,84 @@
+"""
+    struct ProbabilisticLabelling{
+        R <: Real, 
+        MR <: AbstractMatrix{R}
+    }
+
+A type representing the Probabilistic labelling of IMDP states into DFA inputs. Each labelling is assigned a probability.
+
+Formally, let ``L : S \\times 2^{AP} \\to [0, 1]`` be a labelling function, where 
+- ``S`` is the set of IMDP states, and
+- ``2^{AP}`` is the power set of atomic propositions
+
+Then the ```ProbabilisticLabelling``` type is defined as matrix which stores the mapping. 
+
+### Fields
+- `map::MT`: mapping function encoded as matrix with labels on the rows, IMDP states on the columns, and valid probability values for the destination.
+
+The choice to have labels on the rows is due to the column-major storage of matrices in Julia and the fact that we want the inner loop over DFA target states 
+in the Bellman operator `bellman!`.
+
+"""
+struct ProbabilisticLabelling{R <: Real, MR <: AbstractMatrix{R}} <: AbstractLabelling
+    map::MR
+
+    function ProbabilisticLabelling(map::MR) where {R <: Real, MR <: AbstractMatrix{R}}
+        checklabellingprobs(map)
+
+        return new{R, MR}(map)
+    end
+end
+
+function checklabellingprobs(map::AbstractMatrix{<:Real})
+
+    # check for each state, all the labels probabilities sum to 1
+    if any(sum(map; dims=1) .!= 1)
+        throw(
+            ArgumentError(
+                "For each IMDP state, probabilities over label states must sum to 1",
+            ),
+        )
+    end
+end
+
+"""
+    mapping(pl::ProbabilisticLabelling)
+
+Return the mapping matrix of the probabilistic labelling function. 
+"""
+mapping(pl::ProbabilisticLabelling) = pl.map
+
+Base.size(pl::ProbabilisticLabelling) = size(pl.map)
+Base.size(pl::ProbabilisticLabelling, i) = size(pl.map, i)
+
+"""
+    getindex(pl::ProbabilisticLabelling, s, l)
+
+Return the probabilities for labelling l from state s.
+"""
+Base.getindex(pl::ProbabilisticLabelling, s, l) = pl.map[l, s]
+
+"""
+    getindex(pl::ProbabilisticLabelling, s)
+
+Return the probabilities over labels from state s.
+"""
+Base.getindex(pl::ProbabilisticLabelling, s) = @view(pl.map[:, s])
+
+"""
+    num_labels(pl::ProbabilisticLabelling)
+Return the number of labels (DFA inputs) in the probabilistic labelling function.
+"""
+num_labels(pl::ProbabilisticLabelling) = size(pl.map, 1)
+
+"""
+    state_values(pl::ProbabilisticLabelling)
+Return a tuple with the number of states for each state variable of the labeling function ``L : S \\to 2^{AP}``, which can be multiple dimensions in case of factored IMDPs. 
+"""
+state_values(pl::ProbabilisticLabelling) = Base.tail(size(pl.map))
+
+"""
+    num_states(pl::ProbabilisticLabelling)
+Return the number of states of the labeling function ``L : S \\to 2^{AP}``
+"""
+num_states(pl::ProbabilisticLabelling) = prod(state_values(pl))

src/probabilities/TransitionFunction.jl

@@ -77,3 +77,23 @@ num_states(tf::TransitionFunction) = size(tf.transition, 2)
 Return the number of labels (DFA inputs) in the transition function.
 """
 num_labels(tf::TransitionFunction) = size(tf.transition, 1)
+
+"""
+    onehot(tf:TransitionFunction)
+Return the transition function matrix one hot encoded.
+"""
+function onehot(tf::TransitionFunction)
+    nz = num_states(tf)
+    nl = num_labels(tf)
+
+    W = zeros(nz, nl, nz) #(|z| x |l| x |z'|)
+
+    for z in 1:nz
+        for l in 1:nl
+            y = getindex(tf, z, l)
+            W[z, l, y] = 1
+        end
+    end
+
+    return W
+end

src/probabilities/probabilities.jl

@@ -58,4 +58,9 @@ include("TransitionFunction.jl")
 export TransitionFunction, transition
 
 include("Labelling.jl")
-export LabellingFunction, mapping, num_labels
+
+include("DeterministicLabelling.jl")
+export DeterministicLabelling, mapping, num_labels, state_values, num_states
+
+include("ProbabilisticLabelling.jl")
+export ProbabilisticLabelling, mapping, num_labels, state_values, num_states