space

app.py

@@ -0,0 +1,135 @@
+import numpy as np
+import torch
+import gym
+from models.attention_model_wrapper import Agent
+device = 'cpu'
+ckpt_path = './runs/tsp-v0__ppo_or__1__1678160003/ckpt/12000.pt'
+agent = Agent(device=device, name='tsp').to(device)
+agent.load_state_dict(torch.load(ckpt_path))
+
+from wrappers.syncVectorEnvPomo import SyncVectorEnv
+from wrappers.recordWrapper import RecordEpisodeStatistics
+
+env_id = 'tsp-v0'
+env_entry_point = 'envs.tsp_vector_env:TSPVectorEnv'
+seed = 0
+
+gym.envs.register(
+    id=env_id,
+    entry_point=env_entry_point,
+)
+
+def make_env(env_id, seed, cfg={}):
+    def thunk():
+        env = gym.make(env_id, **cfg)
+        env = RecordEpisodeStatistics(env)
+        env.seed(seed)
+        env.action_space.seed(seed)
+        env.observation_space.seed(seed)
+        return env
+    return thunk
+
+
+def inference(data):
+    envs = SyncVectorEnv([make_env(env_id, seed, dict(n_traj=1, 
+                                                    max_nodes = len(data),
+                                                    eval_data = 'from_input',
+                                                    eval_data_from_input = data))])
+
+    trajectories = []
+    agent.eval()
+    obs = envs.reset()
+    done = np.array([False])
+    while not done.all():
+      # ALGO LOGIC: action logic
+        with torch.no_grad():
+            action, logits = agent(obs)
+        obs, reward, done, info = envs.step(action.cpu().numpy())
+        trajectories.append(action.cpu().numpy())
+    nodes_coordinates = obs['observations'][0]
+    final_return = info[0]['episode']['r']
+    resulting_traj = np.array(trajectories)[:,0,0]
+    return resulting_traj, final_return
+
+default_data = np.array([[0.5488135 , 0.71518937],
+       [0.60276338, 0.54488318],
+       [0.4236548 , 0.64589411],
+       [0.43758721, 0.891773  ],
+       [0.96366276, 0.38344152],
+       [0.79172504, 0.52889492],
+       [0.56804456, 0.92559664],
+       [0.07103606, 0.0871293 ],
+       [0.0202184 , 0.83261985],
+       [0.77815675, 0.87001215],])
+
+#@title Helper function for plotting
+# colorline taken from https://nbviewer.org/github/dpsanders/matplotlib-examples/blob/master/colorline.ipynb
+import matplotlib.pyplot as plt
+from matplotlib.collections import LineCollection
+from matplotlib.colors import ListedColormap, BoundaryNorm
+
+def make_segments(x, y):
+    '''
+    Create list of line segments from x and y coordinates, in the correct format for LineCollection:
+    an array of the form   numlines x (points per line) x 2 (x and y) array
+    '''
+
+    points = np.array([x, y]).T.reshape(-1, 1, 2)
+    segments = np.concatenate([points[:-1], points[1:]], axis=1)
+    
+    return segments
+
+def colorline(x, y, z=None, cmap=plt.get_cmap('copper'), norm=plt.Normalize(0.0, 1.0), linewidth=1, alpha=1.0):
+    '''
+    Plot a colored line with coordinates x and y
+    Optionally specify colors in the array z
+    Optionally specify a colormap, a norm function and a line width
+    '''
+    
+    # Default colors equally spaced on [0,1]:
+    if z is None:
+        z = np.linspace(0.3, 1.0, len(x))
+           
+    # Special case if a single number:
+    if not hasattr(z, "__iter__"):  # to check for numerical input -- this is a hack
+        z = np.array([z])
+        
+    z = np.asarray(z)
+    
+    segments = make_segments(x, y)
+    lc = LineCollection(segments, array=z, cmap=cmap, norm=norm, linewidth=linewidth, alpha=alpha)
+    
+    ax = plt.gca()
+    ax.add_collection(lc)
+    
+    return lc
+
+def plot(coords):
+    fig = plt.figure()
+    x,y = coords.T
+    lc = colorline(x,y,cmap='Reds')
+    plt.axis('square')
+    return fig
+
+import gradio as gr
+
+def run_inference(data):
+    data = data.astype(float).to_numpy()
+    resulting_traj, final_return = inference(data)
+    result_text = f'Planned Tour:\t{resulting_traj}\nTotal tour length:\t{final_return[0]:.2f}'
+    return [plot(data[resulting_traj]),result_text]
+
+demo = gr.Interface(run_inference, gr.Dataframe(
+                                              label = 'Input',
+                                              headers=['x','y'], 
+                                              row_count=10,
+                                              col_count=(2, "fixed"),
+                                              max_rows = 10,
+                                              value = default_data.tolist(),
+                                              overflow_row_behaviour = 'show_ends'
+                                              ), 
+                    [gr.Plot(label= 'Results Visualization'),
+                     gr.Code(label= 'Results',
+                                interactive=False)])
+demo.launch(share = True)
+

envs/tsp_vector_env.py

@@ -25,6 +25,9 @@ class TSPVectorEnv(gym.Env):
         self.eval_data = False
         self.eval_partition = "test"
         self.eval_data_idx = 0
+
+        self.eval_data_from_input = None
+
         assign_env_config(self, kwargs)
 
         obs_dict = {"observations": spaces.Box(low=0, high=1, shape=(self.max_nodes, 2))}
@@ -50,7 +53,9 @@ class TSPVectorEnv(gym.Env):
         self.last = np.zeros(self.n_traj, dtype=int)  # idx of the first elem
         self.first = np.zeros(self.n_traj, dtype=int)  # idx of the first elem
 
-        if self.eval_data:
+        if self.eval_data == 'from_input':
+            self._load_orders_from_input()
+        elif self.eval_data:
             self._load_orders()
         else:
             self._generate_orders()
@@ -59,6 +64,9 @@ class TSPVectorEnv(gym.Env):
         self.done = False
         return self.state
 
+    def _load_orders_from_input(self):
+        self.nodes = self.eval_data_from_input.copy()
+
     def _load_orders(self):
         self.nodes = np.array(TSPDataset[self.eval_partition, self.max_nodes, self.eval_data_idx])
 

requirements.txt

@@ -0,0 +1,5 @@
+matplotlib
+torch==1.13.0
+torchvision==0.14.0
+gym==0.23.1
+numpy==1.23.4