ATSP Insertion bugfix; return tour for revisions

eval_atsp/ASHPPEnv.py

@@ -0,0 +1,181 @@
+
+"""
+The MIT License
+
+Copyright (c) 2021 MatNet
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+"""
+
+from dataclasses import dataclass
+import torch
+import warnings
+
+from ATSProblemDef import get_random_problems
+
+
+@dataclass
+class Reset_State:
+    problems: torch.Tensor
+    # shape: (batch, node, node)
+
+
+@dataclass
+class Step_State:
+    BATCH_IDX: torch.Tensor
+    POMO_IDX: torch.Tensor
+    # shape: (batch, pomo)
+    current_node: torch.Tensor = None
+    # shape: (batch, pomo)
+    ninf_mask: torch.Tensor = None
+    # shape: (batch, pomo, node)
+
+
+class ASHPPEnv:
+    def __init__(self, **env_params):
+
+        # Const @INIT
+        ####################################
+        self.env_params = env_params
+        self.node_cnt = env_params['node_cnt']
+        self.pomo_size = env_params['pomo_size'] # pomo size if sample size here
+
+        # Const @Load_Problem
+        ####################################
+        self.batch_size = None
+        self.BATCH_IDX = None
+        self.POMO_IDX = None
+        # IDX.shape: (batch, pomo)
+        self.problems = None
+        # shape: (batch, node, node)
+
+        # Dynamic
+        ####################################
+        self.selected_count = None
+        self.current_node = None
+        # shape: (batch, pomo)
+        self.selected_node_list = None
+        # shape: (batch, pomo, 0~)
+
+        # STEP-State
+        ####################################
+        self.step_state = None
+
+    def load_problems(self, batch_size):
+        self.batch_size = batch_size
+        self.BATCH_IDX = torch.arange(self.batch_size)[:, None].expand(self.batch_size, self.pomo_size)
+        self.POMO_IDX = torch.arange(self.pomo_size)[None, :].expand(self.batch_size, self.pomo_size)
+
+        problem_gen_params = self.env_params['problem_gen_params']
+        self.problems = get_random_problems(batch_size, self.node_cnt, problem_gen_params)
+        # shape: (batch, node, node)
+
+    def load_problems_manual(self, problems):
+        # problems.shape: (batch, node, node)
+
+        self.batch_size = problems.size(0)
+        self.BATCH_IDX = torch.arange(self.batch_size)[:, None].expand(self.batch_size, self.pomo_size)
+        self.POMO_IDX = torch.arange(self.pomo_size)[None, :].expand(self.batch_size, self.pomo_size)
+        self.problems = problems
+        # shape: (batch, node, node)
+
+    def reset(self):
+        self.selected_count = 2 # Add starting and terminating ndoes
+        # Set current nodes as 0
+        self.current_node = torch.zeros((self.batch_size, self.pomo_size), dtype=torch.long)
+        # Set the last node as node - 1
+        self.last_node = torch.ones((self.batch_size, self.pomo_size), dtype=torch.long) * (self.node_cnt - 1)
+
+        # shape: (batch, pomo)
+        self.selected_node_list = self.current_node[:, :, None]
+        # shape: (batch, pomo, 0~)
+
+        self._create_step_state()
+
+        reward = None
+        done = False
+        return Reset_State(problems=self.problems), reward, done
+
+    def _create_step_state(self):
+        self.step_state = Step_State(BATCH_IDX=self.BATCH_IDX, POMO_IDX=self.POMO_IDX)
+        self.step_state.ninf_mask = torch.zeros((self.batch_size, self.pomo_size, self.node_cnt))
+        # shape: (batch, pomo, node)
+
+    def pre_step(self):
+        reward = None
+        done = False
+
+        # Set the starting and terminating nodes to -inf
+        self.step_state.ninf_mask[self.BATCH_IDX, self.POMO_IDX, 0] = float('-inf')
+        self.step_state.ninf_mask[self.BATCH_IDX, self.POMO_IDX, -1] = float('-inf')
+
+        # Set current node to 0
+        self.step_state.current_node = self.current_node
+        # Set last node to node - 1
+        self.step_state.last_node = self.last_node
+
+
+        return self.step_state, reward, done
+
+    def step(self, node_idx):
+        # node_idx.shape: (batch, pomo)
+
+        self.selected_count += 1
+        self.current_node = node_idx
+        # shape: (batch, pomo)
+        self.selected_node_list = torch.cat((self.selected_node_list, self.current_node[:, :, None]), dim=2)
+        # shape: (batch, pomo, 0~node)
+
+        self._update_step_state()
+
+        # returning values
+        done = (self.selected_count == self.node_cnt)
+        if done:
+            # Concat the terminating node (the last node) to the selected node list
+            self.current_node = torch.ones((self.batch_size, self.pomo_size), dtype=torch.long) * (self.node_cnt - 1)
+            self.selected_node_list = torch.cat((self.selected_node_list, self.current_node[:, :, None]), dim=2)
+            reward = -self._get_total_distance()  # Note the MINUS Sign ==> We MAXIMIZE reward
+            # shape: (batch, pomo)
+        else:    
+            reward = None
+        return self.step_state, reward, done
+
+    def _update_step_state(self):
+        self.step_state.current_node = self.current_node
+        # shape: (batch, pomo)
+        self.step_state.ninf_mask[self.BATCH_IDX, self.POMO_IDX, self.current_node] = float('-inf')
+        # shape: (batch, pomo, node)
+
+    def _get_total_distance(self):
+
+        node_from = self.selected_node_list[:, :, :-1]
+        # shape: (batch, pomo, node - 1)
+        node_to = self.selected_node_list.roll(dims=2, shifts=-1)[:, :, :-1]
+        # shape: (batch, pomo, node - 1)
+        batch_index = self.BATCH_IDX[:, :, None].expand(self.batch_size, self.pomo_size, self.node_cnt - 1)
+        # shape: (batch, pomo, node - 1)
+
+        selected_cost = self.problems[batch_index, node_from, node_to]
+        # shape: (batch, pomo, node - 1)
+        total_distance = selected_cost.sum(2)
+        # shape: (batch, pomo)
+
+        return total_distance