G AUGE E QUIVARIANT D EEP Q-L EARNING ON D ISCRETE M ANIFOLDS
A BSTRACT
Data, at any point on a manifold, can be represented on the tangent plane at that point with respect to a basis, called a gauge.
But the choice of gauge is not unique for arbitrary manifolds.
Hence, for agents traversing an environment embedded on a manifold, the same environment may appear differently if the choice of gauge changes or when moving to a different point that has a different gauge.
This may be deleterious to an agent’s learning, as compared to learning on, say, a flat grid where it is easy to choose a fixed gauge for each point.
To this end, we provide a formulation of deep Q-learning that learns policies (and Q-values) that are equivariant (invariant) to changes in choice of gauge.
This leads to an efficient learning algorithm independent of the choice of gauge
Our experimental results demonstrate significant improvement in learning on novel environments embedded in arbitrary manifolds such as spheres , hills, and urns , compared to naive approaches.
I NTRODUCTION
Consider the problem of deep Q-learning on discrete manifolds.
Data can be represented on the tangent plane of any manifold using a basis, called a gauge .
But for general manifolds, the choice of gauge is not unique.
Hence, the same geometric data may appear differently based on the choice of the gauge, as shown in Fig.
1a, where a spider (agent) crawls on the surface of a cube with each surface having a different choice of gauge.
The observations made by the spider change with the choice of gauge.
Moreover, Fig.
1b shows the challenge of parallel transporting (De Haan et al., 2020) data along different paths between the same points, resulting in different features being transported.
Hence, an agent moving between two points along different paths on a manifold may observe the data differently.
These two challenges do not arise on flat surfaces, where gauges can be fixed easily and moving parallelly along different paths to the same point yields the same observation.
We encounter both these challenges by first arbitrarily fixing gauges at each point and then providing an efficient gauge equivariant framework of deep Q-learning.
[Figure]
[Figure]
P RELIMINARIES
Equivalence and equivariance: Two points x, x ′ ∈ X are called f -equivalent for some function f : X (cid:55)→ Y if f ( x ) = f
( x ′ ) .
A function f : X (cid:55)→ Y is equivariant to sets of transformations { L g : X (cid:55)→ X| g ∈ G } and { K g : Y (cid:55)→ Y| g ∈ G } if K g ( f ( x )) =
f ( L g ( x )) for all g ∈ G , for some group G .
The function f : X (cid:55)→ Y is called invariant to the set of transformations { L g : X (cid:55)→ Y| g ∈ G } if K g is the identity function, i.e. f ( x ) = f
( L g ( x )) for all g ∈ G .
Gauge and gauge transformations We define a discrete mesh M D by a set of vertices in R 3 , with a set of faces F of tuples of vertices describing their corners.
Further, for M D to be a discrete 2D manifold, we require each edge to be connected to two faces and the neighborhood of each vertex to be homeomorphic to a disk.
At each node v ∈ V , we define a vertex normal N v , as the weighted sum of the normals of the faces that include v .
We define the plane perpendicular to N v as the tangent plane, T v M D , at v .
In this discrete case, defining the gauge at any vertex v simplifies to choosing a reference neighbor vertex, say, v q .
The frames on the tangent plane becomes the line joining v to the projection of v q on the tangent plane at v , and the line perpendicular to it.
Features at any point v on M D are represented on the tangent plane, T v M D .
Gauge on M D is defined as a position-dependent invertible linear map w v : R 2 (cid:55)→ T v M D .
Thus, if { e 1 , e 2 } is the standard basis of R 2 , then { w v ( e 1 ) , w v ( e 2 ) } defines the basis of T v M D .
Gauge transformations, g v , are point-dependent transformations of gauges on the manifold, which are described as d × d invertible matrices, GL ( d, R ) .
Depending on the transformation of matrices we consider, we can restrict the allowed g v matrices, also called reduction of the structure group
GL ( d, R ) .
In our case, we consider orientable manifolds, for which we use g v ∈ SO(2) , i.e. the set of orthogonal matrices with positive determinant in GL (2 , R ) .
MDP and MDP homomorphism A Markov Decision Process (MDP), M , is given by a tuple ( S , A , R , T , γ ) , where S is the state space, A is the action space, T gives the transition probabilities
T : S ×
A × S (cid:55)→ R ≥ 0 , and γ ∈ [0 , 1] is the discount factor.
The goal of an MDP is to find a policy π ∈ Π , π : S × A (cid:55)→ R ≥ 0 , such that (cid:80) a ∈A π ( a | s ) =
1 for all s ∈ S , which maximizes the expected reward R t =
E π [ (cid:80) Tk =
γ k r t + k +1 ] .
Two important associated terms with a policy π are its value function V π ( s ) and Q π ( s, a ) , where V π ( s ) is the expected reward at state s , and Q π ( s, a ) is the expected reward on playing a at state s under policy π .
V π ( s ) and Q π ( s, a ) are related by the
Bellman equations (Bellman, 1957) and optimal policies π ∗ result in optimal V ∗ ( s ) and Q ∗ ( s, a ) .
An MDP homomorphism, h , consists of a tuple of surjective maps ( σ, { α s | s ∈ S} ) from the state-action space S
× A of the MDP M to the abstract state-action space ¯ S × ¯ A of an MDP ¯
M = ( ¯ S , ¯ A , ¯ R , ¯ T , ¯ γ ) , where σ : S (cid:55)→ ¯ S and α s : A (cid:55)→ ¯ A .
Further, these maps must satisfy the following conditions:
[Equation]
M ETHOD
Our main insight is that the outcome of a policy for an MDP with state-action space defined on a manifold should be independent of the gauge of its representation.
That is, the policy is an intrinsic property of the environment, irrespective of the choice of gauge.
Consider the example in Fig.
1a, where a spider moves on the faces of a cube.
The same direction of motion with respect to the ambient space, R 3 , appears differently with different choices of gauges, but, intrinsically it is the same environment and the policy learned should not be independent of the choice of gauge.
Hence, we introduce policy networks (and Q-networks) that are equivariant (invariant) to the change of gauge on a manifold.
Gauge-equivariant MDP homomorphism
First we look at some equivalence relations for MDP on manifolds.
[Equation]
where ρ g ∈ GL (2 , R ) is a representation of g ∈ G , denoting a transformation of gauge.
Note that the equivalence relations obtained are similar to the ones obtained by van der Pol et al.
but, we do not assume any symmetry in our environment as in (van der Pol et al., 2020).
Symmetries are often hard to find in any system let alone an RL problem.
Our relation comes from the insight that change in representation of the system does not change its dynamics.
We define gauge equivariant MDPs by defining h : S × A (cid:55)→ O ( S × A ) that maps each state-action pair to their respective orbits under a set of gauge transformations, say, G .
From Ravindran & Barto (2001), we know that h -equivalent state-action pairs share the same optimal Q and V functions.
Moreover, there exist corresponding abstract ¯ Q and ¯ V functions that obtain these optimal values in the abstract MDP, ¯ M , obtained from the map h .
Further, policies learnt in ¯ M , can be lifted to M as shown in equation 5, which is optimal in M .
[Equation]
any
[Equation]
where ¯ a = α s ( a ) .
Using the above definition of lifting, it is easy to see that π ↑ ( a | s ) =
π ↑ ( a ′ | s ′ ) , where s ′ =
ρ g ( s ) , a ′ =
ρ s ′ g ( a ) for g ∈ G .
Thus, similar to van der Pol et al.
(2020), we have the relation π ( ρ g ( s ))
=
ρ g ( π ( s )) .
But, in our framework of meshes, the action space may be irregular, unlike in (van der Pol et al., 2020), where the problems considered mostly have fixed action space.
E.g., in a gridworld-like problem where the agent moves on a mesh to find a goal by traversing on the mesh, at every node, the possible neighbors are different and at different angles with respect to each other.
The possible orientations of the neighbors at each node are continuous and infinite.
Hence, instead of constructing equivariant policy networks like in (van der Pol et al., 2020), we construct gauge-invariant Q-networks that output the Q-values for any pair of state-action.
For Q-values, the equivariance relation on policy is the same as invariance relation on the Q-function, i.e. Q ( s, a ) =
Q ( ρ g ( s ) , ρ g ( a )) for g ∈ G .
Thus, we need Q to be invariant to possible change in gauges on a manifold, e.g. for environments lying on the surface of a cube we would need equivariance to 90 ◦ rotations, whereas for general manifolds with state-actions lying on the tangent plane we would need SO (2) -equivariance, i.e. equivariance to transformations by arbitrary 2D angles.
R ELATED W ORKS
Group equivariant networks are very well studied (Cohen & Welling, 2016; Cohen et al., 2018;
Ravanbakhsh et al., 2017).
In RL, group equivariant MDP was proposed exploiting symmetries in
RL environments (van der Pol et al., 2020).
But, searching for symmetries in data is a non-trivial problem and is an active area of research (Zhou et al., 2020; Dehmamy et al., 2021; Basu et al., 2021;
Finzi et al., 2021).
In contrast, we do not assume any symmetry in the environment.
We focus on learning independent of representation of the environment based on gauges by developing a gauge equivariant deep Q-learning framework.
E XPERIMENTS
We first describe the environments followed by the results.
Details of the equivariant network construction method and hyperparameters used are given in Appendix A.
RL Environments on a manifold We consider the basic problem of gridworld embedded on different manifolds like a sphere, hills, or an urn.
We call this environment meshworld .
Here, an agent starting at some random point on a mesh wants to find the goal, which is another randomly picked point on the mesh.
The observation of the agent is the vector starting from its location pointing to the goal, but projected on the tangent plane corresponding to the location of the agent.
The gauge at each location is chosen arbitrarily but fixed for training.
For evaluation, similarly, we choose arbitrarily the gauges at each location, but they are not necessarily the same as in training.
This is because evaluation maybe performed separately from training and there is no unique choice of gauges.
The meshes chosen are shown in Fig.
2, which were all created using the PyVista software (Sullivan &
Kaszynski, 2019) ensuring that each one has 100 nodes.
The action space at any node is the set of neighbors and the reward r ( s, a ) for any state s and action a is inversely proportional to the geodesic distance of the next state to the goal.
For updating states, we also need to take care of gauge change and avoid the problem of parallel transport illustrated in
Fig.
1b.
Going from a state s 1 at node p to a state s 2 at a neighbor q , we need to ensure that we also take care of the change in gauge from node p to q .
This can either be done by projecting the state on the new gauge after updating the 3D state vector or by changing the magnitude of the current state vector appropriately and multiplying by a change of gauge matrix, ρ ( p −→ q ) .
Results
The results of our experiments on meshworld in Fig.
2 (d), (e), (f) show huge gains in average time taken by the agent to reach the goal in each of the manifolds considered: sphere, hills, and urn.
The results indicate the advantage of using gauge equivariance in our deep Q-learning formulation.
All plots shown are results averaged over 10 runs with fixed seeds.
Moreover, the gains obtained from equivariance can be seen very early in the training, indicating better sample efficiency of our method.
We also conduct more experiments on meshes of different sizes showing similar gains, as illustrated in Appendix B.
[Figure]
C ONCLUSION
We propose a natural framework of deep Q-learning on meshes by addressing the problem of choice of gauge and parallel transport on meshes.
We release a novel RL environment called meshworld , where the environment is embedded on meshes.
We show that our framework of gauge equivariant deep Q-learning addresses the above-mentioned challenges and experimental evidence confirms huge gains over traditional deep Q-learning frameworks. A CKNOWLEDGEMENT
A E QUIVARIANT N ETWORK C ONSTRUCTION D ETAILS
For constructing gauge equivariant networks, we use SO(2)-equivariant kernels from (De Haan et al., 2020).
The input features to our network is of type 2 × ρ 1 with both state and the action represented as ρ 1 features.
We use two layers of SO(2) equivariant kernels with intermediate layers with 3 channels of representation type ( ρ 0 + ρ 1 + ρ 2 ) .
The output has 6 channels of representation type ( ρ 0 + ρ 1 ) .
We take outputs of the norms of each of the six channels and pass it through a fully connected layer with a single output, hence, making the output invariant to change in gauges.
For comparison with non-equivariant models, we use fully connected networks with two layers and nearly equal number of parameters.
For all cases, we use the Adam optimizer (Kingma & Ba, 2015) with a learning rate of 1 e -4 and weight decay of 1 e -2.
For Q-learning, we use a discount factor of 0 . 99 and for training we use ϵ -greedy strategy, where epsilon is set to 1 in the beginning and is decreased exponentially to 1 e -2 using an exponential decay of 1 e -3. [Figure]
B E XPERIMENTS WITH VARYING SIZE OF NODES IN MESHWORLD
In this section, we generalize the experiments we conduct in Sec.
5 to varying number of nodes showing that the gain obtained from gauge equivariance shows similar advantages across varying sizes of meshes.