Summary
This paper maximizes the expected reward wrt the policy by rewriting it as a log marginal likelihood in a probabilistic model where latents comprise of states, actions, and “optimality” observations.
The prior over states is the world model and the prior over actions is the policy which being optimized. The likelihood is the exponentiated reward which can be normalized to make it a distribution over auxiliary “optimality” variables.
The marginal likelihood can be optimized via EM where in the E step, we maximize the evidence lower bound (ELBO) in (3) with respect to the variational distribution q and in the M step, we fix q and maximize the objective with respect to the policy π.
In the E-step (Sec. 4.2), the paper derives an analytic expression for q via the Lagrange multipliers method (eq. 6).
In the M-step (Sec. 4.3), the paper uses this q to maximize the objective with respect to π. The key difficulty is that the q derived in the E-step is only known up to a normalizing constant and so the paper proposes to get approximate samples from q using SMC (Sec. 4.2.1).
The experiments compare the proposed method to baselines on several tasks with discrete and continuous action spaces, and show wins. There are ablations showing the tradeoff between using more particles vs unrolling for longer.
In addition to performance, the paper argues that SPO is preferable to the popular MCTS method because it parallelizes more easily and performs well across more domains.
Strengths
Combining SMC and RL is interesting and the experiments show strong empirical gains. However, theoretically, I am not fully satisfied with the method itself—or at the very least, I found it difficult to understand the method fully from the description. See questions / comments below.
Therefore, I’m giving this paper a reject in its current form but I am willing to raise the score if the method or its description is improved.
Questions
Major
- SMC (Sec. 4.2.1): From (6) it seems like there is just one distribution we want to sample from instead of a sequence of them. What’s the sequence of distributions that SMC targets? In other words, in the background section about SMC, it’s unclear what is $p(x)$ and what is $p(x_t | x_{1:t - 1})$; how do terms in (8) relate to this $p$?
- If I reverse-engineer a sequence of targets distribution from (8), it should be $p(a_{1:t}, s_{1:t}) = \prod_{j = 1}^t \mathcal T(s_t | …) \pi(a_t | …) \exp(A^{\bar \pi}(a_t, s_t) / \eta^*)$. That is, the final target distribution is the transition times the policy times the exponent of the *sum* of advantage functions. I think this is incorrect since it doesn’t correspond to the sum of per-step rewards. Something that would make more sense to me is to have the incremental weights contain the exponentiated rewards for all time steps instead of the last one, where it’s the exponentiated advantage function. This is obviously a completely different algorithm.
- I don’t understand why (5) follows from (4). We have an unconstrained optimization of a regularized objective in (4) which is turned into a constrained optimization problem in (5). I also don’t understand how $\alpha$ turns into $\epsilon$.
- Why does (10) rewrite (9) as a constrained objective where it seems like in practice, we’re actually just optimizing (9).
- How is the value function and the Q-function learned? These seem critical for the method but I couldn’t find a description.
- E-step: I agree that it is difficult to optimize with respect to $q$ since it’s the distribution of the expectation and the term inside the expectation, $Q^q$ depends on it (eq. 4). The paper chooses to fix $Q^q$. I would have expected that it is the $q$ in $\mathbb E_q$ that makes optimization difficult since typically we’d have to resort to the reparameterization trick or REINFORCE. If we fix the $q$ in $\mathbb E_q$, optimizing the expectation is simple: just sample from $q$ and directly optimize $Q^q$.
Minor
- $\gamma$ in (3) is not introduced before.
- Should there be a likelihood term in the incremental weight in (2)? Or is $p(x_t | x_{1:t - 1})$ an unnormalized distribution? $p$ is not defined in this paragraph.
- Swapping the prior term $\log p(\theta)$ in (9) to the KL term in (10) is a little odd. Would it simply not work otherwise?
- In App. G.1, why is $\mu_q$ dependent on $q$?
Other relevant work
- [AESMC](https://arxiv.org/abs/1705.10306), [FIVO](https://arxiv.org/abs/1705.09279), [VSMC](https://arxiv.org/abs/1705.11140)
- [DVRL](https://arxiv.org/abs/1806.02426)
- [SIXO](https://arxiv.org/abs/2206.05952), [NAS-X](https://arxiv.org/abs/2308.14864)