Delete actor.py

actor.py

@@ -1,751 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from Actor.graph import Graph
-from Actor.fleet import Fleet
-from utils.actor_utils import widen_data, select_data
-from Actor.normalization import Normalization
-
-
-#This is the updated version of the actor
-class Actor(nn.Module):
-
-    def __init__(self, model=None, num_movers=5, num_neighbors_encoder=5,
-                 num_neighbors_action=5, normalize=False, use_fleet_attention=True,
-                 device='cpu'):
-        super().__init__()
-
-        self.device = device
-        self.num_movers = num_movers
-        self.num_neighbors_encoder = num_neighbors_encoder
-        self.num_neighbors_action = num_neighbors_action
-
-        self.apply_normalization = normalize
-        self.normalization = None
-        self.use_fleet_attention = use_fleet_attention
-
-
-        if model is None:
-            self.mode = 'nearest_neighbors'
-        else:
-            self.encoder = model.encoder.to(self.device)
-            self.decoder = model.decoder.to(self.device)
-            self.projections = model.projections.to(self.device)
-            self.fleet_attention = model.fleet_attention.to(self.device)
-            self.mode = 'train'
-
-
-        self.sample_size = 1
-
-
-    def train_mode(self, sample_size=1):
-        self.train()
-        self.sample_size = sample_size
-        self.mode = 'train'
-
-
-    def greedy_search(self):
-        self.eval()
-        self.mode = 'greedy'
-
-
-    def nearest_neighbors(self):
-        self.eval()
-        self.mode = 'nearest_neighbors'
-
-
-    def sample_mode(self, sample_size=10):
-        self.sample_size = sample_size
-        self.eval()
-        self.mode = 'sample'
-
-
-    def beam_search(self, sample_size=10):
-        self.sample_size=sample_size
-        self.eval()
-        self.mode = 'beam_search'
-
-
-    def update_batch_size(self):
-        self.batch_size = self.fleet.time.shape[0]
-        self.fleet.batch_size = self.fleet.time.shape[0]
-        self.graph.batch_size  = self.fleet.time.shape[0]
-
-
-
-    def forward(self, batch, *args, **kwargs):
-
-        graph_data, fleet_data = batch
-
-        self.original_batch_size = graph_data['distance_matrix'].shape[0]
-        self.batch_size = self.original_batch_size
-
-        self.num_nodes = graph_data['distance_matrix'].shape[1]
-        self.num_cars = fleet_data['start_time'].shape[1]
-
-
-        self.graph = Graph(graph_data, device=self.device)
-        self.fleet = Fleet(fleet_data, num_nodes=self.num_nodes, device=self.device)
-        self.num_nodes = self.graph.distance_matrix.shape[1]
-        self.update_batch_size()
-
-        self.normalization = Normalization(self, normalize_position=True)
-        if self.apply_normalization:
-            self.normalization.normalize(self)
-
-        self.num_depots = self.fleet.num_depots.max().item()
-        self.num_movers_corrected = int(min(max(self.num_movers, self.num_depots), self.num_cars))
-
-
-        if self.mode != 'nearest_neighbors':
-            encoder_input = self.graph.construct_vector()
-            encoder_mask = self.compute_encoder_mask()
-            self.node_embeddings = self.encoder(encoder_input, encoder_mask)
-            self.node_projections = self.projections(self.node_embeddings)
-
-
-        if self.mode == 'sample':
-            widen_data(self, include_embeddings=True, include_projections=True)
-            self.update_batch_size()
-
-
-        self.log_probs = torch.zeros(self.batch_size).to(self.device)
-        self.counter = 0
-        while self.loop_condition() and (self.counter < self.num_nodes*4):
-
-            unavailable_moves = self.check_non_depot_options(use_time=True)
-            mover_indices = self.get_mover_indices(unavailable_moves=unavailable_moves)
-            action_mask = self.compute_action_mask(mover_indices=mover_indices,
-                                                   unavailable_moves=unavailable_moves)
-
-            if self.mode != 'nearest_neighbors':
-                decoder_input = self.construct_decoder_input(mover_indices=mover_indices)
-                decoder_mask = self.compute_decoder_mask(mover_indices=mover_indices,
-                                                         unavailable_moves=unavailable_moves)
-
-                decoder_output = self.decoder(decoder_input=decoder_input,
-                                            projections=self.node_projections,
-                                            mask=decoder_mask)
-            else:
-                decoder_output = None
-
-            next_node, car_to_move, log_prob = self.compute_action(decoder_output, action_mask, mover_indices)
-
-            if self.mode == 'beam_search':
-                widen_data(self, include_embeddings=True, include_projections=True)
-                self.update_batch_size()
-
-            self.log_probs += log_prob
-
-            self.update_time(next_node, car_to_move)
-            self.update_distance(next_node, car_to_move)
-            self.update_node_path(next_node, car_to_move)
-            self.update_traversed_nodes()
-
-            #self.return_to_depot()
-            #self.update_traversed_nodes()
-
-            if (self.mode == 'beam_search') and (self.counter > 0):
-                self.consolidate_beams()
-
-            self.update_batch_size()
-            self.counter += 1
-
-        self.return_to_depot_1()
-
-
-        if self.mode in {'beam_search', 'sample'}:
-            # we now must select out the best k elments
-
-            incomplete = self.check_complete().float()
-            cost = self.compute_cost()
-
-            max_cost = cost.max()
-            masked_cost = (1 - incomplete)*cost + incomplete*max_cost*10
-
-            p = masked_cost.reshape(self.original_batch_size, self.sample_size)
-            a = torch.argmin(p, dim=1)
-
-            b = torch.arange(self.original_batch_size).to(self.device)
-            b = b * self.sample_size
-
-            index = a + b
-            select_data(self, index=index, include_projections=True, include_embeddings=True)
-
-
-        self.update_batch_size()
-        if self.apply_normalization:
-            self.normalization.inverse_normalize(self)
-
-        total_distance = self.fleet.distance.sum(dim=1).squeeze(1).detach()
-        total_time = self.fleet.time.sum(dim=1).squeeze(1).detach()
-        total_late_time = self.fleet.late_time.sum(dim=1).squeeze(1).detach()
-
-        output = {
-            'distance': total_distance.detach(),
-            'total_time': total_time.detach(),
-            'log_probs': self.log_probs,
-            'late_time': total_late_time.detach(),
-            'incomplete': self.check_complete(),
-            'path': self.fleet.path,
-            'arrival_times': self.fleet.arrival_times
-        }
-
-        return output
-
-
-    def consolidate_beams(self):
-        p = self.log_probs.reshape(self.original_batch_size, self.sample_size * self.sample_size)
-        a = torch.topk(p, dim=1, k=self.sample_size, largest=True)[1]
-
-        b = torch.arange(self.original_batch_size).unsqueeze(1).repeat(1, self.sample_size).to(self.device)
-        b = b * self.sample_size * self.sample_size
-
-        ind = (a + b).reshape(-1)
-        select_data(self, index=ind, include_projections=True, include_embeddings=True)
-
-
-    def adjust_arrival_times(self):
-
-        if self.apply_normalization and (self.normalization_params is not None):
-            num_steps = self.fleet.arrival_times.shape[2]
-            a = self.normalization_params['earliest_start_time'].reshape(self.batch_size, 1, 1).repeat(1, self.num_cars, num_steps)
-            b = self.normalization_params['greatest_drive_time'].reshape(self.batch_size, 1, 1).repeat(1, self.num_cars, num_steps)
-            self.fleet.arrival_times = self.fleet.arrival_times*b + a
-
-            a = self.normalization_params['earliest_start_time']
-            b = self.normalization_params['greatest_drive_time']
-            self.fleet.late_time = self.fleet.late_time*b + a
-
-
-    def compute_encoder_mask(self):
-        #check for drive times being too long
-        time_window_non_compatibility = 1 - self.graph.time_window_compatibility
-
-        #compute diag mask
-        diag = torch.diag(torch.ones(self.num_nodes)).unsqueeze(0).repeat(self.batch_size, 1, 1).to(self.device)
-
-        # compute neighbors mask
-        m = (time_window_non_compatibility + diag > 0).float()
-
-        dist = self.graph.distance_matrix*(1 - m) + m*self.graph.max_dist*10
-
-        K = min(self.num_nodes, self.num_neighbors_encoder)
-        neighbors_index = torch.topk(dist, k=K, dim=2, largest=False)[1] # ~ [batch, num_nodes, num_neighbors]
-
-        a = torch.arange(self.num_nodes).reshape(1, 1, -1, 1).repeat(self.batch_size, self.num_nodes, 1, K).to(self.device)
-        b = neighbors_index.unsqueeze(2).repeat(1, 1, self.num_nodes, 1)
-        neighbors_mask = (a == b).float().sum(dim=3)
-
-
-        m = (time_window_non_compatibility == 0).float()
-        neighbs_time_mask = neighbors_mask*m
-
-        #compute depot mask
-        v = self.graph.depot.reshape(self.batch_size, 1, self.num_nodes).repeat(1, self.num_nodes, 1)
-        w = self.graph.depot.reshape(self.batch_size, self.num_nodes, 1).repeat(1, 1, self.num_nodes)
-        depot_mask = (v + w > 0).float()
-
-        diag = torch.diag(torch.ones(self.num_nodes)).unsqueeze(0).repeat(self.batch_size, 1, 1).to(self.device)
-
-        mask = (neighbs_time_mask + depot_mask + diag > 0).float()
-        return mask
-
-
-    def compute_depoyment_priority_score(self):
-        #number of available nodes
-        available_nodes = (self.check_non_depot_options().float() == 0).float()
-        num_available_nodes = available_nodes.sum(dim=2)
-
-        #excess
-        ind = self.fleet.node.reshape(self.batch_size, self.num_cars, 1).repeat(1, 1, self.num_nodes)
-        distances = torch.gather(self.graph.distance_matrix, dim=1, index=ind)
-
-        max_distance = self.graph.distance_matrix.reshape(self.batch_size, -1).max(dim=1)[0].reshape(
-                                    self.batch_size, 1, 1).repeat(1, self.num_cars, self.num_nodes)
-        excess = ((max_distance - distances)*available_nodes).sum(dim=2)
-
-
-        max_excess = excess.max()
-
-        score = (num_available_nodes + max_excess)*10 + excess
-        return score
-
-
-
-    def check_complete(self):
-        has_untraversed_nodes = ((self.fleet.traversed_nodes == 0).float().sum(dim=1) > 0)
-        #has_cars_in_field = ((self.fleet.node != self.fleet.depot).sum(dim=1) > 0)
-        #incomplete = (has_untraversed_nodes | has_cars_in_field)
-        incomplete = has_untraversed_nodes
-        return incomplete
-
-
-    def loop_condition(self):
-        a = self.check_complete().sum().item()
-        if a == 0:
-            return False
-        else:
-            return True
-
-
-    def compute_cost(self):
-
-        dist = self.fleet.distance.sum(dim=1).squeeze(1)
-        time = self.fleet.time.sum(dim=1).squeeze(1)
-        lateness = self.fleet.late_time.sum(dim=1).squeeze(1)
-
-        #normally we compute the cost as a linear combination of these three quantities
-        return time
-
-
-
-    def compute_total_distance(self):
-        p_2 = self.fleet.path
-        p_1 = torch.cat([p_2[:,:,-1:], p_2[:,:,:-1]], dim=2)
-
-        mat = self.graph.distance_matrix.reshape(self.batch_size, 1, self.num_nodes, self.num_nodes).repeat(1, self.num_cars, 1, 1)
-        ind_1 = p_1.unsqueeze(3).repeat(1, 1, 1, self.num_nodes)
-
-        d = torch.gather(mat, dim=2, index=ind_1)
-        ind_2 = p_2.unsqueeze(3)
-        pairwise_distances = torch.gather(d, dim=3, index=ind_2).squeeze(3)
-
-        self.car_distances = pairwise_distances.sum(dim=2)
-        self.total_distance = self.car_distances.sum(dim=1)
-
-        if self.apply_normalization and (self.normalization_params is not None):
-            distance_multiplier = self.normalization_params['greatest_distance']
-            self.total_distance = self.total_distance*distance_multiplier.reshape(self.batch_size)
-
-        return self.total_distance
-
-
-    def update_traversed_nodes(self):
-
-        path_length = self.fleet.path.shape[2]
-        a = self.fleet.path.reshape(self.batch_size, self.num_cars, path_length, 1).repeat(1, 1, 1, self.num_nodes)
-        b = torch.arange(self.num_nodes).reshape(
-            1, 1, 1, self.num_nodes).repeat(
-            self.batch_size, self.num_cars, path_length, 1).to(self.device)
-
-        s = (a == b).float().reshape(self.batch_size, self.num_cars*path_length, self.num_nodes).sum(dim=1)
-        self.fleet.traversed_nodes = (s > 0).float()
-
-
-    def compute_action(self, decoder_output, action_mask, mover_indices):
-
-        if self.mode == 'nearest_neighbors':
-
-            assert decoder_output is None
-
-            num_movers = mover_indices.shape[1]
-            mover_nodes = torch.gather(self.fleet.node, dim=1, index=mover_indices)
-
-            ind = mover_nodes.reshape(self.batch_size, num_movers, 1).repeat(1, 1, self.num_nodes)
-            distances = torch.gather(self.graph.distance_matrix, dim=1, index=ind)
-
-            a = mover_nodes.reshape(self.batch_size, num_movers, 1).repeat(1, 1, self.num_nodes)
-            b = torch.arange(self.num_nodes).reshape(1, 1, self.num_nodes).repeat(self.batch_size, num_movers, 1).to(self.device)
-            current_node_indicator = (a == b).float()
-
-            a = action_mask.reshape(self.batch_size, -1)
-            no_options = (a.sum(dim=1) == 0).reshape(-1, 1, 1).repeat(1, num_movers, self.num_nodes).float()
-
-            mask = action_mask * (1 - no_options) + current_node_indicator * no_options
-            masked_distances = distances*mask + self.graph.max_dist*(1 - mask)*10
-
-            x = masked_distances.reshape(self.batch_size, -1)
-            ind = torch.argmin(x, dim=1).unsqueeze(1)
-
-            mover_index = ind // self.num_nodes
-            next_node = ind % self.num_nodes
-            car_to_move = torch.gather(mover_indices, dim=1, index=mover_index)
-
-            log_prob = torch.zeros(self.batch_size).to(self.device)
-            return next_node, car_to_move, log_prob
-
-        else:
-
-            assert decoder_output is not None
-
-            num_movers = mover_indices.shape[1]
-            assert (num_movers == action_mask.shape[1]) and (num_movers == decoder_output.shape[1])
-
-            a = action_mask.reshape(self.batch_size, -1)
-            no_options = (a.sum(dim=1) == 0).reshape(-1, 1, 1).repeat(1, num_movers, self.num_nodes).float()
-
-            mover_nodes = torch.gather(self.fleet.node, dim=1, index=mover_indices)
-            a = torch.arange(self.num_nodes).reshape(1, 1, -1).repeat(self.batch_size, num_movers, 1).to(self.device)
-            b = mover_nodes.reshape(self.batch_size, num_movers, 1).repeat(1, 1, self.num_nodes)
-            default_option = (a == b).float()
-
-            mask = action_mask*(1 - no_options) + default_option*no_options
-            masked_decoder_output = decoder_output + mask.log()
-
-            x = masked_decoder_output.reshape(self.batch_size, -1)
-            probs = torch.softmax(x, dim=1)
-
-            if self.mode == 'greedy':
-                prob, ind = torch.max(probs, dim=1)[0].unsqueeze(1), torch.argmax(probs, dim=1).unsqueeze(1)
-
-            elif self.mode in {'train', 'sample'}:
-                ind = torch.multinomial(probs, num_samples=1)
-                prob = torch.gather(probs, dim=1, index=ind)
-
-            elif self.mode == 'beam_search':
-                prob, ind = torch.topk(probs, dim=1, k=self.sample_size, largest=True)
-
-            mover_index = ind // self.num_nodes
-            next_node = ind % self.num_nodes
-            car_to_move = torch.gather(mover_indices, dim=1, index=mover_index)
-
-            next_node = next_node.reshape(-1)
-            car_to_move = car_to_move.reshape(-1)
-            prob = prob.reshape(-1)
-
-            return next_node, car_to_move, prob.log()
-
-
-    def update_distance(self, next_node, car_to_move):
-        ind = car_to_move.reshape(self.batch_size, 1)
-        n = self.fleet.node.reshape(self.batch_size, self.num_cars)
-        current_node = torch.gather(n, dim=1, index=ind).squeeze(1)
-
-        #compute distance to next node
-        ind_1 = current_node.reshape(-1, 1, 1).repeat(1, 1, self.num_nodes)
-        distances = torch.gather(self.graph.distance_matrix, dim=1, index=ind_1).squeeze(1)
-        ind_2 = next_node.reshape(-1, 1)
-        dist_to_next_node = torch.gather(distances, dim=1, index=ind_2).squeeze(1)
-
-        #compute current distance
-        t = self.fleet.distance.reshape(self.batch_size, self.num_cars)
-        current_distance = torch.gather(t, dim=1, index=car_to_move.reshape(self.batch_size, 1)).squeeze(1)
-
-        #compute updated_distance
-        updated_distance = current_distance + dist_to_next_node
-
-        #compute mover mask
-        a = torch.arange(self.num_cars).reshape(1, -1).repeat(self.batch_size, 1).to(self.device)
-        b = car_to_move.reshape(self.batch_size, 1).repeat(1, self.num_cars)
-        update_mask = (a == b).float().unsqueeze(2)
-
-        #update distance
-        t = updated_distance.reshape(self.batch_size, 1, 1).repeat(1, self.num_cars, 1)
-        self.fleet.distance = self.fleet.distance * (1 - update_mask) + t * update_mask
-
-
-    def update_time(self, next_node, car_to_move):
-        #get current node of mover
-        ind = car_to_move.reshape(self.batch_size, 1)
-        n = self.fleet.node.reshape(self.batch_size, self.num_cars)
-        current_node = torch.gather(n, dim=1, index=ind).squeeze(1)
-
-        #compute time to next node
-        ind_1 = current_node.reshape(-1, 1, 1).repeat(1, 1, self.num_nodes)
-        drive_times = torch.gather(self.graph.time_matrix, dim=1, index=ind_1).squeeze(1)
-        ind_2 = next_node.reshape(-1, 1)
-        time_to_next_node = torch.gather(drive_times, dim=1, index=ind_2).squeeze(1)
-
-        #compute start time at next node
-        start_time = self.graph.start_time.reshape(self.batch_size, self.num_nodes)
-        ind = next_node.reshape(self.batch_size, 1)
-        next_start_time = torch.gather(start_time, dim=1, index=ind).squeeze(1)
-
-        #compute current time
-        t = self.fleet.time.reshape(self.batch_size, self.num_cars)
-        current_node_time = torch.gather(t, dim=1, index=car_to_move.reshape(self.batch_size, 1)).squeeze(1)
-
-        #compute updated_time
-        a = time_to_next_node + current_node_time
-        b = next_start_time
-        updated_time = (a > b).float()*a + (a <= b).float()*b
-
-        #compute end_time at next node
-        ind = next_node.reshape(-1, 1)
-        end_times = self.graph.end_time.reshape(self.batch_size, self.num_nodes)
-        end_time_next_node = torch.gather(end_times, dim=1, index=ind).squeeze(1)
-        late_time = F.relu(updated_time - end_time_next_node)
-
-
-        #compute mover mask
-        a = torch.arange(self.num_cars).reshape(1, -1).repeat(self.batch_size, 1).to(self.device)
-        b = car_to_move.reshape(self.batch_size, 1).repeat(1, self.num_cars)
-        update_mask = (a == b).float().unsqueeze(2)
-
-        #update time
-        t = updated_time.reshape(self.batch_size, 1, 1).repeat(1, self.num_cars, 1)
-        self.fleet.time = self.fleet.time*(1 - update_mask) + t*update_mask
-
-        # update_late_time
-        l = late_time.reshape(self.batch_size, 1, 1).repeat(1, self.num_cars, 1)
-        self.fleet.late_time = self.fleet.late_time*(1 - update_mask) + l*update_mask
-
-
-
-    def update_node_path(self, next_node, car_to_move):
-        #compute mover mask
-        a = torch.arange(self.num_cars).reshape(1, -1).repeat(self.batch_size, 1).to(self.device)
-        b = car_to_move.reshape(self.batch_size, 1).repeat(1, self.num_cars)
-        update_mask = (a == b).long()
-
-        new_node = next_node.reshape(self.batch_size, 1).repeat(1, self.num_cars)
-        self.fleet.node = update_mask*new_node + self.fleet.node*(1 - update_mask)
-
-        L = [self.fleet.path, self.fleet.node.unsqueeze(2)]
-        self.fleet.path = torch.cat(L, dim=2)
-
-        t = self.fleet.time.reshape(self.batch_size, self.num_cars, 1)
-        H = [self.fleet.arrival_times, t]
-        self.fleet.arrival_times = torch.cat(H, dim=2)
-
-
-
-    def check_non_depot_options(self, use_time=True):
-        '''
-        Output is a byte tensor of shape [batch, num_cars, num_nodes] with entry (i,j) = 1 if
-        the move of car i to node j is invalid
-        '''
-
-        if use_time:
-            #check for arrival times
-            too_far = 1 - self.check_arival_times()
-
-        #is depot
-        a = self.fleet.depot.reshape(self.batch_size, self.num_cars, 1).repeat(1, 1, self.num_nodes)
-        b = torch.arange(self.num_nodes).reshape(1, 1, -1).repeat(self.batch_size, self.num_cars, 1).to(self.device)
-        is_depot = (a == b)
-
-        #check traversed nodes
-        a = self.fleet.traversed_nodes.reshape(self.batch_size, 1, self.num_nodes).repeat(1, self.num_cars, 1)
-        traversed_nodes = (a == 1)
-
-        # has value of 1 if the move to that node is NOT possible
-        if use_time:
-            unavailable_moves = (too_far | is_depot | traversed_nodes)
-        else:
-            unavailable_moves = (is_depot | traversed_nodes)
-
-        return unavailable_moves
-
-
-    def compute_decoder_mask(self, mover_indices, unavailable_moves):
-
-        num_movers = mover_indices.shape[1]
-
-        mover_nodes = torch.gather(self.fleet.node, dim=1, index=mover_indices)
-
-        ind = mover_indices.reshape(self.batch_size, num_movers, 1).repeat(1, 1, self.num_nodes)
-        unavailable_moves = torch.gather(unavailable_moves, dim=1, index=ind)
-
-        no_options = ((1 - unavailable_moves.float()).sum(dim=2) == 0).unsqueeze(2).repeat(1, 1, self.num_nodes).float()
-
-        a = mover_nodes.reshape(self.batch_size, num_movers, 1).repeat(1, 1, self.num_nodes)
-        b = torch.arange(self.num_nodes).reshape(1, 1, -1).repeat(self.batch_size, num_movers, 1).to(self.device)
-        default_move = (a == b).float()
-
-        decoder_mask = (1 - unavailable_moves.float())*(1 - no_options) + default_move*no_options
-        return decoder_mask
-
-
-    def compute_action_mask(self, mover_indices, unavailable_moves):
-
-        if self.mode != 'nearest_neighbors':
-            ######################################################
-            d = torch.diag(torch.ones(self.num_nodes)).unsqueeze(0).repeat(self.batch_size, 1, 1).to(self.device)
-            ind = self.fleet.node.reshape(self.batch_size, self.num_cars, 1).repeat(1, 1, self.num_nodes)
-
-            diag = torch.gather(d, dim=1, index=ind)
-            distances = torch.gather(self.graph.distance_matrix, dim=1, index=ind)
-            m = (unavailable_moves.float() + diag > 0).float()
-
-            masked_distances = distances*(1 - m) + m*self.graph.max_dist*100
-            K = min(self.num_neighbors_action, self.num_nodes)
-            neighbor_indices = torch.topk(masked_distances, dim=2, k=K, largest=False)[1]
-
-            a = torch.arange(self.num_nodes).reshape(1, 1, -1, 1).repeat(self.batch_size, self.num_cars, 1, K).to(self.device)
-            b = neighbor_indices.reshape(self.batch_size, self.num_cars, 1, K).repeat(1, 1, self.num_nodes, 1)
-
-            non_neighbor_mask = ((a == b).float().sum(dim=3) == 0)
-            mask = 1 - (non_neighbor_mask | unavailable_moves).float()
-            ######################################################
-        else:
-            mask = 1 - unavailable_moves.float()
-
-        num_movers = mover_indices.shape[1]
-        ind = mover_indices.reshape(self.batch_size, num_movers, 1).repeat(1, 1, self.num_nodes)
-        action_mask = torch.gather(mask, dim=1, index=ind)
-        return action_mask
-
-
-
-    def construct_decoder_input(self, mover_indices):
-
-        embedding_size = self.node_embeddings.shape[2]
-        mask = self.check_non_depot_options(use_time=True)
-        mask = mask.permute(0, 2, 1).unsqueeze(3).repeat(1, 1, 1, embedding_size).float()
-
-        node_vectors = self.node_embeddings.reshape(self.batch_size, self.num_nodes, 1, embedding_size).repeat(1, 1, self.num_cars, 1)
-
-        #depot vector
-        start_ind = self.fleet.depot.reshape(self.batch_size, 1, self.num_cars, 1).repeat(1, 1, 1, embedding_size)
-        a = torch.gather(node_vectors, dim=1, index=start_ind)
-        depot_vector = a.squeeze(1)
-
-        #current node vector
-        current_ind = self.fleet.node.reshape(self.batch_size, 1, self.num_cars, 1).repeat(1, 1, 1, embedding_size)
-        b = torch.gather(node_vectors, dim=1, index=current_ind)
-        current_node_vector = b.squeeze(1)
-
-        #mean graph vector
-        mean_graph_vector = node_vectors.mean(dim=1)
-
-
-        #current feature values
-        feature_values = self.fleet.construct_vector()
-
-
-        #other cars in field
-        num_movers = mover_indices.shape[1]
-        if num_movers == 1:
-            movers_vector = current_node_vector*0
-        elif not self.use_fleet_attention:
-            a = mover_indices.reshape(self.batch_size, num_movers, 1).repeat(1, 1, self.num_cars)
-            b = torch.arange(self.num_cars).reshape(1, 1, self.num_cars).repeat(self.batch_size, num_movers, 1).to(self.device)
-            c = ((a == b).float().sum(dim=1) > 0).float()
-            movers_mask = c.reshape(self.batch_size, self.num_cars, 1).repeat(1, 1, embedding_size)
-            movers_vector = ((current_node_vector*movers_mask).sum(dim=1).unsqueeze(1) - current_node_vector)/(num_movers - 1)
-        else:
-            x = torch.cat([current_node_vector, feature_values], dim=2)
-            movers_vector = self.fleet_attention(x)
-
-
-        L = [current_node_vector, depot_vector, mean_graph_vector, movers_vector, feature_values]
-        pre_output = torch.cat(L, dim=2)
-
-
-        num_movers = mover_indices.shape[1]
-        ind = mover_indices.reshape(self.batch_size, num_movers, 1).repeat(1, 1, pre_output.shape[2])
-        output = torch.gather(pre_output, dim=1, index=ind)
-
-        return output
-
-
-    def get_mover_indices(self, unavailable_moves):
-        depot = self.fleet.depot.reshape(self.batch_size, self.num_cars).long()
-        current_node = self.fleet.node.reshape(self.batch_size, self.num_cars).long()
-
-        # find all cars with "no options"
-        has_option = ((1 - unavailable_moves.float()).sum(dim=2) > 0)
-
-        at_depot = (current_node == depot)
-        in_field = (current_node != depot)
-        active_in_field = in_field & has_option
-        active_at_depot = at_depot & has_option
-
-        deployment_score = self.compute_depoyment_priority_score()
-        max_deployment_score = deployment_score.max()
-
-        A = max_deployment_score * 100
-        B = deployment_score * 10
-
-        score = active_in_field.float() * A + active_at_depot.float() * B
-        score = score.reshape(self.batch_size, self.num_cars, 1).repeat(1, 1, self.num_depots)
-
-        a = self.fleet.depot.reshape(self.batch_size, self.num_cars, 1).repeat(1, 1, self.num_depots)
-        b = torch.arange(self.num_depots).reshape(1, 1, self.num_depots).repeat(self.batch_size, self.num_cars, 1).to(self.device)
-        m = (a == b).float()
-
-        score = score*m
-
-        K = self.num_movers_corrected
-        indices = torch.topk(score, k=K, dim=1, largest=True)[1]
-        indices = indices.reshape(self.batch_size, self.num_depots*K)
-        return indices
-
-
-    def check_arival_times(self):
-        #check for arrival times
-        d = self.graph.time_matrix.unsqueeze(1).repeat(1, self.num_cars, 1, 1)
-        ind = self.fleet.node.reshape(self.batch_size, self.num_cars, 1, 1).repeat(1, 1, 1, self.num_nodes)
-        drive_times = torch.gather(d, dim=2, index=ind).squeeze(2)
-        t = self.fleet.time.reshape(self.batch_size, self.num_cars, 1).repeat(1, 1, self.num_nodes)
-        arival_times = drive_times + t
-        b = self.graph.end_time.reshape(self.batch_size, 1, self.num_nodes).repeat(1, self.num_cars, 1)
-        attainable = (arival_times <= b)
-        return attainable
-
-
-    def return_to_depot(self):
-        unavailable_moves = 1 - self.check_non_depot_options(use_time=False).float()
-        return_to_depot = (unavailable_moves.reshape(self.batch_size, -1).sum(dim=1) == 0).float()
-        return_to_depot = return_to_depot.reshape(self.batch_size, 1).repeat(1, self.num_cars)
-
-
-        if return_to_depot.sum().item() > 0:
-
-            depot = self.fleet.depot.reshape(self.batch_size, self.num_cars).long()
-            node = self.fleet.node.reshape(self.batch_size, self.num_cars).long()
-
-            #compute next node
-            return_to_depot = return_to_depot.long()
-            next_node = return_to_depot*depot + (1 - return_to_depot)*node
-
-            #update time
-            return_to_depot = return_to_depot.unsqueeze(2).float()
-
-            current_node = self.fleet.node
-            ind_1 = current_node.reshape(self.batch_size, self.num_cars, 1).repeat(1, 1, self.num_nodes)
-            drive_times = torch.gather(self.graph.time_matrix, dim=1, index=ind_1)
-            ind_2 = next_node.reshape(self.batch_size, self.num_cars, 1)
-            time_to_next_node = torch.gather(drive_times, dim=2, index=ind_2)
-            self.fleet.time = self.fleet.time + time_to_next_node*return_to_depot
-
-            #update node
-            self.fleet.node = next_node
-
-            #update path
-            n = self.fleet.node.reshape(self.batch_size, self.num_cars, 1)
-            self.fleet.path = torch.cat([self.fleet.path, n], dim=2)
-
-            #update arrival time
-            t = self.fleet.time.reshape(self.batch_size, self.num_cars, 1)
-            self.fleet.arrival_times = torch.cat([self.fleet.arrival_times, t], dim=2)
-
-
-    def return_to_depot_1(self):
-
-        depot = self.fleet.depot.reshape(self.batch_size, self.num_cars).long()
-        node = self.fleet.node.reshape(self.batch_size, self.num_cars).long()
-
-        #compute next node
-        next_node = depot
-
-        #update time
-        current_node = self.fleet.node
-        ind_1 = current_node.reshape(self.batch_size, self.num_cars, 1).repeat(1, 1, self.num_nodes)
-        drive_times = torch.gather(self.graph.time_matrix, dim=1, index=ind_1)
-        ind_2 = next_node.reshape(self.batch_size, self.num_cars, 1)
-        time_to_next_node = torch.gather(drive_times, dim=2, index=ind_2)
-        self.fleet.time = self.fleet.time + time_to_next_node
-
-
-        #update distance
-        current_node = self.fleet.node
-        ind_1 = current_node.reshape(self.batch_size, self.num_cars, 1).repeat(1, 1, self.num_nodes)
-        drive_distances = torch.gather(self.graph.distance_matrix, dim=1, index=ind_1)
-        ind_2 = next_node.reshape(self.batch_size, self.num_cars, 1)
-        distance_to_next_node = torch.gather(drive_distances, dim=2, index=ind_2)
-        self.fleet.distance = self.fleet.distance + distance_to_next_node
-
-
-        #update node
-        self.fleet.node = next_node
-
-        #update path
-        n = self.fleet.node.reshape(self.batch_size, self.num_cars, 1)
-        self.fleet.path = torch.cat([self.fleet.path, n], dim=2)
-
-        #update arrival time
-        t = self.fleet.time.reshape(self.batch_size, self.num_cars, 1)
-        self.fleet.arrival_times = torch.cat([self.fleet.arrival_times, t], dim=2)