SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Aoyang Qin\*,1,2 Feng Gao\* Qing Li2 Song-Chun Zhu1,2,4 Sirui Xie\*,5
$^{1}$  Department of Automation, Tsinghua University  $^{2}$  Beijing Institute for General Artificial Intelligence (BIGAI)  $^{3}$  Department of Statistics, UCLA  $^{4}$  School of Artificial Intelligence, Peking University  $^{5}$  Department of Computer Science, UCLA
Conventional imitation learning assumes access to the actions of demonstrators, but these motor signals are often non-observable in naturalistic settings. Additionally, sequential decision-making behaviors in these settings can deviate from the assumptions of a standard Markov Decision Process (MDP). To address these challenges, we explore deep generative modeling of state-only sequences with non-Markov Decision Process (nMDP), where the policy is an energy-based prior in the latent space of the state transition generator. We develop maximum likelihood estimation to learn both the transition and the policy, which involves short-run MCMC sampling from the prior and importance sampling for the posterior. The learned model enables decision-making as inference: model-free policy execution is equivalent to prior sampling, model-based planning is posterior sampling initialized from the policy. We demonstrate the efficacy of the proposed method in a prototypical path planning task with non-Markovian constraints and show that the learned model exhibits strong performances in challenging domains from the MuJoCo suite.
Imitation from others is a prevalent phenomenon in humans and many other species, where individuals learn by observing and mimicking the actions of others. An intriguing aspect of this process is the brain's ability to extract motor signals from sensory input. This remarkable capability is facilitated by mirror neurons [1, 2], which respond to observations as if the imitator is performing the actions themselves. In conventional imitation learning [3, 4] and offline reinforcement learning [5], action labels have served as proxies for mirror neurons. But it is important to recognize that they are actually productions of human interventions. Given the recent advancements in AI, now is probably an opportune time to explore imitation learning in a more naturalistic setting.
While the setting of state-only demonstrations is not common, there are certain exceptions. For example, Inverse Reinforcement Learning (IRL) initially formulated the problem as state visitation matching [6], where demonstrations consist solely of state sequences. Subsequently, this state-only setting was rebranded as Imitation Learning from Observations (ILfO), which introduced the generalized formulation of matching marginal state distributions [7, 8]. These methods typically rely on the Markov assumption and Temporal Difference (TD) learning techniques [9]. One consequence of this assumption, previously believed to be advantageous, is that sequences with different state orders are treated as equivalent. However, the success of general sequence modeling [10] has challenged
* indicates equal contribution. Correspondence: Sirui Xie (srxie@ucla.edu). Code and data are available at https://github.com/qayqaq/LanMDP
(a) Policy and transition in MDP
(b) Policy and transition in nMDP
Figure 1: Graphical model of policy and transition in standard Markov Decision Process and non-Markov Decision Process. Reward variables are omitted in the probabilistic graph to emphasize the difference in dependency between actions and states. nMDP is a natural generalization of standard MDP.
this belief, leading to deep reflections. Notable progresses since then include an analysis of the expressivity of Markovian rewards [11] and a series of sequence models tailored for decision-making problems [12-15]. Aligning with this evolving trend, we extend the state-only imitation learning problem to encompass non-Markovian domains.
In this work, we propose a generative model based on non-Markov Decision Process (nMDP), in which states are fully observable and actions are latent. Unlike existing monolithic sequence models, we factorize the joint state-action distribution into policy and causal transition according to the standard Markov Decision Process (MDP). To further extend to the non-Markovian domain, we condition the policy on sequential contexts. The density families of policy and transition are consistent with conventional IRL [4]. We refer to this model as Latent-action non-Markov Decision Process (LanMDP). Because the actions are latent variables following Boltzmann distribution, the present model is closely related to the Latent-space Energy-Based Model (LEBM) [16]. To learn the latent policy by Maximum Likelihood Estimation (MLE), we need to sample from the prior and the posterior. We sample the prior using short-run Markov Chain Monte Carlo (MCMC) [17], and the posterior using importance sampling. Specifically, the proposed importance sampling sidesteps back-propagation through time in posterior MCMC with a single-step lookahead of the Markov transition. The transition is learned from self-interaction.
Once the LanMDP is learned, it can be used for policy execution and planning through prior and posterior sampling, or in other words, policy as prior, planning as posterior inference [18, 19]. In our analysis, we derive an objective of the non-Markovian decision-making problem induced from the MLE. We show that the prior sampling at each step can indeed lead to optimal expected returns. Almost surprisingly, we find that the entire family of maximum entropy reinforcement learning [4, 20-24] naturally emerges from the algebraic structures in the MLE of latent policies. This formulation avoids the peculiarities of maximizing state transition entropy in prior arts [20, 24]. We also show that when a target goal state is in-distribution, the posterior sampling is optimizing a conditional variant of the objective, realizing model-based planning. In our experiments, we validate the necessity and efficacy of our model in learning to sequentially plan cubic curves, and illustrate an over-imitation phenomenon [25, 26] when the learned model is repurposed for goal-reaching. We also test the proposed modeling, learning, and computing method in MuJoCo, a domain with higher-dimensional state and action spaces, and achieve performance competitive to existing methods, even those that learn with action labels.
2 Non-Markov Decision Process
The most well-known sequence model of a decision-making process is Markov Decision Process. A MDP is a tuple  $\mathcal{M} = \langle S, A, T, R, \rho, H \rangle$  that contains a set  $S$  of states, a set  $A$  of actions, a transition  $T: S \times A \mapsto \Pi(S)$  that returns for every state  $s_t$  and action  $a_t$  a distribution over the next state  $s_{t+1}$ ; a reward function  $R: S \times A \mapsto \mathbb{R}$  that specifies the real-valued reward received by the agent when taking action  $a_t$  in state  $s_t$ ; an initial state distribution  $\rho: \Pi(S)$ ; and a horizon  $H$  that is the maximum number of actions/steps the agent can execute in one episode. A solution to an MDP is a policy that maps states to actions,  $\pi: S \mapsto \Pi(A)$ . The value of policy  $\pi$ ,  $V^{\pi}(s) = \mathbb{E}_{T,\pi}[\sum_{t=0}^{H} R(s_t)|s_0 = s]$  is the expected cumulative reward (i.e., return) when executing with this policy starting from state  $s$ . The state-action value of policy  $\pi$  is  $Q^{\pi}(s_t, a_t) = R(s_t, a_t) + \mathbb{E}_{T(s_{t+1}|s_t, a_t)}[V^{\pi}(s_{t+1})]$ . The optimal policy  $\pi^*$  can maximize either  $E_{\rho(s_0)}[V^{\pi}(s_0)]$ , or the same objective plus the policy entropy [27, 4, 22]. The Markovian assumption supports the convergence of a series of TD-learning methods [9], whose reliability in non-Markovian domains is still an open problem.
A non-Markov Decision Process is also a tuple  $\mathcal{M} = \langle S, A, T, R, \rho, H \rangle$ . It generalizes MDP by allowing for non-Markovian transitions and rewards [28]. Notably, assuming Markovian transition and non-Markovian reward is usually sufficient since a state space with non-Markovian transition can be represented with its Markov abstraction [29]. Markov abstraction can be done either by treating the original space as observations generated from the latent belief state in a Partially Observable Markov Decision Process (POMDP) [30], or by projecting historic contexts into an embedding space for sequence pattern detection [31, 32, 28]. Presumably, it is statistically more interesting in deep learning to focus our attention on non-Markovian domains where the temporal dependencies in transition and reward differ. Therefore, without loss of generality, we assume that the state transition is Markovian  $T: S \times A \mapsto \Pi(S)$ , while the reward is not [33, 11], i.e.  $R: S^{+} \mapsto \mathbb{R}$ , with  $S^{+}$  denotes the set of all finite non-empty state sequences with length smaller than  $H$ . Obviously, the policy should also be non-Markovian  $\pi: S^{+} \mapsto \Pi(A)$ . Check Figure 1 for a probabilistic graphical model of the generation process of state sequences from a policy.
3 Learning and Sampling
3.1 Latent-action nMDP
A complete trajectory is denoted by
$$
\zeta = \left\{s _ {0}, a _ {0}, s _ {1}, a _ {1}, \dots , a _ {T - 1}, s _ {T} \right\}, \tag {1}
$$
where  $T$  is the maximum length of all observed trajectories and  $T \leqslant H$ . The joint distribution of state and action sequences can be factorized according to the causal assumptions in nMDP:
$$
\begin{array}{l} p _ {\theta} (\zeta) = p \left(s _ {0}\right) p _ {\alpha} \left(a _ {0} \mid s _ {0}\right) p _ {\beta} \left(s _ {1} \mid s _ {0}, a _ {0}\right) \dots p _ {\alpha} \left(a _ {T - 1} \mid s _ {0: T - 1}\right) p _ {\beta} \left(s _ {T} \mid s _ {T - 1}, a _ {T - 1}\right) \\ = p \left(s _ {0}\right) \prod_ {t = 0} ^ {T - 1} p _ {\alpha} \left(a _ {t} \mid s _ {0: t}\right) p _ {\beta} \left(s _ {t + 1} \mid s _ {t}, a _ {t}\right), \tag {2} \\ \end{array}
$$
where  $p_{\alpha}(a_t|s_{0:t - 1})$  is the policy model with parameter  $\alpha$ ,  $p_{\beta}(s_t|s_{t - 1},a_{t - 1})$  is the transition model with parameter  $\beta$ , both of which are parameterized with neural networks,  $\theta = (\alpha ,\beta)$ .  $p(s_0)$  is the initial state distribution, which can be sampled as a black box.
The density families of policy and transition are consistent with the conventional setting of IRL [4], where the transition describes the predictable change in state as a single-mode Gaussian,  $s_{t + 1} \sim \mathcal{N}(g_\beta (s_t,a_t),\sigma^2)$ , and the policy accounts for bounded rationality as a Boltzmann distribution with state-action value as the unnormalized energy:
$$
p _ {\alpha} \left(a _ {t} \mid s _ {0: t}\right) = \frac {1}{Z \left(\alpha , s _ {0 : t}\right)} \exp \left(f _ {\alpha} \left(a _ {t}; s _ {0: t}\right)\right), \tag {3}
$$
where  $f_{\alpha}(a_t; s_{0:t})$  is the negative energy,  $Z(\alpha, s_{0:t}) = \int \exp(f_{\alpha}(a_t; s_{0:t})) da_t$  is the normalizing constant given the contexts  $s_{0:t}$ . We discuss a general push-forward transition in Appx A.3.
Since we can only observe state sequences, the aforementioned generative model can be understood as a sequential variant of LEBM [16], where the transition serves as the generator and the policy is a history-conditioned latent prior. The marginal distribution of state sequences and the posterior distribution of action sequences are:
$$
p _ {\theta} \left(s _ {0: T}\right) = \int p _ {\theta} \left(s _ {0: T}, a _ {0: T - 1}\right) d a _ {0: T - 1}, \quad p _ {\theta} \left(a _ {0: T - 1} \mid s _ {0: T}\right) = \frac {p _ {\theta} \left(s _ {0 : T} , a _ {0 : T - 1}\right)}{p _ {\theta} \left(s _ {0 : T}\right)}. \tag {4}
$$
3.2 Maximum likelihood learning
We need to estimate  $\theta = (\alpha, \beta)$ . Suppose we observe offline training examples:  $\{\xi^i\}, i = 1, 2, \dots, n$ ,  $\xi^i = [s_0^i, s_1^i, \dots, s_T^i]$ . The log-likelihood function is:
$$
L _ {o f f} (\theta) = \sum_ {i = 1} ^ {n} \log p _ {\theta} \left(\xi^ {i}\right). \tag {5}
$$
Denote posterior distribution of action sequence  $p_{\theta}(a_{0:T-1}|s_{0:T})$  as  $p_{\theta}(A|S)$  for convenience where  $A$  and  $S$  means the complete action and state sequences in a trajectory. The full derivation of the learning method can be found in Appx A.2, which results in the following gradient:
$$
\nabla_ {\theta} \log p _ {\theta} (\xi) = \mathbb {E} _ {p _ {\theta} (A | S)} [ \sum_ {t = 0} ^ {T - 1} (\underbrace {\nabla_ {\alpha} \log p _ {\alpha} \left(a _ {t} \mid s _ {0 : t}\right)} _ {\text {p o l i c y / p r i o r}}, \underbrace {\nabla_ {\beta} \log p _ {\beta} \left(s _ {t + 1} \mid s _ {t} , a _ {t}\right)} _ {\text {t r a n s i t i o n}}) ]. \tag {6}
$$
Due to the normalizing constant  $Z(\alpha, s_{0:t})$  in the energy-based prior  $p_{\alpha}$ , the gradient for the policy term involves both posterior and prior samples:
$$
\delta_ {\alpha , t} (S) = \mathbb {E} _ {p _ {\theta} (A | S)} \left[ \nabla_ {\alpha} \log p _ {\alpha} \left(a _ {t} \mid s _ {0: t}\right) \right] = \mathbb {E} _ {p _ {\theta} (A | S)} \left[ \nabla_ {\alpha} f _ {\alpha} \left(a _ {t}; s _ {0: t}\right) \right] - \mathbb {E} _ {p _ {\alpha} \left(a _ {t} \mid s _ {0: t}\right)} \left[ \nabla_ {\alpha} f _ {\alpha} \left(a _ {t}; s _ {0: t}\right) \right], \tag {7}
$$
where  $\delta_{\alpha ,t}(S)$  denotes the expected gradient of policy term for time step  $t$ . Intuition can be gained from the perspective of adversarial training [34, 35]: On one hand, the model utilizes action samples from the posterior  $p_{\theta}(A|S)$  as pseudo-labels to supervise the unnormalized prior at each step. On the other hand, it discourages action samples directly sampled from the prior. The model converges when prior samples and posterior samples are indistinguishable.
To ensure the transition model's validity, it needs to be grounded in real-world dynamics  $Tr$  when jointly learned with the policy. Otherwise, the agent would be purely hallucinating based on the demonstrations. Throughout the training process, we allow the agent to periodically collect on-policy data  $\{(s_t^i,a_t^i,s_{t + 1}^i)\}$ ,  $i = 1,2,\dots ,m$ ,  $t = 1,2,\dots ,T$  with  $p_{\alpha}(a_{t}|s_{0:t})$  and update the transition with a composite likelihood [36]
$$
L _ {c o m p} (\beta) = L _ {o f f} (\theta) + L _ {o n} (\beta), \quad L _ {o n} (\beta) = \sum_ {i = 1} ^ {m} \sum_ {t = 1} ^ {T} \log p _ {\beta} \left(s _ {t + 1} ^ {i} \mid s _ {t} ^ {i}, a _ {t} ^ {i}\right). \tag {8}
$$
3.3 Prior and posterior sampling
The maximum likelihood estimation requires samples from the prior and the posterior distributions of actions. It would not be a problem if the action space is quantized. However, since we target general latent action learning, we proceed to introduce sampling techniques for continuous actions.
When sampling from a continuous energy space, short-run Langevin dynamics [17] can be an efficient choice. For a target distribution  $\pi(a)$ , Langevin dynamics iterates  $a_{k+1} = a_k + s\nabla_{a_k}\log\pi(a_k) + \sqrt{2s}\epsilon_k$ , where  $k$  indexes the number of iteration,  $s$  is a small step size, and  $\epsilon_k$  is the Gaussian white noise.  $\pi(a)$  can be either the prior  $p_\alpha(a_t|s_{0:t})$  or the posterior  $p_\theta(A|S)$ . One property of Langevin dynamics that is particularly amenable for EBM is that we can get rid of the normalizing constant. So for each  $t$  the iterative update for prior samples is
$$
a _ {t, k + 1} = a _ {t, k} + s \nabla_ {a _ {t, k}} f _ {\alpha} \left(a _ {t, k}; s _ {0: t}\right) + \sqrt {2 s} \epsilon_ {k}. \tag {9}
$$
Given a state sequence  $s_{0:T}$  from the demonstrations, the posterior samples at each time step  $a_t$  come from the conditional distribution  $p(a_t | s_{0:T})$ . Notice that with Markov transition, we can derive
$$
p _ {\theta} \left(a _ {0: T - 1} \mid s _ {0: T}\right) = \prod_ {t = 0} ^ {T - 1} p _ {\theta} \left(a _ {t} \mid s _ {0: T}\right) = \prod_ {t = 0} ^ {T - 1} p _ {\theta} \left(a _ {t} \mid s _ {0: t + 1}\right). \tag {10}
$$
Eq. (10) reveals that given the previous and the next subsequent state, the posterior can be sampled at each step independently. So the posterior iterative update is
$$
a _ {t, k + 1} = a _ {t, k} + s \nabla_ {a _ {t, k}} \left(\underbrace {\log p _ {\alpha} \left(a _ {t , k} \mid s _ {0 : t}\right)} _ {\text {p o l i c y / p r i o r}} + \underbrace {\log p _ {\beta} \left(s _ {t + 1} \mid s _ {t} , a _ {t , k}\right)} _ {\text {t r a n s i t i o n}}\right) + \sqrt {2 s} \epsilon_ {k}. \tag {11}
$$
Intuitively, action samples at each step are updated by back-propagation from its prior energy and a single-step lookahead. While gradients from the transition term are analogous to the inverse dynamics in Behavior Cloning from Observations (BCO) [37], it may lead to poor training performance due to non-injectiveness in forward dynamics [38].
We develop an alternative posterior sampling method with importance sampling to overcome this challenge. Leveraging the learned transition, we have
$$
p _ {\theta} \left(a _ {t} \mid s _ {0: t + 1}\right) = \frac {p _ {\beta} \left(s _ {t + 1} \mid s _ {t} , a _ {t}\right)}{\mathbb {E} _ {p _ {\alpha} \left(a _ {t} \mid s _ {0 : t}\right)} \left[ p _ {\beta} \left(s _ {t + 1} \mid s _ {t} , a _ {t}\right) \right]} p _ {\alpha} \left(a _ {t} \mid s _ {0: t}\right). \tag {12}
$$
Let  $c(a_{t}; s_{0:t+1}) = \mathbb{E}_{p_{\alpha}(a_{t}|s_{0:t})} [p_{\beta}(s_{t+1}|s_{t}, a_{t})]$ , posterior sampling from  $p_{\theta}(a_{0:T-1}|s_{0:T})$  can be realized by adjusting importance weights of independent samples from the prior  $p_{\alpha}(a_{t}|s_{0:t})$ , in which the estimation of weights involves another prior sampling. In this way, we avoid back-propagating through non-injective dynamics and save some computation overhead.
To train the policy, Eq. (7) can now be rewritten as
$$
\delta_ {\alpha , t} (S) = \mathbb {E} _ {p _ {\alpha} \left(a _ {t} \mid s _ {0: t}\right)} \left[ \frac {p _ {\beta} \left(s _ {t + 1} \mid s _ {t} , a _ {t}\right)}{c \left(a _ {t} ; s _ {0 : t + 1}\right)} \nabla_ {\alpha} f _ {\alpha} \left(a _ {t}; s _ {0: t}\right) \right] - \mathbb {E} _ {p _ {\alpha} \left(a _ {t} \mid s _ {0: t}\right)} \left[ \nabla_ {\alpha} f _ {\alpha} \left(a _ {t}; s _ {0: t}\right) \right]. \tag {13}
$$
4 Decision-making as Inference
In Section 3, we present our method within the framework of probabilistic inference, providing a self-contained description. However, from a decision-making perspective, the learned policy may appear arbitrary. In this section, we establish a connection between probabilistic inference and decision-making, contributing a novel analysis that incorporates the latent action setting, the non-Markovian assumption, and maximum likelihood learning. This analysis is inspired by, but distinct from, previous studies on the relationship between these two fields [39, 20, 40, 41, 24].
4.1 Policy execution with prior sampling
Let the ground-truth distribution of demonstrations be  $p^*(s_{0:T})$ , and the learned marginal distributions of state sequences be  $p_\theta(s_{0:T})$ . Eq. (5) in Section 3.2 is an empirical estimate of
$$
\mathbb {E} _ {p ^ {*} (s _ {0: T})} \left[ \log p _ {\theta} (s _ {0: T}) \right] = \mathbb {E} _ {p ^ {*} (s _ {0})} \left[ \log p ^ {*} (s _ {0}) + \mathbb {E} _ {p ^ {*} (s _ {1: T} | s _ {0})} \left[ \log p _ {\theta} (s _ {1: T} | s _ {0}) \right] \right]. \tag {14}
$$
We can show that a sequential decision-making problem can be constructed to maximize the same objective. Our main result is summarized as Theorem 1.
Theorem 1. Assuming the Markovian transition  $p_{\beta^*}(s_{t+1}|s_t, a_t)$  is known, the ground-truth conditional state distribution  $p^*(s_{t+1}|s_{0:t})$  for demonstration sequences is accessible, we can construct a sequential decision-making problem, based on a reward function  $r_{\alpha}(s_{t+1}, s_{0:t}) := \log \int p_{\alpha}(a_t|s_{0:t}) p_{\beta^*}(s_{t+1}|s_t, a_t) da_t$  for an arbitrary energy-based policy  $p_{\alpha}(a_t|s_{0:t})$ . Its objective is
$$
\sum_ {t = 0} ^ {T} \mathbb {E} _ {p * (s _ {0: t})} [ V ^ {p _ {\alpha}} (s _ {0: t}) ] = \mathbb {E} _ {p * (s _ {0: T})} \left[ \sum_ {t = 0} ^ {T} \sum_ {k = t} ^ {T} r _ {\alpha} (s _ {k + 1}; s _ {0: k}) \right],
$$
where  $V^{p_{\alpha}}(s_{0:t}) \coloneqq E_{p^{*}(s_{t+1:T}|s_{0:t})}[\sum_{k=t}^{T} r_{\alpha}(s_{k+1}; s_{0:k})]$  is the value function for  $p_{\alpha}$ . This objective yields the same optimal policy as the Maximum Likelihood Estimation  $\mathbb{E}_{p^{*}(s_{0:T})}[\log p_{\theta}(s_{0:T})]$ .
If we further define a reward function  $r_{\alpha}(s_{t+1}, a_t, s_{0:t}) \coloneqq r_{\alpha}(s_{t+1}, s_{0:t}) + \log p_{\alpha}(a_t | s_{0:t})$  to construct a  $Q$  function for  $p_{\alpha}$
$$
Q ^ {p _ {\alpha}} (a _ {t}; s _ {0: t}) := \mathbb {E} _ {p ^ {*} (s _ {t + 1} | s _ {0: t})} \left[ r _ {\alpha} (s _ {t + 1}, a _ {t}, s _ {0: t}) + V ^ {p _ {\alpha}} (s _ {0: t + 1}) \right].
$$
The expected return of  $Q^{p_{\alpha}}(a_t; s_{0:t})$  forms an alternative objective
$$
\mathbb {E} _ {p _ {\alpha} (a _ {t} | s _ {0: t})} [ Q ^ {p _ {\alpha}} (a _ {t}; s _ {0: t}) ] = V ^ {p _ {\alpha}} (s _ {0: t}) - \mathcal {H} _ {\alpha} (a _ {t} | s _ {0: t}) - \sum_ {k = t + 1} ^ {T - 1} \mathbb {E} _ {p ^ {*} (s _ {t + 1: k} | s _ {0: t})} [ \mathcal {H} _ {\alpha} (a _ {k} | s _ {0: k}) ]
$$
that yields the same optimal policy, for which the optimal  $Q^{*}(a_{t};s_{0:t})$  can be the energy function.
Only under certain conditions, this sequential decision-making problem is solvable through non-Markovian extensions of the maximum entropy reinforcement learning algorithms.
Proof. See Appx B.
The constructive proof above offers profound insights. By starting with the hypothesis of latent actions and MLE, and then considering known transition and accessible ground-truth conditional state distribution, we witness the automatic emergence of the entire family of maximum entropy (inverse) RL. This includes prominent algorithms such as soft policy iteration [20], soft Q learning [22] and soft Actor-Critic (SAC) [23]. Among them, SAC is the best-performing off-policy RL algorithm in practice. Unlike the formulation with joint state-action distribution [20, 24], our formulation avoids the peculiarities associated with maximizing state transition entropy. The choice of the maximum entropy policy aligns naturally with the objective of capturing uncertainty in latent actions, and it offers inherent advantages for exploration in model-free learning [42, 22].
4.2 Model-based planning with posterior sampling
Lastly, with the learned model, we can do posterior sampling given any complete or incomplete state sequences. The computation involved is analogous to model-based planning. In Section 3.3, we introduce posterior sampling with short-run MCMC and importance sampling when we have the target next state, which generalizes all cases where the targets of immediate subsequent states are given. Here we introduce the complementary case, where the goal state  $s_T$  is given as the target.
The posterior of actions given the sequential context  $s_{0:t}$  and a target goal state  $s_T$  is
$$
\begin{array}{l} p _ {\theta} \left(a _ {t: T} \mid s _ {0: t}, s _ {T}\right) \propto p _ {\theta} \left(a _ {t: T}, s _ {T} \mid s _ {0: t}\right) \\ = \int \prod_ {k = 0} ^ {T - t - 1} \left[ p _ {\beta} \left(s _ {t + k + 1} \mid a _ {t + k}, s _ {t + k}\right) p _ {\alpha} \left(a _ {t + k} \mid s _ {0: t + k}\right) \right] p _ {\beta} \left(s _ {T} \mid a _ {T - 1}, s _ {T - 1}\right) d s _ {t + 1: T - 1}, \tag {15} \\ \end{array}
$$
in which all Gaussian expectation  $\mathbb{E}_{p_{\beta}}[\cdot ]$  can be approximated with the mean [43]. Therefore,  $a_{t:T}$  can be sampled via short-run MCMC with  $\nabla_{a_{t:T}}\log p_{\theta}(a_{t:T},s_T|s_{0:t})$  back propagated through time. The learned prior can be used to initialize these samples and facilitate the MCMC mixing.
5.1 Cubic curve planning