Control Based Implementations

[rsenne]

Hi @gdalle — happy new year!

I wanted to get your thoughts on a design question that came up while working on emission models that require exogenous control inputs. The most obvious example is GLMs (and more generally regression-based emissions), but the same issue arises for a broader class of conditional or input-driven emissions.

I really like the way HiddenMarkovModels.jl already exposes control information via the control_seq interface. That said, I’ve found that under the current setup, emission models that depend on control can be a bit awkward to implement in a way that feels fully “first-class,” especially when it comes to integration with Baum–Welch and the emission update interface.

As far as I can see, there are a few possible directions one could take here:

1.) Specialized controlled HMMs
Implement control-aware HMMs with the emission model effectively baked in (e.g. Gaussian GLM-HMM, Bernoulli GLM-HMM, etc.), similar in spirit to the current examples in the documentation. This is straightforward conceptually, but it doesn’t scale particularly well as the space of emission models grows.

2.) Wrapping observations to carry control information
Make emission models operate on observation–control pairs by introducing a lightweight wrapper, e.g

struct GLMObs{T}
    X::Vector{T}
    y::T
end

and then pass this wrapper as the observation type throughout the HMM machinery. This works with the existing interfaces, but it pushes control information into the observation itself, which may not always be desirable from a modeling or ergonomics standpoint.

3.) Extending HiddenMarkovModels.jl to natively support control-aware emissions via multiple dispatch
Introduce an additional (optional) control-aware path upstream, so that when a control_seq is provided, both likelihood evaluation and emission updates can dispatch to methods of the form

logdensityof(emission, obs; control_seq=...)
fit!(emission, obs_seq, weights; control_seq=...)

while preserving the current behavior for emissions that do not require control. This would keep emissions as first-class objects and avoid baking control into the observation representation.

I’m curious what your perspective is on this design space, and whether you see a preferred direction (or an alternative I may be missing). I’d be happy to explore a concrete proposal or PR if you think it’s worth pursuing. I have a good amount of time to luckily spend on this given I'm on break from Ph.D. life and now have some time to pursue passion projects (another PR is on my list for HiddenMarkovModels.jl lol). Let me know what you think! And again happy new year! :)

[gdalle]

Hi Ryan, happy new year to you too!
Can you elaborate on what you mean by controlled models not "feeling first-class"? Which parts of the HMMs API seem most clunky to you?
As for your suggested options:

1. Doesn't scale, and probably not worth pursuing in a general package (that's what StateSpaceModels.jl is for?).
2. Semantically and causally dubious, if you think in terms of graphical models.
3. I don't understand what you mean here. You can already do either logdensityof(model, obs_seq) or logdensityof(model, obs_seq, control_seq), so I'm not sure what you want to add?

[rsenne]

Sorry, for the lack of clarity on my part! My basic thought is that for many controlled-HMMs I can imagine, creating a ControlledHMM custom object feels like it could be unnecessary. IIUC from your code, the only thing that prevent us from fitting say a GLM-HMM or analogous model "out-of-the-box" in the basic HMM type is because of the defined fit! method and fit_in_sequence.

Concretely, baum_welch already accepts a control_seq and passes it into fit!(hmm, fb_storage, obs_seq, control_seq; ...), but the default StatsAPI.fit!(hmm::HMM, ...) only has the signature without control_seq and calls fit_in_sequence!(dists, i, obs_seq, γ_i) which likewise cannot receive controls.

To my eye, it looks like we could hook into these rather easily by some simple extensions. Specifically:

• keep the existing StatsAPI.fit!(::HMM, ..., obs_seq; ...) and fit_in_sequence!(..., x, w) methods intact, and introduce overloads that accept control_seq and forward it through:

• add fit_in_sequence!(..., x, w, control_seq) with a default implementation that just calls the existing 4-arg method (so controls are ignored unless a model opts in), and

add StatsAPI.fit!(::HMM, ..., obs_seq, control_seq; ...) that is identical to the current method except it calls the 5-arg hook in the emission M-step.

This would be backward compatible and would let controlled emission models (e.g., GLM emissions) plug into the standard Baum–Welch training pipeline via dispatch, without requiring a separate ControlledHMM type. To my mind, this simplifies many controlled models as its the same workflow for the majority of other models: all one nees to do is implement rand, fit!, and logdensityof.

Does this make more sense?

[gdalle]

Sorry, for the lack of clarity on my part! My basic thought is that for many controlled-HMMs I can imagine, creating a ControlledHMM custom object feels like it could be unnecessary. IIUC from your code, the only thing that prevent us from fitting say a GLM-HMM or analogous model "out-of-the-box" in the basic HMM type is because of the defined fit! method and fit_in_sequence.

Well, not quite, because the struct you use as your HMM model also needs to store the parameters of the relationship from control variables to transitions and emissions. This seems hard to do in general because such relationships can be arbitrary (unlike the transitions which are always given by a matrix, and the emissions which are always a vector of distribution-like objects). Maybe you have something like this in mind

@kwdef struct HMM
    init
    trans
    dists
    control_stuff = nothing
end

but then people would basically have to define their own type for whatever they put inside the control_stuff field?

[rsenne]

Okay, I see the light now. It seemed at first glance that this would me simpler for models where control only influenced the emission model and not the transition structure but after looking through the abstractHMM file its clearer to me why that's a problem.

My motivation for this discussion was to just make these types of models dead simple for the modeling-curious experimentalists I work/collaborate with. But I think the best solution is creating a subtype of the AbstractHMM and host it either here, in StateSpaceDynamics.jl, or in HiddenMarkovModels.jl that is deigned to work with models where control only interacts with the emission model and does not affect the transition matrix?

[gdalle]

Even then, what would such a subtype look like? What are the attributes?

[rsenne]

using DensityInterface
using Random
using StatsAPI
using HiddenMarkovModels
using LinearAlgebra
import HiddenMarkovModels:
    initialization, log_initialization,
    transition_matrix, log_transition_matrix,
    obs_distributions
using Base.Threads: @threads

struct ControlledHMM{
    V<:AbstractVector,
    M<:AbstractMatrix,
    VD<:AbstractVector,
    Vl<:AbstractVector,
    Ml<:AbstractMatrix,
} <: HiddenMarkovModels.AbstractHMM
    init::V
    trans::M
    dists::VD
    loginit::Vl
    logtrans::Ml

    function ControlledHMM(init::AbstractVector, trans::AbstractMatrix, dists::AbstractVector)
        loginit  = HiddenMarkovModels.elementwise_log(init)
        logtrans = HiddenMarkovModels.elementwise_log(trans)
        hmm = new{typeof(init), typeof(trans), typeof(dists), typeof(loginit), typeof(logtrans)}(
            init, trans, dists, loginit, logtrans
        )
        HiddenMarkovModels.@argcheck HiddenMarkovModels.valid_hmm(hmm)
        return hmm
    end
end

HiddenMarkovModels.initialization(hmm::ControlledHMM) = hmm.init
HiddenMarkovModels.log_initialization(hmm::ControlledHMM) = hmm.loginit

initialization(hmm::ControlledHMM, control) = hmm.init
initialization(hmm::ControlledHMM, ::Nothing) = hmm.init
log_initialization(hmm::ControlledHMM, control) = hmm.loginit
log_initialization(hmm::ControlledHMM, ::Nothing) = hmm.loginit

HiddenMarkovModels.transition_matrix(hmm::ControlledHMM) = hmm.trans
HiddenMarkovModels.log_transition_matrix(hmm::ControlledHMM) = hmm.logtrans

transition_matrix(hmm::ControlledHMM, control) = hmm.trans
transition_matrix(hmm::ControlledHMM, ::Nothing) = hmm.trans
log_transition_matrix(hmm::ControlledHMM, control) = hmm.logtrans
log_transition_matrix(hmm::ControlledHMM, ::Nothing) = hmm.logtrans

HiddenMarkovModels.obs_distributions(hmm::ControlledHMM) = hmm.dists
obs_distributions(hmm::ControlledHMM, control) = hmm.dists
obs_distributions(hmm::ControlledHMM, ::Nothing) = hmm.dists

Base.length(hmm::ControlledHMM) = length(hmm.dists)


# (a) eltype must reflect control-aware emission logdensities
function Base.eltype(hmm::ControlledHMM, obs, control)
    init_type = eltype(HiddenMarkovModels.initialization(hmm))
    trans_type = eltype(HiddenMarkovModels.transition_matrix(hmm, control))
    dist = HiddenMarkovModels.obs_distributions(hmm, control)[1]
    logdensity_type = typeof(DensityInterface.logdensityof(dist, obs, control))
    return promote_type(init_type, trans_type, logdensity_type)
end

# (b) control-aware observation logdensities
function HiddenMarkovModels.obs_logdensities!(
    logb::AbstractVector{T},
    hmm::ControlledHMM,
    obs,
    control;
    error_if_not_finite::Bool=true,
) where {T}
    dists = HiddenMarkovModels.obs_distributions(hmm, control)
    @inbounds for i in eachindex(logb, dists)
        logb[i] = DensityInterface.logdensityof(dists[i], obs, control)
    end
    error_if_not_finite && HiddenMarkovModels.@argcheck maximum(logb) < typemax(T)
    return nothing
end

function HiddenMarkovModels.obs_logdensities!(logb::AbstractVector, hmm::ControlledHMM, obs, ::Nothing; kwargs...)
    throw(ArgumentError("ControlledHMM requires control at every time step for emission logdensity."))
end

# (c) control-aware sampling
function Random.rand(rng::AbstractRNG, hmm::ControlledHMM, control_seq::AbstractVector)
    T = length(control_seq)
    dummy_log_probas = fill(-Inf, length(hmm))

    init = HiddenMarkovModels.initialization(hmm)
    state_seq = Vector{Int}(undef, T)
    state1 = rand(rng, HiddenMarkovModels.LightCategorical(init, dummy_log_probas))
    state_seq[1] = state1

    @inbounds for t in 1:(T - 1)
        trans = HiddenMarkovModels.transition_matrix(hmm, control_seq[t + 1]) # ignores control by design
        state_seq[t + 1] = rand(rng, HiddenMarkovModels.LightCategorical(trans[state_seq[t], :], dummy_log_probas))
    end

    dists1 = HiddenMarkovModels.obs_distributions(hmm, control_seq[1])
    obs1 = rand(rng, dists1[state1], control_seq[1])
    obs_seq = Vector{typeof(obs1)}(undef, T)
    obs_seq[1] = obs1

    @inbounds for t in 2:T
        dists = HiddenMarkovModels.obs_distributions(hmm, control_seq[t])
        obs_seq[t] = rand(rng, dists[state_seq[t]], control_seq[t])
    end
    return (; state_seq, obs_seq)
end

function Random.rand(hmm::ControlledHMM, T::Integer)
    throw(ArgumentError("ControlledHMM requires a control sequence; call rand(hmm, control_seq)."))
end

# helper: controlled emission fitting
function HiddenMarkovModels.fit_in_sequence!(
    dists::AbstractVector, i::Integer, x, control_seq, w
)
    StatsAPI.fit!(dists[i], x, control_seq, w)
    return nothing
end

# controlled fit! for the HMM itself
function StatsAPI.fit!(
    hmm::ControlledHMM,
    fb_storage::HiddenMarkovModels.ForwardBackwardStorage,
    obs_seq::AbstractVector,
    control_seq::AbstractVector;
    seq_ends::HiddenMarkovModels.AbstractVectorOrNTuple{Int},
)
    (; γ, ξ) = fb_storage

    # Same state fitting logic as HMM.fit!
    if seq_ends isa NTuple
        for k in eachindex(seq_ends)
            t1, t2 = HiddenMarkovModels.seq_limits(seq_ends, k)
            scratch = ξ[t2]
            fill!(scratch, zero(eltype(scratch)))
            for t in t1:(t2 - 1)
                scratch .+= ξ[t]
            end
        end
    else
        @threads for k in eachindex(seq_ends)
            t1, t2 = HiddenMarkovModels.seq_limits(seq_ends, k)
            scratch = ξ[t2]
            fill!(scratch, zero(eltype(scratch)))
            for t in t1:(t2 - 1)
                scratch .+= ξ[t]
            end
        end
    end

    fill!(hmm.init, zero(eltype(hmm.init)))
    fill!(hmm.trans, zero(eltype(hmm.trans)))
    for k in eachindex(seq_ends)
        t1, t2 = HiddenMarkovModels.seq_limits(seq_ends, k)
        hmm.init .+= view(γ, :, t1)
        hmm.trans .+= ξ[t2]
    end
    HiddenMarkovModels.sum_to_one!(hmm.init)
    foreach(HiddenMarkovModels.sum_to_one!, eachrow(hmm.trans))

    # Fit emissions with control
    for i in 1:length(hmm)
        HiddenMarkovModels.fit_in_sequence!(hmm.dists, i, obs_seq, control_seq, view(γ, i, :))
    end

    # Update logs
    hmm.loginit .= log.(hmm.init)
    HiddenMarkovModels.mynonzeros(hmm.logtrans) .= log.(HiddenMarkovModels.mynonzeros(hmm.trans))

    HiddenMarkovModels.@argcheck HiddenMarkovModels.valid_hmm(hmm)
    return nothing
end

function StatsAPI.fit!(
    hmm::ControlledHMM,
    fb_storage::HiddenMarkovModels.ForwardBackwardStorage,
    obs_seq::AbstractVector;
    seq_ends::HiddenMarkovModels.AbstractVectorOrNTuple{Int},
)
    throw(ArgumentError("ControlledHMM requires control_seq; call fit!(hmm, fb_storage, obs_seq, control_seq; seq_ends=...)."))
end

function StatsAPI.fit!(
    hmm::ControlledHMM,
    fb_storage::HiddenMarkovModels.ForwardBackwardStorage,
    obs_seq::AbstractVector;
    control_seq=nothing,
    seq_ends::HiddenMarkovModels.AbstractVectorOrNTuple{Int},
)
    control_seq === nothing && throw(ArgumentError("ControlledHMM requires control_seq."))
    return StatsAPI.fit!(hmm, fb_storage, obs_seq, control_seq; seq_ends=seq_ends)
end


mutable struct GLMNormalModel{T}
    β0::T
    β1::T
    logσ::T
end

DensityInterface.DensityKind(::GLMNormalModel) = DensityInterface.HasDensity()

# Control-aware logdensity: accept any Real types and promote
function DensityInterface.logdensityof(mod::GLMNormalModel, obs::Real, control::Real)
    T = promote_type(typeof(mod.β0), typeof(obs), typeof(control))
    μ = T(mod.β0) + T(mod.β1) * T(control)
    s = exp(T(mod.logσ))
    z = (T(obs) - μ) / s
    return -T(0.5) * log(T(2) * T(pi)) - log(s) - T(0.5) * z * z
end

# Control-aware sampling: accept any Real control
function Random.rand(rng::AbstractRNG, mod::GLMNormalModel, control::Real)
    T = promote_type(typeof(mod.β0), typeof(control))
    μ = T(mod.β0) + T(mod.β1) * T(control)
    s = exp(T(mod.logσ))
    return μ + s * randn(rng, T)
end

# Adapter to match ControlledHMM's expected order
function StatsAPI.fit!(
    mod::GLMNormalModel,
    data::AbstractVector{<:Real},
    control_seq::AbstractVector{<:Real},
    weights::AbstractVector{<:Real},
)
    n = length(data)
    @assert length(control_seq) == n
    @assert length(weights) == n

    # Build X with consistent element type
    T = promote_type(typeof(mod.β0), eltype(data), eltype(control_seq), eltype(weights))
    X = Matrix{T}(undef, n, 2)
    @inbounds for i in 1:n
        X[i, 1] = one(T)
        X[i, 2] = T(control_seq[i])
    end
    y = T.(data)
    w = T.(weights)

    # Weighted normal equations
    W = Diagonal(w)
    β_est = (X' * W * X) \ (X' * W * y)

    residuals = y .- X * β_est
    σ_est = sqrt(sum(w .* (residuals .^ 2)) / sum(w))

    mod.β0 = β_est[1]
    mod.β1 = β_est[2]
    mod.logσ = log(σ_est)
    return nothing
end

#=
Example usage
=#
true_init = [0.5, 0.5]
true_trans = [0.8 0.2; 0.3 0.7]
true_dists = [GLMNormalModel(0.0, 1.0, log(1.0)), GLMNormalModel(5.0, -1.0, log(2.0))]
hmm = ControlledHMM(true_init, true_trans, true_dists)

control = collect(0.0:9.0)
sample = rand(hmm, control)
x = sample.obs_seq

hmm_est, _ = baum_welch(hmm, x, control)

This seems to work?

[gdalle]

Thanks for cooking up a prototype! This definitely makes sense.
So we're basically extending the DensityInterface convention by allowing a third argument. Not sure how I feel about that, maybe the dispatch-based alternative

DensityInterface.logdensityof(dists[i], ControlledObservation(obs, control))

would be conceptually cleaner. Either way, I feel like this could be part of HiddenMarkovModels or StateSpaceModels, whichever is most appropriate. What do you think?

[rsenne]

I think this is generic enough that it probably belongs in HiddenMarkovModels. I can start a PR there and we can iterate on it? I'm open to any changes; this was cooked up pretty quick on my end.

In regards to the dispatch idea:

Make emission models operate on observation–control pairs by introducing a lightweight wrapper, e.g

struct GLMObs{T}
  X::Vector{T}
   y::T
end

This feels similar to my thinking above, but tried to avoid this approach given your thoughts. Do you see this as being more an appropriate approach now as we are wrapping in a new subtype of AbstractHMM? Just want to make sure i follow.

Related question: would you want HiddenMarkovModels to export ControlledObservation and then expect users to implement their emission models via something like logdensityof(dist, ControlledObservation(y,u)) (and similarly for sampling / fitting)? That seems like a reasonable and doable contract if we document it clearly.

[gdalle]

Upon further consideration, I think the three-argument version is cleaner, even though it is not strictly speaking part of DensityInterface's original API. Feel free to open a PR adding it here!