Continuous MDP Homomorphisms and Homomorphic Policy Gradient
Abstract
Abstraction has been widely studied as a way to improve the efficiency and generalization of reinforcement learning algorithms.
In this paper, we study abstraction inthe continuous-control setting.
We extend the definition of MDP homomorphismsto encompass continuous actions in continuous state spaces.
We derive a policygradient theorem on the abstract MDP, which allows us to leverage approximatesymmetries of the environment for policy optimization.
Based on this theorem, wepropose an actor-critic algorithm that is able to learn the policy and the MDP homomorphism map simultaneously, using the lax bisimulation metric.
We demonstratethe effectiveness of our method on benchmark tasks in the DeepMind Control Suite.
Our method’s ability to utilize MDP homomorphisms for representation learningleads to improved performance when learning from pixel observations.
Introduction
[Figure]
HPG image of M obtained π be from
For reinforcement learning from high-dimensional observations, such as images, learning a simpler problemby abstraction from the original problem can be critical [2, 49].
The coupling between states, actions andrewards complicates learning RL abstractions.
MDPhomomorphisms [64, 65, 67, 59] define a concept thatallows one to exploit symmetries, yielding behavioralequivalence and preserving values, while giving the potential to arrive at a substantially smaller MDP.
Recentworks [81, 83, 11] have shown that using MDP homomorphisms is effective to guide learning in discrete problems.
This paper is one of the first to consider MDPhomomorphisms in the continuous-control setting andto develop actor-critic algorithms with a tightly integrated state-action abstraction.
To that end, we identifyand answer a series of key challenges:
Can MDP homomorphisms be defined on continuousstate and action spaces?
Our first contribution is to define continuous MDP homomorphisms oncontinuous state and action spaces, which requires more intricate proofs that do not follow in anydirect way from the finite 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 theabstract MDP into the policy gradient.
Importantly, we rigorously prove that performing HPG onthe 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 approximatesymmetries for improved sample efficiency.
To derive these results, we prove that continuous MDPhomomorphisms 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 practicaldeep 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 unifiesstate and action abstractions.
DHPG is able to simultaneously learn the policy and the homomorphismmap using the lax bisimulation metric [78], a metric for measuring the equivalence of state-actionpairs under an MDP homomorphism relation.
We empirically show that state-action abstractionslearned through MDP homomorphisms provide a natural inductive bias for representation learning.
Despite the existence of well-studied abstraction notions, learning state abstractions in a scalablefashion 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, andcompared to previous works guided by bisimulations [91, 29, 46], our algorithm leads to more robustsolutions, as suggested by our empirical results.
Most importantly, the key difference between priorworks [81, 83, 91, 29, 46] and our method is the homomorphic policy gradient theorem which allowsfor a tight integration of the abstraction notion into the policy gradient.
Our contributions are:
Defining continuous MDP homomorphisms on continuous state and action spaces, usingtools from measure theory and differential geometry.
Proving that continuous MDP homomorphisms preserve value and optimal value functions.3.
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 representationsthat can preserve values and represent the minimal MDP image.
To the best of our knowledge, thisis the first homomorphic policy gradient derivation and the first work to define and scale up MDPhomomorphisms to continuous visual control problems.
Our code will be publicly available.
Background2.1 Markov Decision Processes
We consider the standard MDP that is defined by a 5-tuple (S , A ,τ a ,R,γ ) , with state space S , actionspace A , transition dynamics τ a ∶S×A→ Dist (S) , reward function R ∶S×A→ R , and discount factorγ ∈( 0 , 1 ] .
The goal is to find a policy π ∶ S → Dist (A) that maximizes the expected sum of discountedrewards, the expected return , defined as E π
[ R t ]= E π
[∑ Tk = 0 γ k r t + k + 1 ] .
Value function V π ( s ) denotesthe expected return from s under policy π , and action-value function
Q π ( s,a ) denotes the expectedreturn from s after taking action a under π .
Value functions are fixed 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 fixed points of the Bellman optimality equation [8].
2.2 MDP Homomorphisms
MDP homomorphisms are formally defined for finite MDPs by Ravindran and Barto [65] as:
Definition 1 (MDP Homomorphism) .
An MDP homomorphism h =(
f,g s )∶M→M is a surjectivemap from a finite MDP M=(S , A ,R,τ a ,γ ) onto an abstract finite MDP M=(S , A ,R,τ a ,γ ) wheref ∶S →S
and g s ∶A→A are surjective maps onto the abstract state and action spaces:
Invariance of reward: R ( f ( s ) ,g s ( a ))
= R ( s,a ) ∀ s ∈ S ,a ∈ A (1)
Equivariance of transitions: τ g s ( a ) ( f ( s ′ ) ∣ f ( s )) =
[Equation]
where B h is the partition of S induced by the equivalence relation of homomorphism h , B h ∣S isthe 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 homomorphismmap h ∶M→M/ B h .
As MDP homomorphisms are sensitive with respect to changes in rewards ortransitions, approximate MDP homomorphisms [67] allow equations (1-2) to hold approximately.
The significance 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 liftingit 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 twopolicies must satisfy ∑ a ∈ g − 1 s ( a ) π ↑ ( a ∣ s )
=
π ( a ∣ f ( s )) for all s ∈ S and a ∈ A .
2.3 Bisimulation and Lax Bisimulation Metrics
Bisimulation for finite MDPs [20, 30] defines an equivalence relation on S where two states s iand 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 94 every equivalence class C defined by the equivalence relation.
The rigidity of bisimulation limits itsapplications.
Bisimulation metrics [26, 25, 27] measure the equivalence as an approximation:
[Equation]
where the first term measures reward similarity and K is the Kantorovich (Wasserstein) metricmeasuring the distance between the transition probabilities.
However, bisimulation metrics can stillbe brittle as they require the behaviour to match for all actions.
This may be problematic particularlyin the case of continuous actions in which small changes to actions may not drastically changethe 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 thestate equivalence relation to state-action equivalence.
Taylor et al.
[78] show that lax bisimulation isprecisely the same relation as the MDP homomorphism and define the lax bisimulation metric as:
[Equation]
Furthermore, Taylor et al.
[78] show that minimizing the lax bisimulation metric corresponds tofinding 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 first prove their value equivalence property in the finite case as the generalization of the priorresult 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 finite 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, weextend their definition to continuous state and action spaces, and derive results analogous to the finitecase.
First, we define continuous MDPs and state our underlying assumptions.
Importantly, thecorrect definitions of continuous MDPs and continuous MDP homomorphisms require care regardingmeasurability and differentiability of spaces, and our formulation is chosen to fit the HPG derivation;see Appendix A.2 for an overview of the tools we used from measure theory and differential geometry.
Definition 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 ofactions , is a locally compact metric space, usually taken to be a subset of R n , τ a is the transitionprobability kernel for each possible action a , for each fixed 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 toactions a , which is needed in Section 5 for deriving the HPG theorem.
Given the continuous MDPs described above, we define continuous MDP homomorphisms.
Theequivariance condition on the transition dynamics, Equation (2), can no longer be expressed in termsof a discrete sum over partitions, and instead we use the σ -algebra structure on the different statespaces.
Definition 3 (Continuous MDP Homomorphism) .
A continuous MDP homomorphism is a maph =
( f,g s )
∶ M → M where f ∶ S → S and for every s in S , g s ∶ A → A are measurable, surjective 132 maps such that the following hold:
Invariance
[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 Definition 3, we prove that optimal value functions are preservedby the continuous MDP homomorphism as in the finite 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 measureof τ 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 finite case, we also require a lifting process to define π ↑ = lift ( π ) given a policy π on theabstract 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 policiesexist based on the disintegration theorem, the full proof and design of a computationally tractablealgorithm for this process is left for future work.
Therefore, here and in the subsequent sectionswe assume the policy is deterministic in which case the lifted policy can be simply obtained bychoosing one representative for the preimage g − 1 s ( π ( f ( s ))) .
If we select g s to be a bijection, thelifted policy can be uniquely defined as π ↑ ( s ) =
g − 1 s ( π ( f ( s ))) .
The assumption on deterministicpolicies is not limiting, as in general the optimal policy of a given MDP is deterministic [9].
Withthis lifting definition, 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 fromthe 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 improvingsample efficiency.
Equipped with continuous MDP homomorphisms from Definition 3 and theirvalue 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 ) , parameterizedby 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 conditionsof Definition 2 and Appendix B. Similarly to prior works on policy gradients [76, 70], we definethe performance measure as J ( θ ) =
E π
[ V π ( s )] where the expectation is over the uncertainty intransitions, rewards, and initial states.
Finally, as detailed in Section 4.2, our results are derived fordeterministic policies and a bijective g s .
Notably, this choice allows us to parameterize one of thepolicies 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 onmanifolds, which in turn raises the need for g s to be a bijection and local diffeomorphism.
Theorem4 highlights that the gradient of the abstract MDP is equivalent to that of the original, despite theunderlying spaces being abstracted.
This implies that performing HPG on the abstract MDP isequivalent to performing DPG on the actual MDP, allowing us to use them synergistically to updatethe same parameters θ , as shown in Figure 1.
While one can naively use Theorem 4 to substitute gradients of the standard DPG, theoretically thisdoes not produce any useful result as the expectation remains estimated with respect to the stationarystate distribution of the actual MDP M under π ↑ θ ( s ) .
However, using properties of continuous MDPhomomorphisms, we can change the integration space from S to S , and consequently estimate thepolicy 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 defined on M .
Then the gradient of the performance measure J ( θ ) , defined 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 variablesformula of the pushforward measure on the state space.
The significance of Theorem 5, which formsthe basis of our proposed homomorphic actor-critic algorithm, is twofold.
First, we can get anotherestimate 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 anotheras they are tied through the homomorphism map, HPG will potentially have less variance at theexpense 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 simplified problem.
In other words, each abstract state-action pair ( s,a ) usedto train Q π θ represents all ( s,a ) pairs that are equivalent under the MDP homomorphism relation,thus improving sample efficiency.
However, the amount of complexity reduction is dependent on theapproximate 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 ourmethod as Deep Homomorphic Policy Gradient (DHPG).
Although HPG is applicable to any otherdeterministic 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 suchas image reconstruction [29, 34, 90, 36], as the notion of MDP homomorphism provides a naturalinductive 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, specifically 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 updatea 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 actora =
π θ ( s ) , homomorphism map h ϕ,η =(
f ϕ ( s ) ,g η ( s,a )) , reward predictor R ρ ( s ) , and probabilistic 211 transition model τ ν ( s ′ ∣ s,a ) which outputs a Gaussian distribution.
We use target networks and avanilla 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 afaster reward propagation [7].
The loss function for each critic is therefore defined as the expectationof 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 policysmoothing and delayed actor updates [28].
The pseudo-code of DHPG is presented in Appendix E.1.
Learning Continuous MDP Homomorphisms.
We now address the problem of learning continuous MDP homomorphisms.
While few methods have been proposed for learning finite MDPhomomorphisms [81, 83], these are not readily extendable to continuous actions.
In this work, we usethe lax bisimulation metric [78], Equation (5), to propose a loss function that encodes lax bisimilarstates closer together in the abstract space.
The lax bisimulation metric is applicable to continuousactions and as a (pseudo-)metric, it can naturally represent approximate MDP homomorphisms.
Following the same intuition as prior works on bisimulations [91], we define our proposed lax 230 bisimulation loss over pairs of transition tuples sampled from the replay buffer.
We permute samplesto compute their pairwise distance in the abstract space and their pairwise lax bisimilarity distance.
Consequently, we minimize the distance between these two terms:
[Equation]
[Figure]
⋅∣ f ϕ
[Figure]
s
[Figure]
s
[Figure]
[Figure]
⋅∣
[Figure]
ϕ
[Figure]
j
[Figure]
Similarly to Zhang et al.
[91], we replaced the Kantorovich ( W 1 ) metric in Equation (5) with the W 2metric as there is an explicit formula for it for Gaussian distributions.
Finally, we apply the conditionsof a continuous MDP homomorphism map from Definition 3 via the loss function of:
[Equation]
13) where s ′ i ∼ τ ν (⋅∣ f ϕ ( s i ) ,g η ( s i ,a i )) .
The final 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 thedistribution of results, we follow the best practices proposed by Agarwal et al.
[5] and report theinterquartile mean (IQM) and performance profiles aggregated on all tasks over 10 random seeds.
While our baseline results are obtained using the official code, when possible 2 , some of the resultsmay differ from the originally reported ones due to the difference in the seed numbers and our goal topresent a faithful representation of the true performance distribution [5].
7.1 State Observations
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
∗ = g s ( a ∗ ) = π ∗ ( s ) . The relation g s 1 ( a 1 )
=
g s 2 ( a 2 ) holds for equivalent state-action pairs.
We compare DHPG on state observationsagainst three commonly-used off-policy modelfree algorithms: DDPG [51], TD3 [28], and
SAC [35].
All methods use n -step returns, andshare the same hyperparameters presented in Appendix E.2.
For a fair comparison with DDPGand a better study of the impact of HPG, wehave improved our DDPG by adding delayedpolicy updates and target policy [28].
Thus, theonly difference between our DDPG and TD3 isthe clipped double Q-learning present in TD3,which appears to be hurting the performance insome tasks of DMC as also observed in [60, 50].
DHPG outperforms or matches other algorithms on state observations and has a bettersample efficiency.
Results are presented in Figure 2, and full results are in Appendix D.1 .
Expectedly, performance gains are larger on tasks withsymmetries, as DHPG is able to learn a compressed abstract MDP.
The learned mapping h =( f,g s ) demonstrates properties of an MDP homomorphism.
We usethe pendulum swingup task to visualize its learned MDP homomorphism, as its symmetries areperfectly intelligible.
Two state-action pairs ( s 1 =( θ 1 , ˙ θ 1 ) ,a 1 ) and ( s 2 =( θ 2 , ˙ θ 2 ) ,a 2 ) are equivalentif a 1 = − a 2 , θ 1 = − θ 2 , and ˙ θ 1 = − ˙ θ 2 .
Therefore, the learned action representations are expected toreflect this by setting g s 1 ( a 1 )=
g s 2 ( a 2 ) .
Figure 3a shows contours of optimal actions over S , while 273 Figure 3b shows action representations a
= g s ( a ) of optimal actions over S .
Clearly, abstract actionsadhere to the aforementioned relation for equivalent state-action pairs, indicating g s ( a ) is in factrepresenting 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
[Figure] 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]
DHPG outperforms or matches other algorithms on pixel observations, demonstratingits effectiveness in representation learning.
Results are presented in Figure 4 and full resultsare in Appendix D.2 .
Interestingly, DHPG without image augmentation outperforms DrQ-v2on 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 isparticularly important when learning from pixelobservations.
To demonstrate this ability, wehave 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 learnthe 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 qualitativelydemonstrate the significance 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 [Figure] state encoder f ϕ has mapped corresponding legs to the same abstract latent state (Figure 6b) as theyare 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 withinthe same domain.
Importantly, one consideration with representation learning methods relying onrewards is the transferability of the learned representations to a new reward setting within the samedomain.
To ensure that our method does not hinder such transfer, we conduct experiments in whichthe actor, critics, and the learned MDP homomorphism map are transferred to a new task in the samedomain.
Results, given in Appendix D.3 show that our method has not compromised transfer abilities.
Additional Experiments.
We study the value equivalence property as a measure for the quality of thelearned 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.
Itis extended to continuous state spaces [12, 20] and extended to MDPs [30].
Bisimulation metrics[21, 26, 25, 27] define a pseudometric to quantify the degree of behavioural similarity.
Recently, Zhang et al.
[91] defined 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 methodto approximate the bisimulation metric for deterministic MDPs with discrete actions.
van der Polet al.
[81] have defined a contrastive loss based on MDP homomorphisms for learning an abstract MDP for planning, however, their method is only applicable to finite MDPs.
Another approach is todirectly embed the MDP homomorphic relation in the network architecture [83, 82].
Other recentlyproposed methods seek to learn representations that preserve values [33, 32] or policies [4], or viaa sampling-based similarity metric [15].
Finally, state abstractions can in principle help improvetransferring 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 actionspaces [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-actionembedding for continuous control is learned.
In contrast, we use the notion of homomorphisms tolearn the state-dependent action representations, while preserving values.
Action representations canbe 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 rawobservations 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-supervisedlearning [6, 71, 38, 39, 24].
From another view, representation learning can be effectively decoupledfrom the RL problem [23, 74].
Symmetries of the environment can also be used for representationlearning [58, 55, 61, 84, 40, 41, 62, 13].
In fact, MDP homomorphisms are specializations of suchapproaches for RL.
A key distinguishing factor of MDP homomorphisms is their ability to take actionsinto account for representation learning in the same premises as Thomas et al.
Recently, simpleimage augmentation methods have shown significant 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 measuretheory, and we rigorously proved their value and optimal value equivalence properties.
We derivedthe 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 isequivalent 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 MDPhomomorphism map using the lax bisimulation metric.
Our algorithm improves upon strong baselines.
The visualization of the latent space demonstrates the strong potential of MDP homomorphisms inlearning structured representations that can preserve value functions.
We believe that our work willopen-up future possibilities for MDP homomorphisms in challenging continuous control problems.