Update app.py

app.py

@@ -2,16 +2,17 @@ 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))
+
+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, map_location=torch.device("cpu")))
 
 from wrappers.syncVectorEnvPomo import SyncVectorEnv
 from wrappers.recordWrapper import RecordEpisodeStatistics
 
-env_id = 'tsp-v0'
-env_entry_point = 'envs.tsp_vector_env:TSPVectorEnv'
+env_id = "tsp-v0"
+env_entry_point = "envs.tsp_vector_env:TSPVectorEnv"
 seed = 0
 
 gym.envs.register(
@@ -19,6 +20,7 @@ gym.envs.register(
     entry_point=env_entry_point,
 )
 
+
 def make_env(env_id, seed, cfg={}):
     def thunk():
         env = gym.make(env_id, **cfg)
@@ -27,109 +29,124 @@ def make_env(env_id, seed, cfg={}):
         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))])
+    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
+        # 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]
+    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
+
+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):
-    '''
+
+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')
+    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)
-
+    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()