implement rao-blackwellised particle filtering using dynamax and blackjax
murphyk opened this issue · 9 comments
combine Kalman Filter from dynamax and sequential monte carlo from Blackjax to implement RBPF for switching linear Gaussian SSM. For details, see sec 13.4.1 of https://probml.github.io/pml-book/book2.html
Is this still of interest? If so I can have a go at it. Also I was thinking that maybe this could be part of a larger submodule of dynamax
with other SLDS inference and learning functions?
yes, this is still of interest. You may want to lookinto the SMC support in blackjax, and see if we can pass a dynamax model to it for vanilla PF, then use the marginalization trick from dynamax.kalman_filter to do RBPF.
I have looked into the SMC support of blackjax
and as far as I can tell blackjax.smc.base.py
is an implementation of waste-free SMC, for approximating a sequence of distributions. Correct me if I am missing something, but in the RBPF we are sampling a discrete Markov chain (so no need for the MCMC steps of the waste-free sampler) ; wouldn't it be sufficient to use our own custom proposals?
Also, a link to dynamax.kalman_filter
would be much appreciated. I cannot find the file in the current version of the repo.
It's true that for RBPF we can just use a simple discrete proposal, so blackjax may be irrelevant.
Here is the KF code:
https://github.com/probml/dynamax/blob/main/dynamax/linear_gaussian_ssm/inference.py#L370
Sounds good! I'll get started on it then.
I have a question regarding implementation of the optimal resampling method:
Is there an efficient way to determine the threshold
The best that I have come up with is the naive way, i.e., to try each value of
Once this has been addressed I can submit a PR for this feature. Thanks!
Thanks for your effort @kostastsa! I am still wrapping my head around application of RBPF to SLDSs, but so far I've been able to filter the state in my model quite successfully. I am now wondering if I understand it correctly that the current implementation of SLDS
does not yet support parameter estimation (e.g. dynamic and emission covariances). And if so, what would be the most reasonable way to extend the model? The PML2 book hints at SMC for parameter inference, but I, frankly, cannot quite see the exact way to combine the tempering techniques from section 13.6 with this RBPF approach.
Hi @Sibgatulin, thank you for your interest and for your question. You understand correctly, the current SLDS
class does not support parameter estimation at the moment, although it's definitely something we're interested in adding. I think that probably the most popular way of doing parameter estimation in this model is using the EM algorithm. See for example @murphyk's books (or this technical report which is the most detailed description that I have found). The SMC approach would be useful for computing the posterior of the parameters. It is also a matter of preference, but I would begin with an implementation of EM, rather than SMC. Hope this addresses some of your questions! Best, Kostas