Spaces:
Runtime error
Runtime error
| 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) | |