Original: A S IMPLE A PPROACH FOR S TATE -A CTION A BSTRACTION USING A L EARNED MDP H OMOMORPHISM
Revised: A S IMPLE A PPROACH FOR S TATE -A CTION A BSTRACTION USING A L EARNED MDP H OMOMORPHISM

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Animals are able to rapidly infer, from limited experience, when sets of state-action pairs have equivalent reward and transition dynamics.
Revised: Animals are able to rapidly infer, from limited experience, when sets of state-action pairs have equivalent reward and transition dynamics.

Original: 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.
Revised: 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.

Original: MDP homomorphisms have been proposed that reduce the observed MDP of an environment to an abstract
Revised: MDP homomorphisms have been proposed that reduce the observed MDP of an environment to an abstract

Original: MDP, which can enable more sample efficient policy learning.
Revised: MDP, which can enable more sample efficient policy learning.

Original: Consequently, impressive improvements in sample efficiency have been achieved when a suitable
Revised: Consequently, impressive improvements in sample efficiency have been achieved when a suitable

Original: MDP homomorphism can be constructed a priori —usually by exploiting a practitioner’s knowledge of environment symmetries.
Revised: MDP homomorphism can be constructed a priori —usually by exploiting a practitioner’s knowledge of environment symmetries.

Original: 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.
Revised: 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.

Original: We call this method equivalent effect abstraction.
Revised: We call this method equivalent effect abstraction.

Original: On MDP homomorphism benchmarks, we demonstrate improved sample efficiency over previous attempts to learn MDP homomorphisms, while achieving comparable sample efficiency to approaches that rely on prior knowledge of environment symmetries
Revised: On MDP homomorphism benchmarks, we demonstrate improved sample efficiency over previous attempts to learn MDP homomorphisms, while achieving comparable sample efficiency to approaches that rely on prior knowledge of environment symmetries.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Modern reinforcement learning (RL) agents outperform humans in previously impregnable benchmarks such as Go (Silver et al., 2016) and Starcraft (Vinyals et al., 2019).
Revised: Modern reinforcement learning (RL) agents outperform humans in previously impregnable benchmarks such as Go (Silver et al., 2016) and Starcraft (Vinyals et al., 2019).

Original: However, the computational expense of RL hinders its deployment in promising real world applications.
Revised: However, the computational expense of RL hinders its deployment in promising real world applications.

Original: In environments with large state spaces, RL agents demand hundreds of millions of samples (or even hundreds of billions) to learn a policy—either within an environment model (Hafner et al., 2020) or by direct interaction (Mnih et al., 2013).
Revised: In environments with large state spaces, RL agents demand hundreds of millions of samples (or even hundreds of billions) to learn a policy—either within an environment model (Hafner et al., 2020) or by direct interaction (Mnih et al., 2013).

Original: Function approximation can enable some generalisation within a large state space but still most RL agents struggle to extrapolate value judgements to equivalent states.
Revised: Function approximation can enable some generalisation within a large state space but still most RL agents struggle to extrapolate value judgements to equivalent states.

Original: In contrast, animals can easily abstract away details about states that do not effect their values.
Revised: In contrast, animals can easily abstract away details about states that do not effect their values.

Original: For example, a foraging mouse understands that approaching a goalstate (shown as a piece of cheese in Figure 1) while travelling East will have the same value as approaching the same goalstate from the West.
Revised: For example, a foraging mouse understands that approaching a goal while travelling east will have the same value as approaching the same goal from the west.

Original: These sort of state abstractions have been defined in RL as Markov decision process (MDP) homomorphisms Ravindran & Barto (2001); van der Pol et al.
Revised: These sort of state abstractions have been defined in RL as Markov decision process (MDP) homomorphisms Ravindran & Barto (2001); van der Pol et al.

Original: (2020b).
Revised: (2020b).

Original: MDP homomorphisms reduce large state-action spaces to smaller abstract state-action spaces by collapsing equivalent state-action pairs in an observed MDP onto a single abstract state-action pair in an abstract MDP van der Pol et al.
Revised: MDP homomorphisms reduce large state-action spaces to smaller abstract state-action spaces by collapsing equivalent state-actions in an observed MDP onto a single abstract state-action pair in an abstract MDP van der Pol et al.

Original: (2020b).
Revised: (2020b).

Original: Given a mapping between an abstract MDP and an experienced MDP, policies can be learned efficiently in the smaller abstract space and then mapped back to the experienced MDP when interacting with the environment (Ravindran & Barto, 2001).
Revised: Given a mapping between an abstract MDP and an experienced MDP, policies can be learned efficiently in the smaller abstract space and then mapped back to the experienced MDP when interacting with the environment (Ravindran & Barto, 2001).

Original: However, obtaining mappings toand from the abstract state-action space is challenging.
Revised: However, obtaining mappings to the abstract state-action space is challenging.

Original: Success has been achieved by hard coding homomorphisms into  policynetworks (van der Pol et al., 2020b) but learning homomorphic mappings online is an unsolved problem.
Revised: Previous works hard code homomorphisms into a policy (van der Pol et al., 2020b) but learning homomorphic mappings online is an unsolved problem.

Original: We develop equivalent effect abstraction, a method that constructs MDP homomorphisms from experience via a dynamics model—leveraging the fact that state-action pairs leading to the same next state usually have equivalent value.
Revised: We develop equivalent effect abstraction, a method that constructs MDP homomorphisms from experience via a dynamics model—leveraging the fact that state-action pairs leading to the same next state usually have equivalent value.

Original: Consider the example of a gridworld maze, moving to a given cell has the same value whether you approached from the right or the left.
Revised: Consider the example of a gridworld maze, moving to a given cell has the same value whether you approached from the right or the left.

Original: Therefore, if we know the value for approaching from the right then we also know the value of moving the approaching from the left.
Revised: Therefore, if we know the value for approaching from the right then we also know the value of moving the approaching from the left.

Original: By extrapolating value judgements between equivalent state-action pairs we can use this equivalence to reduce the amount of experience required to learn a policy.
Revised: By extrapolating value judgements between equivalent state-action pairs we can use this equivalence to reduce the amount of experience required to learn a policy.

Original: An important distinction from previous works is that we do not use environment symmetries to reduce the size of the state-action space.
Revised: An important distinction from previous works is that we do not use environment symmetries to reduce the size of the state-action space.

Original: We exploit a separate redundancy common to many MDPs—for a given state there are often multiple state-action pairs that also lead to that state.
Revised: We exploit a separate redundancy common to many MDPs—for a given state there are often multiple state-action pairs that also lead to that state.

Original: While we do use a 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.
Revised: While we do use a 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.

Original: Additionally, unlike model-based RL, equivalent abstraction can be plugged into model-free algorithms without a reward model.
Revised: Additionally, unlike model-based RL, equivalent abstraction can be plugged into model-free algorithms without a reward model.

Original: Our contributions are as follows:
Revised: Our contributions are as follows:

Original: We develop a novel approach for constructing MDP homomorphisms (equivalent effect abstraction) that requires no prior knowledge from a practitioner 2.
Revised: We develop a novel approach for constructing MDP homomorphisms (equivalent effect abstraction) that requires no prior knowledge from a practitioner 2.

Original: 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.
Revised: 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.

Original: In the deep RL setting, we show equivalent effect abstraction can be learned from experience and then leveraged to improve sample efficiency
Revised: In the deep RL setting, we show equivalent effect abstraction can be learned from experience and then leveraged to improve sample efficiency

Original: In Section 2 we formally describethe MDP homomorphism framework and then introduce equivalent effect abstraction in Section 3.
Revised: In Section 2 we formally describe MDP homomorphisms and then introduce equivalent effect abstraction in Section 3.

Original: In Section 4 weempirically validate our approach using benchmarks from the MDP homomorphism literature.
Revised: In Section 4 we validate our approach using benchmarks from the MDP homomorphism literature.

Original: An overview of related work is found in Section 5 and we finish with limitations in Section 6 as well conclusions  in Section 7.
Revised: Related work, limitations and conclusions are in Sections 5, 6 and 7.

Original:  MDP H OMOMORPHISMS
Revised:  MDP H OMOMORPHISMS

Original: Using the definition from (Silver, 2015), an MDP M can be described by a tuple ⟨S , A , P , R , γ ⟩ where S is the set of all states, A is the set of all actions, R =
Revised: Using the definition from (Silver, 2015), an MDP M can be described by a tuple ⟨S , A , P , R , γ ⟩ where S is the set of all states, A is the set of all actions, R =

Original: E
Revised: E

Original: [ R t +1 | S t
Revised: [ R t +1 | S t

Original: = s, A t = a ] is the reward function that determines the scalar reward received at each state,
Revised: = s, A t = a ] is the reward function that determines the scalar reward received at each state,

Original: P = P [ S t +1 = s ′
Revised: P = P [ S t +1 = s ′

Original: | S t = s, A t = a ] is the transition function of the environment describing the probability of moving from state to another for a given action and γ ∈ [0 , 1] is the discount factor describing how much an agent should favour immediate rewards over those in future states.
Revised: | S t = s, A t = a ] is the transition function of the environment describing the probability of moving from state to another for a given action and γ ∈ [0 , 1] is the discount factor describing how much an agent should favour immediate rewards over those in future states.

Original: An agent interacts with an environment through its policy π ( a | s ) =
Revised: An agent interacts with an environment through its policy π ( a | s ) =

Original: P [ A t = a | S t = s ] (Silver, 2015) which maps the current state to a given action.
Revised: P [ A t = a | S t = s ] (Silver, 2015) which maps the current state to a given action.

Original: To solve an MDP, an RL agent must develop a policy that maximises the return G , which is equal to the sum of discounted future rewards G =
Revised: To solve an MDP, an RL agent must develop a policy that maximises the return G , which is equal to the sum of discounted future rewards G =

Original: (cid:80) Tt =

Revised: (cid:80) Tt =

Original: γ t R t +1 (where t is the current timestep and
Revised: γ t R t +1 (where t is the current timestep and

Original: T is the number of timesteps in a learning episode).
Revised: T is the number of timesteps in a learning episode).

Original: It is worth briefly mentioning that equivalent effect abstraction assumes an MDP definition where the reward function is defined by a given state rather than how an agent travels to that state (i.e. reward functions are defined as R ( s ′ ) rather than
Revised: It is worth briefly mentioning that equivalent effect abstraction assumes an MDP definition where the reward function is defined by a given state rather than how an agent travels to that state (i.e. reward functions are defined as R ( s ′ ) rather than

Original: R ( s, a, s ′ ) —in the vast majority of RL benchmarks this is a safe assumption.
Revised: R ( s, a, s ′ ) —in the vast majority of RL benchmarks this is a safe assumption.

Original: (Ravindran & Barto, 2001) introduced the concept of a homomorphism which, using the notation and definitions from (van der Pol et al., 2020b), is a homomorphism h = ⟨ σ, { α s | s ∈ S⟩} from an agent’s experienced MDP M to an abstract MDP ¯ M = (cid:10) ¯ S , ¯ A , ¯ P , ¯ R , γ (cid:11) , which satisfies the following. 
Revised: (Ravindran & Barto, 2001) introduced the concept of a homomorphism which, using the notation and definitions from (van der Pol et al., 2020b), is a homomorphism h = ⟨ σ, { α s | s ∈ S⟩} from an agent’s experienced MDP M to an abstract MDP ¯ M = (cid:10) ¯ S , ¯ A , ¯ P , ¯ R , γ (cid:11) , which satisfies the following. [Equation]

Original: [Equation]
Revised: [Equation]

Original: Where σ is a mapping from experienced states to abstract states, while α s is a state dependent mapping between experienced and abstract actions.
Revised: Where σ is a mapping from experienced states to abstract states, while α s is a state dependent mapping between experienced and abstract actions.

Original: Equations (1) and (2) demonstrate that a homomorphic map of states and actions must maintain the transition dynamics and reward functions of the experienced
Revised: Equations (1) and (2) demonstrate that a homomorphic map of states and actions must maintain the transition dynamics and reward functions of the experienced

Original: MDP within the constructed abstract MDP.
Revised: MDP within the constructed abstract MDP.

Original: 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 — ¯
Revised: 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 — ¯

Original: Q ∗ ( σ ( s ) , α s ( a )) —then equipped with a homomorphism h , the true Q-values in the experienced
Revised: Q ∗ ( σ ( s ) , α s ( a )) —then equipped with a homomorphism h , the true Q-values in the experienced

Original: MDP can be “lifted” (Ravindran & Barto, 2001) from their abstract counterparts.
Revised: MDP can be “lifted” (Ravindran & Barto, 2001) from their abstract counterparts.

Original: [Equation]
Revised: [Equation]

Original: Where Q is the value of observing state s and selecting action a and ∗ signifies the value function is optimal (van der Pol et al., 2020b).
Revised: Where Q is the value of observing state s and selecting action a and ∗ signifies the value function is optimal (van der Pol et al., 2020b).

Original: This is useful because the size of abstract state-action space ¯ S × ¯ A is usually much smaller than thesize of the experienced state-action space S × A .
Revised: The size of abstract state-action space ¯ S × ¯ A is often smaller than the experienced state-action space S × A .

Original: Consequently, if an agent is equipped with a homomorphic map, a policy can be learned efficiently in thereduced abstract state-action space and then “lifted” to the experienced state-action spacewhen interacting with the environment (Ravindran & Barto, 2001).
Revised: Consequently, if an agent is equipped with a homomorphic map, a policy can be learned efficiently in the abstract state-action space and then “lifted” to the experienced state-action space (Ravindran & Barto, 2001).

Original: Unfortunately, without knowing an environment’s symmetries, obtaining a homomorphic map is difficult.
Revised: In general, obtaining a homomorphic map is difficult.

Original: In the next section we present a novel method for inferring a homomorphic map from experience.
Revised: Next, we present a novel method for inferring a homomorphic map from experience.

Original:  E QUIVALENT E FFECT A BSTRACTION
Revised:  E QUIVALENT E FFECT A BSTRACTION

Original: Equivalent effect abstraction is based on a few simple observations that are guaranteed to be true  if R and P are deterministic.
Revised: Equivalent effect abstraction is based on a few simple observations that are guaranteed to be true in most RL tasks if R and P are deterministic.

Original: Firstly, state-action pairs that lead to the same next state have equivalent reward functions (by definition).
Revised: Firstly, state-action pairs that lead to the same next state have equivalent reward functions in the majority of RL benchmarks (e.g. Deepmind control Tassa et al. (2018), gym control environments Brockman et al. and the Atari learning environment Bellemare et al.

Revised: 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.

Original: Secondly, state-action pairs that lead to the same next state also have equivalent transition functions under an unknown state homomorphism function σ .
Revised: Secondly, state-action pairs that lead to the same next state also have equivalent transition functions under an unknown state homomorphism function σ .

Original: To leverage these properties to learn policies in an abstract MDP, we must develop a method to learn the mapping σ between state-action pairs that lead to the same state.
Revised: To leverage these properties to learn policies in an abstract MDP, we must develop a method to learn the mapping σ between state-action pairs that lead to the same state.

Original: Consider the normal RL scenario where an agent takes an action a in a staterepresented by a vector s and ends up in state s ′ .
Revised: Consider the normal RL scenario, where an agent takes an action a in a state s and ends up in state s ′ .

Original: Once in state s ′ , how can we infer what other state-action pairs could also lead to state s ′ ?
Revised: Once in state s ′ , how can we infer what other state-action pairs could also lead to state s ′ ?

Original: We could make this inference if we had a forwards and backwards model of the environment.
Revised: We could make this inference if we had a forwards and backwards model of the environment.

Original: To train our forward model F and backwards model B , we collect experience tuples ⟨ s , a, r ′ , s ′ ⟩ from the environment and then optimise the parameters of our models θ and ϕ with an MSE loss function.
Revised: To train our forward model F and backwards model B , we collect experience tuples ⟨ s , a, r ′ , s ′ ⟩ from the environment and then optimise the parameters of our models θ and ϕ with an MSE loss function.

Original: [Equation]
Revised: [Equation]

Original: [Figure]

Original: With a learned forwards model F θ and backwards model B ϕ we are ready to reduce the size of our
Revised: With a learned forwards model F θ and backwards model B ϕ we are ready to reduce the size of our

Original: MDP.
Revised: MDP.

Original: First, we select an action ¯ a to be our canonical action.
Revised: First, we select an action ¯ a to be our canonical action.

Original: Next for a given state-action 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 ¯ a . 
Revised: Next for a given state-action 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 ¯ a . [Equation]

Original: [Equation]
Revised: [Equation]

Revised: class EquivalentEffectQNetwork (nn.

Revised: Module): def __init__ ( self , forward_model, backward_model, num_of_actions, canonical_action): self .forward_model =

Revised: forward_model # pretrained forward model self .backward_model

Revised: =

Revised: backward_model

Original: 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—effectively reducing the size of the space of values from S × A to S × 1 .
Revised: # pretrained backward model self .number_of_actions = num_of_actions self .canonical_action = canonical_action self .q_net

Revised: = nn.

Revised: Sequential(nn.

Revised: Linear(2, 250), nn.ReLU(), nn.Linear(250, 1))

Revised: def forward ( self , state): values = [] # list to store q − values for each action for i, an_action in enumerate ( self .number_of_actions):

Revised: next_state = self .forward_model(state, an_action) equivalent_state = self .backward_model(state, canonical_action) value_for_ith_action = self .q_net(equivalent_state) values.append(value_for_ith_action) return values

Revised: We describe the combination of the forwards and backwards model as σ solely to be consistent with the previous literature van der Pol et al.

Revised: (2020b) (i.e. it is a state homomorphism).

Revised: In simple terms, σ is a mapping to another state action pair that leads to the same next state.

Revised: To construct σ , we combine a forwards model prediction and a backwards model prediction.

Original: 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 our  equivalent action (i.e. α s = { ¯ a, ∀ a ∈ A} ).
Revised: The first model is conditioned on the action a you would like to know the value for and the reverse action ¯ a is the canonical action. In theory, one could calculate multiple equivalent state-action pairs for a given state-action pair—by querying the backwards model with multiple different actions. However, our aim is to simplify the state-action space that our policy learns 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 right—meaning our Q-network only has to learn about moving right and we can map all right values to their equivalent left values with our forwards and backwards models.

Revised: The canonical action is a hyperparameter.

Revised: In some environments such as Cartpole, the canonical action can be selected randomly.

Revised: However, in other environments backwards state predictions are not always possible for a given state-action pair—meaning canonical actions need to be selected more carefully.

Original: Algorithm 1 describes the full modified Q-learning algorithm.
Revised: Next, we discuss situations where backwards predictions cannot be made. 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— reducing the size of the space of values from S × A to S × 1 + N , where N ∈ S × A are the states where the canonical action is not computable because of the edge cases described in Section 2 (e.g. 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). It is important to stress that with forwards and backwards models we have obtained a mapping σ that satisfies Equation 3.

Revised: This mapping σ ( s , ¯ a ) can map any state-action pair into a smaller state-action space—matching it with a state-action pair that has equivalent value Ravindran & Barto (2001).

Revised: 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.

Revised: (2020b)

Revised: and used as a policy in a given experienced environment.

Revised: To leverage this property in Q-learning, for every state-action pair we want to know a value for we

Revised: [Figure]

Revised: apply Equation (6) to find the equivalent state and then predict its value with a Q-network.

Original: Our definition of an MDP homomorphisms is slightly different from that introduced by (Ravindran & Barto, 2001).
Revised: Our definition of an MDP homomorphisms is slightly different from that introduced by (Ravindran & Barto, 2001).

Original: Equivalent effect abstraction considers equivalent transitions that collide into the same future state with the same reward.
Revised: Equivalent effect abstraction considers equivalent transitions that collide into the same future state with the same reward.

Original: This means it is difficult to map equivalent effect abstraction onto the definitions in Equations (1) and (2).
Revised: This means it is difficult to map equivalent effect abstraction onto the definitions in Equations (1) and (2).

Original: 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.
Revised: 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.

Original: 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.
Revised: We have assumed that for every state-action pair, at least one equivalent canonical state-action can be computed. Given a model of the environment, this is often true but it is not guaranteed for a variety of reasons.

Original: 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 discuss this in more detail in
Revised: 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.

Original: Sections 4.4 and 6. We only found this to significantly effect results in the a tabular setting, which we mitigate by reverting to vanilla Q-learning table when necessary (that is leaned in parallel). 
Revised: We only found this to significantly effect results in the a tabular setting, which we mitigate by reverting to vanilla Q-learning table when necessary (that is learned in parallel). Furthermore, stochasticity could be problematic with the current formulation.

Revised: If a transition is multimodal and the models predict the mean of the two modes then the value network will be presented with states that do not actually exist in the environment.

Revised: In Section 6, we discuss how recurrent state space models could be used to sample from a distribution of predicted transitions in future work (Hafner et al., 2019).

Revised: Lastly, if forward dynamics are deterministic the previous state could still be unpredictable as a unique previous state for a given previous action may not exist—again the model based RL provides some ideas for addressing this in future work (Yu et al., 2021).

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We validate our approach on awide range of RL tasks, from tabular reinforcement learning with no function approximation to using a convolutional DQN Silver et al.
Revised: We test our approach on a range of RL tasks, from tabular RL with no function approximation to using a convolutional DQN Silver et al.

Original: (2016) to do control from pixels.
Revised: (2016) to do control from pixels.

Original: Where possible, we try to have overlap in our experiments with the MDP homomorphism literature—using Cartpole to overlap with the well known van der Pol et al.
Revised: Where possible, we try to have overlap in our experiments with the MDP homomorphism literature—using Cartpole to overlap with the well known van der Pol et al.

Original: (2020b)
Revised: (2020b)

Original: and van der Pol et al.
Revised: and van der Pol et al.

Original: (2020a) as well as Predator Prey to overlap with van der Pol et al.
Revised: (2020a) as well as Predator Prey to overlap with van der Pol et al.

Original: (2020a).
Revised: (2020a).

Original: Hyperparameter settings and sweep configurations can be found in Appendix A.2.
Revised: Hyperparameter settings and sweep configurations can be found in Appendix A.1.

Original: Shaded regions indicate standard error of the mean.
Revised: Shaded regions indicate standard error of the mean.

Original: 4. T ABULAR M AZE
Revised: 4. T ABULAR M AZE

Original: We begin with the maze environment from (Sutton & Barto, 2018, p. 165).
Revised: We begin with the maze environment from (Sutton & Barto, 2018, p. 165).

Original: The maze consists of 6 × 9 cells.
Revised: The maze consists of 6 × 9 cells.

Original: 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.
Revised: 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.

Original: After passing the borders, the agent must then travel further East to reach a corridor, before moving North to reach a goal location.
Revised: After passing the borders, the agent must then travel further East to reach a corridor, before moving North to reach a goal location.

Original: In this particular experiment, we assume a model of the environment is known beforehand—which we move through forwards and then backwards to create a homomorphic map.
Revised: In this particular experiment, we assume a model of the environment is known beforehand—which we move through forwards and then backwards to create a homomorphic map.

Original: We use an open-source Q-learning (Watkins & Dayan, 1992) implementation as a baseline(along with the maze code from the same authors).
Revised: We use an open-source Q-learning (Watkins & Dayan, 1992) implementation as a baseline(along with the maze code from the same authors).

Original: We then adapt the Q-learning implementation to make use of a homomorphic map as shown in Algorithm 1.
Revised: We then adapt the Q-learning implementation to make use of a homomorphic map as shown in Algorithm 1.

Original: We show in Figure 3a that our approach compares favorably to vanilla Q-learning, converging much faster to the optimal policy.
Revised: We show in Figure 4a that our approach compares favorably to vanilla Q-learning, converging much faster to the optimal policy.

Revised: In the gridworld environment, we reduce the size of the state action space and hence the number of

Revised: Q-values (when going left) from 216 to 67 .

Revised: In this case, the number of states in the reduced action space is not reduced by the cardinality of the action space exactly.

Revised: This is because of the 13 states where it is not possible to reach them by traveling left because there is a border to the right.

Original: To demonstrate the potential for improved planning efficiency, we performed a further experiment on the gridworld with model-based Q-planning (Sutton & Barto, 2018, p. 161).
Revised: To test planning efficiency, we performed a further experiment on the gridworld with model-based Q-planning (Sutton & Barto, 2018, p. 161).

Original: Our Q-planning method is equivalent to the Q-learning implementation but altered to learn by randomly sampling model transitions.
Revised: Our Q-planning method is equivalent to the Q-learning implementation but altered to learn by randomly sampling model transitions.

Original: At each Q-planning update we pause training and evaluate performance in a 100 step episode.
Revised: At each Q-planning update we pause training and evaluate performance in a 100 step episode.

Original: Planning performance is shown in Figure 3b.
Revised: Planning performance is shown in Figure 4b.

Original: Similarly to the Q-learning results, leveraging knowledge of equivalent actions is able to significantly improve learning efficiency.
Revised: Similarly to the Q-learning results, leveraging knowledge of equivalent actions is able to significantly improve learning efficiency.

Original: 4. C ARTPOLE
Revised: 4. C ARTPOLE

Original: Next, we apply our approach to the Cartpolebenchmark from Brockman et al.
Revised: Next, we test our approach on Cartpole Brockman et al.

Original: In Cartpole an agent controls a cart with a pendulum attached to it.
Revised: In Cartpole, an agent controls a cart with a pendulum attached to it.

Original: The goal of the task is to learn to balance the pendulum upright by moving the cart horizontally left or right.
Revised: The goal of the task is to learn to balance the pendulum upright by moving the cart horizontally left or right.

Original: The environment is formulated with inputs represented by the four-dimensional vector ( position , velocity , angle , angular velocity ) and  a discrete action space of moving the cart left or right.
Revised: The states are a four-dimensional vector ( position , velocity , angle , angular velocity ) and Carpotle has a discrete action space (move the cart left or right).

Original: 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.
Revised: To learn a model, we use 3 episodes of learning experience at the beginning of training 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.

Original: These initial training steps are included for equivalent effect abstraction in Figure 4b.
Revised: These initial training steps are included for equivalent effect abstraction in Figure 5b.

Original: After these initial episodes we freeze the learned mapping.

Original: We integrate the learned homomorphism into a DQN implementation Silver (2015) 2 , with Q-values only being learned for an arbitrarily chosen canonical action.
Revised: We integrate the learned homomorphism into a DQN implementation Silver (2015) 2 , with Q-values only being learned for an arbitrarily chosen canonical action.

Original: We compare our approach to three baselines: vanilla DQN, MDP Homomorphic Networks van der Pol et al.
Revised: We compare our approach to three baselines: vanilla DQN, MDP Homomorphic Networks van der Pol et al.

Original: (2020b)
Revised: (2020b)

Original: and PRAE van der
Revised: and PRAE van der

Original: Pol et al.
Revised: Pol et al.

Original: (2020a).
Revised: (2020a).

Original: For this experiment and those in future sections, we use the baselines’ open source implementations and the libraries they build upon Stooke & Abbeel (2019).
Revised: For all experiments we use the baselines’ open source implementations and the libraries they build upon Stooke & Abbeel (2019).

Original: MDP Homomorphic networks uses specially constructed weights that are equivariant to environment symmetriesto improve sample efficiency (requiring prior knowledge of symmetries).
Revised: MDP Homomorphic networks uses specially constructed weights that are equivariant to environment symmetries (requiring prior knowledge of symmetries).

Original: PRAE trains a contrastive model to learn a mapping to a latent “plannable” MDP, that satisfies the definitions of an MDP homomorphism.
Revised: PRAE trains a contrastive model to learn a mapping to a latent “plannable” MDP, that satisfies the definitions of an MDP homomorphism.

Original: PRAE then performs planning on the learned abstract MDP.
Revised: PRAE then performs planning on the learned abstract MDP.

Original: For Cartpole, we tried to learn a mapping for PRAE with datasets of 10, 100, 10000 episodes of random experience.
Revised: For Cartpole, we tried to learn a mapping for PRAE with datasets of 10, 100, 10000 episodes of random experience.

Original: We plot the average convergence performance of PRAE’s planning algorithm.
Revised: We plot the average convergence performance of PRAE’s planning algorithm.

Revised: In the cartpole task, we did not need to avoid any canonical state action pairs that are not computable—meaning a reduction of the size of the state action space from S×A to S× 1 .

Original: As shown in Figure 4b, in the low sample regime we improve upon both PRAE and vanilla
Revised: As shown in Figure 5b, in the low sample regime we improve upon both PRAE and vanilla

Original: DQN—with our approach converging at around episode 12 while vanilla Q-learning takes aroundepisodes to converge—note that this improvement includes the number of episodes required to learn our mapping.
Revised: DQN—with our approach converging at around episode 12 while vanilla Q-learning takes aroundepisodes to converge—note that this improvement includes the number of episodes required to learn our mapping.

Original: We found that in many cases MDP homomorphic networks were able to achieve good performance but worst case runs brought down mean performance significantly.
Revised: We found that in many cases MDP homomorphic networks were able to achieve good performance but worst case runs brought down mean performance significantly.

Original: To demonstrate this, we plot the mean and median for MDP homomorphic networks (only the mean is plotted for other methods).
Revised: To demonstrate this, we plot the mean and median for MDP homomorphic networks (only the mean is plotted for other methods).

Original: The median MDP homomorphic network is also able to similarly improve upon the vanilla
Revised: The median MDP homomorphic network is also able to similarly improve upon the vanilla

Original: DQN baseline, by leveraging a practitioner’s prior knowledge of environment symmetries.
Revised: DQN baseline, by leveraging a practitioner’s prior knowledge of environment symmetries.

Original: MDP 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.
Revised: MDP 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.

Original: 4. P REDATOR P REY
Revised: 4. P REDATOR P REY

Original: Following the seminal work of van der Pol et al.
Revised: Following the seminal work of van der Pol et al.

Original: (2020b), we also benchmark our method on the predator prey environment, where an agent much chase a stochastically moving prey in a 2D world van der Pol et al.
Revised: (2020b), we also benchmark our method on the predator prey environment, where an agent must chase a stochastically moving prey in a 2D world van der Pol et al.

Original: (2020a).
Revised: (2020a).

Original: The observations are a 7 × 7 × 3 tensor encoding agent position and prey
Revised: The observations are a 7 × 7 × 3 tensor encoding agent position and prey

Original: [Figure]
Revised: [Figure]

Original: position.
Revised: position.

Original: The objective of the agent is to catch the prey in fewest steps possible.
Revised: The objective of the agent is to catch the prey in fewest steps possible.

Original: Rewards are set to -0.1 unless the predator catches the prey, in which case the episode ends and a reward of 1 is provided.
Revised: Rewards are set to -0.1 unless the predator catches the prey, in which case the episode ends and a reward of 1 is provided.

Original: 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 10environment steps (equivalent to around 170 episodes of random experience).
Revised: 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 10environment steps (equivalent to around 170 episodes of random experience).

Original: An interesting point to note here is that the environment is stochastic, so learning a perfect model of the environment is impossible.
Revised: An interesting point to note here is that the environment is stochastic, so learning a perfect model of the environment is impossible.

Original: We compare to the same baselines introduced in Section 4.2.
Revised: We compare to the same baselines introduced in Section 4.2.

Original: For PRAE, we benchmark with both 10000 and 100 episodes of experience data and then perform planning to convergence.
Revised: For PRAE, we benchmark with both 10000 and 100 episodes of experience data and then perform planning to convergence.

Original: Equivalent effect abstraction delivers an improvement over vanilla DQN, consistent with the Cartpole results.
Revised: If a model can be obtained before online policy learning, then equivalent effect abstraction can deliver an improvement over vanilla DQN on the stochastic predator prey environment. However, for this environment, the number of episodes to learn a policy is extremely small ( ∼ 5 ), so learning a model of the environment online is not practical. As a policy can be learned very quickly, it is not a particularly useful environment to test sample efficiency—it instead serves as a useful benchmark (as demonstrated in van der Pol et al. (2020b))

Revised: to address the possible limitations of MDP homomorphisms in stochastic environments.

Original: 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.
Revised: 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.

Original: 4. M INATAR A STERIX
Revised: 4. M INATAR A STERIX

Original: MinAtar Young & Tian (2019) emulates a subset of games from the Atari learning environment
Revised: MinAtar Young & Tian (2019) emulates a subset of games from the Atari learning environment

Original: Bellemare et al.
Revised: Bellemare et al.

Original: but at a lower dimensionality—allowing policies similar to those required in the full Atari learning environment to be trained at lower computational cost. This is useful because we can do similar benchmarking to previous works on MDP homomorphisms (i.e. van der Pol et al.
Revised: but at a lower dimensionality—allowing policies similar to those required in the full Atari learning environment to be trained at lower computational cost—we can do similar benchmarking to previous works on MDP homomorphisms (i.e. van der Pol et al.

Original: (2020b)’s Atari experiments)
Revised: (2020b)’s Atari experiments)

Original: but within tasks that are easy to reproduce.
Revised: but within tasks that are easy to reproduce.

Original: We use Asterix to benchmark equivalent effect abstraction.
Revised: We use Asterix to benchmark equivalent effect abstraction.

Original: In Asterix, the agent must explore its environment and capture moving treasure pots, whilst avoiding being struck by moving enemies.
Revised: In Asterix, the agent must explore its environment and capture moving treasure pots, whilst avoiding being struck by moving enemies 3 .

Original: Visualisations of the environments can be found here 3 .

Original: For Asterix, the effective actions are left, right, up, down and “do nothing”.
Revised: For Asterix, the effective actions are left, right, up, down and “do nothing”.

Original: It is important to note that Asterix has no obvious global symmetries, making it difficult for approaches that rely on symmetries van der Pol et al.
Revised: It is important to note that Asterix has no obvious global symmetries, making it difficult for approaches that rely on symmetries van der Pol et al.

Original: (2020b).
Revised: (2020b).

Original: In this section, we assume a model of the environment is known, which we construct programmatically to perform equivalent effect abstraction (see supplementary code).
Revised: In this section, we assume a model of the environment is known, which we construct programmatically to perform equivalent effect abstraction (see supplementary code).

Original: 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).
Revised: 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).

Original: Instead we seek to understand, given a model of the environment, what approach can we take to perform value based deep RL in the smallest number of learning steps.
Revised: Instead we seek to understand, given a model of the environment, what approach can we take to perform value based deep RL in the smallest number of learning steps.

Original: We use the convolutional DQN Silver et al.
Revised: We use the convolutional DQN Silver et al.

Original: (2016) implementation from the Minatar repository Young & Tian (2019) as the backbone for equivalent effect abstraction, leaving their hyperparameters unchanged.
Revised: (2016) implementation from the Minatar repository Young & Tian (2019) as the backbone for equivalent effect abstraction, leaving their hyperparameters unchanged.

Original: As is the case in many games (and RL environments more generally),
Revised: As is the case in many games (and RL environments more generally),

Original: MinAtar allows a “do nothing” action.
Revised: MinAtar allows a “do nothing” action.

Original: Interestingly, it is always possible to construct our 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).
Revised: Interestingly, it is always possible to construct our 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).

Revised: In Asterix it is always possible to reach every state by doing the canonical action—meaning we reduce the size of the state-action space by a factor of 6 (the number of actions in the action space).

Original: We compare equivalent effect abstraction to the vanilla DQN from the 
Revised: We compare equivalent effect abstraction to the vanilla DQN from the MinAtar

Original: MinAtar Young & Tian (2019) repository (again leaving hyperparameters unchanged).
Revised: Young & Tian (2019) repository (again leaving hyperparameters unchanged).

Original: As a further baseline, we construct a DQN with a equivariant feature extractor, similar to the network used to learn Pong in van der Pol et al.
Revised: As a further baseline, we construct a DQN with a equivariant feature extractor, similar to the network used to learn Pong in van der Pol et al.

Original: (2020b).
Revised: (2020b).

Original: We implement feature extractors that are equivariant to both horizontal and vertical reflections.
Revised: We implement feature extractors that are equivariant to both horizontal and vertical reflections.

Original: The hyperparameter selection process for this baseline is described in A.2.4.
Revised: The hyperparameter selection process for this baseline is described in A.1.4.

Original: Asterix is a relatively complex game requiring an agent to explore for moving treasures while simultaneously avoiding enemies (who are also moving).
Revised: Asterix is a relatively complex game requiring an agent to explore for moving treasures while simultaneously avoiding enemies.

Original: During this exploration for treasure, we show empirically that traversing the state space is much more efficient when the equivalences between different state-action pairs are known.
Revised: During this exploration, we show empirically that traversing the state space is much more efficient when the equivalences between different state-action pairs are known.

Original: Equivalent effect abstraction converges 3 × faster than vanilla Q-learning, consistent with out results in previous sections.
Revised: Equivalent effect abstraction converges 3 × faster than vanilla Q-learning, consistent with out results in previous sections.

Original: As there is no obvious symmetry—as there is for Pong for example—homomorphic DQN does not yield an advantage over the vanilla DQN. 
Revised: As there is no obvious symmetry—as there is for Pong for example—homomorphic DQN does not yield an advantage over the vanilla DQN. 4.4. S TICKY A CTIONS

Original: R ELATED W ORK
Revised: The original Minatar Asterix has some elements of stochasticity for one step predictions. For example, with one step predictions one cannot predict when spirtes will appear at the side of the screen. Nevertheless, once the sprites appear, their movements are deterministic. This is an issue also present in the full high-dimensional atari emulator Machado et al.

Original: (Ravindran & Barto, 2001) developed a theoretical framework for state-action abstraction under symmetry, while also deriving theoretical guarantees when lifting abstract policies to the experienced
Revised: (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 policy are delivered to the emulator, while 25% of the time the emulator uses the agents previous action.

Original:  MDPs.
Revised: We apply 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.  R ELATED W ORK (Ravindran & Barto, 2001) developed a framework for abstraction under symmetry—deriving theoretical guarantees when lifting abstract policies to the experienced MDPs.

Original: (Givan et al., 2003) proposed MDP model minimisation using bisimulation (Larsen & Skou, 1991), where states could be mapped to a common abstract state if their expected transition dynamics and reward sequences were indistinguishable.
Revised: (Givan et al., 2003) proposed MDP model minimisation using bisimulation (Larsen & Skou, 1991), where states could be mapped to an abstract state if their expected transition dynamics and reward sequences were indistinguishable.

Original: Despite their nice theoretical properties, early applications of MDP homomorphismswere limited because they required  prior knowledge from a practitioner and there was also no obvious method for integrating homomorphisms into function approximation.
Revised: These early applications of MDP homomorphisms required of prior knowledge from a practitioner.

Original: (Biza et al., 2021)also use bisimulation to inform structured Hidden Markov Model priors to infer an explicit discrete reduced state space from learnt embeddings in deep RL.
Revised: (Biza et al., 2021) use bisimulation to inform structured Hidden Markov Model priors to infer an explicit discrete reduced state space from learnt embeddings in deep RL.

Original: They show that polices resulting fromplanning in the reduced abstract space are viable in the original state space.
Revised: They show policies from the reduced abstract space are viable in the original state space.

Original: It has also been shown that bisimulation can be relaxed from a binary mapping to a similarity metric in the works of (Ferns et al., 2011; Ferns & Precup, 2014)—thesebisimulation metrics have been used to condition embeddings using contrastive losses to generate invariant representations of states in deep RL.
Revised: It has also been shown that bisimulation can be relaxed from a binary mapping to a similarity metric in the works of (Ferns et al., 2011; Ferns & Precup, 2014)—these metrics have been used to condition embeddings using contrastive losses to generate invariant representations of states.

Original: (Zhang et al., 2020) augment the embedding encoder of their model with an additional contrastive loss term that, when used in conjunction with RL, results in an encoder in which the distance between processed observations in embedding space corresponds to an approximate bisimulation metric. This allows agents to perform control tasks in the presence of distractor cues being inserted into their observations.
Revised: (Zhang et al., 2020) augment the embedding encoder of their model with an additional contrastive loss term that, when used in conjunction with RL, results in an encoder in which the distance between processed observations in embedding space corresponds to an approximate bisimulation metric—allowing agents to perform control tasks in the presence of distractor cues.

Original: A more general similarity metric for conditioning embedding in deep RL was proposed in Agarwal et al.
Revised: A more general similarity metric for conditioning embedding in deep RL was proposed in Agarwal et al.

Original: (2021), where they outlined a new policy similarity metric.
Revised: (2021), where they outlined a new policy similarity metric.

Original: This metric is also used in conjunction with a contrastive loss, and pulls together embeddings that result in similar policies and long term behaviours.
Revised: This metric is also used in conjunction with a contrastive loss, and pulls together embeddings that result in similar policies and long term behaviours.

Original: While using contrastive learning and bisimulations has achieved impressive results, contrastive methods can be sample hungry—making their deployment in the low sample regime difficult van der Pol et al.
Revised: Contrastice learning achieves impressive results but can be sample hungry—making their deployment in the low sample regime difficult van der Pol et al.

Original: (2020a)).
Revised: (2020a)).

Original: More recently, homomorphisms in MDPs have been promoted as a means to improve sample efficiency in tasks with symmetries.
Revised: More recently, homomorphisms in MDPs have been promoted as a means to improve sample efficiency in tasks with symmetries.

Original: PRAE (van der Pol et al., 2020a) trained a contrastive world model and accompanying loss function that satisfies the transition dynamics definition (Equation (1)) by design—with good results in Cartpole when 1000 episode of experience are available for training.
Revised: PRAE (van der Pol et al., 2020a) trained a contrastive world model and accompanying loss function that satisfies the transition dynamics definition (Equation (1)) by design—with good results in Cartpole when 1000 episode of experience are available for training.

Original: However, in the relatively low sample regime (e.g. 100 episodes) their homomorphism is less effective as we show in Figure 4b.
Revised: However, in the relatively low sample regime (e.g. 100 episodes) their homomorphism is less effective as we show in Figure 5b.

Original: (van der Pol et al., 2020b) derived a 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.
Revised: (van der Pol et al., 2020b) derived a 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.

Original: Similar approaches have been adopted for environments with continuous symmetries Wang et al.
Revised: Similar approaches have been adopted for environments with continuous symmetries Wang et al.

Original: Another recent approach to using MDP homomorphisms is detailed in (Biza & Platt, 2019).
Revised: Another recent approach to using MDP homomorphisms is detailed in (Biza & Platt, 2019).

Original: Here the problem of finding MDP homomorphisms is approached using Online Partition
Revised: Here the problem of finding MDP homomorphisms is approached using Online Partition

Original: 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.
Revised: 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.

Original: Perhaps the most relevant related approach is the use of afterstates Sutton & Barto (2018)[p.
Revised: Perhaps the most relevant related approach is the use of afterstates Sutton & Barto (2018)[p.

Original: Afterstates slightly shift an MDP out of phase with conventional state transitions—creating environment states that are in-between the initial effect of a policy’s action and the reaction of the environment to said action.
Revised: Afterstates slightly shift an MDP out of phase with conventional state transitions—creating environment states that are in-between the initial effect of a policy’s action and the reaction of the environment to said action.

Original: However, applications of the afterstates framework have generally been constrained to board games Tesauro et al.
Revised: However, applications of the afterstates framework have generally been constrained to board games Tesauro et al.

Original: and usually focus on dealing with the stochasticity of an opponent rather than improved sample efficiency Antonoglou et al.
Revised: and usually focus on dealing with the stochasticity of an opponent rather than improved sample efficiency Antonoglou et al.

Revised: (Misra et al., 2020) use contrastive learning to learn a similar abstraction—grouping together state-actions that will pass

Original: More broadly, model-based RL has enabled superhuman performance in Atari with arelatively small amount of experience (Hafner et al., 2020), but the number of planning steps required to learn a policy is still very large.
Revised: More broadly, model-based RL has enabled superhuman performance in Atari with a small amount of experience (Hafner et al., 2020), but the number of planning steps required to learn a policy is still very large.

Original: Backwards models haverecently 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).
Revised: Backwards models have 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).

Revised: A related exploration method shows simple autoencoders can learn binary abstract representations useful for logging of what states an agent has experienced (Tang et al., 2017).

Original: Our work is related to approaches that learn data augmentation policies for RL (Raileanu et al., 2021), however instead of trying to augment experience with a more diverse distribution of states, we take a different approach of trying to learn values for a narrower distribution of canonical states.
Revised: Our work is related to approaches that learn data augmentation policies for RL (Raileanu et al., 2021), however instead of augmenting experience with a more diverse distribution of states, we are trying to learn values for a narrower distribution of canonical states.

Original:  L IMITATIONS AND F UTURE W ORK
Revised:  L IMITATIONS AND F UTURE W ORK

Original: Equivalent effect abstraction requires a competent model of environment dynamics.
Revised: Equivalent effect abstraction requires a dynamics model.

Original: For equivalent effect abstraction to be practical, the number of environment steps required to learn the environment model must be smaller than the number steps needed to learn a policy—luckily this is true in many environments (e.g. Cartpole) but it cannot be guaranteed.
Revised: For equivalent effect abstraction to be practical, the number of transitions required to learn the environment model must be smaller than those required to learn a policy—this is true in many environments (e.g. Cartpole) but it is not guaranteed.

Original: What’s more, despite recent progressin model based learning in high dimensional environments Hafner et al.
Revised: What’s more, despite recent progress Hafner et al.

Original: (2019; 2020); Saxena et al.
Revised: (2019; 2020); Saxena et al.

Original: ; Ye et al.
Revised: ; Ye et al.

Original: (2021),  a dynamics model can be difficult to learn for a number of reasons.
Revised: (2021), learning a dynamics model can be difficult to learn for a number of reasons.

Original: Firstly, stochastic dynamics can make environment transitions multimodal, meaning point predictions of future or previous states are not adequate.
Revised: Firstly, stochastic dynamics can make environment transitions multimodal, meaning point predictions of future or previous states are not adequate.

Original: In theory, this can be dealt with using stochastic models
Revised: In theory, this can be dealt with using stochastic models

Original: Hafner et al.
Revised: Hafner et al.

Original: (2019) of forwards and backwards dynamics that could be sampled when abstracting state-action pairs.
Revised: (2019) of forwards and backwards dynamics that could be sampled when abstracting state-action pairs.

Original: Even worse, dynamics can be completely unpredictable Kendall & Gal (2017);
Revised: Even worse, dynamics can be completely unpredictable Kendall & Gal (2017);

Original: Mavor-Parker et al.
Revised: Mavor-Parker et al.

Original: (2021), in which casealeatoric uncertainty predictions could be used to signal that equivalent effect abstraction should temporarily revert to vanilla Q-learning.
Revised: (2021), in which case uncertainty predictions could be used to signal that equivalent effect abstraction should temporarily revert to vanilla Q-learning.

Original: Furthermore, backwards transitions could also be impossible to predict because an equivalent state does not exist—as described in Section 1.
Revised: Furthermore, backwards transitions may be unpredictable because an equivalent state does not exist—see Section 1.

Original: In this case, uncertainty predictions could also be used to avoid using mappings that are out of distributionwith the model training data (because they are not possible) Kendall & Gal (2017).
Revised: In this case, uncertainty predictions could also be used to avoid using mappings that are out of distribution (because they are not possible) Kendall & Gal (2017).

Original: Integrating the described probabilistic transition models into equivalent effect abstraction is an interesting direction for future work.
Revised: Integrating the described probabilistic transition models into equivalent effect abstraction is an interesting direction for future work.

Original: Lastly,although the challenges surrounding model learning are significant, it is worth mentioning that we only require one step transition models, avoiding the significant harder challenge of learning multi-step prediction models Saxena et al.
Revised: Lastly, it is worth mentioning that we only require one step transition models, avoiding the significant harder challenge of learning multi-step prediction models Saxena et al.

Original: So far, we have only formulated equivalent effect abstraction  within the framework of value based RL in discrete action spaces.
Revised: Equivalent effect abstraction is formulated within the framework of value based RL using discrete actions.

Original: Whilevalue based methods are behind many of the recent breakthroughs indeep RL  (e.g. Silver et al.
Revised: While many of the recent breakthroughs in RL are value based (e.g. Silver et al.

Original:  Badia et al.
Revised: (2017), Badia et al.

Original: , actor-critic approaches Schulman et al.
Revised: , actor-critic approaches Schulman et al.

Original: ; Mnih et al.
Revised: ; Mnih et al.

Original: (2016) are often the natural choice for control tasks.
Revised: (2016) are often the natural choice for control tasks.

Original: Embedding equivalent effect abstraction into actor-critic architectures is apotentially fruitful avenue for future research.
Revised: Embedding equivalent effect abstraction into actor-critic architectures is a fruitful avenue for future research.

Original: It is also conceivable to formulate equivalent effect abstraction within a continuous action space by simply discretising the actionspace, which has allowed discrete action space  methods  to obtain state of the art performance on RL tasks Banino et al.
Revised: It is also conceivable to formulate equivalent effect abstraction within a continuous action space by simply discretising the action space Banino et al. 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.

Original: C ONCLUSION
Revised: (2020b).

Original: Equivalent effect abstraction is a simple method that reduces the size of the state-action space considerably for discrete action space MDPs.
Revised: In 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

Original: The approach is easy to implement and provided with a 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 spaces, stochastic deep RL environments and also within an image based game playing task.
Revised: Algorithm 1.

Original: 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.
Revised: Alternatively, planning could be done with an abstract homomorphic representation van der Pol et al. (2020a), which is then reduced further with equivalent effect abstraction.

Original:  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). A A PPENDIX A. E QUIVALENT E FFECT A BSTRACTION P SEUDOCODE Algorithm [Table]
Revised:  C ONCLUSION

Original: , 1] , small ϵ Select a canonical action ¯ a arbitrarily from A Randomly initialise Q ( s, a ) , ∀ s ∈ S , a ∈ A
Revised: Equivalent effect abstraction is a simple method that decreases the size of a state-action space.

Original: Obtain state homomorphism σ function G REEDY P OLICY ( s, a ) Initialise values = [ ] for a = 1 , |A| do if ¯ s = σ ( s t , a ) ∈ S then values.append ( Q ( σ ( s, a ) , ¯ a )) else values.append ( Q ( s, a )) end if end for return a = arg max a ( values ) end function for episode = 1 , M do Initialise sfor t = 1 , T do
Revised: It is easy to implement and provided with a backwards dynamics model of the environment it requires no prior knowledge of environment symmetries from a practitioner.

Original: With probability ϵ pick random action a t otherwise compute a t = GreedyPolicy ( s t ) Take action a t , observe r t , s t + Set a ∗ =
Revised: We have demonstrated that equivalent effect abstraction improves sample efficiency in tabular environments, control tasks with continuous state spaces, stochastic deep RL environments and also within an image based game playing task.

Original: GreedyPolicy ( s t +1 ) Set Q ( s t , a t ) += α [ r t + γQ ( s t +1 , a ∗ ) − Q ( s t , a t )] Set s t = s t +1 end for end for
Revised: 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.

Original:  A. H YPERPARAMETER S EARCH
Revised: A A PPENDIX A. H YPERPARAMETER S EARCH

Original: Below we show the hyperparameters swept through for homomorphic MDP, DQN and Equivalent
Revised: Below we show the hyperparameters swept through for homomorphic MDP DQN and Equivalent

Original: Effect Abstraction agents, broken down by environment.
Revised: Effect Abstraction agents, broken down by environment.

Original: The PRAE architectures were generally kept the same as the hyperparameters provided in (van der Pol et al., 2020a), with the exception of the learning rate which we evaluate at 0 . 0001 , 0 . 001 and 0 .
Revised: The PRAE architectures were generally kept the same as the hyperparameters provided in (van der Pol et al., 2020a) with the exception of the learning rate which we evaluate at 0 . 0001 , 0 . 001 and 0 .

Original: A.
Revised: A.1. S UTTON AND B ARTO T ABULAR G RIDWORLD

Original: . S UTTON AND B ARTO T ABULAR G RIDWORLD

Original: We use the hyperparameters specified in (Sutton & Barto, 2018, p. 165): namely, Learning Rate = 0 .
Revised: We use the hyperparameters specified in (Sutton & Barto, 2018, p. 165): namely, Learning Rate = 0 .

Original: 1 , γ = 0 .
Revised: 1 , γ = 0 .

Original: 95 and ϵ = 0 .
Revised: 95 and ϵ = 0 .

Original: A.
Revised: A.1. C ARTPOLE

Original: . C ARTPOLE

Original: [Table]
Revised: [Table]

Revised: 001 , 0 .ϵ decay schedule

Revised: No decay, exponential τ = − 1γ 0 . 8 , 0 .

Revised: Activation

Revised: ReLU, tanh

Revised: Learning Rate decay

Revised: No decay, 5, 10, 15,

Original: Homomorphic MDP best hyperparameters: Learning Rate = 0 .
Revised: Homomorphic MDP best hyperparameters: Learning Rate = 0 .

Original: 001 , ϵ decay schedule =
Revised: 001 , ϵ decay schedule =

Original: No decay, γ = 0 .
Revised: No decay, γ = 0 .

Original: 8 , activation = tanh, Learning Rate decay =
Revised: 8 , activation = tanh, Learning Rate decay =

Original: No decay
Revised: No decay

Original: Equivalent Effect Abstraction best hyperparameters: Learning Rate = 
Revised: Equivalent Effect Abstraction best hyperparameters: Learning Rate = 0 .

Original: [Table] [Table] schedule =
Revised: 001 , ϵ decay schedule =

Original: No decay, γ = 0 .
Revised: No decay, γ = 0 .

Original: 8 , activation = tanh, Learning Rate decay =
Revised: 8 , activation = tanh, Learning Rate decay =

Original: [Table]

Original:  DQN best hyperparameters: Learning Rate = [Table]
Revised: Vanilla DQN best hyperparameters: Learning Rate = 0 .

Original: 001 , ϵ decay schedule =
Revised: 001 , ϵ decay schedule =

Original: No decay, γ = 0 .
Revised: No decay, γ = 0 .

Original: 8 , activation = tanh, Learning Rate decay =
Revised: 8 , activation = tanh, Learning Rate decay =

Original: [Table]
Revised: Homomorphic MDP best hyperparameters: Learning Rate = 0 .

Original: [Table]
Revised: 001 , γ = 0 .

Original: Equivalent Effect Abstraction best hyperparameters: Learning Rate = 0 .
Revised: Equivalent Effect Abstraction best hyperparameters: Learning Rate = 0 .

Original: 01 , γ = 0 .
Revised: 01 , γ = 0 .

Original:  Vanilla DQN besthyperparameters: Learning Rate = 0 .
Revised: [Table] Vanilla DQN best

Original: 001 , γ = 0 .
Revised: [Table]

Original: A.2. M INATAR
Revised: A.1. M INATAR

Original: All hyperparameters were kept at the tuned DQN values provided in Young & Tian (2019).
Revised: All hyperparameters were kept at the tuned DQN values provided in Young & Tian (2019).

Original: For the
Revised: For the

Original: Homomorphic MDP network, we use a 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 axes.
Revised: Homomorphic MDP network, we use a 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 axes.

Original: A. M ODEL A RCHITECTURES A.3. C ARTPOLE
Revised: A. M ODEL A RCHITECTURES A.2. C ARTPOLE

Original: Listing 1: Homomorphic MDP Network (van der Pol et al., 2020b) 1 BasisLinear*(repr_in=4, channels_in=1, repr_out=2, channels_out=64) 2 ReLU() / tanh() 3 BasisLinear(repr_in=2, channels_in=64, repr_out=2, channels_out=64) 4 ReLU() / tanh() 5 BasisLinear(repr_in=2, channels_in=64, repr_out=2, channels_out=1)
Revised: Listing 1 Homomorphic MDP Network (van der Pol et al., 2020b) [Table]

Original: * BasisLinear refers to the symmeterised layers used in (van der Pol et al., 2020b) to create homomorphic networks.
Revised: * BasisLinear refers to the symmeterised layers used in (van der Pol et al., 2020b) to create homomorphic networks.

Original: This network is identical to the Cartpole network presented in that paper, but with only one output head that outputs state-action values.
Revised: This network is identical to the Cartpole network presented in that paper, but with only one output head that outputs state-action values.

Original: Listing 2: Value Network architecture for DQN and Equivalent Effect Abstraction1 Linear(input_size=4, output_size=1024)2 tanh()3 Linear(input_size=1024, output_size=1024)4 tanh()5 Linear(input_size=8, output_size=1024)6 tanh()7 Linear(input_size=1024, output_size=2)
Revised: Listing 2 Value Network architecture for DQN and Equivalent Effect Abstraction Linear(input_size=4, output_size=1024) tanh() Linear(input_size=1024, output_size=1024) tanh() Linear(input_size=8, output_size=1024) tanh()

Original:  Listing 3: Transition Model Architecture for Equivalent Effect Abstraction1 Linear(input_size=2, output_size=2)
Revised: Linear(input_size=1024, output_size=2) Listing 3 Transition Model Architecture for Equivalent Effect Abstraction

Original:  Listing 4: PRAE Architectures (van der Pol et al., 2020a)1 # state encoder 2 Linear(input_size=4 ,output_size=64) 3 ReLU() 4 Linear(input_size=64, output_size=32) 5 ReLU() 6 Linear(input_size=32, output_size=50)
Revised: Linear(input_size=2, output_size=2) Listing 4 PRAE Architectures (van der Pol et al., 2020a)

Original: 7 #action encoder 8 Linear(input_size=54 ,output_size=100) 9 ReLU() 10 Linear(input_size=100, output_size=2) 11 # reward prediction network 12 Linear(input_size=50 ,output _size=64) 13 ReLU() 14 Linear(input_size=64, output_size=1)
Revised: [Table]

Original: [-A Listing 1 BasisConv2d(repr_in=1, channels_in=1, repr_out=4, channels_out=4, 2 filter_size=(7,7), stride=2, padding=0) 3 ReLU() 
 BasisConv2d(repr_in=4, channels_in=4, repr_out=4, channels_out=8, 5 filter_size=(5,5), stride=1, padding=0) 6 ReLU() 7 GlobalMaxPool() 8 BasisLinear(repr_in=4, channels_in=8, repr_out=4, channels_out=128) 9 ReLU() 10 BasisLinear(repr_in=4, channels_in=8, repr_out=4, channels_out=128) 11 ReLU() 12 BasisLinear(repr_in=4, channels_in=128, repr_out=5, channels_out=1)-]
Revised: [-A Listing 1 BasisConv2d(repr_in=1, channels_in=1, repr_out=4, channels_out=4, 2 filter_size=(7,7), stride=2, padding=0) 3 ReLU() 
 BasisConv2d(repr_in=4, channels_in=4, repr_out=4, channels_out=8, 5 filter_size=(5,5), stride=1, padding=0) 6 ReLU() 7 GlobalMaxPool() 8 BasisLinear(repr_in=4, channels_in=8, repr_out=4, channels_out=128) 9 ReLU() 10 BasisLinear(repr_in=4, channels_in=8, repr_out=4, channels_out=128) 11 ReLU() 12 BasisLinear(repr_in=4, channels_in=128, repr_out=5, channels_out=1)-][Table]

Original: Listing 6: Value Network Architecture for DQN and Equivalent Effect Abstraction 1 Linear(input_size=441, output_size=1024) 2 ReLU() 3 Linear(input_size=1024, output_size=8) 4 ReLU() 5 Linear(input_size=8, output_size=1024) 6 ReLU() 7 Linear(input_size=1024, output_size=5) Listing 7: Transition Model Architecture for Equivalent Effect Abstraction 1 Linear(input_size=882, output_size=512) 2 ReLU() 3 Linear(input_size=512, output_size=8) 4 ReLU() 5 Linear(input_size=8, output_size=512) 6 ReLU() 7 Linear(input_size=512, output_size=441)
Revised: A.2. P REDATOR P REY

Original: Listing 8: PRAE Architectures (van der Pol et al., 2020a)
Revised: Listing 5 Homomorphic MDP Network (van der Pol et al., 2020b)

Original: 1 # state encoder 2 Linear(input_size=441, output_size=64) 3 ReLU() 4 Linear(input_size=64, output_size=32) 5 ReLU() 6 Linear(input_size=32, output_size=50) 7 #action encoder 8 Linear(input_size=54, output_size=100) 9 ReLU() 10 Linear(input_size=100, output_size=2) 11 # reward prediction network 12 Linear(input_size=50, output_size=64) 13 ReLU() 14 Linear(input_size=64, output_size=1) A.3. M IN A TAR
Revised: [Table]

Original: Listing 9: Homomorphic MDP Network (van der Pol et al., 2020b) 1 BasisConv2d(repr_in=1, channels_in=4, repr_out=2, channels_out=16, 2 filter_size=(3,3), stride=1, padding=0) 3 ReLU() 4 Linear(input_size=1408, output_size=256) 5 ReLU() 6 Linear(input_size=256, output_size=
Revised: 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(

Original: [Table]
Revised: [Table]

Original: [Table] 
Revised: [Table] [Table] [Table] [Table]

Revised: Listing 7 Transition Model Architecture for Equivalent Effect Abstraction

Original: [Table]
Revised: [Table]

Revised: Listing 8 PRAE Architectures (van der Pol et al., 2020a

Original: [Table] 
Revised: [Table] [Table] [Figure] B S YNCING L OG ARRAYS

Original: [Figure]
Revised: 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—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 ( ? ) linear interpolator to estimate the NaN values in our final plots. This process does not effect our results significantly as the points were logged at a very high frequency.