L EARNING PROPOSALS FOR SEQUENTIAL IMPORTANCE SAMPLERS USING REINFORCED VARIATIONAL INFERENCE
A BSTRACT
The problem of inferring unobserved values in a partially observed trajectory from a stochastic process can be considered as a structured prediction problem.
Traditionally inference is conducted using heuristic-based Monte Carlo methods.
This work considers learning heuristics by leveraging a connection between policy optimization reinforcement learning and approximate inference.
In particular, we learn proposal distributions used in importance samplers by casting it as a variational inference problem.
We then rewrite the variational lower bound as a policy optimization problem similar to Weber et al. allowing us to transfer techniques from reinforcement learning.
We apply this technique to a simple stochastic process as a proof-of-concept and show that while it is viable, it will require more engineering effort to scale inference for rare observations 1 .
I NTRODUCTION
Stochastic processes are an important model of sequential data in the natural sciences (Allen, 2010)
Learning to ﬁt theoretical models or sampling new sequential data that is consistent with the observed data is crucial for a deeper mechanistic understanding.
Given that data from a stochastic process is ordered 2 in a trajectory, predicting the values of unobserved variables in a partially observed trajectory can be considered as a structured prediction problem.
For example, disease progression during an outbreak can be modeled by considering how the number of healthy, sick and recovered individuals changes every day (Allen, 2008).
Given observations from an epidemic where we counted individuals towards the end of the outbreak, could we infer how the disease progressed at the start of the outbreak?
Approximating these distributions directly from data can be difﬁcult and multiple methods based on heuristics have been proposed.
In particular, many techniques use handcrafted Monte Carlo methods based on importance samplers, Markov Chain Monte Carlo or rejection sampling schemes (Gilks et al., 1995; Geyer & Thompson, 1995; Jiao et al., 2014; Sharma, 2017; Nelson et al., 2018).
The promise of automatically learning heuristics from raw and complex data is one of the tantalizing goals of machine learning: It might allow us to use Monte Carlo methods for problems in which heuristics are not obvious.
Variational inference (VI, Blei et al.
(2006; 2017)) is an alternative technique to approximate probability distributions that has seen increasing success when combined with deep learning (Kingma &
Welling, 2013; Goodfellow et al., 2016; Gulrajani et al., 2016;
Miao et al., 2016) in high dimensional datasets.
Our work adds to existing techniques that enhance Monte Carlo methods using variational inference (Wexler & Geiger, 2007; Salimans et al., 2015; Naesseth et al., 2017; M¨uller et al., 2018).
Speciﬁcally, we consider the sampling of trajectories from the posterior distribution of a stochastic process using importance samplers.
We formulate learning of proposal distributions as a VI problem.
We then leverage the connection between VI and policy optimization reinforcement learning (Sutton et al., 2000; Bachman & Precup, 2015; Weber et al., 2015) to learn these proposal distributions.
We apply reinforced variational inference (RVI, Weber et al.
to a simpliﬁed stochastic process as a proof-of-concept where we aim to sample trajectories, τ , from a posterior over trajectories, p ( τ | x T ) , given a single terminal observation, x T .
Our contributions are:
Empirical validation of the RVI framework (Weber
et al.,
suggests that RL is a viable approach for VI and more advanced tools from policy optimization might carry over to VI. 2.
Learning of proposal distributions in sequential importance samplers can result in performance that is close to a hand-crafted importance sampler for most observed data.
Qualitative evaluation demonstrates the reasonability of proposals found but quantitative evaluation identiﬁes room for improvement.
B ACKGROUND
In this section we describe technical background information regarding importance sampling, variational inference and reinforcement learning.
These notions will be crucial to derive a variational approach to importance sampling, trained using policy optimization.
2. I MPORTANCE SAMPLING
The expectation of function f under a distribution p ( τ | x T ) can be estimated by drawing a large number of samples or particles , τ i ∼ p ( τ | x T ) and computing the sample mean: E p ( τ | x T )
[ f ( τ )]
= (cid:90) p ( τ | x T ) f
( τ ) d τ ≈ 1 N (cid:88) i f
( τ i ) .
In cases where p ( τ | x T ) is difﬁcult to sample from or the integral requires evaluating rare samples, we can rewrite the integral and expectation by drawing instead from a proposal distribution, τ i ∼ q ( τ ) , that is easy to sample from:
[Equation]
This approach is known as importance sampling .
We can think of the importance weight , w i = p ( τ | x T ) q ( τ ) , as a way to correct for the relative probability of drawing τ according to q ( τ ) instead of p ( τ | x T ) .
Though Equation 1 is valid for any distribution q , the speed and quality of the approximation depends on q .
In general, designing a q ( τ ) that is better than random is difﬁcult but can be done for some simple posteriors (See Section S.1.2).
2. V ARIATIONAL I NFERENCE
Variational inference (VI) is an alternate approach to sample from a difﬁcult distribution, p ( τ ) , by sampling from a learnable approximate distribution q ( τ ; θ ) instead.
The variational parameters θ are learned such that they minimize a divergence metric like the Kullback-Leibler or χ 2 divergence with the true posterior: q ∗ ( τ ) =
arg min θ D ( q ( τ ; θ ) || p ( τ | x T )) .
This procedure effectively turns the inference problem into a optimization problem (Blei et al., 2017; Dieng et al., 2017).
2. R EINFORCEMENT L EARNING AND P OLICY G RADIENT M ETHODS
In reinforcement learning (RL) (Sutton & Barto, 1998), we search for a policy π that dictates how an agent should interact with the environment in order to maximize the cumulative reward.
Policybased methods (Sutton et al., 2000) involve parameterizing the policy π θ ( a | s ) and then directly optimizing: O ER =
E τ ∼ π
θ
[ (cid:80) Tt =
γ t r t ] where r t is the reward obtained at time t by executing π θ in the environment and γ is a discount factor that prioritizes rewards to be sooner rather than later.
The gradient of O ER can be approximated using the R EINFORCE (Williams, 1992) estimator: [Equation] ≈ (cid:88) [Equation]
[Equation]
) where Ψ t is a target that describes how good the sampled action, a t , was on the long term reward.
We use the generalized advantage estimate (GAE) (Schulman et al., 2015b) as the target: Ψ t = (cid:80) ∞
l =0 ( γλ ) l δ t + l where δ t + l = r t + γV π
( s t +1 )
− V π ( s t ) is the temporal difference residual (Sutton & Barto, 1998) and V π is a parameterized value function that estimates the long term discounted reward.
We also experiment with an entropy regularized version of this algorithm where the entropy of the policy, H ( π ) , is added to the objective to prevent learned policies from becoming deterministic too quickly (Schulman et al., 2017a).
A PPROACH USING R EINFORCED V ARIATIONAL I NFERENCE (RVI)
We begin by casting the problem of learning a proposal distribution as a variational inference problem in Section 3.1.
We then draw the connection with policy optimization in Section 3.2 which results in a similar formulation to Weber et al.
3. T HE V ARIATIONAL I NFERENCE P ROBLEM
Recall that we are interested in sampling from the posterior over trajectories, p ( τ | x T ) given a single terminal observation, x T .
We consider Markov processes where the recursive relationship x t = x t − 1 + ∆ x t holds between random variables.
Only a ﬁxed number of transitions, ∆ x t , are possible and we model this as a discrete random variable with probability distribution p (∆ x | x t − 1 ) .
The probability of observing a trajectory from the stochastic process τ = ( x 0 , . . .
, x T ) can be decomposed as a product of the transitions: p ( τ ) = p ( x 0 ) (cid:81) T − 1 t =0 p (∆ x t | x t − 1 ) .
Recall that variational inference introduces a simpler approximating posterior.
Let q ( (cid:126) τ | x T ) be the approximate posterior over trajectories that can be written as a product of local parameterized conditionals called proposals :
q ( (cid:126) τ | x T ) =
(cid:81) T − 1 t =0 q θ (∆ x t | x t +1 ) . Since we are given the ground-truth terminal observation, x T , the trajectory is created by going backwards in time starting the sampling from q θ (∆ x T − 1 | x T ) , computing the next state x T − 1 and repeatedly sampling from q θ (∆ x t − 1 | x t ) until reaching t = 0 .
We can construct the trajectory by using x t = x t +1 − ∆ x t where ∆ x t is one of the outcomes possible under the forward transition probabilities.
The KLD cannot be minimized exactly, and instead we maximize the evidence lower bound (Blei et al., 2017):
[Equation]
In the next section, we will see how to transform the maximization of the ELBO into a policy optimization problem and draw the connection with rewards in reinforcement learning.
3. T HE C ONNECTION TO P OLICY O PTIMIZATION
To demonstrate the link with RL, we substitute the full forms of the variational distribution, q ( (cid:126) τ | x T ) and trajectory probability, p ( τ ) into the ELBO (Equation 3):
[Equation]
By comparing with R EINFORCE (Equation 2), Weber et al.
(2015) view this as a RL problem where at state x t +1 the policy q takes an action ∆ x t in an environment and receives a reward r t = log p (∆ x t | x t ) q θ (∆ x t | x t +1 ) and moves to state x t .
In the ﬁnal step, we also obtain r f = log p ( x 0 ) as a reward.
The reward has an appealing interpretation: It is the relative log likelihood of executing the action ∆ x t in the stochastic process versus that of the policy.
It encourages the policy to choose actions that would be likely under the generative stochastic process.
We write Equationas E q θ [ (cid:80) t r t + r f ] , to show the equivalence to the RL objective with no discounting 3 , i.e. γ = 1 .
Equation 5 suggests that we can use off-the-shelf reinforcement learning algorithms to optimize the parameters of the proposal distribution q ( τ ) by learning the one step proposal q θ (∆ x t | x t +1 ) as the policy.
An algorithm to obtain a trajectory from proposal and subsequently update it would proceed in three phases:
Collect a trajectory: Sample ∆ x t ∼ q θ (∆ t | x t ) and update the state x t − 1 = x t − ∆ x t .
Now x t − 1 is used to obtain the next sample and is repeated until t = 0 .
Update the proposal: For each ∆ x t that is obtained, rewards are calculated as r t = max( R min , log p (∆ x t | x t ) q θ (∆ x t | x t +1 ) ) and r f = log p ( x 0 ) and used to compute the empirical return 4 .
The gradient is then estimated using R EINFORCE and a GAE target and used to update the parameters via stochastic optimization.
Update the posterior: If the ﬁnal position x 0 is valid under the prior,
i
e. p ( x 0 ) > 0 , then the trajectory is saved with the weight w i = p ( x 0 ) (cid:81)
t p (∆ x t | x t ) q θ (∆ x t | x t +1 ) .
3. U SING THE χ 2 DIVERGENCE
In this section, we will motivate the χ 2 divergence as alternative to the KLD.
We then demonstrate the ﬂexibility of the RVI framework by extending it to minimize this divergence and making a connection to sparse reward reinforcement learning.
Assume that we want to estimate the likelihood of an observation x T , which requires taking an intractable integral over all possible paths: p ( x T )
= (cid:82) p ( x T , τ ) dτ .
We can use importance sampling with a parameterized proposal distribution q θ ( τ ) so that we can estimate p ( x T ) with fewer samples and lower variance.
Speciﬁcally,
[Equation]
where τ ∼ q θ ( · ) and w τ = p ( x T ,τ ) q θ ( τ ) are the importance weights.
Burda et al.
(2015) used such an importance sampling estimate of the log-likelihood to improve variational autoencoders.
One property we might desire from the importance sampling proposal is that it minimizes the variance of ˆ p ( x T ) :
[Equation]
which corresponds to minimizing the χ 2 -divergence : D χ 2 ( p [Equation] =
[Equation]
p
[Equation]
q
[Equation]
(cid:105) .
In Dieng et al.
(2017), the likelihood ratio estimate for the gradient is derived for the exponentiated “ χ 2 Upper Bound” or CUBO,
12 log E q
[ (cid:0) p ( τ ) q ( τ ) (cid:1) 2 ] , which is an upper bound on D χ 2 ( p ( τ ) ||
q ( τ )) :
[Equation]
Unlike the KLD, the χ 2 -divergence cannot be broken down into a summation.
Instead, we can treat − (cid:16) p ( τ ) q
( τ ) (cid:17) 2 as a sparse reward that only appears at the end of a trajectory.
Since we still build trajectories by sampling from the local conditionals q θ (∆ x t | x t +1 ) , we can never assign credit to any particular action.
Sparse reward problems are generally quite difﬁcult as there are no intermediate signals to guide learning (Ng et al., 1999).
Our work will only evaluate the feasibility of replacing the KLD with the χ 2 and we will leave investigation of densiﬁcation to future work.
R ELATED W ORK
Many efforts have been made to combine variational optimization and Monte Carlo methods.
We cover a few approaches here and, where relevant, describe connections and deviations with existing literature.
See Appendix S.2 for a wider discussion on the connections between policy optimization, variational inference and importance sampling.
We note that the formulation provided here was also provided in Weber et al.
for more generic stochastic computation graphs.
Our work grounds this framework in a concrete approximate inference problem.
Grover et al.
(2018) used rejection sampling to minimize the KL divergence using a resampled
ELBO: Proposed samples are rejected if they have a low likelihood under the model resulting in a resampled proposal distribution.
This is particularly similar to our work as their accept-reject step is differentiable which means it can be adapted over time.
The most similar to our work is Gu et al.
and Naesseth et al.
(2017) who learned proposal distributions in the sequential Monte Carlo setting.
Our work deviates in one crucial way: we learn discrete transitions between states in the sequence.
This allows some increased ﬂexibility at the cost of being potentially high variance due to the gradient estimator.
Indeed Naesseth et al.
(2017) use the reparameterization technique (Kingma & Welling, 2013) that provides a low variance gradient estimator for continuous distributions 5 .
In particular, Naesseth et al.
(2017) derive a differentiable ﬁltering algorithm that combines variational optimization and an expectation-maximization algorithm to automatically learn both the proposal and forward model.
Recently M ¨ uller et al.
(2018) used variational inference by optimizing for proposals represented by normalizing ﬂows (Rezende & Mohamed, 2015) to learn sampling distributions for constructing light paths .
Like this work, the authors optimize the χ 2 divergence and obtained similar qualitative results.
Markov Chain Monte Carlo (MCMC) methods are another family of techniques to approximate probability distributions.
Salimans et al.
(2015) used variational optimization to learn the transition operator of the Markov chain that is used in MCMC algorithms.
Though this kind of algorithm can be extended to the sequential setting it would require sampling many times from the Markov
Chain.
Our work avoids the use of a transition operator for the Markov chain by directly requiring the proposal distribution to produce a valid transition within the stochastic process itself.
R ESULTS
In this section we will describe some quantitative and qualitative empirical results that show the viability of the problem formulation and approach using RVI.
We will also highlight failure modes of the method.
5. E XPERIMENTAL S ETUP
Environments: We consider stochastic processes based on random walks as environments (Section S.1).
They are chosen speciﬁcally to demonstrate the usefulness of the technique and not because they are particularly complex.
Random walks form the basis for more complex processes.
[Figure] Speciﬁcally, we use the simple 1D random walk whose initial point is described by a discrete uniform distribution with a window (-5, 5) and evolves by taking a step +
or -1 with equal probability for T = 50 timesteps (See complete details in Section S.1.1).
We control the ending point, x T , to modulate the difﬁculty of the posterior estimation problem, i.e. as x T → 50 the observation becomes rarer and the posterior over starting positions becomes skewed towards values near the edges of the window. The proposal, q θ , is a categorical distribution over possible transitions, ∆ x t .
Quantitative Evaluations: To measure the online performance of importance samplers, we report the KL Divergence between the approximated and true posterior 6 over starting positions (KLpq).
Note that the KLpq is only one downstream metric of interest and presents some limitations.
To avoid overﬁtting to one metric, we also consider qualitative behaviours to better understand the different components of the learned proposals and the posteriors.
Qualitative Evaluations: To understand the adaptive behaviour of the proposal, we observe how trajectories sampled from the proposal change during the optimization procedure.
We investigate the reasonability of proposal learned by visualizing the expected direction a particle will move at any point x t at any time t .
Finally, we use a histogram to compare the approximated posterior over starting positions to the true posterior to understand the sensitivity in KLpq.
Baselines: We consider two baselines: (1) Naive Monte Carlo (MC) where the transition probability p ( x t | x t − 1 ) is used as the proposal q ( x t − 1 | x t ) ; and (2) Hand-crafted importance sampler with a soft proposal (ISSoft) where the backward probability is biased toward the window depending on the average number of steps needed to be within the range of the window 7 .
5. Q UANTITATIVE P ERFORMANCE
We ﬁrst investigate the viability of learning proposals online that are speciﬁc to each ending observation.
Unsurprisingly, ISSoft performs well for all ending observations and is therefore a simple but strong baseline to beat.
After sampling and training with 250,000 trajectories RVI-KL performs as well as ISSoft for observations that are close to the window x T ∈ { 0 , 12 } .
For observations that are far away from the window, x T ∈ { 24 , 36 } , RVI-KL performs better than the random MC method but not as good as the hand-crafted ISSoft (Figure 1).
5. Q UALITATIVE ANALYSIS
In this section we take a look at qualitative aspects of learning proposals as well as ﬁnal posteriors learned.
To understand if the posterior samples better trajectories after successive steps of optimization, we measure the proportion of trajectories that end within the window during optimization.
Compared to the start of the optimization, a larger number of trajectories sampled from the proposal end within the window (Figure S1). [Figure]
[Figure] To understand the poor quantitative performance (Figure 1), we visualize the posterior over starting positions, p ( x 0 | x T ) , obtained by RVI-KL (Figure 2) compared to the true analytics posterior.
For ending points that are close to the window, the posterior recovered is near-perfect for both 1D (Figure 2a and 2b) and 2D random walks (Figure S2).
Even though we were unable to record quantitative performance for x T = 48 (Figure 1), the posteriors found captured the general shape, and only missed the low probability mass values of x 0 (Figure 2c).
Visualizing the proposal reveals sufﬁcient specialization for each ending observation (Figure 3).
Particularly interesting is the progressively skewed proposals obtained when the ending point changes from x T = 0 to x T = 36 .
Indeed the proposal learned when x T = 48 is not completely unreasonable as it prevents the particle from moving into the region where x i < 0 .
[Table]
5. A N ALTERNATE OBJECTIVE FUNCTION
In this section, we devise an experiment to probe whether we are using the right objective function.
In particular, we pre-train a proposal by using the algorithm outlined in Section 3.2 and changing the ending observation, x T , after each trajectory (Figure S3a) 8 .
Evaluation of the pre-trained proposal by measuring KLpq after 100,000 trajectories reveals that it performs on-par with ISSoft and better than training from scratch:
The model is of sufﬁcient representational capacity and the pre-training scheme can produce a proposal that generalizes well to different end points (Table 1,
RVI-KL+pre+eval).
Next, we repeat the experiment but this time allowing the proposal to be ﬁnetuned (Figure S3b and S3c).
Surprisingly, performance becomes worse compared to no ﬁne-tuning at all: Even though the initialization is good, the training scheme resulted in worse performance.
One possible reason could be the use of the wrong objective function (Table 1, RVI-KL+pre+ﬁne).
In Section 3.3 we showed how minimizing the χ 2 divergence is equivalent to minimizing the variance of the estimator of the log-likeihood.
Furthermore, we showed how it could be used in the RVI framework.
We ﬁnd that this technique does not perform well (Table 1, RVI-C2) compared to using
KLD (Table 1, RVI-KL).
Despite not having good quantitative performance, the correct shape of the approximate posterior has been captured (Figure S4).
D ISCUSSION AND F UTURE W ORK
Our experiments have shown that policy optimization can be a viable approach to learn interesting proposals that, in some instances, performs on par with hand-crafted ones.
This result shows us the viability of two things: (1) Framing sequential variational inference problems as policy optimization problems; (2) Learning proposals that can approximate posteriors for stochastic processes.
These observations are encouraging and suggest that more advanced techniques from policy optimization, like generalized advantage estimation (Schulman et al., 2015b) or asynchronous advantage actor critic (Mnih et al., 2016), can be brought to variational inference 9 .
Therefore, RVI can learn proposals for simple stochastic processes.
Despite the approach being easily extended to different stochastic processes given the forward dynamics, it might require signiﬁcant future work to be able to learn posterior distributions for more advanced stochastic processes.
We identify at least two areas for improvement.
In particular, poor results on rare instances of x T that are far from the window highlights the need for better data sampling strategies 10 .
Secondly, it is somewhat paradoxical that RVI cannot learn to improve upon pre-trained proposals:
The local optimum obtained by training on the distribution of ending points, i.e. pre-trained, is worse than if we trained from scratch or specialized the pre-trained distribution. This means that the KLD objective function being optimized is not reﬂective of the kind of downstream tasks we want to evaluate the model on. One way to check this would be to directly optimize for the downstream metric and verify the quality of the KLD objective and will be left for future work. The latter part of this work considered alternative objectives such as the χ 2 divergence, but is inconclusive as to its merits.
One possible solution to make transfer more feasible is to impose a penalty that ensures that proposals do not deviate from the pre-trained proposal unless there is a signiﬁcant beneﬁt and is a known approach in policy optimization (Schulman et al., 2015a).
Perhaps the gradients are too noisy to ﬁnd a better optimum and a better multi-sample gradient estimator (Mnih & Rezende, 2016) or variance reduction techniques should be used.
In this work we have investigated and empirically validated that we can leverage connections between policy optimization and variational inference to learn proposals in importance sampling.
Our results are encouraging but also highlight the amount of engineering needed to get these methods to work for more difﬁcult problems.
We observed progressively improved optimization by using more sophisticated targets, reinforcement
A CKNOWLEDGMENTS
The authors would like to thank Riashat Islam for detailed feedback on this work as well as the helpful comments from anonymous workshop reviewers. This work relied on funding from the Canada Graduate Scholarship program via the Canadian Institutes of Health Research and Google Summer of Code.
Computational resources were provided from Compute Canada and its regional partners.
D ERIVING THE ANALYTIC POSTERIOR FOR THE RANDOM WALK
Though sampling from the posterior over trajectories, p ( τ | x T ) , can be difﬁcult for arbitrary stochastic processes, in the case of the simple random walk described in Section S.1, the posterior over starting states, p ( x 0 | x T ) , can be calculated exactly.
We use this result to compute a downstream measure for the performance of the approximate distribution obtained by different methods considered in the experimental section.
For the purposes of clarity, we will override notation for this section: Let S t = ∆ X t as deﬁned in Equation 11; X 0 the random variable describing the starting position; and D =
X T the random variable describing the ﬁnal position.
Our goal is to derive the posterior over starting states, p ( X 0 | D ) .
Using Bayes Rule:
[Equation]
Therefore, we can ﬁrst derive p D | X
( D | X 0 ) and then apply Bayes rule to obtain the desired posterior distribution.
[Table]
method
ISSoft 0 .
000033 ± 0 .0 .
000039 ± 0 .0 .
000058 ± 0 .0 .
000068 ± 0 .0 .
000036 ± 0 .
RVI-KL
[Table]
pre+eval 0 . 000034 ± 0 .0 .
000055 ± 0 .0 .
000055 ± 0 .0 .
000207 ± 0 .
NaN RVI
[Table]
KL
[Table]
pre+ﬁne 0 .
000034 ± 0 .0 .
000108 ± 0 .0 . 003550 ± 0 .
NaN
NaN RVI
[Table]
KL 0 .
000030 ± 0 .0 .
000118 ± 0 .0 . 000720 ± 0 .0 . 005527 ± 0 .
NaN RVI-KL+ent 0 . 000031 ± 0 .0 . 000093 ± 0 .0 .
000661 ± 0 .
[Table]
Since we are taking steps of size +1 or -1, the random variable describing our step, S t , is a shifted
Bernoulli, i.e. S t = 2 X t − 1 where X t ∼ Bern ( p ) .
Explicitly, S t has the following distribution:
[Equation]
If the walk is of length T , the distribution of the total displacement due to the steps taken, S T = (cid:80)
Ti =1 S t , is given by:
[Equation]
We know that (cid:80) T i =
1 X t is distributed according to the Bernoulli distribution.
Therefore:
[Equation]
We now have all the terms necessary to compute the numerator of Equation 12.
We ﬁrst consider p X
( X 0 ) .
Since it is drawn uniformly within a window between − c and + c , i.e. X 0 ∼
DiscreteUniform ( − c, c ) , the probability mass function is given by
[Equation]
Next, we focus on p D | X
( D | X 0 ) .
Note that the total distance moved is D =
X 0 + S T .
Therefore, we can write:
[Equation]
We then have that P X 0 | D ( d, x ) ∝ P S T ( d − x ) P X 0 ( x ) .
This lets us write the posterior as:
[Equation]
if x ∈ [ − c, c ] and (cid:4) d − x + T(cid:5) ∈ [0 , T ] and 0 otherwise.
(cid:0) ab (cid:1)
=
a !
( a − b )!
b ! is the binomial coefﬁcient.
S.1.
T HE DIFFICULTY OF DESIGNING GOOD PROPOSALS
Consider the simple example of sampling trajectories from the posterior of the random walk introduced in Section S.1.
Since the random walk is unbiased 11 , we could use a Naive MC method: construct a trajectory backward, (cid:126) τ 12 , by considering a new walk starting from x T .
Speciﬁcally, we sample x t according to q ( x t | x t +1 )
= p ( x t +1 | x t ) .
The process described works remarkably well but poses a problem when x T has drifted far from the window of possible starting positions.
Sampling backward according to the forward probabilities will likely result in the ﬁnal sampled position, ˆ x 0 , being close to x T 13 .
To combat this, we can use an importance sampling scheme to sample from a q that is different from p , but ensures that we obtain trajectories that are biased towards the window.
This will result in x 0 ’s that are more likely under the prior and therefore a better estimate of the posterior.
The question now becomes How would we design a proposal q ?
For the problem of the random walk, designing a q that is better than random is easy: Given that we know particles started in a window deﬁned between − c and c , we can design a proposal that pushes x t toward the window.
We ﬁrst calculate a bias term, b , which measures, on average, how many steps are needed to move toward the window at any given time t during the walk.
We use a softness coefﬁcient, s , to deﬁne how strong the bias should be:
[Equation]
where T is the total length of the walk.
If the particle at time t , x t , drifts too far from the window, i.e. the average number of steps needed is larger than the step size , | b | > 1 , it is not possible to return to a position where the window is in reach within a reasonable time frame and a random action is executed instead.
Otherwise, with probability p = 1 − | b | we take a random step and probability 1 − (1 − | b | ) a step in the bias direction.
Speciﬁcally, the probability of each action is:
[Equation]
where 1 is the indicator function: It is 1 when the argument is true and 0 otherwise.
S.
C ONNECTIONS BETWEEN IMPORTANCE SAMPLING , VARIATIONAL INFERENCE AND POLICY OPTIMIZATION I MPLICATIONS FOR IMPORTANCE SAMPLING
In this work we have viewed sequential importance sampling as a variational inference problem.
We demonstrated the ﬂexibility of the approach in learning proposals to situations where they are not obvious.
The implication of this result is encouraging:
Applying importance sampling to nontraditional problems where proposals are not obvious.
Given that VI works well with high dimensional real world sequential data (Bachman & Precup, 2015; Goyal et al., 2017; M ¨ uller et al., 2018)
we conjecture that it can be used to learn complex proposal distributions in domains beyond vision and text.
Given that the sequential KLD breaks down into sums of rewards, we leveraged the fact that we did not need to optimize the joint probability of the whole trajectory: The transition at every step could be weighted by the action probability similar to the per-decision importance sampler in Precup (2000).
I MPLICATIONS FOR VARIATIONAL INFERENCE
Our work has shown the viability of using policy optimization algorithms for sequential variational inference problems.
This suggests that state-of-the-art methods from reinforcement learning can be brought into variational inference (Kakade, 2002; Schulman et al., 2015a; 2017b) for learning approximating distributions.
For example, we transferred the concept of “ λ -return” using generalized advantage estimation to sequential VI (Schulman et al., 2015b).
Beyond the obvious scheme of transferring methods, we can also consider more subtle concepts from reinforcement learning.
For example, we can consider the KLD terms in the RVI framework as a generalized value function and combine these with auxiliary tasks (Sutton et al., 2011) that can stabilize learning.
Concepts like eligibility traces for credit assignment might be able to allow us to use a more diverse set of reward speciﬁcations (Singh & Sutton, 1996).
I MPLICATIONS FOR REINFORCEMENT LEARNING
Though we have explored probabilistic inference as reinforcement learning, there has been a renewed interest in viewing reinforcement learning as probabilistic inference (Levine, 2018).
All these connections are under-explored and will be the subject of new research in the next few years.
The KLD makes an appearance in a variety of policy optimization algorithms to regularize policy improvement steps.
Perhaps some work should be put into exploring different divergences.
For example, Neumann (2011) uses the reverse KLD to obtain policies with more cautious behaviour.
More directly related to this work is the concept of learning a backward model and was successfully introduced as recall traces in Goyal et al.
Most recently, Piche et al.
(2019) have successfully used a (learned) sequential Monte Carlo algorithm to learn promising future trajectories to plan in continuous control reinforcement learning problems.
Signiﬁcant research will surface in the next few years from exploring this connection more closely.
Policy optimization is a special case of gradient descent on an importance sampled loss function (Jie & Abbeel, 2010): When estimating importance sampled objectives, we could re-use trajectories collected at any time before a given update to estimate a gradient.
Mnih & Rezende (2016) shows an approach to multi-sample R EINFORCE -like objectives that may be useful in policy optimization when using multiple samples.
The idea further generalizes by thinking of off-policy learning where we have a target policy, π , that we wish to learn and a behaviour policy µ from which data is collected.
The forward model serves the same purpose as µ and the proposal of π .
Revisiting these notions of off-policy policy optimization by using tools from variational inference is a promising direction (Fellows et al., 2018).
S. I MPLEMENTATION D ETAILS
We represent policy and value functions using two independent three-layer neural networks with
RELU non-linearities (Glorot et al., 2011) implemented in PyTorch (Paszke et al., 2017).
Each layer has 32 hidden units.
We use the RMSProp update rule during optimization (Tieleman & Hinton, 2012).
We make our code open source at https://github.com/zafarali/better-sampling.
To pre-train the proposal distribution we train q on a range of ﬁnal positions, x T , obtained by sampling an initial point, x 0 ∼ p ( x 0 ) , and doing a rollout. This allows q to be trained from the natural distribution of ending points. To ﬁne-tune the model, we take the pre-trained model from above and continue training with x T held ﬁxed to the end point of interest.