

\subsection{Equivalence, Invariance, and Equivariance}
\label{sec:equivalence}

\textbf{Equivalence}
If a function $f: \mathcal{X} \to \mathcal{Y}$ maps two inputs $x, x' \in \mathcal{X}$ to the same value, that is $f(x) = f(x')$, then we say that $x$ and $x'$ are $f$-equivalent. For instance, two states $s, s'$ leading to the same optimal value $V^*(s) = V^*(s')$ would be $V^*$-equivalent or \emph{optimal value equivalent}~\cite{ravindran2001symmetries}. An example of two optimal value equivalent states would be states $s$ and $L[s]$ in the CartPole example of Figure \ref{fig:symmetry_example}. The set of all points $f$-equivalent to $x$ is called the \emph{equivalence class} of $x$.

\textbf{Invariance and Symmetries}
Typically there exist very intuitive relationships between the points in an equivalence class. In the CartPole example of Figure \ref{fig:symmetry_example} this relationship is a horizontal flip about the vertical axis. This is formalized with the transformation operator $L_g: \mathcal{X} \to \mathcal{X}$, where $g \in G$ and $G$ is a mathematical group. If $L_g$ satisfies 
\begin{align}
    f(x) = f(L_g[x]), \qquad \text{for all } g \in G, x \in \mathcal{X},
\end{align}
then we say that $f$ is \emph{invariant} or \emph{symmetric} to $L_g$ and that \updated{$\{L_g\}_{g \in G}$} is a set of \emph{symmetries} of $f$ .
We can see that for the invariance equation to be satisfied, it must be that $L_g$ can only map $x$ to points in its equivalence class. Note that in abstract algebra for $L_g$ to be a true transformation operator, $G$ must contain an identity operation; that is $L_g[x] = x$ for some $g$ and all $x$. An interesting property of transformation operators which leave $f$ invariant, is that they can be composed and still leave $f$ invariant, so $L_g \circ L_h$ is also a symmetry of $f$ for all $g, h \in G$. In abstract algebra, this property is known as a \emph{semigroup property}. If $L_g$ is always invertible, this is called a \emph{group property}. In this work, we experiment with group-structured transformation operators. For more information, see~\cite{dummit2004abstract}. One extra helpful concept is that of \emph{orbits}. If $f$ is invariant to $L_g$, then it is invariant along the orbits of $G$. The orbit ${\mathcal{O}}_x$ of point $x$ is the set of points reachable from $x$ via transformation operator $L_g$:
\begin{align}
    {\mathcal{O}}_x \triangleq \{ L_g[x] \in \mathcal{X} | g \in G \}.
\end{align}

\textbf{Equivariance}
A related notion to invariance is \emph{equivariance}. Given a transformation operator $L_g: \mathcal{X} \to \mathcal{X}$ and a mapping $f: \mathcal{X} \to \mathcal{Y}$, we say that $f$ is equivariant \cite{cohen2016group, worrall2017harmonic} to the transformation if there exists a second transformation operator $K_g: \mathcal{Y} \to \mathcal{Y}$ in the output space of $f$ such that
\begin{align}
    K_g [f(x)] = f(L_g[x]), \qquad \text{for all } g \in G, x \in \mathcal{X}.
\end{align}
The operators $L_g$ and $K_g$ can be seen to describe the same transformation, but in different spaces. In fact, an equivariant map can be seen to map orbits to orbits. We also see that invariance is a special case of equivariance, if we set $K_g$ to the identity operator for all $g$. Given $L_g$ and $K_g$, we can solve for the collection of equivariant functions $f$ satisfying the equivariance constraint. Moreover, for linear transformation operators and linear $f$ a rich theory already exists in which $f$ is referred to as an \emph{intertwiner} \citep{cohen2016steerable}. In the equivariant deep learning literature, neural networks are built from interleaving intertwiners and equivariant nonlinearities. As far as we are aware, most of these methods are hand-designed per pair of transformation operators, with the exception of \cite{diaconu2019learning}. In this paper, we introduce a computational method to solve for intertwiners given a pair of transformation operators.

\subsection{Markov Decision Processes}
A Markov decision process (MDP) is a tuple $({\mathcal{S}}, {\mathcal{A}}, {R}, {T}, \gamma)$, with \emph{state space} ${\mathcal{S}}$, \emph{action space} ${\mathcal{A}}$, \emph{immediate reward function} ${R}: {\mathcal{S}} \times {\mathcal{A}} \rightarrow \mathbb{R}$, \emph{transition function} ${T}: {\mathcal{S}} \times {\mathcal{A}} \times {\mathcal{S}} \rightarrow \mathbb{R}_{\geq 0}$, and \emph{discount factor} $\gamma \in [0, 1]$. The goal of solving an MDP is to find a policy $\pi \in \Pi$, $\pi: {\mathcal{S}} \times {\mathcal{A}} \rightarrow \mathbb{R}_{\geq 0}$ (written $\pi(a | s)$), where $\pi$ normalizes to unity over the action space, that maximizes the expected return $R_t = \mathbb{E}_\pi[\sum^T_{k=0} \gamma^k r_{t+k+1}]$. The expected return from a state ${s}$ under a policy $\pi$ is given by the \textit{value function} $V^\pi$. A related object is the $Q$-value $Q^\pi$, the expected return from a state ${s}$ after taking action ${a}$ under $\pi$.  $V^\pi$ and $Q^\pi$ are governed by the well-known Bellman equations~\cite{bellman1957dynamic} (see Supplementary). In an MDP, optimal policies $\pi^*$ attain an optimal value $V^*$ and corresponding $Q$-value given by $V^*({s}) = \underset{\pi \in \Pi}{\text{ max }} V^\pi (s)$ and $Q^*({s}) = \underset{\pi \in \Pi}{\text{ max }} Q^\pi (s)$.


\begin{figure}
   
    \begin{center}
    \includegraphics[width=0.65\textwidth]{figures/mdp_hom.pdf}
    \end{center}
    \caption{Example of a reduction in an MDP's state-action space under an MDP homomorphism $h$. Here `equivalence' is represented by a reflection of the dynamics in the vertical axis. This equivalence class is encoded by $h$ by mapping all equivalent state-action pairs to the same abstract state-actions.}
    \label{fig:mdphom}
    
\end{figure}

\paragraph{MDP with Symmetries}
\updated{Symmetries can appear in MDPs. For instance, in Figure \ref{fig:mdphom} CartPole has a reflection symmetry about the vertical axis. Here we define an \emph{MDP with symmetries}. In an MDP with symmetries there is a set of transformations on the state-action space, which leaves the reward function and transition operator invariant. We define a state transformation and a state-dependent action transformation as ${L_g}: {\mathcal{S}} \to {\mathcal{S}}$ and ${K_g^s}: {\mathcal{A}} \to {\mathcal{A}}$ respectively. Invariance of the reward function and transition function is then characterized as
\begin{align}
    R(s, a) &= R({L_g}[s], {K_g^s}[a]) && \text{for all } g \in G, s \in {\mathcal{S}}, a \in {\mathcal{A}} \\
    T(s' | s, a) &= T({L_g}[s'] | {L_g}[s], {K_g^s}[a]) && \text{for all } g \in G, s \in {\mathcal{S}}, a \in {\mathcal{A}}.
\end{align}
Written like this, we see that in an MDP with symmetries the reward function and transition operator are invariant along orbits defined by the transformations $({L_g}, {K_g^s})$. }

\paragraph{MDP Homomorphisms}
\updated{MDPs with symmetries are closely related to MDP homomorphisms, as we explain below. First we define the latter. An \emph{MDP homomorphism} $h$~\cite{ravindran2004approximate, ravindran2001symmetries} is a mapping from one MDP ${M} = ({\mathcal{S}}, {\mathcal{A}}, {R}, {T}, \gamma)$ to another ${\bar{M}} = ({\mathcal{\bar{S}}}, {\mathcal{\bar{A}}}, {\bar{R}}, {\bar{T}}, \gamma)$ defined by a surjective map from the state-action space ${\mathcal{S}} \times {\mathcal{A}}$ to an \emph{abstract state-action space} ${\mathcal{\bar{S}}} \times {\mathcal{\bar{A}}}$. In particular, $h$ consists of a tuple of surjective maps ${(\smap, \amapset)}$, where we have the state map ${\sigma}: {\mathcal{S}} \to {\mathcal{\bar{S}}}$ and the state-dependent action map ${\alpha_s}: {\mathcal{A}} \to {\mathcal{\bar{A}}}$. These maps are built to satisfy the following conditions}
\begin{align}
    {\bar{R}}({\sigma}(s), {\alpha_s}(a)) &\triangleq {R}({s}, {a}) \hspace{1em} &&\text{for all } {s} \in {\mathcal{S}},  {a} \in {\mathcal{A}}, \label{eq:reward-map} \\
    {\bar{T}}({\sigma}({s}') | {\sigma}({s}), {\alpha_s}(a)) &\triangleq \sum_{{s}'' \in {\sigma}^{-1}(s')} {T}({s}'' | {s}, {a}) \hspace{1em} &&\text{for all } {s}, {s}' \in {\mathcal{S}}, {a} \in {\mathcal{A}}. \label{eq:transition-map}
\end{align}
\updated{
An exact MDP homomorphism provides a model equivalent abstraction~\cite{li2006towards}. Given an MDP homomorphism $h$, two state-action pairs $(s, a)$ and $(s', a')$
are called \emph{$h$-equivalent} if ${\sigma}(s) = {\sigma}(s')$ and ${\alpha_s}(a) = \alpha_{s'}(a')$.
Symmetries and MDP homomorphisms are connected in a natural way:  }
\updated{If
an MDP has symmetries $L_g$ and $K_g$, the above equations (4) and (5) hold.
This means that we can define a corresponding MDP homomorphism, which we define next.}
\paragraph{Group-structured MDP Homomorphisms}
\updated{Specifically, for an MDP with symmetries, we can define an abstract state-action space, by mapping $(s,a)$ pairs to (a representative point of) their equivalence class $(\sigma(s), \alpha_s(a))$. That is, state-action pairs and their transformed version are mapped to the same abstract state in the reduced MDP:
\begin{align}
    \left (\sigma(s), \alpha_s(a) \right ) &= \left (\sigma(L_g[s]), \alpha_{L_g[s]}(K_g^s[a]) \right ) \hspace{1em} \forall g \in G, s \in \mathcal{S}, a \in \mathcal{A}
\end{align}
In this case, we call the resulting MDP homomorphism \emph{group structured}.
In other words, all the state-action pairs in an orbit defined by a group transformation are mapped to the same abstract state by a group-structured MDP homomorphism.
}

\textbf{Optimal Value Equivalence and Lifted Policies}
\updated{
$h$-equivalent state-action pairs share the same optimal
$Q$-value and optimal value function~\cite{ravindran2001symmetries}.
}
Furthermore, there exists an abstract
optimal $Q$-value $\bar{Q}^*$ and abstract optimal value function $\bar{V}^*$,
such that $Q^*({s}, {a}) = \bar{Q}^*({\sigma}({s}), {\alpha_s}({a}))$ and $V^*({s}) =
\bar{V}^*({\sigma}({s}))$. This is known as \emph{optimal value
equivalence}~\cite{ravindran2001symmetries}. Policies can thus be optimized in
the simpler abstract MDP. The optimal abstract policy $\bar{\pi}(\bar{a} |
\sigma(s))$ can then be pulled back to the original MDP using a procedure
called \emph{lifting}~\footnote{Note that we use the terminology \emph{lifting}
to stay consistent with~\cite{ravindran2001symmetries}.}. The lifted policy is
given in \updated{Equation~\ref{eq:lifted-policy_1}}. A lifted optimal abstract
policy is also an optimal policy in the original
MDP~\cite{ravindran2001symmetries}. Note that while other lifted policies
exist, we follow~\cite{ravindran2001symmetries, ravindran2004approximate} and
choose the lifting that divides probability mass uniformly over the preimage:
\begin{align}
    {\pi^\uparrow}(a | s) \triangleq \frac{\bar{\pi}(\bar{a} | \sigma(s))}{|\{a \in {\alpha_s^{-1}}(\bar{a})\}|}, \qquad \text{for any } s \in {\mathcal{S}} \text{ and } 
    a \in {\alpha_s^{-1}}(\bar{a}). \label{eq:lifted-policy_1}
\end{align}



\subsection{Environments}
For each environment we show ${\mathcal{S}}$ and ${\mathcal{A}}$ with respective representations of the group transformations.

\textbf{CartPole}
In the classic pole balancing task~\cite{barto1983neuronlike}, we used a two-element group of reflections about the $y$-axis. We used OpenAI's Cartpole-v1~\cite{brockman2016openai} implementation, which has a 4-dimensional observation vector: (cart position $x$, pole angle $\theta$, cart velocity $\dot{x}$, pole velocity $\dot{\theta}$). The (discrete) action space consists of applying a force left and right $(\leftarrow, \rightarrow)$. We chose this example for its simple symmetries.

\textbf{Grid world}
We evaluated on a toroidal 7-by-7 predator-prey grid world with agent-centered coordinates. The prey and predator are randomly placed at the start of each episode, lasting a maximum of 100 time steps. The agent's goal is to catch the prey, which takes a step in a random compass direction with probability $0.15$ and stands still otherwise. Upon catching the prey, the agent receives a reward of +1, and -0.1 otherwise. The observation is a $21\times21$ binary image identifying the position of the agent in the center and the prey in relative coordinates. See Figure~\ref{fig:gridworld}. This environment was chosen due to its four-fold rotational symmetry.

\begin{table}[t]
    \centering
    \caption{\textsc{Environments and Symmetries}: We showcase a visual guide of the state and action spaces for each environment along with the effect of the transformations. Note, the symbols should not be taken to be hard mathematical statements, they are merely a visual guide for communication.}
    \label{tab:my_label}
    \resizebox{\columnwidth}{!}{%
    \begin{tabular}{rlll}
    \toprule
        Environment  && Space  & Transformations \\
    \midrule
        CartPole & ${\mathcal{S}}$ &$(x,\theta,\dot{x},\dot{\theta})$    & $(x,\theta,\dot{x},\dot{\theta}),(-x,-\theta,-\dot{x},-\dot{\theta})$ \\
        & ${\mathcal{A}}$ &($\leftarrow$, $\rightarrow$)      &($\leftarrow$, $\rightarrow$), ($\rightarrow$, $\leftarrow$) \\
        Grid World & ${\mathcal{S}}$ & $\{0,1\}^{21 \times 21}$    & Identity, \qquad\;\;\; $\curvearrowright 90^\circ$, \qquad\qquad$\curvearrowright 180^\circ$, \qquad\quad $\curvearrowright 270^\circ$ \\
        & ${\mathcal{A}}$ & $(\varnothing, \uparrow, \rightarrow, \downarrow, \leftarrow)$      &$(\varnothing, \uparrow, \rightarrow, \downarrow, \leftarrow), (\varnothing, \rightarrow, \downarrow, \leftarrow, \uparrow),(\varnothing, \downarrow, \leftarrow, \uparrow, \rightarrow),(\varnothing, \leftarrow, \uparrow, \rightarrow, \downarrow)$ \\
        Pong & ${\mathcal{S}}$ & $\{0,..., 255\}^{4\times 80\times 80}$    & Identity, \qquad\qquad reflect
        \\
        & ${\mathcal{A}}$ &($\varnothing$, $\varnothing$, $\uparrow$, $\downarrow$, $\uparrow$, $\downarrow$)      &($\varnothing$, $\varnothing$, $\uparrow$, $\downarrow$, $\uparrow$, $\downarrow$), ($\varnothing$, $\varnothing$, $\downarrow$, $\uparrow$, $\downarrow$, $\uparrow$)  \\
    \bottomrule
    \end{tabular}
    }
\end{table}

\textbf{Pong}
We evaluated on the RLPYT~\cite{stooke2019rlpyt} implementation of Pong. In our experiments, the observation consisted of the 4 last observed frames, with upper and lower margins cut off and downscaled to an $80\times80$ grayscale image. In this setting, there is a flip symmetry over the horizontal axis: if we flip the observations, the up and down actions also flip.
A curious artifact of Pong is that it has duplicate (up, down) actions, which means that to simplify matters, we mask out the policy values for the second pair of (up, down) actions. We chose Pong because of its higher dimensional state space.
\updated{Finally, for Pong we additionally compare to two data augmentation baselines: stochastic data augmentation, where for each state, action pair we randomly transform them or not before feeding them to the network, and the second an equivariant version of~\cite{kostrikov2020image} and similar to~\cite{silver2016mastering}, where both state and transformed state are input to the network. The output of the transformed state is appropriately transformed, and both policies are averaged.}

\subsection{Models}
We implemented MDP homomorphic networks on top of two base architectures: MLP and CNN (exact architectures in Supplementary). We further experimented with an equivariant feature extractor, appended by a non-equivariant network, to isolate where equivariance made the greatest impact. 

\textbf{Basis Networks}
We call networks whose weights are linear combinations of basis weights \emph{basis networks}. As an ablation study on all equivariant networks, we sought to measure the effects of the basis training dynamics. We compared an \emph{equivariant} basis against a pure \emph{nullspace} basis\updated{, i.e. an explicitly non-symmetric basis} using the right-null vectors from the equivariant layer construction, and a \emph{random} basis, where we skip the symmetrization step in the layer construction and use the full rank basis. Unless stated otherwise, we reduce the number of `channels' in the basis networks compared to the regular networks by dividing by the square root of the group size, ending up with a comparable number of trainable parameters.

\newlength{\heightofhw}
\begin{figure}[t]
    \centering
    \begin{subfigure}[b]{0.325\textwidth}
        \includegraphics[width=\textwidth]{figures/cartpole_versus_bases.pdf}
        \caption{Cartpole-v1: Bases}
        \label{fig:cartpole_bases}
    \end{subfigure}
    \hspace{-5pt}
    \begin{subfigure}[b]{0.325\textwidth}
        \includegraphics[width=0.92\textwidth, trim=3em 0 0 0, clip]{figures/cartpole_versus_mlp.pdf}
        \caption{Cartpole-v1: MLPs}
        \label{fig:cartpole_mlp}
    \end{subfigure}
    \begin{subfigure}[b]{0.325\textwidth}
        \includegraphics[width=\textwidth]{figures/pong_versus_all.pdf}
        \caption{Pong}
        \label{fig:pong_conv}
    \end{subfigure}
    \caption{\textsc{Cartpole}: Trained with PPO, all networks fine-tuned over 7 learning rates. 25\%, 50\% and 75\% quantiles over 25 random seeds shown. a) Equivariant, random, and nullspace bases. b) Equivariant basis, and two MLPs with different degrees of freedom. \textsc{Pong}: Trained with A2C, all networks tuned over 3 learning rates. 25\%, 50\% and 75\% quantiles over 15 random seeds shown c) Equivariant, nullspace, and random bases, and regular CNN for Pong.}
    \label{exp:cartpole}
\end{figure}

\subsection{Results and Discussion}
We show training curves for CartPole in~\ref{fig:cartpole_bases}-\ref{fig:cartpole_mlp}, Pong in Figure~\ref{fig:pong_conv} and for the grid world in Figure~\ref{fig:grid}. Across all experiments we observed that the MDP homomorphic network outperforms both the non-equivariant basis networks and the standard architectures, in terms of convergence speed.  
\begin{wrapfigure}{r}{0.3\textwidth}
\vspace{-0.3cm}
\includegraphics[width=0.3\textwidth]{figures/crop_pongdataaug.pdf}
\caption{Data augmentation comparison on Pong.}
\label{fig:dataaug}
\vspace{-0.5cm}
\end{wrapfigure}
This confirms our motivations that building symmetry-preserving policy networks leads to faster convergence.
\updated{Additionally, when compared to the data augmentation baselines in Figure~\ref{fig:dataaug},
using equivariant networks is more beneficial. This is consistent with other results in the equivariance literature~\cite{bekkers2018roto,weiler_2018,winkels20183d,worrall2017harmonic}. While data augmentation can be used to create a larger dataset by exploiting symmetries, it does not directly lead to effective parameter sharing (as our approach does).} Note, in Pong we only train the first 15 million frames to highlight the difference in the beginning; in constrast, a typical training duration is 50-200 million frames~\cite{mnih2016asynchronous, stooke2019rlpyt}.

For our ablation experiment, we wanted to control for the introduction of bases. It is not clear \emph{a priori} that a network with a basis has the same gradient descent dynamics as an equivalent `basisless' network. We compared equivariant, non-equivariant, and random bases, as mentioned above. We found the equivariant basis led to the fastest convergence. Figures~\ref{fig:cartpole_bases} and~\ref{fig:pong_conv} show that for CartPole and Pong the nullspace basis converged faster than the random basis. In the grid world there was no clear winner between the two. This is a curious result, requiring deeper investigation in a follow-up. 

For a third experiment, we investigated what happens if we sacrifice complete equivariance of the policy. This is attractive because it removes the need to find a transformation operator for a flattened output feature map. Instead, we only maintained an equivariant feature extractor, compared against a basic CNN feature extractor. The networks built on top of these extractors were MLPs. The results, in Figure~\ref{fig:pong_conv}, are two-fold: 1) Basis feature extractors converge faster than standard CNNs, and 2) the equivariant feature extractor has fastest convergence. We hypothesize the equivariant feature extractor is fastest as it is easiest to learn an equivariant policy from equivariant features. \\
\updated{We have additionally compared an equivariant feature extractor to a regular convolutional network on the Atari game Breakout, where the difference between the equivariant network and the regular network is much less pronounced. For details, see Appendix C.}

\begin{figure}[t]
    \centering
    \begin{subfigure}[b]{0.21\textwidth}
        \includegraphics[width=\textwidth]{figures/gridworld.pdf}
        \vspace{0pt}
        \caption{Symmetries}
        \label{fig:gridworld}
    \end{subfigure}
    ~
    \begin{subfigure}[b]{0.35\textwidth}
        \includegraphics[width=\textwidth]{figures/grid_versus_bases_crop.pdf}
        \caption{Grid World: Bases}
        \label{fig:pong_bases}
    \end{subfigure}
    \hspace{-3pt}
    \begin{subfigure}[b]{0.35\textwidth}
        \includegraphics[width=0.92\textwidth, trim=3em 0 0 0, clip]{figures/grid_versus_conv_crop.pdf}
        \caption{Grid World: CNNs}
        \label{fig:grid_world}
    \end{subfigure}
    \caption{\textsc{Grid World}: Trained with A2C, all networks fine-tuned over 6 learning rates. 25\%, 50\% and 75\% quantiles over 20 random seeds shown. a) showcase of symmetries, b) Equivariant, nullspace, and random bases c) plain CNN and equivariant CNN.}
    \label{fig:grid}
\end{figure}




\section{Introduction}
\input{introduction}

\section{Background}
\input{background} 

\section{Method}
\input{method}

\section{Experiments}
\input{experiments}

\section{Related Work}
\input{related}

\section{Conclusion}
\input{conclusion}

\section{Acknowledgments and Funding Disclosure }
\input{acknowledgments}


\section{Broader Impact}
\input{impact_statement}


\medskip
\small




\subsection{Lifted Policies Are Invariant}
Lifted policies in symmetric MDPs with group-structured symmetries are invariant under the group of symmetries. \updated{Consider the following: Take an MDP with symmetries defined by transformation operators $({L_g}, {K_g^s})$ for $g\in G$. Now, if we take $s' = L_g[s]$ and $a' = K_g^{s}[a]$ for any $g\in G$, $(s', a')$ and $(s, a)$ are h-equivalent under the corresponding MDP homomorphism $h = {(\smap, \amapset)}$.} So
\begin{align}
    {\pi^\uparrow}(a | s) = \frac{\bar{\pi}({\alpha_s}(a) | {\sigma}({s}))}{|\{a \in {\alpha_s^{-1}}(\bar{a})\}|} = \frac{\bar{\pi}(\alpha_{s'}(a') | {\sigma}(s'))}{|\{a' \in \alpha_{s'}^{-1}(\bar{a})\}|} = {\pi^\uparrow}(a' | s'), \label{eq:lifted-policy2}
\end{align}
for all $s \in {\mathcal{S}}, a \in {\mathcal{A}}$ and $g \in G$. In the first equality we have used the definition of the lifted policy. In the second equality, we have used the definition of $h$-equivalent state-action pairs, where ${\sigma}(s) = {\sigma}({L_g}(s))$ and ${\alpha_s}(a) = \alpha_{s'}(a')$. In the third equality, we have reused the definition of the lifted policy. Thus we see that, written in this way, the lifted policy is invariant under state-action transformations $({L_g}, {K_g^s})$. This equation is very general and applies for all group-structured state-action transformations. For a finite action space, this statement of invariance can be re-expressed as a statement of equivariance, by considering the vectorized policy.

\textbf{Invariant Policies On Finite Action Spaces Are Equivariant Vectorized Policies}
For convenience we introduce a vector of probabilities for each of the discrete actions under the policy
\begin{align}
    \bm{\pi}(s) \triangleq \begin{bmatrix}
        \pi(a_1 | s), & \pi(a_2 | s), & ..., & \pi(a_N | s)
    \end{bmatrix}^\top,
\end{align}
where $a_1, ..., a_N$ are the $N$ possible discrete actions in action space ${\mathcal{A}}$. The action transformation ${K_g^s}$ maps actions to actions invertibly. Thus applying an action transformation to the vectorized policy permutes the elements. We write the corresponding permutation matrix as ${\textbf{K}}_g$. Note that
\begin{align}
    {\textbf{K}}_g^{-1} \bm{\pi} (s) \triangleq \begin{bmatrix}
        \pi({K_g^s}[a_1] | s), & \pi({K_g^s}[a_2] | s), & ..., & \pi({K_g^s}[a_N] | s)
    \end{bmatrix}^\top,
\end{align}
where writing the inverse ${\textbf{K}}_g^{-1}$ instead of ${\textbf{K}}_g$ is required to maintain the property ${\textbf{K}}_g {\textbf{K}}_h = {\textbf{K}}_{gh}$. The invariance of the lifted policy can then be written as $\bm{\pi}^\uparrow(s) = {\textbf{K}}_g^{-1}\bm{\pi}^\uparrow({L_g}[s])$, which can be rearranged to the equivariance equation
\begin{align}
    {\textbf{K}}_g\bm{\pi}^\uparrow(s) = \bm{\pi}^\uparrow({L_g}[s]) \qquad \text{for all } g \in G, s \in {\mathcal{S}}, a \in {\mathcal{A}}. \label{eq:constraint}
\end{align}
This equation shows that the lifted policy must satisfy an equivariance constraint. In deep learning, this has already been well-explored in the context of supervised learning~\cite{cohen2016group, cohen2016steerable, worrall2017harmonic, worrall2019deep, weiler2018learning}. Next, we present a novel way to construct such networks.


\subsection{Building MDP Homomorphic Networks}
Our goal is to build neural networks that follow Eq.~\ref{eq:constraint}; that is, we wish to find neural networks that are \emph{equivariant} under a set of state and policy transformations. Equivariant networks are common in supervised learning~\cite{cohen2016group, cohen2016steerable, worrall2017harmonic, worrall2019deep, weiler2018learning, weiler2019general}. For instance, in semantic segmentation shifts and rotations of the input image result in shifts and rotations in the segmentation. 
A neural network consisting of only equivariant layers and non-linearities is equivariant as a whole, too\footnote{\updated{See Appendix B for more details.}}~\cite{cohen2016group}. Thus, once we know how to build a single equivariant layer, we can simply stack such layers together. \updated{Note that this is true regardless of the representation of the group, i.e. this works for spatial transformations of the input, feature map permutations in intermediate layers, and policy transformations in the output layer. For the experiments presented in this paper, we use the same group representations for the intermediate layers as for the output, i.e. permutations.} For finite groups, such as cyclic groups or permutations, pointwise nonlinearities preserve equivariance \cite{cohen2016group}. 

In the past, learnable equivariant layers were designed by hand for each transformation group individually~\cite{cohen2016group,cohen2016steerable,worrall2017harmonic,worrall2019deep,winkels20183d,weiler2018learning,weiler2019general}. This is time-consuming and laborious. Here we present a novel way to build learnable linear layers that satisfy equivariance automatically.

\textbf{Equivariant Layers}
\begin{algorithm}[b]
  \caption{Equivariant layer construction}
  \label{alg:symmetrization}
  \begin{algorithmic}[1]
    \State Sample $N$ weight matrices ${\textbf{W}}_1, {\textbf{W}}_2, ..., {\textbf{W}}_N \sim \mathcal{N}({\textbf{W}}; \textbf{0}, \textbf{I})$ for $N \geq \text{dim}({\mathcal{W}_\text{total}})$
    \State Symmetrize samples: $\bar{{\textbf{W}}}_i = S({\textbf{W}}_i)$ for $i=1,...,N$
    \State Vectorize samples and stack as $\bar{{\textbf{W}}} = [{\text{vec}}(\bar{{\textbf{W}}}_1), {\text{vec}}(\bar{{\textbf{W}}}_2), ...]$
    \State Apply SVD: $\bar{{\textbf{W}}} = {\textbf{U}} {\bm{\Sigma}} {\textbf{V}}^{\top}$
    \State Keep first $r = \text{rank}(\bar{{\textbf{W}}})$ right-singular vectors (columns of ${\textbf{V}}$) and unvectorize to  shape of ${\textbf{W}}_i$
  \end{algorithmic}
\end{algorithm}
We begin with a single linear layer ${\textbf{z}}' = {\textbf{W}}{\textbf{z}} + {\textbf{b}}$, where ${\textbf{W}} \in \mathbb{R}^{D_\text{out} \times D_\text{in}}$ and ${\textbf{b}} \in \mathbb{R}^{D_\text{in}}$ is a bias. To simplify the math, we merge the bias into the weights so ${\textbf{W}} \mapsto [{\textbf{W}}, {\textbf{b}}]$ and ${\textbf{z}} \mapsto [{\textbf{z}}, 1]^\top$. We denote the space of the augmented weights as ${\mathcal{W}_\text{total}}$. For a given pair of linear group transformation operators in matrix form $({\textbf{L}_g}, {\textbf{K}_g})$, where ${\textbf{L}_g}$ is the input transformation and ${\textbf{K}_g}$ is the output transformation, we then have to solve the equation
\begin{align}
    {\textbf{K}_g} {\textbf{W}} {\textbf{z}} = {\textbf{W}} {\textbf{L}_g} {\textbf{z}}, \qquad \text{for all } g \in G, {\textbf{z}} \in \mathbb{R}^{D_\text{in} + 1}. \label{eq:equivariant-layer}
\end{align}
Since this equation is true for all ${\textbf{z}}$ we can in fact drop ${\textbf{z}}$ entirely. Our task now is to find all weights ${\textbf{W}}$ which satisfy Equation \ref{eq:equivariant-layer}. We label this space of equivariant weights as ${\mathcal{W}}$, defined as
\begin{align}
    {\mathcal{W}} \triangleq \{{\textbf{W}} \in {\mathcal{W}_\text{total}} \mid {\textbf{K}_g} {\textbf{W}} = {\textbf{W}} {\textbf{L}_g}, \text{ for all } g\in G \}, \label{eq:equivariant-subspace}
\end{align}
again noting that we have dropped ${\textbf{z}}$. To find the space ${\mathcal{W}}$ notice that for each $g \in G$ the constraint ${\textbf{K}_g} {\textbf{W}} = {\textbf{W}} {\textbf{L}_g}$ is in fact linear in ${\textbf{W}}$. Thus, to find ${\mathcal{W}}$ we need to solve a set of linear equations in ${\textbf{W}}$. For this we introduce a construction, which we call a \emph{symmetrizer} $S({\textbf{W}})$.  The symmetrizer is
\begin{align}
    S({\textbf{W}}) \triangleq \frac{1}{|G|}\sum_{g \in G} {\textbf{K}_g^{-1}} {\textbf{W}} {\textbf{L}_g}. \label{eq:symmetrizer}
\end{align}

$S$ has three important properties, of which proofs are provided in Appendix A. First, $S({\textbf{W}})$ is \emph{symmetric} ($S({\textbf{W}})\in {\mathcal{W}}$). Second, $S$ \emph{fixes} any symmetric ${\textbf{W}}$: (${\textbf{W}} \in {\mathcal{W}} \implies S({\textbf{W}}) = {\textbf{W}}$). These properties show that $S$ projects arbitrary ${\textbf{W}} \in {\mathcal{W}_\text{total}}$ to the equivariant subspace ${\mathcal{W}}$.

Since ${\mathcal{W}}$ is the solution set for a set of simultaneous linear equations, ${\mathcal{W}}$ is a linear subspace of the space of all possible weights ${\mathcal{W}_\text{total}}$. Thus each ${\textbf{W}} \in {\mathcal{W}}$ can be parametrized as a linear combination of basis weights $\{{\textbf{V}}_i\}_{i=1}^r$, where $r$ is the rank of the subspace and $\text{span}(\{{\textbf{V}}_i\}_{i=1}^r) = {\mathcal{W}}$. To find as basis for ${\textbf{W}}$, we take a Gram-Schmidt orthogonalization approach. We first sample weights in the total space ${\mathcal{W}_\text{total}}$ and then project them into the equivariant \begin{wrapfigure}{r}{0.3\textwidth}
        \includegraphics[width=0.3\textwidth]{figures/filters_in.pdf}
        \caption{Example of 4-way rotationally symmetric filters.}
        \label{fig:grid_filters}
\end{wrapfigure}subspace with the symmetrizer. We do this for multiple weight matrices, which we then stack and feed through a singular value decomposition to find a basis for the equivariant space. This procedure is outlined in Algorithm \ref{alg:symmetrization}. Any equivariant layer can then be written as a linear combination of bases
\begin{align}
    {\textbf{W}} = \sum_{i=1}^r c_i {\textbf{V}}_i,
\end{align}
where the $c_i$'s are learnable scalar coefficients, $r$ is the rank of the equivariant space, and the matrices ${\textbf{V}}_i$ are the basis vectors, formed from the reshaped right-singular vectors in the SVD. An example is shown in Figure~\ref{fig:grid_filters}. To run this procedure, all that is needed are the transformation operators ${\textbf{L}_g}$ and ${\textbf{K}_g}$. Note we do not need to know the explicit transformation matrices, but just to be able to perform the mappings ${\textbf{W}} \mapsto {\textbf{W}}{\textbf{L}_g}$ and ${\textbf{W}} \mapsto {\textbf{K}_g^{-1}}{\textbf{W}}$. For instance, some matrix ${\textbf{L}_g}$ rotates an image patch, but we could equally implement ${\textbf{W}}{\textbf{L}_g}$ using a built-in rotation function. Code is available~\footnote{\url{https://github.com/ElisevanderPol/symmetrizer/}}.

