SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The density families of policy and transition are consistent with the conventional setting of IRL [4]. The transition describes the predictable change in state space, which is often possible to express the random variable  $s_{t+1}$  as a deterministic variable  $s_{t+1} = g_{\beta}(s_t, a_t, \epsilon)$ , where  $\epsilon$  is an auxiliary variable with independent marginal  $p(\epsilon)$ , and  $g_{\beta}(.)$  is some vector-valued function parameterized by  $\beta$ . 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 (\alpha , s _ {0 : t})} \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 history  $s_{0:t}$ .
A.2 Maximum likelihood learning
We need to estimate  $\theta = (\alpha, \beta)$ . Suppose we observe 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 (\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 gradient of log-likelihood is:
$$
\begin{array}{l} \nabla_ {\theta} \log p _ {\theta} (\xi) = \nabla_ {\theta} \log p _ {\theta} (s _ {0}, s _ {1}, \dots , s _ {T}) \\ = \mathbb {E} _ {p _ {\theta} (A | S)} [ \nabla_ {\theta} \log p _ {\theta} (s _ {0}, s _ {1}, \dots , s _ {T}) ] \\ = \mathbb {E} _ {p _ {\theta} (A | S)} [ \nabla_ {\theta} \log p _ {\theta} (s _ {0}, s _ {1}, \dots , s _ {T}) ] + \mathbb {E} _ {p _ {\theta} (A | S)} [ \nabla_ {\theta} \log p _ {\theta} (A | S) ] \\ = \mathbb {E} _ {p _ {\theta} (A | S)} \left[ \nabla_ {\theta} \log p _ {\theta} \left(s _ {0}, a _ {0}, s _ {1}, a _ {1}, \dots , a _ {T - 1}, s _ {T}\right) \right] \\ = \mathbb {E} _ {p _ {\theta} (A | S)} \left[ \nabla_ {\theta} \log p \left(s _ {0}\right) p _ {\alpha} \left(a _ {0} \mid s _ {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) \right] \tag {6} \\ = \mathbb {E} _ {p _ {\theta} (A | S)} \big [ \nabla_ {\theta} \sum_ {t = 0} ^ {T - 1} (\log p _ {\alpha} (a _ {t} | s _ {0: t}) + \log p _ {\beta} (s _ {t + 1} | s _ {t}, a _ {t})) \big ] \\ = \mathbb{E}_{p_{\theta}(A|S)}\bigl[\sum_{t = 0}^{T - 1}\bigl(\underbrace{\nabla_{\alpha}\log p_{\alpha}(a_{t}|s_{0:t})}_{\text{policy/prior}} + \underbrace{\nabla_{\beta}\log p_{\beta}(s_{t + 1}|s_{t},a_{t})}_{\text{transition}}\bigr)\bigr], \\ \end{array}
$$
where the third equation is because of a simple identity  $\mathrm{E}_{\pi_{\theta}(a)}[\nabla_{\theta}\log \pi_{\theta}(a)] = 0$  for any probability distribution  $\pi_{\theta}(a)$ . Applying this simple identity, we also have:
$$
\begin{array}{l} 0 = \mathbb {E} _ {p _ {\alpha} (a _ {t} | s _ {0: t})} [ \nabla_ {\alpha} \log p _ {\alpha} (a _ {t} | s _ {0: t}) ] \\ = \mathbb {E} _ {p _ {\alpha} \left(a _ {t} \mid s _ {0: t}\right)} \left[ \nabla_ {\alpha} f _ {\alpha} \left(a _ {t}; s _ {0: t}\right) - \nabla_ {\alpha} \log Z (\alpha , s _ {0: t}) \right] \tag {7} \\ = \mathbb {E} _ {p _ {\alpha} (a _ {t} | s _ {0: t})} \left[ \nabla_ {\alpha} f _ {\alpha} (a _ {t}; s _ {0: t}) \right] - \nabla_ {\alpha} \log Z (\alpha , s _ {0: t}). \\ \end{array}
$$
$$
\begin{array}{l} \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) - \nabla_ {\alpha} \log Z (\alpha , s _ {0: t}) \right] \tag {8} \\ = \mathbb {E} _ {p _ {\theta} (A | S)} \left[ \nabla_ {\alpha} f _ {\alpha} \left(a _ {t}; s _ {0: t}\right) - \mathbb {E} _ {p _ {\alpha} \left(a _ {t} \mid s _ {0: t}\right)} \left[ \nabla f _ {\alpha} \left(a _ {t}; s _ {0: t}\right) \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], \\ \end{array}
$$
where  $\delta_{\alpha ,t}(S)$  denotes the surrogate loss 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  $p_{\alpha}(a_t|s_{0:t})$ . 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 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 self-interaction data with  $p_{\alpha}(a_t|s_{0:t})$  and mix transition data from two sources with weight  $w_{\beta}$ :
$$
\delta_ {\beta , t} (S) = w _ {\beta} \mathbb {E} _ {p _ {\theta} (A | S)} \left[ \nabla_ {\beta} \log p _ {\beta} \left(s _ {t + 1} \mid s _ {t}, a _ {t}\right) \right] + \left(1 - w _ {\beta}\right) \mathbb {E} _ {p _ {\alpha} \left(a _ {t} \mid s _ {0: t}\right), T r} \left[ \nabla_ {\beta} \log p _ {\beta} \left(s _ {t + 1} \mid s _ {t}, a _ {t}\right) \right]. \tag {9}
$$
A.3 General transition model
We need to compute the gradient of  $\beta$  for the logarithm of transition probability in Equation 9, and as we will see in section 3.3, we also need to compute the gradient of the action during sampling actions. The reparameterization [43] is useful since it can be used to rewrite an expectation w.r.t  $p_{\beta}(s_{t + 1}|s_t,a_t)$  such that the Monte Carlo estimate of the expectation is differentiable, so we use delta function  $\delta (.)$  to rewrite probability as an expectation:
$$
\begin{array}{l} p _ {\beta} \left(s _ {t + 1} \mid s _ {t}, a _ {t}\right) = \int_ {a} \delta \left(s _ {t + 1} - s _ {t + 1} ^ {\prime}\right) p _ {\beta} \left(s _ {t + 1} ^ {\prime} \mid s _ {t}, a _ {t}\right) d s _ {t + 1} ^ {\prime} \tag {10} \\ = \int \delta \left(s _ {t + 1} - g _ {\beta} \left(s _ {t}, a _ {t}, \epsilon\right)\right) p (\epsilon) d \epsilon . \\ \end{array}
$$
Taking advantage of the properties of  $\delta (.)$ :
$$
\int f (x) \delta (x) d x = f (0), \quad \delta (f (x)) = \Sigma_ {n} \frac {1}{| f ^ {\prime} (x _ {n}) |} \delta (x - x _ {n}), \tag {11}
$$
where  $f$  is differentiable and have isolated zeros, which is  $x_{n}$ , we can rewrite the transition probability as:
$$
\begin{array}{l} p _ {\beta} \left(s _ {t + 1} \mid s _ {t}, a _ {t}\right) = \int \sum_ {n} \frac {1}{\left| \frac {\partial}{\partial \epsilon} g _ {\beta} \left(s _ {t} , a _ {t} , \epsilon\right) \right| _ {\epsilon = \epsilon_ {n}}} \delta \left(\epsilon - \epsilon_ {n}\right) p (\epsilon) d \epsilon \tag {12} \\ = \sum_ {n} \frac {p (\epsilon_ {n})}{\left| \frac {\partial}{\partial \epsilon} g _ {\beta} (s _ {t} , a _ {t} , \epsilon) \right| _ {\epsilon = \epsilon_ {n}}}, \\ \end{array}
$$
where  $\epsilon_{n}$  is the zero of  $s_{t + 1} = g_{\beta}(s_t,a_t,\epsilon)$ . Therefore, if we have a differentiable simulator  $\nabla_{a_t}\log p_\beta (s_{t + 1}|s_t,a_t)$  and the analytical form of  $p(\epsilon)$ , then gradient of both  $a_{t}$  and  $\beta$  for  $\log p_{\beta}(s_{t + 1}|s_t,a_t)$  can be computed.
The simplest situation is:
$$
s _ {t + 1} = g _ {\beta} \left(s _ {t}, a _ {t}\right) + \epsilon , \epsilon \sim p (\epsilon) = \mathcal {N} \left(0, \sigma^ {2}\right). \tag {13}
$$
In this case, there is only one zero  $\epsilon^{*}$  for the transition function,  $s_{t + 1} = g_{\beta}(s_t,a_t) + \epsilon^*$ , and the gradient of log probability is:
$$
\begin{array}{l} \nabla \log p _ {\beta} (s _ {t + 1} | s _ {t}, a _ {t}) = \nabla \log \frac {p (\epsilon^ {*})}{| \frac {\partial}{\partial \epsilon} (g _ {\beta} (s _ {t} , a _ {t}) + \epsilon) | _ {\epsilon = \epsilon^ {*}}} \\ = \nabla \log p \left(\epsilon^ {*}\right) \tag {14} \\ = \nabla \log p \left(s _ {t + 1} - g _ {\beta} \left(s _ {t}, a _ {t}\right)\right) \\ = \frac {1}{\sigma^ {2}} \left(s _ {t + 1} - g _ {\beta} \left(s _ {t}, a _ {t}\right)\right) \nabla g _ {\beta} \left(s _ {t}, a _ {t}\right). \\ \end{array}
$$
A.4 Prior and posterior sampling
$$
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 {15}
$$
$$
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 {16}
$$
The point is, given the previous and the next subsequent state, the posterior can be sampled at each step independently. So the posterior iterative update is
$$
\begin{array}{l} a _ {t, k + 1} = a _ {t, k} + s \nabla_ {a _ {t, k}} \log p _ {\theta} \left(a _ {t, k} \mid s _ {0: t + 1}\right) + \sqrt {2 s} \epsilon_ {k} \\ = a _ {t, k} + s \nabla_ {a _ {t, k}} \log p _ {\theta} \left(s _ {0: t}, a _ {t, k}, s _ {t + 1}\right) + \sqrt {2 s} \epsilon_ {k} \tag {17} \\ = a _ {t, k} + s \nabla_ {a _ {t, k}} (\underbrace {\log p _ {\alpha} (a _ {t , k} | s _ {0 : t})} _ {\text {p o l i c y / p r i o r}} + \underbrace {\log p _ {\beta} (s _ {t + 1} | s _ {t} , a _ {t})} _ {\text {t r a n s i t i o n}}) + \sqrt {2 s} \epsilon_ {k}. \\ \end{array}
$$
Intuitively, action samples at each step are updated with the energy of all subsequent actions and a single-step forward by back-propagation. However, while gradients from the transition term are analogous to the inverse dynamics in BCO [37], it may lead to poor training performance due to non-injectiveness in forward dynamics [38].
$$
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 {18}
$$
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 in Eq. (17).
To train the policy, Eq. (8) 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 _ {\theta} \left(a _ {t} \mid s _ {0: t}\right)} \left[ \nabla_ {\alpha} f _ {\alpha} \left(a _ {t}; s _ {0: t}\right) \right]. \tag {19}
$$
A.5 Algorithm
The learning and sampling algorithms with MCMC and with importance sampling for posterior sampling are described in Algorithm 1 and Algorithm 2.
Algorithm 1: LanMDP without importance sampling
Input: Learning iterations  $N$ , learning rate for energy-based policy  $\eta_{\alpha}$ , learning rate for transition model  $\eta_{\beta}$ , initial parameters  $\theta_0 = (\alpha_0, \beta_0)$ , expert demonstrations  $\{s_{0:H}\}$ , context length  $L$ , batch size  $m$ , number of prior and posterior sampling steps  $\{K_0, K_1\}$ , prior and posterior sampling step sizes  $\{s_0, s_1\}$ .
Output:  $\theta_N = (\alpha_N, \beta_N)$ .
Reorganize  $\{s_{0:H}\}$  to state sequenc segments  $(s_{t-L+1}, \dots, s_{t+1})$  with length  $L+1$ .
Use energy-based policy with  $\alpha_0$  collect transitions to fill in the replay buffer.
Use transitions in replay buffer to pre-train transition model  $\beta_0$ .
**Demo sampling** Sample observed examples  $(s_{t - L + 1},\dots ,s_{t + 1})_{i = 1}^{m}$ .
Posterior sampling: Sample  $\{a_{t}\}_{i = 1}^{m}$  using Eq. (17) with  $K_{1}$  iterations and stepsize  $s_1$ .
Prior sampling: Sample  $\{\hat{a}_t\}_{i = 1}^m$  using Eq. (15) with  $K_{0}$  iterations and stepsize  $s_0$ .
Policy learning: Update  $\alpha_{t}$  to  $\alpha_{t + 1}$  by Eq. (8) with learning rate  $\eta_{\alpha}$ .
Transition learning: Update replay buffer with trajectories from current policy model  $\alpha_{t + 1}$  then update  $\beta_{t}$  to  $\beta_{t + 1}$  by Eq. (9) with learning rate  $\eta_{\beta}$ .
Algorithm 2: LanMDP with importance sampling
Input: Learning iterations  $N$ , learning rate for energy-based policy  $\eta_{\alpha}$ , learning rate for transition model  $\eta_{\beta}$ , initial parameters  $\theta_0 = (\alpha_0, \beta_0)$ , expert demonstrations  $\{s_{0:H}\}$ , context length  $L$ , batch size  $m$ , number of prior sampling steps  $K$  and step sizes  $s$ .
**Demo sampling** Sample observed examples  $(s_{t - L + 1},\dots ,s_{t + 1})_{i = 1}^{m}$
Policy learning: Update  $\alpha_{t}$  to  $\alpha_{t + 1}$  by Eq. (19) with learning rate  $\eta_{\alpha}$ .
Transition learning: Update replay buffer with trajectories from current policy model  $\alpha_{t + 1}$
then update  $\beta_{t}$  to  $\beta_{t + 1}$  by Eq. (9) with learning rate  $\eta_{\beta}$ .
B A Decision-making Problem in MLE
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) is an empirical estimate of
$$
\mathbb {E} _ {p ^ {*} (s _ {0: T})} \left[ \log p _ {\theta} \left(s _ {0: T}\right) \right] = \mathbb {E} _ {p ^ {*} (s _ {0})} \left[ \log p ^ {*} (s _ {0}) + \mathbb {E} _ {p ^ {*} (s _ {1: T} | s _ {0})} \left[ \log p _ {\theta} \left(s _ {1: T} \mid s _ {0}\right) \right] \right]. \tag {20}
$$
We can show that a sequential decision-making problem can be constructed to maximize the same objective. To start off, suppose the MLE yields the maximum, we will have  $p_{\theta *} = p^*$ .
Define  $V^{*}(s_{0}) \coloneqq \mathbb{E}_{p^{*}(s_{1:T}|s_{0})}[\log p^{*}(s_{1:T}|s_{0})]$ , we can generalize it to have a  $V$  function
$$
V ^ {*} (s _ {0: t}) := \mathbb {E} _ {p ^ {*} (s _ {t + 1: T} | s _ {0: t})} \left[ \log p ^ {*} (s _ {t + 1: T} | s _ {0: t}) \right], \tag {21}
$$
which comes with a Bellman optimality equation:
$$
V ^ {*} \left(s _ {0: t}\right) = \mathbb {E} _ {p ^ {*} \left(s _ {t + 1} \mid s _ {0: t}\right)} \left[ r \left(s _ {t + 1}, s _ {0: t}\right) + V ^ {*} \left(s _ {0: t + 1}\right) \right], \tag {22}
$$
with  $r(s_{t+1}, s_{0:t}) \coloneqq \log p^*(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$ ,  $V^*(s_{0:T}) \coloneqq 0$ . It is worth noting that the  $r$  defined above involves the optimal policy, which may not be known a priori. We can resolve this by replacing it with  $r_\alpha$  for an arbitrary policy  $p_\alpha(a_t|s_{0:t})$ . All Bellman identities and updates should still hold. Anyways, involving the current policy in the reward function should not appear to be too odd given the popularity of maximum entropy RL [20, 24].
The entailed Bellman update, value iteration, for arbitrary  $V$  and  $\alpha$  is
$$
V \left(s _ {0: t}\right) = \mathbb {E} _ {p ^ {*} \left(s _ {t + 1} \mid s _ {0: t}\right)} \left[ r _ {\alpha} \left(s _ {0: t}, s _ {t + 1}\right) + V \left(s _ {0: t + 1}\right) \right]. \tag {23}
$$
We then define  $r(s_{t + 1},a_t,s_{0:t})\coloneqq r(s_{t + 1},s_{0:t}) + \log p_{\alpha}* (a_t|s_{0:t})$  to construct a  $Q$  function:
$$
Q ^ {*} \left(a _ {t}; s _ {0: t}\right) := \mathbb {E} _ {p ^ {*} \left(s _ {t + 1} \mid s _ {0: t}\right)} \left[ r \left(s _ {t + 1}, a _ {t}, s _ {0: t}\right) + V ^ {*} \left(s _ {0: t + 1}\right) \right], \tag {24}
$$
which entails a Bellman update,  $Q$  backup, for arbitrary  $\alpha$ ,  $Q$  and  $V$
$$
Q \left(a _ {t}; s _ {0: t}\right) = \mathbb {E} _ {p * \left(s _ {t + 1} \mid s _ {0: t}\right)} \left[ r _ {\alpha} \left(s _ {0: t}, a _ {t}, s _ {t + 1}\right) + V \left(s _ {0: t + 1}\right) \right]. \tag {25}
$$
Also note that the  $V$  and  $Q$  in identities Eq. (23) and Eq. (25) respectively are not necessarily associated with the policy  $p_{\alpha}(a_t|s_{0:t})$ . Slightly overloading the notations, we use  $Q^{\alpha}$ ,  $V^{\alpha}$  to denote the expected returns from policy  $p_{\alpha}(a_t|s_{0:t})$ .
By now, we finish the construction of atomic algebraic components and move on to check if the relations between them align with the algebraic structure of a sequential decision-making problem [9].
We first prove the construction above is valid at optimality.
Lemma 1. When  $f_{\alpha}(a_t; s_{0:t}) = Q^*(a_t; s_{0:t}) - V^*(s_{0:t})$ ,  $p_{\alpha}(a_t | s_{0:t})$  is the optimal policy. Proof. Note that the construction gives us
$$
\begin{array}{l} Q ^ {*} (a _ {t}; s _ {0: t}) = \mathbb {E} _ {p ^ {*} (s _ {t + 1} | s _ {0: t})} [ r (s _ {t + 1}, s _ {0: t}) + \log p _ {\alpha} * (a _ {t} | s _ {0: t}) + V ^ {*} (s _ {0: t + 1}) ] \\ = \log p _ {\alpha} * \left(a _ {t} \mid s _ {0: t}\right) + \mathbb {E} _ {p ^ {*} \left(s _ {t + 1} \mid s _ {0: t}\right)} \left[ r \left(s _ {t + 1}, s _ {0: t}\right) + V ^ {*} \left(s _ {0: t + 1}\right) \right] \tag {26} \\ = \log p _ {\alpha} * (a _ {t} | s _ {0: t}) + V ^ {*} (s _ {0: t}). \\ \end{array}
$$
Obviously,  $Q^{*}(a_{t}; s_{0:t})$  lies in the hypothesis space of  $f_{\alpha}(a_{t}; s_{0:t})$ .
Lemma 1 indicates that we need to either parametrize  $f_{\alpha}(a_t; s_{0:t})$  or  $Q(a_t; s_{0:t})$ .
While  $Q^{\alpha}$  and  $V^{\alpha}$  are constructed from the optimality, the derived  $Q^{\alpha}$  and  $V^{\alpha}$  measure the performance of an interactive agent when it executes with the policy  $p_{\alpha}(a_t|s_{0:t})$ . They should be consistent with each other.
Lemma 2.  $V^{\alpha}(s_{0:t})$  and  $\mathbb{E}_{p_{\alpha}(a_t|s_{0:t})}[Q^{\alpha}(a_t;s_{0:t})]$  yield the same optimal policy  $p_{\alpha *} (a_t | s_{0:t})$ . Proof.
$$
\begin{array}{l} \mathbb {E} _ {p _ {\alpha} (a _ {t} | s _ {0: t})} [ Q ^ {\alpha} (a _ {t}; s _ {0: t}) ] := \mathbb {E} _ {p _ {\alpha} (a _ {t} | s _ {0: t})} \left[ \mathbb {E} _ {p ^ {*} (s _ {t + 1} | s _ {0: t})} \left[ r (s _ {t + 1}, a _ {t}, s _ {0: t}) + V ^ {\alpha} (s _ {0: t + 1}) \right] \right] \\ = \mathbb {E} _ {p _ {\alpha} (a _ {t} | s _ {0: t})} \left[ \mathbb {E} _ {p ^ {*} (s _ {t + 1} | s _ {0: t})} \left[ \log p _ {\alpha} (a _ {t} | s _ {0: t}) + r (s _ {t + 1}, s _ {0: t}) + V ^ {\alpha} (s _ {0: t + 1}) \right] \right] \\ = \mathbb {E} _ {p * (s _ {t + 1} | s _ {0: t})} [ r (s _ {t + 1}, s _ {0: t}) - \mathcal {H} _ {\alpha} (a _ {t} | s _ {0: t}) + V ^ {\alpha} (s _ {0: t + 1}) ] \tag {27} \\ = V ^ {\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}) ], \\ \end{array}
$$
where the last line is derived by recursively applying the Bellman equation in the line above until  $s_{0:T}$  and then applying backup with Eq. (23). As an energy-based policy,  $p_{\alpha}(a_t|s_{0:t})$ 's entropy is inherently maximized [66]. Therefore, within the hypothesis space,  $p_{\alpha*}(a_t|s_{0:t})$  that optimizes  $V^{\alpha}(s_{0:t})$  also leads to optimal expected return  $\mathbb{E}_{p_{\alpha}(a_t|s_{0:t})}[Q^{\alpha}(a_t;s_{0:t})]$ .
If we parametrize the policy as  $p_{\alpha}(a_t|s_{0:t}) \propto \exp(Q^{\alpha}(a_t; s_{0:t}))$ , the logarithmic normalizing constant  $\log Z^{\alpha_k}(s_{0:t})$  will be the soft V function in maximum entropy RL [21-23]
$$
V _ {s o f t} ^ {\alpha} \left(s _ {0: t}\right) := \log \int \exp \left(Q ^ {\alpha} \left(a _ {t}; s _ {0: t}\right)\right) d a _ {t}, \tag {28}
$$
even if the reward function is defined differently. We can further show that Bellman identities and backup updates above can entail RL algorithms that achieve optimality of the decision-making objective  $V^{\alpha}$ , including soft policy iteration [20]
$$
p _ {\alpha_ {k + 1}} \left(a _ {t} \mid s _ {0: t}\right) \leftarrow \frac {\exp \left(Q ^ {\alpha_ {k}} \left(a _ {t} ; s _ {0 : t}\right)\right)}{Z ^ {\alpha_ {k}} \left(s _ {0 : t}\right)}, \forall s _ {0: t}, k \in [ 0, 1, \dots M ]; \tag {29}
$$
and soft  $Q$  iteration [21]
$$
\begin{array}{l} Q ^ {\alpha_ {k + 1}} \left(a _ {t}; s _ {0: t}\right) \leftarrow \mathbb {E} _ {p * \left(s _ {t + 1} \mid s _ {0: t}\right)} \left[ r _ {\alpha} \left(s _ {0: t}, a _ {t}, s _ {t + 1}\right) + V _ {s o f t} ^ {\alpha_ {k}} \left(s _ {0: t + 1}\right) \right], \forall s _ {0: t}, a _ {t}, \\ V _ {\text {s o f t}} ^ {\alpha_ {k + 1}} \left(s _ {0: t}\right) \leftarrow \log \int \exp \left(Q ^ {\alpha_ {k}} \left(a; s _ {0: t}\right)\right) d a, \forall s _ {0: t}, k \in [ 0, 1, \dots M ]. \tag {30} \\ \end{array}
$$
Lemma 3. If  $p^*(s_{t+1}|s_{0:t})$  is accessible and  $p_{\beta^*}(s_{t+1}|s_t, a_t)$  is known, soft policy iteration and soft  $Q$  learning both converge to  $p_{\alpha^*}(a_t|s_{0:t}) = p_{Q^*}(a_t|s_{0:t}) \propto \exp(Q^*(a_t; s_{0:t}))$  under certain conditions.
Proof. See the convergence proof by Ziebart [20] for soft policy iteration and the proof by Fox et al. [21] for soft  $Q$  learning. The latter requires Markovian assumption. But under some conditions, it can be extended to non-Markovian domains in the same way as proposed by Majeed and Hutter [67].
Lemma 3 means given  $p^*(s_{t+1}|s_{0:t})$  and  $p_{\beta^*}(s_{t+1}|s_t, a_t)$ , we can recover  $p_{\alpha^*}$  through reinforcement learning methods, instead of the proposed MLE. So  $p_{\alpha}(a_t|s_{0:t})$  is a viable policy space for the constructed sequential decision-making problem.
Together, Lemma 1, Lemma 2 and Lemma 3 provide constructive proof for a valid sequential decision-making problem that maximizes the same objective of MLE, described by Theorem 1.
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})]$ .
$$
Q ^ {p _ {\alpha}} \left(a _ {t}; s _ {0: t}\right) := \mathbb {E} _ {p ^ {*} \left(s _ {t + 1} \mid s _ {0: t}\right)} \left[ r _ {\alpha} \left(s _ {t + 1}, a _ {t}, s _ {0: t}\right) + V ^ {p _ {\alpha}} \left(s _ {0: t + 1}\right) \right].
$$
C More results on Curve Planning
The energy function is parameterized by a small MLP with one hidden layer and  $4 * L$  hidden neurons, where  $L$  is the context length. In short-run Langevin dynamics, the number of samples, the number of sampling steps, and the stepsize are 4, 20 and 1 respectively. We use Adam optimizer with a learning rate 1e-4 and batch size 64. Here we present the complete result in Fig. A1 with different training steps under context length 1 2 4 6, the acceptance rate and residual error of the testing trajectories, as well as the behavior cloning results. We can see that even with sufficient context, BC performs worse than LanMDP. Also, from the result of context length 6 we can see that excessive expressivity does not impair performance, it just requires more training.
Figure A1: More results for cubic curve generation
D Implementation Details of MuJoCo Environment