Jump Operator Planning: Goal-Conditioned
Policy Ensembles and Zero-Shot Transfer

We introduce a novel extension of the LMDP, that operates on state-action pairs. This extension will be employed in two places throughout the algorithm, for both the low-level goal-conditioned policy ensemble, and the high-level tLMDP task policy that “stitches” policies from the ensemble together. Reintroducing the action variable will have non-trivial consequences for the tLMDP, as we can construct more expressive cost functions including ordering constraints on sub-goals, and the cost of following a given policy to complete a sub-goal, all while retaining a largest eigenvector solution.

We formulate a new objective function by defining the state in saLMDP as a state-action pair ${y}=(x,a)$, and then substitute ${y}$ for $x$ in (1). After substitution, we can decompose the resulting joint distributions $u_{xa}(x^{\prime},a^{\prime}|x,a)$ and $p_{xa}(x^{\prime},a^{\prime}|x,a)$ into $u_{xa}:=u_{x}(x^{\prime}|x,a)u_{a}(a^{\prime}|x^{\prime},x,a)$ and $p_{xa}:=p_{x}(x^{\prime}|x,a)p_{a}(a^{\prime}|x^{\prime},x,a)$ via the chain rule, where $u_{x}:\mathcal{X}\times\mathcal{A}\times\mathcal{X}\rightarrow[0,1]$ is a state-transition policy, and $u_{a}:\mathcal{X}\times\mathcal{A}\times\mathcal{X}\times\mathcal{A}\rightarrow[0,1]$ is an action transition policy. We enforce an optimization constraint that $u_{x}(x^{\prime}|x,a)=p_{x}(x^{\prime}|x,a)$ so the controller is not free to optimize $u_{x}$, which would modify a fixed physics model encoded in $p_{x}$. Thus, we only allow the controller to modify the distribution over actions, $u_{a}$, giving rise to the saLMDP objective and solution:

	$$\displaystyle v(x,a)=\min\limits_{\begin{subarray}{c}u_{a}\end{subarray}}\Bigg{[}q(x,a)+\mathop{\mathbb{E}}\limits_{\begin{subarray}{c}a^{\prime}\sim u_{a}\end{subarray}}\Bigg{[}log\Bigg{(}\frac{u_{a}(a^{\prime}|x^{\prime},x,a)}{p_{a}(a^{\prime}|x^{\prime},x,a)}\Bigg{)}\Bigg{]}+\mathop{\mathbb{E}}\limits_{\begin{subarray}{c}a^{\prime}\sim u_{a}\\ x^{\prime}\sim u_{x}=p_{x}\end{subarray}}[v(x^{\prime},a^{\prime})]\Bigg{]},$$		(2)
	$$\displaystyle u_{a}^{*}(a^{\prime}|x^{\prime},x,a)=\frac{p_{a}(a^{\prime}|x^{\prime},x,a)z(x^{\prime},a^{\prime})}{G[z](x^{\prime},x,a)},$$		(3)
	$$\displaystyle\mathbf{z}=QP_{xa}\mathbf{z},$$		(4)

where $G[z](x^{\prime},x,a)=\mathop{\mathbb{E}}_{a^{\prime}\sim p_{a}}z(x^{\prime},a^{\prime})$ is normalization term and $\mathbf{z}$ is the desirability vector on $\mathcal{X}\times\mathcal{A}$, i.e. for each $y=(x,a)$ we have $z(x,a)=exp(-v(x,a))$, and $u_{a}^{*}$ is the optimal policy–we will refer to $u^{*}_{a}$ as $u^{*}$ for compactness¹¹1$v(x,a)$ is the same as the ”Q-function”, $Q(x,a)$, in RL [28], however, since we mostly work with $z(x,a)$, we keep consistent with the original LMDP notation in which $q$ is a cost-function and $Q$ is a cost-matrix.. When $p_{x}$ is a deterministic distribution, same as the LMDP, the optimal policy $u^{*}$ can be solved for by computing the optimal desirability function $z(x,a)$ which is the largest eigenvector of $QP_{xa}$ (see Appx. 5.2 for derivation), where $Q=diag(exp(-\mathbf{q}))$ and $P_{xa}$ is the matrix form of the joint passive dynamics $p_{xa}(x^{\prime},a^{\prime}|x,a)$. In the case where $p_{x}$ is stochastic, the solution does not reduce down to a linear equation, but there exists an iterative non-linear mapping for computing the desirability function (25). Henceforth, the following theory we develop will be linear, only considering $p_{x}$ as deterministic.

Note that as the cost of a state-action in $\mathbf{q}$ approaches infinity, its negatively exponentiated counterpart, represented as a diagonal entry in $Q$, approaches zero. Thus, zero-entries on the diagonal effectively edit the transition structure of $QP_{xa}$, shaping the backwards propagation of desirability when computing (4) with power-iteration, suppressing state-actions which would violate the infinite-cost transitions. Todorov [7] showed that $\mathbf{z}$ could approximate shortest-paths with arbitrary precision in the limit of driving the internal state costs arbitrarily high towards infinity, i.e. $S(x)=\lim_{c\rightarrow\infty}\frac{v_{c}(x)}{c}$, where $S(x)$ is the shortest-path distance from $x$ to the goal state and $c$ is the internal state cost. In this paper, obstacles are represented with infinite cost, internal states have an arbitrarily high constant cost, and the boundary (goal) state-action has zero cost. Therefore, the saLMDP objective with this cost structure is a discrete reach-avoid shortest-path first-exit control problem, avoiding obstacles states while controlling to a boundary state-action. These polices are said to be goal-conditioned, as we can explicitly interpret the policy by the goal’s grounded state-action, rendering the polices remappable for new tasks.

We define an ordered goal task, which specifies a set of necessary sub-goals to be completed under precedence constraints.

Here, $\mathcal{O}_{\Psi}$ is a set of precedence relations, $o(\cdot,\cdot)$, on types, e.g. $o(\psi_{1},\psi_{2})\in\mathcal{O}_{\Psi}\implies\psi_{1}\prec\psi_{2},\leavevmode\nobreak\ \psi_{1},\psi_{2}\in\Psi$, and $\Sigma$ is a binary vector state-space which tracks the progress of goal achievement. For example: $\sigma=[0,1,0,0,1]\in\Sigma$ indicates that $\text{g}_{2},\textsl{g}_{5}\in\mathcal{G}$ have been completed and three other goals $\textsl{g}_{1},\textsl{g}_{3},\textsl{g}_{4}\in\mathcal{G}$ remain to be completed. With these primary objects, we can derive a transition operator $T_{\sigma}(\sigma^{\prime}|\sigma,\textsl{g})$ which specifies the dynamics of task-space under given sub-goals, with the restriction that $\textsl{g}_{i}$ can only flip the $i^{th}$ bit in the vector, $\sigma(i)$. The task space transition operator $T_{\sigma}$ encodes the task-space cube in the example of Fig. 1. Likewise, a cost function on this space $q_{\sigma}(\sigma,\textsl{g})$ can be derived directly from the ordering constraints. That is, if we would like to specify that a sub-goal of one type precedes a sub-goal of another type, e.g. $\psi_{1}\prec\psi_{2}$, then these constraints are imposed directly on to the cost function. We can define an induced precedence relation set on goals

$$\mathcal{O}_{\mathcal{G}}:=\{o(\textsl{g}_{i},\textsl{g}_{j})\leavevmode\nobreak\ |\leavevmode\nobreak\ o(\psi_{1},\psi_{2})\in\mathcal{O}_{\Psi},\leavevmode\nobreak\ \exists(\psi_{1},\psi_{2})\in f_{\Psi}(\textsl{g}_{i})\times f_{\Psi}(\textsl{g}_{j})\}.$$

Type relations add an additional level of symbolic reasoning into the architecture, enabling us to make statements such as: "complete the purple sub-goals before the orange sub-goals", or "chop wood with an axe in addition to collecting water", as depicted in Fig. 2.A and 5.A by red arrows. The cost can then be defined as $q=q_{\sigma\textsl{g}}+q_{\sigma}$:

$$q_{\sigma\textsl{g}}(\sigma,\textsl{g}_{i})=\{\infty\leavevmode\nobreak\ \leavevmode\nobreak\ \mathbf{if}:\leavevmode\nobreak\ \exists\textsl{g}_{j}\in\mathcal{G}\leavevmode\nobreak\ s.t.\leavevmode\nobreak\ o(\textsl{g}_{i},\textsl{g}_{j})\in\mathcal{O}_{\mathcal{G}}\land\leavevmode\nobreak\ \sigma(j)=1;\leavevmode\nobreak\ \mathbf{else}:\leavevmode\nobreak\ 0.\},$$

$$q_{\sigma}(\sigma)=\{0\leavevmode\nobreak\ \leavevmode\nobreak\ \mathbf{if}:\leavevmode\nobreak\ \sigma=\sigma_{f};\leavevmode\nobreak\ \mathbf{else}:\leavevmode\nobreak\ const.\}$$

where $q_{\sigma\textsl{g}}$ penalizes ordering violations on the bit-flips of $\sigma$, and $q_{\sigma}$ is a constant cost on task states with the exception of the task-space boundary state $\sigma_{f}=[1,1,...]$.

Since the objects of an OG-task $\mathscr{T}$ can be directly used to define a transition operator and a cost function, much like Definition 3.1, we can introduce a passive transition dynamics $p_{\textsl{g}}(\textsl{g}^{\prime}|\sigma^{\prime},\sigma,\textsl{g})$ and solve for a policy in the task space, $u_{\textsl{g}}^{*}(\textsl{g}^{\prime}|\sigma^{\prime},\sigma,\textsl{g})$, which satisfies our constraints, as illustrated in Fig. 2.A. However, synthesizing a policy only in the task space is not sufficiently useful for an agent. To make task-states and actions useful and reusable representations, we show how to map goal-actions to the agent’s underlying controllable space, in this case, the base space $\mathcal{X}\times\mathcal{A}$.

In order handle this mapping, we introduce specialized functions called grounding and ungrounding functions which mediate the interaction between high-level goals and the agent’s low-level dynamics. A grounding function, $h:\mathcal{G}\rightarrow\mathcal{X}\times\mathcal{A}$ is a function which associates a goal variable with a state-action pair in the base space. An ungrounding function $\bar{h}:\mathcal{X}\times\mathcal{A}\rightarrow\mathcal{G}$ is a left-inverse of $h$ which maps state-actions back to the goal variable. For grid-world examples in Fig. 2, the base action-space $\mathcal{A}=\{a_{u},a_{d},a_{r},a_{l},a_{n},a_{g}\}$ includes an action $a_{g}$, which is a special neutral action for completing a sub-goal in addition to five standard directional actions: up, down, right, left and neutral. Consider a saLMDP transition operator $p_{xa}(x^{\prime},a^{\prime}|x,a)$. If we compose the Markov chain with the ungrounding function we obtain

$$p_{\textsl{g}xa}(\textsl{g}^{\prime},x^{\prime},a^{\prime}|\textsl{g},x,a):=\bar{h}(\textsl{g}^{\prime}|x^{\prime},a^{\prime})p_{xa}(x^{\prime},a^{\prime}|x,a).$$

Since we have specified transition dynamics $T_{\sigma}$ over $\Sigma$ in Section 3.2, we can also obtain the following joint dynamics

$$p_{\sigma\textsl{g}xa}(\sigma^{\prime},\textsl{g}^{\prime},x^{\prime},a^{\prime}|\sigma,\textsl{g},x,a):=T_{\sigma}(\sigma^{\prime}|\sigma,\textsl{g}^{\prime})p_{\textsl{g}xa}(\textsl{g}^{\prime},x^{\prime},a^{\prime}|\textsl{g},x,a).$$

We call $p_{\sigma\textsl{g}xa}$ the full dynamics. Note that this transition operator is specified over a space of size $|\Sigma\times\mathcal{G}\times\mathcal{X}\times\mathcal{A}|$. Since $\Sigma$ grows exponentially in the number of sub-goals, i.e. $2^{|\mathcal{G}|}$, this operator over the full product space can introduce considerable space and time complexity into the computation of the optimal policy. However, the added complexity can be mitigated by creating specialized transition operators called jump operators which summarize the resulting dynamics of the a policy from its initial state to the state-action of sub-goal completion. In doing so, we can avoid representing internal (non-subgoal) states of $\mathcal{X}$ and compute a solution in a lower dimensional subspace of the full space, to be discussed in 3.6.

Let $\textsl{g}\in\mathcal{G}$ be a goal in a given OG-task $\mathscr{T}$. Recall the joint distribution $u_{xa}$ over the base space $\mathcal{X}\times\mathcal{A}$ in Section 3.1 that expresses the control policy. Once a control policy is optimized with a saLMDP, a Markov chain is induced. Consider an optimal shortest-path first-exit Markov chain matrix $U_{g}$, induced by applying the distribution ${u_{xa}}_{g}^{*}(x^{\prime},a^{\prime}|x,a)={u^{*}_{a}}_{g}(a^{\prime}|x^{\prime};x,a)p(x^{\prime}|x,a)$ that controls to the goal g’s state-action pair $(x,a)_{g}$²²2We use stylized ‘g’ to refer to goal variables and ‘$g$’ to refer to the goal index on a state-action (Fig. 2.A). By slight notation abuse we might use $(x_{g},a_{g})$ instead of $(x,a)_{g}$. With this Markov chain, we can construct a matrix

$$C=\begin{bmatrix}U_{\bar{g}}&H_{g}\\ 0&I\end{bmatrix},$$

which segregates the transient state dynamics $U_{\bar{g}}$ from the absorbing state dynamics $H_{g}$. $U_{\bar{g}}$ is a sub-matrix of $U_{g}$ which excludes the $g^{th}$ row and column, and $H_{g}$ is a rank-one matrix of zeros except for the $g^{th}$ column, which is taken from $g^{th}$ column of $U_{g}$, excluding the $g^{th}$ element. By taking an arbitrary power, $C^{\tau}$, the elements represent the probability of remaining in the transient states or absorbing into the goal state-action after $\tau$ time steps. The power of $C$ in the infinite limit can be written as [29]:

$$\lim_{\tau\rightarrow\infty}C^{\tau}=\begin{bmatrix}\lim_{\tau\rightarrow\infty}U_{\bar{g}}^{\tau}&(\sum_{\tau=0}^{\infty}U_{\bar{g}}^{\tau})H_{g}\\ 0&I\end{bmatrix}=\begin{bmatrix}\lim_{\tau\rightarrow\infty}U_{\bar{g}}^{\tau}&(I-U_{\bar{g}})^{-1}H_{g}\\ 0&I\end{bmatrix},$$

(5)

where $(I-U_{\bar{g}})^{-1}$ is the fundamental matrix of the absorbing chain³³3Invertibility is generally not a concern if U is derived from a first-exit problem. See 5.6 for discussion., which is the solution to the Neumann series $\sum_{\tau=0}^{\infty}U_{\bar{g}}^{\tau}$. Note that the upper right sub-matrix $(I-U_{\bar{g}})^{-1}H_{g}$ is a rank-one matrix, which represents the probability of starting in state-action $(x,a)_{i}$, indexing the row space, and absorbing into state-action $(x,a)_{j}$, indexing the column space. There can only be positive probability in $(I-U_{\bar{g}})^{-1}H_{g}$ in the case that $j=g$, corresponding to the single defined absorbing goal state-action $(x_{g},a_{g})$. Let $\mathscr{M}^{N}$, be a set of $N$ first-exit saLMDPs, one for every grounded state-action. Consider a set of policies for each saLMDP, $\mathcal{U}_{N}=\{{u^{*}_{xa}}_{g_{1}},{u^{*}_{xa}}_{g_{2}},...\}$ and an associated set of policy indices $\Pi=\{\pi_{\textsl{g}_{1}},\pi_{\textsl{g}_{2}},...\}$, which indexes corresponding members of $\mathcal{U}_{N}$ to avoid overlapping notation for high- and low-level polices. Specifically, $\pi_{\textsl{g}}$ points to a low-level policy: $\pi_{\textsl{g}}\xrightarrow{}{u^{*}_{xa}}_{g}$. We can define a new operator using the upper-right sub-matrix for any $g$ and its associated ${u^{*}_{xa}}_{g}$ as:

$$\displaystyle J((x,a)_{j}|(x,a)_{i},\pi_{\textsl{g}})=e_{i}^{T}(I-U_{\bar{g}})^{-1}H_{g}e_{j},$$

(6)

where $e_{i}$ and $e_{j}$ are one-hot unit vectors corresponding to $(x,a)_{i}$ and $(x,a)_{j}$. The probability of starting and ending at the boundary state-action $(x,a)_{g}$ is defined to be 1. $J$ can also be written as $J((x,a)_{j}|(x,a)_{i},u_{xa_{g}})$, given that $\pi_{\textsl{g}}$ is simply an index variable–this will be important when we remap low-level solutions to new problems. The new operator $J$ is a jump operator which maps the agent from an initial state-action to the goal state-action with the probability that it will eventually reach it. We now show how $J$ can be combined with $\bar{h}$ to construct joint-state jump dynamics.

Recall that in Section 3.3 we defined a new passive dynamics $p_{\textsl{g}xa}(\textsl{g}^{\prime},x^{\prime},a^{\prime}|\textsl{g},x,a)$ by composing $p_{xa}$ with $\bar{h}$. Likewise, now that we have defined the jump operator we can define a new representation called a feasibility function, $\kappa$, which composes the jump operator with the ungrounding function:

	$$\displaystyle\kappa(\textsl{g}^{\prime},x^{\prime},a^{\prime}|x,a,\pi_{\textsl{g}})=\bar{h}(\textsl{g}^{\prime}|x^{\prime},a^{\prime})J(x^{\prime},a^{\prime}|x,a,\pi_{\textsl{g}}),$$		(7)
	$$\displaystyle T_{\kappa}(\sigma^{\prime},\textsl{g}^{\prime},x^{\prime},a^{\prime}|\sigma,x,a,\pi_{\textsl{g}})=T_{\sigma}(\sigma^{\prime}|\sigma,\textsl{g}^{\prime})\kappa(\textsl{g}^{\prime},x^{\prime},a^{\prime}|x,a,\pi_{\textsl{g}}),$$		(8)

The feasibility function $\kappa:(\mathcal{X}\times\mathcal{A}\times\Pi)\times(\mathcal{G}\times\mathcal{X}\times\mathcal{A})\rightarrow[0,1]$, originally conceived of in [12] in a non-stationary context, is a function in which the transition of the action variable (goal) of the higher level space is parameterized in terms of the policy on the lower level space. In (8), we use the feasibility function with the high-level operator $T_{\sigma}$ to create coupled dynamics, $T_{\kappa}$, summarizing the initial-to-final dynamics over both spaces.

With the definition of the new joint dynamics, $T_{\kappa}$, we are in a position to repurpose the saLMDP to optimize a task-policy over the joint space, stitching together polices from the ensemble to complete the task, an algorithm called the task-LMDP (tLMDP).

From these three components, we can derive a desirability ensemble $\mathcal{Z}_{N}=\{z_{1},z_{2},...\}$, policy ensemble $\mathcal{U}_{N}=\{u_{1},u_{2},...\}$, jump operator $J$, ungrounding function $\bar{h}$, and feasibility function, $\kappa$, in order to construct the coupled dynamics $T_{\kappa}$ in (8). $N$ can either be $N_{All}$, for a complete ensemble of all non-obstacle states, or $N_{GS}$, the number of goal states defined by the grounding function. The choice depends on whether or not the solution is intended for task-transfer, where policies in the complete ensemble can be remapped to different problems, as we will soon explain. By defining a uniform passive and controlled dynamics distribution $p_{\pi_{\textsl{g}}},u_{\pi_{\textsl{g}}}$, we arrive at the tLMDP objective:

$\displaystyle v(s,\pi_{\textsl{g}})=\min\limits_{\begin{subarray}{c}u_{\pi_{\textsl{g}}}\end{subarray}}\Bigg{[}q(s,\pi_{\textsl{g}})$

$\displaystyle+\mathop{\mathbb{E}}\limits_{\begin{subarray}{c}\pi_{\textsl{g}}\sim u_{\pi_{\textsl{g}}}\end{subarray}}\Bigg{[}log\Bigg{(}\frac{u_{\pi_{\textsl{g}}}(\pi_{\textsl{g}}^{\prime}|s^{\prime},s,\pi_{\textsl{g}})}{p_{\pi_{\textsl{g}}}(\pi_{\textsl{g}}^{\prime}|s^{\prime},s,\pi_{\textsl{g}})}\Bigg{)}\Bigg{]}+\mathop{\mathbb{E}}\limits_{\begin{subarray}{c}\pi_{\textsl{g}}^{\prime}\sim u_{\pi_{\textsl{g}}}\\ s^{\prime}\sim T_{\kappa}\end{subarray}}[v(s^{\prime},\pi_{\textsl{g}}^{\prime})]\Bigg{]}$

(9)

where $s=(\sigma,\textsl{g},x,a)$. This objective function is similar to the saLMDP, however, the cost function $q(s,\pi_{\textsl{g}})$ is an additive mixture of functions $q_{\sigma\textsl{g}}(\sigma,\pi_{\textsl{g}})+q_{\sigma}(\sigma)+q_{\pi}(x,a,\pi_{\textsl{g}})$. Just like in Section 3.2, $q_{\sigma\textsl{g}}$ and $q_{\sigma}$ encodes the ordering constraints and a constant state-cost, while $q_{\pi}(x,a,\pi_{\textsl{g}})=v_{\pi_{\textsl{g}}}(x,a)$ encodes the cost of following the policy from $(x,a)$ until the ungrounding of g, which is the value function of the policy. This is significant because Jump Operators combined with a value-function cost form compact representations which preserve low-level trajectory information. We motivate this approach with a proof that, in the case of deterministic transition operators and polices, value functions based on the full operator are equivalent to value functions based on the jump operator when the jump-operator cost-function equals the value function of the low-level policy (Appx.5.9). The optimal desirability function and policy for this coupled-state system is:

	$$\displaystyle\mathbf{z}_{GS}=Q_{\sigma\textsl{g}}Q_{\sigma}Q_{\pi}P_{\kappa}\mathbf{z}_{GS}$$		(10)
	$$\displaystyle u_{\pi_{\textsl{g}}}^{*}(\pi_{\textsl{g}}^{\prime}|s^{\prime},s,\pi_{\textsl{g}})=\frac{p_{\pi_{\textsl{g}}}(\pi_{\textsl{g}}^{\prime}|s^{\prime},s,\pi_{\textsl{g}})z_{GS}(s^{\prime},\pi_{\textsl{g}}^{\prime})}{G[z](s^{\prime},s,\pi_{\textsl{g}})}$$		(11)

where the $Q$-matrices are diagonal matrices of each negatively exponentiated cost function, and $P_{\kappa}$ is the matrix for the joint state-policy passive dynamics of $p_{\kappa}(s^{\prime},\pi^{\prime}|s,\pi):=T_{\kappa}(s^{\prime}|s,\pi)p_{\pi_{\textsl{g}}}(\pi^{\prime}|s^{\prime},s,\pi)$, analogous to $P_{xa}$ in (4). In (10), all matrices have the indexing scheme $(\sigma,\textsl{g},x_{g},a_{g},\pi_{\textsl{g}})$. However, the grounding function is one-to-one, so there are only $N_{\textsl{g}}$ $(\textsl{g},x,a)$ tuples, and there are $N_{\textsl{g}}$ polices $\pi_{\textsl{g}}$, which advance $\sigma$ by ungrounding from these state-action pairs, called the Grounded Subspace (GS). Hence, the transition matrix $P_{\kappa}$ can be represented as a spare matrix of size $|\Sigma|N_{\textsl{g}}^{2}\times|\Sigma|N_{\textsl{g}}^{2}$. We can compute the first-exit problem restricted to this sub-space, then compute the desirability for the agent’s best initial state policy pair $(s,\pi_{\textsl{g}})_{0}$ which lies outside the GS (for all possible $\pi_{\textsl{g}}$ from the state). This is the desirability to enter grounded subspace (and then complete the task) at a given grounded state-action under a policy from the initial $(x,a,\pi_{\textsl{g}})_{0}$. Henceforth, we call this the Desirability-To-Enter (DTE): $\mathbf{z}_{DTE}=\bar{Q}_{\sigma\textsl{g}}\bar{Q}_{\sigma}\bar{Q}_{\pi}\bar{P}_{\kappa}\mathbf{z}_{GS}$, which is a single matrix-vector multiplication where $\mathbf{z}_{GS}$ is the result of (10) computed for the grounded subspace (see 3 for visualization). $\bar{P}_{\kappa}$ is wide matrix of size $N_{\textsl{g}}\times|\Sigma|N_{\textsl{g}}^{2}$ where rows correspond to the agent’s current state $s$ paired with members of the set of GS-polices $\{\pi_{\textsl{g}}\}_{GS}$, and $\bar{Q}$-matrices are $N_{\textsl{g}}\times N_{\textsl{g}}$ diagonal cost matrices specific to the starting state. By computing $\mathbf{z}_{GS}$ separately from $\mathbf{z}_{DTE}$ we effectively break the problem into two parts, a one-step finite horizon problem for interior states (where the "step" is a single policy call, not one discrete time-step) and a first-exit problem for the GS. This enables the agent to be anywhere in the state-space while requiring only one additional linear transformation to act optimally.

Note that the ordering constraints in $Q_{\sigma\textsl{g}}$ are acting on the row space of $P_{\kappa}$, where $\kappa$ determines the structure of $P_{\kappa}$ using global knowledge of sub-goal connectivity. Also, since the cost function $q_{\pi}$ equals the value function, i.e. $q_{\pi}(\pi_{\textsl{g}},x,a)=v_{\pi_{\textsl{g}}}(x,a)$, the cost matrix $Q_{\pi}$ has diagonal entries which are the desirability functions of each sub-problem concatenated into a vector: $Q_{\pi}=diag(exp(-\mathbf{q}_{\pi}))=diag(exp(-\mathbf{v}_{\pi}))=diag(\mathbf{z}_{\pi})$. (see Appx.5.1 for architecture visualization.). Once $u_{\pi_{\textsl{g}}}^{*}$ has been computed, the agent follows the policy indexed by each sampled $\pi_{\textsl{g}}$, i.e. $\pi_{\textsl{g}}\xrightarrow{}u_{xa_{g}}^{*}(x^{\prime},a^{\prime}|x,a)$, until the goal state-action has been reached, followed by a newly sampled $\pi_{\textsl{g}}$. Fig. 5.A shows an example of coupled state trajectory and jump dynamics, drawn from the optimal policy which satisfies the task of getting an axe before chopping wood, and collecting water. Because the tLMDP solves order-constrained multi-goal tasks, it is structurally analogous to (precedence-constrained) Traveling Salesman Problems [30] given a final goal that grounds to the agent’s initial state. Fig. 5.A shows a problem with nine sub-goals with precedence relations on sub-goal colors as types, $\psi$.

The tLMDP solution (10) has three key properties that make it particularly flexible in the context of efficient task transfer:

1. $P1$

   Policy Remappability: Given a complete policy-desirability ensemble and jump operator, under a new grounding function we only need to re-index the polices (where $\pi_{\textsl{g}}$ references different low-level elements of $\mathcal{U}_{All}$ and $\mathcal{Z}_{All}$) to derive new matrices for equation (10). Consequently, for $N_{\textsl{g}}$ goals there is a space of ${N_{x}\choose N_{\textsl{g}}}2^{N_{\textsl{g}}^{2}}$ possible tasks one could formulate⁴⁴4For $N_{\textsl{g}}$ goals grounded onto $N_{x}$ possible states along with a set of ordering constraints which would not require recomputation of $\mathcal{U}_{All},\mathcal{Z}_{All}$ or $J_{All}$.

2. $P2$

   Grounding Invariance: Since the optimal desirability function is a largest eigenvector, there is an invariance of the solution to scalar multiples of the diagonal matrices, assuming $P_{\kappa}$ remains constant under the new grounding.

3. $P3$

   Minimal Solution Distance: The agent can compute the optimal desirability (the desirability-to-enter, DTE) for any given initial state, $s_{0}$, outside the GS with one matrix-vector product. Importantly, this also holds true after the zero-shot regrounding of a previous grounded-subspace solution.

These properties all allow for the fast transfer of task solutions within the same environment, or across different environments with distinct obstacle configurations and transition operators. The following sections will elaborate on the empirical and theoretical implications of these properties.

A key advantage of computing the solution in the GS pertains to Grounding Invariance (property $\mathbf{P2}$), conferred by Policy Remappability (property $\mathbf{P1}$). Consider a solution for the tLMDP problem $\mathscr{L}_{1}$ under grounding $h_{1}$. $\mathbf{P1}$ means that under a new grounding function $h_{2}$ for $\mathscr{L}_{2}$, the GS-index $\pi_{\textsl{g}}$ now references a different policy and desirability vector, and also conditions different dynamics in $J_{All}$. Since the matrices $Q_{\pi}$ and $P_{\kappa}$ are the only two matrices whose elements are functions of $h$, if these matrices are exactly the same under $h_{1}$ and $h_{2}$, then (10) is the same, holding the task-matrices $Q_{\sigma\textsl{g}}$ and $Q_{\sigma}$ constant. Furthermore, (10) being a largest eigenvector solution implies that if $Q_{\pi,h_{2}}$ is a scalar multiple of $Q_{\pi,h_{1}}$, it has the same solution. This means that the solution for $\mathscr{L}_{1}$ will preserve cost-minimization for the new grounding $h_{2}$, which can change distances between sub-goals. This invariance is called a task and cost preserving Grounding Invariant Equation (tc-GIE), as in (12):

	$\displaystyle\text{tc-GIE:}\leavevmode\nobreak\ \mathbf{z}_{GS}=\gamma Q_{\sigma\textsl{g}}Q_{\sigma}Q_{\pi}P_{\kappa}\mathbf{z}_{GS}\leavevmode\nobreak\ \leavevmode\nobreak\ :\leavevmode\nobreak\ \leavevmode\nobreak\ \exists\gamma\in\mathbb{R}^{+},\leavevmode\nobreak\ Q_{\pi,h_{1}}=\gamma Q_{\pi,h_{2}},\leavevmode\nobreak\ P_{\kappa_{1}}=P_{\kappa_{2}},$		(12)
	$\displaystyle\text{t-GIE:}\leavevmode\nobreak\ \mathbf{z}_{GS}=Q_{\sigma\textsl{g}}Q_{\sigma}P_{\kappa}\mathbf{z}_{GS}\leavevmode\nobreak\ \leavevmode\nobreak\ :\leavevmode\nobreak\ \leavevmode\nobreak\ Q_{\pi,h_{1}},Q_{\pi,h_{2}}\leftarrow I,\leavevmode\nobreak\ P_{\kappa_{1}}=P_{\kappa_{2}}.$		(13)

If, however, $Q_{\pi,h_{2}}$ is not a scalar multiple of $Q_{\pi,h_{1}}$, it is still possible for the two problems to share the same solution if $Q_{\pi}$ is set to the identity matrix for both problems (policy cost equal to zero). This means that the task will be solved, however, the policy will not minimize the low-level trajectory costs on $\mathcal{X}$. This is a relaxed version of grounding-invariance called the task preserving Grounding Invariant Equation (t-GIE), (13). Fig.5.E shows the relationship between $\mathscr{L}_{GS}$ (from 5.D) and $\mathscr{L}_{tGIE}$, a t-GIE problem with the same task specification. Whereas $\mathscr{L}_{GS}$ sees a slight additional time increase for optimizing (10) for every new grounding, $\mathscr{L}_{tGIE}$ solves the task with no further optimization, resulting in zero-shot transfer. Furthermore, regardless of whether two problem formulations are related by (12) or (13), the ordering violation costs in $Q_{\sigma\textsl{g}}$ are hard constraints (infinite cost), so as long as $P_{\kappa}$ and $Q_{\sigma\textsl{g}}$ remains constant between $\mathscr{L}_{1}$ and $\mathscr{L}_{2}$, a policy $u^{*}_{\pi,1}$ for $\mathscr{L}_{1}$ can be used to solve $\mathscr{L}_{2}$ however, the sub-goal dynamics will retain the biases of $\mathscr{L}_{1}$’s low-level policy costs, dictated by the obstacle configuration.

Figure 6 demonstrates both GIE equations for an environment with a cliff that the agent can descend, but not ascend, resulting in different $P_{\kappa}$-matrices depending on the grounding. Since $P_{\kappa}$ is a function of $h$ indirectly through $\kappa$, differences can only arise in $P_{\kappa}$ under different sub-goal connectivity structures. These connectivity structures are stored in the matrix $K(i,g)=\kappa((x,a)_{g},\textsl{g}|(x,a)_{i},\pi_{\textsl{g}})$, which represents the global connectivity of the sub-goals. Furthermore, recall that $Q_{\pi}=diag(\mathbf{z}_{\pi})$ where $\mathbf{z}_{\pi}$ is a concatenated state-action vector of all low-level desirability functions in $\mathcal{Z}_{GS}$. Because the desirability function is defined as $z(x)=e^{-v(x)}$, constant multiples of $z$ translate into constant additives of $v(x)$, (i.e. $\gamma z(x)=e^{-(v(x)+\alpha)}$ where $\gamma=e^{-\alpha}$). Since we are working with shortest path polices, to better illustrate the tc-GIE relationship in Figure 6, we construct a matrix $S$ of shortest path distances which are derived from the value function divided by the internal state cost, which gives us the shortest path distance, i.e. $S(i,g)=\frac{v_{g}((x,a)_{i})}{c}$, where $c$ is the constant internal state-action cost used to define $q(x,a)$, and $v_{g}=-log(z_{g})$ is the value function corresponding to the low-level policy with boundary state $(x,a)_{g}$. The new grounding functions preserve the solution under an additive constant $\alpha$ to inter-subgoal distances, which translates into a multiplicative constant $e^{-a}e^{1/c}$ between desirability functions used to define $Q_{\pi}$. Therefore, even between environments with different obstacle constraints, we can compute that two tasks in the grounded subspace are exactly the same because they share the same underlying objective function. The t-GIE is also depicted in Figure 6: task $T_{3}$ in $E_{3}$ does not share a policy cost matrix with a multiplicative constant with $T_{1},T_{2}$, however, the feasibility structure remains the same and we can therefore set $Q_{\pi}\leftarrow I$ for both problems to obtain the same objective function.

We begin by creating a state vector ${y}=(x,a)$ and substitute this state in the equation for an LMDP:

$\displaystyle v^{*}({y})=v^{*}(x,a)=\min_{u_{xa}}\Big{[}q(x,a)+KL(u_{xa}(x^{\prime},a^{\prime}|x,a)||p_{xa}(x^{\prime},a^{\prime}|x,a))\,\,+\,\,\mathop{\mathbb{E}}_{\mathclap{x^{\prime},a^{\prime}\sim u_{xa}}}\,\,[v^{*}(x^{\prime},a^{\prime})]\Big{]}$

(14)

We now decompose the joint distributions using the product rule and impose a dynamics constraint. We would like the low level dynamics of the saLMDP, $p_{x}(x^{\prime}|x,a)$, to behave like an uncontrollable physical system and control to only be achieved over the state-space by manipulating the action distribution, $u_{a}$. To enforce this condition we introduce the constraint that $u_{x}=p_{x}$, and we minimize over $u_{a}$.

	$\displaystyle=\min\limits_{\begin{subarray}{c}u_{a}\\ u_{x}\end{subarray}}\Bigg{[}q(x,a)+\mathop{\mathbb{E}}\limits_{\begin{subarray}{c}a^{\prime}\sim u_{a}\\ x^{\prime}\sim u_{x}\end{subarray}}\Bigg{[}log\Bigg{(}\frac{u_{a}(a^{\prime}|x^{\prime};x,a)u_{x}(x^{\prime}|x,a)}{p_{a}(a^{\prime}|x^{\prime};x,a)p_{x}(x^{\prime}|x,a)}\Bigg{)}\Bigg{]}+\mathop{\mathbb{E}}\limits_{\begin{subarray}{c}a^{\prime}\sim u_{a}\\ x^{\prime}\sim u_{x}\end{subarray}}[v(x^{\prime},a^{\prime})]\Bigg{]}\quad\stackrel{{\scriptstyle s.t.}}{{u_{x}=p_{x}}}$		(15)
	$\displaystyle=\min\limits_{\begin{subarray}{c}u_{a}\end{subarray}}\Bigg{[}q(x,a)+\mathop{\mathbb{E}}\limits_{\begin{subarray}{c}a^{\prime}\sim u_{a}\end{subarray}}\Bigg{[}log\Bigg{(}\frac{u_{a}(a^{\prime}|x^{\prime};x,a)}{p_{a}(a^{\prime}|x^{\prime};x,a)}\Bigg{)}\Bigg{]}+\mathop{\mathbb{E}}\limits_{\begin{subarray}{c}a^{\prime}\sim u_{a}\\ x^{\prime}\sim u_{x}=p_{x}\end{subarray}}[v(x^{\prime},a^{\prime})]\Bigg{]}$		(16)

We consider (16) to be the canonical form of the saLMDP objective function.

After transforming the cost-to-go, $v(x,a)$, into desirability via $v(x,a)=-log(z(x,a))$ the derivation proceeds as follows:

	$\displaystyle=\min\limits_{\begin{subarray}{c}u_{a}\end{subarray}}\Bigg{[}q(x,a)+\mathop{\mathbb{E}}\limits_{\begin{subarray}{c}a^{\prime}\sim u_{a}\end{subarray}}\Bigg{[}log\Bigg{(}\frac{u_{a}(a^{\prime}|x^{\prime};x,a)}{p_{a}(a^{\prime}|x^{\prime};x,a)}\Bigg{)}\Bigg{]}+\mathop{\mathbb{E}}\limits_{\begin{subarray}{c}a^{\prime}\sim u_{a}\\ x^{\prime}\sim p_{x}\end{subarray}}[-log(z(x^{\prime},a^{\prime}))]\Bigg{]}$		(17)
	$\displaystyle=\min\limits_{\begin{subarray}{c}u_{a}\end{subarray}}\Bigg{[}q(x,a)+\mathop{\mathbb{E}}\limits_{x^{\prime}\sim p_{x}}\Bigg{[}\mathop{\mathbb{E}}\limits_{a^{\prime}\sim u_{a}}\Bigg{[}log\Bigg{(}\frac{u_{a}(a^{\prime}|x^{\prime};x,a)}{p_{a}(a^{\prime}|x^{\prime};x,a)z(x^{\prime},a^{\prime})}\Bigg{)}\Bigg{]}\Bigg{]}\Bigg{]}$		(18)
	$\displaystyle=\min\limits_{\begin{subarray}{c}u_{a}\end{subarray}}\Bigg{[}q(x,a)+\mathop{\mathbb{E}}\limits_{x^{\prime}\sim p_{x}}\Bigg{[}\mathop{\mathbb{E}}\limits_{a^{\prime}\sim u_{a}}\Bigg{[}log\Bigg{(}\frac{u_{a}(a^{\prime}|x^{\prime};x,a)}{p_{a}(a^{\prime}|x^{\prime};x,a)z(x^{\prime},a^{\prime})\frac{G[z](x^{\prime},x,a)}{G[z](x^{\prime},x,a)}}\Bigg{)}\Bigg{]}\Bigg{]}\Bigg{]}$		(19)
	$\displaystyle=\min\limits_{\begin{subarray}{c}u_{a}\end{subarray}}\Bigg{[}q(x,a)+\mathop{\mathbb{E}}\limits_{x^{\prime}\sim p_{x}}\Bigg{[}-log\big{(}G[z](x^{\prime},x,a)\big{)}+\mathop{\mathbb{E}}\limits_{a^{\prime}\sim u_{a}}\Bigg{[}log\Bigg{(}\frac{u_{a}(a^{\prime}|x^{\prime};x,a)}{\frac{p_{a}(a^{\prime}|x^{\prime};x,a)z(x^{\prime},a^{\prime})}{G[z](x^{\prime},x,a)}}\Bigg{)}\Bigg{]}\Bigg{]}\Bigg{]}$		(20)

Where $G[\cdot](\cdot)$ is a normalization term that satisfies the KL divergence requirement that the terms sum to 1.

$\displaystyle G[z](x^{\prime},x,a)=\sum_{a^{\prime}\in\mathcal{A}}p_{a}(a^{\prime}|x^{\prime};x,a)z(x^{\prime},a^{\prime})$

(21)

By setting the policy $u_{a}^{*}$ to the denominator (20), the third term goes to zero.

$\displaystyle u_{a}^{*}(a^{\prime}|x^{\prime};x,a)=\frac{p_{a}(a^{\prime}|x^{\prime};x,a)z(x^{\prime},a^{\prime})}{G[z](x^{\prime},x,a)}$

(22)

This leaves two terms that do not depend on the minimization argument $u_{a}$. The resulting equation is:

	$\displaystyle-log(z(x,a))=q(x,a)-\mathop{\mathbb{E}}\limits_{x^{\prime}\sim p_{x}}\big{[}log\big{(}G[z](x^{\prime},x,a)\big{)}\big{]}$		(23)
	$\displaystyle z(x,a)=exp(-q(x,a))exp\Big{(}\mathop{\mathbb{E}}\limits_{x^{\prime}\sim p_{x}}\big{[}log\big{(}G[z](x^{\prime},x,a)\big{)}\big{]}\Big{)}$		(24)