<filename>tests/ParticleGNN/particle_graph.py
import torch
import numpy as np
import torch_geometric.nn as pyg_nn
import torch_geometric.utils as pyg_utils
import rlkit.torch.pytorch_util as ptu

class ParticleGraphBuilder(torch.nn.Module):
	def __init__(self, 
				num_agents,
				num_landmarks,
				batch_size,
				append_action=False,
				contain_self_loop=True,
				):
		super(ParticleGraphBuilder, self).__init__()

		self.num_agents = num_agents
		self.num_landmarks = num_landmarks
		self.batch_size = batch_size
		self.append_action = append_action
		self.contain_self_loop = contain_self_loop

		if append_action:
			self.output_node_dim = 5+3+2
			# agents: (vx, vy, x, y, id, 0, 0, 0, ax, ay); 
			# landmarks: (0, 0, 0, 0, 0, x, y, id, 0, 0)
		else:
			self.output_node_dim = 5+3
			# agents: (vx, vy, x, y, id, 0, 0, 0);
			# landmarks: (0, 0, 0, 0, 0, x, y, id)
	
		self.edges = self.get_full_edges(num_agents+num_landmarks)
		self.batch_edges = self.get_batch_edges(self.edges, self.batch_size, num_agents+num_landmarks)

	def forward(self, obs):
		# obs: batch x (num_agent*|O|)
		# x: batch x num_agent x output_node_dim
		# edge_index: 2 x node_edge

		batch_size = obs.shape[0]
		batch_agents = ptu.zeros(batch_size, self.num_agents, self.output_node_dim)
		batch_landmarks = ptu.zeros(batch_size, self.num_landmarks, self.output_node_dim)

		batch_agents[:,:,:5] = obs[:,:(5*self.num_agents)].reshape(batch_size,self.num_agents,5)
		if self.append_action:
			batch_agents[:,:,-2:] = obs[:,-(2*self.num_agents):].reshape(batch_size,self.num_agents,2)
		batch_landmarks[:,:,5:8] = obs[:,(5*self.num_agents):(5*self.num_agents+3*self.num_landmarks)].reshape(batch_size,self.num_landmarks,3)

		x = torch.cat([batch_agents,batch_landmarks],dim=1)
		x = x.reshape(-1,self.output_node_dim)

		edge_index = self.batch_edges[:,:self.edges.shape[1]*batch_size]

		return x, edge_index

	def get_full_edges(self, num_node):
		edges = ptu.zeros(2, num_node*num_node, dtype=int)
		edges[0,:] = torch.arange(num_node).repeat(num_node,1).transpose(0,1).reshape(num_node*num_node)
		edges[1,:] = torch.arange(num_node).repeat(1,num_node).reshape(num_node*num_node)
		if not self.contain_self_loop:
			edges = pyg_utils.remove_self_loops(edges)[0]
		return edges

	def get_batch_edges(self, edges, batch_size, num_node):
		edge_index = edges[None,:,:] # 1 x 2 x |E|
		index_offsets = torch.arange(batch_size).to(ptu.device)
		index_offsets = index_offsets * num_node
		edge_index = edge_index + index_offsets[:,None,None] # batch x 2 x |E|
		edge_index = edge_index.transpose(0,1).contiguous() # 2 x batch x |E|
		edge_index = edge_index.reshape(2,-1) # 2 x (batch * |E|)
		return edge_index
