I am trying to determine if Pyro is the right tool to perform model inversion on a behavioural model I developed in PyTorch.

Specifically, I have modelled an agent that performs Bayesian filtering (e.g. with a particle filter) on a simple generative model of the environment. It uses the inferred latent variables at each time step to take actions in the environment. Actions depend both on the inferred latent variables and a set of fixed action parameters.

I would like to invert this model to predict the internal state of the agent. That is, I would like to find the action parameters that best fit observations of the true hidden state of the environment.

Can someone tell me if my intuition that Pyro might be suitable for this kind of problem is correct and outline how to start?

Thanks!

Simon