Trajectory Coupling

Today I want to share a rabbit hole I went down where I made an (in my view) interesting mistake when learning about bisimulation metrics. (*1)

Bisimulation Metrics

Bisimulation metrics can be viewed as a measure of behavioral equivalence/similarities between states in Markov chains. As a starter, consider a Markov chain over state space $\mathcal{X}$ with a reward function $R: \mathcal{X} \to \mathbb{R}.$ Roughly speaking, if two states $x_A$ and $x_B$ admit similar rewards (*2)—the rewards are the “behavior” of the states that actually matter—we are willing to say the states are similar, even if $x_A$ and $x_B$ themselves “look” very different (e.g., $x_A$ and $x_B$ are real vectors that are far away from each other in Euclidean distance).

Of course things are not that easy: these states will transition to random next-states $x_A'$ and $x_B'$ , who will produce their own rewards $R(x_A')$ and $R(x_B')$ . So when we say $x_A$ and $x_B$ are behaviorally similar, we want to make sure not only they share similar immediate rewards $R(x_A) \approx R(x_B)$ , but they also transition to distributions of states that admit similar rewards, i.e., $R(x_A')$ and $R(x_B')$ are similar distributed, and $x_A'$ and $x_B'$ further transition to states whose reward similarities need to be accounted for… and so on and so forth. (*3)

The bisimulation metric captures the above intuition mathematically. In fact, it allows you to measure state similarity even between different chains. Consider two Markov chains defined over state spaces $\mathcal{X}$ and $\mathcal{Y}$ , respectively, with transition kernels $P_X$ and $P_Y$ . (If you want to measure state similarity in a single chain $P_X$ , just invoke the definition with $\mathcal{Y}=\mathcal{X}$ and $P_X = P_Y$ .) Let $R_X$ and $R_Y$ be the reward functions. Then, the bisimulation metric $d(\cdot,\cdot): \mathcal{X}\times\mathcal{Y} \to \mathbb{R}_{\ge 0}$ measures the distance between a state $x \in \mathcal{X}$ and $y\in \mathcal{Y}$ as:

d(x,y) = (1-\gamma) |R_X(x) - R_Y(y)| + \gamma \inf_{p \in \Pi(P_X(\cdot|x), P_Y(\cdot|y))} \mathbb{E}_{(x', y') \sim p}[d(x',y')], \quad (1)

where $\Pi(\mu, \nu)$ is a coupling between two distributions as often seen in the definition of Wasserstein’s distance, i.e., for $\mu\in\Delta(\mathcal{X})$ and $\nu\in\Delta(\mathcal{Y})$ , $p \in \Pi(\mu, \nu)$ is any joint distribution whose $X$ marginal matches $\mu$ and whose $Y$ marginal matches $\nu$ .

Trajectory Coupling

A remarkable observation, made here and here, is that Eq.(1) is the Bellman optimality equation of an “augmented MDP” with state space $\mathcal{X}\times\mathcal{Y}$ . (*4) Viewed this way, $d(.,.)$ is the $V^\star$ in this augmented MDP, the LHS is simply $V^\star$ at state $(x,y)$ , and $(1-\gamma) \vert R_X(x) - R_Y(y)\vert$ on the RHS is the immediate reward. $\inf_p$ is minimizing over actions (so we are doing cost minimization instead of reward maximization), and the transition distribution over the next $(x',y')$ is given by the action $p$ itself.

And my down-the-rabbit-hole journey starts here: what does it look like to rollout Monte-Carlo trajectories in the MDP described by Eq.(1), under some policy for choosing $p \in \Pi(P_X(\cdot\mid x), P_Y(\cdot\mid y))$ ? Well, we start with some $(x,y)$ as initial state, i.e., $(X_0,Y_0)=(x,y)$ . In each step $t$ , a policy chooses a $p$ to draw $(X_t, Y_t)$ in a way that respects the transition laws of $P_X$ and $P_Y$ , i.e., $P(X_t \mid X_{t-1}) = P_X(X_t \mid X_{t-1})$ , and similarly for $Y$ . So at the end of the day, this defines a joint distribution $q$ over

X_0, X_1, X_2, \ldots, Y_0, Y_1, Y_2, \ldots

where the marginal of $q$ over $X_0, X_1, X_2, \ldots$ is clearly identically distributed as a trajectory drawn from $P_X$ (let’s denote the distribution of trajectories starting from $x$ as $P_{\{X_t\}}^{x}$ ), and similarly for $Y$ . Thus, we have

q \in \Pi(P_{\{X_t\}}^{x}, P_{\{Y_t\}}^{y}).

Moreover, the random return along this trajectory is

\sum_{t\ge 0} \gamma^t |R_X(X_t) - R_Y(Y_t)|. \quad \quad (2)

Seeing this, a guess immediately came to my mind: if a policy for choosing $p$ defines a coupling between $X$ and $Y$ trajectories, perhaps the optimal policy in Eq.(1) simply induces the optimal coupling between $P_{\{X_t\}}^{x}$ and $P_{\{Y_t\}}^{y}$ for minimizing Eq.(2)? If this is true, we will have a remarkable consequence that, $d(x,y)$ is the Wasserstein distance between $P_{\{X_t\}}^{x}$ and $P_{\{Y_t\}}^{y}$ , where the base metric is Eq.(2), i.e.,

d(x,y) = \min_{q \in \Pi(P_{\{X_t\}}^{x}, P_{\{Y_t\}}^{y})} \mathbb{E}_{\{X_t\}, \{Y_t\} \sim q}[\sum_{t\ge 0} \gamma^t |R_X(X_t) - R_Y(Y_t)|]. \quad (3)

Believing my intuition was correct, I started to sketch a proof based on the following claims (my mistake is hidden somewhere; see if you can spot it before I reveal the answer!):

Claim 1: A policy for choosing $p$ in the augmented MDP defines a coupling between $P_{\{X_t\}}^{x}$ and $P_{\{Y_t\}}^{y}$ .

Justification: Evident from the description above.

Claim 2: Conversely, any coupling between $P_{\{X_t\}}^{x}$ and $P_{\{Y_t\}}^{y}$ is a history-dependent policy in the augmented MDP.

Justification: The joint distribution $q$ between $\{X_t\}$ and $\{Y_t\}$ specified by the coupling can always be written as the product of $q(X_{t+1}, Y_{t+1} \mid X_{1:t}, Y_{1:t})$ , which maps history of the augmented MDP to a distribution over next $X_{t+1}, Y_{t+1}$ . Since the coupling respects the marginals of $P_{\{X_t\}}^{x}$ and $P_{\{Y_t\}}^{y}$ , it also respects $P_X(X_{t+1} \mid X_{t})$ and $P_Y(Y_{t+1} \mid Y_{t})$ , which is consistent with the “ $p \in \Pi(P_X(\cdot\mid x), P_Y(\cdot\mid y))$ ” specification in Eq.(1).

Claim 3: Among all history-dependent policies, the optimal one is Markov, i.e., the form of policy considered in Eq.(1).

Justification: Basic property of an infinite-horizon discounted MDP.

Combining the claims, it follows that the optimal policy implied in Eq.(1) minimizes Eq.(2) over all trajectory couplings.

Weird Implications

For a few weeks(!) I believed my reasoning was correct and was pretty excited about Eq.(3), as it implies an embarrassingly simple way (compared to what’s known in the literature) to estimate $d(\cdot, \cdot)$ : roll out multiple trajectories from $x$ and $y$ , resp., and simply calculate the empirical Wasserstein according to Eq.(2) as the base metric! “There is an elegant paper to write,” I said to myself, “only if I could make sense of a few confusing implications…”

Implication 1: The rate of estimating Wasserstein’s suffers a curse of dimensionality, i.e., if the base metric is a standard distance (say $\ell_1$ or $\ell_2$ ) in $\mathcal{R}^{\mathsf{d}}$ , typically you’d expect a poor rate of $n^{-1/{\mathsf{d}}}$ . Now let’s look at Eq.(2): if we assume all trajectories end in $H$ steps (or perform the standard truncation), Eq.(2) is an $\ell_1$ distance in $\mathbb{R}^H$ , resulting in $n^{-1/H}$ rate.

Now what’s really peculiar is this: if the process only has terminal rewards, i.e., $R(X_t)=R_Y(Y_t) = 0$ unless $t=H$ , then Eq.(2) is $\ell_1$ in $\mathbb{R}$ , improving the estimation rate to $1/n$ ! This is super weird and strongly counterintuitive, as terminal rewards alone as an assumption almost never makes a problem fundamentally easier in RL.

Implication 2: In the hierarchy of state abstractions, model-irrelevance (which is closely related to bisimulation metrics (*3)) is at the finest level. Value-irrelevance (such as $Q^\pi$ - or $Q^\star$ -irrelevance) is a significantly coarser notion, allowing the aggregation of states with similar expected values. Looking at Eq.(3), however, it is clearly closer to the value-irrelevance notion of abstraction than model-irrelevance. In fact, it is like a distributional variant of the former, and they coincide when the Markov chain is deterministic and there are only terminal rewards. This is also counterintuitive as we have not made any significant assumptions that should collapse the abstraction hierarchy.

Understanding the Mistake

The more I thought about the weirdness, the more I believed Eq.(3) is wrong. Eventually, I was able to come up with a simple counterexample:

counterexample

This is a finite-horizon example and I only consider $d(\cdot, \cdot)$ across states at the same “level”, and it can be easily converted to the discounted setting. Basically we have two chains that both produce a Bernoulli-distributed terminal reward, but their bifurcations occur at different time steps. If we calculate $d(A,B)$ according to Eq.(3), clearly we have $d(A,B)=0$ . However, if we perform dynamic programming according to Eq.(1), we see that $d(C, D)>0$ and $d(C, E) >0$ , so it must follow that $d(A,B)>0$ .

So what went wrong? (Last chance to make your own guess before the answer is revealed!)

Let’s examine the optimal trajectory coupling in Eq.(3), which is obviously

q(\{X_t\} = ACF, \{Y_t\} = BDH) = q(\{X_t\} = ACG, \{Y_t\} = BEI) = 0.5.

Now, the conditionals like $q(F\mid C) = 0.5$ are consistent with the transition rule of each chain. However, when you do the dynamic programming step in Eq.(1) by hand, you realize that what you need is not $q(F \mid C) = 0.5$ , but $q(F\mid C, D) = q(F\mid C, E) = 0.5$ ! So among my claims, Claim 2 is incorrect: not any trajectory coupling can be realized by a valid policy in the augmented MDP. This is because an arbitrary trajectory coupling $q$ only satisfies:

q(X_{t+1} \mid X_t) = P_X(X_{t+1} \mid X_t),

but for the $q$ to be valid in the augmented MDP (i.e., it can be induced as the trajectory distribution of some policy), we must have:

q(X_{t+1} \mid X_t, Y_t = y) = P_X(X_{t+1} \mid X_t), ~ \forall y.

In fact, the latter condition is discussed in Sec 2.1 of Calo et al., but without going through all of this I won’t be able to tell why conditioning on $Y_t$ is necessary. If I have to find an excuse for myself, I’d say that $p \in \Pi(P_X(\cdot\mid x), P_Y(\cdot\mid y))$ in Eq.(1) is quite deceptive—it makes you feel that only $P_X(\cdot\mid x)$ needs to be respected, while the dependence on $y$ is secretly hidden on the LHS of the equation.

To conclude, the valid trajectory distributions in the augmented MDP is generally a small subset of all possible trajectory couplings $\Pi(P_{\{X_t\}}^x, P_{\{Y_t\}}^y)$ . I want to develop more intuitions about this subset, and for now I only have some vague ideas like it looks “nested” as transition mass from $(x,y)$ cannot be coupled with that from a different pair of states (if this means anything). (*5) In any case, the weird implications that bothered me are all wrong, so I can finally make peace with myself. Phew.

Footnote Remarks

(*1) This was inspired by a conversation with Necmiye Ozay and her student Alperen Tercan at L4DC.

(*2) More generally, you don’t need to have a reward, but just some way to measure the “instantaneous” similarity between states.

(*3) This is also the intuition I use for bisimulation abstractions, a highly related topic that is very close to my heart and I teach in my RL courses (I might be one of the very few who does that). The difference and connection between bisimulation abstraction and metric is a topic for another day, but very briefly: for the former, you need to first specify an abstraction, then you can calculate the violation of the bisimulation criteria. Bisimulation metrics are more straightforward and generally tighter as a result: you directly compute $d(\cdot, \cdot)$ from Eq.(1), and can decide how to aggregate states based off it.

(*4) One can extend bisimulation metrics to MDPs, where we need max over $a\in\mathcal{A}$ on the RHS of Eq.(1). The corresponding “augmented MDP” will actually become a two-player zero-sum game, where the two players take actions in turns (like Chess or Go): the max player chooses an action $a$ , and the min player chooses $p$ afterwards. The max player here plays a (what I call) “synchronized” policy $\mathcal{X}\times\mathcal{Y}\to\mathcal{A}$ , that is, choose an action for a $(x,y)$ pair that is supposedly “coupled” (in the sense they are drawn from a coupling created by the min player). In contrast, when we aggregate states according to bisimulation metrics (see (*3)), we may end up learning a policy that takes different actions for similar states; if you loosely think of these states as “coupled”, then the policy is not “synchronized” which can cause issues. In this sense, on-policy bisimulation metrics, which is basically bisimulation metrics for chains induced by an arbitrary fixed policy, is not necessarily “less pessimistic” than the $\max_a$ version, as two states close in the latter can have large distance in the former if the policy takes different actions in them.

(*5) I guess what I am looking for is some way to encode the subset in a clean(?) way. For example, a general trajectory coupling $q$ in Eq.(3) only needs to satisfy $q(\{X_t\}) = P_{\{X_t\}}^x(\{X_t\})$ and similarly for $Y$ , which provides some “trajectory-level” understanding (that is, if it were right) compared to the very “local” form of the transition coupling in Eq.(1). Now, Calo et al. has something really close (their Eq.(6-8)), which is an exact characterization of the valid occupancies in the augmented MDP. (So this entire blogpost is really just me learning their paper in my own way!) Still, I’d be curious to know if there is a way to characterize the valid distributions at the trajectory level.