SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In our implementation of MDPF and A-MDPF, we define the dynamics and measurement models as a set of small learnable neural networks, Fig. 2. The dynamics model uses two encoder networks which encode the particles and the actions into latent dimensions. The particle encoder is applied individually to each particle, allowing for parallelization. The encoded actions and particles are used as input, along with zero-mean unit variance Gaussian noise, into a residual network to produce new particles.
The measurement model is also comprised of several small neural networks. Particles and observations are encoded into latent representations before being used as input into a simple feed-forward network. Our method makes no restrictions on the network architectures allowing modern architectures to be used. For high dimensional observations, such as images, or for more complex problems a sophisticated architecture such as Convolutional Neural Networks [43] or Spatial Transformer Networks [44] can be implemented.
Since particle states are interpretable, they are often in a representation not suitable for neural networks. For example angles contain a discontinuities (at 0 and  $2\pi$ ) which hinders their direct use with neural networks. To address this, non-learned transforms are applied to the particles, converting them into representations that are better suited for neural networks. In the case of angles transforms convert angles to and from unit vectors:
$$
T (\theta) = \left(\cos (\theta), \sin (\theta)\right) \quad T ^ {- 1} (u, v) = \operatorname {a t a n 2} (u, v) \tag {19}
$$
B Synthetic Learning Problem Details
In section 4, a simple synthetic PF learning problem was presented to illustrate the smooth and stable gradient estimates produced by our Importance Weighted Samples Gradient (IWSG) estimator while also highlighting the instability of reparameterization methods when applied to mixture models. Here we give additional details about this learning problem.
In this learning problem we define several regularized PFs with linear Gaussian dynamics and bimodal observation likelihoods:
$$
p \left(x _ {t} \mid x _ {t - 1}, a _ {t}\right) = \operatorname {N o r m} \left(A x _ {t - 1} + B a _ {t}, \sigma^ {2}\right) \tag {20}
$$
$$
p \left(y _ {t} \mid x _ {t}\right) = w _ {1} \operatorname {N o r m} \left(C _ {1} x _ {t} + c _ {1}, \gamma^ {2}\right) + w _ {2} \operatorname {N o r m} \left(C _ {2} x _ {t} + c _ {2}, \gamma^ {2}\right). \tag {21}
$$
For each PF method, we set the dynamics and measurement models to the true linear Gaussian dynamics and bimodal observation with perturbed parameters with the goal of learning the true values of the parameters. In this problem we hold  $\sigma$  and  $\gamma$  fixed and only perturb and learn the other parameters:  $A$ ,  $B$ ,  $C_1$ ,  $C_2$ ,  $c_1$ ,  $c_2$ ,  $w_1$  and  $w_2$ .
We train three variants of the regularized PF with differing gradient estimation procedures for resampling. IWSG-PF uses our IWSG estimator, IRG-PF uses the implicit reparameterization gradients estimator, and TG-PF does not use any gradient estimator for mixture resampling; instead it simply truncates gradients to zero at resampling steps. We compare the stability and convergence of these three methods when learning the parameter values. For this problem, we decouple the resampling and posterior bandwidth parameters as in our A-MDPF model. For resampling we set  $\tilde{\beta} = 0.05$ , and for the posterior we set  $\beta = 0.5$ .
We also train a mixture Kalman Filter (KF) [11, 40] with the same true dynamics and measurement models, also initializing with perturbed parameters. Since this model has linear dynamics with Gaussian noise, and each component of the observation distribution mixture is also linear and Gaussian, we can apply mixture KFs to compute the exact posterior as a mixture of Gaussians. Dynamics are applied to each Gaussian component of the mixture using the KF update equations [11], keeping the number of components the same. When applying the measurement model, we simply multiply the KF mixture with the observation likelihood mixture, increasing the number of mixture components in the KF posterior. The number of components grows exponentially with time, so the mixture KF is not a practical general inference algorithm, but the computation remains tractable for short sequences. Special care must be taken when computing the weights of each mixture component after applying the observation likelihoods; we adapt the approach of Frei and Kunsch [51].
To make the learning problem identifiable and to enforce the constraints  $w_{1} + w_{2} = 1$ ,  $w_{1} \geq 0$ ,  $w_{2} \geq 0$ , we parameterize  $w_{1}$  and  $w_{2}$  within the PF and KF models via the logistic of a single learnable parameter  $v$ :
$$
w _ {1} = \frac {1}{1 + \exp (v)}, \quad w _ {2} = \frac {1}{1 + \exp (- v)}. \tag {22}
$$
To learn the dynamics and measurement model parameters, we setup a synthetic training problem with a dataset of 1000 sequences, each of length 5 time-steps. Using stochastic gradient descent with batch size 64, we train to minimize the negative log-likelihood loss of the true state given the posterior distribution at only the final time-step. For PF models, we use kernel density estimation [19] to estimate a mixture distribution using the final weighted particle set. For the mixture KF, the posterior distribution is simply a mixture model and no further processing is needed to compute the loss.
Fig. 9 shows the convergence and gradient estimates for all of the learned parameters when using the different gradient estimators. IRG's unstable gradients prevents convergence of all parameters, while IWSG converges smoothly and more rapidly than truncated gradients. Further we see that for parameters  $c_{1}$  and  $c_{2}$ , the TG-PF method does not converge to the true value. TG-PF truncates gradients at the resampling, resulting in biased gradients. When the loss is computed at the final time, only information gained from that final time-step is present when updating the parameters, resulting in biased updates.
Figure 9: A simple temporal prediction problem where IRG has highly unstable gradient estimates. We use  $N = 25$  particles when training IRG-PF, IWSG-PF, and truncated gradients (TG-PF). We show the learned values for all learnable parameters and their gradients during training. IRG produces unstable gradients which prevent parameters from converging at all, but IWSG allows for smooth convergence of all parameters. IWSG is faster than biased TG, and nearly as effective as (expensive, and for general models intractable) mixture KF. TG-PF fails to learn correct values for  $c_{1}$  and  $c_{2}$  due to biases from truncating gradients at resampling.
B.1 Computation Requirements
The computational requirements for our MDPF and A-MDPF methods are quadratic with the number of particle during training,  $O(N^2)$ , due to gradient computation. During inference, the computation
Figure 10: Computational time for various particle filters versus the number of particles on the Deepmind Maze task. A batch size of 24 is used. During training, the quadratic run-time of MDPF and A-MDPF is due to gradient computation. During inference MDPF and A-MDPF have linear run-time requirements. In contrast OT-PF must solve a complex optimal transport problem at both training and inference resulting in overall slow run-times.
requirement is linear in the number of particles,  $O(N)$ , since we nearly have to sample from the resampling mixture distribution rather than computing gradients for the samples. This is highlighted in fig. 10. Other methods except OT-PF are have computational complexity  $O(N)$  during both inference and training. OT-PF requires solving an optimal transport problem with complexity  $O(N^2)$  during both training and inference [8]. In practice we find that OT-PF has significant computational requirements resulting in slow run-times both at inference and testing.
In this section we give additional qualitative experiment results for the various tracking tasks.
C.1 Bearingsions Only Tracking Task
In fig. 11 we show additional qualitative results for various PF methods where it is apparent that A-MDPF and MDPF both track the true state well whereas other methods have wide posteriors due to their poor tracking ability. Figs. 12, 13, 14, 15, 16, 17 and 18 show additional trajectories for the various PF methods.
Figure 11: Example trajectory with various PF methods shown for the bearings only tracking task. The radar station (green circle) produces bearings observations (green line) from the radar station to the car (red arrow). Note that observations are noisy and occasionally totally incorrect. Shown is the posterior distribution of the current state (blue cloud, with darker being higher probability), the mean particle (yellow arrow) and the N = 25 particle set (black dots). A-MDPF and MDPF both have tight posterior distributions over the true state while maintaining tracking. Other methods such as DIS-PF, OT-PF and C-PF have wide posteriors showing their uncertainty due to their poor ability to track the true state.
Figure 12: Example trajectory using A-MDPF shown for the bearings only tracking task. The radar station (green circle) produces noisy (and occasionally incorrect) bearings observations (green line) from the radar station to the car (red arrow). A-MDPF maintains a tight posterior distribution (blue cloud, with darker being higher probability) over the true state (red arrow) over time showing the effectiveness of our method. The mean particle (yellow arrow) and the  $\mathrm{N} = 25$  particle set (black dots) are also shown.
Figure 13: Another example trajectory using A-MDPF shown for the bearings only tracking task.
Figure 14: Example trajectory using MDPF shown for the bearings only tracking task. Similar to A-MDPF in figs. 12 and 13, MDPF accurately tracks the true state (red arrow) with a tight posterior density (blue cloud).
Figure 15: Example trajectory using DIS-PF shown for the bearings only tracking task. Due to the poorly learned dynamics and measurement models, DIS-PF often loses track of the true state (red arrow) as seen in time-steps  $18 - 28$ . This results in a large estimated posterior density (blue cloud) to account for the incorrectly placed particles (black dots) caused by poor dynamics and particle weighting.
Figure 16: Example trajectory using SR-PF shown for the bearings only tracking task.
Figure 17: Example trajectory using OT-PF shown for the bearings only tracking task. Like DIS-PF in fig. 15, OT-PF has trouble tracking the true state (red) due to poorly learned dynamics and measurement models.
Figure 18: Example trajectory using C-PF shown for the bearings only tracking task. C-PF is unable to accurately track the true state (red) as shown by the relatively non-moving mean particle (yellow arrow).
C.2 Deepmind Maze Tracking Task
In figs. 19, 20 and 21 we show additional qualitative results for various PF methods for mazes 1, 2 and 3, respectively. Again it is apparent that A-MDPF and MDPF both track the true state well whereas other methods have wide posteriors due to their poor tracking ability.
Figure 19: Example trajectory with various PF methods shown for the Deepmind maze tracking task for maze 1. Shown is the current true state (red arrow), the posterior estimate of the current state (blue cloud, with darker being higher probability), the mean particle (yellow arrow) and the  $\mathrm{N} = 25$  particle set (black dots).
Figure 20: Example trajectory with various PF methods shown for the Deepmind maze tracking task for maze 2. Shown is the current true state (red arrow), the posterior estimate of the current state (blue cloud, with darker being higher probability), the mean particle (yellow arrow) and the  $\mathrm{N} = 25$  particle set (black dots).
Figure 21: Example trajectory with various PF methods shown for the Deepmind maze tracking task for maze 3. Shown is the current true state (red arrow), the posterior estimate of the current state (blue cloud, with darker being higher probability), the mean particle (yellow arrow) and the  $\mathrm{N} = 25$  particle set (black dots).
C.3 House3D Tracking Task
Figs. 22 and 23 show example trajectories for the House3D tracking task using A-MDPF and MDPF respectively. We see that our A-MDPF and MDPF methods both successfully track the true state over time as evident by the tight posterior over the true state as well as the closeness of the mean state to the true state. Fig. 24 shows an example trajectory using the LSTM model. It is clear that the LSTM model is ignoring observations and instead is blind propagating its state estimate using good dynamics and the actions (dead-reckoning). In later time-steps, the LSTM estimate drifts away from the true state however the predicted state trajectory looks similar to that of the true state, but with a growing offset, indicating blind propagation.
Fig. 25 shows qualitative results when using OT-PF. OT-PF is unable to learn good dynamics and measurement models and quickly diverges away from true state. Figs. 26, 27 and 28 give example trajectories for TG-PF, DIS-PF and SR-PF respectively. These methods fail to learn usable dynamics or measurement models with their trajectories essentially being random.
Figure 22: Example trajectory for the House3D tracking task using A-MDPF. The posterior density of the state is tight showing the confidence A-MDPF has in its estimate. The estimated posterior density (blue cloud, with darker being higher probability) also converges on the true state (red arrow) showing its accuracy. The mean particle (yellow arrow) closely tracks the true state. The observations are shown as inset images in the plots.
Figure 23: Example trajectory for the House3D tracking task using MDPF. Similar to fig. 22, MDPF accurately tracks the true state.
Figure 24: Example trajectory for the House3D tracking task using the LSTM model. The LSTM model does not use observations but rather blind-propagates the estimated state using the available noisy actions. This is seen by the true state (shown in red) and the estimated state (shown in yellow) having similar shapes but with increasing error.
Figure 25: Example trajectory for the House3D tracking task using the OT-PF. OT-PF is unable to learn good dynamics or measurement models and thus the estimated posterior (blue cloud) diverges away from the true state (shown in red) after a couple time-steps.
Figure 26: Example trajectory for the House3D tracking task using the TG-PF. TG-PF is unable to learn a usable dynamic or measurement model resulting in fairly random output. This is reflected in the large uncertainty of the estimated posterior (blue cloud) as well as the random nature of the mean state (shown in yellow).
Figure 27: Example trajectory for the House3D tracking task using the DIS-PF. Similar to TG-PF in fig. 26, DIS-PF is unable to learn a usable dynamic or measurement model and therefore is unable to collapse the estimated posterior density (blue cloud) onto the true state (shown in red).
True State Mean Particle Posterior Density
Figure 28: Example trajectory for the House3D tracking task using the SR-PF. Similar to TG-PF in fig. 26 and DIS-PF in fig. 27, SR-PF is unable to learn a usable dynamic or measurement model.
In this section give additional details for the various experiments (tasks). We also provide specific neural architectures used in each task.
D.1 Mean Squared Error Loss
When training the comparison PF methods, we use the Mean-Squared Error (MSE) loss. This is computed as the mean squared error between the sparse true states  $s_t$  and the mean particle at each time  $t$ :
$$
\mathcal {L} _ {\mathrm {M S E}} = \frac {1}{| \mathcal {T} |} \sum_ {t \in \mathcal {T}} \left(x _ {t} - \sum_ {i = 1} ^ {N} w _ {t} ^ {(i)} x _ {t} ^ {(i)}\right) ^ {2} \tag {23}
$$
Here  $\mathcal{T}$  are the indices of the labeled true states. During training, sparsely labeled true states are used and thus  $\mathcal{T} = \{4,8,12,\ldots\}$ . During evaluation we use dense labels,  $\mathcal{T} = \{1,2,3,4,\ldots\}$
For angular dimensions of  $x_{t}^{(i)}$ , we compute the mean particle by converting the angle into a unit vector, averaging the unit vectors (with normalization) and converting the mean unit vector back into an angle. Further we take special care when computing square differences for angles due to the discontinuity present at 0 and  $2\pi$ .
D.2 Pre-Training Measurement Model
We pre-train the measurement model for each task using the sparse true states. To pre-train, a random true state from the dataset is selected. A particle set is constructed by drawing samples from a distribution centered around the select true state and weights for this particle set are computed using the measurement model and the observation associated with the selected true state. A mixture distribution is then fit using KDE with a manually specified bandwidth. Negative Log-Likelihood loss of the true state is used as the training loss. We find pre-training not to be sensitive to choice of bandwidth so long as it is not too small.
The dynamics model cannot be pre-trained because of the lack of true state pairs due to the sparse true state training labels and thus must be trained for the first time within the full PF algorithm.
D.3 Additional General Task and Training Details
Table 2: Hyper-parameters use for various PF methods.
Parameter	Value
SR-PF [7] λ	0.1
OT-PF [8] λ (Bearings Only)	0.5
OT-PF [8] λ (Deepmind)	0.01
OT-PF [8] λ (House3D)	0.01
OT-PF [8] scaling	0.9
OT-PF [8] threshold	1e-3
OT-PF [8] max iterations	500
C-PF [31, 32] λ	0.5
For all tasks we wish to track the 3D (translation and angle) state of a robot,  $s = (x,y,\theta)$ , over time. This state space contains an angle dimension which is not suitable for neural networks. To address this issue, we convert particles into 4D using a transform  $T(s) = (x,y,\sin(\theta),\cos(\theta))$  before using them as input into a neural network. Further we convert back into the original 3D representation to retrieve the a 3D particle after applying neural networks to the particles.
We use the Adam [52] optimizer with default parameters, tuning the learning rate for each task. We use a batch size of 64 during training and decrease the learning rate by a factor of 10 when the validation loss reaches a plateau. We use Truncated-Back-Propagation-Through-Time (T-BPTT) [47], truncating the gradients every 4 time-steps.
For all tasks we use the Gaussian distribution for the positional components of the state and the Von Mises distribution for the angular components when applying KDE. Here  $\beta$  is a dimension-specific bandwidth corresponding to the two standard deviations for the Gaussians and a concentration for the Von Mises.
Some methods have non-learnable hyper-parameters that must be set. The values of the hyperparameters are chosen via brief hyper-parameter search with the final chosen values shown in table 2. The hyper-parameter values are the same for all tasks unless otherwise specified. We find the regularization parameter  $\lambda$  for OT-PF to be sensitive to the specific datasets.
We choose to make the dynamics position invariant. This prevents the learned PF models from memorizing the state space based on position but rather forces the learned dynamics to the true dynamics. Position invariance can be achieved by masking out the translational dimensions of the particle before input into the dynamics model.
D.4 Bearings Only Tracking Task
(a) Particle encoder for dynamics Model.
(c) Observation encoder.
(d) Dynamics model residual network, with actions omitted. Output scaling all dimensions to be within [-0.99, 0.99].
Figure 29: Neural architecture for the networks within the dynamics and measurement models for the bearings only tracking task.
(b) Particle encoder for measurement model.
(e) Measurement model weights network.
In this section we give additional task specific information and network architecture details for the bearings only tracking task.
D.4.1 Additional Experiment Details
When training MDPF using the Negative Log-Likelihood (NLL) loss, we set the learning rate (LR) of the neural networks to be 0.0005 and the bandwidth learning rate to 0.00005. When training
Figure 30: Example trajectories and observations from each of the three mazes from the Deepmind maze tracking task.
using MSE loss, we use 0.0001 and 0.00001 for the neural network and bandwidth learning rates respectively. For the discrete PF methods, we use  $\mathrm{LR} = 0.0005$  for training the neural network using the MSE loss. We use the same bandwidth for optimizing the bandwidth using the NLL loss when the neural networks are held fixed. For the LSTM [45], we use 0.001 as the learning rate for training the networks and  $\mathrm{LR} = 0.0005$  when optimizing the bandwidth parameter using NLL.
We apply gradient norm clipping, clipping the norm of the gradients to be  $\leq 100.0$ .
In addition to using the Gaussian distribution for KDE, we also train with the Epanechnikov distribution applied to the translational dimensions of the state. To do this we first train using Gaussian kernels before re-training with the Epanechnikov, using the already trained dynamics and measurement models as a starting point.
The observations  $o_{t}$  are generated as noisy bearings from the radar station to the car:
$$
o _ {t} \sim \alpha \cdot \operatorname {U n i f o r m} (- \pi , \pi) + (1 - \alpha) \cdot \operatorname {V o n M i s e s} (\psi_ {t}, \kappa)
$$
where  $\psi_t$  is the true bearing,  $\alpha = 0.15$  and  $\kappa = 50$ . Observations  $o_t$  are angles and not suitable as input into a neural network. As such we convert the angular representation into a unit vector as is done with the angle dimensions of the particles.
D.4.2 Network Architectures
Fig. 29 gives details on the neural network architectures used for the dynamics and measurement models within the PF for the bearings only tracking task. Since this task does not provide actions, the action encoder has been omitted.
D.5 Deepmind Maze Tracking Task
In this section we give additional task specific information and network architecture details for the Deepmind maze tracking task.
D.5.1 Additional Experiment Details
The Deepmind maze tracking task was first proposed in [5]. In this task, a robot moves through modified versions of the maze environments from Deepmind Lab [48]. Three distinct mazes are provided and we train and test on each maze individually. Example robot trajectories as well as example observations for each maze are shown in fig. 30. We modify this task by increasing the noise applied to the actions by 5 fold as well as using sparsely labeled true states during training.
When training MDPF using the Negative Log-Likelihood (NLL) loss, we set the learning rate (LR) of the neural networks to be 0.0005 and the bandwidth learning rate to 0.00005. When training using MSE loss, we use 0.0005 and 0.00005 for the neural network and bandwidth learning rates respectively. For the discrete PF methods, we use  $\mathrm{LR} = 0.0005$  for training the neural network using the MSE loss. We use the same bandwidth for optimizing the bandwidth using the NLL loss when the neural networks are held fixed. For the LSTM, we use 0.001 as the learning rate for training the networks and  $\mathrm{LR} = 0.0005$  when optimizing the bandwidth parameter using NLL.
(a) Observation encoder.
(b) Particle encoder for dynamics Model.
(c) Particle encoder for measurement model.
(d) Action encoder for dynamics Model.
(e) Dynamics model residual network. Output scaling scales the transnational dimensions to be within [-7, 7] and the unit vector representing angles to be within [-1, 1].
(f) Measurement model weights network.
Figure 31: Neural architecture for the networks within the dynamics and measurement models for the Deepmind maze tracking task.
We apply gradient norm clipping, clipping the norm of the gradients to be  $\leq 500.0$ . Further we limit the bandwidth for the angle dimension such that  $\beta \leq 4.0$ . We do this only for this task as non-MDPF methods were prone to learning an unreasonably large bandwidth resulting in a near uniform distribution for the angular part of the state.
D.5.2 Network Architectures