Equivariant MuZero
Authors: Andreea Deac
Theophane Weber
George Papamakarios
Venue: Accepted by TMLR ()

Abstract: Deep reinforcement learning has shown lots of success in closed, well-defined domains such as games (Chess, Go, StarCraft). The next frontier is real-world scenarios, where setups are numerous and varied. For this, agents need to learn the underlying environment dynamics, so as to robustly generalise to conditions that differ from those they were trained on. Model-based reinforcement learning algorithms, such as MuZero or Dreamer, aim to accomplish this by learning a world model. However, leveraging a world model has not yet consistently shown greater generalisation capabilities compared to model-free alternatives. In this work, we propose improving the data efficiency and generalisation capabilities of MuZero by explicitly incorporating the \emph{symmetries} of the environment in its world-model architecture. We prove that, so long as the neural networks used by MuZero are equivariant to a particular symmetry group acting on the environment, the entirety of MuZero's action-selection algorithm will also be equivariant to that group. As such, Equivariant MuZero is guaranteed to behave symmetrically in symmetrically-transformed states, and will hence be more data-efficient when learning its world models. We evaluate Equivariant MuZero on procedurally-generated MiniPacman and on Chaser from the ProcGen suite: training on a set of mazes, and then testing on unseen rotated versions, demonstrating the benefits of equivariance. We verify that our improvements hold even when only some of the components of Equivariant MuZero obey strict equivariance, which highlights the robustness of our construction.

Published in Transactions on Machine Learning Research (12/2023)
Equivariant MuZero
Andreca Deac Mila, Université de Montréal
deacandr@vmilll.. quebee
Théophane Weber
theophane Srrophane@google.oom
Google DeepMind
George Papamakarios
spenparak@@goongle.com
Google Deep
Reviewed on OpenReview: https: // openreview. net/ forum? id= EadGarTbLE
Abstract
Deep reinforcement learning has shown lots of success in closed, well-defined domains such as games (Chess, Go, StarCraft). The next frontier is real-world scenarios, where setups are numerous and varied. For this, agents need to learn the underlying environment dynamics, so robustly generalise to conditions that differ from those they were trained on. Model- based reinforcement learning algorithms, such as MuZero Dreamer, aim to accomplish this by learning world model. However, leveraging world model has not yet consistently shown greater generalisation capabilities compared to model-free alternatives. In this work, we propose improving the data efficiency and generalisation capabilities MuZero by explicitly incorporating the symmetries of the environment in its world-model architecture. We prove that, so long as the neural networks used by MuZero are cquivariant particular symmetry group acting on the enviromment, the entirely of MuZero's action-selection algorithm will also be equivariant to that group. As such, Equivariant MuZero is guaranteed to behave symmetrically in
states, and will hence be more data-cfficient when learning its world models. We evaluate Equivariant MuZero on procedurally-geenraae MiniPacman and on Chaser from the ProcGen suite: training on a set of mazes, and then testing on unseen rotated versions, demonstrating the benefits of equivariance. We verify that our improvements hold even when only some of the components of Equivariant MuZero obey striet equivariance, which highlights the robustness of construction.
1 Introduction
Reinforcement learning (RL) is a potent paradigm for solving sequential decision making problems a dynamically changing environment. Successful examples of its include game playing (Vinyals et al. 2019), drug design (Segler et al. 2018), robotics (Ibarz et al. 2021 and theoretical computer science (Fawzi ct al. 2022) However, the generality of RL often leads to data inefficiency, poor generalisation and lack of safety guarantees. This is an issue especially in domains where data is scarce or difficult to obtain, such as medicine or human-in-the-loop scenarios.
Most RL approaches do not directly attempt to capture the regularities present in the environment. As an example, consider a grid-world: moving down a maze equivalent moving left in the 90º clock-wise rotation of the same maze. Such equivalences
be formalised via Markov Decision Process homomorphisms [Ravindran] 2004 Ravindran & Barto 2004), and while some works incorporate them (e.g, van der Pol et al. 2020 Rewaei-Shoshne et. al.
2022 most. deep reinforcement learning agents would act differently in such equivalent states if they do not observe enough data. This becomes even more problematic when the number
Work performed while the author was Google DeepMind.

---

Published in Transactions on Machine Learning Research (12/2023)
of equivalent states is large. One common example is 3D regularities, such changing camera angles in robotic tasks.
In recent years, there has been significant progress in building deep neural networks that explicitly obey such regularities, often termed geometric deep learning (Bronstein et al., 2021). In this context, the regularities are formalised using symmetry groups and architectures are built by composing transformations that are equivariant to these symmetry groups (e.g. convolutional neural networks for the translation group, graph neural networks and transformers for the permutation group).
As we are looking to capture the symmetries present in an enviromment, fitting place is within the framework of model-based RL (MBRL). MBRL leverages explicit world-models to forecast the effect of action sequences, either in the form of next-state or immediate reward predictions. These imagined trajectories are used to construct plans that optimise the forecasted returns. In the context state-of- tatte-of-the-art the-art MBRL agent MuZero (Schritticser et al. 2020), a Monto-Carlo tree scarch exccuted over these world-models in order to perform aclion selection.
In this paper, demonstrate that equivariance and MBRL can be effectively combined by proposing Equivariant MuZero (EqMuZero, shown in Figure 2 variant of MuZero where equivariance constraints are enforced by design in its constituent neural networks. As MuZero does not use these networks directly to act, but rather executes a search algorithm on top of their predictions, is not immediately obvious that the actions taken by the EqMuZero agent would obey the same constraints is guaranteed to produce a rotated action when given rotated maze? One of our key contributions is proof that guarantees this: as long as all neural networks are equivariant to symmetry group, all actions taken will also be equivariant to that symmetry group. Consequently, EqMuZero can be more data-cfficient than standard MuZero, it knows by construction how to act in states it has never seen before. On the practical side, present specific implementation of EqMuZero that is equivariant to 90º rotations by design, and we empirically verify its generalisation capabilities in two rotationally symmetric grid-worlds: procedurally-generaaee MiniPacman and the Chaser game in the ProcGen suite.
2 Background
2.1
Reinforcement Learning
The reinforcement. learning problem is lypically formalised as a Markov Decision Process A, P, R, 7) formed from set of states S, a set of actions A, discount factor y [0, 1], and two functions that model the outcome of taking action a in state the transition distribution P(s'ls, a) specifying the next state probabilities, and the reward function R(s, a) specifying the reward. The aim is to learn a policy, x <als), a function specifying (probabilities of) actions to take in state 8, such that the agent maximises the (expected) cumulative reward G(tr)- = >t= 7 R(R, ax), where tr = (so- ag, 81, Ai...., ST, aT) is the trajectory taken by the agent starting in the initial state and following the policy to decide a based on
2.2 MuZero
Reinforcement learning agents broadly fall into two categories: model-free and model-based. The specific agent we extend here, MuZero (Schrittwieser et al. 2020), is a model-based agent for deterministic environments (where P(s'|s, a) = 1 for exactly one for all s E S and a E A). MuZero relies on several neural-network components that are composed to create a world model, which estimates the unseen processes in the MDP.
The neural networks
The slate encoder, h S + Z, which embeds states into latent space (e.g. Z = R*).
The action encoder, bΛ : A + ZA, which embeds actions into latent space ZA (e.g. Za = R2))
The transition function
Z x A + Z, which predicts embeddings of next states.
I Note that both the transition and reward functions implicitly call the action encoder, hA, in order appropriately embed the input action. We expand this in Section 3.1
2

---

Published in Transactions on Machine Learning Research (12/2023)
The reward function, p: Z x + R, which predicts the immediate expected reward after taking action in particular state.
The value function, v: Z - R, which predicts the value (expected cumulative reward) from a given state.
The policy function, p: 2 + 1] which predicts the probability of taking each action from the current state. These probabilities should add up to 1: EaEa p(a|z) 1.
Using these, MuZero computes its correspouding transition, reward, value and policy models respectively through composition of neural networks.
To plan its next action, MuZero execules Monte Carlo tree search (MCTS) over many simulated trajectories, generated using the above models. We use superscript notation to denote information belonging to planning such as 2 being the hidden representation of the root state in MCTS step, and subscripts for interacting with the environment such as at being the action taken at time t. Moreover, please note that we use s for input observation and for latent space representations, while the original MuZero paper (Schrittwieser ct al. 2020) uses o and s respectively.
MuZero has demonstrated state-of-thee capabilities environments, such as Go, Chess, Shogi and Atari, and has been successfully applied to real-world domains such as video compression (Mandhane et al., 2022). Although here we focus on MuZero for deterministic environments, we note that extensions to stochastic environments also exist. (Antonoglou et al. 2021 and are an interesting target for future work.
variety of deterministic or near- determinitii
2.3 Groups, Representations and Symmetries
A group o) set @ equipped with a composition operation g gh), satisfying the following axioms:
6 - 6 (written concisely
Associativity: (gh)l = s(h() for all g, I.
Identity: there exists a unique e E @ satisfying eg ge for all E B.
Inverse: for every E G there exists unique g E 6 such that gg- -
Groups are a natural way to describe symmetries: object transformations that leave the object unchanged.
Since group elements are just abstract set elements, in order to reason about them transformations, we define the notion of group action act G f2 + Ω where Ω is the space of the input data. For example, if Ω z refers to the pixels of an image and @ = Z2 refers to circular translations, the group action act translates each pixel accordingly:
act((a, b). (u,v)) ((u a) mod n, (v + b) mod n)
(1)
where (a,b) E specifies the translation operation and (u, v) E Ω the pixel being translated.
In most cases of interest, the group actions will be lincar transformations of the data Q. We can thus reason about them in the context of linear algebra by using their real represenlations: functions py 6 N that give, for every group element g 6 real matrix demonstrating how this element acts on vector space V. For example, for the rotation group @ = SO(n), the representation would provide an appropriate n rotation matrix for each rotation g. We perform the group action by multiplying the matrices corresponding to the representation and our data.
Note that we use P to denote both a group reprosentation and the reward function of MuZero. However, they should be easy distinguish as the representation function will always be applied to group elements such g. whereas the reward function will be applied to neural network embeddings
3

---

Published in Transactions on Machine Learning Research (12/2023)
2.4 Equivariance and Invariance
As symmetries are assumed to not change the essence of the data they act on, would like to construct neural networks that adequately represent such tramsery-transformed inputs. Assume we have a neural network f : >> y, mapping between vector spaces X' and V, and that we would like this network to respoct the symmetries within a group 6. Then we can impose the following condition, for all group elements g E B and inputs x E X::
f(x((0)x) - (x())
(2)
This condition known as G-equivariance: for any group element, it does not matter whether we act with it on the input or on the output of the function f f-the end result is the same. A special case of this, (6-ivarriance, is when the output representation is trivial (py(g) II)
f(x(())x) = fxx).
(3)
In geometric deep learning, equivariance to reflections, rotations, translations and permutations has been of particular interest (Bronstein al. 2021). For the specific context of reinforcement learning and MuZero, the possible inputs to our neural networks may be both states and actions. Therefore, when further discussing these conditions, note that may refer to either state or action representations, or both at once.
Generally speaking, there are three ways to obtain an equivariant model: a) data augmentation, b) data canonicalisation and c) specialised architectures. Data augmentation creates additional training data by applying group elements g to input/output pairs (x. (ar=))-quivariann is encouraged by training on the trans- formed data and/or minimising auxiliary losses such as (g))()) ((x(())x)] Data augmentation can be simple to apply, but it results in only approximately cquivariant models. Data canonicalisation requires method to standardise the input, such as breaking the translation symmetry for molecular representation by centering the atoms around the origin (Musil et al. 2021) however, in many cases such a canonical transformation may not cxist. Specialised architectures have the downside of being harder to build, but they can guarantee exact equivariance- as such, they reduce the search space functions, potentially reducing the number of parameters and increasing training efficiency.
Broustein al. (2021) discuss many ways in which one can construct state-of-the-ard htp.oo-ttee-attsssiplilieeeeee architectures, from rudimentary approaches relying on exhaustive view generation (Cohen & WWeling 2016a), to more efficient approaches like Steorable CNNs (Cohen & Welling, 2016b Weiler & Cosa 2019) all the way to
advanced approaches preserving gauge equivariance over manifolds and meshes (De Haan et al. 2020). All of these approaches can be leveraged to build constituent networks in our Equivariant MuZero framework.
2.5 Equivariance in RL and Related Work
There has been previous work at the intersection of reinforcement learning and equivariance. While leveraging multi-agent symmetries was repeatedly shown to hold promise (van der Pol et al. 2021 Muglich et al. 2022). of particular interest to us are the symmetrics emerging from the environment,
single-agent scenario.
Related work in this space can be summarised by the commulative diagram in Figure D This diagram represents various transformations Some of these arrows are neural network layers (such as Enc 5 + 2), some are action executions (such as a : S + S), while others are symmetry transformations (such as p2(g) Z + Z). In a commutative diagram, each pair of paths connecting the same start and endpoint specifies a mathematical constraint that the compositions of arrows in these two paths must be the same transformation.
might want to make data that originates from RL, as arrows.
S ps(B) S Enc, Z and S Enc, p2(0), Z. As these two paths both start in the same point and end in the same point, we recover the previously discussed G-equivariance condition on Enc:
For example, in Figure] consider the front-facing square of the lower cube, i.e., paths
Enc(ps(g)s) =
(4)
When considering only the cube the bottom, we recover Park al. 2022 supervised learning task where latent transition model T learns to predict the next state embedding They show that if T is
4

---

Published in Transactions on Machine Learning Research (12/2023)
Py(g)
R
I
a(a)
(<ga)
Dec
RR
Py(
IR
Dec
pz(g)
Z
Z
T(a)
T(ga)
Dec
Dec
Pz
Z
Enc
Enc
ps(g)
s
S
Enc a
ga
Enc
ps(g)
S
S
Figure 1: Commutative diagram of symmetries in RL. State transitions due to action are back-to-front transformations due to symmetry g are left-to-right, state encoding and decoding by the model is bottom- to-top.
equivariant, the encoder can pick up the symmetries of the environment even if it is not fully equivariant encoder
by design. MMndal et al. (2022) build a model-free agent by combining an squivariant-by-dssie and enforcing the remaining cquivariances via regularisation losses. They also consider the invariance of the reward, captured in Figure by taking the decoder be the reward model and 1. The work of van der Pol et al. (2020) can be described by having the value model the decoder, while the work of Wang et al. (2022) has the policy model the decoder and (= |A|-
While our work is the first to study the interplay of equivariance and Monte-Carbe tree search (MCTS)- backed MBRL algorithms such as MuZero, we would like to point out three recent related works that are generally inspired by exploiting symmetries model-based planning. Firstly, SymPlan (Zhao et al. (022)) takes inspiration from the symmetries present in the Value Iteration algorithm (Bellman) 1966), extend Value Iteration Networks (Tamar et al. 2016) with additional equivariance constraints. Note that, while there exists MBRL inspiration in this work, it is still implicăt planner, i.c. end-to-end agent model without any explicit invocation of mosel-based methods. Secondly, EDGI (Brehmer el al. 2023 is recent extension of diffusion-based planning (Janner et al. 2022) that incorporates group equivariance constraints. As planning with diffusion naturally occurs in a high-dimensional space with rich geometry, imposing additional spatial constraints on it is a natural direction to follow. Lastly, in work concurrent to ours, Zhao al. (2023) study how various path-planning algorith e.g., model predictive cestrol-behaee under symmetries, and derive arguments that are similar in spirit our theoretical analysis of MuZero under equivariance.
3 Equivariant MuZero
This scction describes Equivariant MuZero (EqMuZero), variant of MuZero whose constituent neural networks (encoder, transition model, reward model, value model and policy model) are equivariant by design. Section 3.1 describes a specific implementation of EqMuZero, where the equivariance is with respect to the 4-clement cyclic group C4, that is, the group of 90º rotations. Scction 3.2 formally proves that if the components of EqMuZero are equivariant with respect to any symmetry group 6, the action selection will also be equivariant with respect to 6.
5

---

Published in Transactions on Machine Learning Research (12/2023)
3.1 Implementation of -q-equvvariant MuZero
We present a concrete instance of EqMuZero for the 4-element cyclic group Ca, and we describe how its various components can be designed to obey C--equivariance (Figure 2). The C₁ group is applicable, for example, in rotationally symmetric grid-worlds, where the elements of the group represent rotating the grid- world by all four possible multiples of 90", that is, C1 = {c, 1ooo, V180e, 12700 }. Such grid-worlds are symmetric with respect to 90" rotations, the sense that moving down in a grid-world map is the same as moving left in the 90" clock-wise rotated version of the same map. As our exemplar, we will use the G-CNN architecture from Cohen& & Welling (2016a), which is among the carliest prominent architectures that directly supports equivariance to small finite groups, such as C₁-
In what follows we assume the enviromment is a grid-world for concreteness, however we note that our EqMuZero implementation is directly applicable to any environment with the same symmetry structure. Let s denote the state of the grid-world represented as a 2D array (i.e. an image). For simplicity, we assume there are only four directional movement actions in the environment, denoted by A 1...,,,,1 Any additional non-movement actions (such as the "do nothing" action) can be included without difficully. We states and actions in the obvious way: ps((0o-) )s Rgo- is the grid-world
assume that the C4 group state rotated clock-wise by 90º, and PAA (1gg.)) -
State and action encoders To enforce Cy-equivariance in the state encoder, we first need to specify the effect of rotations on the latent state z. In our implementation, the latent state consists of 4 equally shaped arrays, z = (Z), B2 23, 24), and we prescribe that a 90º clock-wise rotation manifests as a cyclical permutation: P2 (T90= )z (Z2. Z3, Z1, Z1). Then, our equivariant encoder embeds state follows:
(s) = (h(s), (Rg (8)188))) h h(2271)) = (Z): Z2, Zz. Z+)
(5)
where h is an arbitrary neural network (we use a CNN). Similarly, our equivariant action-encoder embeds action a as follows:
hA-aq (a) = (oll())) (al = (Z,:1, Za;2; Ze:3, Zaz4)
(6)
where ha is an MLP.
Transition model We propose two ways of building a C--equivariant transition model. The first one works by maintaining the structure in the latent space:
Teg (z. (eq a) = (T(z1, (a)-) T(Z2, ha-eg(a))), r(zs, (a)) ha-o
(7)
where ha-aq(a), represents the first element of the tuple obtained from encoding action The network 7 is arbitrary; in our implementation, broadcast the output of hA across all pixels of h's output, followed by ResNet. This equation satisfies Ca-oquivariance, that is, Tag (Rgee P.A )a) pz (1oo>) (s. ( a).
Our second transition model is less constrained but more involved, as it allows components of z to interact, while still retaining C--quivariiace:
Teg (2g a) (T (7,),, 22, Rg, Ra); T (22, Rz, 24, 21), / (23, 24, 21, 22), 7 (Z., 21, R2, Rz)).
For brevity, 21 denotes the pair (z), ha-- (a)1). Note that, in this version, T takes more arguments than before. In our experiments, we use the more constrained variant for MiniPacman, and the less constrained variant for Chaser, as more data is available for the latter.
Policy model We make Ca-oquivariant policy by combining state and action embeddings from all four latents follows:
(e 2) p(alz)) PlPA )a 22) + p(p.(tiso0 2a) + (( )a 24)
(9)
4
It is straightforward verify that EvEA Peg(a z) 1, ie. Peq(-1z) is properly normalisod, and that Pag (PA (190)) )a| p )z) (a 22, Lee. it satisfies Ca-equivariance.
6

---

Published in Transactions on Machine Learning Research (12/2023)
E
Pioo (z', -+)
Rggs
Zinv
Vinv (2')
21
2
s
Σ
Z2
2/2
R
23
2/3
peal lz))
2,
2/4
Pegl 1 22
Peal (z2)
pegit x.)
AA
Figure 2: Architecture of Equivariant MuZero, where h, hA are encoders, T is the transition model, P the reward model, is the value model and p is the policy predictor. Each colour represents an element of the C4 group fe, Egns, U100., U270>) applicd to the input (environment state and action).
Reward and value models Lastly, the reward and value networks (Pinv: Vanv) should be Ca-invariant. We can satisfy this constraint by aggregating the latent space with any C Cq-invariant function, such as sum, average or max. Here we use summation:
(2, = a( + plzz, Pa )a) + p(23, PA ( )a) + p(zs, P.A P.A (1270=))) )a)
(10)
Uinv (z) = v(z1) + v(2)) + u uZ33) + v(z)).
(11)
Composing the equivariant components described above (Equations 10 we construct the end-to-end equivariant EqMuZero agent., displayed in Figure
While the specific implementation presented here represents a valid EqMuZero architecture for the C4 sym- metry group, we reiterate that it is only an exemplar. If we were interested in expanding to a broader symmetry group such as which includes eight, elements, spanning both rotations and reflections we could extend our representations to include all eight transformed views, or utilise a more memory-efficient approach of Stcerable CNNs (Cohen & Welling 2016b Weiler & Cesa 2019 Zhao et al. (022)) or explicit symmetrisation (van der Pol al., (2020). In fact, we now prove theorem stating that, no matter the symmetry group of choice or the method used to ensure equivariance to it (Bronstein al. 2021), so long equivariance is guaranteed, Eq will select actions in a way that respects the underlying symmetries.
3.2 Proof of Action Equivariance for Arbitrary Symmetry Groups
As discussed in Section 2.2 MuZero does not use its constituent networks directly for acting, but it uses them to perform MlonteCCarlo tree search (MCTS), from which an action is selected. Therefore, even if MuZero's constituent networks are fully equivariant, it is not obvious that the action selection is too. In other words, are we guaranteed that EqMuZero will behave in an equivariant manner, that is, will it always perform a rotated act in a rotated word- The following theorem shows that the answer is yes.
Thcorem If all the relevant neural networks used by MuZero are G-equivoriant, the proposed EqMuZero agent will select actions in B- c--quivariant monner, that is for every state s E S and for every g E 6, if EqMuZero selects action while in 8, then it must select pa(g)a while in ps(g)s.
7

---

Published in Transactions on Machine Learning Research (12/2023)
Proof. Assume that the EqMuZero agent uses the following equivariant neural networks: heg for the encoder (which also includes the action encoder for simplicity), Teq for the transition model, Peq for the policy model, Vinv for the value model and Piav for the reward function. By assumption, heq Tag and Peq are G-equivariant, and Vinv and Pinv are O-invariant.
MuZero relies on multi-step rollouts using its internal models. Specifically, requires simulating the reward, value, policy and transition models n steps in the future, starting from state and simulating the action sequence a a ..... a". We can simulate these rollouts the following compositions of neural networks:
r Finr(Teg( ... Teq (hag (s), (a) ........""" ),")),
y" (iar .. Tog (Tg heg (8), a²) ,)))))))
(12)
p" Peq (Teg( Teq To hog (s).
aa)))
z" Tegl.-. Teq (Teq (heg (5),
aa)).
As these rollout simulations are computed by composing several G-equivariant and G-invariant. functions, they are themselves (--quivariant (in the case of p" and z") and G-invariant. (in the case of r" and v²).
MuZero also computes returns starting from intermediate simulated state emboddings, G(z*). where the depth of this state, assuming the actions simulated after this depth are *+1 k+2 ,..., a a'11, The return is
also G-invariant function, as it is linear combination of several G-invariant functions:
1-1-
G( = Σ
7 Pinol (aao===*"attttei artltT) + " 1.
(13)
7=
To prove that planning step is equivariant, we need to show that the action selection is G-cquivariant.
Since the outcome of MuZero's MCTS function is based on the initial observation, we denote the internal state of MCTS as (z, a), (z) for each node embedding x that we simulate as we expand the tree, and each action a. These values correspond respectively to Q-values, visit counts, expansion policy estimates, etc. We will use identical notation as (Shrittwieser et al. (2020) for these states, even though we denote the MuZero models, hag Tequ Pag: Pinr and Uiav: somewhat differently.
First, MuZero simulates rollout by repeatedly selecting the next action to simulate, at, as follows:
a = argmax Q* P a) VE, N2 c1 + log Es log(
(b) + c2 +
(14)
(2*
C2
where 6₁ and 02 are constant hyperparameters. Once the rollout is completed, the intermediate MCTS states are updated accordingly:
Qi(2--1 ,at)= No No ( at 12 (2 1,0 )+ + G(z+-1)
No (2 -1 at) +1
(15)
Na .aa = No Ni
1
Note that we will, from now on, use shortened notation to represent group-transformed states, actions and embeddings. Specifically, we will use g.s as shorthand for ps(g)s, a.a as shorthand for pA(g)a, and g as shorthand for Pz(g)2, in order to maintain brevity.
As discussed previously, we need to show that, for each MCTS internal state (e.g. No), if we assume hog: Teq: Peq: Pinv and Viav be (-oquivariant functions, the resulting state would also be G-oquivariant under transformations of the initial observation. That is, for all 8, 8,
N No. paa) (2,a)
(16)
To prove this, we will use induction on the number of backups performed by MCTS, t. We proceed:
Base case (t = 0) 110."(1.2 gaa)= = N6(2, a) = 0
(17)
Q8 (g=z, gaa) = Qī(2,a) = 0.
8

---

Published in Transactions on Machine Learning Research (12/2023)
Assume:
Case t: ANg0-(0s0 0. (2,
(18)
(g.2, Q)(2,a).
We will start by showing that the states and actions expanded by MCTS under initial 6-transformed observation 9s8, (3", a² 3 _2. ......., would exactly correspond (g.20, Da 0s=2 Daa" ....), where (=", aa, =z,,,a,_.. ..?.... are states expanded under the non-transformed observation,
By equivariance of h, 2" = h(g.s) = grh(s) g2²", as expected.
Next, we show that the actions selected by MCTS also obey G-equivariance constraint, in the sense that: if zh-1 = g then 7 = gaa²
As we assumed N2 to be G-equivariant (Case 1), must hold that Σ, (2, b) is G-invariant (as a sum- reduction of equivariant functions). Hence, we can rewrite Equation 14 as:
((= at = arg max 0i(s" P ,0), 1+ (2
(19)
where is 6-invariant, Pa is G-cquivariant by composition of functions that are G-cquivariant by assumption, and Qa is G-equivariant by assumption of Case I.
Hence, using this formula to define at, we recover:
(
at arg
B
8 (2h-1 a)) A.P.P. P...(2* ,0)
1+ No (zk-1)
<l0=
(0z2 k-1 ,a) po+* [0= 1-1
k-1 a)
arg max
+ No. (g= a)
((
Q 8. (0=2 k-1 10000 a) + p.a. (g= k-1 0a0m a)
arg max
No. " #-
c(*x-1)
Q7 0. (m)) a) P.. (2* ga a)
arg
1 1 N₂¹ (3
Da a)
E(2 (11
Qi(= (a) aa)+ + (2 k-1 a)
Qa arg max
1 + Nf(22-1,0
a
We note that we have taken the Sa out of the arg max, which is an unambiguous operation only if there is a unique action a that maximises the expression in Equation 19 To avoid breaking the symmetry in practice, propose that liebreaks for a are resolved in purely randomised fashion.
Showing this, we now only need to verify that the updates to N, and Qi (in Equation 15 are equivariant for all state-action pairs along the trajectory. Values of N and Q for all other state-action pairs will be unchanged from Nr, and therefore trivially still 6-equivariant.
First show this for N:
N NB (2k-) _ = (B aaaat)
No.6
NN. (g.2" h-1 gaa*) +1
(=*-1 ) +1
= (2 a
Hence, Case + 1 still holds for N. Now turn attention to Q.

---

Published in Transactions on Machine Learning Research (12/2023)
First, by invariance of and v, we can show that G(**) is sum of G-invariant functions and therefore also invariant. Plugging into the Q update:
Q 0 (3*-1,
_a*) "aa))""a"))))
Q44i (0224-1 BBaa))
No (g=0h-1 a 208 - (g,2 Baat) + G(gg==-1)
No (022 0aak) )+++ +1
Ni(22-1 a 20) aa G(2
Nf aaa) +1
1 ).
Henco, Case t ++1 also holds for Q. The other states stored by MCTS, such as P, are computed by di- rectly evaluating expressions in equation 12 As discussed before, these expressions are (-equivariant by composition.
Having proved that all internal state of of MCTS consistently remains transformed by & under transformed input observations, conclude that the final policy given by MCTS will be exactly G-equivariant. ]
4 Experiments and results
4.1 Environments
We consider two 2D grid-world environments, MiniPacman (Guez al. 2019 and Chaser (Cobbe et al. 2020), that feature an agent navigating in a 2D maze. In both environments, the grid-world state is rep- resented as a 2D array (i.e. an image) and an action is a direction to move (one of (-,,,,,,1)) of these grid-worlds are symmetric with respect to 90" rotations, in the sense that moving down in some map is the same as moving left the 90º clock-wise rotated version of the same map. Hence, we take the symmetry group to be C1 {e, Ug0o, U180", U2700 }. the 4-element cyclic group, which in our case represents rotating the grid-world state by all four possible multiples of 90".
Both
4.2 Model hyperparameters
The architecture of all agents we study here based on MuZero (Schrittwieser et al. 2020 with similar hyperparameters as in Hamrick al. (2020). The encoder and the transition model are implemented standard ResNet modules (He et al., 2016) with hidden dimension of 128. Each ResNet starts with 3 x convolutional layer, followed by layer normalisation (Ba et al.] 2016) and five residual blocks, each of them consisting of the following sequence of operations: layer normalisation, ReLU, 3 x 3 convolution, layer normalisation, ReLU and another convolution, to which the input of the residual block is added. After the residual blocks are applied, a final ReLU activation is used. For training the models used by MuZero, we maintain a prioritisod experience replay buffer (Schaul al.) 2015),, with batch size 512, discount factor 7 = 0.97 and trajectory length of 10 for computing n-step returns. All models are trained using the Adam SGD optimiser (Kingma & Ba 2014 with learning rate 10-3, 3, To prototype the use of equivariance in various parts of the MuZero models, augment these ResNets in the way which is described in Section
3.1 Implementationally, this augmentation reduces to applying the ResNet on several Cy-transformed views of the input, and appropriately aggregating or concatenating the results as specified.
4.3 Results
We compare our C-equivariann implementation of EqMuZero (as described in section 3.1 with a standard MuZero that uses non-equivariant components: ResNet-style networks for the encoder and transition models, and MLP- policy, value and reward models, following Hamrick et al. (2020). As the encoder and the policy of Muzeo are the only two components which require knowledge of how the symmetry group acts
Note that, conform with the API of ProcGen, Chascr also contains though they also have no direct effect the game state if played. We nate that such actions will have equivariant effects equivariant states, and hence framework as presented can support their inclusion.
actions which do transform equivariantly,
10

---

Published in Transactions on Machine Learning Research (12/2023)
Same maps
Rotated maps
Different maps
160
160
160
Equivariant MuZero
140
140
140
Equivariant MuZero, standard policy
Equivariant MuZero, standard encoder
120
120
120
Standard MuZero
100
100
100
Standard MuZero, equivariant encoder
80
80
80
40
40
8.0
0.2
0.4
0.6
0.8 updates
1.0
1.2
1.4
8.0
0.2
0.4
0.6
0.8 updates
1.0
1.2
1.4
8.0
0.2
0.4
0.6
0.8 updates
1.0
1.2
1.4
1e4
1e4
1e4
Same maps
Rotated maps
Different maps
Equivariant MuZero, standard policy
16
Equivariant standard encoder
Standard MuZero
Standard MuZero, equivarlant encoder
8.0
0.5
1.0
1.5
2.0
8.0
0.5
1.0
1.5
2.0
8.0
0.5
1.0
1.5
2.0
Learner updates
1e5
Learner updates
1e5
Learner updales
1e5
Figure 3: Results procedurally--eneraa.. MiniPacman (lop) and Chaser from ProcGen (bottom). The results are aggregated over three random seeds each model variant. The dark shaded area is one standard deviation from the mean, and the light shaded area is the minimum/ minimun/maximum /maximum value across the seeds.
on the environment, we include the following ablations in order evaluate the trade-off between end-to-end equivariance and general applicability: Standard MuZero with equivariant encoder, equivariant MuZero with standard encoder and equivariant MuZero with standard policy model.
We train each agent on a set of maps, X. To test for generalisation, we measure the agent's performance on three, progressively harder, settings. Namely, we evaluate the agent on X, with randomisod initial agent position (denoted by same our results), on the set of rotated maps R.X, where E [Roo=, Ris0s, R270- } (denoted by rolated) and, lastly, on a set of maps Y, such that Y n X = and Y RX (denoted by different).
Figure X (top) presents the results of the agents MiniPacman. First, empirically confirm that the average reward on layouts X, seen during training, matches the average reward gathered on the rotations of the same mazes, RX,, for EqMuZero. Second, we notice that changing the equivariant policy with a non- equivariant one does not significantly impact performance. However, the same swap in the encoder brings the performance of the agent down that of Standard MuZero this suggests that the structure the latent space of the transition model, when not combined with some explicit method of imposing equivariance in the encoder, does not provide noticeable benefits. Third, we notice that Equivariant MuZero is generally robust to layout variations, as the learnt high-reward behaviours also transfer to Y. At the same time, Standard MuZero significantly drops in performance for both Y and RX. We note that experiments on MiniPacman were done in a low-data scenario, using 5 maps of size 14 x 14 for training; we observed that the differences between agents diminished when all agents were trained with at least 20 times more maps.
Figure (bottom) comparcs the performance of the agents on the ProcGen game, Chaser, which has similar dynamics to MiniPacman, but larger mazes of size 64 64 and more complex action space. Due the complexity of the action space, we only use EqMuZero with standard policy, rather than fully
11

---

Published in Transactions on Machine Learning Research (12/2023)
equivariant version. We use 500 maze instances for training. Our results demonstrate that, even when the problem complexity is increased in such way, Equivariant MuZero still consistently outperforms the other agents, leading to more robust plans being discovered.
Our theoretical analysis proves that the internal decision-making of MuZero's MCTS algorithm and, hence, derived visit counts will be E-equivariant, provided its constitutent networks are G-equivariant. How- ever, as discussed in MuZero (Schrittwieser al., 2020). the actions actually executed by the MuZero agent depend on stochastically sampling the MCTS policy, whose logits are based on these visit counts. As the action sampling is stochastic, executions of MuZero with different random seeds will yield different exact ac- tion soquences, leading to different downstream performance, even when the input maps are exactly rotated. Regardless, the simplification of the function search space for MuZero's neural models arising by equivariance constraints can indeed bring about clear improvements in performance, both for rotated maps and entirely unseen ones. This is especially the case simpler, lower-data settings such as the MiniPacman environment. In the Chaser environment, the agent is exposod to a significantly visually richer input and a much larger dataset of maps. We posit that, in such setting, rotated maps do not pose substantially easier challenge than any other unseen map, yet still the equivariant versions of MuZero are capable of outperforming all other basclines we tested.
5 Limitations and future work
While the theory of Equivariant MuZero generalizes to any symmetry group, in this work we test an instance where the component neural networks satisfy the criteria for the C2 group. Leveraging neural networks equivariant to continuous rotations, such the SO(3) group, makes EqMuZero applicable to different problems, such molecular tasks. However, for parts such as the encoder, the transition model and the policy, different strategy would be required, as it is impossible to transform the input state and action with every element of an infinite group. Future work could consider enforcing the equivariance constraints via additional losses, possibly combining with an approach such as that of Park et al. (2022). keeping in mind that the theoretical guarantees will no longer apply in their current form. More generally, this work can also be composed with a module that discovers symmetries, such as in the work of Yang et al. (2023).
It is important to note the key benefit of enforcing equivariance within world model: it reduces the hypothesis space of our world model to only functions respecting the relevant symmetries. This can, clearly, only improve the data efficiency of learning the world model, and when the complexities of the RT problem align nicely with these symmetric behaviours, then our approach can be highly beneficial for downstream performance. That being said, we acknowledge that there cxist tasks requiring hand explovation (such as Montezurna's revenge), where the agent is tasked with discovering long sequences of actions, before any reward can be observed. Furthermore, these action sequences may not be amenable the same symmetries as the underlying environment (horizontal reflections in the case of Montezuma's revenge). EqMuZero's 's innovations are not sufficient to resolve such challenges, but they are likely to be complementary with any exploration-based research that is conducted in the future.
6 Conclusions
We present Equivariant MuZero, a model- based agent that is, by construction, equivariant. We theoretically verify its propertics with respect to general symmetry groups, proving the agent's 's overall oquivariance given the appropriate conditions are met by its constituent neural networks. Moreover, we empirically demostrate that an Equivariant MuZero agent that is C--equivariant generalizes to unseen rotations of the training data, and also performs more robustly on test mazes, with diminishing returns when presented with 100 times more data.
References
Joannis Antonoglou, Julian Schrittwieser, Sherjil Ozair, Thomas K Hubert, and David Silver. Planning stochastic environments with a learned model. In International Conference on Learning Representations, 2021.
12

---

Published in Transactions on Machine Learning Research (12/2023)
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv:1607.064 100.006770660, 2016.
arXiv preprint
Richard Bellman. Dynamic programming. Science, 153(3731):34333 37, 1966.
Johann Brehmer, Joey Bose, Pim De Haan, and Taco Cohen. Edgi: Equivariant diffusion for planning with embodied agents. arXiv preprint or Xiv:2500. 14,0, 2023.
Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličkovié. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXie:2104.15475 2021.
Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation benchmark reinforcement learning. In International Conference on Machine Learning, pp. 2048 2056. PMLR, 2020.
Taco Cohen and Max Welling. Group equivariant convolutional networks. In International conference on machine learning, pp. 2990-2999. PMLR, 2016a.
Taco Cohen and Max Welling. Steerable cmns. arXiv preprint arXiv:1612.08400 2016b.
Pim De Haan, Maurice Weiler, Taco Cohen, and Max Welling, Gauge equivariant mesh cmns: Anisotropic convolutions on geometric graphs. arXiv preprint orXi::2003.0422, 2020.
Alhussein Fawzi, Matej Balog, Aja Huang, Thomas Hubert, Bernardino Romera-Paredes, Mohammadamin Barekatain, Alexander Novikov, Francisco Ruiz, Julian Schritlwieser, Grzegorz Swirszez, David Silver, Demis Hassabis, and Pushmeet Kohli. Discovering faster matrix multiplication algorithms with reinforce- ment learning. Nature, 610(7930):47-53, 2022.
Arthur Guez, Mchdi Mirza, Karol Gregor, Rishabh Kabra, Sébastien Racanière, Théophanc Weber, David Raposo, Adam Santoro, Laurent Orseau, Tom Eccles, Greg Wayne, David Silver, and Timothy Lillicrap. An investigation of model-free planning. In International Conference on Machine Learning, pp. PMLR, 2019.
Jessica B Hamrick, Abram L Friesen, Foryal Behbahani, Arthur Guez, Fabio Viola, Sims Witherspoon, Thomas Anthony, Lars Buesing, Petar Veličković, and Théophane Weber. On the role of planning in model- based deep reinforcement learning. arXiv preprint arXiv:2011.04021 2020.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Julian Ibarz, Jie Tan, Chelsea Finn, Mrinal Kalakrishnan, Peter Pastor, and Sergey Levine. How to train your robot with deep reinforcement learning: lessons we have learned. The International Journal of Robotics Research, 40(4-5):698-72) 721, 2021.
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXio:2205.09991 2022.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. ar preprint
arXiv:1412.m 6980, 2014.
Amol Mandhane, Anton Zhernov, Maribeth Rauh, Chenjie Gu, Miaosen Wang, Flora Xue, Wendy Shang, Derek Pang. Rene Claus, Ching- Chiang, et al. MuZero with self-competition for rate control in VP9 video compression. arXie preprint orXiv:2202.00 06626, 2022.
Arnab Kumar Mondal, Vineet Jain, Kaleem Siddiqi, and Siamak Ravanbakhsh. EqR: Equivariant represen- tations for data-efficient reinforcement learning. In International Conference on Machine Learning, pp. 15908-15926. PMLR, 2022.
Darius Muglich, Christian Schroeder de Witt, Elise van der Pol, Shimon Whiteson, and Jakob Foerster. Equivariant networks for zero-shot coordination. arXio preprint 22..2.00.112,, 2022.
13

---

Published in Transactions on Machine Learning Research (12/2023)
Felix Musil, Andrea Grisafi, Albert P Bartók, Christoph Ortner, Gábor Csányi, and Michele Ceriotti. Physics-inspired structural representations for molecules and materials. Chemical Reviews, (16) 9815, 2021.
Jung Yeon Park, Ondrej Biza, Linfeng Zhao, Jan Willem de Meent, and Robin Walters. Learning symmetric emboddings for equivariant world models. arXiv preprint or orXiw:2004.1197 2022.
Balaraman Ravindran. An algebroic approach to abstraction in reinforcement learning. University of Mas- sachusetts Amherst, 2004.
Balaraman Ravindran and Andrew G Barto. Approximate homomorphisms: A framework for non-exact minimization in Markov decision processes. 2004.
Sahand Rezaci-Shoshtari, Rosic Zhao, Prakash Panangaden, David Meger, and Doina Procup. Continuous MDP homomorphisms and homomorphic policy gradient. orXiv preprint arXiv:2299.0740 07364, 2022.
Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver. Prioritized experience replay. arXiv preprinl arXiv:1511.05952, 2015.
Julian Schrittwieser, Toannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy Lillicrap, and David Silver. Mastering Atari, Go, chess and shogi by planning with a learned model. Nature, 588(7839):604 609, 2020.
Marwin HS Segler, Mike Preuss, and Mark P Waller. Planning chemical syntheses with deep neural networks and symbolic AI. Nature, (76986))60004 610, 2018.
Aviv Tamar, Yi Wu, Garrett Thomas, Sergey Levine, and Pieter Abbeel. Value iteration networks. Advances in neural information processing systems, 29, 2016.
Elise van der Pol, Daniel Worrall, Herke van Hoof, Frans Oliehoek, and Max Welling. MDP homomor- phic networks: Group symmetries in reinforcement learning. Advances in Neural Information Processing Systems, 33:4199-422 4210, 2020.
Elise van der Pol, Herke van Hoof, Frans A Olichock, and Max Welling. Multi-agent MDP homomorphic networks. Xiv preprint arXiv:2110.04495 2021.
Oriol Vinyals, Igor Babuschkin, Wojeicch M Czarnocki, Michaöl Mathicu, 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((782):350-333 354, 2019.
Dian Wang, Robin Walters, and Robert Platt. (0())-equvvariant reinforcement learning. arXiv preprint arXiv:2203004499 2022.
Maurice Weiler and Gabriele Cesa. General e (2)-oquivariant steerable cnns. Advances neural information processing systems, 32, 2019.
Jianke Yang, Robin Walters, Nima Dehmany, and Rose Yu. Generative adversarial symmetry discovery, 2023.
Linfeng Zhao, Xupeng Zhu, Lingzhi Kong, Robin Walters, and Lawson LS Wong, Integrating symmetry into differentiable planning with stecrable convolutions. In The Eleventh International Conference on Learning Representations, 2022.
Linfeng Zhao, Owen Lewis Howell, Jung Yeon Park, Xupeng Zhu, Robin Walters, and Lawson LS Wong. Can euclidean symmetry help in reinforcement learning and planning. 2023.
14


Reviews:
[{"review_id": "NdSmJf9XnW", "reviewer": "Reviewer_GAxJ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- The paper explores an important idea: improving data efficiency in MuZero using equivariance, which is a new and relevant approach.\n- A clear plan is proposed to achieve equivariance in MuZero, which hasn't been done before for MuZero-based approaches, making this a notable contribution.\n- The paper provides a strong guarantee that the new EqMuZero will act consistently in similar situations, backing up its claims to improve data efficiency and generalization.\n- The specific way EqMuZero is designed to handle 90-degree rotations is smart, especially for the chosen testing environments. This shows a good attempt to prove their claims.\n- The paper mentions that their approach stays strong even when not all parts of Equivariant MuZero follow strict equivariance, hinting at potential usefulness in real-world scenarios with imperfect symmetry.\n\nWeaknesses and Concerns:\n\n- The paper falls short in acknowledging an entire spectrum of work on equivariant neural networks. Using the latest techniques, the paper will be much more interesting, instead of limited to $C_4$.\n    - They have been used quite intensively in most of the equivariant model-free/model-based RL work. For this paper to bring enough interests to the field, it would be convincing to do any of the following (1) give convincing reasons why missing the entire body of works on equivariant NNs, (2) give empirical comparison that the proposed approach is better in some sense, (3) show this approach is technically interesting via other evidences.\n    - The approach in Sec 3.1 seems a special case of G-CNN and Steerable CNN — it can only consider regular representation of $C_4$ (cyclic permutation of feature channels), while it didn’t consider other representations (standard, other irreducible representations) or groups (dihedral groups). This specific implementation resembles the group equivariant convolution network (G-CNN, Cohen et al. 2016), which should be acknowledged. It is the backbone of most current equivariant network architecture. Cohen et al. later developed Steerable CNN and extended G-CNN, which can implement all these features.\n    - The proposed method seems to not be able to handle this case: Equivariant mapping from regular representation (4 channels indicating 4 directions) to standard representation (output is a 2D vector, e.g., using (0,1) to represent north). Such equivariance constraint needs to solve the space of intertwiner / steeable constraint of $W_{2 \\times 4}\\rho_{reg}(g) = \\rho_{std}(g) W_{2 \\times 4}$, where the existing work has already provided solutions, see E(2)-steerable CNN by Weiler et al. (ICLR 2021).\n- There is absence of closely related work concerning symmetry in planning/search and model-based reinforcement learning (RL).\n    - For instance, the paper \"Integrating Symmetry into Differentiable Planning with Steerable Convolutions\" by Zhao et al. (ICLR 2023) delves into $D_4$ symmetry, which could have offered a richer contextual understanding.\n    - The paper \"EDGI: Equivariant Diffusion for Planning with Embodied Agents\" by Brehmer et al. (arxiv 2023), could have provided a broader perspective on equivariant diffusion models for equivariant model–based planning.\n    - Furthermore, the paper \"Can Euclidean Symmetry Help in Reinforcement Learning and Planning?\" (arxiv 2023) discuss how a model-based RL algorithm can be made equivariant and could serve as a pertinent reference, although it's uncertain to what extent it’s related.\n- The paper is not using the full symmetry of a planar board $D_4$ dihedral group (4 rotations times 2 reflections = 8 transformations), as shown in Zhao et al. 2023. Can the author’s approach handle D_4?\n- There are some other technical concerns. For example, the policy probability $\\pi(a|s) = \\pi(g.a|g.s)$  should be invariant, while $\\pi(s)=a$ is equivariant. Also, why not using fully equivariant policy? The theoretical part only showed for equivariant version.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces a new version of the MuZero algorithm, called Equivariant MuZero (EqMuZero), which explicitly incorporates the symmetries of the environment into its world-model architecture. The main claims of the paper are:\n\n1. The authors propose a new approach to improve the data efficiency and generalization capabilities of MuZero by incorporating the symmetries of the environment into its world-model architecture. This approach is based on the concept of equivariance, which ensures that the action-selection algorithm of the EqMuZero agent will behave consistently across equivalent states.\n2. They provide a proof that guarantees that if all neural networks used by MuZero are equivariant to a particular symmetry group acting on the environment, all actions taken by the agent will also be equivariant to that same symmetry group. This means that EqMuZero can be more data-efficient than standard MuZero, as it knows how to act in states it has never seen before.\n3. The paper also presents a specific implementation of EqMuZero that is equivariant to 90-degree rotations by design. The authors claim they will test this implementation on two rotationally symmetric grid-worlds: procedurally-generated MiniPacman and the Chaser game in the ProcGen suite, to verify the generalization capabilities of EqMuZero, demonstrating the benefits of incorporating equivariance into the world-model architecture.\n4. The authors also claim that their performance improvements hold even when only some of the components of Equivariant MuZero obey strict equivariance, which would show the robustness of their construction.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Add related works and discuss them: use aforementioned papers to provide a more holistic understanding of the current landscape of equivariant neural networks, symmetry in planning/search, and model-based RL.\n- Attempt to employ existing equivariant neural network techniques (e.g., E(2)-steerable CNN by Weiler et al.) as done in Zhao et al. (ICLR 2023), or convincing explanation/comparison.\n- A thorough examination on how $D_4$ symmetry could be handled within the author’s proposed framework is useful. It could provide improve sample efficiency, which is the core claim of using equivariance in the paper.\n- A more comprehensive analysis of the equivariance in MuZero is urged, especially on using equivariant components in different places, given it's the key of the paper’s contribution.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "bZxAgguEB3", "reviewer": "Reviewer_QDcZ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths:\n* The paper has a clear hypothesis: incorporate symmetries in the environment into the MuZero architecture, and show that this modification leads to faster learning and better generalization.\n* The modifications made to MuZero to take advantage of a particular kind of symmetry ($C_4$-equivariance), are well thought out.\n* The above modifications ground the theoretical analysis that shows that MuZero will take equivariant actions in equivariant states if corresponding modifications are made for any equivariance, which in itself seems to be correct (see comments below)\n* The experiments validate the above hypothesis that MuZero learns faster and generalizes to new instances of the environment better when the equivariant components are incorporated.\n* The ablations also make sense and give confidence that all the architectural changes made are valuable, and which ones are essential.\n\n## Weaknesses\n* While the experiments show that EqMuZero can take advantage of symmetries in the environment, a question that remains is how well it would deal with scenarios where the symmetries are not a major component of what makes a problem difficult (such as left-right symmetry in say Montezuma's revenge, with a lot of actions that are not part of the equivariance, and where the exploration problem is hard regardless of the symmetries)\n* The other points of contention are fairly minor, and I mention them here with questions I have:\n  * In the background section, functions and sets are notated with the same format (capital letters).\n  * An action encoder is defined to map actions from $A \\longmapsto Z_A$, but the transition and reward functions do not take advantage of this encoding.\n  * In Figure 1, shouldn't the space after decoding be the same space as before  encoding? $S$ instead of $\\mathbb{R}^l$ in this case?\n  * In proof of Theorem 1, on page 8 when deriving $\\tilde{a}^k = g_a a^k$, I am unsure of how the last two steps are arrived at. Specifically, how $g_a$ is taken outside the $argmax$.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the problem of improving sample complexity and generalization of model-free method -- specifically MuZero -- by taking advantage of symmetries in the environment. To do so, it introduces architectural modifications in the networks that MuZero is composed of, such that the internal computations of state embeddings, transitions, and the policy, are equivariant, and the reward and Q-value computation is invariant, to the symmetries in the environment. \nThe paper proves that these architectural modifications lead to the actions taken by MuZero itself being equivariant to symmetries in the environment.\nThe experiments to validate this equivariant architecture are conducted in two domains, MiniPacman, and Chaser from the ProcGen suite of domains. These experiments seek to show that MuZero learns faster with the equivariant architecture, and the subsequent policies generalize better to symmetrically transformed maps in these games as well as novel maps that the agents have never trained on before.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Some minor changes that would be good to have:\n  * The paper specifies that equivariant actions are also needed for the symmetries to hold very late, in Section 3.1. Perhaps introduce or clarify it earlier?\n  * The paper mentions that the Chaser environment includes some non-equivariant actions. Could you mention how those are handled explicitly? I assume you do not need to do anything special since the same actions should have equivariant effects in equivariant states. But explicitly stating this effect will be helpful.", "broader_impact_concerns": "Principled ways to improve sample complexity and generalization of RL methods such as these would be helpful to ensure consistent and expected behavior of RL agents. I would imagine the broader impacts of these investigations to be mostly positive.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "TpSyIfFGYK", "reviewer": "Reviewer_XLFJ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- The submission is technically correct.\n\nWeaknesses\n- The submission is limited in applicability to toy environments and is way behind state of the art, which uses more general theoretical treatment (van der Pol'20, Bronstein'21)", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The submission proposes a version of the MuZero algorithm equivariant to the C4 group, e.g. the group of the 4 90-degree rotations. This is useful in gridworlds which are often equivariant in this way (except e.g. when gravity is present such as in atari games). Experiments show the method was implemented correctly.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The submission claims that\n>  ... agents need to learn the underlying rules governing the environment, so as to robustly generalise to conditions that differ from those they were trained on.\n\nHowever, the proposed agent fails to do this. Instead of learning the underlying rules of the environment, the rules are hard-coded. Please clarify how this sentence relates to the proposed agent.", "broader_impact_concerns": "\\-", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies how symmetries can be exploited to improve the efficiency of model-based RL methods. 

All the reviewers agree that the contribution is interesting, albeit narrow in its scope, as it only focuses on the MuZero algorithm and provides empirical evidence for one specific group symmetry.

Whereas the experimental evalaution looks to be sound, it is somewhat limited and could be improved (see requested changes below).

I report below some changes that are necessary to clear the acceptance of the paper and some additional suggestions to further improve the experimental analysis.

Requested changes (**necessary for acceptance**):
- It is not imemdiately clear from the text how many indipendent runs have been considered to generate the plots nor the meaning of the shaded area around the curves. Please specify those in the text and/or the figure captions;
- Include a deeper discussion of the reported results, e.g., why the performance for *different* maps looks to be better than *same* in the Chaser domain? Instead, why the performance declines slightly in *rotated* maps despite the equivariance?
- Report detailed experiments setup and how to reproduce the results (possibly in the Appendix).

Suggestions (**optional**):
- Consider reporting the "number of samples" instead of the "learner updates" in the x-axis of Figure 3;
- Consider releasing the code used to run the experiments, which can benefit future research on the same topic (e.g., similar experiments with more general equivariant architectures).
Decision: Accept (poster)