Package: graph-conv-memory
Graph Convolutional Memory using Topological PriorsDescriptionGraph convolutional memory (GCM) is graph-structured memory that may be applied to reinforcement learning to solve POMDPs, replacing LSTMs or attention mechanisms. GCM allows you to embed your domain knowledge in the form of connections in a graph. See thefull paperfor further details. This repo contains the GCM library implementation for use in your projects. To replicate the experiments from the paper, please seethis repository instead.If you use GCM, please cite the paper!@article{morad2021graph,
  title={Graph Convolutional Memory using Topological Priors},
  author={Morad, Steven D and Liwicki, Stephan and Kortvelesy, Ryan and Mecca, Roberto and Prorok, Amanda},
  journal={arXiv preprint arXiv:2106.14117},
  year={2021}
}InstallationGCM is installed usingpip. The dependencies must be installed manually, as they target your specific architecture (with or without CUDA).Conda installFirst installtorch >= 1.8.0andtorch-geometricdependencies, thengcm:conda install torch
conda install pytorch-geometric -c rusty1s -c conda-forge
pip install graph-conv-memoryPip installPlease follow thetorch-geometric install guide, thenpip install graph-conv-memoryQuickstartBelow is a quick example of how to use GCM in a basic RL problem:importtorchimporttorch_geometricfromgcm.gcmimportDenseGCMfromgcm.edge_selectors.temporalimportTemporalBackedge# graph_size denotes the maximum number of observations in the graph, after which# the oldest observations will be overwritten with newer observations. Reduce this number to# reduce memory usage.graph_size=128# Define the GNN used in GCM. The following is the one used in the paper# Make sure you define the first layer to match your observation spaceclassGNN(torch.nn.Module):"""A simple two-layer graph neural network"""def__init__(self,obs_size,hidden_size=32):super().__init__()self.gc0=torch_geometric.nn.DenseGraphConv(obs_size,hidden_size)self.gc1=torch_geometric.nn.DenseGraphConv(hidden_size,hidden_size)self.act=torch.nn.Tanh()defforward(self,x,adj,weights,B,N):x=self.act(self.gc0(x,adj))returnself.act(self.gc1(x,adj))# Build GNN that GCM uses internallyobs_size=8gnn=GNN(obs_size)# Create the GCM using our GNN and edge selection criteria. TemporalBackedge([1]) will link observation o_t to o_{t-1}.# See `gcm.edge_selectors` for different kinds of priors suitable for your specific problem. Do not be afraid to implement your own!gcm=DenseGCM(gnn,edge_selectors=TemporalBackedge([1]),graph_size=graph_size)# If the hidden state m_t is None, GCM will initialize one for you# only do this at the beginning, as GCM must track and update the hidden# state to function correctly## You can inspect m_t, as it is just a graph of observations# the first element is the node feature matrix and the second is the adjacency matrixm_t=Nonefortintrain_timestep:# Obs at timestep t should be a tensor of shape (batch_size, obs_size)# obs = my_env.step()belief,m_t=gcm(obs,m_t)# GCM provides a belief state -- a combination of all past observational data relevant to the problem# What you likely want to do is put this state through actor and critic networks to obtain# action and value estimatesaction_logits=logits_nn(belief)state_value=vf_nn(belief)We provide a few edge selectors, which we briefly detail here:gcm.edge_selectors.temporal.TemporalBackedge# Connections to the past. Give it [1,2,4] to connect each# observation t to t-1, t-2, and t-4.gcm.edge_selectors.dense.DenseEdge# Connections to all past observations# observation t is connected to t-1, t-2, ... 0gcm.edge_selectors.distance.EuclideanEdge# Connections to observations within some max_distance# e.g. if l2_norm(o_t, o_k) < max_distance, create an edgegcm.edge_selectors.distance.CosineEdge# Like euclidean edge, but using cosine similarity insteadgcm.edge_selectors.distance.SpatialEdge# Euclidean distance, but only compares slices from the observation# this is useful if you have an 'x' and 'y' dimension in your observation# and only want to connect nearby entries## You can also implement the identity priors using this by setting# max_distance to something like 1e-6gcm.edge_selectors.learned.LearnedEdge# Learn an edge function from the data# Will randomly sample edges and train thru gradient descent# call the constructor with the output size of your GNNRay Quickstart (WIP)We provide a ray rllib wrapper around GCM, see the example below for how to use itimportunittestimporttorchimporttorch_geometricimportrayfromrayimporttunefromgcm.ray_gcmimportRayDenseGCMfromgcm.edge_selectors.temporalimportTemporalBackedgeclassGNN(torch.nn.Module):"""A simple two-layer graph neural network"""def__init__(self,obs_size,hidden_size=32):super().__init__()self.gc0=torch_geometric.nn.DenseGraphConv(obs_size,hidden_size)self.gc1=torch_geometric.nn.DenseGraphConv(hidden_size,hidden_size)self.act=torch.nn.Tanh()defforward(self,x,adj,weights,B,N):x=self.act(self.gc0(x,adj))returnself.act(self.gc1(x,adj))ray.init(local_mode=True,object_store_memory=3e10,)input_size=16hidden_size=32cfg={"framework":"torch","num_gpus":0,"env":"CartPole-v0","num_workers":0,"model":{"custom_model":RayDenseGCM,"custom_model_config":{"graph_size":20,# GCM Ray wrapper will automatically convert observation# to gnn_input_size using a linear layer"gnn_input_size":input_size,"gnn_output_size":hidden_size,"gnn":GNN(input_size),"edge_selectors":TemporalBackedge([1]),"edge_weights":False,}}}tune.run("A2C",config=cfg,stop={"info/num_steps_trained":100})