TuringLang / SSMProblems.jl

Common abstractions for state-space models
http://turinglang.org/SSMProblems.jl/
MIT License
2 stars 2 forks source link

Create prototype Rao-Blackwellised Particle Filter #33

Closed THargreaves closed 1 month ago

THargreaves commented 7 months ago

This PR introduces a prototype for Rao-Blackwellised particle filtering following the SSMProblems.jl interface.

The implementation is verified using a 2D linear Gaussian SSM with ground truth filtering from Kalman.jl and compared to the RBPF where we use a particle filter for the first hidden state.

As expected, the inner (conditionally Gaussian) state is perfectly filtered whilst the outer (particle filter) filtered state is close to the ground truth filtered mean.

Comparison

The outer/inner (a.k.a conditioning/conditional, non-analytic/analytic) dynamics are both represented as their own SSM.

The outer model has no emission_logprobability since the observations only depend on the hidden state in a Rao-Blackwellised model.

The inner model is conditioned on the outer model through the introduction of control variables.

function transition!!(
    rng::AbstractRNG,
    model::FullyLinearGaussianSubsetSSM,
    state::Gaussian,
    step::Int
    control::Vector{Float64},
)

This is a common inclusion in SSMs so we would have needed to add this to the interface anyway.

I don't think this interface is ideal—it's a bit strange that the outer dynamics are defined through a SSM despite having no observations—but it's not too bad at all.

THargreaves commented 7 months ago

A cleaner interface could perhaps be inspired by Kalman.jl. In that package the latent dynamics and observation model are defined separately.

E = LinearEvolution(Φ, Gaussian(b, Q))
Obs = LinearObservationModel(H, R)
M = LinearStateSpaceModel(E, Obs)

A similar approach could be taken here where we define

struct StateSpaceModel <: AbstractStateSpaceModel
    latent_dynamics
    observation_dynamics
end

struct NestedStateSpaceModel <: AbstractStateSpaceModel
    conditioning_latent_dynamics
    conditional_latent_dynamics
    observation_dynamics
end
THargreaves commented 6 months ago

@yebai I would appreciate your thoughts on how best to make this model compatible with the SSMProblems.jl interface. I have an implementation of Rao-Blackwellised particle filtering with backward simulation ready to add to this package as soon the interface is ready.

We discussed the above comment in our meeting last week and concluded that defining a SSM as a composition of a latent dynamics object and a emission model was a bit too complicated for what is meant to be a simple package.

The trouble is, if we don't take this approach, it's not clear to me how one would cleanly write a Rao-Blackwellisable model.

The approach I took in this PR (two full SSMs) feels very sloppy.

The alternative would be to define a new abstract type, say, AbstractNestedStateSpaceModel which has methods transition_outer!! and transition_inner!!. I'm not completely opposed to this but it doesn't feel the most elegant.

I could also see benefits to this structure if you are wanting to compare different models for the latent dynamics whilst maintaining the same (say, isotropic Gaussian) emission model.

What do you think?

yebai commented 6 months ago

I don't understand your point, Tim. Under the current SSMProblems interface, if we want a composable version of SSM, consisting a dynamics model and emission model. We can simply do what you suggested above, then forward the transition!! and emission_logdensity calls to their corresponding member fields, i.e.


struct StateSpaceModel <: AbstractStateSpaceModel
    latent_dynamics
    observation_dynamics
end

transition!!(m::StateSpaceModel, ...) = transition!!(m.latent_dynamics, ...)
emission_logdensity(m::StateSpaceModel, ...) = logdensity(m.observation_dynamics, ...)

# Example model 
E = LinearEvolution(Φ, Gaussian(b, Q))
Obs = LinearObservationModel(H, R)

model = StateSpaceModel(E, Obs)

Am I missing any obvious points?