Continuous MDP Homomorphisms and Homomorphic Policy Gradient
Abstract
Abstraction has been widely studied as a way to improve the efﬁciency and generalization of reinforcement learning algorithms.
In this paper, we study abstraction in the continuous-control setting.
We extend the deﬁnition of MDP homomorphisms to encompass continuous actions in continuous state spaces.
We derive a policy gradient theorem on the abstract MDP, which allows us to leverage approximate symmetries of the environment for policy optimization.
Based on this theorem, we propose an actor-critic algorithm that is able to learn the policy and the MDP homomorphism map simultaneously, using the lax bisimulation metric.
We demonstrate the effectiveness of our method on benchmark tasks in the DeepMind Control Suite.
Our method’s ability to utilize MDP homomorphisms for representation learning leads to improved performance when learning from pixel observations
Introduction
[Figure]
H o m o m o r ph i c
P o li cy
G r a d i e n t
D e t e r m i n i s t
i c P
[Figure]
li cy G r a d i
[Figure]
n
[Figure]
with HPG.
M is the MDP homomorphic image of M obtained by learning the homomorphism map h =( f,g s ) .
Policies π ↑ and π can be derived from each other.
For reinforcement learning from high-dimensional observations, such as images, learning a simpler problem by abstraction from the original problem can be critical [2, 49].
The coupling between states, actions and rewards complicates learning RL abstractions.
MDP homomorphisms [64, 65, 67, 59] deﬁne a concept that allows one to exploit symmetries, yielding behavioral equivalence and preserving values, while giving the potential to arrive at a substantially smaller MDP.
Recent works [81, 83, 11] have shown that using MDP homomorphisms is effective to guide learning in discrete problems.
This paper is one of the ﬁrst to consider MDP homomorphisms in the continuous-control setting and to develop actor-critic algorithms with a tightly integrated state-action abstraction.
To that end, we identify and answer a series of key challenges:
Can MDP homomorphisms be deﬁned on continuous state and action spaces?
Our ﬁrst contribution is to deﬁne continuous MDP homomorphisms on continuous state and action spaces, which requires more intricate proofs that do not follow in any direct way from the ﬁnite case and requires tools from measure theory and differential geometry.
Can MDP homomorphisms be tightly integrated into the policy gradient?
Our second contribution is the derivation of the homomorphic policy gradient (HPG) theorem to closely integrate the 36th Conference on Neural Information Processing Systems (NeurIPS 2022).
abstract MDP into the policy gradient.
Importantly, we rigorously prove that performing HPG on the abstract MDP is equivalent to performing the deterministic policy gradient (DPG) on the actual
MDP.
Therefore, HPG can act as an additional gradient estimator capable of utilizing approximate symmetries for improved sample efﬁciency.
To derive these results, we prove that continuous MDP homomorphisms preserve value functions [33], which in turn enables their use for policy evaluation.
Can MDP homomorphisms be learned simultaneously with the optimal policy in a practical deep reinforcement learning algorithm?
We propose a deep actor-critic algorithm, depicted in
Figure 1, based on HPG, referred to as Deep Homomorphic Policy Gradient (DHPG), that uniﬁes state and action abstractions.
DHPG is able to simultaneously learn the policy and the homomorphism map using the lax bisimulation metric [78], a metric for measuring the equivalence of state-action pairs under an MDP homomorphism relation.
We empirically show that state-action abstractions learned through MDP homomorphisms provide a natural inductive bias for representation learning.
Despite the existence of well-studied abstraction notions, learning state abstractions in a scalable fashion for continuous control remains a key challenge.
In contrast to previous works on learning
MDP homomorphisms [81, 83, 11], our algorithm is readily applicable to continuous actions, and compared to previous works guided by bisimulations [91, 29, 46], our algorithm leads to more robust solutions, as suggested by our empirical results.
The bisimulation relation [56, 47, 12, 20, 30] and bisimulation metrics [21, 26, 25, 27] do not allow abstracting actions as they require exact matching of actions, whereas MDP homomorphisms and equivalently the lax bisimulation metric [78] remove this strong limitation, giving them greater modeling ﬂexibility.
Most importantly, the key difference between prior works [81, 83, 91, 29, 46] and our method is the homomorphic policy gradient theorem which allows for a tight integration of the abstraction notion into the policy gradient that theoretically motivates using the abstract MDP for policy optimization in the actual MDP.
Our contributions are:
Deﬁning continuous MDP homomorphisms on continuous state and action spaces, using tools from measure theory and differential geometry.
Proving that continuous MDP homomorphisms preserve value and optimal value functions.
Deriving the homomorphic policy gradient theorem.
Developing a deep actor-critic algorithm for learning the optimal policy simultaneously with the MDP homomorphism map in challenging continuous control problems.
DHPG improves upon strong baselines on pixel observations [88, 91] on DM Control, and our visualizations demonstrate the potential of MDP homomorphisms in learning structured representations that can preserve values and represent the minimal MDP image.
To the best of our knowledge, this is the ﬁrst homomorphic policy gradient derivation and the ﬁrst work to deﬁne and scale up
MDP homomorphisms to continuous visual control problems.
Our code is publicly available at https://github.com/sahandrez/homomorphic_policy_gradient .
Background 2.1 Markov Decision Processes
We consider the standard MDP that is deﬁned by a 5-tuple (S , A ,τ a ,R,γ ) , with state space S , action space A , transition dynamics τ a ∶S×A→ Dist (S) , reward function R ∶S×A→ R , and discount factor γ ∈( 0 , 1 ] .
The goal is to ﬁnd a policy π ∶ S → Dist (A) that maximizes the expected sum of discounted rewards, the expected return , deﬁned as E π
[ R t ]= E π
[∑ Tk = 0 γ k r t + k + 1 ] .
Value function V π ( s ) denotes the expected return from s under policy π , and action-value function Q π ( s,a ) denotes the expected return from s after taking action a under π .
Value functions are ﬁxed points of the Bellman equation [8] and can be computed iteratively through a process referred to as policy evaluation [75].
Similarly, optimal value functions V ∗ ( s ) and Q ∗ ( s,a ) are ﬁxed points of the Bellman optimality equation [8].
2.2 MDP Homomorphisms
MDP homomorphisms are formally deﬁned for ﬁnite MDPs by Ravindran and Barto [65] as:
Deﬁnition 1 (MDP Homomorphism) .
An MDP homomorphism h =(
f,g s )∶M→M is a surjective map from a ﬁnite MDP M=(S , A ,R,τ a ,γ ) onto an abstract ﬁnite MDP M=(S , A ,R,τ a ,γ ) where f ∶S →S
and g s ∶A→A are surjective maps onto the abstract state and action spaces:
Invariance of reward: R ( f (
[Equation]
) ,g
[Equation]
=
[Equation]
Equivariance of transitions: τ g s ( a ) ( f ( s ′
[Equation]
∣ f
[Equation]
=
[Equation]
where B h is the partition of S induced by the equivalence relation of homomorphism h ,
B h ∣S is the projection of B h onto S , and [ s ′ ]
B h ∣S denotes the block of B h ∣S to which s ′ belongs.
Thus, when applying action a in state s , the right-hand side is the probability that the resulting state is in [ s ′ ]
B h ∣S .
The abstract MDP M is in fact the quotient MDP M/ B h based on the homomorphism map h ∶M→M/ B h .
As MDP homomorphisms are sensitive with respect to changes in rewards or transitions, approximate MDP homomorphisms [67] allow equations (1-2) to hold approximately.
The signiﬁcance of MDP homomorphisms is the optimal value equivalence between M and M [65]:
[Equation]
which in turn allows for learning the optimal policy π ∗ in the abstract MDP and consequently lifting it to obtain the optimal policy in the actual MDP, using:
[Equation]
where g − 1 s ( a ) denotes the set of actions that have the same image a under g s .
Equivalently, the two policies must satisfy ∑ a ∈ g − 1 s ( a ) π ↑ ( a ∣ s )
=
π ( a ∣ f ( s )) for all s ∈ S and
a
[Equation]
A .
2.3 Bisimulation and Lax Bisimulation Metrics
Bisimulation for ﬁnite MDPs [20, 30] deﬁnes an equivalence relation on S where two states s i and s j are equivalent or bisimilar if R ( s i ,a )=
R ( s j ,a ) and τ a ( C ∣ s i )=
τ a ( C ∣ s j ) for all a ∈A and every equivalence class C deﬁned by the equivalence relation.
The rigidity of bisimulation limits its applications.
Bisimulation metrics [26, 25, 27] measure the equivalence as an approximation:
[Equation]
where the ﬁrst term measures reward similarity and K is the Kantorovich (Wasserstein) metric measuring the distance between the transition probabilities.
However, bisimulation metrics can still be brittle as they require the behaviour to match for all actions.
This may be problematic particularly in the case of continuous actions in which small changes to actions may not drastically change the outcome.
Additionally, bisimulation metrics are not able to represent environment symmetries.
Instead, lax bisimulation [78] waives the requirement on action matching in favor of extending the state equivalence relation to state-action equivalence.
Taylor et al.
[78] show that lax bisimulation is precisely the same relation as the MDP homomorphism and deﬁne the lax bisimulation metric as:
[Equation]
Furthermore, Taylor et al.
[78] show that minimizing the lax bisimulation metric corresponds to ﬁnding approximate MDP homomorphisms and bound the value error.
Value Equivalence Property
To motivate the use of MDP homomorphisms for policy evaluation and consequently policy optimization, we ﬁrst prove their value equivalence property in the ﬁnite case as the generalization of the prior result of the optimal value equivalence [65], stated in Equation (3).
The proof is in Appendix C.1.
Theorem 1 (Value Equivalence) .
Let M be the image of an MDP homomorphism h from a ﬁnite M .
Then any two corresponding policies π ↑ = lift ( π ) have equivalent values:
[Equation]
Continuous MDP Homomorphisms
To concretely lay out the foundations of using MDP homomorphisms for continuous control, we extend their deﬁnition to continuous state and action spaces, and derive results analogous to the ﬁnite case.
First, we deﬁne continuous MDPs and state our underlying assumptions.
Importantly, the correct deﬁnitions of continuous MDPs and continuous MDP homomorphisms require care regarding measurability and differentiability of spaces, and our formulation is chosen to ﬁt the HPG derivation; see Appendix A.2 for an overview of the tools we used from measure theory and differential geometry.
Deﬁnition 2 (Continuous MDP) .
A continuous Markov decision process (MDP) is a 6 -tuple:
[Equation]
where S , the state space is assumed to be a Polish space, Σ is a σ -algebra on S 1 , A , the space of actions , is a locally compact metric space, usually taken to be a subset of R n , τ a is the transition probability kernel for each possible action a , for each ﬁxed s , τ a (⋅∣ s ) is a probability distribution on Σ while R is the reward function, and γ is the discount factor.
Furthermore, for all s ∈ S and B ∈ Σ the map a ↦ τ a ( B ∣ s ) is smooth.
The last assumption is required for differentiability with respect to actions a , which is needed in Section 5 for deriving the HPG theorem.
Given the continuous MDPs described above, we deﬁne continuous MDP homomorphisms.
The equivariance condition on the transition dynamics, Equation (2), can no longer be expressed in terms of a discrete sum over partitions, and instead we use the σ -algebra structure on the different state spaces.
Deﬁnition 3 (Continuous MDP Homomorphism) .
A continuous MDP homomorphism is a map
h =
( f,g s )
∶ M → M where f ∶ S → S and for every s in S , g s ∶ A → A are measurable, surjective maps such that the following hold:
Invariance of reward: R ( f ( s ) ,g
[Equation]
Equivariance of transitions
[Equation]
Note that if g s is the identity map, the second condition reduces to τ a ( B ∣ f ( s )) =
τ a ( f − 1 ( B )∣ s ) which is simply the condition for preservation of transition probabilities as used in bisimulation [20].
4.1 Optimal Value Equivalence
Assuming the conditions given in Deﬁnition 3, we prove that optimal value functions are preserved by the continuous MDP homomorphism as in the ﬁnite case.
Theorem 2 (Optimal Value Equivalence) .
Let M =
(S , Σ , A ,τ a ,R ) be the image of a continuous
MDP homomorphism h
= ( f,g s ) from M =
(S , Σ , A ,τ a ,R ) .
Then:
[Equation]
The proof, given in Appendix C.2, uses the change of variable formula of the pushforward measure of τ a (⋅∣ s ) with respect to f to change the integration space from S to S .
4.2 Value Equivalence for Lifting Deterministic Policies
As in the ﬁnite case, we also require a lifting process to deﬁne π ↑ = lift ( π ) given a policy π on the abstract MDP.
In general, the lifted policy needs to satisfy the relation π ↑ ( g − 1 s ( β )∣ s ) =
π ( β ∣ f ( s )) for every Borel set β ⊆
A and s ∈ S .
While our initial progress shows that lifted stochastic policies exist based on the disintegration theorem, the full proof and design of a computationally tractable algorithm for this process is left for future work.
Therefore, here and in the subsequent sections we assume the policy is deterministic in which case the lifted policy can be simply obtained by choosing one representative for the preimage g − 1 s ( π ( f ( s ))) .
If we select g s to be a bijection, the lifted policy can be uniquely deﬁned as π ↑ ( s ) =
g − 1 s ( π ( f ( s ))) .
The assumption on deterministic policies is not limiting, as in general the optimal policy of a given MDP is deterministic [9].
With this lifting deﬁnition, we state and prove the following value equivalence result:
Theorem 3 (Value Equivalence for Deterministic Policies) .
Let M be the image of a continuous
MDP homomorphism h
= ( f,g s ) from M , then any two deterministic policies π ↑ ∶ S
→ A and π ∶ S → A where π ↑ = lift ( π ) have equivalent value functions on their domain:
[Equation]
The proof, given in Appendix C.3, uses the change of variable formula of the pushforward measure of τ a (⋅∣ s ) with respect to f to change the integration space from S to S and assumes g s to be bijective.
Homomorphic Policy Gradient
The next goal of this work is to derive a policy gradient estimator using samples obtained from the abstract MDP.
Intuitively, this allows for direct incorporation of state-action abstraction as an inductive bias for policy optimization, thereby reducing the variance of actor updates and improving sample efﬁciency.
Equipped with continuous MDP homomorphisms from Deﬁnition 3 and their value equivalence property, we now derive the homomorphic policy gradient (HPG) theorem.
In this section, we assume having access to an MDP homomorphism map h =( f,g s ) , parameterized by differentiable functions.
The problem of learning such mapping from samples is addressed in
Section 6.
Additionally, we assume the MDP and the homomorphism map adhere to the conditions of Deﬁnition 2 and Appendix B. Similarly to prior works on policy gradients [76, 70], we deﬁne the performance measure as J ( θ ) =
E π
[ V π ( s )] where the expectation is over the uncertainty in transitions, rewards, and initial states.
Finally, as detailed in Section 4.2, our results are derived for deterministic policies and a bijective g s .
Notably, this choice allows us to parameterize one of the policies and to uniquely derive the other policy.
In practice, we parameterize the actual policy as π ↑ θ and obtain the abstract policy as π θ = g s ( π ↑ θ ( s )) .
First, we show the equivalence of policy gradients :
Theorem 4 (Equivalence of Deterministic Policy Gradients) .
Let M be the image of a continuous
MDP homomorphism h from M , and let π ↑ θ ∶ S
→ A be the lifted deterministic policy corresponding to the abstract deterministic policy π θ ∶ S
→ A .
Then for any ( s,a ) ∈ S
× A we have:
[Equation]
The proof is given in Appendix C.4 and uses the chain rule and the inverse function theorem on manifolds, which in turn raises the need for g s to be a bijection and local diffeomorphism.
Theorem 4 highlights that the gradient of the abstract MDP is equivalent to that of the original, despite the underlying spaces being abstracted.
This implies that performing HPG on the abstract MDP is equivalent to performing DPG on the actual MDP, allowing us to use them synergistically to update the same parameters θ , as shown in Figure 1.
While one can naively use Theorem 4 to substitute gradients of the standard DPG, theoretically this does not produce any useful result as the expectation remains estimated with respect to the stationary state distribution of the actual MDP M under π ↑ θ ( s ) .
However, using properties of continuous MDP homomorphisms, we can change the integration space from S to S , and consequently estimate the policy gradient with respect to the stationary distribution of the abstract MDP M under π θ ( s ) :
Theorem 5 (Homomorphic Policy Gradient Theorem) .
Let M be the image of a continuous MDP homomorphism h from M , and let π θ ∶ S
→ A be a deterministic abstract policy deﬁned on M .
Then the gradient of the performance measure J ( θ ) , deﬁned on the actual MDP M , w.r.t.
θ is:
[Equation]
where ρ π θ ( s ) is the discounted state distribution of M following the deterministic policy
π θ ( s ) .
The proof is given in Appendix C.5 and applies the result of Theorem 4 and the change of variables formula of the pushforward measure on the state space.
The signiﬁcance of Theorem 5, which forms the basis of our proposed homomorphic actor-critic algorithm, is twofold.
First, we can get another estimate for the policy gradient based on the approximate MDP homomorphic image in addition to
DPG.
Although the two policy gradient estimates are not statistically independent from one another as they are tied through the homomorphism map, HPG will potentially have less variance at the expense of some bias due to the approximation of the MDP homomorphism.
Second, since the minimal image of an MDP is the MDP homomorphic image [65], the abstract critic
Q π θ is trained on a simpliﬁed problem.
In other words, each abstract state-action pair ( s,a ) used to train Q π θ represents all ( s,a ) pairs that are equivalent under the MDP homomorphism relation, thus improving sample efﬁciency.
However, the amount of complexity reduction is dependent on the approximate symmetries of the environment, as also supported by our empirical results.
Homomorphic Actor-Critic Algorithms
We propose a deep actor-critic algorithm based on HPG by adapting DDPG [51].
We refer to our method as Deep Homomorphic Policy Gradient (DHPG).
Although HPG is applicable to any other deterministic actor-critic, we chose DDPG as its simplicity compared to modern choices [7, 43] allows for a better study on the impact of MDP homomorphisms.
Notably, a strong advantage of
DHPG is that it is readily applicable to pixel observations without the need for extra mechanisms such as image reconstruction [29, 34, 90, 36], as the notion of MDP homomorphism provides a natural inductive bias for learning representations that preserve values and optimal values.
Since DHPG is learning the MDP homomorphism map h online and concurrently with the policy, using the actual MDP for training the critic, speciﬁcally at the early stages of training, is helpful.
Therefore, we utilize a separate critic for each MDP M and M .
Ultimately, critics are used to update a single set of parameters, as shown in Figure 1; see Appendix D.5 for the ablation study on this.
Thus, the components of DHPG are: actual critic Q ψ ( s,a ) , abstract critic Q ψ ( s,a ) , deterministic actor a = π θ ( s ) , homomorphism map h φ,η =( f φ ( s ) ,g η ( s,a )) , reward predictor R ρ ( s ) , and probabilistic transition model τ ν ( s ′ ∣ s,a ) which outputs a Gaussian distribution.
We use target networks and a vanilla replay buffer [57, 51].
As discussed in Section 5, an abstract actor is obtained as a = g s ( π θ ( s )) .
In case of pixel observations, a single image encoder E µ is shared among all components.
Training the Policy and Critic.
Actual and abstract critics are trained using n -step TD error for a faster reward propagation [7].
The loss function for each critic is therefore deﬁned as the expectation of the n -step Bellman error estimated over transitions samples from the replay buffer B :
[Equation]
where s t = f φ ( s t ) and a t =
g η ( s t ,a t ) are computed using the learned MDP homomorphism, ψ ′ and ψ ′ denote parameters of target networks, and R ( n ) t =
∑ n − 1 i =
γ i r t
+ i is the n -step return.
Consequently, we train the policy using DPG [70] and HPG from Theorem 5 by backpropagating the following loss:
[Equation]
Here, the two gradients are added together and a single policy update is conducted; see Appendix
D.5 for the ablation study on other combinations of HPG and DPG.
Finally, we utilize target policy smoothing and delayed actor updates [28].
The pseudo-code of DHPG is presented in Appendix E.1.
[Figure]
Learning Continuous MDP Homomorphisms.
We now address the problem of learning continuous MDP homomorphisms.
While few methods have been proposed for learning ﬁnite MDP homomorphisms [81, 83], these are not readily extendable to continuous actions.
In this work, we use the lax bisimulation metric [78], Equation (5), to propose a loss function that encodes lax bisimilar states closer together in the abstract space.
The lax bisimulation metric is applicable to continuous actions and as a (pseudo-)metric, it can naturally represent approximate MDP homomorphisms.
Following the same intuition as prior works on bisimulations [91], we deﬁne our proposed lax bisimulation loss over pairs of transition tuples sampled from the replay buffer.
We permute samples to compute their pairwise distance in the abstract space and their pairwise lax bisimilarity distance.
Consequently, we minimize the distance between these two terms:
[Equation]
Similarly to Zhang et al.
[91], we replaced the Kantorovich ( W 1 ) metric in Equation (5) with the Wmetric as there is an explicit formula for it for Gaussian distributions.
Finally, we apply the conditions of a continuous MDP homomorphism map from Deﬁnition 3 via the loss function of:
[Equation] [Equation]
where s ′ i ∼ τ ν (⋅∣ f φ ( s i ) ,g η ( s i ,a i )) .
The ﬁnal loss function is obtained as L lax ( φ,η ) + L h ( φ,η,ρ,ν ) .
Experiments
In our experiments, we aim to answer the following key questions:
Does the homomorphic policy gradient improve policy optimization?
What are the qualitative properties of the learned representations and the abstract MDP? 3.
Can DHPG learn and recover the minimal MDP image from raw pixel observations?
We evaluate DHPG on continuous control tasks from DM Control on state and pixel observations.
Importantly, to reliably evaluate our algorithm against the baselines and to correctly capture the distribution of results, we follow the best practices proposed by Agarwal et al.
[5] and report the interquartile mean (IQM) and performance proﬁles aggregated on all tasks over 10 random seeds.
While our baseline results are obtained using the ofﬁcial code, when possible 2 , some of the results may differ from the originally reported ones due to the difference in the seed numbers and our goal to present a faithful representation of the true performance distribution [5].
7.1 State Observations
[Figure]
We compare DHPG on state observations against three commonly-used off-policy modelfree algorithms: DDPG [51], TD3 [28] and
SAC [35].
All methods use n -step returns, and share the same hyperparameters presented in Appendix E.2.
For a fair comparison with DDPG and a better study of the impact of HPG, we have improved our DDPG by adding delayed policy updates and target policy [28].
Thus, the only difference between our DDPG and TD3 is the clipped double Q-learning present in TD3, which appears to be hurting the performance in some tasks of DMC as also observed in [60, 50].
DHPG outperforms or matches other algorithms on state observations and has a better sample efﬁciency.
Results are presented in Figure 2, and full results are in Appendix D.1 .
Expectedly, performance gains are larger on tasks with symmetries, as DHPG is able to learn a compressed abstract MDP.
[Figure]
The learned mapping h =( f,g s ) demonstrates properties of an MDP homomorphism.
We use the pendulum swingup task to visualize its learned MDP homomorphism, as its symmetries are perfectly intelligible.
Two state-action pairs ( s 1 =( θ 1 , ˙ θ 1 ) ,a 1 ) and ( s 2 =( θ 2 , ˙ θ 2 ) ,a 2 ) are equivalent if a 1 = − a 2 , θ 1 = − θ 2 , and ˙ θ 1 = − ˙ θ 2 .
Therefore, the learned action representations are expected to reﬂect this by setting g s 1 ( a 1 )=
g s 2 ( a 2 ) .
Figure 3a shows contours of optimal actions over S , while Figure 3b shows action representations a
= g s ( a ) of optimal actions over S .
Clearly, abstract actions adhere to the aforementioned relation for equivalent state-action pairs, indicating g s ( a ) is in fact representing the action encoder of an MDP homomorphism mapping.
7.2 Pixel Observations
We compare the effectiveness of DHPG on pixel observations against DBC [91], DeepMDP [29],
SAC-AE [90], and state-of-the-art performing DrQ-v2 [88].
All methods use n -step returns, share the same hyperparameters in Appendix E.2 and all hyperparameters are adapted from DrQ-v2 without any further tuning .
Importantly, for a fair comparison with DrQ-v2 which uses image augmentation, we present two variations of DHPG and other baselines, with and without image augmentation .
[Figure]
(b) Curves.
[Figure]
DHPG
[Figure]
Latent dim
[Figure] [Figure]
Latent dim
[Figure]
Latent dim
[Figure]
-AE (Latent dim: 4) (c) Curves.
DHPG outperforms or matches other algorithms on pixel observations, demonstrating its effectiveness in representation learning.
Results are presented in Figure 4 and full results are in Appendix D.2 .
Interestingly, DHPG without image augmentation outperforms DrQ-v on domains with easily learnable MDP homomorphism maps, such as cartpole and pendulum, showing its power of representation learning.
DHPG can learn and recover a lowdimensional MDP image.
A key strength of
MDP homomorphisms is their ability to represent the minimal MDP image [65], which is particularly important when learning from pixel observations.
To demonstrate this ability, we have limited the latent space dimensions to the dimension of the real system and compared DHPG (without image augmentation) with baselines in Figure 5.
While other methods are not able to learn the tasks, DHPG can successfully learn the policy and the minimal low-dimensional latent space.
Surprisingly, trajectories of the latent states resemble that of the real states as shown in Figure 5a.
The abstract MDP demonstrates properties of an MDP homomorphic image.
To qualitatively demonstrate the signiﬁcance of learning joint state-action representations, Figure 6 shows visualizations of latent states for quadruped-walk, a task with symmetries around movements of its four legs.
Interestingly, while the latent space of DHPG (Figure 6a) shows distinct states for each leg, abstract state encoder f φ has mapped corresponding legs (e.g., left forward leg and right backward leg) to the same abstract latent state (Figure 6b) as they are some homomorphic image of one another.
Clearly,
DBC and DrQ-v2 are not able to achieve this.
The learned representations and the MDP homomorphism map transfer to new tasks within the same domain.
Importantly, one consideration with representation learning methods relying on rewards is the transferability of the learned representations to a new reward setting within the same domain.
To ensure that our method does not hinder such transfer, we have carried out experiments in which the actor, critics, and the learned MDP homomorphism map are transferred to another task from the same domain.
Results, given in Appendix D.3 show that our method has not compromised transfer abilities.
Accounting for the larger network capacity of DHPG compared to the baselines.
Since our
DHPG algorithm contains additional networks, such as the parameterized MDP homomorphism map and the abstract critic, it may have a higher network capacity compared to the baselines.
To control for the effect of the network capacity and for a fair evaluation, we compare DHPG with higher-capacity variants of DBC and DrQ-v2 that have a larger critic networks, selected such that the total number of parameters are considerably more than that of DHPG.
Results are presented in Figure 7, while full results and a detailed description of the total number of parameters are in Appendix D.7 .
As suggested by the results, DHPG outperforms or matches the performance of the higher-capacity baselines, demonstrating the improved performance is rather due to the use of the abstract MDP homomorphic image for representation learning and performing HPG updates.
Additional Experiments.
We study the value equivalence property as a measure for the quality of the learned MDP homomorphisms in Appendix D.4, and we present ablation studies on DHPG variants, and the impact of n -step return on our method in Appendices D.5 and D.6, respectively.
Related Work
State Abstraction.
Bisimulation [56, 47] is a notion of behavioral equivalence between systems.
It was extended to continuous state spaces by Blute et al.
[12, 20] and extended to MDPs by Givan et al.
Bisimulation metrics [21, 26, 25, 27] deﬁne a pseudometric to quantify the degree of behavioural similarity.
Recently, Zhang et al.
[91] deﬁned a loss function for learning representations via bisimilarity of latent states, and Kemertas et al.
[46] have further improved its robustness.
Castro [14] has proposed a method to approximate the bisimulation metric for deterministic MDPs with continuous states but discrete actions.
van der Pol et al.
[81] have deﬁned a contrastive loss based on MDP homomorphisms for learning an abstract MDP for planning, however, their method is only applicable to ﬁnite MDPs.
Another approach is to directly embed the MDP homomorphic relation in the network architecture [83, 82].
Other recently proposed methods seek to learn representations that preserve values [33, 32] or policies [4], or via a sampling-based similarity metric [15].
Finally, state abstractions can in principle help improve transferring of policies [1, 16, 72, 73, 63], or learning temporally extended actions [17, 87, 86, 77].
Action Abstraction.
Action representations are often studied in the context of large discrete action spaces [68] as a form of a look-up embedding that is known a-priori [22], factored representations [69], or policy decomposition [18].
Action representations can also be learned from expert demonstrations [79].
More related to our work is dynamics-aware embeddings [85] where a combined state-action
[Figure]
[Figure]
embedding for continuous control is learned.
In contrast, we use the notion of homomorphisms to learn the state-dependent action representations, while preserving values.
Action representations can be combined with temporal abstraction [77] for discovering extended actions [66, 3, 16, 17].
State Representation Learning.
Extant methods for learning the underlying state space from raw observations often use latent models [29, 36, 37, 34, 10], auxiliary prediction tasks [44, 53, 54], physics-inspired inductive biases [45, 19, 31], unsupervised learning [42, 52], or self-supervised learning [6, 71, 38, 39, 24].
From another point of view, representation learning can be effectively decoupled from the RL problem [23, 74].
Symmetries of the environment can also be used for representation learning [58, 55, 61, 84, 40, 41, 62, 13].
In fact, MDP homomorphisms are specializations of such approaches for RL.
A key distinguishing factor of MDP homomorphisms is their ability to take actions into account for representation learning in the same premises as Thomas et al.
Recently, simple image augmentation methods have shown signiﬁcant improvements in RL performance [89, 48].
Since these approaches are in general orthogonal to state abstractions, they can be combined together.
Conclusion
In this paper, we developed the novel theory of continuous MDP homomorphisms using measure theory, and we rigorously proved their value and optimal value equivalence properties.
We derived the homomorphic PG in order to directly use a joint state-action abstraction for policy optimization.
Importantly, we rigorously proved that applying our homomorphic PG on the abstract MDP is equivalent to applying the standard DPG on the actual MDP.
Based on our novel theoretical results, we developed a deep actor-critic algorithm that can simultaneously learn the policy and the MDP homomorphism map using the lax bisimulation metric.
Our algorithm improves upon strong baselines in both learning from state and pixel observations.
The visualization of the latent space demonstrates the strong potential of MDP homomorphisms in learning structured representations that can preserve value functions.
We believe that our work will open-up future possibilities for the application of MDP homomorphisms in challenging continuous control problems.
Acknowledgements
SRS is supported by an NSERC CGS-D scholarship.
RZ was supported by an NSERC CGS-M scholarship at the time this work was completed.
We would like to thank Juan Camilo Gamboa
Higuera, Harley Wiltzer, and Scott Fujimoto for insightful discussions.
The computing resources for this research were provided by Calcul Quebec and the Digital Research Alliance of Canada.