Lab 5

Lab 5: Distributed Collaborative Systems (11/09 – 11/20)

Graph neural networks (GNNs) explore the irregular structure of graph signals, and exhibit superior performance in various applications of recommendation systems, wireless networks and control. A key property GNNs inherit from graph filters is the distributed implementation. This property makes them suitable candidates for distributed learning over large-scale networks, where global information is not available at individual agents. Each agent must decide its own actions from local observations and communication with immediate neighbors. In this lab assignment, we focus on the distributed learning with graph neural networks.

Download assignment.

2.1 System simulation

To simulate the system described by (12)-(14), we create the file $\p{consensus}$ that contains the Python class $\p{ConsensusSimulator}$ as:

from abc import ABC
import torch
import numpy as np

class ConsensusSimulator:
    def __init__(
            self,
            N,
	    nSamples,
            initial_ref_vel=1.0,
            ref_vel_delta=1.0,
            ref_vel_bias=4.0,
	    ref_est_delta=1.0,
            max_agent_acceleration=3.0,
            sample_time=0.1,
	    duration=10,
            energy_weight=1,
            is_bias_constant=True
    ):
        self.N = N
        self.nSamples = nSamples
        self.initial_ref_vel = initial_ref_vel
        self.ref_vel_delta = ref_vel_delta
        self.ref_vel_bias = ref_vel_bias
        self.ref_est_delta = ref_est_delta
        self.max_agent_acceleration = max_agent_acceleration
        self.sample_time = sample_time
        self.tSamples = int(duration / self.sample_time)
        self.energy_weight = energy_weight
        self.is_bias_constant = is_bias_constant

    def random_vector(self, expected_norm, N, samples):
        return np.random.normal(0, (2 / 3.14) ** 0.5 * expected_norm, (samples, 2, N))

    def bias(self):
        return self.random_vector(self.ref_vel_bias, self.N, self.nSamples)

    def initial(self):
        ref_vel = self.random_vector(self.initial_ref_vel, 1, self.nSamples)
        est_vel = ref_vel + self.random_vector(self.ref_est_delta, self.N, self.nSamples)
        return ref_vel, est_vel

    def simulate(self, steps, controller: ConsensusController):
	
        ref_vel, est_vel = self.initial()
        bias = self.bias()

        ref_vels = np.zeros((self.nSamples, self.tSamples, 2, 1))
        est_vels = np.zeros((self.nSamples, self.tSamples, 2, self.N))
        biased_vels = np.zeros((self.nSamples, self.tSamples, 2, self.N))
        accels = np.zeros((self.nSamples, self.tSamples, 2, self.N))
		
        ref_vels[:,0,:,:] = ref_vel
        est_vels[:,0,:,:] = est_vel
        biased_vels[:,0,:,:] = bias + ref_vel
        accels[:,0,:,:] = controller(est_vels, biased_vels, 0)
		
        this_accel = accels[:,0,:,:].copy()
        this_accel[accels[:,0,:,:] > self.max_agent_acceleration] = self.max_agent_acceleration
        this_accel[accels[:,0,:,:] < -self.max_agent_acceleration] = -self.max_agent_acceleration
        accels[:,0,:,:] = this_accel

        for step in range(1, steps):
			
            ref_accel = self.random_vector(self.ref_vel_delta, 1, self.nSamples)
            ref_vels[:,step,:,:] = ref_vels[:,step-1,:,:] + ref_accel * self.sample_time
            est_vels[:,step,:,:] = accels[:,step-1,:,:] * self.sample_time + est_vels[:,step-1,:,:]
            biased_vels[:,step,:,:] = bias + ref_vels[:,step,:,:]
			
            accels[:,step,:,:] = controller(est_vels, biased_vels, step)
			
            this_accel = accels[:,step,:,:].copy()
            this_accel[accels[:,step,:,:] > self.max_agent_acceleration] = self.max_agent_acceleration
            this_accel[accels[:,step,:,:] < -self.max_agent_acceleration] = -self.max_agent_acceleration
            accels[:,step,:,:] = this_accel
		
        return ref_vel, est_vel, est_vels, biased_vels, accels

    def cost(self, est_vels, biased_vels, accels):
	
        biased_vel_mean = np.mean(biased_vels, axis = 3)
        biased_vel_mean = np.expand_dims(biased_vel_mean, 3)
        vel_error = est_vels - biased_vel_mean
        vel_error_norm = np.square(np.linalg.norm(vel_error, ord=2, axis=2, keepdims=False))
        vel_error_norm_mean = np.sum(np.mean(vel_error_norm, axis = 2), axis = 1) 
		
        accel_norm = np.square(np.linalg.norm(self.sample_time * accels, ord=2, axis=2, keepdims=False))
        accel_norm_mean = np.sum(np.mean(accel_norm, axis = 2), axis = 1) 
		
        cost = np.mean(vel_error_norm_mean/2 + accel_norm_mean/2)
	
        return cost

Here nSamples is the number of trajectroies generated, i.e., the number of simulation times, and the controller is to be determined.

2.2 Optimal controller

To implement the optimal controller, we create the python class $\p{Optimal Controller}$ in the file $\p{consensus}$ as:

class ConsensusController(ABC):
    def __call__(self, est_vel, biased_vel):
        pass

class OptimalController(ConsensusController):
    def __init__(self, sample_time=0.1, energy_weight=1):
        self.sample_time = sample_time
        self.energy_weight = energy_weight

    def __call__(self, est_vel, biased_vel, step):
	
        mean_biased_vel  = np.mean(biased_vel[:,step,:,:], axis=2)
        mean_biased_vel = np.expand_dims(mean_biased_vel, 2)
	
        return -(est_vel[:,step,:,:] - mean_biased_vel) / (self.sample_time * (1 + self.energy_weight))

2.3 Communication network

To implement the decentralized controller, we first create the file $\p{networks}$ to create the function that computes the communication graph as:

import numpy as np
from scipy.spatial import distance_matrix
import copy
import heapq
rng = np.random

def agent_communication(wx, wy, n_agents, degree):
    positions = np.vstack((
        rng.uniform(0, wx, n_agents),
        rng.uniform(0, wy, n_agents)))
    distance = distance_matrix(positions.T, positions.T)
    network = copy.deepcopy(distance)    
    for i in range(n_agents):        
        re2 = heapq.nsmallest((degree+1), distance[i,:])        
        for j in range(n_agents):       
            if distance[i,j] > re2[degree]:                
                network[i,j] = 0
    network = network + np.transpose(network, axes = [1,0])    
    network[np.arange(0,n_agents),np.arange(0,n_agents)] = 0.
    network = (network > 0).astype(distance.dtype)
    W = np.linalg.eigvalsh(network)
    maxEigenvalue = np.max(np.real(W))
    normalized_network = network / maxEigenvalue
    
    return network, normalized_network

2.4 Decentralized controller

With the communication network available, we then create the python class $\p{DistributedController}$ for decentralized controller in the file $\p{consensus}$ as:

class DistributedController(ConsensusController):
    def __init__(self, adjacency, K=1, sample_time=0.1, energy_weight=1):
	
        self.nSamples = adjacency.shape[0]
        self.nAgents = adjacency.shape[1]
        self.K = K
        self.neighbor_network = np.zeros((self.nSamples, K, self.nAgents, self.nAgents))
        self.num_neighbor = np.zeros((self.nSamples, K, 2, self.nAgents))
	
        for i in range(self.nSamples):
            for k in range(1, (K+1)):
                inter_matrix = np.linalg.matrix_power(adjacency[i,:,:], k)
                inter_matrix[inter_matrix>=1] = 1
                self.num_neighbor[i,k-1,0,:] = np.sum(inter_matrix,1)
                self.num_neighbor[i,k-1,1,:] = np.sum(inter_matrix,1)
                self.neighbor_network[i,k-1,:,:] = inter_matrix
				
        self.sample_time = sample_time
        self.energy_weight = energy_weight

    def __call__(self, est_vel, biased_vel, step):
        # find the averages along 0,K neighborhoods and average them
        biased_means = np.zeros((self.nSamples, self.K, 2, self.nAgents))
        
        if step < (self.K-1):
            for k in range(step):
                biased_means[:,k,:,:] = np.matmul(biased_vel[:,step-k,:,:], self.neighbor_network[:,k,:,:])
                biased_means[:,k,:,:] = biased_means[:,k,:,:] / self.num_neighbor[:,k,:,:]            
        else:
            for k in range(self.K):
                biased_means[:,k,:,:] = np.matmul(biased_vel[:,step-k,:,:], self.neighbor_network[:,k,:,:])
                biased_means[:,k,:,:] = biased_means[:,k,:,:] / self.num_neighbor[:,k,:,:]  
        biased_mean = np.mean(biased_means, 1)
        return (biased_mean - est_vel[:,step,:,:]) / (self.sample_time * (1 + self.energy_weight))

Together, we use above classes to perform system simulations in the file $\p{main}$ as:

from importlib import reload
import numpy as np
import consensus
import networks
consensus = reload(consensus)
networks = reload(networks)

# %% Parameters
D = 2
K = 1
nSamples = 200
N = 100

com_network = np.zeros((nSamples, N, N))

# %% Networks
for i in range(nSamples):
    com_network[i,:,:], _ = networks.agent_communication(100, 200, N, D)

# %% Dataset generation
optimal_controller = consensus.OptimalController()
decentralized_controller = consensus.DistributedController(com_network, K)
simulator = consensus.ConsensusSimulator(N, nSamples)

opt_ref_vel, opt_est_vel, opt_est_vels, opt_biased_vels, opt_accels = simulator.simulate(100, optimal_controller)
dec_ref_vel, dec_est_vel, dec_est_vels, dec_biased_vels, dec_accels = simulator.simulate(100, decentralized_controller)
opt_cost = simulator.cost(opt_est_vels, opt_biased_vels, opt_accels)
dec_cost = simulator.cost(dec_est_vels, dec_biased_vels, dec_accels)

3.1 Learning with a GNN

To learn a decentralized controller with a GNN, we first create the GNN architecture as the file $\p{modules}$:

import math
import torch
import torch.nn as nn
import numpy as np

def LSIGF_DB(h, S, x, b=None):

    F = h.shape[0]
    E = h.shape[1]
    K = h.shape[2]
    G = h.shape[3]
    B = S.shape[0]
    T = S.shape[1]
    N = S.shape[3]
    
    x = x.reshape([B, T, 1, G, N]).repeat(1, 1, E, 1, 1)    
    z = x.reshape([B, T, 1, E, G, N])
    for k in range(1,K):
        x, _ = torch.split(x, [T-1, 1], dim = 1)
        zeroRow = torch.zeros(B, 1, E, G, N, dtype=x.dtype,device=x.device)
        x = torch.cat((zeroRow, x), dim = 1)
        x = torch.matmul(x, S)
        xS = x.reshape(B, T, 1, E, G, N)
        z = torch.cat((z, xS), dim = 2)

    z = z.permute(0, 1, 5, 3, 2, 4) 
    z = z.reshape(B, T, N, E*K*G)
    h = h.reshape(F, E*K*G)
    h = h.permute(1, 0)     
    y = torch.matmul(z, h) 
    y = y.permute(0, 1, 3, 2) 
    if b is not None:
        y = y + b
    
    return y
		
class GraphFilter_DB(nn.Module):

    def __init__(self, G, F, K, E = 1, bias = True):

        super().__init__()
        
        self.G = G
        self.F = F
        self.K = K
        self.E = E
        self.S = None 
        self.weight = nn.parameter.Parameter(torch.Tensor(F, E, K, G))
        if bias:
            self.bias = nn.parameter.Parameter(torch.Tensor(F, 1))
        else:
            self.register_parameter('bias', None)
        # Initialize parameters
        self.reset_parameters()

    def reset_parameters(self):

        stdv = 1. / math.sqrt(self.G * self.K)
        self.weight.data.uniform_(-stdv, stdv)
        if self.bias is not None:
            self.bias.data.uniform_(-stdv, stdv)

    def addGSO(self, S):

        assert len(S.shape) == 5
        assert S.shape[2] == self.E
        self.N = S.shape[3]
        assert S.shape[4] == self.N
        self.S = S

    def forward(self, x):

        assert len(x.shape) == 4
        B = x.shape[0]
        assert self.S.shape[0] == B
        T = x.shape[1]
        assert self.S.shape[1] == T
        assert x.shape[3] == self.N
        u = LSIGF_DB(self.weight, self.S, x, self.bias)
        return u
				
class LocalGNN_DB(nn.Module):

    def __init__(self,
                 dimNodeSignals, nFilterTaps, bias,
                 nonlinearity,
                 dimReadout,
                 dimEdgeFeatures):

        super().__init__()
        
        self.L = len(nFilterTaps) 
        self.F = dimNodeSignals 
        self.K = nFilterTaps 
        self.E = dimEdgeFeatures 
        self.bias = bias 
        self.sigma = nonlinearity
        self.dimReadout = dimReadout

        gfl0 = []
        for l in range(self.L):
            gfl0.append(GraphFilter_DB(self.F[l], self.F[l+1], self.K[l],
                                              self.E, self.bias))
            gfl0.append(self.sigma())
        self.GFL0 = nn.Sequential(*gfl0)
        
        fc0 = []
        if len(self.dimReadout) > 0: 
            fc0.append(nn.Linear(self.F[-1], dimReadout[0], bias = self.bias))
            for l in range(len(dimReadout)-1):
                fc0.append(self.sigma())
                fc0.append(nn.Linear(dimReadout[l], dimReadout[l+1],
                                    bias = self.bias))

        self.Readout0 = nn.Sequential(*fc0)

    def splitForward(self, x, S):

        assert len(S.shape) == 4 or len(S.shape) == 5
        if len(S.shape) == 4:
            S = S.unsqueeze(2)
        for l in range(self.L):
            self.GFL0[2*l].addGSO(S)
        yGFL0 = self.GFL0(x)
        y = yGFL0.permute(0, 1, 3, 2)
        y = self.Readout0(y)
        
        return y.permute(0, 1, 3, 2), yGFL0
        # B x T x dimReadout[-1] x N, B x T x dimFeatures[-1] x N
    
    def forward(self, x, S):

        output, _ = self.splitForward(x, S)
        
        return output

To perform the imitation learning, we need the validation to evaluate the cost in the training process. We create a function in the class $\p{ConsensusSimulator}$ (in the file $\p{consensus}$) named $\p{ComputeTrajectory}$ that simulates the trajectory by using the GNN controller for cost evaluation as:

import networks
import copy

    def computeTrajectory(self, archit, ref_vel, est_vel, step):
        
        batchSize = est_vel.shape[0]
        nAgents = est_vel.shape[2]
        ref_vel = ref_vel.squeeze(2)
        architDevice = list(archit.parameters())[0].device
        est_vels = np.zeros((batchSize, step, 2, nAgents), dtype = np.float)
        ref_vels = np.zeros((batchSize, step, 2), dtype = np.float)
        biased_ref_vels = np.zeros((batchSize, step, 2, nAgents), dtype = np.float)
        bias = np.zeros((batchSize, step, 2, nAgents))
        accels = np.zeros((batchSize, step, 2, nAgents), dtype=np.float)
        states = np.zeros((batchSize, step, 4, nAgents), dtype=np.float)
        graphs = np.zeros((batchSize, step, nAgents, nAgents), dtype = np.float)
            
        est_vels[:,0,:,:] = est_vel.copy()
        ref_vels[:,0,:] = ref_vel.copy()
        bias[:,0,:,:] = np.random.normal(loc=0,scale=(4 * np.sqrt(2/np.pi)),size=(batchSize, 2, nAgents))
        ref_vel_temp = np.expand_dims(ref_vels[:,0,:], 2)
        biased_ref_vels[:,0,:,:] = ref_vel_temp + bias[:,0,:,:]

        for i in range(20):
            _, graphs[i,0,:,:] = networks.agent_communication(100, 200, self.N, 2)
        
        for t in range(1, step):
            
            thisbiased_ref_vels = np.expand_dims(biased_ref_vels[:,t-1,:,:], 1)
            thisest_vels = np.expand_dims(est_vels[:,t-1,:,:], 1)
            thisState = np.concatenate((thisest_vels, thisbiased_ref_vels), axis = 2)
            states[:,t-1,:,:] = thisState.squeeze(1)

            x = torch.tensor(states[:,0:t,:,:], device = architDevice)
            S = torch.tensor(graphs[:,0:t,:,:], device = architDevice)
            with torch.no_grad():
                thisaccels = archit(x, S)
            thisaccels = thisaccels.cpu().numpy()[:,-1,:,:]
            thisaccels[thisaccels > 3] = 3
            thisaccels[thisaccels < -3] = -3
            accels[:,t-1,:,:] = thisaccels
            est_vels[:,t,:,:] = accels[:,t-1,:,:] * 0.1 +est_vels[:,t-1,:,:]
            ref_accel = np.random.normal(loc=0,scale=(self.ref_vel_delta * np.sqrt(2/np.pi)),size=(batchSize, 2))
            ref_vels[:,t,:] = ref_vels[:,t-1,:] + ref_accel * 0.1
            ref_vel_temp = np.expand_dims(ref_vels[:,t,:], 2)
            biased_ref_vels[:,t,:,:] = ref_vel_temp + bias[:,0,:,:]
            graphs[:,t,:,:] = copy.deepcopy(graphs[:,0,:,:])
        
        thisbiased_ref_vels = np.expand_dims(biased_ref_vels[:,-1,:,:], 1)
        thisest_vels = np.expand_dims(est_vels[:,-1,:,:], 1)
        thisState = np.concatenate((thisest_vels, thisbiased_ref_vels), axis = 2)
        states[:,-1,:,:] = thisState.squeeze(1)
        x = torch.tensor(states, device = architDevice)
        S = torch.tensor(graphs, device = architDevice)
        with torch.no_grad():
            thisaccels = archit(x, S)
        thisaccels = thisaccels.cpu().numpy()[:,-1,:,:]
        thisaccels[thisaccels > 3] = 3
        thisaccels[thisaccels < -3] = -3
        accels[:,-1,:,:] = thisaccels
            
        return ref_vel, est_vel, est_vels, biased_ref_vels, accels

We follow to build the imitation learning process for the GNN as the file $\p{training}$:

import torch
import numpy as np
import datetime

class Trainer:
    
    def __init__(self, model, simulator, ref_vel, est_vel, est_vels, biased_vels, accels, S, nEpochs, batchSize):
        
        self.model = model
        self.est_vel = est_vel
        self.ref_vel = ref_vel
        self.est_vels = est_vels
        self.biased_vels = biased_vels
        self.accels = accels
        self.S = S
        self.nTrain = 200
        self.nTest = 20
        self.nValid = 20
        self.nEpochs = nEpochs
        self.simulator = simulator
        self.validationInterval = self.nTrain//batchSize

        if self.nTrain < batchSize:
            self.nBatches = 1
            self.batchSize = [self.nTrain]
        elif self.nTrain % batchSize != 0:
            self.nBatches = np.ceil(self.nTrain/batchSize).astype(np.int64)
            self.batchSize = [batchSize] * self.nBatches
            while sum(batchSize) != self.nTrain:
                self.batchSize[-1] -= 1
        else:
            self.nBatches = np.int(self.nTrain/batchSize)
            self.batchSize = [batchSize] * self.nBatches
            
        self.batchIndex = np.cumsum(self.batchSize).tolist()
        self.batchIndex = [0] + self.batchIndex

class TrainerFlocking(Trainer):
    
    def __init__(self, model, simulator, ref_vel, est_vel, est_vels, biased_vels, accels, S, nEpochs, batchSize):
        
        super().__init__(model, simulator, ref_vel, est_vel, est_vels, biased_vels, accels, S, nEpochs, batchSize)

    def train(self):
        
        nTrain = self.nTrain
        thisArchit = self.model.archit
        thisLoss = self.model.loss
        thisOptim = self.model.optim
        thisDevice = self.model.device
        epoch = 0 
        xTrainAll = np.concatenate((self.est_vels[0:self.nTrain].copy(), self.biased_vels[0:self.nTrain].copy()), axis = 2)
        yTrainAll = self.accels[0:self.nTrain].copy()
        StrainAll = self.S[0:self.nTrain].copy()

        while epoch < self.nEpochs:

            randomPermutation = np.random.permutation(nTrain)
            idxEpoch = [int(i) for i in randomPermutation]
            batch = 0 

            while batch < self.nBatches:

                thisBatchIndices = idxEpoch[self.batchIndex[batch] : self.batchIndex[batch+1]]
                xTrain = xTrainAll[thisBatchIndices]
                yTrain = yTrainAll[thisBatchIndices]
                Strain = StrainAll[thisBatchIndices]
                    
                xTrain = torch.tensor(xTrain, device = thisDevice)
                Strain = torch.tensor(Strain, device = thisDevice)
                yTrain = torch.tensor(yTrain, device = thisDevice)

                startTime = datetime.datetime.now()
                thisArchit.zero_grad()
                yHatTrain = thisArchit(xTrain, Strain)
                lossValueTrain = thisLoss(yHatTrain, yTrain)
                lossValueTrain.backward()
                thisOptim.step()
                endTime = datetime.datetime.now()
                timeElapsed = abs(endTime - startTime).total_seconds()
                del xTrain
                del Strain
                del yTrain
                del lossValueTrain

                #\\\\\\\
                #\\\ VALIDATION
                #\\\\\\\

                if (epoch * self.nBatches + batch) % self.validationInterval == 0:
                    
                    startTime = datetime.datetime.now()
                    ref_vel = self.ref_vel[220:240]
                    est_vel = self.est_vel[220:240]
                    _, _, est_vels_valid, biased_vels_valid, accels_valid = self.simulator.computeTrajectory(thisArchit, ref_vel, est_vel, 100)
                    accValid = self.simulator.cost(est_vels_valid, biased_vels_valid, accels_valid)
                    endTime = datetime.datetime.now()
                    timeElapsed = abs(endTime - startTime).total_seconds()

                    print("\t(E: %2d, B: %3d) %8.4f - %6.4fs" % (
                            epoch+1, batch+1,
                            accValid, 
                            timeElapsed), end = ' ')
                    print("[VALIDATION] \n", end = '')

                    if epoch == 0 and batch == 0:
                        bestScore = accValid
                        bestEpoch, bestBatch = epoch, batch
                        self.model.save(label = 'Best')
                        # Start the counter
                    else:
                        thisValidScore = accValid
                        if thisValidScore < bestScore:
                            bestScore = thisValidScore
                            bestEpoch, bestBatch = epoch, batch
                            print("\t=> New best achieved: %.4f" % (bestScore))
                            self.model.save(label = 'Best')
                    del ref_vel
                    del est_vel
                batch += 1
            epoch += 1

        self.model.save(label = 'Last')
        #################
        # TRAINING OVER #
        #################
        if self.nEpochs == 0:
            self.model.save(label = 'Best')
            self.model.save(label = 'Last')
            print("\nWARNING: No training. Best and Last models are the same.\n")
        self.model.load(label = 'Best')
        if self.nEpochs > 0:
            print("\t=> Best validation achieved (E: %d, B: %d): %.4f" % (
                    bestEpoch + 1, bestBatch + 1, bestScore))

and the file $\p{model}$ that contains the Python class $\p{Model}$ aggregating all training components with the GNN architecture as:

import os
import torch

class Model:
    
    def __init__(self,
                 architecture,
                 loss,
                 optimizer,
                 trainer,
                 device, name, saveDir):
        
        self.archit = architecture
        self.nParameters = 0
        for param in list(self.archit.parameters()):
            if len(param.shape)>0:
                thisNParam = 1
                for p in range(len(param.shape)):
                    thisNParam *= param.shape[p]
                self.nParameters += thisNParam
            else:
                pass
        self.loss = loss
        self.optim = optimizer
        self.trainer = trainer
        self.device = device
        self.name = name
        self.saveDir = saveDir
        
    def train(self, simulator, ref_vel, est_vel, est_vels, biased_vels, accels, com_network, nEpochs, batchSize):
        
        self.trainer = self.trainer(self, simulator, ref_vel, est_vel, est_vels, biased_vels, accels, com_network, nEpochs, batchSize)
        
        return self.trainer.train()
    
    def save(self, label = '', **kwargs):
        if 'saveDir' in kwargs.keys():
            saveDir = kwargs['saveDir']
        else:
            saveDir = self.saveDir
        saveModelDir = os.path.join(saveDir,'savedModels')
        # Create directory savedModels if it doesn't exist yet:
        if not os.path.exists(saveModelDir):
            os.makedirs(saveModelDir)
        saveFile = os.path.join(saveModelDir, self.name)
        torch.save(self.archit.state_dict(), saveFile+'Archit'+ label+'.ckpt')
        torch.save(self.optim.state_dict(), saveFile+'Optim'+label+'.ckpt')

    def load(self, label = '', **kwargs):
        if 'loadFiles' in kwargs.keys():
            (architLoadFile, optimLoadFile) = kwargs['loadFiles']
        else:
            saveModelDir = os.path.join(self.saveDir,'savedModels')
            architLoadFile = os.path.join(saveModelDir,
                                          self.name + 'Archit' + label +'.ckpt')
            optimLoadFile = os.path.join(saveModelDir,
                                         self.name + 'Optim' + label + '.ckpt')
        self.archit.load_state_dict(torch.load(architLoadFile))
        self.optim.load_state_dict(torch.load(optimLoadFile))

We can then use aforementioned classes and files to train a decentralized GNN controller in the file $\p{main}$ as:

from importlib import reload
import numpy as np
import consensus
import networks
import os
import torch; torch.set_default_dtype(torch.float64)
import torch.nn as nn
import torch.optim as optim
import modules
import training
import model
import datetime
from copy import deepcopy

consensus = reload(consensus)
networks = reload(networks)

# %% Parameters
D = 2
K = 3
nSamples = 240
N = 100
step = 100

network = np.zeros((nSamples, N, N))
com_network = np.zeros((nSamples, step, N, N))

# %% Networks
for i in range(nSamples):
    network[i,:,:], com_network[i,0,:,:] = networks.agent_communication(100, 200, N, D)
    for t in range(1, step):
        com_network[i,t,:,:] = deepcopy(com_network[i,0,:,:])

# %% Dataset generation
optimal_controller = consensus.OptimalController()
decentralized_controller = consensus.DistributedController(network, K)
simulator = consensus.ConsensusSimulator(N, nSamples)

ref_vel, est_vel, est_vels, biased_vels, accels = simulator.simulate(100, optimal_controller)
    
# %% Training
thisFilename = 'flockingGNN' 
saveDirRoot = 'experiments' 
saveDir = os.path.join(saveDirRoot, thisFilename) 
today = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
saveDir = saveDir + '-%03d-' + today
# Create directory
if not os.path.exists(saveDir):
    os.makedirs(saveDir)

#   PyTorch seeds
torchState = torch.get_rng_state()
torchSeed = torch.initial_seed()
#   Numpy seeds
numpyState = np.random.RandomState().get_state()
    
nEpochs = 50 # Number of epochs
batchSize = 20 # Batch size
learningRate = 0.00025
beta1 = 0.9
beta2 = 0.999

hParamsLocalGNN = {} # Hyperparameters (hParams) for the Local GNN
hParamsLocalGNN['name'] = 'LocalGNN'
hParamsLocalGNN['archit'] = modules.LocalGNN_DB
hParamsLocalGNN['device'] = 'cuda:0' \
                                if (True and torch.cuda.is_available()) \
                                else 'cpu'
hParamsLocalGNN['dimNodeSignals'] = [4, 64] # Features per layer
hParamsLocalGNN['nFilterTaps'] = [4] # Number of filter taps
hParamsLocalGNN['bias'] = True 
hParamsLocalGNN['nonlinearity'] = nn.Tanh 
hParamsLocalGNN['dimReadout'] = [2] 
hParamsLocalGNN['dimEdgeFeatures'] = 1 

if torch.cuda.is_available():
    torch.cuda.empty_cache()
hParamsDict = deepcopy(hParamsLocalGNN)
thisName = hParamsDict.pop('name')
callArchit = hParamsDict.pop('archit')
thisDevice = hParamsDict.pop('device')

##############
# PARAMETERS #
##############

thisArchit = callArchit(**hParamsDict)
thisArchit.to(thisDevice)
thisOptim = optim.Adam(thisArchit.parameters(), lr = learningRate, betas = (beta1, beta2))
thisLossFunction = nn.MSELoss()
thisTrainer = training.TrainerFlocking
GraphNNs = model.Model(thisArchit, thisLossFunction, thisOptim, thisTrainer, thisDevice, thisName, saveDir)

#Train
GraphNNs.train(simulator, ref_vel, est_vel, est_vels, biased_vels, accels, com_network, nEpochs, batchSize)

#Test
ref_vel_test = ref_vel[200:220]
est_vel_test = est_vel[200:220]
_, _, est_vels_valid, biased_vels_valid, accels_valid = simulator.computeTrajectory(GraphNNs.archit, ref_vel_test, est_vel_test, 100)
cost = simulator.cost(est_vels_valid, biased_vels_valid, accels_valid)
print(cost)

4.1-4.3 Agent Mobility

We proceed to account for the agent mobility. In this context, we first create the function in the file $\p{networks}$ that computes the communication graph based on agents positions as:

def agent_communication_pos(pos, n_agents, degree):
    
    distance = distance_matrix(pos.T, pos.T)
    network = copy.deepcopy(distance)
    for i in range(n_agents):        
        re2 = heapq.nsmallest((degree+1), distance[i,:])        
        for j in range(n_agents):       
            if distance[i,j] > re2[degree]:                
                network[i,j] = 0
    network = network + np.transpose(network, axes = [1,0]) 
    network[np.arange(0,n_agents),np.arange(0,n_agents)] = 0.
    network = (network > 0).astype(distance.dtype)
    W = np.linalg.eigvalsh(network)
    maxEigenvalue = np.max(np.real(W))
    normalized_network = network / maxEigenvalue

    return network, normalized_network

To consider agent mobility in the system simulation, we create the Python class $\p{simulatepos}$ in the class $\p{ConsensusSimulator}$ (in the file $\p{consensus}$) to incorporate communication graph changes induced by time-varying agent positions as:

    def simulate_pos(self, wx, wy, degree, steps, controller: ConsensusController):
	
        ref_vel, est_vel = self.initial()
        bias = self.bias()
        x_pos = np.random.uniform(0, wx, (self.nSamples, 1, self.N))
        y_pos = np.random.uniform(0, wy, (self.nSamples, 1, self.N))
        pos = np.concatenate((x_pos, y_pos), axis = 1)
        ref_vels = np.zeros((self.nSamples, self.tSamples, 2, 1))
        est_vels = np.zeros((self.nSamples, self.tSamples, 2, self.N))
        biased_vels = np.zeros((self.nSamples, self.tSamples, 2, self.N))
        accels = np.zeros((self.nSamples, self.tSamples, 2, self.N))
        positions = np.zeros((self.nSamples, self.tSamples, 2, self.N))
        networkss = np.zeros((self.nSamples, self.tSamples, self.N, self.N))		
        ref_vels[:,0,:,:] = ref_vel
        est_vels[:,0,:,:] = est_vel
        biased_vels[:,0,:,:] = bias + ref_vel
        accels[:,0,:,:] = controller(est_vels, biased_vels, 0)
        positions[:,0,:,:] = pos
        for i in range(self.nSamples):
            _, networkss[i,0,:,:] = networks.agent_communication_pos(positions[i,0,:,:], self.N, degree)	
        this_accel = accels[:,0,:,:].copy()
        this_accel[accels[:,0,:,:] > self.max_agent_acceleration] = self.max_agent_acceleration
        this_accel[accels[:,0,:,:] < -self.max_agent_acceleration] = -self.max_agent_acceleration
        accels[:,0,:,:] = this_accel

        for step in range(1, steps):
            ref_accel = self.random_vector(self.ref_vel_delta, 1, self.nSamples)
            ref_vels[:,step,:,:] = ref_vels[:,step-1,:,:] + ref_accel * self.sample_time
            est_vels[:,step,:,:] = accels[:,step-1,:,:] * self.sample_time + est_vels[:,step-1,:,:]
            positions[:,step,:,:] = accels[:,step-1,:,:] * (self.sample_time ** 2)/2 + est_vels[:,step-1,:,:] * self.sample_time + positions[:,step-1,:,:]
            for i in range(self.nSamples):
                _, networkss[i,step,:,:] = networks.agent_communication_pos(positions[i,step,:,:], self.N, degree)	
            biased_vels[:,step,:,:] = bias + ref_vels[:,step,:,:]			
            accels[:,step,:,:] = controller(est_vels, biased_vels, step)			
            this_accel = accels[:,step,:,:].copy()
            this_accel[accels[:,step,:,:] > self.max_agent_acceleration] = self.max_agent_acceleration
            this_accel[accels[:,step,:,:] < -self.max_agent_acceleration] = -self.max_agent_acceleration
            accels[:,step,:,:] = this_accel
		
        return networkss, positions, ref_vel, est_vel, est_vels, biased_vels, accels    

To evaluate the cost that accounts for agent mobility, we shall similarly create the Python function $\p{computeTrajectorypos}$ in the class $\p{ConsensusSimulator}$ (in the file $\p{consensus}$) as:

    def computeTrajectory_pos(self, archit, position, ref_vel, est_vel, step):
        
        batchSize = est_vel.shape[0]
        nAgents = est_vel.shape[2]
        ref_vel = ref_vel.squeeze(2)
        architDevice = list(archit.parameters())[0].device
        est_vels = np.zeros((batchSize, step, 2, nAgents), dtype = np.float)
        ref_vels = np.zeros((batchSize, step, 2), dtype = np.float)
        positions = np.zeros((batchSize, step, 2, nAgents), dtype = np.float)
        biased_ref_vels = np.zeros((batchSize, step, 2, nAgents), dtype = np.float)
        bias = np.zeros((batchSize, step, 2, nAgents))
        accels = np.zeros((batchSize, step, 2, nAgents), dtype=np.float)
        states = np.zeros((batchSize, step, 4, nAgents), dtype=np.float)
        graphs = np.zeros((batchSize, step, nAgents, nAgents), dtype = np.float)            
        est_vels[:,0,:,:] = est_vel.copy()
        ref_vels[:,0,:] = ref_vel.copy()
        bias[:,0,:,:] = np.random.normal(loc=0,scale=(4 * np.sqrt(2/np.pi)),size=(batchSize, 2, nAgents))
        ref_vel_temp = np.expand_dims(ref_vels[:,0,:], 2)
        biased_ref_vels[:,0,:,:] = ref_vel_temp + bias[:,0,:,:]
        positions[:,0,:,:] = position[:,0,:,:]

        for i in range(20):
            _, graphs[i,0,:,:] = networks.agent_communication_pos(positions[i,0,:,:], nAgents, 2)
        
        for t in range(1, step):
            
            thisbiased_ref_vels = np.expand_dims(biased_ref_vels[:,t-1,:,:], 1)
            thisest_vels = np.expand_dims(est_vels[:,t-1,:,:], 1)
            thisState = np.concatenate((thisest_vels, thisbiased_ref_vels), axis = 2)
            states[:,t-1,:,:] = thisState.squeeze(1)

            x = torch.tensor(states[:,0:t,:,:], device = architDevice)
            S = torch.tensor(graphs[:,0:t,:,:], device = architDevice)
            with torch.no_grad():
                thisaccels = archit(x, S)
            thisaccels = thisaccels.cpu().numpy()[:,-1,:,:]
            thisaccels[thisaccels > 3] = 3
            thisaccels[thisaccels < -3] = -3
            accels[:,t-1,:,:] = thisaccels
            est_vels[:,t,:,:] = accels[:,t-1,:,:] * 0.1 +est_vels[:,t-1,:,:]
            positions[:,t,:,:] = accels[:,t-1,:,:] * (0.1 ** 2)/2 + est_vels[:,t-1,:,:] * 0.1 + positions[:,t-1,:,:]
            ref_accel = np.random.normal(loc=0,scale=(self.ref_vel_delta * np.sqrt(2/np.pi)),size=(batchSize, 2))
            ref_vels[:,t,:] = ref_vels[:,t-1,:] + ref_accel * 0.1
            ref_vel_temp = np.expand_dims(ref_vels[:,t,:], 2)
            biased_ref_vels[:,t,:,:] = ref_vel_temp + bias[:,0,:,:]
            for i in range(20):
                _, graphs[i,t,:,:] = networks.agent_communication_pos(positions[i,t,:,:], nAgents, 2)   
          
        thisbiased_ref_vels = np.expand_dims(biased_ref_vels[:,-1,:,:], 1)
        thisest_vels = np.expand_dims(est_vels[:,-1,:,:], 1)
        thisState = np.concatenate((thisest_vels, thisbiased_ref_vels), axis = 2)
        states[:,-1,:,:] = thisState.squeeze(1)
        x = torch.tensor(states, device = architDevice)
        S = torch.tensor(graphs, device = architDevice)
        with torch.no_grad():
            thisaccels = archit(x, S)
        thisaccels = thisaccels.cpu().numpy()[:,-1,:,:]
        thisaccels[thisaccels > 3] = 3
        thisaccels[thisaccels < -3] = -3
        accels[:,-1,:,:] = thisaccels
            
        return ref_vel, est_vel, est_vels, biased_ref_vels, accels

The modified simulate class $\p{simulatepos}$ additionally outputs time-varying agent positions and communication graphs. The function $\p{train}$ in the file $\p{model}$ is modified to input both values as:

    def train(self, simulator, positions, ref_vel, est_vel, est_vels, biased_vels, accels, com_network, nEpochs, batchSize):        
        self.trainer = self.trainer(self, simulator, positions, ref_vel, est_vel, est_vels, biased_vels, accels, com_network, nEpochs, batchSize)        
        return self.trainer.train()

We then need to replace the function $\p{computeTrajectory}$ with $\p{computeTrajectorypos}$ in the file $\p{training}$ to consider agent movements as:

import torch
import numpy as np
import datetime

class Trainer:
    
    def __init__(self, model, simulator, positions, ref_vel, est_vel, est_vels, biased_vels, accels, S, nEpochs, batchSize):
        
        self.model = model
        self.est_vel = est_vel
        self.ref_vel = ref_vel
        self.est_vels = est_vels
        self.biased_vels = biased_vels
        self.accels = accels
        self.S = S
        self.nTrain = 200
        self.nTest = 20
        self.nValid = 20
        self.nEpochs = nEpochs
        self.simulator = simulator
        self.positions = positions
        self.validationInterval = self.nTrain//batchSize

        if self.nTrain < batchSize:
            self.nBatches = 1
            self.batchSize = [self.nTrain]
        elif self.nTrain % batchSize != 0:
            self.nBatches = np.ceil(self.nTrain/batchSize).astype(np.int64)
            self.batchSize = [batchSize] * self.nBatches
            while sum(batchSize) != self.nTrain:
                self.batchSize[-1] -= 1
        else:
            self.nBatches = np.int(self.nTrain/batchSize)
            self.batchSize = [batchSize] * self.nBatches
            
        self.batchIndex = np.cumsum(self.batchSize).tolist()
        self.batchIndex = [0] + self.batchIndex

class TrainerFlocking(Trainer):
    
    def __init__(self, model, simulator, positions, ref_vel, est_vel, est_vels, biased_vels, accels, S, nEpochs, batchSize):
        
        super().__init__(model, simulator, positions, ref_vel, est_vel, est_vels, biased_vels, accels, S, nEpochs, batchSize)

    def train(self):
        
        nTrain = self.nTrain
        thisArchit = self.model.archit
        thisLoss = self.model.loss
        thisOptim = self.model.optim
        thisDevice = self.model.device
        epoch = 0        
        xTrainAll = np.concatenate((self.est_vels[0:self.nTrain].copy(), self.biased_vels[0:self.nTrain].copy()), axis = 2)
        yTrainAll = self.accels[0:self.nTrain].copy()
        StrainAll = self.S[0:self.nTrain].copy()

        while epoch < self.nEpochs:

            randomPermutation = np.random.permutation(nTrain)
            idxEpoch = [int(i) for i in randomPermutation]
            batch = 0 
            while batch < self.nBatches:

                thisBatchIndices = idxEpoch[self.batchIndex[batch] : self.batchIndex[batch+1]]
                xTrain = xTrainAll[thisBatchIndices]
                yTrain = yTrainAll[thisBatchIndices]
                Strain = StrainAll[thisBatchIndices]                    
                xTrain = torch.tensor(xTrain, device = thisDevice)
                Strain = torch.tensor(Strain, device = thisDevice)
                yTrain = torch.tensor(yTrain, device = thisDevice)
                startTime = datetime.datetime.now()
                thisArchit.zero_grad()
                yHatTrain = thisArchit(xTrain, Strain)
                lossValueTrain = thisLoss(yHatTrain, yTrain)
                lossValueTrain.backward()
                thisOptim.step()
                endTime = datetime.datetime.now()
                timeElapsed = abs(endTime - startTime).total_seconds()
                del xTrain
                del Strain
                del yTrain
                del lossValueTrain

                #\\\\\\\
                #\\\ VALIDATION
                #\\\\\\\

                if (epoch * self.nBatches + batch) % self.validationInterval == 0:
                    
                    startTime = datetime.datetime.now()
                    ref_vel = self.ref_vel[220:240]
                    est_vel = self.est_vel[220:240]
                    position = self.positions[220:240]                   
                    _, _, est_vels_valid, biased_vels_valid, accels_valid = self.simulator.computeTrajectory_pos(thisArchit, position, ref_vel, est_vel, 100)
                    accValid = self.simulator.cost(est_vels_valid, biased_vels_valid, accels_valid)
                    endTime = datetime.datetime.now()
                    timeElapsed = abs(endTime - startTime).total_seconds()

                    print("\t(E: %2d, B: %3d) %8.4f - %6.4fs" % (
                            epoch+1, batch+1,
                            accValid, 
                            timeElapsed), end = ' ')
                    print("[VALIDATION] \n", end = '')

                    if epoch == 0 and batch == 0:
                        bestScore = accValid
                        bestEpoch, bestBatch = epoch, batch
                        self.model.save(label = 'Best')
                        # Start the counter
                    else:
                        thisValidScore = accValid
                        if thisValidScore < bestScore:
                            bestScore = thisValidScore
                            bestEpoch, bestBatch = epoch, batch
                            print("\t=> New best achieved: %.4f" % (bestScore))
                            self.model.save(label = 'Best')
                    del ref_vel
                    del est_vel
                batch += 1
            epoch += 1

        self.model.save(label = 'Last')
        #################
        # TRAINING OVER #
        #################
        if self.nEpochs == 0:
            self.model.save(label = 'Best')
            self.model.save(label = 'Last')
            print("\nWARNING: No training. Best and Last models are the same.\n")
        self.model.load(label = 'Best')
        if self.nEpochs > 0:
            print("\t=> Best validation achieved (E: %d, B: %d): %.4f" % (
                    bestEpoch + 1, bestBatch + 1, bestScore))

We can then implement the simulation with agent mobility by aggregating all above modifications in the $\p{main}$ file as:

from importlib import reload
import numpy as np
import consensus
import networks
import os
import torch; torch.set_default_dtype(torch.float64)
import torch.nn as nn
import torch.optim as optim
import modules
import training
import model
import datetime
from copy import deepcopy
consensus = reload(consensus)
networks = reload(networks)

# %% Parameters
D = 2
K = 3
nSamples = 240
N = 100
step = 100
network = np.zeros((nSamples, N, N))
com_network = np.zeros((nSamples, step, N, N))

# %% Networks
for i in range(nSamples):
    network[i,:,:], com_network[i,0,:,:] = networks.agent_communication(100, 200, N, D)
    for t in range(1, step):
        com_network[i,t,:,:] = deepcopy(com_network[i,0,:,:])

# %% Dataset generation
optimal_controller = consensus.OptimalController()
decentralized_controller = consensus.DistributedController(network, K)
simulator = consensus.ConsensusSimulator(N, nSamples)
networkss, positions, ref_vel, est_vel, est_vels, biased_vels, accels = simulator.simulate_pos(100, 200, D, 100, optimal_controller)
    
# %% Training
thisFilename = 'flockingGNN' 
saveDirRoot = 'experiments' 
saveDir = os.path.join(saveDirRoot, thisFilename) 
today = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
saveDir = saveDir + '-%03d-' + today
# Create directory
if not os.path.exists(saveDir):
    os.makedirs(saveDir)

#   PyTorch seeds
torchState = torch.get_rng_state()
torchSeed = torch.initial_seed()
#   Numpy seeds
numpyState = np.random.RandomState().get_state()
    
nEpochs = 50 # Number of epochs
batchSize = 20 # Batch size
learningRate = 0.00025
beta1 = 0.9
beta2 = 0.999
hParamsLocalGNN = {} # Hyperparameters (hParams) for the Local GNN
hParamsLocalGNN['name'] = 'LocalGNN'
hParamsLocalGNN['archit'] = modules.LocalGNN_DB
hParamsLocalGNN['device'] = 'cuda:0' \
                                if (True and torch.cuda.is_available()) \
                                else 'cpu'
hParamsLocalGNN['dimNodeSignals'] = [4, 64] # Features per layer
hParamsLocalGNN['nFilterTaps'] = [4] # Number of filter taps
hParamsLocalGNN['bias'] = True 
hParamsLocalGNN['nonlinearity'] = nn.Tanh 
hParamsLocalGNN['dimReadout'] = [2] 
hParamsLocalGNN['dimEdgeFeatures'] = 1 

if torch.cuda.is_available():
    torch.cuda.empty_cache()
hParamsDict = deepcopy(hParamsLocalGNN)
thisName = hParamsDict.pop('name')
callArchit = hParamsDict.pop('archit')
thisDevice = hParamsDict.pop('device')

##############
# PARAMETERS #
##############

thisArchit = callArchit(**hParamsDict)
thisArchit.to(thisDevice)
thisOptim = optim.Adam(thisArchit.parameters(), lr = learningRate, betas = (beta1, beta2))
thisLossFunction = nn.MSELoss()
thisTrainer = training.TrainerFlocking
GraphNNs = model.Model(thisArchit, thisLossFunction, thisOptim, thisTrainer, thisDevice, thisName, saveDir)

#Train
GraphNNs.train(simulator, positions, ref_vel, est_vel, est_vels, biased_vels, accels, networkss, nEpochs, batchSize)

#Test
ref_vel_test = ref_vel[200:220]
est_vel_test = est_vel[200:220]
pos_test = positions[200:220]
_, _, est_vels_valid, biased_vels_valid, accels_valid = simulator.computeTrajectory_pos(GraphNNs.archit, pos_test, ref_vel_test, est_vel_test, 100)
cost = simulator.cost(est_vels_valid, biased_vels_valid, accels_valid)
print(cost)

5.1-5.2 Collision Avoidance

To consider the collision avoidance, we first create the Python function in the file $\p{networks}$ that computes the relative position difference between agents as:

def computeDifferences(u):
        #take input as: nSamples x tSamples x 2 x nAgents
    if len(u.shape) == 3:
        u = np.expand_dims(u, 1)
        hasTimeDim = False
    else:
        hasTimeDim = True
    nSamples = u.shape[0]
    tSamples = u.shape[1]
    nAgents = u.shape[3]    
    uCol_x = u[:,:,0,:].reshape((nSamples, tSamples, nAgents, 1))
    uRow_x = u[:,:,0,:].reshape((nSamples, tSamples, 1, nAgents))
    uDiff_x = uCol_x - uRow_x 
    uCol_y = u[:,:,1,:].reshape((nSamples, tSamples, nAgents, 1))
    uRow_y = u[:,:,1,:].reshape((nSamples, tSamples, 1, nAgents))
    uDiff_y = uCol_y - uRow_y
    uDistSq = uDiff_x ** 2 + uDiff_y ** 2
    uDiff_x = np.expand_dims(uDiff_x, 2)
    uDiff_y = np.expand_dims(uDiff_y, 2)
    uDiff = np.concatenate((uDiff_x, uDiff_y), 2)
    if not hasTimeDim:
        uDistSq = uDistSq.squeeze(1)
        uDiff = uDiff.squeeze(1)
        
    return uDiff, uDistSq

Based on (28)-(29), we create the python class $\p{OptimalControllerCollision}$ in the file $\p{consensus}$ for the optimal centralized controller with collision avoidance as:

class OptimalControllerCollision(ConsensusController):
    def __init__(self, sample_time=0.1, energy_weight=1):
        self.sample_time = sample_time
        self.energy_weight = energy_weight

    def __call__(self, gama, d0, pos, est_vel, biased_vel, step):
	
        ijDiffPos, ijDistSq = networks.computeDifferences(pos[:,step,:,:])
        gammaMask = (ijDistSq < (gama**2)).astype(ijDiffPos.dtype)	
        ijDiffPos = ijDiffPos * np.expand_dims(gammaMask, 1)
        ijDistSqInv = ijDistSq.copy() 
        ijDistSqInv[ijDistSq < 1e-9] = 1. 
        ijDistSqInv = 1./ijDistSqInv 
        ijDistSqInv[ijDistSq < 1e-9] = 0.  
        ijDistSqInv = np.expand_dims(ijDistSqInv, 1)
        mean_biased_vel  = np.mean(biased_vel[:,step,:,:], axis=2)
        mean_biased_vel = np.expand_dims(mean_biased_vel, 2)
        accel = -(est_vel[:,step,:,:] - mean_biased_vel) / (self.sample_time * (1 + self.energy_weight)) + 1/self.sample_time/2*2* np.sum(ijDiffPos * (d0*d0*(ijDistSqInv ** 2) + d0*d0*ijDistSqInv), axis = 3)
	
        return accel

which takes agent positions as additional inputs. We create the python function $\p{simulateposcollision}$ in the class $\p{ConsensusSimulator}$ to simulate agent trajectories as:

    def simulate_pos_collision(self, wx, wy, degree, steps, controller: ConsensusController):
	
        ref_vel, est_vel = self.initial()
        bias = self.bias()
        x_pos = np.random.uniform(0, wx, (self.nSamples, 1, self.N))
        y_pos = np.random.uniform(0, wy, (self.nSamples, 1, self.N))
        pos = np.concatenate((x_pos, y_pos), axis = 1)
        ref_vels = np.zeros((self.nSamples, self.tSamples, 2, 1))
        est_vels = np.zeros((self.nSamples, self.tSamples, 2, self.N))
        biased_vels = np.zeros((self.nSamples, self.tSamples, 2, self.N))
        accels = np.zeros((self.nSamples, self.tSamples, 2, self.N))
        positions = np.zeros((self.nSamples, self.tSamples, 2, self.N))
        networkss = np.zeros((self.nSamples, self.tSamples, self.N, self.N))		
        ref_vels[:,0,:,:] = ref_vel
        est_vels[:,0,:,:] = est_vel
        biased_vels[:,0,:,:] = bias + ref_vel
        positions[:,0,:,:] = pos
        accels[:,0,:,:] = controller(10, 2, positions, est_vels, biased_vels, 0)
        
        for i in range(self.nSamples):
            _, networkss[i,0,:,:] = networks.agent_communication_pos(positions[i,0,:,:], self.N, degree)	
        this_accel = accels[:,0,:,:].copy()
        this_accel[accels[:,0,:,:] > self.max_agent_acceleration] = self.max_agent_acceleration
        this_accel[accels[:,0,:,:] < -self.max_agent_acceleration] = -self.max_agent_acceleration
        accels[:,0,:,:] = this_accel

        for step in range(1, steps):
            ref_accel = self.random_vector(self.ref_vel_delta, 1, self.nSamples)
            ref_vels[:,step,:,:] = ref_vels[:,step-1,:,:] + ref_accel * self.sample_time
            est_vels[:,step,:,:] = accels[:,step-1,:,:] * self.sample_time + est_vels[:,step-1,:,:]
            positions[:,step,:,:] = accels[:,step-1,:,:] * (self.sample_time ** 2)/2 + est_vels[:,step-1,:,:] * self.sample_time + positions[:,step-1,:,:]
            for i in range(self.nSamples):
                _, networkss[i,step,:,:] = networks.agent_communication_pos(positions[i,step,:,:], self.N, degree)	
            biased_vels[:,step,:,:] = bias + ref_vels[:,step,:,:]			
            accels[:,step,:,:] = controller(10, 2, positions, est_vels, biased_vels, step)			
            this_accel = accels[:,step,:,:].copy()
            this_accel[accels[:,step,:,:] > self.max_agent_acceleration] = self.max_agent_acceleration
            this_accel[accels[:,step,:,:] < -self.max_agent_acceleration] = -self.max_agent_acceleration
            accels[:,step,:,:] = this_accel
		
        return networkss, positions, ref_vel, est_vel, est_vels, biased_vels, accels   

To learn a GNN controller that mimics this optimal centralized controlle, we now have an additional feature as the relative mass, and the file $\p{training}$ is modified as:

import torch
import numpy as np
import datetime
import networks
import copy

class Trainer:
    
    def __init__(self, model, simulator, positions, ref_vel, est_vel, est_vels, biased_vels, accels, S, nEpochs, batchSize):
        
        self.model = model
        self.est_vel = est_vel
        self.ref_vel = ref_vel
        self.est_vels = est_vels
        self.biased_vels = biased_vels
        self.accels = accels
        self.S = S
        self.nTrain = 200
        self.nTest = 20
        self.nValid = 20
        self.nEpochs = nEpochs
        self.simulator = simulator
        self.positions = positions
        self.validationInterval = self.nTrain//batchSize

        if self.nTrain < batchSize:
            self.nBatches = 1
            self.batchSize = [self.nTrain]
        elif self.nTrain % batchSize != 0:
            self.nBatches = np.ceil(self.nTrain/batchSize).astype(np.int64)
            self.batchSize = [batchSize] * self.nBatches
            while sum(batchSize) != self.nTrain:
                self.batchSize[-1] -= 1
        else:
            self.nBatches = np.int(self.nTrain/batchSize)
            self.batchSize = [batchSize] * self.nBatches
            
        self.batchIndex = np.cumsum(self.batchSize).tolist()
        self.batchIndex = [0] + self.batchIndex

class TrainerFlocking(Trainer):
    
    def __init__(self, model, simulator, positions, ref_vel, est_vel, est_vels, biased_vels, accels, S, nEpochs, batchSize):
        
        super().__init__(model, simulator, positions, ref_vel, est_vel, est_vels, biased_vels, accels, S, nEpochs, batchSize)

    def train(self):
        
        nTrain = self.nTrain
        thisArchit = self.model.archit
        thisLoss = self.model.loss
        thisOptim = self.model.optim
        thisDevice = self.model.device
        epoch = 0         
        pos_train = self.positions[0:self.nTrain]
        posDiff, posDistSq = networks.computeDifferences(pos_train)
        com_networks = copy.deepcopy(self.S[0:self.nTrain])
        com_networks = (np.abs(com_networks) > 1e-9).astype(self.positions.dtype)
        com_networks = np.expand_dims(com_networks, 2)
        posDiff = posDiff * com_networks
        statePos = np.sum(posDiff, axis = 4)     
        xTrainAll = np.concatenate((self.est_vels[0:self.nTrain].copy(), self.biased_vels[0:self.nTrain].copy(), statePos), axis = 2)
        yTrainAll = self.accels[0:self.nTrain].copy()
        StrainAll = self.S[0:self.nTrain].copy()

        while epoch < self.nEpochs:

            randomPermutation = np.random.permutation(nTrain)
            idxEpoch = [int(i) for i in randomPermutation]

            batch = 0 
            while batch < self.nBatches:

                thisBatchIndices = idxEpoch[self.batchIndex[batch] : self.batchIndex[batch+1]]
                xTrain = xTrainAll[thisBatchIndices]
                yTrain = yTrainAll[thisBatchIndices]
                Strain = StrainAll[thisBatchIndices]                    
                xTrain = torch.tensor(xTrain, device = thisDevice)
                Strain = torch.tensor(Strain, device = thisDevice)
                yTrain = torch.tensor(yTrain, device = thisDevice)
                startTime = datetime.datetime.now()
                thisArchit.zero_grad()
                yHatTrain = thisArchit(xTrain, Strain)
                lossValueTrain = thisLoss(yHatTrain, yTrain)
                lossValueTrain.backward()
                thisOptim.step()
                endTime = datetime.datetime.now()
                timeElapsed = abs(endTime - startTime).total_seconds()
                del xTrain
                del Strain
                del yTrain
                del lossValueTrain

                #\\\\\\\
                #\\\ VALIDATION
                #\\\\\\\

                if (epoch * self.nBatches + batch) % self.validationInterval == 0:
                    
                    startTime = datetime.datetime.now()
                    ref_vel = self.ref_vel[220:240]
                    est_vel = self.est_vel[220:240]
                    position = self.positions[220:240]               
                    _, _, est_vels_valid, biased_vels_valid, accels_valid = self.simulator.computeTrajectory_pos_collision(thisArchit, position, ref_vel, est_vel, 100)
                    accValid = self.simulator.cost(est_vels_valid, biased_vels_valid, accels_valid)
                    endTime = datetime.datetime.now()
                    timeElapsed = abs(endTime - startTime).total_seconds()

                    print("\t(E: %2d, B: %3d) %8.4f - %6.4fs" % (
                            epoch+1, batch+1,
                            accValid, 
                            timeElapsed), end = ' ')
                    print("[VALIDATION] \n", end = '')

                    if epoch == 0 and batch == 0:
                        bestScore = accValid
                        bestEpoch, bestBatch = epoch, batch
                        self.model.save(label = 'Best')
                        # Start the counter
                    else:
                        thisValidScore = accValid
                        if thisValidScore < bestScore:
                            bestScore = thisValidScore
                            bestEpoch, bestBatch = epoch, batch
                            print("\t=> New best achieved: %.4f" % (bestScore))
                            self.model.save(label = 'Best')
                    del ref_vel
                    del est_vel
                batch += 1
            epoch += 1

        self.model.save(label = 'Last')
        #################
        # TRAINING OVER #
        #################
        if self.nEpochs == 0:
            self.model.save(label = 'Best')
            self.model.save(label = 'Last')
            print("\nWARNING: No training. Best and Last models are the same.\n")
        self.model.load(label = 'Best')
        if self.nEpochs > 0:
            print("\t=> Best validation achieved (E: %d, B: %d): %.4f" % (
                    bestEpoch + 1, bestBatch + 1, bestScore))

Here, for the validation, we similarly create the Python function $\p{computeTrajectoryposcollision}$ in the class $\p{ConsensusSimulator}$ as:

    def computeTrajectory_pos_collision(self, archit, position, ref_vel, est_vel, step):
        
        batchSize = est_vel.shape[0]
        nAgents = est_vel.shape[2]
        ref_vel = ref_vel.squeeze(2)
        architDevice = list(archit.parameters())[0].device
        est_vels = np.zeros((batchSize, step, 2, nAgents), dtype = np.float)
        ref_vels = np.zeros((batchSize, step, 2), dtype = np.float)
        positions = np.zeros((batchSize, step, 2, nAgents), dtype = np.float)
        biased_ref_vels = np.zeros((batchSize, step, 2, nAgents), dtype = np.float)
        bias = np.zeros((batchSize, step, 2, nAgents))
        accels = np.zeros((batchSize, step, 2, nAgents), dtype=np.float)
        states = np.zeros((batchSize, step, 6, nAgents), dtype=np.float)
        graphs = np.zeros((batchSize, step, nAgents, nAgents), dtype = np.float)
        com_networks = np.zeros((batchSize, step, nAgents, nAgents), dtype = np.float)           
        est_vels[:,0,:,:] = est_vel.copy()
        ref_vels[:,0,:] = ref_vel.copy()
        bias[:,0,:,:] = np.random.normal(loc=0,scale=(4 * np.sqrt(2/np.pi)),size=(batchSize, 2, nAgents))
        ref_vel_temp = np.expand_dims(ref_vels[:,0,:], 2)
        biased_ref_vels[:,0,:,:] = ref_vel_temp + bias[:,0,:,:]
        positions[:,0,:,:] = position[:,0,:,:]      

        for i in range(20):
            com_networks[i,0,:,:], graphs[i,0,:,:] = networks.agent_communication_pos(positions[i,0,:,:], nAgents, 2)
        
        for t in range(1, step):
            
            posDiff, posDistSq = networks.computeDifferences(positions[:,t-1,:,:])
            com_network = copy.deepcopy(com_networks[:,t-1,:,:])
            com_network = np.expand_dims(com_network, 1)
            posDiff = posDiff * com_network
            statePos = np.sum(posDiff, axis = 3)
            statePos = np.expand_dims(statePos, 1)               
            thisbiased_ref_vels = np.expand_dims(biased_ref_vels[:,t-1,:,:], 1)
            thisest_vels = np.expand_dims(est_vels[:,t-1,:,:], 1)
            thisState = np.concatenate((thisest_vels, thisbiased_ref_vels, statePos), axis = 2)
            states[:,t-1,:,:] = thisState.squeeze(1)

            x = torch.tensor(states[:,0:t,:,:], device = architDevice)
            S = torch.tensor(graphs[:,0:t,:,:], device = architDevice)
            with torch.no_grad():
                thisaccels = archit(x, S)
            thisaccels = thisaccels.cpu().numpy()[:,-1,:,:]
            thisaccels[thisaccels > 3] = 3
            thisaccels[thisaccels < -3] = -3
            accels[:,t-1,:,:] = thisaccels
            est_vels[:,t,:,:] = accels[:,t-1,:,:] * 0.1 +est_vels[:,t-1,:,:]
            positions[:,t,:,:] = accels[:,t-1,:,:] * (0.1 ** 2)/2 + est_vels[:,t-1,:,:] * 0.1 + positions[:,t-1,:,:]
            ref_accel = np.random.normal(loc=0,scale=(self.ref_vel_delta * np.sqrt(2/np.pi)),size=(batchSize, 2))
            ref_vels[:,t,:] = ref_vels[:,t-1,:] + ref_accel * 0.1
            ref_vel_temp = np.expand_dims(ref_vels[:,t,:], 2)
            biased_ref_vels[:,t,:,:] = ref_vel_temp + bias[:,0,:,:]
            for i in range(20):
                com_networks[i,t,:,:], graphs[i,t,:,:] = networks.agent_communication_pos(positions[i,t,:,:], nAgents, 2)   

        posDiff, posDistSq = networks.computeDifferences(positions[:,-1,:,:])
        com_network = copy.deepcopy(com_networks[:,-1,:,:])
        com_network = np.expand_dims(com_network, 1)
        posDiff = posDiff * com_network
        statePos = np.sum(posDiff, axis = 3)
        statePos = np.expand_dims(statePos, 1)                  
        thisbiased_ref_vels = np.expand_dims(biased_ref_vels[:,-1,:,:], 1)
        thisest_vels = np.expand_dims(est_vels[:,-1,:,:], 1)
        thisState = np.concatenate((thisest_vels, thisbiased_ref_vels, statePos), axis = 2)
        states[:,-1,:,:] = thisState.squeeze(1)
        x = torch.tensor(states, device = architDevice)
        S = torch.tensor(graphs, device = architDevice)
        with torch.no_grad():
            thisaccels = archit(x, S)
        thisaccels = thisaccels.cpu().numpy()[:,-1,:,:]
        thisaccels[thisaccels > 3] = 3
        thisaccels[thisaccels < -3] = -3
        accels[:,-1,:,:] = thisaccels
            
        return ref_vel, est_vel, est_vels, biased_ref_vels, accels

By aggregating all above modifications, we can perform the GNN training in the file $\p{main}$ as:

from importlib import reload
import numpy as np
import consensus
import networks
import os
import torch; torch.set_default_dtype(torch.float64)
import torch.nn as nn
import torch.optim as optim
import modules
import training
import model
import datetime
from copy import deepcopy
consensus = reload(consensus)
networks = reload(networks)

# %% Parameters
D = 2
K = 3
nSamples = 240
N = 100
step = 100
network = np.zeros((nSamples, N, N))
com_network = np.zeros((nSamples, step, N, N))

# %% Networks
for i in range(nSamples):
    network[i,:,:], com_network[i,0,:,:] = networks.agent_communication(100, 200, N, D)
    for t in range(1, step):
        com_network[i,t,:,:] = deepcopy(com_network[i,0,:,:])

# %% Dataset generation
optimal_controller = consensus.OptimalController()
decentralized_controller = consensus.DistributedController(network, K)
simulator = consensus.ConsensusSimulator(N, nSamples)
optimal_controller_collision = consensus.OptimalControllerCollision()
networkss, positions, ref_vel, est_vel, est_vels, biased_vels, accels = simulator.simulate_pos_collision(100, 200, D, 100, optimal_controller_collision)

# %% Training
thisFilename = 'flockingGNN' 
saveDirRoot = 'experiments' 
saveDir = os.path.join(saveDirRoot, thisFilename) 
today = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
saveDir = saveDir + '-%03d-' + today
# Create directory
if not os.path.exists(saveDir):
    os.makedirs(saveDir)

#   PyTorch seeds
torchState = torch.get_rng_state()
torchSeed = torch.initial_seed()
#   Numpy seeds
numpyState = np.random.RandomState().get_state()
    
nEpochs = 50 # Number of epochs
batchSize = 20 # Batch size
learningRate = 0.00025
beta1 = 0.9
beta2 = 0.999

hParamsLocalGNN = {} # Hyperparameters (hParams) for the Local GNN
hParamsLocalGNN['name'] = 'LocalGNN'
hParamsLocalGNN['archit'] = modules.LocalGNN_DB
hParamsLocalGNN['device'] = 'cuda:0' \
                                if (True and torch.cuda.is_available()) \
                                else 'cpu'
hParamsLocalGNN['dimNodeSignals'] = [6, 64] # Features per layer
hParamsLocalGNN['nFilterTaps'] = [4] # Number of filter taps
hParamsLocalGNN['bias'] = True 
hParamsLocalGNN['nonlinearity'] = nn.Tanh 
hParamsLocalGNN['dimReadout'] = [2] 
hParamsLocalGNN['dimEdgeFeatures'] = 1 

if torch.cuda.is_available():
    torch.cuda.empty_cache()
hParamsDict = deepcopy(hParamsLocalGNN)
thisName = hParamsDict.pop('name')
callArchit = hParamsDict.pop('archit')
thisDevice = hParamsDict.pop('device')

##############
# PARAMETERS #
##############

thisArchit = callArchit(**hParamsDict)
thisArchit.to(thisDevice)
thisOptim = optim.Adam(thisArchit.parameters(), lr = learningRate, betas = (beta1, beta2))
thisLossFunction = nn.MSELoss()
thisTrainer = training.TrainerFlocking
GraphNNs = model.Model(thisArchit, thisLossFunction, thisOptim, thisTrainer, thisDevice, thisName, saveDir)

#Train
GraphNNs.train(simulator, positions, ref_vel, est_vel, est_vels, biased_vels, accels, networkss, nEpochs, batchSize)

#Test
ref_vel_test = ref_vel[200:220]
est_vel_test = est_vel[200:220]
pos_test = positions[200:220]
_, _, est_vels_valid, biased_vels_valid, accels_valid = simulator.computeTrajectory_pos_collision(GraphNNs.archit, pos_test, ref_vel_test, est_vel_test, 100)
cost = simulator.cost(est_vels_valid, biased_vels_valid, accels_valid)
print(cost)

6 plot

Plotting figure is straightforward. Given the estimation velocities and observed biased reference velocities obtained from the simulation (Python function $\p{simulate}$, $\p{simulatepos}$ or $\p{simulateposcollision}$), we draw the velocity trajectories for some representative agents as:

trajec_num = 5
Agent_num1 = 5
Agent_num2 = 25
Agent_num3 = 45
Agent_num4 = 65
Agent_num5 = 85
Mean = np.mean(biased_vels[trajec_num, :, 0, :], axis = 1)
A1 = est_vels[trajec_num, :, 0, Agent_num1]
A2 = est_vels[trajec_num, :, 0, Agent_num2]
A3 = est_vels[trajec_num, :, 0, Agent_num3]
A4 = est_vels[trajec_num, :, 0, Agent_num4]
A5 = est_vels[trajec_num, :, 0, Agent_num5]
plt.plot(A1,label='Agent1')
plt.plot(A2,label='Agent2')
plt.plot(A3,label='Agent3')
plt.plot(A4,label='Agent4')
plt.plot(A5,label='Agent5')
plt.plot(Mean,label='Mean')
plt.legend()
plt.show()

where "trajecnum" is the number of trajectory, "agentnum" is the number of representative agent, and "0" in the third dimension represents the velocity along x-axis, which can be "1" as well representing the velocity along y-axis.