Continuous MDP Homomorphisms and Homomorphic Policy Gradient

Main Idea

The paper extends MDP homomorphisms to continuous state and action spaces and shows how the resulting abstraction can be used directly for policy optimization.

The motivation is that an MDP may contain state-action pairs that look different but are equivalent for decision making. For example, in pendulum swing-up, the states $(\theta, \dot\theta)$ and $(-\theta,-\dot\theta)$ are symmetric if the action is also reflected from $a$ to $-a$. Learning each case independently wastes experience. An MDP homomorphism maps both cases to the same abstract control problem while preserving rewards, transition behavior, and values.

The main contribution is the Homomorphic Policy Gradient (HPG) theorem: under an exact homomorphism, the deterministic policy gradient computed in the abstract MDP is equivalent to the deterministic policy gradient in the original MDP. Therefore, the agent can optimize its policy using a simpler abstract representation without changing the underlying optimization objective.

Continuous MDP Homomorphism

Let $\mathcal{M}$ be the original MDP and $\bar{\mathcal{M}}$ its abstract image. A homomorphism is composed of

$$h=(f,g_s)$$

where

• $f:\mathcal{S}\rightarrow\bar{\mathcal{S}}$ maps original states to abstract states
• $g_s:\mathcal{A}\rightarrow\bar{\mathcal{A}}$ maps original actions to abstract actions and may depend on the current state

The mapping must preserve rewards:

$$R(s,a)=\bar{R}(f(s),g_s(a))$$

and preserve transition probabilities after states are mapped into the abstract space:

$$\bar{P}(\bar{B}\mid f(s),g_s(a)) =P(f^{-1}(\bar{B})\mid s,a).$$

This differs from ordinary state-only bisimulation because the action map can align different actions before comparing their outcomes. This makes homomorphisms more flexible for representing environment symmetries.

Value Preservation

If $h$ is an exact homomorphism, the optimal value functions are preserved:

$$V^*(s)=\bar{V}^*(f(s))$$ $$Q^*(s,a)=\bar{Q}^*(f(s),g_s(a)).$$

Thus, states and actions that collapse to the same abstract pair have the same role in optimal decision making. The abstract MDP removes redundancy but retains all information required to obtain an optimal policy.

An abstract policy must ultimately be converted back into an action that can be executed in the original environment. For a deterministic abstract policy $\bar{\pi}$, the lifted policy is

$$\pi^{\uparrow}(s)=g_s^{-1}(\bar{\pi}(f(s))).$$

The paper assumes that $g_s$ is bijective so this inverse is uniquely defined. The abstract policy and its lifted policy also have equivalent value functions.

Homomorphic Policy Gradient

The key theoretical result is that the policy-gradient contribution from an original state-action pair is equal to the contribution from its abstract representation:

$$\nabla_a Q^{\pi^{\uparrow}}(s,a)\nabla_\theta\pi^{\uparrow}_\theta(s) = \nabla_{\bar a}\bar Q^{\bar\pi}(\bar s,\bar a)\nabla_\theta\bar\pi_\theta(\bar s),$$

where $\bar{s}=f(s)$ and $\bar{a}=g_s(a)$.

The expected policy gradient can therefore be written entirely in terms of the abstract MDP:

$$\nabla_\theta J(\theta) =\mathbb{E}_{\bar{s}\sim\rho^{\bar{\pi}}} \left[ \nabla_{\bar a}\bar Q^{\bar\pi}(\bar{s},\bar{a}) \nabla_\theta\bar{\pi}_\theta(\bar{s}) \right].$$

Under the exact assumptions, HPG is not a different optimization objective: it is the original deterministic policy gradient expressed in the quotient MDP. Its potential advantage is that one abstract state-action pair represents many equivalent original pairs.

Deep Homomorphic Policy Gradient

The practical algorithm, DHPG, does not assume that the homomorphism is known. It learns the abstraction and policy simultaneously using:

• an original critic $Q(s,a)$
• an abstract critic $\bar{Q}(f(s),g_s(a))$
• a state encoder $f$
• a state-dependent action encoder $g_s$
• an abstract reward predictor and transition model

The actor is updated using both critics:

$$\mathcal{L}_{actor}\approx-\mathbb{E}_{s\sim\mathcal{B}} \left[Q(s,\pi(s))+\bar{Q}(f(s),g_s(\pi(s)))\right].$$

The original critic provides a reliable learning signal while the abstraction is still inaccurate. The abstract critic provides an additional gradient that shares information between behaviorally equivalent state-action pairs.

Learning the Abstraction

DHPG learns approximate homomorphisms using the lax bisimulation metric. Two state-action pairs should be represented closely when they have

• similar immediate rewards
• similar distributions over abstract next states

Unlike standard bisimulation, lax bisimulation does not require the same action to be taken in both states. The learned action map can discover that different actions have equivalent effects in different states. The transition distributions are compared with a Wasserstein distance, and the state encoder is trained so that latent distances reflect these reward and transition differences.

For two transitions $(s_i,a_i,r_i,s_i')$ and $(s_j,a_j,r_j,s_j')$ sampled from the replay buffer, the lax bisimulation loss is

$$\begin{aligned} \mathcal{L}_{\mathrm{lax}}(\phi,\eta) =\mathbb{E}_{\mathcal{B}}\Big[& \lVert f_\phi(s_i)-f_\phi(s_j)\rVert_1 -\lVert r_i-r_j\rVert_1 \\ &-\alpha W_2\Big( \bar{\tau}_\nu(\cdot\mid f_\phi(s_i),g_\eta(s_i,a_i)), \bar{\tau}_\nu(\cdot\mid f_\phi(s_j),g_\eta(s_j,a_j)) \Big) \Big]. \end{aligned}$$

Here, $\bar{\tau}_\nu$ is the learned abstract transition model and $W_2$ is the Wasserstein distance between the two predicted next-state distributions. The loss makes distance in the abstract state space reflect differences in rewards and transition dynamics.

The homomorphism conditions are additionally encouraged through a transition and reward prediction loss:

$$\mathcal{L}_{h}(\phi,\eta,\nu,\rho) =\mathbb{E}_{(s_i,a_i,s_i',r_i)\sim\mathcal{B}} \left[ \left\lVert f_\phi(s_i')-\bar{s}_i'\right\rVert_2^2 +\left(r_i-\bar{R}_\rho(f_\phi(s_i))\right)^2 \right],$$

where the predicted abstract next state is sampled from

$$\bar{s}_i'\sim \bar{\tau}_\nu\left( \cdot\mid f_\phi(s_i),g_\eta(s_i,a_i) \right).$$

The final abstraction objective is

$$\mathcal{L}_{\mathrm{abstraction}} =\mathcal{L}_{\mathrm{lax}}(\phi,\eta) +\mathcal{L}_{h}(\phi,\eta,\nu,\rho).$$

Pseudocode

Limitations

• the policy-gradient theory in this paper is restricted to deterministic policies
• the action map must be bijective and differentiable for the gradient-equivalence result
• learning the homomorphism requires additional critics and reward/transition models
• improvements depend on the environment having an abstraction that is both useful and learnable