MDP Homomorphic Networks: Group Symmetries in Reinforcement Learning

MDP with Symmetries

In a MDP with symmetries, there exists transformations on the state-action space that leave the reward and state transition invariant. Consider a state transformation $L_g : \mathcal{S}\rightarrow \mathcal{S}$ and a state-dependent action transformation $K_g^s:\mathcal{A} \rightarrow \mathcal{A}$, then invariance of the reward and state transition in a MPD with symmetries are defined as
$R(s,a)=R(L_g[s], K^s_g[a])$
$T(s' \vert s, a) = T( L_g[s'] \vert L_g[s], K^s_g[a])$

Main Problem

In traditional deep RL, symmetry is not inherently applied into the problem. If the agent collects a transition with $s$ and $a$, it does not affect the learning of $L_g[s]$ and $K^a_G[a]$. The agent would have to explore these state action pairs independently, and even then, it is not guaranteed that the trained policy network would produce a symmetric behavior. This affects the sample efficiency and learning speed of the RL algorithm as it would have to explore the whole state-action space instead of learning from equivalent transitions.

Paper's Idea

The paper aims to build the symmetry directly into the policy and value neural networks. In other words, they want the networks to be equivariant, i.e., to satisfy:
$\pi(L_g[s]) = K^s_a[\pi(s)]$
Which means that: transform the state and pass it through the network = pass the state through the network then transform action

The idea is to constrain the weight matrices of ordinary linear layers such that the network satisfies the relation above.
Starting from a single MLP layer $z'=Wz+b$, the bias can be merged into the weights: $W \leftarrow [W, b]$ and $z \leftarrow [z,1]$ to obtain $z'=Wz$. Given the symmetry transformations $(L_g, K_g)$, then the layer must satisfy:
$K_gWz=WL_gz$
By dropping $z$:
$K_gW=WL_g$
This is a linear constrain in $W$.

They propose a numerical method to find the set of matrices $W$ that satisfy the constraint:

where the symmetrizer is defined as
$S(W) \triangleq \frac{1}{|G|} \sum_{g \in G} K^{-1}_gWL_g$
Algorithm 1 produces the matrices $V_i$ that form the basis of $W$. In other words, $W$ is a linear combination of the basis weights $\{V_i\}_{i=1}^r$:
$W=\sum_{i=1}^rc_iV_i$
where the $c_i$'s are the learnable parameters of the layer.

Because each $V_i$ is equivariant, then their linear combination is equivariant. Also, the composition of equivariant functions is equivariant, then, having having multiple equivariant layers results in an equivariant network.