Observation Adaptation via Annealed Importance Resampling for Partially Observable Markov Decision Processes

Partially observable Markov decision processes (POMDPs) are a framework for modeling decision-making under uncertainty, defined by the tuple $(S,A,T,R,\Omega,O)$. Here, $S$ is the set of states, $A$ the set of actions, $T(s^{\prime}|s,a)$ the transition function, $R(s,a)$ the reward function, $\Omega$ the set of observations, and $O(o|s^{\prime},a)$ the observation function. In POMDPs, at time t, the agent maintains a belief state $b_{t}$, a probability distribution over states. This distribution is updated after action $a_{t}$ and observation $o_{t+1}$ using the equation: $b_{t+1}(s^{\prime})=\eta O(o_{t+1}|s^{\prime},a_{t})\sum_{s\in S}T(s^{\prime}|s,a_{t})b_{t}(s)$, where $\eta$ is a normalizing factor. This belief update process enables the agent to estimate its current state despite partial observability, facilitating informed decision-making in uncertain environments.

Particle filters  (Arulampalam et al. 2002), also known as sequential Monte Carlo (SMC) methods, are computational algorithms used to estimate the state of a dynamic system when it is observed through noisy measurements. In the context of Partially observable Markov decision processes (POMDPs), particle filters offer an efficient way to approximate the belief state by maintaining a set of weighted samples (particles) that represent the probability distribution over possible states. The widely-known bootstrap particle filter algorithm  (Gordon, Salmond, and Smith 1993) works by recursively propagating these particles through the state transition model and updating their weights based on new observations. In POMDPs, at each time step $t$, the particle filter aims to approximate the optimal posterior distribution:

where $s_{t}$ is the current state, $o_{t}$ is the current observation, $s_{t-1}$ is the previous state, and $a$ is the action taken.

A crucial step in bootstrap particle filter is resampling (Kitagawa 1996), which addresses the problem of particle degeneracy. Over time, some particles may have negligible weights, leading to poor representation of the state space. Resampling involves randomly drawing particles with replacement from the current set, with probabilities proportional to their weights, and then resetting all weights to $1/N$, where $N$ is the number of particles. However overly frequent resampling can lead to sample impoverishment, where the particle set loses diversity and fails to adequately represent the full state space. To determine when to resample, the Effective Sample Size (ESS) (Forthofer, Lee, and Hernandez 2007) is often used:

where $w_{i}^{n}$ are the normalized weights. When the ESS falls below a predefined threshold (typically $N/2$), resampling is triggered.

Although the bootstrap particle filter is relatively easy to implement, its accuracy can be significantly compromised when there is a substantial discrepancy between the optimal posterior distribution described in equation (1) and the state-transition distribution. This limitation becomes particularly pronounced in scenarios where the one-step received observation is highly informative about the true state, leading to a sharply peaked posterior distribution. In such cases, the bootstrap filter’s reliance on the state-transition model for proposal generation may lead to particles being proposed in regions where the posterior has significant mass but the proposal density is low, resulting in many particles having negligible weights. Consequently, this can exacerbate the problem of sample degeneracy, necessitating more frequent resampling and potentially leading to sample impoverishment.

Importance Sampling (IS) is a technique for estimating properties of a target distribution by sampling from a different, easier-to-sample proposal distribution and reweighting the samples (Kloek and van Dijk 1978). However, its major limitation lies in its inefficiency when the proposal distribution differs significantly from the target distribution — in such cases, only a few samples that happen to fall in the high-probability regions of the target distribution receive high importance weights, while most samples have negligible weights. This leads to high variance in the estimates and poor sample efficiency.

Annealed Importance Sampling (AIS) constructs a sequence of intermediate distributions that gradually bridge the gap between the proposal and target distributions (Neal 1998). Given a target density $p(x)$ and a proposal density $q(x)$, AIS defines a sequence of intermediate distributions:

where $0=\beta_{0}\leq\cdots\leq\beta_{K}=1$ represents a sequence of inverse temperatures. The method proceeds by first drawing an initial sample $x_{0}\sim\pi_{0}(x)$, then evolving this sample through a series of transition kernels:

where each $T_{k}$ is constructed to leave $\pi_{k}$ invariant. The final importance weight is computed as the product of ratios between successive distributions:

This gradual transition through intermediate distributions helps AIS overcome the limitations of standard importance sampling by maintaining better overlap between successive distributions, resulting in more reliable estimates.

Monte Carlo tree search (MCTS) has demonstrated success in solving large POMDPs online. POMCP (Silver and Veness 2010) pioneered this approach by combining a UCT-based tree search with particle filtering for belief updates, but faces challenges with continuous observation spaces. DESPOT and its variant AR-DESPOT (Ye et al. 2017) improve upon POMCP by focusing the search on a fixed set of scenarios and using dual bounds to guide exploration, making it more robust for large discrete problems. POMCPOW (Sunberg and Kochenderfer 2018) extends these ideas to continuous state-action-observation spaces by incorporating progressive widening and weighted particle filtering.

A common thread among these approaches is their reliance on particle filtering techniques, specifically bootstrap particle filtering or Sequential Importance Resampling (SIR), to update belief states. While these filtering methods are computationally efficient, they can sometimes lead to particle degeneracy issues, especially in continuous observation spaces or when observations are unlikely under the current belief.

Most recently, AdaOPS (Wu et al. 2021) attempts to address this challenge using KLD-sampling, which dynamically adapts the number of particles based on the Kullback-Leibler divergence between the true and approximated distributions. However, KLD-sampling exhibits two major limitations: First, it requires discretizing the state space into bins for divergence calculation, which becomes computationally prohibitive especially in high-dimensional spaces (Li, Sun, and Sattar 2013). Second, and more critically, while it adjusts particle quantities, it does not modify the particle values themselves, leading to potential underestimation of distribution variance and failure to capture multimodal aspects of the target distribution. In contrast, Annealed Importance Sampling (AIS) offers several compelling advantages: it constructs a sequence of intermediate distributions that gradually bridge the prior and posterior distributions, effectively maintaining particle diversity while preventing degeneracy. This annealing process enables particles to adaptively migrate toward regions of high posterior probability, making it particularly effective when dealing with concentrated observation likelihoods or significant disparities between prior and posterior distributions. The gradual transition through intermediate distributions allows AIS to better capture the full structure of multimodal distributions and provide more accurate representations of distribution tails compared to KLD-sampling approaches.

Previous work has explored adding tempering iterations to particle filters, where the optimal posterior distribution is constructed adaptively through a sequence of Monte Carlo steps. For example, M. Johansen (2015) developed a block-tempered particle filter that uses bridge distributions to gradually adapt particle values, while (Herbst and Schorfheide 2019) proposed a tempered particle filter that sequentially reduces inflated measurement error variance in Dynamic Stochastic General Equilibrium (DSGE) models. However, these tempering approaches have been primarily studied in the context of state estimation and system identification, but not yet explored within the POMDP planning literature. To our knowledge, we are the first to investigate using bridge distributions to connect target and proposal distributions in particle filtering specifically for POMDP planning, combining the strengths of both sequential Monte Carlo methods and POMDP planning algorithms.

Similar to (Wu et al. 2021; Ye et al. 2017; Smith and Simmons 2004), during online planning, AIROAS maintains both upper and lower bounds, $u(\bar{b},a)$ and $l(\bar{b},a)$, for each action node $a$ at a belief node $\bar{b}$. These bounds estimate the optimal value that can be achieved by taking action $a$ at belief $\bar{b}$. The action selection follows an optimistic strategy - at each belief node $\bar{b}$, AIROAS chooses the action $a^{*}$ that maximizes the upper bound(Line 1 in algorithm 2):

After selecting an action, we need to choose which observation branch to explore next. This choice is guided by the probability-weighted excess uncertainty (EU) criterion (Ye et al. 2017; Smith and Simmons 2004) (Line 2 in algorithm 2). For each possible observation $o$ after taking action $a^{*}$, we compute:

where $\hat{p}(o|\bar{b},a^{*})$ is an estimation of the probability $p(o|\bar{b},a^{*})$ i.e estimates the probability of receiving observation $o$ after taking action $a^{*}$ at belief $\bar{b}$, $\tau(\bar{b},a^{*},o)$ represents the updated belief after executing $a^{*}$ and receiving observation $o$, and $EU(\cdot)$ measures the excess uncertainty at a belief node  (Wu et al. 2021). The excess uncertainty quantifies the degree of value uncertainty at each belief node:

where $u(\bar{b})$ and $l(\bar{b})$ are the upper and lower bounds at belief $\bar{b}$, $\bar{b}_{0}$ is the root belief, $d(\bar{b})$ denotes the depth of belief $\bar{b}$ in the tree, $\gamma$ is the discount factor, and $\xi\in(0,1)$ controls the target uncertainty reduction at the root. The observation $o^{*}$ that maximizes the weighted excess uncertainty is selected for exploration, effectively focusing the search on belief states with high uncertainty and high probability of being reached (Wu et al. 2021).

When encountering a leaf node, we first perform annealed importance resampling to adjust the particle states and weights (detailed in the next section). Then, AIROAS expands this leaf node by creating child nodes for each possible action $a$. For each action node, the expansion process involves propagating the parent belief node’s particles forward using a simulator $\mathcal{G}$ that generates state transitions (Line 6 in algorithm 3):

The expansion continues by creating new belief nodes for each unique observation encountered during particle propagation. Specifically, after propagating each particle through $\mathcal{G}$, we obtain a set of updated particle states, initially assigning each particle a weight of 1. For each observation $o$ generated during this process, we create a new belief node where the particle weights are updated according to the observation likelihood $p(o|s,a)$ — the probability of receiving observation $o$ given the particle state $s$ and action $a$. This results in a tree structure where each action node branches into multiple belief nodes based on possible observations, with each belief node containing weighted particles that represent the updated belief state.

After expanding belief nodes, we perform a backup operation to update the bounds of their ancestor nodes. Similar as (Ye et al. 2017) and (Wu et al. 2021), this backup process is implemented by recursively applying the Bellman equation, which calculates the optimal value of a belief state in two parts. First, we consider the expected immediate reward for taking an action in the current belief state, calculated by looking at each possible state, weighing its reward by how likely we think we are in that state according to our current belief. Second, we consider the long-term value by looking at all possible observations we might receive after taking an action. For each possible observation, we calculate how likely we are to see that observation, determine what our new belief state would be after seeing it, consider the value of that resulting belief state, and weight this value by the probability of getting that observation. The total value is the sum of these immediate and future components, and we choose the action that maximizes this total value. In practice, since we can’t compute exact values over continuous states and observations, we approximate this calculation using discrete sums over our particle-based belief representation.

As discussed above, the belief of each belief node is generated using a state-transition function from its parent belief node. Suppose the leaf node represents belief at timestep t and its parent belief node represents belief at timestep t-1, while our simulator $\mathcal{G}$ implements the state-transition function $p(s_{t}|s_{t-1},a)$ during expansion, our ultimate objective is to have the final particle distribution approximate the optimal posterior distribution $p(s_{t}|o_{t},s_{t-1},a)$ referenced in eq. 1.

The primary challenge in this context stems from the impossibility of directly sampling from the optimal posterior distribution. Previous work (Silver and Veness 2010; Ye et al. 2017; Sunberg and Kochenderfer 2018; Wu et al. 2021) approximates sampling from this optimal posterior distribution using Sequential Importance Resampling, treating the state-transition distribution as proposal distribution and the optimal posterior distribution as target distribution. However, this approach suffers from particle degeneracy issues, particularly when observations are highly informative about the true state, leading to a sharply peaked posterior distribution and a significant divergence between the optimal posterior and the proposal distribution derived from the current belief (Wu et al. 2021).

To address the issue, Annealed Importance Sampling (Neal 1998) introduced in Background offers an effective solution by constructing intermediate bridging distributions between the proposal distribution and target distribution, gradually guiding particles toward high-probability regions of the target distribution by constructing a sequence of intermediate distributions parameterized by tempering parameters $0=\beta_{0}\leq\beta_{1}\leq\cdots\leq\beta_{K}=1$.

Building upon these principles, we integrate Annealed Importance Resampling (AIR), which extends AIS with a resampling step, into our AIROAS algorithm to address the particle degeneracy issue inherent in POMDP planning. Our approach adapts this concept to the POMDP planning context by designing intermediate distributions that account for observations. By gradually bridging the proposal and target distributions, particles are adaptively reweighted and resampled, allowing them to migrate toward regions of high posterior probability. This gradual transition through intermediate distributions not only addresses degeneracy but also enables our approach to capture the multimodal characteristics of the optimal posterior distribution more effectively.

These tempering parameters are generated using a sigmoid function applied to a linearly spaced sequence between $10^{-3}$ and 1. Specifically, for a sequence $\{x_{1},x_{2},\cdots x_{K}\}$ linearly spaced between $10^{-3}$ and 1, each tempering parameter $\beta_{i}$ is computed as:

This sigmoid transformation ensures a smooth progression of $\beta$ values from near zero to one, with denser sampling in the middle range.

The intermediate distributions are characterized by their probability densities:

When $k=K$, this formulation exactly matches the optimal posterior distribution. By progressively transforming particles sampled from the prior distribution $p(s_{t}|s_{t-1},a)$ through the sequence $\pi_{1},\pi_{2},\ldots,\pi_{K}$, we can effectively approximate particle states drawn from the optimal posterior distribution.

The overall procedure of annealed importance resampling is described in algorithm 4.

To update weights of the particles at iteration $k$ and timestep t, for every $\beta_{k}$ in the sequence $\{\beta_{0},\ldots,\beta_{K}\}$, we transform the particle approximation from $\pi_{k-1}$ to $\pi_{k}$ by recalculating the weights. Assume the particle set at iteration $k$ and timestep t is $S_{tk}$, for each particle $s_{tk}^{j}$ in the particle set $S_{tk}$ at iteration $k$, the weight is updated as:

This update reflects the incremental adjustment of weights based on the annealing parameter $\beta_{k}$. Specifically, at iteration $k$, the importance weight of each particle $s_{tk}^{j}$ is updated by multiplying its previous weight $w_{t(k-1)}^{j}$ with an incremental weight. This incremental weight is computed as the ratio of two observation likelihood terms: the observation density function raised to the current tempering parameter $\beta_{k}$ in the numerator, and the same function raised to the previous tempering parameter $\beta_{k-1}$ in the denominator. Specifically, the incremental weight is $\frac{p(o_{t}|s_{tk}^{j},a_{t})^{\beta_{k}}}{p(o_{t}|s_{tk}^{j},a_{t})^{\beta_{k-1}}}$, where $o_{t}$ is the observation received at timestep $t$ and $a_{t}$ is the action taken.

After updating the particle weights, we calculate their inefficiency score in a similar way as (Herbst and Schorfheide 2019). For a particle set containing $M$ particles, the inefficiency score is computed as:

This score reflects the variance in particle weights. As in (Herbst and Schorfheide 2019) and specified in algorithm 4, we maintain a pre-defined target inefficiency ratio $r^{*}$. If the inefficiency score exceeds $r^{*}$, indicating that the variance of the resulting weights remains high, we continue adjusting the particle states and weights. Otherwise, we exit the current iteration.

For each particle $s_{t(k-1)}^{j}$, the algorithm employs a Markov transition kernel $T_{k}(s_{k-1}^{j},s_{k}^{j})$ to evolve the state from time $k-1$ to $k$. This transition kernel is constructed to preserve $\pi_{k}$ as its invariant distribution, as specified in eq. 4.

The transition mechanism implements a Metropolis-Hastings kernel (Hastings 1970) with a multivariate Gaussian proposal distribution. Given a state vector $s_{t(k-1)}^{j}$ that can be decomposed into components $s_{t(k-1)}^{j}=<x,y,\cdots>$, the proposal distribution is centered at $s_{t(k-1)}^{j}$ with covariance structure. The forward proposal distribution i.e the transition kernel generates new states according to:

where $I$ denotes the identity matrix and $\sigma^{2}$ is a scaling parameter calibrated proportionally to the L1 distance between the current state $s_{t(k-1)}^{j}$ and the observation $o_{t}$ in the state space. Then we define its reverse transition kernel $\tilde{T_{k}}(s_{k}^{j})$ as:

The acceptance rate shown in algorithm 5 is computed as:

Then we update the particle state:

The iterative process continues until either the particle approximation converges to the optimal posterior distribution, or the variance of the particle weights becomes less than or equal to the threshold $r^{*}$.

We evaluate our approach against four state-of-the-art baselines: POMCP (Silver and Veness 2010), ARDESPOT (Ye et al. 2017), POMCPOW (Sunberg and Kochenderfer 2018), and AdaOPS (Wu et al. 2021) on four domains: Light Dark (LD), Tag, Laser Tag, and Rock Sample (RS) (Platt et al. 2010; Egorov et al. 2017; Smith and Simmons 2004). Detailed descriptions of these domains are provided in section A.1. For ARDESPOT and AdaOPS, we employ different initialization strategies: in Light Dark, Tag, and Laser Tag domains, we use domain-specific independent bounds based on maximum and minimum achievable discounted rewards, while for RockSample, bounds are initialized using heuristics. For POMCPOW, we optimize the maxUCB parameter—which governs action selection at each node—by testing values in {1.0, 10.0, 20.0} and selecting the best-performing configuration for each domain. The performance comparison across all domains is presented in table 1. Our experimental results for the baselines may differ from previously published results due to modifications in implementation code, differences in computing infrastructure, and alternative bound initialization methods. Specifically, while (Wu et al. 2021) uses heuristics for bound initialization in Light Dark, Tag, and Laser Tag, we initialize bounds based on maximum and minimum achievable discounted rewards. However, we maintain consistent initialization methods across all approaches within each domain to ensure fair comparison. Complete details of hyperparameter selection for each domain are provided in appendix B.

The main configurations for all baselines are detailed in Baseline Approaches. For fair comparison, we ensure AIROAS’s maximum allowed time per decision matches those of comparable baselines with similar parameters. While POMCP and POMCPOW use iterations rather than time limits, we maintain consistency by keeping the number of iterations identical during testing. All these parameters are documented in appendix B. AIROAS uses the same bound initialization methods as ARDESPOT and AdaOPS. For AIROAS-specific configurations, we generate 100 tempering parameters $\beta_{k}$ ranging from 0 to 1 using the sigmoid transformation described in eq. 9. We evaluate AIROAS using different target inefficiency ratios $r^{*}\in\{2.0,3.0,5.0,10.0\}$ and select the best-performing value for each domain. All experiments were conducted on a computer equipped with an Intel(R) Core(TM) i9-14900KS processor with 32 CPUs.

The average discounted return and its standard error of mean (SEM) are presented in table 1. In all these domains, AIROAS outperforms the other solvers with suitable target inefficiency ratio and tempering parameters.

In the Light Dark Problem, we evaluate performance with step sizes $\alpha=0.5$ and $\alpha=1.0$. The state consists of position coordinates and termination status (terminated or not), while observations provide noisy measurements of the current position. For particle state mutation, we generate a Gaussian distribution over the position components, where the variance is proportional to the distance between the state’s position and the received observation. Through this mutation strategy in bridging distributions, combined with weight adjustments, we observe significant mitigation of the sample impoverishment problem, leading to improved performance.

In the Tag problem, the state consists of the robot’s position, target’s position, and tagging status. While the agent’s position is fully observable, the target’s position is only observed when both actors occupy the same cell. For particle state mutation, we generate a Gaussian distribution exclusively over the target’s position, as the agent’s position is known. Given this limited observability structure, we did not anticipate significant performance improvements from Annealed Importance Resampling. The observed improvements may be attributed to the compact state and observation spaces, where the mutation step effectively explores different states within this limited domain.

The Laser Tag problem represents a more challenging variant of the Tag problem. The state space comprises the agent’s position, target’s position, and terminal signal, with neither position directly observable. The observation consists of noisy range readings from laser sensors in eight different directions. For particle state mutation, we construct a multivariate Gaussian distribution over both agent and target positions. Given these high-dimensional observations, which induce multimodal distributions over the state space, we anticipate Annealed Importance Sampling to demonstrate significant advantages in this domain.

In the RockSample problem, the state space comprises the robot’s position and the quality (good or bad) of each rock. With scenarios containing 11 or 15 rocks, this results in an extremely high-dimensional state space with a large number of states. In contrast, the observation structure is simple, providing information about only one rock’s quality at a time (or no observation). For particle state mutation, we focus specifically on mutating the quality of the rock for which the observation provides information, rather than mutating the entire state vector.

We conduct an ablation study to evaluate the impact of Annealed Importance Resampling (AIR) on our approach. Figure fig. 2 presents a comparative analysis between AIROAS with and without AIR on the Light Dark domain, using a step size of $\alpha=1.0$. To ensure a comprehensive evaluation, we vary the number of particles from the set $\{100,200,500,1000,2000\}$. The results demonstrate that the performance advantage of incorporating AIR becomes more pronounced as the number of particles increases. This suggests that AIR’s particle state and weight mutation strategy becomes increasingly effective with larger particle populations, while maintaining all other components of the algorithm unchanged.