Particle filter collapse

[cooijmanstim]

I was able to move past my issue by using UKF, but I think there is still something to be explained here. Basically I'm using particle filter in a loop, feeding the weighted particles back into the next call as an initial condition. The particles seem to collapse in a single step, and also fail to move in time after that. I thought it was something to do with my mujoco setup, but it also happens in a straightforward system:

import jax
import jax.numpy as jnp
import jax.random as rax

jax.config.update("jax_traceback_filtering", "off")

import dynestyx as dsx
import numpy as np
import numpyro.distributions as dist
from dynestyx import DynamicalModel, Filter
from dynestyx.inference.filters import PFConfig
from numpyro.infer import Predictive

from dynestyx.inference.integrations.utils import particles_to_delta_mixtures

state_size = 13
observation_size = 17
control_size = 5

def fuzz(x):
    return dist.Normal(jnp.astype(x, "float32"), jnp.array(0.1, "float32"))

def step(state, rng):
        rng1, rng2 = rax.split(rng)
        obs = rax.normal(rng1, shape=[observation_size])
        ctrl = rax.normal(rng2, shape=[control_size])

        def state_evolution(x, u, t_now, t_next):
            return fuzz(2 * x + 1 + u.sum())

        def observation_model(x, u, t):
            return fuzz(np.ones([observation_size]) + x.sum() + u.sum())

        def dsx_model(**kwargs):
            dynamics = DynamicalModel(
                control_dim=control_size,
                initial_condition=state,
                state_evolution=state_evolution,
                observation_model=observation_model,
            )
            return dsx.sample("f", dynamics, **kwargs)

        with Filter(PFConfig(n_particles=11)):
            predictive = Predictive(dsx_model, num_samples=1)
            t1 = jnp.ones([1])
            pred = predictive(
                rng,
                obs_times=t1,
                obs_values=obs[None],
                ctrl_times=t1,
                ctrl_values=ctrl[None],
                predict_times=t1,
            )
        # squeeze Predictive num_samples axis (batch_ndims=0 doesn't seem to do it)
        x = jnp.squeeze(pred["f_filtered_particles"], axis=0)
        w = jnp.squeeze(pred["f_filtered_log_weights"], axis=0)
        [px1] = particles_to_delta_mixtures(x, w)
        return px1

state0 = fuzz(np.zeros([state_size]))
states = [state0]
for t in range(7):
     states.append(step(states[-1], rax.key(t)))

particles = np.stack([state.particles for state in states[1:]])  # states[0] is Normal

from matplotlib import pyplot as plt
plt.figure()
plt.plot(jnp.diff(particles, axis=0)[:, 0])
plt.title("particle 0 time derivative")
plt.figure()
plt.plot(particles.var(axis=1))
plt.title("variance across particles")

[cooijmanstim]

@DanWaxman so from our chat and looking a bit closer it does make sense that this doesn't call the state evolution, because obs_times, ctrl_times are 1 while predict_times is also 1, which I didn't mean to do. But I got the impression you were telling me it wouldn't work even if these were different.

[DanWaxman]

Thanks Tim!! Indeed, I think this is not possible with a UKF. This is because UKF's sample propagation is actually a specialized integration rule, which is not easy for us to emulate with the predict_times interface. To be honest, this is a kind of low-level thing to do, and so I don't think we plan to support it (instead, I'd rather us do the low-level things internally, a la #154).

I did write a version that uses the PF + predict_times interface. This is a somewhat high-variance version of the particle filter (it's doing resampling every step), but I think should be correct.

[DanWaxman]

I'll close this as not planned, since I think we don't want to support this kind of low-level breaking down of the inference loop (@mattlevine22 feel free to re-open if you feel to the contrary). Definitely want to handle this weirdness at the root cause though via #154.

[mattlevine22]

• I agree that addressing Enable control_model feature #154 ought to more seamlessly support what @cooijmanstim is trying to do here, but
• I do think he should be able to do this stuff (including with UKF) in this current (messy) way.

However, I'm not clear on what the corrected example would look like...

• Sounds like you want predict_times and obs_times to be different...what would that look like and why?
• Regardless of time-specification, predict_times would be ignored in this example because there is no call to a Simulator (further supports thoughts in Make Simulator the default interpretation of dsx.sample #151...user should not need to write Simulator here, as it is implied!)
• Separately from the above 2 things, "predictions" would not be stored/used in this code (step only returns the filtered states anyway)

It sounds like you guys had a more thorough conversation, so I'm probably behind on all of this.

I would like to re-open this issue just to make sure we clarify it for posterity, if nothing else.

[cooijmanstim]

Yes I made a big mess of the time subscripts when I initially wrote this and never re-examined it, let me clarify. I have obs[t] and ctrl[t-1] and want to infer state[t], and then I'll use that estimated state to compute ctrl[t]. So yes I do actually want obs_times to equal predict_times, but I need ctrl_times to be one step before that. I don't know if ctrl is useful in that case, but those are the correct time points on the quantities I have.

Is there a time point associated with the initial condition?

[DanWaxman]

I don't think there was much more discussion, just some thoughts I hadn't written down clearly, sorry about that.

If we conceptualize the filters as predict(...) and update(...) the point is that there is always a single data point, so there is update(...), but never predict(...). So dynamics are never called, there is just a single Bayesian update from $p(x_0)$ to $p(x_0 | y_0)$.

By appending another predict_times time point, we can get another point via Simulator(...), i.e., $p(y_1 | y_0)$. This is tenable for the PF, as the corresponding pattern is indeed (a higher variance version of) the predict(...) step for a bootstrap PF. However, the predict(...) step for a UKF entails an (unscented) integration step that we can't replicate via Simulator. Without support for NaN observations or a closed-loop controller, I thus don't think there's a clear way to obtain the UKF predict step in this data-point-by-data-point pattern.

I guess my reasoning for saying this is not planned is that we don't really want users having to think about predict and update directly. There's a few different ways to address this (e.g., #154, #138, and handling a custom t0 properly), but all of them are a distinct high-level feature change.