Asymmetric Norms to Approximate the Minimum Action Distance

In discrete and finite MDPs we can compute the state-transition graph $G=(V,E)$ of an MDP. In this section, we will revise how to learn the minimum action distance from the graph adjacency matrix.

A state-transition graph $G=(V,E)$ of an MDP $\mathcal{M}=\langle\mathcal{S},\mathcal{A},\mathcal{P},r\rangle$ is a graph with nodes representing the states in the MDP and the edges representing state adjacency in the MDP. More precisely, $V=\mathcal{S}$, $e\left(s_{i},s_{j}\right)\in E\,$ iff $\,\exists a\>\mathcal{P}\left(s_{i},a,s_{j}\right)>0$. An adjacency matrix represents a graph with a square matrix of size $|\mathcal{S}|\times|\mathcal{S}|$ with $(i,j)$-value being 1 if $e\left(s_{i},s_{j}\right)\in E$ and 0 otherwise.

Having access to the adjacency matrix $A^{G}$ we can simply compute the minimum action distance by using the Floyd-Warshall algorithm [Floyd, 1962, Roy, 1959, Warshall, 1962].

The Floyd-Warshall algorithm compares all possible paths through the graph between each pair of vertices. It is able to do this with $\Theta\left(|V|^{3}\right)$ comparisons in a graph, even though there may be up to $\Omega\left(|V|^{2}\right)$ edges in the graph, and every combination of edges is tested. It does so by incrementally improving an estimate on the shortest path between two vertices until the estimate is optimal.

This dynamic programming procedure relies on having access to the edge weights, which in the case of MAD reduces to having access to the adjacency matrix $A^{G}$ of $a_{i,j}=1$ when $s_{i},s_{j}$ are connected by an edge.

Thanks to this we can define the shortest path on the $A^{G}$ by just first computing the base cases:

and subsequently computing the recursive case leveraging the triangle inequality property:

Note that in case we do not have access to the adjacency matrix $A^{G}$ this can be retrieved by interacting with the environment by visiting all the one-step transitions $(s,s^{\prime})$.

The Minimum Action Distance between states is a priori unknown and is not directly observable in continuous and/or noisy state spaces where we cannot simply enumerate the states and construct the adjacency matrix of the MDP. Instead, we will approximate it using the distances between states observed on trajectories. We introduce the notion of Trajectory Distance (TD) as follows:

We can observe that given any state trajectory $\tau=\{s_{0},...,s_{n}\}$, choosing any pair of states $(s_{i},s_{j})\in\tau$ with $0\leq i\leq j\leq n$, their distance along the trajectory represents an upper bound of the MAD.

Given a dataset of trajectories $\mathcal{D}$ collected by any unknown behavior policy, we can retrieve the MAD $d_{MAD}$ by solving the following constrained optimization problem:

As a first step we can leverage the triangle inequality to simplify the constrain in 6 and reduce the dependency on the quality of the trajectories in the dataset $\mathcal{D}$.

where $(s,s^{\prime})\in\tau$ refers to a one-step transition (i.e. $d_{TD}(s,s^{\prime}|\tau)=1$) in the trajectory $\tau\in\mathcal{D}$ while $(s_{i},s_{j},s_{k})\in\mathcal{S}_{\mathcal{D}}^{3}$ indicates all the combinations of $3$ states contained in the trajectory dataset $\mathcal{S}_{\mathcal{D}}$.

Note that the first constraint in 7 imposes an upper bound on one-step transitions, i.e. it says that two states $(s,s^{\prime})$ at distance one along a trajectory $d_{TD}(s,s^{\prime}|\tau)=1$ are either the same state $s=s^{\prime}$ or they must satisfy $d_{MAD}=1$. This allows us to approximate the MAD without having to identify whether two states along a trajectory are the same state or not.

Note that the second constraint in 7 corresponds to the triangle inequality, which thus is a property that holds for the MAD. Moreover, this second constraint implies that we have to calculate it for all the combinations of $3$ states contained in $\mathcal{S}_{\mathcal{D}}$ which can become intractable for large state spaces.

To address this issue we proposed an alternative formulation based on embedding the MAD in a parametric embedding space $\phi_{\theta}:\mathcal{S}\rightarrow\mathbb{R}^{dim}$ where a chosen distance metric that respects the triangle inequality (e.g. any norm $||\cdot||_{p}$) can be used to enforce the triangle inequality constraint.

The goal is to learn a parametric state embedding $\phi_{\theta}:\mathcal{S}\rightarrow\mathbb{R}^{dim}$ such that the distance $d$ between any pair of embedded states approximates the Minimum Action Distance.

Steccanella and Jonsson [2022] used this formulation to favour symmetric embeddings, where they use norms as distance functions, e.g. the L1 norm $d(z,y)=||z-y||_{1}$. If we use symmetric embeddings we will have that for any pair of states $(s_{i},s_{j})\in\mathcal{S}$,

The problem of learning this embedding can then be formulated as a constrained optimization problem:

Intuitively, the objective is to make the embedded distance between pairs of states as close as possible to the observed trajectory distance while respecting the upper bound constraints. Without constraints, the objective is minimized when the embedding matches the expected Trajectory Distance $\mathbb{E}\left[d_{TD}\right]$ between all pairs of states observed on trajectories in the dataset $\mathcal{D}$. In contrast, constraining the solution to match the minimum TD with the upper-bound constraints $\left\|\phi_{\theta}(s)-\phi_{\theta}(s^{\prime})\right\|_{p}\leq d_{TD}(s,s^{\prime}\mid\tau)$ allows us to approximate the MAD. The precision of this approximation depends on the quality of the given trajectories.

To make the constrained optimization problem tractable, we can relax the hard constraints in (9) and convert them into a penalty term to retrieve a simple unconstrained formulation. Moreover, we rely on sampling $(s_{i},s_{j},d_{TD}(s_{i},s_{j}\mid\tau))$ and $(s,s^{\prime},d_{TD}(s,s^{\prime}\mid\tau))$ from the dataset of trajectories $\mathcal{D}$ making this formulation amenable for gradient descent and to fit within the optimization scheme of neural networks.

where $C$ is our penalty term defined as

The penalty term $C$ introduces a quadratic penalization of the objective for violating the upper-bound constraints $\left\|\phi_{\theta}(s)-\phi_{\theta}(s^{\prime})\right\|_{p}<=d_{TD}(s,s^{\prime}\mid\tau)$.

In the previous section, we have seen how it is possible to define the MAD embedding problem with the use of norms $||\cdot||_{p}$. While the formulation is useful to understand how it is possible to remove the triangle inequality constraint in 7 the Minimum Action Distance is naturally asymmetric and we would like embedding that preserves this asymmetry.

A norm is a function $\|\cdot\|:\mathcal{X}\rightarrow\mathbb{R}$ satisfying, $\forall x,y\in\mathcal{X},\alpha\in\mathbb{R}^{+}$:

• •

  $\mathbf{N1}$ (Pos. def.). $\|x\|>0$, unless $x=0$.

• •

  $\mathbf{N2}$ (Pos. homo.). $\alpha\|x\|=\|\alpha x\|$, for $\alpha\geq 0$.

• •

  $\mathbf{N3}$ (Subadditivity). $\|x+y\|\leq\|x\|+\|y\|$.

• •

  $\mathbf{N4}$ (Symmetry). $\|x\|=\|-x\|$.

An asymmetric semi-norm satisfies $\mathbf{N2}$, $\mathbf{N3}$ but not necessarily $\mathbf{N1}$, $\mathbf{N4}$.

A convex function $f:\mathcal{X}\rightarrow\mathbb{R}$ is a function satisfying $\mathbf{C1}:\forall x,y\in\mathcal{X},\alpha\in[0,1]:f(\alpha x+(1-\alpha)y)\leq$ $\alpha f(x)+(1-\alpha)f(y)$. The commonly used ReLU activation, $\operatorname{relu}(x)=\max(0,x)$, is convex.

Is easy to observe that any $\mathbf{N2}$ and any $\mathbf{N3}$ function is convex and thus that any asymmetric semi-norm is convex.

Motivated by this relationship between convex functions and norms,Pitis et al. [2020] introduced Wide Norms, a parametric distance that models symmetric and asymmetric norms.

A Wide Norm is any combination of symmetric/asymmetric semi-norms. They are based on the Mahalanobis norm of $x\in\mathbb{R}^{dim}$, parametrized by $W\in\mathbb{R}^{m\times n}$, defined as $||x||_{W}=||Wx||_{2}$.

Asymmetric Wide Norms are defined as:

We can then use the parametrized Wide Norm distance to constrain the triangular inequality on the embedding space:

where $C$ is our penalty term defined as