Continuous MDP Homomorphisms and Homomorphic Policy Gradient
Authors: Sahand Rezaei-Shoshtari
Rosie Zhao
Prakash Panangaden
David Meger
Doina Precup
Venue: NeurIPS 2022 Accept (2022)
Keywords: reinforcement learning
state abstraction
policy optimization
continuous control
representation learning in reinforcement learning

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 in the continuous-control setting. We extend the definition 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.

Continuous MDP Homomorphisns and Homomorphic Policy Gradient
Sahand Rezaei-Shoshtari
Rosie Zhao
Prakash Panangaden McGill University and Mila
McGill University and Mila
McGill University and Mila
David Meger
Doina Precup
McGill University and Mila
McGill University, Mila, and DeepMind
Abstract
Abstraction has been widely studied way to improve the efficiency and general- ization of reinforcement learning algorithms. In this paper, we study abstraction the continuous-contee setting. We extend the definition of MDP homomorphisms to encompass continuous actions in continuous state spaces. We derive a policy gradient theorem on the abstract MDP, which allows us leverage approximate symmetries of the environment for policy optimization. Based on this theorem, we propose an actor-critic algorithm that able learn the policy and the MDP homo- morphism 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 's ability to utilize MDP homomorphisms for representation learning leads to improved performance when learning from pixel observations.
Introduction
For reinforcement learning from dithensieniii ob- servations, such as images, learning simpler problem by abstraction from the original problem can be criti-
Actual Policy
Abstrect Policy
Determincctie Polacy Gradient
Policy Liting
x
T(s)
Inverse Pollcy Litinng
cal 2 49 The coupling between states, actions and rewards complicates learning RL abstractions. MDP
A= (f.g.)
Homomorphii Pollyy
define concept that allows one to exploit symmetries, yielding behavioral equivalence and preserving values, while giving the po- tential to arrive at a substantially smaller MDP. Recent works [81 83 [I] have shown that using MDP homo- morphisms is effective to guide learning in discrete prob- lems. This paper is one of the first to consider MDP homomorphisms the continuous-contre setting and to develop actor-critic algorithms with tightly inte- grated state-action abstraction. To that end, we identify and answer series of key challenges:
homomorphisms [64, 65 67 59)
Ecuialent
Q (s,a)
Figure 1: Schematics of our method. The actual MDP M is used to train Q=' and update n¹ with DPG, while the abstract MDP M is used to train Q and update क with HPG. M is the MDP homomorphic image of M obtained by learning the ho- momorphism map h= (f (f.gg g.). Policies s! and ग can be derived from each other.
Can MDP homomorphisms be defined on continu- ous state and action spaces? Our first contribution is to define continuous MDP homomorphisms on continuous state and action spaces, which requires more intricate proofs that do not follow in any direct 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 contribu- tion the derivation of the homomorphic policy gradient (HPG) theorem closely integrate the
36th Conference on Neural Information Processing Systems (NcurIPS 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 efficiency. To derive these results, we prove that continuous MDP homomorphisms preserve value functions [3], which in turn enables their use for policy evaluation.
Can MDP homomorphisms be learned simultaneously with the optimal policy in practical deep reinforcement learning algorithm? We propose a deep actor-critic algorithm, depicted in Figure I based on HPG, referred to as Deep Homomorphic Policy Gradient (DHPG), that unifies state and action abstractions, DHPG is able to simultaneously learn the policy and the homomorphism map using the lax bisimulation metric 77), 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 previous works on learning MDP homomorphisms [8] 83 []]. 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 metries [2] 26 25 27 do not allow abstracting actions as they require exact matching of actions, whereas MDP homomorphisms and equivalently the lax bisimulation metrie [78] remove this strong limitation, giving them greater modeling flexibility. 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:
1. Defining continuous MDP homomorphisms on continuous state and action spaces, using tools from measure theory and differential geometry.
2. Proving that continuous MDP homomorphisms preserve value and optimal value functions.
3. Deriving the homomorphic policy gradient theorem.
4. 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 izations 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 first homomorphic policy gradient derivation and the first work to define and scale up MDP homomorphisns to continuous visual control problems. Our code is publicly available at https:// https://github.coo httpb//////ssaaaaee////ooooopppp prtproooorrpccooollccggrrreeee
88 91
DM Control, and visual-
2 Background
2.1 Markov Decision Processes
We consider the standard MDP that is defined by a 5-tuple (S,A,T. To. R,y)), with state space S. action space A, transition dynamics Ta:Sx_ A Dist(S), reward function R:Sx T.SxAA= - R, and discount factor 7 (0, 1]. The goal is to find a policy π : :S - Dist (AA) that maximizes the expected sum of discounted rewards, the expected return, defined as the expected return from under policy æ. and action value function Q" (s. a) denotes the expected return from s after taking action a under т. Value functions are fixed points of the Bellman equation g and can be computed iteratively through 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 [
Value functiin
(*)(s denotes
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= (J.g.):MM :MM is surjective map from a finite MDP M=(S, A, R, Tax 7) onto abstract finite MDP M (5, A, R, Ta, 7) where
2

---

f:S+S and gs:A-A are surjective maps onto the abstract state and action spaces:
Invariance of reward: R(((s)) = (),a)
VsES,
(1)
Equivariance of transitions: (a) (I())))(()))
Σ Ta(o"1s) VsES,aEA
(2)
where B₁ is the partition of S induced by the equivalence relation of homomorphism h, B is the projection of B, onto S, and [8'] JBBis denotes the block of B to which belongs. Thus, when applying action in state s, the right-hand side is the probability that the resulting state is in [s'),, The abstract MDP M is in fact the quotient MDP M/BB, Br based on the homomorphism map hh:MMM/aa As MDP homomorphisms are sensitive with respect to changes in rewards or transitions, approximate MDP homomorphisms [67 allow equations ( hold approximately. The significance of MDP homomorphisms is the optimal value equivalence between M and M 66]
S which in turn allows for learning the optimal policy मं in the abstract MDP and consequently lifting to obtain the optimal policy the actual MDP, using:
f())(s)) (a g3 (a))l'
म' (als)
VsES, (a)
where (ā) denotes the set of actions that have the same image under 9.. Equivalently, the two policies must satisfy Eaeg,1(a) (z) T² (als) क (s) for all s S and ā E A.
2.3 Bisimulation and Lax Bisimulation Metrics
defines an equivalence relation on S where two states s; and s are equivalent bisimilar if R(s, st, a) = R(s).a) a) and Ta (CC(,,)==aal (C(s;) for all a A and every equivalence class defined by the equivalence relation. The rigidity of bisimulation limits its approximation:
Bisimulation for finite MDPs (20 30
applications. Bisimulation metrics 26 25 27 measure the equivalence
drisim (si, 85) max Gr R( 8; a) R(s;. a) e, (TTT(((ss) De.A
(4)
where the first term measures reward similarity and is the Kantorovich (Wasserstein) metric measuring the distance between the transition probabilities. However, bisimulation metrics still be brittle as they require the behaviour to match for all actions. This may be problematic particularly actions may not drastically change
in the case of continuous actions in which small changes the outcome. Additionally, bisimulation metrics are not able to represent environment symmetries. Instead, lax bisimulation [78] waives the requirement on action matching 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 define the lax bisimulation metric as:
dus (ax((s))))) a:); ()),,))) c, [c a,) R(53. ag) c (-1se), (isx),aa, Tas (13;)).
Furthermore, Taylor al. 78 show that minimizing the lax bisimulation metric corresponds to finding approximate MDP homomorphisms and bound the value error.
3 Value Equivalence Property
To motivate the use of MDP homomorphisms for policy evaluation and consequently policy optimiza- tion, we first prove their value equivalence property in the finite case as the generalization of the prior result of the optimal value equivalence [65], stated in Equation (3). The proof is in Appendix C.I
Theorem (Value Equivalence). Let M be the image of an MDP homomorphism h from finite M. Then any two corresponding policies I1 lifi(T) have equivalent values:
J (s)== =)=(((((s)) VsES,
VsES,aEA
3

---

4 Continuous MDP Homomorphisms
To concretely lay out the foundations of using MDP homomorphisms for continuous control, we extend their definition to continuous state and action spaces, and derive results analogous to the finite case. First, we define continuous MDPs and state our underlying assumptions. Importantly, the correct definitions of continuous MDPs and continuous MDP homomorphisms require care regarding measurability and differentiability of spaces, and our formulation is chosen to fit the HPG derivation; see Appendix A.2 for 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: M (S,E,A Ya E A S [0.1) A
1R,
where S, the state space is assumed to be a Polish space, Σ is a o-algebra on S A. the space of actions, is locally compact metric space, usually taken to be a subset of RR, Ta is the transition probability kernel for each possible action for each fixed s, (-(s)) is probability distribution on E while R the reward function, and y is the discount factor. Furthermore, for all S and B Σ the map - (a BB|s) is smooth. The last assumption is required for differentiability with respect to actions a, which needed in Section for deriving the HPG theorem.
Given the continuous MDPs described above, we define continuous MDP homomorphisms. The cquivariance 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 o-algebra structure on the different state spaces.
Definition 3 (Continuous MDP Homomorphism). A continuous MDP homomorphism is a map h = (f.g.): : .M- + M where f: f:S + S and for every s in S. Is A A are measurable, surjective maps such that the following hold:
Invariance of reward:
(a() = R(s,a)
V
(6)
Equivariance of transitions: (a) =a) (B)ss) S BE
(7)
Note that if Is is the identity map, the second condition reduces to F. (b)) (B)})
which is simply the condition for preservation of transition probabilities as used in bisimulation 20
4.1 Optimal Value Equivalence
3 prove that optimal value functions are preserved by the continuous MDP homomorphism as in the finite case.
Assuming the conditions given in Definition
Theorem 2 (Optimal Value Equivalence). Let M = (S,E,,, Rt) be the image of a continuous MDP homomorphism h (J, from M = (S, E, A, Tay R). Then:
V*(s)=V (f())) Vs (s, (0))===(((((sss))))))))=>=)) (f(s),g((a)) V(s,a) E SxA
The proof, given in Appendix C2 uses the change of variable formula of the pushforward measure of ra (-Is) 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 lifting process to define -' (ift given a policy on the abstract MDP. In general, the lifted policy needs to satisfy the relation n. (g," (3)]s) (z) for every Borel set 3 c A and é S. While our initial progress shows that lifted stochastic policies exist based on the disintegration theorem, the full proof and design of 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- (=)(ss))) If select gs to be bijection, the lifted policy can be uniquely defined as (s)) g. (x)())) The assumption on deterministic policies not limiting, as in general the optimal policy of a given MDP deterministic (9. With this lifting definition, we state and prove the following value equivalence result:
Usually the Borel algebra.

---

Theorem 3 (Value Equivalence for Deterministic Policies), Let M be the image of a continuous MDP homomorphism h (f.gs) from M, then any two deterministic policies T² S A and T - A where nr = -iif((T)) have equivalent value functions on their domain:
(s)= = (F((s)) V (5.0) - (T(()))
V(s,a)
The proof, given in Appendix uses the change of variable formula of the pushforward measure of Ta (-Is) with respect to f change the integration space from S S and assumes 9s to be bijective.
5 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 efficiency. Equipped with continuous MDP homomorphisms from Definition 3 and their value equivalence property, we now derive the homomorphic policy gradiens (HPG) theorem.
In this section, we assume having access to an MDP homomorphism map h= g,), 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 Definition 2 and Appendix B Similarly to prior works on policy gradients [76] 70,, we define the performance measure as (() = Ex [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 ys- 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 n. and obtain the abstract policy as F. 9. (x)(s)) First, show the equivalence of policy gradients:
Theorem 4 (Equivalence of Deterministic Policy Gradients). Let M be the image of continuous MDP homomorphism h from M, and let | S + A be the lifted deterministie policy corresponding the abstract deterministic policy no S A. Then for any (8,a) S x A we have:
V VoQ (s,
(oT =
=V.))=V[[=" V. (3,0) (o)))())) VOT
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 9s to be bijection and local diffeomorphism. Theorem hhighlights 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 0, as shown in Figure
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 m (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 the stationary distribution of the abstract MDP M under To()):
Theorem (Homomorphic Policy Gradient Theorem). Let M be the image of a continuous MDP homomorphism h from M, and let Te S + À be a deterministic abstract policy defined on M. Then the gradient of the performance measure J (0), defined on the actual MDP M, w.r.. 0 is:
VoJ(O)= = S (3)V (B, (=)
V (3)d
where fa* is the discounted state distribution of M following the deterministic policy To(3) (3).
The proof is given in Appendix C and applies the result of Theorem and the change of variables formula of the pushforward measure on the state space. The significance of Theorem 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
5

---

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 655 the abstract critic Q is trained to train Q represents all (s, d) pairs that are equivalent under the MDP homomorphism relation, thus improving sample efficiency. However, the amount of complexity reduction is dependent on the approximate symmetries of the environment, as also supported by our empirical results.
simplified problem. In other words, each abstract state-action pair (8,a) used
6 Homomorphic Actor-Criti Algorithms
We propose deep actor-critic algorithm based on HPG by adapting DDPG [1 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 better study on the impact of MDP homomorphisms. Notably, a strong advantage of DHPG is that it readily applicable to pixel observations without the need for extra mechanisms such
We refer to our
as image reconstruction [29] 34, 90, (6), as the notion of MDP homomorphism provides natural inductive bias for learning representations that preserve values and optimal values.
Since DHPG is learning the MDP homomorphism map / online and concurrently with the policy, using the actual MDP for training the critic, specifically the early stages of training, is helpful. Therefore, we utilize a separate critic for each MDP M and M. Ultimately, critics used to update a single set of parameters, as shown in Figure ! see Appendix D.5 for the ablation study on this.
Thus, the components of DHPG are: actual critic Qu(s, abstract critic (3) a). deterministic actor (s)) homomorphism map h (fe(s). gn (s))),), reward predictor R3) and probabílistic transition model F, 3 3, a which outputs a Gaussian distribution. We use target networks and a vanilla replay buffer [57 (11 As discussed in Section 5 an abstract actor is obtained ā= (Te
In case of pixel observations, single image encoder E,, is shared among all components.
Training the Policy and Critic. Actual and abstract crities are trained using n-step TD error for faster reward propagation Z7. The loss function for each critic is therefore defined as the expectation of the n-step Bellman error estimated over transitions samples from the replay buffer B:
Lactual critic (v²) B(..,,
(r) (R( (1) + +7 +7"Qu" (Stoni atnn) Q (s,, a:)) ]
(9)
Cabstruct critic v.) = E(s.,, .))BB (R + +7" (3 a..n) (3,,
(10)
where s; = (s:) and a= gy (5t, at) are computed using the learned MDP homomorphism, e and T denote parameters of target networks, and = y're-i is the n-step return. Consequently, we train the policy using DPG (70 and HPG from TheoremmSby backpropagating the following loss:
Lacur (0) - Qo To(s)) + (T((((s), (n {.]
(11)
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.
10000
Quadiuped Iolk
Finger Turn hard
0.4
0.30
02
15:95
0.0
0.00
500
6.2
0.6
0.6
(1:60
3.0
8-4 6.0
(a) Sample efficiency.
(b) Performance profiles.
(c) Learning curves.
(d) Learning curves.
Figure 2: Results of DM Control tasks with state observations obtained on seeds. RLiable metrics are aggregated over 17 tasks. (a) RLiable IQM scores as a function of number of steps for comparing sample efficiency, (b) RLiable performance profiles at 500k steps, (c)-(d) examples of learning curves. Full results are in Appendix D.I Shaded regions represent 95% confidence intervals.

---

Learning Continuous MDP Homomorphisms. We now address the problem of learning con- tinuous MDP homomorphisms. While few methods have been proposed for learning finite MDP homomorphisms 81 83 these are not readily extendable to continuous actions. In this work, we use the lax bisimulation metric 778,, Equation ( 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-)
can naturally represent approximate MDP homomorphisms.
Following the same intuition as prior works on bisimulations I9I, we define 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, minimize the distance between these terms:
(ax( n) Rs (s:) fo(s;) Il.- -151 oW2 (F (8t Fr(f)()))) (n),)
(12)
Similarly to Zhang et al. [9T]. we replaced the Kantorovich (W1) metric in Equation (5) with the W2 metric there is an explicit formula for it for Gaussian distributions. Finally, apply the conditions of a continuous MDP homomorphism map from Definition 3 via the loss function of:
L 1, v, = (f (otp:onnnnB((()))) (fo(s)) - x)²² + (r.- (=)((()))s)))))
(13)
where (/) 9n, az)). The final loss function is obtained as Lax (Ø. 11) La( n. p.P.)).
7 Experiments
In our experiments, we aim to answer the following key questions:
1. Does the homomorphic policy gradient improve policy optimization?
2. 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 profiles aggregated on all tasks over 10 random seeds.
While our baseline results are obtained using the official code, when possible 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 the true performance distribution [5].
7.1 State Observations
We compare DHPG on state observations
peddsan requle wocks db
against three commonly-used off-policy model- free algorithms: DDPG [5I], TD3 [28.. and SAC (35). All methods use n-step returns, and share the same hyperparameters presented in Ap- pendix[E22 For fair comparison with DDPG
cc
and a better study of the impact of HPG, we have improved our DDPG by adding delayed policy updates and target policy
Penduuun Angle
(a) Actual optimal policy. (b) Abstract optimal policy.
(28). Thus, the only difference between our DDPG and TD3 is the clipped double Q-learning present in TD3,
Figure Contours of actual and abstract opti- mal actions over the state space of the pendulum- swingup task. Colors represent action values, and states are s (0, (a) Actual optimal policy; contours of optimal actions a* = (x). (b) Ab- stract optimal policy; contours of abstract opti- mal actions ã 9. (a(")) . (3).. The relation Io (a((a)) (az) holds for equivalent state-action pairs, and the abstract optimal policy is symmetric.
which appears to hurting the performance in some tasks of DMC as also observed in (0).
DHPG outperforms or matches other algo- rithms on state observations and has better sample efficiency. Results are presented in Fig- 2 and full results are in Appendix [.7] Ex- pectedly, performance gains are larger on tasks with symmetries, as DHPG is able to learn compressed abstract MDP.
We use the official implementations of DrQv2, DBC, and SAC-AE, while we re-implement DeepMDP due the unavailability of the official code. See Appendix E.3 for full details.
7

---

Pandulr Swhpus
Quadrupd
Desiiga
Number Time Steps (1e4)
48
02
8.4
Try
24
24
Seprident
(a) Sample efficiency.
(b) Performance profiles.
(c) Learning curves.
(d) Learning curves.
Figure 4: Results of DM Control tasks with pixel observations obtained on 10 seeds. RLiable metrics are aggregated over 14 tasks. (a) RLiable IQM scores as a function of number of steps for comparing sample efficiency, (b) RLiable performance profiles 500k steps, (c)-(d) examples of learning curves. Full results are in Appendix D.2 Shaded regions represent 95% confidence intervals.
The learned mapping h= ga) demonstrates properties of MDP homomorphism. We use the pendulum swingup task to visualize its learned MDP homomorphism, as its symmetries are perfectly intelligible. Two state-action pairs (01.6))) a1) and ($2 (82=(02, 82), az) are equivalent if a1 = =22 01 = =-22, and 0: = -62. Therefore. the Iearned action representations are expected to reflect this by setting 9e, (a1) = (a2). Figure shows contours of optimal actions over S, while Figure batree3bjshoo shoshoww action representations == gs (a(a) of optimal actions over S. Clearly, abstract actions adhere the aforementioned relation for equivalent state-action pairs, indicating 9, (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 991] DeepMDD 22)] SAC-AE [90), and state-of-the-art performing DrQ-v2 [8). All methods use ntstep returns, share the same hyperparameters in Appendix E.22 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.
DHPG outperforms or matches other algo- rithms on pixel observations, demonstrating its effectiveness in representation learning. Results are presented in Figure and full results are in Appendix [.2] Interestingly, DHPG with- out image augmentation outperforms DrQ-v2 on domains with easily learnable MDP homo- morphism maps, such as cartpole and pendulum, showing its power of representation learning.
Carpaio
Rpace
Calyabi Kamppe
(a) Trajectorics.
(b) Curves.
(c) Curves.
Figure 5: Effectiveness of DHPG in recovering the minimal MDP from pixels. All methods are lim- ited to a 4-dimensional latent space which is equal to the dimensions pole. (a) Trajectories of real states obtained from Mujoco and trajectories of latent states of DHPG. (b, e) Learning curves averaged on 10 seeds.
DHPG can learn and recover a low- dimensional MDP image. A key strength of MDP homomorphisms is their ability to repre- sent 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 basclines in Figure While other methods are not able to learn the tasks, DHPG can successfully learn the policy and the minimal -ineeneniiooall latent space. Surprisingly, trajectories of the latent states resemble that of the real states as shown in Figure Sa
the real state space of cart-
The abstract MDP demonstrates properties of an MDP homomorphic image. To qualitatively demonstrate the significance of learning joint state-action representations, Figure shows visualiza- tions of latent states for quadruped-wall 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 So has mapped corresponding legs (e.g., left forward leg and right backward leg) to the same abstract latent state (Figure (bh) as they are some homomorphic image of one another. Clearly, DBC and DrQ-v2 are not able to achieve this.
8

---

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 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 fair evaluation, we compare DHPG with higher-capacity variants of DBC and DrQ-v2 that have 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 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 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 Appendices D.5 and D.6 respectively.
8 Related Work
State Abstraction. Bisimulation [56 47 is notion of behavioral equivalence between systems. was extended to continuous state spaces by Blute al. [12 [0 and extended to MDPs by Givan et al. 330 Bisimulation metrics [21 26 25 27 define a pseudometric to quantify the degree of behavioural similarity. Recently, Zhang al. 91 defined loss function for learning representations via bisimilarity of latent states, and Kemertas et al. [46] have further improved its robustness. Castro [4] has proposed method approximate the bisimulation metric for deterministic MDPs with continuous states but discrete actions. van der Pol et al. [8] 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 to directly embed the MDP homomorphic relation in the network architecture that preserve values [33 32 or policies [4), or via a sampling-basee -based similarity metric [15]. Finally,
82 Other recently proposed methods seek learn representations
state abstractions can principle help improve transferring policies temporally extended actions
or learning
86 77].
Action Abstraction. Action representations often studied in the context of large discrete action spaces [68] as a form a look-up embedding that is known a-priori [22], factored representations [69].
or policy decomposition [18). Action representations also be learned from expert demonstrations 79). More related to our work is dynamics-aware embeddings g5] where a combined state-action
120
(a) Latent states, DHPG. (b) Abstract states, DHPG. (c) Latent states, DBC. (d) Latent States, DrQ-v2.
Figure 6: PCA projection of learned representations for quadruped-walk with pixel observations. (a) Latent states s= E (o), (b) abstract latent states 3= (o), for DHPG, (c) latent states s= (o) for DBC, and (d) DrQ-v2. Color of each point denotes its value learned by Q(s, a) or Q(x,a), a). Points are also projected onto each main plane. The homomorphism map DHPG has mapped the latent states of corresponding legs (c.g.. left forward leg and right backward leg) (a) on to the same abstract latent states (b), indicating clear structure in S.

---

Pandllr Swhgus
Quadruped
0.8
0.6
0.4
0.2
Number Time Steps (1e4)
Normataed
(a) Sample efficiency.
(b) Performance profiles.
(c) Learning curves.
(d) Learning curves.
Figure 7: Results of DM Control tasks with pixel observations for higher-capacity variants of DBC and DrQ-v2 obtained on 10 seeds. RLiable metrics are aggregated over 14 tasks. (a) RLiable IQM scores function of number of steps for comparing sample efficiency, (h) RLiable performance profiles at 500k steps, (c)-(d) examples of learning curves. Full results are in Appendix D.7 Shaded regions represent 95% confidence intervals.
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 22
36
37
34 10], auxiliary prediction tasks 44 53 54]
physics-inspired inductive biases 45
19
unsupervised learning [42] 52 or self-supervised
learning 66 71 38
39 24 From another point of view, representation learning be effectively Symmetries of the environment can also be used for
decoupled from the RL problem [2 74
62 [3]. In fact, MDP homomorphisms are special- izations of such approaches for RL. A key distinguishing factor of MDP homomorphisms is their ability to take actions into account for representation learning the same premises as Thomas et al. [8) Recently, simple image augmentation methods have shown significant improvements in RL performance [89 48). Since these approaches are in general orthogonal state abstractions, they can be combined together.
representation learning (58) 55
61
84 40
41
Conclusion
In this paper, 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 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 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.
References
[1] David Abel, Dilip Arumugam, Kavosh Asadi, Yuu Jinnai, Michael L. Littman, and Lawson LS Wong. State abstraction as compression in apprenticeship learning. In Proceedings of the AAAI Conference Artificial Intelligence, volume 33, pages 3134-3142, 2019.
10

---

[2] David Abel, David Hershkowitz, and Michael Littman. Near optimal behavior via approximate state abstraction. In International Conference on Machine Learning, pages 2915-2922 2923, PMLR, 2016.
[3] David Abel, Nate Umbanhowar, Khimya Khetarpal, Dilip Arumugam, Doina Precup, and Michael Littman. Value preserving state-action abstractions. In International Conference on Artificial Intelligence and Statistics, pages 1639-1650. PMLR, 2020.
[4] Rishabh Agarwal, Marlos Machado, Pablo Samuel Castro, and Marc Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. In International Conference on Learning Representations, 2020.
[5] Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Belle- mare. Deep reinforcement learning at the edge of the statistical precipice. Advances in Neural Information Processing Systems, 34, 2021.
[6] Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Marc-Alexandre Côté, and R Devon Hjelm. Unsupervised state representation learning Atari. Advances in Neural Information Processing Systems, 32:8769-877 8782, 2019.
[7] Gabriel Barth-Maron, Matthew W Hoffman, David Budden, Will Dabney, Dan Horgan, Dhruva Tb, Alistair Muldal, Nicolas Heess, and Timothy Lillicrap. Distributed distributional determin- istic policy gradients. arXiv preprint 1xXiv:104.08611 2018.
[8] Richard Bellman. Dynamic Programming. Princeton University Press, 1957.
[9] Dimitri Bertsekas. Dynamic programming and optimal control: Volume I, volume 1. Athena scientific, 2012.
[10] Biza, R Platt, JW van de Meent, and L Wong. Learning discrete state abstractions with deep variational inference. Advances in Approximate Bayesian Inference, 2021.
[11] Ondrej Biza and Robert Platt. Online abstraction with mdp homomorphisms for deep learning. In Proceedings of the 18th International Conference on Autonomous Agents and MultiAgent Systems, 2019.
[12] Richard Blute, Josée Desharnais, Abbas Edalat, and Prakash Panangaden. Bisimulation for labelled markov processes. In Proceedings of Twelfth Annual IEEE Symposium on Logic in Computer Science, pages 149-158. IEEE, 1997.
[13] Hugo Caselles-Dupré, Michael Garcia Ortiz, and David Filliat. Symmetry- based disentangled representation learning requires interaction with environments. Advances in Neural Information Processing Systems, 32:4606-4615, 2019.
[14] Pablo Samuel Castro. Scalable methods for computing state similarity in deterministic markov decision processes. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 10069-10076. 2020.
[15] Pablo Samuel Castro, Tyler Kastner, Prakash Panangaden, and Mark Rowland. Mico: Improved representations via sampling-basse based state similarity for Markov decision processes. Advances in Neural Information Processing Systems, 34, 2021.
[16] Pablo Samuel Castro and Doina Precup. Using bisimulation for policy transfer in MDPs. In Twenty-Fourth AAAI Conference on Artificial Intelligence, 2010.
[17] Pablo Samuel Castro and Doina Precup. Automatic construction of temporally extended actions for MDPs using bisimulation metries. In European Workshop on Reinforcement Learning, pages 140-152. Springer, 2011.
[18) Yash Chandak, Georgios Theocharous, James Kostas, Scott Jordan, and Philip Thomas. Learning action representations for reinforcement learning. In International Conference on Machine Learning, pages 941-950. PMLR, 2019.
[19] Miles Cranmer, Sam Greydanus, Stephan Hoyer, Peter Battaglia, David Spergel, and Shirley Ho. Lagrangian neural networks. arXiv preprint arXiv:2003.046630 2020.
[20] J. Desharnais, A. Edalat, and P. Panangaden. Bisimulation for labeled Markov processes. Information and Computation, 179(2):663-193, Dec 2002.
[21] J. Desharnais, V. Gupta, R. Jagadeesan, and P. Panangaden. Metrics for labeled Markov systems. In Proceedings of CONCUR99, number 1664 in Lecture Notes Computer Science. Springer- Vprilger-Verllg, 1999.

---

[22] Gabriel Dulac-Amoki Richard Evans, Hado van Hasselt. Peter Sunehag, Timothy Lillicrap. Jonathan Hunt, Timothy Mann, Theophane Weber, Thomas Degris, and Ben Coppin. Deep reinforcement learning in large discrete action spaces. arXiv preprint arXiv: 15l..07679, 2015.
[23] SM Ali Eslami, Danilo Jimenez Rezende, Frederic Besse, Fabio Viola, Ari S Morcos, Marta Garnelo, Avraham Ruderman, Andrei A Rusu, Ivo Danihelka, Karol Gregor, et al. Neural scene representation and rendering. Science, 16046394) 2018.
[24] Linxi Fan, Guanzhi Wang. De-An Huang, Zhiding Yu, Li Fei-Fei, Yuke Zhu, and Animashree Anandkumar. Secant: Self-expert cloning for zero-shot generalization of visual policies. In International Conference on Machine Learning, pages 3088-3099. PMIR, 2021.
[25] Norm Ferns, Pablo Samuel Castro, Doina Precup, and Prakash Panangaden. Methods for computing state similarity in Markov decision processes. In Proceedings of the Twenty-Second Conference on Uncertainty in Artificial Intelligence, pages 174-181, 2006.
[26] Norm Ferns, Prakash Panangaden, and Doina Precup. Metrics for Markov decision processes with infinite state spaces. Proceedings of the Twenty-Fine Conference on Uncertainty in Artificial Intelligence, pages 201-208, 2005.
[27] Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous Markov decision processes. SIAM Journal on Computing, 10(6):1662-11) 2011.
[28] Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In International Conference on Machine Learning, pages 1587-1596. PMLR, 2018.
[29] Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare G Bellemare. Deep- mdp: Learning continuous latent space models for representation learning. In International Conference on Machine Learning, pages 2170-2179, PMLR. 2019.
[30] Robert Givan, Thomas Dean, and Matthew Greig. Equivalence notions and model minimization in Markov decision processes. Artificial Intelligence, 147(1-2):163-223 2003.
[31] Samuel Greydamus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks. Advances in neural information processing systems, 32, 2019.
[32] Christopher Grimm, André Barreto, Gregory Farquhar, David Silver, and Satinder Singh. Proper value equivalence. arXiv preprint arXiv:2106. 10316, 2021.
[33] Christopher Grimm, André Barreto, Satinder Singh, and David Silver. The value equivalence principle for model-based reinforcement learning. arXiv preprint arXiv:2011.035000 2020.
[34] David Ha and Jürgen Schmidhuber. World models. arXiv e-prints, pages arXiv-1803, 2018.
[35] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pages 1861-1870. PMLR, 2018
[36] Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representa- tions, 2019.
[37] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pages 2555-2565. PMLR, 2019.
[38] Nicklas Hansen, Rishabh Jangir, Yu Sun, Guillem Alenyå, Pieter Abbeel, Alexei A Efros, Lerrel Pinto, and Xiaolong Wang. supervidd policy adaptation during deployment. In International Conference on Learning Representations, 2020.
[39] Nicklas Hansen and Xiaolong Wang. Generalization reinforcement learning by soft data augmentation. In 2021 IEFE International Conference on Robotics and Automation (ICRA),
pages 13611 13611- IEEE, 2021.
[40] Irina Higgins, David Amos, David Pfau, Sebastien Racaniere, Loie Mathhey, Danilo Rezende, and Alexander Lerchner. Towards a definition of disentangled representations. arXiv preprint 1rXiv:1812.02223. 2018.
[41] Irina Higgins, Peter Wirnsberger, Andrew Jaegle, and Aleksandar Botev. Symetric: Measuring the quality of learnt hamiltonian dynamics inferred from vision. Advances in Neural Information Processing Systems, 34, 2021.
12

---

[42] R Devon Hjelm, Alex Fedorov, Samuel Lavoie-Marchilee Marchildon, Karan Grewal, Phil Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. In International Conference on Learning Representations, 2018.
[43] Dan Horgan, John Quan, David Budden, Gabriel Martn, Matteo Hessel, Hado van Hasselt, and David Silver. Distributed prioritized experience replay. In International Conference Learning Representations, 2018.
[44] Max Jaderberg. Volodymyr Mnih. Wojeiech Marian Czarnecki, Tom Schaul. Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint arXiv:1611.05339 2016.
[45] Rico Jonschkowski and Oliver Brock. Learning state representations with robotic priors. Autonomous Robots, 39(3):00744777 2015.
[46] Mete Kemertas and Tristan Aumentado-Armstronnn Towards robust bisimulation metrie learning. Advances in Neural Information Processing Systems, 34, 2021.
[47] Kim G Larsen and Ame Skou. Bisimulation through probabilistic testing. Information and computation, 94(1) 94(1):1-28, 1991.
[48] Kimin Lee, Kibok Lee, Jinwoo Shin, and Honglak Lee. Network randomization: A simple technique for generalization in deep reinforcement learning. arXiv preprint arXiv:1910.05336. 2019.
[49] Lihong Li, Thomas J Walsh, and Michael Littman. Towards unified theory of state abstraction for MDPs. ISAIM, 4:5, 2006.
[50]
Qing Li. Continuous control benchmark of DeepMind control suite and MuJoCo. https 2021.
//gi wtw.hh....omweee//eeee/ea
[51] Timothy P Lillicrap. Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning, arXiv preprint arXiv:/509.023: 02971, 2015.
[52] Hao Liu and Pieter Abbeel. Aps: Active pretraining with successor features. In International Conference on Machine Learning, pages 6736-6744 6747. PMLR, 2021.
[53] Shikun Liu, Andrew Davison, and Edward Johns. Self-supervised generalisation with meta auxiliary learning. Advances in Neural Information Processing Systems, 32, 2019.
[54] Clare Lyle, Mark Rowland, Georg Ostrovski, and Will Dabney. On the effect of auxiliary tasks on representation dynamics. In International Conference on Artificial Intelligence and Statistics, pages 1-9. PMLR, 2021.
[55] Anuj Mahajan and Theja Tulabandhula. Symmetry learning for function approximation in reinforcement learning. arXiv preprint arXiv:1706.02990 2017.
[56] Robin Milner. Communication and concurrency, volume 84. Prentice hall Englewood Cliffs, 1989.
[57] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing Atari with deep reinforcement learning. arXiv preprint arXiv: 1322.660, 2013.
[58] Amab Kumar Mondal, Vineet Jain, Kaleem Siddiqi, and Siamak Ravanbakhsh. Eqr: Equivariant representations for data-efficient reinforcement learning. In International Conference on Machine Learning, pages 15908-15926. PMLR, 2022.
[59] Shravan Matthur Narayanamurthy and Balaraman Ravindran. On the hardness of finding symmetries in markov decision processes. In Proceedings of the 25th international conference on Machine learning, pages 688-695, 2008
[60] Fabio Pardo. Tonic: A deep reinforcement learning library for fast prototyping and benchmark- ing. arXiv preprint arX1v:2011.07537, 2020.
[61] Jung Yeon Park, Ondrej Biza, Linfeng Zhao, Jan Willem van de Meent, and Robin Wal- ters. Learning symmetric embeddings for equivariant world models. arXiv preprint
arXiv:2204.11377 2022.
[62] Robin Quessard, Thomas D Barrett, and William R Clements. Learning group structure and disentangled representations of dynamical environments. arXiv preprint (tl:22002.0699 2020.
13

---

[63] Srividhya Rajendran and Manfred Huber. Learning to generalize and reuse skills using approxi- mate partial policy homomorphisms. In 2009 IEEE International Conference on Systems, Man and Cybernetics, pages 2239-222200 2244. IEFF, 2009.
[64] Balaraman Ravindran. An algebraic approach to abstraction in reinforcement learning. Univer- sity of Massachusetts Amherst, 2004.
[65] Balaraman Ravindran and Andrew G Barto. Symmetries and model minimization in markov decision processes, 2001.
[66] Balaraman Ravindran and Andrew G Barto, Relativized options: mation. In Proceedings of the 20th International Conference on Machine Learning (ICML-03), pages 608-615, 2003.
Choosing the right transfor-
[67] Balaraman Ravindran and Andrew G Barto. Approximate homomorphisms: A framework for non-exact minimization in Markov Decision Processes, 2004.
[68] Brian Sallans and Geoffrey E Hinton. Reinforcement learning with factored states and actions. The Journal of Machine Learning Research, 1:88, 2004.
[69] Sahil Sharma, Aravind Suresh, Rahul Ramesh, and Balaraman Ravindran. Learning to factor policies and action-value functions: Factored action space representations for deep reinforce- ment learning. arXiv preprint arXiv: (00.::70007000 705.07269, 2017.
[70] David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. In International conference machine learning, pages 387-395. PMLR, 2014.
[71] Samarth Sinha, Ajay Mandlekar, and Animesh Garg. S4rl: Surprisingly simple self-supervision for offline reinforcement learning in robotics. In 5th Annual Conference on Robot Learning, 2021.
[72] Vishal Soni and Satinder Singh. Using homomorphisms to transfer options across continuous reinforcement learning domains. In AAAI, volume 6, pages 494-499, 2006.
[73] Jonathan Sorg and Satinder Singh. Transfer via soft homomorphisms. In Proceedings of The 8th International Conference on Autonomous Agents and Multiagent Systems- Systems-Volume 2. pages 741-748, 2009.
[74] Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In International Conference on Machine Learning, pages 9870-9879, PMLR, 2021.
[75] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[76] Richard S Sutton, David A McAllester, Satinder P Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. In Advances in neural information processing systems, pages 1057-1063, 2000.
[77] Richard S Sutton, Doina Precup, and Satinder Singh. Between MDPs and semi-MDPs: A framework for temporal abstraction in reinforcement learning. Artificial intelligence, 1 112(1- 2):181-211, 1999.
[78] Jonathan Taylor, Doina Precup, and Prakash Panagaden. Bounding performance loss in approx- imate mdp homomorphisms. Advances in Neural Information Processing Systems, 21:1649- 1656, 2008
[79] Guy Tennenholtz and Shie Mannor. The natural language of actions. In International Conference on Machine Learning, pages 6196-6200 PMLR, 2019.
[80] Valentin Thomas, Jules Pondard, Emmanuel Bengio, Mare Sarfati, Philippe Beaudoin, Marie- Jean Meurs, Joelle Pineau, Doina Precup, and Yoshua Bengio. Independently controllable factors. arXiv preprint arXiv: 01289, 2017.
[81] Elise van der Pol, Thomas Kipf, Frans A Oliehoek, and Max Welling. Plannable approximations to mdp homomorphisms: Equivariance under actions. In Proceedings of the 19th International Conference on Autonomous Agents and MultiAgent Systems, pages 1431-1439, 2020.
[82] Elise van der Pol, Herke van Hoof, Frans A Oliehoek, and Max Welling. Multi-agent MDP homomorphic networks. arXiv preprint 27X.v:2110.04495, 2021.
14

---

[83] Elise van der Pol, Daniel Worrall, Herke van Hoof, Frans Oliehoek, and Max Welling. Mdp homomorphic networks: Group symmetries reinforcement learning. Advances in Neural Information Processing Systems, 33, 2020.
[84] Dian Wang, Robin Walters, and Robert Platt. So(2)-cquivariant reinforcement learning. In International Conference Learning Reprexentations, 2021.
[85] William Whitney, Rajat Agarwal, Kyunghyun Cho, and Abhinav Gupta. Dynamics-aware embeddings. In International Conference on Learning Representations. 2019.
[86] Alicia Wolfe and Andrew G Barto. Decision tree methods for finding reusable mdp ho- momorphisms. In PROCEEDINGS OF THE NATIONAL. CONFERENCE ON ARTIFICIAL INTELLIGENCE, volume 21, page 530, Menlo Park, CA; Cambridge, MA; London; AAAI Press; MIT Press; 1999, 2006.
[87] Alicia Peregrin Wolfe and Andrew G Barto. Defining object types and options using mdp homomorphisms. In Proceedings of the ICML-06 Workshop on Structural Knowledge Transfer for Machine Learning. Citeseer, 2006.
[88] Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2107.09645. 2021.
[89] Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regu- larizing deep reinforcement learning from pixels. In International Conference Learning Representations, 2020.
[90] Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improving sample efficiency in model-free reinforcement learning from images. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 10674-10681, 2021.
[91] Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In Inter- national Conference on Learning Representations, 2020.
15

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes] All the theoretical claims are substantiated in Sections
All the empirical claims are substantiated in Section a Appendix [ (b) Did you describe the limitations of your work? [Yes] Theoretical limitations and assumptions are discussed in Section 4 and empirical limitations are discussed in
and
5
Section 7
(c) Did you discuss any potential negative societal impacts of your work? [No] Our work foundational research on reinforcement learning and state abstraction, hence we do not foresee any substantive societal and ethical implications.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of theoretical results? [Yes] The full set of assumptions are detailed in Section 4 and Appendix B
(b) Did you include complete proof's of all theoretical results? [Yes] The proofs of all theoretical results are given in Appendix
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main exper- imental results (either in the supplemental material
URL)? [Yes] Our code, including the instructions, are submitted in the supplemental material.
(b) Did you specify all the training details (e.g.. data splits, hyperparameters, how they were chosen)? [es]] Hyperparameters, implementation, and training details are given in Appendix E
(c) Did you report error bars (e.g., with respect the random seed after running experi- ments multiple times)? [Yes] All experiments are obtained on 10 seeds, and we report confidence intervals, interquartile mean, median, mean, and performance profiles as suggested by Agarwal et al. [5]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] We describe our hardware setup in Appendix E. but we do not include the amount of compute used.
4. If you are using existing assets (c.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] The creators are cited appropriately and information of each asset is included in Appendix E.3
(b) Did you mention the license of the assets? [N/A] All assets are open-source and under permissive open-source licenses.
(c) Did you include any new assets either in the supplemental material or as URL? [Yes] We include the code of our algorithm in the supplemental material.
(d) Did you discuss whether and how consent was obtained from people whose data you're rou're using/curating? [N/A] We are using open open-source simulators and reinforcement learning environments.
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A] We are using open-source simulators and reinforcement learning environments.
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
16


Reviews:
[{"review_id": "y1Y13CW1c6p", "reviewer": "Reviewer_8n1B", "summary": "This paper contributes a) a continuous version of the MDP homomorphism and b) a homomorphic policy gradient. An approach based on lax bisimulation metrics learns a policy and a continuous MDP homomorphism simultaneously. The paper evaluates on DM control tasks from pixel observations and shows improvements compared to reasonable baselines (DrQ-v2, DBC, DeepMDP, etc). ", "questions": "- Could you elaborate on what the effect is of requiring $g_s$ to be a bijection?\n- In Figure 3a, the learned policy looks approximately, but not exactly symmetric. The original problem does have exact symmetry, and in Fig 3b the action representations look like they may have recovered the symmetry. Why is the policy less symmetric than the representations are? (This Figure is hard to read. It would be improved by clarifying which of the plotted values represent the chosen actions)\n- The paper claims that legs are mapped onto each other in Figure 6a/b. Could you elaborate on how one can read this in the Figure and how you know that legs are represented in the abstract/latent state space structure?\n- Suggested references:\n    - Park, Jung Yeon, et al. “Learning Symmetric Embeddings for Equivariant World Models.” ICML 2020 \n    - Narayanamurthy, Shravan Matthur, and Balaraman Ravindran. “On the hardness of finding symmetries in Markov decision processes.” ICML 2008.\n- Appendix C.1 says that $R$ and $\\bar{R}$ are both bounded. Where does this assumption come from? \n- Page 2, line 50-51 claims that the proposed approach leads to more robust solutions. The approach seems similar to an MDP homomorphic version of DBC. Why does the approach outperform DBC?\n", "limitations": "- The checklist says empirical limitations are discussed in Section 7, but I could not find this. Could you point me towards the paragraphs?\n- Section 4 discusses part of the theoretical assumptions/limitations.\n- Code was included with submission. \n", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n- The paper tackles an important problem: learning representations for reinforcement learning problems, especially those that recover structure and/or symmetry in the input MDP.\n- The approach, of mapping into an abstract space and learning a policy + abstraction simultaneously is reasonable and has precedent. \n- According to the paper, the approach is able to learn in smaller latent dimensions and recover symmetric representations, which is a nice result. Discovering symmetries in MDPs is a largely unsolved problem, and the paper takes a step in that direction.\n- It's useful to have an extension to continuous MDP homomorphisms and to have the homomorphic policy gradient theorem. The paper could be of interest to the field of state representation learning and to the field of symmetries in RL. \n\n# Weaknesses\n- Usually in MDP Homomorphisms, $g_s$ is a surjection. In the paper, it is required to be a bijection. This comes with the disadvantage that equivalent actions belonging to a single state can no longer be aggregated/mapped onto the same abstract action. For symmetries this should not be an issue, but making $g_s$ bijective limits the levels of aggregation in action space.\n- The work is evaluated on DMControl tasks. One of the baselines is DBC. In the original DBC paper and follow-up work (e.g. [42]), there are experiments with background distraction videos. Including such experiments in this work, to show robustness to distractors that are irrelevant to the task, would make the paper stronger. \n- The potential of this approach to find symmetries in MDPs is part of its appeal. The paper does allude to this by including plots that seem to show some degree of symmetry recovery. However, the direction is not explored in-depth in the paper. \n- Missing references:\n    - Mahajan, Anuj, and Theja Tulabandhula. “Symmetry learning for function approximation in reinforcement learning.” arXiv preprint arXiv:1706.02999 (2017).\n    - Biza, Ondrej, and Robert Platt. “Online abstraction with MDP homomorphisms for deep learning” AAMAS 2019\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Z6yzKOjxkm", "reviewer": "Reviewer_EAfi", "summary": "The paper extends the theory of MDP homomorphism abstractions to continuous state and action spaces, and includes a theorem that motivates the use of MDP homomorphisms in actor-critic reinforcement learning agents. The paper proposes an agent that learns an abstraction of the ground environment while learning to solve a particular task. The agent is evaluated on continuous control tasks.\n", "questions": "1. The abstraction components (the state and state-action encoders f and g, the abstract critic \n\\bar{Q}, the abstract transition predictor \\tau and the reward predictor \\bar{R}) add additional capacity to your model . What is the number of FLOPS (or the number of parameters if it is difficult to measure FLOPS) in these components compared to a basic actor-critic model? Could the model size be controlled for in the experiments to make the comparison against the baseline fairer?\n2. The abstract encoders f(s) and g(s, a) could capture similar information, since they both have access to the state. In the extreme, the model could use f(s) for reward prediction only and g(s, a) for everything else, since all functions except for the reward predictor are conditioned on the outputs of both f and g. Do f and g output similar features?\n3. Could you clarify lines 305-307 and Figure 6? Please see my comment above.\n4. What is the size of the state and the state-action encoding in Figure 5.\n", "limitations": "Limitations are not explicitly discussed, but shoud be similar to previous works such as DBC and DDPG.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths:\n\n1. The paper extends the theory of MDP homomorphisms to continuous action spaces (in both the ground and the abstract MDP), which I consider to be a significant contribution.\n2. The paper aligns with prior literature on using bisimulation metrics in model-free learning by using the lax bisimulation metric, which includes action abstractions. I find this to be well-motivated.\n3. The proposed method is evaluated on many continuous control tasks, the high variance in the results of deep reinforcement learning agents is controlled for and the source code for the method is included.\n\n# Weaknesses:\n\n1. There is a possibility that the quantitative benefits we see in the proposed method are due to increased model capacity only. I could not find a discussion of the sizes of the evaluated models.\n2. Having an abstraction component should provide additional benefits, such as compactness of representations and robustness to distracting features. These benefits are demonstrated to a limited extent, and some details of the experiments are unclear (please see my detailed comments).\n\n# Detailed comments:\n\n## Writing\n\n– This paper is close to previous works on using bisimulation metrics in model-free reinforcement learning. I am missing a discussion of the similarities between the proposed model and prior works (like DBC [10]) and a discussion of the benefits of adding action abstraction to this framework.\n\n– The literature review is insufficient, many prior works on MDP homomorphisms are missing [1, 2, 3, 4, 5, 6]. Additionally, there are a few more works on bisimulations that could be of interest [7, 8, 9]. This list of papers I provided is non-exhaustive.\n\n– Minor comment: This might come down to personal preference, but I do not like the way the introduction is written. The introduction should tell a story and build interest in the rest of the paper. Instead, the introduction in this paper reads like an extended list of contributions.\n\n– Minor comment: The authors cite one of Ravindran’s original papers on MDP homomorphisms, but not his thesis, which I would consider to be the definitive source [11]. \n\n## Experiments\n\n– Rather than an increase in agent performance, which could be attributed to the increase in its capacity, I see the main benefit of state and action abstractions in (a) learning compact representations, (b) ignoring distractions and (c) facilitating better transfer of the learned representation. (a) is shown in the cartpole environment and I find this experiment very compelling. My only complaint is that I could not find what the size of the state-action encoding is (i.e. we could make the state encoding small, but capture all information in the state-action encoding). (b) Is not included in the paper and I find this to be a missed opportunity. For example, DBC [10], which is a baseline in this paper, shows convincing experiments with distracting backgrounds in their Figures 3 and 4. (c) is shown in the appendix, but I am not sure how to interpret the results of the experiment.\n\n– Lines 305 - 307: Perhaps I misunderstand, but the claim that DHPG “mapped all legs to the same abstract latent state” is problematic. Firstly, I do not know if Figure 6b supports this claim. There are four clusters in Figures 6a and 6b, and they look similar. Secondly, it seems unlikely that having a single abstract latent state to track the state of all four legs is desirable – the legs have to move independently to facilitate locomotion. I suppose there is some division of what is represented in the state and the state-action encoder, but that is not explained in the paper.\n\n## References\n\n[1] Vishal Soni and Satinder Singh. 2006. Using Homomorphisms to Transfer Options Across Continuous Reinforcement Learning Domains. In Proceedings of the 21st National Conference on Artificial Intelligence - Volume 1 (AAAI’06). AAAI Press, 494–499.\n\n[2] Alicia P. Wolfe. 2006. Defining Object Types and Options Using MDP Homomorphisms.\n\n[3] Alicia Peregrin Wolfe and Andrew G. Barto. 2006. Decision Tree Methods for Finding \nReusable MDP Homomorphisms. In Proceedings of the 21st National Conference on Artificial Intelligence - Volume 1 (AAAI’06). AAAI Press, 530–535.\n\n[4] Jonathan Sorg and Satinder Singh. 2009. Transfer via Soft Homomorphisms. In Proceedings of The 8th International Conference on Autonomous Agents and Multiagent Systems - Volume 2 (AAMAS ’09). International Foundation for Autonomous Agents and Multiagent Systems, Richland, SC, 741–748.\n\n[5] S. Rajendran and M. Huber. 2009. Learning to generalize and reuse skills using\napproximate partial policy homomorphisms. In 2009 IEEE International Conference\non Systems, Man and Cybernetics. 2239–2244. \n\n[6] Ondrej Biza and Robert Platt. 2019. Online abstraction with MDP homomorphisms for deep learning. In International Conference on Autonomous Agents and MultiAgent Systems.\n\n[7] Abel, D.; Hershkowitz, D. E.; and Littman, M. L. 2016. Near Optimal Behavior via Approximate State Abstraction. In Proceedings of the International Conference on Machine Learning, 2915–2923.\n\n[8] Abel, D.; Arumugam, D.; Asadi, K.; Jinnai, Y.; Littman, M. L.; and Wong, L. L. S. 2019. State Abstraction as Compression in Apprenticeship Learning. In AAAI.\n\n[9] Ondrej Biza, Robert Platt, Jan-Willem van de Meent, and Lawson LS Wong. 2020. Learning Discrete State Abstractions with Deep Variational Inference. In Proceedings of the 3rd Symposium on Advances in Approximate Bayesian Inference.\n\n[10] https://arxiv.org/abs/2006.10742 \n\n[11] https://all.cs.umass.edu/pubs/2004/ravindran_thesis04.pdf ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "EtM8GorxOD", "reviewer": "Reviewer_c7GW", "summary": "This paper proposes MDP homomorphisms and homomorphic policy gradient for continuous control tasks. The contributions are threefold: 1. introducing the continuous MDP homomorphisms, 2. deriving the homomorphic policy gradient theorem that ties the induced abstract MDP into the policy gradient, 3. developing the deep homomorphic policy gradient algorithm which simutaneously learn the policy and homomorphism.", "questions": "1. In Line 165, the authors mentioned \"we assume having access to an MDP homomorphism map ... . The problem of learning such mapping from samples is addressed in Section 6\". Would the derived policy gradient and the policy optimisation algorithm depend on an accurate MDP homomorphism?\n2. In addition to the first question, what happens to Theorem 4 and 5 if the MDP homomorphism is not an exact equality (as in Equations (6) and (7) where there is an epsilon difference) [1]?\n\n[1] Abel, David, David Hershkowitz, and Michael Littman. \"Near optimal behavior via approximate state abstraction.\" International Conference on Machine Learning. PMLR, 2016.\n", "limitations": "The authors clearly discussed the limitation of assuming deterministic policies in Section 4.2.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "1. The introduction of MDP homomorphisms to continuous control tasks is interesting and significant.\n2. The paper is overall well written, with clear motivations and theoretical explainations. \n3. The authors agree to make their code publicly available and has provided the source code in the appendix.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept