Model Inversion of Behavioural Model

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

I am interested in this question too… This is relevant to active inference I think. The dynamics of Pyro is useful to interact with the black box environment. You can also use Pyro to compute the objective (the loss) of your problem. I think if you already have code in PyTorch, you can switch to Pyro. But I don’t have more to say about how much work it will require. :frowning: