A Simple Approach for State-Action Abstraction using a Learned MDP Homomorphism
Authors: Augustine N. Mavor-Parker
Matthew James Sargent
Andrea Banino
Lewis Griffin
Caswell Barry
Venue: Submitted to ICLR 2023 (2023)

Abstract: Animals are able to rapidly infer from limited experience when sets of state action pairs have equivalent reward and transition dynamics. On the other hand, modern reinforcement learning systems must painstakingly learn through trial and error that sets of state action pairs are value equivalent---requiring an often prohibitively large amount of samples from their environment. MDP homomorphisms have been proposed that reduce the observed MDP of an environment to an abstract MDP, which can enable more sample efficient policy learning. Consequently, impressive improvements in sample efficiency have been achieved when a suitable MDP homomorphism can be constructed a priori---usually by exploiting a practioner's knowledge of environment symmetries. We propose a novel approach to constructing a homomorphism in discrete action spaces, which uses a partial model of environment dynamics to infer which state action pairs lead to the same state---reducing the size of the state-action space by a factor equal to the cardinality of the action space. We call this method equivalent effect abstraction. We demonstrate empirically that equivalent effect abstraction can improve sample efficiency in a model-free setting and planning efficiency for model based approaches.

Under review as a conference paper at ICLR 2023
A SIMPLE APPROACH FOR STATE-ACTION ABSTRAC- TION USING A LEARNED MDP HOMOMORPHISM
Anonymous authors Paper under double-blind review
ABSTRACT
Animals are able to rapidly infer, from limited experience, when sets of state-action pairs have equivalent reward and transition dynamics. On the other hand, modern reinforcement learning systems must painstakingly learn through trial and error that sets of state-action pairs are value equivalent- reqquriig an often prohibitively large amount of samples from their environment. MDP homomorphisms have been proposed that reduce the observed MDP of an environment to an abstract MDP, which can enable more sample efficient policy learning. Consequently, impressive improvements in sample efficiency have been achieved when suitable MDP homomorphism can be constructed priori-usually by exploiting a prac- titioner's knowledge of environment symmetries. We propose novel approach to constructing a homomorphism in discrete action spaces, which uses a partial model of environment dynamics to infer which state-action pairs lead to the same state reduuing the size of the state-action space by a factor equal to the cardinality of the action space. We call this method equivalent effect abstraction. On MDP homomorphism benchmarks, we demonstrate improved sample efficiency over pre- vious attempts to learn MDP homomorphisms, while achieving comparable sample efficiency to approaches that rely on prior knowledge of environment symmetries.
INTRODUCTION
Modern reinforcement learning (RL) agents outperform humans in previously impregnable bench- marks such as Go (Silver et al., 2016) and Starcraft (Vinyals et al., 2019). However, the computational expense of RL hinders its deployment in promising real world applications. In environments with large state spaces, RL agents demand hundreds of millions of samples (or even hundreds of billions) to Iearn a policy- either within an environment model (Hafner et al., 2020) or by direct interaction (Mnih et al., 2013). Function approximation can enable some generalisation within large state space but still most RL agents struggle to extrapolate value judgements to equivalent states.
In contrast, animals can easily abstract away details about states that do not effect their values. For example, a foraging mouse understands that approaching a goal state (shown as a piece of cheese in 1e-ure-l--hhil travelling east will have the same value approaching the same goal state-from the west. These sort of state abstractions have been defined in RL as Markov decision process (MDP) homomorphisms Ravindran & Barto (2001); van der Pol et al. (2020b). MDP homomorphisms reduce large state-action spaces to smaller abstract state-action spaces by collapsing equivalent state-action ptars--tate- actions an observed MDP onto a single abstract state-action pair in an abstract MDP van der Pol et al. (2020b).
Given a mapping between an abstract MDP and an experienced MDP, policies can be learned effi- ciently in the smaller abstract space and then mapped back to the experienced MDP when interacting with the environment (Ravindran & Barto, 2001). However, obtaining mappings to fnd-from the ab- stract atate-action space is challenging. Suecess has been achieved by hard coding homemorphisms inte inte- peticy netwerks- Prevvosss-rreviuu works hard code homomorphisms into a policy (van der Pol et al., 2020b) but learning homomorphic mappings online is an unsolved problem.
We develop equivalent effect abstraction, a method that constructs MDP homomorphisms from experience via dynamics model modee--leverggng the fact that state-action pairs leading to the same next state usually have equivalent value. Consider the example of a gridworld maze, moving to a given cell has the same value whether you approached from the right the left. Therefore, if we know the

---

Under review as a conference paper at ICLR 2023
S
Sa
@(5, -) = QS5,, -)
Figure 1: State-action pairs that lead to the same state often have equivalent values (shown by the purple arrows). Instead of learning these equivalent values individually through reinforcement, instead learn the value of one abstract action that represents both purple arrows. These values are looked up and learned with a backwards dynamics model during training-meaning no prior knowledge is required from practitioner.
value for approaching from the right then we also know the value of moving the approaching from the left. By extrapolating value judgements between equivalent state-action pairs we can use this equivalence reduce the amount of experience required learn policy.
An important distinction from previous works is that we do not use environment symmetries to reduce the size of the state-action space. We exploit separate redundancy common to many MDPs-for given state there are often multiple state-action pairs that also lead to that state. While we do use partial model in our approach, equivalent effect abstraction is different to model based RL because we focus on reducing the size of the state-action space rather than augmenting experience with predicted trajectories. Additionally, unlike model-based RI., equivalent abstraction can be plugged into model-free algorithms without reward model. Our contributions are as follows:
1. We develop a novel approach for constructing MDP homomorphisms (equivalent effect abstraction) that requires no prior knowledge from practitioner
2. In the tabular setting, we show equivalent effect abstraction can improve the planning efficiency of model-based algorithms and the sample efficiency of model-free algorithms
3. In the deep RL setting, show equivalent effect abstraction can be learned from experience and then leveraged to improve sample efficiency
In Section we formally describe the MDP homomorphism framework fraaeerrhhiss--awwwerrk MDD homomorphisms and then introduce equivalent effect abstraction in Section 3. In Section 4 empirieally validate our approach using benchmarks from the MDP homomorphism literature, An-overview related work is found in S and we finifinssh with limitations in Section well weweeenntssssossssss in Section-Related work, limitations and conclusions are in Sections 5, a.d
2
MDP HOMOMORPHISMS
Using the definition from (Silver, 2015), an MDP M can be described by tuple (S, A. R, where S is the set of all states, .A is the set of all actions, R = EE EE[R-11SS = s, A, = al is the reward function that determines the scalar reward received at each state, P P the transition function of the environment describing the probability of moving from state to another for given action and 7 E [0, 1] is the discount factor describing how much agent should favour immediate rewards over those in future states. An agent interacts with an environment through its policy ((]ss) P[A, a S, s] (Silver, 2015) which maps the current state given action. To solve an MDP, an RL agent must develop policy that maximises the return G, which is equal to the sum of discounted future rewards G ΣT
s' = 8, A, a)
o 7 R (where is the current timestep and T' is the number of timesteps a learning episode).
worth briefly mentioning that equivalent effect abstraction assumes MDP definition where the reward function is defined by given state rather than how an agent travels to that state (i.e. reward functions are defined as R(s') rather than R(s, a, in)-in in the vast majority of RL benchmarks this is safe assumption.
(Ravindran & Barto, 2001) introduced the concept of a homomorphism which, using the notation and definitions from (van der Pol al., 2020b), is a homomorphism h (a, {a.g|s E S) from agent's
2

---

Under review as a conference paper at ICLR 2023
the cart is pushed left, what will be the next state?
Experienced S
Next State S
If the cart was just pushed right, what the previous state?
Abstract State 5,
find equivalent states by moving forward and then backwards through model of the environment. Unlike previous approaches to homomorphisms in MDPs (van der Pol et al., 2020b), we are not exploiting the horizontal symmetry of the environment. contrast, we take advantage of the fact that state-action pairs that lead to the same state usually have equivalent values.
Figure 2:
experienced MDP M to an abstract MDP M = (S,A,P,R,,)) which satisfies the following.
P(o(s))o(s),a S0E
Σ
a(s" Vs, ES
(1)
=
VsES,
(2)
Where a is mapping from experienced states to abstract states, while a, is a state dependent mapping between experienced and abstract actions. Equations (1) and (2) demonstrate that homomorphic map of states and actions must maintain the transition dynamics and reward functions of the experienced MDP within the constructed abstract MDP. As proved in (Ravindran & Barto, 2001) and leveraged more recently by (van der Pol et al., 2020b), if the true Q-values in the abstract MDP are known Q" (s)),, (a)) -then equipped with homomorphism h, the true Q-values the experienced MDP can be "lifted" (Ravindran & Barto, 2001) from their abstract counterparts.
(a(s), a. (Q)) Q" a) Vs E S
Where Q is the value of observing state s and selecting action and signifies the value function optimal (van der Pol et al., 2020b). This useful beenuse the-The size of abstract state-action space S usually much--uuhhhheeee much-often smaller than the oize-of the experienced state-action space S A. Consequently, if an agent is equipped with homomorphic map, policy can be learned efficiently in the redseed-abstract state-action space and then "lifted" to the experienced state-action space when interacting with the environment (Ravindran & Barto, 2001). Unffrrtnatell, without knowing an environment's symmerrres,,,,,,, symmetriexsfRavvndaaa & Barto, 2001), In general, obtaining a homomorphic map is difficult. In the next section Next, we present novel method for inferring homomorphic map from experience.
3 EQUIVALENT EFFECT ABSTRACTION
Equivalent effect abstraction is based on few simple observations that are guaranteed to be true in most RL. tasks R. and are deterministic. Firstly, state-action pairs that lead to the same next state have equivalent reward functions (by definition). in the majority of RL benchmarks (e.g. Deepmind Deepmindd control Tassa et al, (2018), gym control environments Brockman et al. (2016) and the Atari learning environment Bellemare al. (2013)). (113))..WW.. While common MDP formulations (like the one used in Section 2) define reward functions as depending on the current state and previous action, in many environments it is possible to drop the dependence on the previous action. Secondly, state-action pairs that lead to the same next state also have equivalent transition functions under an unknown state homomorphism function To leverage these properties to learn policies in abstract MDP, we must develop a method to learn the mapping a between state-action pairs that lead to the same state.
Consider the normal RL scenario, where an agent takes an action in state represented by-a vector s and ends up in state s'. Once in state s', how can we infer what other state-action pairs could also lead state s'? We could make this inference if we had
3

---

Under review as a conference paper at ICLR 2023
Algorithm 1 Equivalent Action Abstraction Q-Network
class
(qulvvllentEffecco@eeeooc....... :odule))
def __init__self, forward_model. backward_moder oum_offactions, canonical_action)
fodwarddewdd forward_model # pretrained forward model
haffbackkwrd_modee backward_mode
pretrained backward model
gelffnumber_of_ace of_aciions num num_of_actions
solf.canonical_cction = canonical_action
nelt.q_net nn. Sequuntnt (nn. (2nerr 250), nn. meRRelU), Linear 1))
def (sel, state):
values = [] # list to store q-values for each action
for i, an_acion in enmmerate(self. ofmerrrtessllf.mmeoooooooooo number_of_actoons
next_state
forward_modell
an_action)
equivalent_state =
mackward_modelsstttt canonical_action) value_for: foreefforrrth__eeeee ith action self, qenettteeettteeeeeee (equivalent_sstate
aplues for f action) return values
Figure 3: Pytorch (Paszke et al., 2017), wyle__pseudoocoe for pee.ivalleenntttt _efeeettfftteeeeeeeeeeeeeee abstraction -ar.e
Q-network. Our Q-network only has an output for one action the canonical action. To transform input states into the canonical reference frame we move forwards and backwards through a dynamics model, with the input action on the forwards model giving the order of the equivalent Q-value outputs.
forwards and backwards model of the environment. To train our forward model F and backwards model B, we collect experience tuples (s, as,,''',')) from the environment and then optimise the parameters of our models and o with MSE loss function.
arg mig (s) ss'll
(4)
arg [B a) s112
(5)
With a learned forwards model Fe and backwards model Bg we are ready to reduce the size of our MDP. First, we select action ā to be our canonical action. Next for a given dttteactiion pair (s, a) the equivalent (canonical) state-action pair can be computed by moving forwards through our model and then backwards with the canonical action ā.
5=0(s,) = (Fo(ss.)
We describe the combination of the forwards and backwards model solely to be consistent with the previous literature van der Pol et al. (2020b) (i.c. is state bomomorphism) In simple terms, o a mapping another state action pair that leads the same next state. To construct a, we combine a forwards model prediction and a backwards model prediction, The first model is conditioned on the action you would like to know the value for and the reverse action is the canonical action.
In theory, one one _ow.eooondeleeeoee calculate multiple multpple_equivalen. equivalent state-action ppirtssnn_pp___________iee for a given state-action pair-by queryiierggggggg. the backwards model with multiple different actions. However, our aim is to simplify the state-action space that our policy Iearss in, so instead we select one action to be our "canonical" action and all other state-action pairs are mapped into its reference frame. In Figure 2. the canonical action is moving Lightmmennnn our Q-network only has learn about moving right and we can map all right values to their equivalent left values with our forwards and backwards models._hre The canonical action is a hyperparameter, In some environments such as Cartpole, the canonical action can selectdd randomly. However, in other gaherr_nnironmeeee environments backwards state predictions are not always possible for a given stavenssttteeeeeeeeeeeeee action pair meaning canonical actions need to be selected more carefully. Next, we discuss situations where backwards predictions cannot be made.
4

---

Under review as a conference paper at ICLR 2023
By finding an equivalent canonical state-action pair for each experienced state-action pair we reduce the number of state-action pairs that an off- policy algorithm must learn to estimate- effeetively the size of the space of values from S A to S-x-1. x_1 + N. where NES. x A are the states where the canonical action is not computable hecause of the edge cases described in Section 2 stochastic environments, no unique backwards action or transitions that do not exist in the environment such as traveling left to a state with a border to its right).
reducing reddaeingessiinaae..o
it important to stress that with forwards and backwards models we have obtained mapping g that satisfies Equation 3. This wappsmg.ppnggggmmmm A.s. can map.aapap.. any state-action pair into a s.alllrr Sataalleeettttteeeeeeee state-action with mathhngg it state-action pair that has equivalent value Rayindran & Barto (2001) 1.011A.A As our mapping satisfies Equation 3, it also assumes the MDP homomorphism properties developed by. Ravindran & Barto (2001), meaning the policy we learn with our canonical state action pairs can be "lifted" van der Pol et al. (2020b) and used as a policy in a given experienced environment,
To leverage this property in Q-learning, for every state-action pair we want to know a value for we simply apply Equation (6) to find the equivalent state and always use the canonical action as equivalent action (i. e. Ot. = {ā, Va E A}). Algorithm ?? describes the full modified Q-learning algerithm, then predict its value with Q-network. Our definition of an MDP homomorphisms is slightly different from that introduced by (Ravindran & Barto, 2001). Equivalent effect abstraction considers equivalent transitions that collide into the same future state with the same reward. This means it is difficult to map equivalent effect abstraction onto the definitions in Equations (1) and (2). Nevertheless, the most pertinent property of MDP homomorphisms introduced by (Ravindran & Barto, 2001) ) is that the state-action homomorphism is invariant to value (see Equation (3)), which we demonstrate with our empirical results.
Lastly, we We have assumed that for every state-action pair, at least one equivalent canonical state- action exists, whw.hhhhhhhcccoc can be computed. Given a model of the environment, this is often true but net guaranteedit is not guaranteed for a variety of reasons. For example, near borders in a grid world there is no way to travel left to a border state that has a border on its right we diseass this in more detail in in Sertieee-te-nsne 4.4 and 6right. We only found this significantly effect results in the a tabular setting, which we mitigate by reverting to vanilla Q-learning table when necessary (that is leaned learned in parallel). Furthermore, stochastictyy coulddcoouuddduooe be problematic with the current formulation. If a transition is multimodal and the models models_predict the mean of the two modes then the value network will be presented discuss how recurrent state space models could be used to sample from a distribution of predicted transitions in future work (Hafner et al., 2019). Lastly, if forward dynamics are deterministic the previous wttessttateeeeeeeeeeeeeeeee state_could still be unpredictable as wetquueereeiouus previous state for a given preyious action may not exist_again again the model based RI. provides some ideas for addressing this in future work (Yu et al., 2021).
-A with states that do not actually exist in the environment, In Section 6, we
4 EXPERIMENTS
We validate test our approach on a wide-range of RL tasks, from tabular reinforeement learning RL with no function approximation to using convolutional DQN Silver et al. (2016) to do control from pixels. Where possible, we try to have overlap in our experiments with the MDP homomorphism literature- Cartpole to overlap with the well known van der Pol et al. (2020b) and van der Pol et al. (2020a) as well as Predator Prey to overlap with van der Pol et al. (2020a). Hyperparameter settings and sweep configurations can be found in Appendix A. Shaded regions indicate standard error of the mean.
4.1 TABULAR MAZE
We begin with the maze environment from (Sutton & Barto, 2018, p. 165). The maze consists of 6 x 9 cells. The agent starts on the far West of the maze and must navigate East and find its way around borders in the middle of the environment. After passing the borders, the agent must then travel further East to reach corridor, before moving North to reach a goal location. In this particular experiment, we assume a model of the environment is known Satorehand-whihe we move through forwards and then backwards to create homomorphic map.
5

---

Under review as a conference paper at ICLR 2023
Tabular Maze Q-Learning
Tabular Maze Q-Planning
1.0
Tost Roturrn
0.5
Return
0.0
Equivalent Effect Abstraction DQN
-0.5
-0.5
DON
25
50
75
100
125
150
175
100
200
300
100
500
Learning Steps
Planning Steps
(a)
(b)
Figure 4: (a), (b) a tabular gridworld abstracting away redundancies in the Q-table improves sample efficiency. For the vanilla Q-learning/planning baselines, agent must learn the value of each state-action pair for all actions in the action space. With equivalent effect abstraction, we reduce the number of Q-values that need to be learned. This yields improvements in both model-free methods (Q-Iearning) and prodel-based approaches (Q-planning). 50 seeds are used for Q-learning while 10 seeds are used for Q-planning-
We use an open-source Q-learning (Watkins & Dayan, 1992) implementation as baseline' (along with the maze code from the same authors). We then adapt the Q-learning implementation to make use of homomorphic map as shown in Algorithm ??- 1. show in Figure 4a that our approach compares favorably to vanilla Q-learning, converging much faster to the optimal policy. To demonstrate the potential for improved In the gridworld environment, we reduce the size of the state action space, and hence the number of Q-yalues (when going left) from 216 to 67, In this case. the number of states in the reduced action space. is not reduced by the cardinality of the action space exactly._Thii is because, of the 13. states where it is not possible to reach them by traveling left because there is border to the right.
To test planning efficiency, we performed a further experiment on the gridworld with model-based Q-planning (Sutton & Barto, 2018, 161). Our Q-planning method is equivalent to the Q-learning implementation but altered to learn by randomly sampling model transitions. At each Q-planning update we pause training and evaluate performance in a 100 step episode. Planning performance is shown in Figure 4b. Similarly to the Q-learning results, leveraging knowledge of equivalent actions is able to significantly improve learning efficiency.
4.2 CARTPOLE
Next, apply our approach the Chr-Chhteeleeeeeeeeeeee benehmark from-test our approach on Cartpole Brock- man et al. (2016). In Cartpole__an an agent controls a cart with a pendulum attached to it. The goal of the task is to learn to balance the pendulum upright by moving the cart horizontally left or right. The enviromment is formulated with imputs represented by states are a four four-dimensional vector (position, velocity, angle, angular velocity) and Carpotle has discrete action space of meving the cart left or right). To learn a model, we use 3 episodes of learning experience at the beginning of training. With these 3 episodes we are able to train the homomorphic map (two simple linear models, one for each action optimised with Adam (Kingma & Ba, 2015)) to generate equivalent state-action pairs for the opposite action to the one experienced. These initial training steps are included for equivalent effect abstraction in Figure 5b. After these tniiiiiiiiiial initial episedes we freeze the leurned mapping-
We integrate the learned homomorphism into DQN implementation Silver (2015)² 2 with Q-values only being learned for an arbitrarily chosen canonical action. We compare our approach to three baselines: vanilla DQN, MDP Homomorphic Networks van der Pol et al. (2020b) and PRAE van der Pol et al. (2020a). For this experiment and those in future sections, all experiments we use the baselines' open source implementations and the libraries they build upon Stooke Abbeel (2019). MDP Homomorphic networks uses specially constructed weights that are equivariant to
Pittps://gite thub. http://ehaaaggrrifiithhhhc wtw.ggi..to//.ooeeoooooooo
Phttps: //t/tt//tttttttt www.giitubbbb..... com hwwpb/ccoo//ccuttcc/33fff. 7786611177777718 122766666

---

Under review as a conference paper at ICLR 2023
Predator Prey
Caripolc
PRAE 10
Delance Time (Detur))
200
FRAE 100
PRAE 10000
1.50
Homemorphic
DQN meon
1000
Homomorphh plin
DQN
Equivalent Effect
50
Abstruction DQN
Episode
Episode
(a)
(h)
Figure 5: (a) In settechhsstie-prddaae preyenvirenmennt, the homomorphic DQN is able learn mapping to improve sample efficiency over vanilla DQN, Note-here thre-that unlike- Unlike in Cartpole the experience used train the homomorphic map is not included in the plot for both the equivalent effect abstraction DQN. In the case of equivalent effect abstraction we reuse one set of pre-trained backwards and forwards models for each RL training run. (b) In Cartpole, the improvement sample efficiency is more dramatic, seattt--eivvieeeete effect abstraction sighiffeaniiiii improving convergence speed. Equivalent effect abstraction learns a model for improved convergence very quickly (the model is learned in the initial 3 episodes of experience that is shown on the x-axis of the plot). We retrain this model for every repeat. For vanilla DQN and equivalent effect abstraction we use 50 seeds, 5 seeds are used for each PRAE result and 30 seeds used for MDP homomorphic networks. For clarity, PRAE 10 refers to using PRAE with 10 episodes of data to construct an environment model, which it then plans in.
environment symmetries to io-impreee PRAE trains a contrastive model to learn mapping latent "plannable" MDP, that satisfies the definitions of an MDP homomorphism. PRAF then performs planning on the learned abstract
prior knowledge of symmetries).
MDP. For Cartpole, we tried to learn a mapping for PRAE with datasets of 10, 100, 10000 episodes of random experience. We plot the average convergence performance of PRAE's planning algorithm. In the cartpole, task, we did not need to avoid any canonical state action pairs that not computable meaning reduction of the size of the state action space from S A to x.1
As shown in Figure 5b, in the low sample regime improve upon both PRAE and vanilla DQN-with our approach converging at around episode 12 while vanilla Q-learning takes around 16 episodes converge- that this improvement includes the number of episodes required to learn our mapping. We found that in many cases MDP homomorphic networks were able to achieve good performance but worst case runs brought down mean performance significantly. To demonstrate this, we plot the mean and median for MDP homomorphic networks (only the mean is plotted for other methods). The median MDP homomorphic network is also able to similarly improve upon the vanilla DQN baseline, by leveraging a practitioner's prior knowledge of environment symmetries.Min MDD homomorphic networks are not necessarily a competing approach to equivalent effect abstraction and future work could conceivably combine these two approaches for even greater sample efficiency.
4.3 STOCHASTIC PREDATOR PREY
Following the seminal work of van der Pol et al. (2020b), we also benchmark our method on the preda- tor prey environment, where agent mueh-mus, must chase stochastically moving prey in 2D world van der Pol et al. (2020a). The observations are 7 x 7 x 3 tensor encoding agent position and prey position. The objective of the agent is catch the prey fewest steps possible. Rewards are set to -0.1 unless the predator catches the prey, in which case the episode ends and reward of is provided.
For equivalent effect abstraction, we train an action dependent forwards and backwards models on an experience dataset created by taking random actions in the environment for 10 environment steps (equivalent around 170 episodes of random experience). An interesting point to note here is that the environment stochastic, so learning a perfect model of the environment is impossible. We compare to the same baselines introduced in Section 4.2. For PRAE, benchmark with both 10000 and 100 episodes of experience data and then perform planning to convergence. Equivalent effeet- abstraction delivers- If model can be obtained before online
7

---

Under review as a conference paper at ICLR 2023
Asterix
Homomorphic DQN
Horizontal
Return
20-
Vertical
lomomorphic DQN
1.5
Equivalent Effect
Abstraction DQN
10
DQN
0.0
0.2 0.4
0.6
0.8 1.0
1.2
Frame
1o6
Figure 6: In the more challenging environment of Asterix. Equivalent effect abstraction allows DQN to reach a return of 20 in around (3.5 x 109) frames, while vanilla DQN requires approximately (1 x 10") frames reach this return. We run each method for five seeds. For readability, both plots have been smoothed with a Gaussian kernel where o =
Policy learning, then eguivalent eeferet effect abstraction can deliver an improvement over vanilla DQN consistent with the Caripole results. on the stochastic predator prey environment, Howeyer, for this environment, the number of episodes to learn policy is extremely small (x 5), so learning model of the environment online is not practical, As policy can be learned very. gerry.y.ckky..omm it is not a particularly useful environment to test, sample efficiency--il instead serves as useful benchmark (as demonstrated in van der Pol et al, (2020b)) to address wi MDP homomorphisms in stochastic environments. Again in a similar vein to Cartpole, PRAE can be very effective, but only when a relatively large amount of model training experience is available.
then
possible limitations of
4.4 MINATAR ASTERIX
MinAtar Young Tian (2019) emulates a subset of games from the Atari learning environment Bellemare al. (2013) but at lower dimensionality- aalllloonn policies similar to those required in the full Atari learning environment to be trained at lower computational eest. e This is useful because vee-cost-we can do similar benchmarking to previous works on MDP homomorphisms (i.e. van der Pol et al. (2020b)'s 's Atari experiments) but within tasks that are easy to reproduce. We use Asterix to benchmark equivalent effect abstraction. In Asterix, the agent must explore its environment and capture moving treasure pots, whilst avoiding being struck by moving enemies. Visualisations of the environments ean be found here'. For Asterix, the effective actions left, right, up, down and "do nothing". It is important to note that Asterix has no obvious global symmetries, making it difficult for approaches that rely on symmetries van der Pol al. (2020b).
In this section, we assume a model of the environment is known, which we construct pro- grammatically to perform equivalent effect abstraction (see supplementary code). Thus the question addressed here is not about learning backwards and forwards models as we did for Cartpole (section 4.2) and Predator Prey (section 4.3). Instead seek understand, given model of the environment, what approach can we take to perform value based deep RL in the smallest number of learning steps.
We use the convolutional DQN Silver et al. (2016) implementation from the Minatar repos- itory Young & Tian (2019) the backbone for equivalent effect abstraction, leaving their hyperparameters unchanged. As is the case in many games (and RI. environments more generally), MinAtar allows a "do nothing" action. Interestingly, homomorphism with "do nothing" as our canonical action, meaning we never have to fall back to vanilla Q-learning (which we did previously in the tabular maze when the homomorphism was impossible to construct). In Asterix it is alwaysspossible possible to reach eeery state by doing the canonical action..,eaning we reduce the size of the state-action space by a factor of 6 (the number of actions in the action space). We compare equivalent effect abstraction to the vanilla DQN from the MinAtar Young & Tian (2019) repository (again leaving hyperparameters unchanged). As a further baseline, construct a DQN with equivariant feature extractor, similar to the network used to learn Pong in van der Pol et al. (2020b). We implement feature extractors that are equivariant to both horizontal
is always possible to construct our
httpp:://gi thub. com jtwp://eenoougg//iiiii wwtp:u//////iitttttaaaaaaaaeaaae http///tt..t////aaaaaaaaaeaaaaea README
8

---

Under review as a conference paper at ICLR 2023
Sticky Action Astarix
Equivalent Effect
25
Abstraction DON
20-
DQN
Retarr
10
5-
0.0
0.2
0.4
0.6
0.0 1.0 1e6
Frame
Figure 7: Eguivalent abbsraction maintains its adyantage over vanilla DQN network in of
the presence of stochastic transitions, Following (aalowingg_MMahhain et al, 2018) we perform additional experiments using the "sticky actions". protocol. We supply our agent with perfect model in the non-stochastic setting, meaning 25% of the time its predictions are incorrect (due to stochasticity). We also perform linear interpolation syne frame and return points due logging quirk with wandb (Biewald, 2020), see Appendix B for details.
and vertical reflections. The hyperparameter selection process for this baseline is described in A.1.4.
Asterix is relatively complex game requiring an agent to explore for moving treasures while simultancously avoiding enemies(whe weeessswhheeeeeeeeeeee are-tee--r---eeeeeeeeeeeeeeeeeee mlving) During this explorationfor treasure, we show empirically that traversing the state space is much more efficient when the equivalences between different state-action pairs are known. Equivalent effect abstraction converges 3x faster than vanilla Q-learning, consistent with out results in previous sections. As there is no obvious symmetry- as there is for Pong for homooorpic DQN does not yield an advantage over the vanilla DQN.
4.4.1 STICKY ACTIONS
The original Minatar Asterix has some elements of stochasticity for one step predictioos For example,_ with one step predictions one cannot, predict when spirtes, will appear at the side of the screen. screen. Nevertheless, once the sprites appear, their movements are deterministic. This is issue also present in the full dime-dina atari emulltor Machado et al. (2018), where an agent can overfit to fully deterministic transitions, As a result, Machado et al. (2018) developed the sticky. actions protocol where 75% of the actions given by the agent's s policy are delivered to the emulator, while 25% of the time the emulator uses the agents previous action. We _ this approach to our Asterix experiments to test how well equivalent effect abstraction can cope with stochastic environments, We use the same environment model as we did in the deterministic environment and show that equivalent effect abstraction can maintain its advantage over the vanilla DQN even when transitions are unpredictable.
cannot_predict
5 RELATED WORK
(Ravindran & Barto, 2001) developed a theoretical framework for state-action abstraction under symmetry, while alse deriving-framewort So guarantees when lifting abstract policies to the experienced MDPs. (Givan et al., 2003) proposed MDP model minimisation using bisimulation (Larsen Skou, 1991). where states could be mapped to a a-common an abstract state if their expected transition dynamics and reward sequences were indistinguishable. Despite their niee theeretieal properties, These early applications of MDP homomorphisms were limited because they required required of prior knowledge from practitionerand there was also no obvious method for integrating homomorphisms into function approximation, (Biza-et 202 2021) alse___ (Biiza et 2021) use bisimulation to inform structured Hidden Markov Model priors to infer an explicit discrete reduced state space from learnt embeddings in deep RL. They show that policies resulting from planning in polices from the reduced abstract space are viable in the original state space.
framework for abstraction under symmetry--derivingg theoretical
9

---

Under review as a conference paper at ICLR 2023
It has also been shown that bisimulation can be relaxed from binary mapping to similar- ity metric in the works of (Ferns et al., 2011; Ferns & Precup, 2014) these bisimulation metrics have been used to condition embeddings using contrastive losses to generate invariant representations of Kattesin-eeep-RA.. (Zhang et al., 2020) augment the embedding encoder of their model with an additional contrastive loss term that, when used in conjunction with RI., results in encoder in which the distance between processed observations in embedding space corresponds to an approximate bisimulation metrie. This aalle-allwws-tte allews-metric--.......o------ -etrre allowing agents to perform control tasks in the presence of distractor cuesbeing inserigdiiedddddde into their ebservations. A more general similarity metric for conditioning embedding in deep RL was proposed in Agarwal et al. (2021), where they outlined a new policy similarity metric. This metric is also used in conjunction with contrastive loss, and pulls together embeddings that result in similar policies and long term behaviours. While using contrastive learning and bisimulations has achieved impressive results contrastive sontrvsne metheds Contrastic learning achieves impressive results but can be sample Sungry-making their deployment in the low sample regime difficult van der Pol et al. (2020a)).
More recently, homomorphisms in MDPs have been promoted a means to improve sam- ple efficiency in tasks with symmetries. PRAE (van der Pol et al., 2020a) trained a contrastive world model and accompanying loss function that satisfies the transition dynamics definition (Equation (I)) by design- wwith good results in Cartpole when 1000 episode of experience are available for training. However, in the relatively low sample regime (e.g. 100 episodes) their homomorphism is less effective as we show in Figure 5b. (van der Pol et al.. 2020b) derived method for learning network weights that are equivariant to environment symmetries (which implicitly creates an MDP homomorphism) but their algorithm requires a practitioner to hardcode symmetric groups beforehand. Similar approaches have been adopted for environments with continuous symmetries Wang et al. (2022). Another recent approach to using MDP homomorphisms is detailed in (Biza & Platt, 2019). Here the problem of finding MDP homomorphisms is approached using Online Partition Iteration the mappings from state to abstract states are learnt by predicting which partition a state should fall into given an action, and refining the partitions through splitting.
Perhaps the most relevant related approach is the use of afterstates Sutton & Barto (p.18) 136]. Afterstates slightly shift an MDP out of phase with conventional state transitions- creating environment states that are in-between the initial effect of policy's action and the reaction of the environment to said action. However, applications of the afterstates framework have generally been constrained to board games Tesauro et al. (1995) and usually focus on dealing with the stochasticity of an opponent rather than improved sample efficiency Antonoglou et al. (2021). (Misra et al., 2020) use contrastive learning to learn similar abstraction abstraction__ grouping together state-actions that will pass through or have passed through same state.
More broadly, model-based RL has enabled superhuman performance in Atari with a relatively small amount of experience (Hafner et al., 2020), but the number of planning steps required to learn policy is still very large. Backwards models have recently been proposed to improve sample efficiency of world model representation learning, which would be interesting if integrated with equivalent effect abstraction (Yu et al., 2021). A related exploration method shows simple autoencoders can learn binary abstract representations .usauul for logging of what states an agent has experienced (Tang et 2017). Our work is related to approaches that learn data augmentation policies for RL (Raileanu et al.. 2021), however instead of trying to augment augmenting experience with a more diverse distribution of states, we take diffrerdttteeeeeeeppeeeeee of are trying to learn values for narrower distribution of canonical states.
LIMITATIONS AND FUTURE WORK
Equivalent effect abstraction requires a competent model of environment dynamics dynamics model. For equivalent effect abstraction to be practical, the number of environment steps transitions reauite to learn the environment model must be smaller than the number steps needed those required to learn a poliey Bediiy--uekiil--eee paiiisspolly...t this is true in many environments (e.g. Cartpole) but it eannot not guaranteed. What's more, despite recent progress in model based learning in high dimensional environments Howeronments-Hfneer et al. (2019; 2020); Saxena et al. (2021); Ye et al. (2021), leaning a dynamics model can be difficult to learn for number of reasons. Firstly, stochastic dynamics can make
10

---

Under review as a conference paper at ICLR 2023
environment transitions multimodal, meaning point predictions of future previous states not adequate. In theory, this can be dealt with using stochastic models Hafner et al. (2019) of forwards and backwards dynamics that could be sampled when abstracting state-action pairs. Even worse, dynamics can be completely unpredictable Kendall & Gal (2017); Mavor-Parke et al. (2021), in which case aleatorie uncertainty predictions could be used signal that equivalent effect abstraction should temporarily revert to vanilla Q-learning.
Furthermore, backwards transitions could also be te may, be unpredictable because an equivalent state does not exist-as as described in exist__see Sectctin 1. In this case, uncertainty predictions could also be used to avoid using mappings that are out of distribution saah-thh--ddeer training daaaa-daaaaaaaae (because they are not possible) Kendall & Gal (2017). Integrating the described probabilistic transition models into equivalent effect abstraction is an interesting direction for future work. Lastly, aithough the cheleehhlleges surronnding model learning are signifieun, it is worth mentioning that only require one step transition models, avoiding the significant harder challenge of learning multi-step prediction models Saxena et al. (2021).
So far, we have only formulated equivalent effect abstraction Equivalent, reffctt abstraction, is formulated within the framework of value based RL in- action spaces. While value based metheds are behind discrete actions. While many of the recent breakthroughs in deep deep-Rt- are value based (e.g. Silver al. (2017). Badia al. (2020)), actor-critic approaches Schulman et al. (2017); Mnih et al. (2016) are often the natural choice for control tasks. Embedding equivalent effect abstraction into actor-critic architectures is petentially- conceivable to formulate equivalent effect abstraction within continuous action space by simply discretising the action space, which has allowed discrete action space methods to obtain state of #he-art performance en RL tasks alttine-et-al--ee (2021)-Banino et al. (202 1). Additionally, equivalent effect abstraction could also be integrated into existing.. homomorphic MDP methods that rely_ on symmetries to further reduce the size of the abstract state-action space van der Pol et al. (2020b). In -
avenue for future research, also
ww w practice, this would mean using homomorphic Q-network van der Pol et al. (2020b) but modified to only have one action output as shown in Algorithm 1. Alternatively, planning could be done with an abstract homomorphic representation van der Pol et al. (2020a), which then reduced further with equivalent effect abstraction.
7 CONCLUSION
Equivalent effect abstraction is simple method that reduces decreases the size of the a atatin spaceconsiderably for diserete action space MDPs. The aapaaapppeeh.... -_ It easy to implement and provided with backwards dynamics model of the environment it requires no prior knowledge of environment symmetries from a practitioner. We have demonstrated that equivalent effect abstraction improves sample efficiency in tabular environments, control tasks with continuous action state spaces, stochastic deep RL environments and also within an image based game playing task. An exciting next step should integrate equivalent effect into popular deep model-based algorithms (Hafner et al., 2020; Yu et al., 2021) to improve planning efficiency. Additionally, equivalent effect abstraction could also be integrated into existing homomorphic MDP methods that rely on symmetries to further reduce the of the abstract state state- action space van der der--er--eeeeeee al. (2020b).
p
REFERENCES
Rishabh Agarwal, Marlos C Machado, Pablo Samuel Castro, and Mare G Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. arXiv preprint 10Xiv:2001.05255 2021.
Ioannis Antonoglou, Julian Schritwieser, Sherjil Ozair, Thomas K Hubert, and David Silver. Planning in stochastic environments with learned model. In International Conference on Learning Representations, 2021.
Adrià Puigdomènech Badia, Bilal Piot, Steven Kapturowski, Pablo Sprechmann, Alex Vitvitskyi. Zhaohan Daniel Guo, and Charles Blundell. Agent57: Outperforming the atari human benchmark. In International Conference on Machine Learning, pp- 507-517. PMLR, 2020.

---

Under review as a conference paper at ICLR 2023
Andrea Banino, Adrià Puidomenech Badia, Jacob Walker, Tim Scholtes, Jovana Mitrovic, and Charles Blundell. Coberl: Contrastive bert for reinforcement learning. arXiv preprint w07.v22107.05431, 2021.
Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013.
Lukas Biewald. Experiment tracking with weights and biases, 2020. URL http:::/t///..... /ttttpp:://ww.... wwnd...wwww.....ooo............. com/. Software available from wandb.com.
Biza, R Platt, JW van de Meent, and L Wong. Learning discrete state abstractions with deep variational inference. Advances in Approximate Bayesian Inference, 2021.
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.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv: 01500, 2016.
Norm Ferms and Doina Precup. Bisimulation metrics are optimal value functions. In Proceedings of the Thirtieth Conference on Uncertainty in Artificial Intelligence, UAI'14, pp. 210-219, Arlington, Virginia, USA, 2014. AUAI Press. ISBN 9780974903910.
Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous markov decision processes. SIAM Journal on Computing, 4066;1662-1114,, 2011.
Robert Givan, Thomas Dean, and Matthew Greig. Equivalence notions and model minimization in markov decision processes. Artificial Intelligence, 147(1-2):163-223 2003..0....
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, pp. 3355-25655. PMLR, 2019.
Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference Learning Representations, 2020.
Alex Kendall and Yarin Gal. What uncertainties do need in bayesian deep learning for computer vision? Advances in neural information processing systems, 30, 2017.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR (Poster), 2015.
Kim G Larsen and Arne Skou. Bisimulation through probabilistic testing. Information and computa- tion, 94(1):1-28, 1991.
Marlos C Machado, Mare G Bellemare, Erik Talvitie, Joel Veness, Matthew Hausknecht, and Michael Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. Journal of Artificial Intelligence Research, 6 61:523-562 562, 2018
Augustine N Mavor-Parker, Kimberly A Young, Caswell Barry, and Lewis D Griffin. Escaping stochastic traps with aleatoric mapping agents. arXiv preprint arXh::2102.04399 2021.
Dipendra Misra, Mikael Henaff, Akshay Krishnamurthy, and John Langford. Kinematic state abstrac- tion and provably efficient rich- observvtbbvvnnnn reinforcement learning. In International conference on machine learning, pp. 6961-6971 PMLR, 2020.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint 132.5223/ 2013.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International conference on machine learning, pp. 1928-1937. PMLR, 2016.
12

---

Under review as a conference paper at ICLR 2023
The pandas development team. pandas-dev/pandas Pandas, February 2020. URL https: /doi. org/10. 5281/ /encdo. 35099333 3509134.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
Roberta Raileanu, Maxwell Goldstein, Denis Yarats, llya Kostrikov, and Rob Fergus. Automatic data augmentation for generalization in reinforcement learning. Advances in Neural Information Processing Systems, 34, 2021.
Balaraman Ravindran and Andrew G Barto. Symmetries and model minimization in markov decision processes, 2001.
Vaibhav Saxena, Jimmy Ba, and Danijar Hafner. Clockwork variational autoencoders. Advances in Neural Information Processing Systems, 34, 2021.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 1rXiv:1707.06447 2017.
David Silver. Lecture notes in reinforcement learning. Computer Science Department, University College London, 2015.
David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, loannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484
2016.
David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7676):35-222 359, 2017.
Adam Stooke and Pieter Abbeel. ripyt: A research code base for deep reinforcement learning in pytorch. arXiv preprint 1909v0199..015500 2019.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Haoran Tang, Rein Houthooft, Davis Foote, Adam Stooke, OpenAl Xi Chen, Yan Duan, John Schulman, Filip DeTurck, and Pieter Abbeel. # exploration: A study of count-based exploration for deep reinforcement learning. Advances in neural information processing systems, 30, 2017.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv:1801.00690, 2018.
Gerald Tesauro et al. Temporal difference learning and td gagmmmo...... Communications of the ACM, 38 (3):58-6) 68, 1995.
Elise van der Pol, Thomas Kipf, Frans A Oliehoek, and Max Welling. Plannable approximations to mdp homomorphisms: Equivariance under actions. In AAMAS, 2020a.
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:4199-4210, 2020b.
Oriol Vinyals, Igor Babuschkin, Wojciech M Czamecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung. David Choi, Richard Powell, Timo Ewalds, Petko Georgiev, et al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 575(7782):350-3355 2019.
Dian Wang, Robin Walters, and Robert Platt. So (2) equivariant reinforcement learning. In Interna- tional conference on learning representations (ICLR), 2022.
Christopher JCH Watkins and Peter Dayan. Q-leariing. Machine learning, 8(3):279-292) 1992.
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. Advances in Neural Information Processing Systems, 34, 2021.
13

---

Under review as a conference paper at ICLR 2023
Kenny Young and Tian Tian. Minatar: An atari-inspired testbed for thorough and reproducible reinforcement learning experiments. arXiv preprint arXiv: 2019.
Changmin Yu, Dong Li, Jianye Hao, Jun Wang. and Neil Burgess. Learning state representations via
retracing in reinforcement learning. arXiv preprint arXiv:2111. 2021.
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Representations, 2020.
A APPENDIX
A.1 HYPERPARAMETER SEARCH
Below we show the hyperparameters swept through for homomorphic MDP, DQN and Equivalent Effect Abstraction agents, broken down by environment. The PRAE architectures were generally kept the same the hyperparameters provided (van der Pol al., 2020a), with the exception of the learning rate which we evaluate at 0.0001, 0.001 and 0.1
A.1.1 SUTTON AND BARTO TABULAR GRIDWORLD
We use the hyperparameters specified (Sutton & Barto, 2018, p. 165): namely, Learning Rate= 0.1, 7 0.95 and 0.1
A.1.2 CARTPOLE
Table 1: Hyperparameters swept through for the Cartpole environment. Learning rate decay refers to decaying the learning rate by factor after a specified number of episodes have elapsed.
Hyperparameter
Values
Learning Rate
0. 0.00001, 0.0001, 0.001,0.00 0.01
decay schedule
No decay, exponential 200
7 Activation
0.8, 0.99 ReLU, tanh No decay, 5, 10, 15, 20
Learning Rate decay
Homomorphic MDP best hyperparameters: Learning Rate 0.001, e decay schedule = No decay, ? 0.8, activation tanh, Learning Rate decay No decay
Equivalent Effect Abstraction best hyperparameters: Learning Rate 0.001, decay schedule No decay, 7 = 0.8, activation = tanh, Learning Rate decay 10
Vanilla DQN best hyperparameters: Learning Rate 0.001, decay schedule No decay, 7 0.8, activation = tanh, Learning Rate decay = 15
A.1.3 STOCHASTIC PREDATOR PREY
Table Hyperparameters swept through for the Predator Prey environment.
Hyperparameter
Values
Learning Rate
0.0001, 0.001, 0.01 0.8, 0.99
Homomorphic MDP best hyperparameters: Learning Rate = 0.001, 7 = 0.99
14

---

Under review as a conference paper at ICLR 2023
Equivalent Effect Abstraction best hyperparameters: Learning Rate 0.01, 0.8
Vanilla DQN best hyperparameters: Learning Rate 0.001, Y 0.99
A.1.4 MINATAR
All hyperparameters were kept at the tuned DQN values provided in Young & Tian (2019). For the Homomorphie MDP network, we use symmetric group for reflection about the vertical axis when training on Breakout; on Asterix we show results for both reflection about the vertical and horizontal
A.2 MODEL ARCHITECTURES
A.2.1 CARTPOLE
Listing Homomorphic MDP Network (van der Pol et al., 2020b)
BasisLinnarr (repr_in=4, Btannels_in=1, repr_out=2, channels_out=60) ReLUO / tanh()
BasisLinear in=2, channels in=64, repr_out=2 channels_out-50
ReLUO / tanh()
Basisl.inear(repr_inn)) channels_in=64, repr_out=2, channels_out=1)
*BasisLinear refers to the symmeterised layers used in (van der Pol et al., 2020b) to create homomorphic networks. This network identical to the Cartpole network presented in that paper, but with only one output head that outputs state-action values.
Listing Value Network architecture for DQN and Equivalent Effect Abstraction
(inusseinut..size= output_size==104) ta
Jinear(inptt_size=1222 output_size-1024) tanh()
Lineut(inpt output_siiee10224 tanh()
Linear(inuut_size-1222 output_size==)
Listing Transition Model Architecture for Equivalent Effect Abstraction simpttiint,, output_size=?)
Listing 4 PRAE Architectures (van der Pol et al., 2020a)
# state encoder
Linear(input_size=/ outputtssize=64
ReLUO
(inaar(iputt.szze-64, output_size=32)
ReLU()
Linear(input_sizee:3, output_size=50)
#action encoder
(inuut(inptt.sizees55 output_size=1100 ReLU()
Linear(inuut_sizeer s00e=100,, output
size=
# reward prediction network
(inuttiinptt .utputtssiz==64) ReLU()
Linear (ineut(iitt_sizeeeeee output_size=1)
15

---

Under review as a conference paper at ICLR 2023
A.2.2 PREDATOR PREY
Listing Homomorphic MDP Network (van der Pol al., 2020b)
BasisConvzd(repr__t in in=l, channels_ing in=l, repr_out=4, channels_outta. filter (11ter_size=:7.7), stride=2, padding=6)
ReLUO
BasisConv2d(reprr.in channels inn4, repr_out=4, chamnels_out=8, filter_size=(5.5), stride=1, padding=6) ReLUO
GlobalMaxPool()
aasisliinerrrrer_innn chamnels_in=8, repr_out=l, chamnels_outt:22)
ReLU()
BasisLinear (repr_ channels_ing repr_out=4, channels_oue caamelssoout12221 s__
ReLU(
BasisLinear(repp in=4, channels in0_1,, repr_out=3, [tannels_out=]]
This is again the same network used in van der Pol et al. (2020b), albeit with a different output head.
Listing 6 Value Network Architecture for DQN and Equivalent Effect Abstraction
Linear(input_size=441 output_size=1024) ReLU()
Limear(inptt_size=1022, output_size=:) ReLUO
Linear(input_sizee:, output_size=1002) ReLUO
Linatrrnnutt_sizee1004,, output_size=5)
Listing 7 Transition Model Architecture for Equivalent Effect Abstraction
Linesa(input.sizee8822 output_size-512) ReLUO
(inuttiinptttssize=511,, output_size=6) ReLU()
Linear(input_siiees,, output_size=512)
ReLU()
(inatsinppttssize=551,, output_size=441)
Listing 8 PRAE Architectures (van der Pol et al., 2020a)
state encoder
Linear(input_ssie=411, ReLU()
output_size=64)
Linear(inpttssizeee6. output_size=32
ReLU()
Linutr(inptt_ssize==, output_size=50)
#action encoder
Lineearinputt_sizee55, output_size=100) ReLU()
(inutrrinput_sizee110 output_size=2)
reward prediction network
Linear(inuttssize=50 sutput_size=64) ReLU()
(inaar(iputt.size-64 sizeut
A.2.3 MINATAR
Listing 9 Homomorphic MDP Network (van der Pol al., 2020b) BasisConvd((repr_Ieg in=1, channels in=4, repr_out=2, channels_out-a filter_size==3,3) stride=l, padding=0) ReLU(
16

---

Under review as a conference paper at ICLR 2023
Linear(input.,sizee14.. sizeut_size=256)
ReLUO)
(inear(ipput.size=2255, output_size=6)
Listing 10 Value Network Architecture for DQN and Equivilent Effect Abstraction
Convid(chamells_ine;) channels_out=16, filter_size=(3,)) stride=1, padding=0) ReLUO
Jimear(input.size=102,, output__ize=122)
ReLUO (inuttrinptt,siiee112,, output_size=6)
17

---

Under review as a conference paper at ICLR 2023
A.3
ADDITIONAL MINATAR RESULTS
Breakout
equounant dgg Etpptwuunt Unrect Abstmettio DON
DQN
Return
3-
0.2
0.4 0.6 0.8
1.0
1e6
Figure 8: Here we get similar results to van der Pol et al. (2020b) showing that homomorphisms do not necessarily improve sample efficiency Breakout. On this particular environment all versions of DQN converge almost instantly to a suboptimal policy. To obtain more meaningful results, future work should integrate equivalent effect abstraction into actor-critic methods that are able to solve Breakout Young & Tian (2019).
B SYNCING LOG ARRAYS
For the sticky MinAtar results we logged frames and return values with wandb while performing
our experiments Biewald (2020), We performed this logging suboptimally as we logged steps for
frames and returns separately y meaning we ended up with two arrays where one was NaN when the other was logging a float and vice versa. We used Panda's (pandas development team, 2020) linear interpolator to estimate the NaN values in our final This significantly as the points were logged at a very high frequency.
one
This process does not effect our results
18


Reviews:
[{"review_id": "sa3FasKOuGh", "reviewer": "Reviewer_zsTX", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "The paper proposes an approach to MDP state abstraction based on the idea of forming equivalence classes of state-action pairs that must have the same value function, and thus can be limped into a single abstract state, thus improving the sample complexity and speed of solving the resulting abstract MDP. The criterion for grouping state-action pairs into an equivalence class is whether they lead to the same successor state. Empirical verification on several well-known benchmark problems demonstrate that indeed learning on the thus abstracted MDP is faster than on the original one. ", "summary_of_the_review": "I find the idea of using the principle of equivalent effect abstraction quite appealing, but am not convinced that it is applicable to general MDPs with stochastic transition functions. It would be helpful if this is investigated empirically. Nevertheless, even if the method is applicable only to problems with deterministic dynamics, it would still be a valuable addition to the toolbox of practitioners. (But, its limitations should be clearly stated.)", "strength_and_weaknesses": "The paper addresses the problem of automatic abstraction of MDPs, which is a form of representation learning, and is thus very appropriate for the conference. The general approach of looking for equivalence classes is certainly sound, with a long tradition in AI and OR, and the empirical verification looks encouraging. However, I am not sure the chosen criterion for state abstraction is very widely applicable. In Section 3, on page 3, the authors state that \"state-action pairs that lead to the same next state have equivalent reward functions (by definition)\". This statement is questionable, and would be true only if the reward depends only on the successor state, as in the example with the mouse given. But, this is not how the reward function is defined in Section 2 - there, it depends on both the starting state and the action applied. One other special circumstance where this statement would be true is when the instantaneous reward (respectively, cost) of every transition is the same. This is true for minimum-time optimal control/decision problems, but definitely not always true. Outside of these special cases, the value function of several state-action pairs with the same successor state is not the same, so how could this approach work in such cases?\n\nFurthermore, the method makes use of forward and backward models, which are trained by minimizing prediction errors, according to Equations 4 and 5. This also looks highly problematic. What happens if the transition probability of the original MDP is multimodal for a given state-action pair? It is well known that by minimizing MSE, the model could learn to predict a successor state that is even impossible according to the original transition distribution, because it lies somewhere between the modes of that distribution. If my understanding is correct, all environment the proposed algorithm has been tested on are deterministic, so this would not be a problem for them, but what happens when they are stochastic, with multimodal transition distributions?\n\nThe same argument applies to the backward model - in case of stochastic environments, it would be possible to have multiple origin states s that, if action a is applied, will result in the same successor state s'. If all such experienced transitions are used for training the backward model for the pair s',a using MSE, it will learn to predict the average of the origin states s, which might not coincide with any of them. Wouldn't that distort significantly the transition function of the abstracted MDP?\n\n          ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "GvQdbOyusP", "reviewer": "Reviewer_DLp5", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "This paper presents a novel approach to abstraction in Reinforcement Learning in discrete action spaces by construction a homomorphic representation through equivalent effect state-action pairs. The approach constructs forward and backward models which allow it to reduce the complexity of the value function representation by mapping equivalent state-action sets to canonical entries. The aim is to increase sample efficiency in subsequent learning tasks and the potential benefits are demonstrated in a number of standard experiments in maze and simple control environments. \nThe main contributions of the paper are the introduction of the new approach which allows the construction of a homomorphic abstraction (although a somewhat limited one) largely on-line from simple transition experiences without prior knowledge of the complete transition dynamics. While somewhat limited in its applicability due to its limitation to deterministic transition tasks, it introduces an interesting concept and illustrates its potential benefits in a number of experiment domains.", "summary_of_the_review": "The paper presents a novel approach to learn a simple (and somewhat limited) homomorphic representation for a problem using equivalent effect state-action pairs in order to reduce small efficiency in RL. This approach allows the abstraction to be learned efficiently on-line form a small number of transition experiences, making it easily applicable to some practical problems. While the underlying assumptions on the problem (reward only depends on state, deterministic transitions, existence of an inverse transition function in most states) somewhat limits the applicability (and thus the significance), the approach points out a potentially useful direction for abstraction work and could thus lead to beneficial follow-up work. This make it, in my view, of sufficient interest to the community to warrant publication.\nThe experiments in the paper illustrate the potential benefit of the technique.\nThe paper also includes a discussion of limitations and potential extensions. This discussion should be slightly further extended (and part of it moved earlier in the paper) to make sure that readers can fully understand the limitations of the presents part but also be able to assess possible extensions in the future to mitigate these limitations.", "strength_and_weaknesses": "The paper presents an approach to constructing a homomorphic abstraction by identifying equivalent effect state-action pairs and using these to map to a canonical element, thus reducing the complexity of the value function and thus potentially accelerating learning of policies on the abstract representations. The strengths of this approach lie in the ability to derive the abstraction form a set of environmental experiences without the need for a prior model. In addition, the presented model only requires one step equivalences for its construction, making it more applicable for on-line learning settings. However, the significance of this is somewhat limited by the prior assumptions on the task, namely that rewards are purely a function of the state (making it difficult to account for variable action costs, except by embedding them in the state through something like a remaining/consumed energy value), the need for deterministic transitions, and the assumption that an inverse model exists in most situations (i.e. that in most situations a state can be reached using every action), which seems to potentially significantly limit the application domain of the method.\nWhile the authors include a discussion of some of the limitations at the end of the paper, it would be useful if some more general discussion could be included earlier on that detailed the effect of these limitations and in particular what it implies for the types of problems that the method is applicable to. Also, it might be worth discussing and considering whether some of the assumptions can be softened to extend the applicability of the approach. \n\nAnother limitation of the homomorphism approach in this paper compared to the more general homomorphism framework seems to be that it can not easily abstract in situations where different regions of the state space have similar dynamics, i.e. where different task instances exist within the overall state space. This is an important property to extract if the learning agent might be placed into different versions of an environment and asked to solve the same task as shown in previous Homomorphism work for policy generalization such as [Ravindran, Balaraman, and Andrew G. Barto. \"Model minimization in hierarchical reinforcement learning.\" International Symposium on Abstraction, Reformulation, and Approximation. Springer, Berlin, Heidelberg, 2002.] of [Rajendran, Srividhya, and Manfred Huber. \"Learning to generalize and reuse skills using approximate partial policy homomorphisms.\" 2009 IEEE International Conference on Systems, Man and Cybernetics. IEEE, 2009.]. While alluded to in the limitations section, a further discussion of how the presented approach might be combined with other homomorphism concepts would be useful to the reader to assess the full potential of the technique.\n\nThere is a typo in the first paragraph of Section 4.3: \"\"...an agent much chase...\" should be \"...an agent must chase...\" ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "6M3ism3NN2c", "reviewer": "Reviewer_eZfk", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "This paper attempts to learn a state abstraction based on MDP homomorphism. The goal is to reduce the size of state-action space and allow for faster RL. The proposed approach learns a backward model $(B)$ and a forward model $(F)$. The forward model $F(s, a)$ is trained to predict the next state given the current state $s$ and action $a$. The backward model $B(s', a)$ is trained to predict the previous state given the current state $s'$ and last action $a$. The state abstraction mechanism maps a state (or a state-action?) pair to a value $\\sigma(s)$ such that:\n\n$\\sigma(s) = B(F(s, a), \\bar{a})$,\n\nwhere $\\bar{a}$ is an action called canonical action. Algorithm 1 shows that this action is chosen arbitrarily.  Using this approach, one can do Q-learning on $\\sigma(s), \\bar{a}$ values instead of $s, a$. A series of experiments on tabular and function approximation settings are presented that show promise of this approach. ", "summary_of_the_review": "Lack of clarity is the main reason for my current score of reject. Based on author response and discussion period, I'll consider changing my score. In particular, clarity on what $\\sigma$ is learning will be most helpful. ", "strength_and_weaknesses": "Strength:\n\n1. State abstraction is an important topic\n2. This paper presents experiments on both tabular and non-tabular settings.\n\nWeakness:\n\n1. It is unclear what is being learned by $\\sigma(s)$? Does it satisfy Equation 3? What is the intuitive idea of $\\bar{a}$? Why is it chosen arbitrarily? What theoretical properties does $\\sigma(s)$ satisfy? There appears to be a conceptual gap in the motivations in the first two pages, and the way $\\sigma$ is extracted from combining the forward and backward model. I hope the discussion period will help make this clearer. \n\nFurther, the paper currently lists handwavy statements such as:\nFigure 2 caption, _\"we take advantage of the fact that state-action pairs that lead to the same state usually have equivalent values.\"_\nand the paragraph before Section 4 says _\"Lastly, we have assumed that for every state-action pair at least one equivalent canonical state-action exists, which is often true but not guaranteed\"_\n\nCan these conditions be expressed formally? and in what realistic settings are these conditions satisfied? \n\n2. The proposed objective performs generative modeling which requires more samples to learn accurately. A more sample-efficient approach would be to do multi-class classification over a small label space, as it done by action-prediction or contrastive learning-style approaches.\n\n3. The paper ignores related work on learning state abstraction with function approximation including autoencoders, contrastive-learning approaches. A few examples are listed below, but there are many more:\n\n- (Autoencoder) #Exploration: A Study of Count-Based Exploration for Deep Reinforcement Learning, Tang et al. NeurIPS 2017\n- (Contrastive Learning) Kinematic state abstraction and provably efficient rich-observation reinforcement learning, Misra et al., ICML 2020\n- (Generative Modeling) FLAMBE: Structural Complexity and Representation Learning of Low-Rank MDPs, Agarwal et al., NeurIPS 2020\n- (Action prediction) Planning from Pixels using Inverse Dynamics Models, Paster et al., ICLR 2021\n\nThe paper also misses related work that learn mapping between a given MDP and an abstract MDP. For example, see\n\n- Provably Efficient Model-based Policy Adaptation, Song et al., ICML 2020\n- Provably Sample-Efficient RL with Side Information about Latent Dynamics, Liu et al, NeurIPS 2022\n\nWhile the proposed approach is new, it will nevertheless benefit readers to situate the work in the literature. \n\n4. Algorithm 1 is very hard to understand. See clarity below.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "4dOyBDUmTg", "reviewer": "Reviewer_MMzC", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "The paper proposes a new way to construct homomorphism for MDPs by a forward model and a backward model. Using the forward and backward models and a canonical action, a state-action pair is mapped to an equivalent canonical state-action pair and the effective size of state-action space is reduced. Empirical experiments demonstrate that the proposed method could learn faster than the baselines.", "summary_of_the_review": "The idea of using forward and backward models to construct MDP homomorphism is interesting, but no details discussion on when is it possible to learn these forward and backward models, and the numerical experiments may not be enough in demonstrating the benefits of the proposed method. ", "strength_and_weaknesses": "Strength\n- The idea of aggregating state-action pairs to an equivalent state using a canonical action is very interesting. For some simple problems this can indeed efficiently reduce the size of the equivalent state space for RL algorithms.\n\n- Experiments show that the proposed method could potentially learn faster.\n\nWeakness\n- Although the idea of using the forward and backward models makes sense in simple deterministic environments, one cannot get these models for typical stochastic environments. Furthermore, even in deterministic environments, since multiple state-action pairs could lead to the same state, the inverse image of the next state is not a singleton in general. When the inverse image of the system dynamics is not a singleton, it is impossible to learn such a backward model as described in the paper. There may be conditions under which the described forward and backward models exist, but there are no discussions about in which cases is it possible to learn the forward and backward models in the paper.\n\n- The main feature of the approach is to learn a MDP homomorphism with a smaller state space, but there is no analysis on whether the resulting learned model actually provides a simpler MDP homomorphism by either analytical or numerical evaluation. For example, in the tabular maze experiment, what is the equivalent size of state space after the homomorphism map compared with original size?\n\n- Half of the numerical experiments are done with the assumption of the given environment model. It is perfectly fine to consider the situation when the model is given and try to learn a good performing policy. But in this case, comparison with prior methods with the knowledge of model is then necessary, and none is provided in the paper.\n\n- For the predator prey environment, since the proposed method uses 170 episodes to learn the forward and backward models, it is actually not fair to claim that the method outperforms vanilla DQN. When compared with prior homomorphism methods like PRAE, since the number of samples used in learning the homomorphism are different for the methods, one cannot draw the conclusion that the proposed method is indeed more sample efficient.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject