Code comments and cleanup

Author: Kent Sommer

README.md

@@ -71,7 +71,7 @@ python test.py --weights trained/vin_28x28.pth --imsize 28 --k 36
 ## Results
 Gridworld | Sample One | Sample Two
 -- | --- | ---
-8x8 | <img src="results/8x8_2.png" width="450"> | <img src="results/8x8_3.png" width="450">
+8x8 | <img src="results/8x8_1.png" width="450"> | <img src="results/8x8_2.png" width="450">
 16x16 | <img src="results/16x16_1.png" width="450"> | <img src="results/16x16_2.png" width="450">
 
 ## Datasets

dataset/dataset.py

@@ -1,40 +1,42 @@
+import numpy as np
+
 import torch
 import torch.utils.data as data
-import numpy as np
 
 
 class GridworldData(data.Dataset):
-    def __init__(self, file, imsize, train=True, transform=None, target_transform=None):
+    def __init__(self, file, imsize, train=True, 
+                    transform=None, target_transform=None):
         assert file.endswith('.npz') # Must be .npz format
         self.file = file
         self.imsize = imsize
         self.transform = transform
         self.target_transform = target_transform
         self.train = train # training set or test set
 
-        self.images, self.S1, self.S2, self.labels = self._process(file, self.train)
+        self.images, self.S1, self.S2, self.labels =  \
+                                self._process(file, self.train)
 
     def __getitem__(self, index):
         img = self.images[index]
         s1 = self.S1[index]
         s2 = self.S2[index]
         label = self.labels[index]
-        
+        # Apply transform if we have one
         if self.transform is not None:
             img = self.transform(img)
         else: # Internal default transform: Just to Tensor
             img = torch.from_numpy(img)
-            
+        # Apply target transform if we have one
         if self.target_transform is not None:
             label = self.target_transform(label)
-            
-        # Ensure labels in naive float type
-        # DataLoader has bug with np.int/float type in default_collate()
         return img, int(s1), int(s2), int(label)
 
+
     def __len__(self):
         return self.images.shape[0]
 
+
     def _process(self, file, train):
         """Data format: A list, [train data, test data]
         Each data sample: label, S1, S2, Images, in this order.
@@ -50,15 +52,14 @@ def _process(self, file, train):
                 S1 = f['arr_5']
                 S2 = f['arr_6']
                 labels = f['arr_7']
-
+        # Set proper datatypes
         images = images.astype(np.float32)
         S1 = S1.astype(int) # (S1, S2) location are integers
         S2 = S2.astype(int)
         labels = labels.astype(int) # labels are integers
-
+        # Print number of samples
         if train:
             print("Number of Train Samples: {0}".format(images.shape[0]))
         else:
             print("Number of Test Samples: {0}".format(images.shape[0]))
-        
         return images, S1, S2, labels

dataset/make_training_data.py

@@ -9,19 +9,24 @@
 sys.path.remove('.')
 
 def extract_action(traj):
+    # Given a trajectory, outputs a 1D vector of 
+    #  actions corresponding to the trajectory. 
     n_actions = 8
-    action_vecs = np.asarray([[-1., 0.],[1.,0.],[0.,1.],[0.,-1.],[-1.,1.],[-1.,-1.],[1.,1.],[1.,-1.]])
+    action_vecs = np.asarray([[-1., 0.],[1.,0.],[0.,1.],[0.,-1.],[-1.,1.],
+                              [-1.,-1.],[1.,1.],[1.,-1.]])
     action_vecs[4:] = 1/np.sqrt(2) * action_vecs[4:]
     action_vecs = action_vecs.T
     state_diff = np.diff(traj, axis=0)
-    norm_state_diff = state_diff * np.tile(1/np.sqrt(np.sum(np.square(state_diff), axis=1)), (2, 1)).T
+    norm_state_diff = state_diff * np.tile(1/np.sqrt(np.sum(np.square(
+                                   state_diff), axis=1)), (2, 1)).T
     prj_state_diff = np.dot(norm_state_diff, action_vecs)
     actions_one_hot = np.abs(prj_state_diff -1)<0.00001
     actions = np.dot(actions_one_hot, np.arange(n_actions).T)
     return actions
 
 
-def make_data(dom_size, n_domains, max_obs, max_obs_size, n_traj, state_batch_size):
+def make_data(dom_size, n_domains, max_obs, 
+                max_obs_size, n_traj, state_batch_size):
 
     X_l = []
     S1_l = []
@@ -43,7 +48,6 @@ def make_data(dom_size, n_domains, max_obs, max_obs_size, n_traj, state_batch_si
             continue
         # Get final map
         im = obs.get_final()
-
         # Generate gridworld from obstacle map
         G = gridworld(im, goal[0], goal[1])
         # Get value prior
@@ -59,7 +63,8 @@ def make_data(dom_size, n_domains, max_obs, max_obs_size, n_traj, state_batch_si
                 image = 1 - im
                 # Resize domain and goal images and concate
                 image_data = np.resize(image, (1,1,dom_size[0],dom_size[1]))
-                value_data = np.resize(value_prior, (1,1,dom_size[0],dom_size[1]))
+                value_data = np.resize(value_prior, (1,1,dom_size[0],
+                                                         dom_size[1]))
                 iv_mixed = np.concatenate((image_data, value_data), axis=1)
                 X_current = np.tile(iv_mixed, (ns, 1, 1, 1))
                 # Resize states
@@ -86,18 +91,19 @@ def make_data(dom_size, n_domains, max_obs, max_obs_size, n_traj, state_batch_si
 
 def main(dom_size=[8,8], n_domains=15000, max_obs=30, max_obs_size=None, 
             n_traj=7, state_batch_size=1):
-
+    # Get path to save dataset
     save_path = "dataset/gridworld_{0}x{1}".format(dom_size[0], dom_size[1])
-
+    # Get training data
     print("Now making training data...")    
-    X_out_tr, S1_out_tr, S2_out_tr, Labels_out_tr = make_data(dom_size, n_domains, max_obs, 
-                                                    max_obs_size, n_traj, state_batch_size)
+    X_out_tr, S1_out_tr, S2_out_tr, Labels_out_tr = make_data(
+        dom_size, n_domains, max_obs, max_obs_size, n_traj, state_batch_size)
+    # Get testing data
     print("\nNow making  testing data...")
-    X_out_ts, S1_out_ts, S2_out_ts, Labels_out_ts = make_data(dom_size, n_domains/6, 
-                                                    max_obs, max_obs_size, n_traj, state_batch_size)
-
-    np.savez_compressed(save_path, X_out_tr, S1_out_tr, S2_out_tr, Labels_out_tr, 
-                        X_out_ts, S1_out_ts, S2_out_ts, Labels_out_ts)
+    X_out_ts, S1_out_ts, S2_out_ts, Labels_out_ts = make_data(
+        dom_size, n_domains/6, max_obs, max_obs_size, n_traj, state_batch_size)
+    # Save dataset
+    np.savez_compressed(save_path, X_out_tr, S1_out_tr, S2_out_tr, 
+        Labels_out_tr, X_out_ts, S1_out_ts, S2_out_ts, Labels_out_ts)
 
 
 if __name__ == '__main__':

domains/gridworld.py

@@ -2,6 +2,7 @@
 from scipy.sparse import csr_matrix
 from scipy.sparse.csgraph import dijkstra
 
+
 class gridworld:
     """A class for making gridworlds"""
     def __init__(self, image, targetx, targety):
@@ -25,6 +26,8 @@ def __init__(self, image, targetx, targety):
 
 
     def set_vals(self):
+        # Setup function to initialize all necessary
+        #  data
         row_obs, col_obs = np.where(self.image == 0)
         row_free, col_free = np.where(self.image != 0)
         self.obstacles = [row_obs, col_obs]
@@ -125,24 +128,28 @@ def set_vals(self):
 
 
     def get_graph(self):
+        # Returns graph
         G = self.G
         W = self.W[self.W != 0]
         return G, W
 
 
     def get_graph_inv(self):
+        # Returns transpose of graph
         G = self.G.T
         W = self.W.T
         return G, W
 
 
     def val_2_image(self, val):
+        # Zeros for obstacles, val for free space
         im = np.zeros((self.n_row, self.n_col))
         im[self.freespace[0], self.freespace[1]] = val
         return im
 
 
     def get_value_prior(self):
+        # Returns value prior for gridworld
         s_map_col, s_map_row = np.meshgrid(np.arange(0,self.n_col), 
             np.arange(0, self.n_row))
         im = np.sqrt(np.square(s_map_col - self.targety) 
@@ -151,30 +158,37 @@ def get_value_prior(self):
 
 
     def get_reward_prior(self):
+        # Returns reward prior for gridworld
         im = -1 * np.ones((self.n_row, self.n_col))
         im[self.targetx, self.targety] = 10
         return im
 
 
     def t_get_reward_prior(self):
+        # Returns reward prior as needed for
+        #  dataset generation
         im = np.zeros((self.n_row, self.n_col))
         im[self.targetx, self.targety] = 10
         return im
 
 
     def get_state_image(self, row, col):
+        # Zeros everywhere except [row,col]
         im = np.zeros((self.n_row, self.n_col))
         im[row, col] = 1
         return im
 
 
     def map_ind_to_state(self, row, col):
+        # Takes [row, col] and maps to a state
         rw = np.where(self.state_map_row == row)
         cl = np.where(self.state_map_col == col)
         return np.intersect1d(rw, cl)[0]
 
 
     def get_coords(self, states):
+        # Given a state or states, returns
+        #  [row,col] pairs for the state(s)
         non_obstacles = np.ravel_multi_index(
             [self.freespace[0], self.freespace[1]], 
             (self.n_row,self.n_col), order='F')
@@ -186,6 +200,7 @@ def get_coords(self, states):
 
 
     def rand_choose(self, in_vec):
+        # Samples 
         if len(in_vec.shape) > 1:
             if in_vec.shape[1] == 1:
                 in_vec = in_vec.T
@@ -197,6 +212,8 @@ def rand_choose(self, in_vec):
 
 
     def next_state_prob(self, s, a):
+        # Gets next state probability for
+        #  a given action (a)
         if hasattr(a, "__iter__"):
             p = np.squeeze(self.P[s, :, a])
         else:
@@ -205,16 +222,22 @@ def next_state_prob(self, s, a):
 
 
     def sample_next_state(self, s, a):
+        # Gets the next state given the
+        #  current state (s) and an 
+        #  action (a)
         vec = self.next_state_prob(s, a)
         result = self.rand_choose(vec)
         return result
 
 
     def get_size(self):
+        # Returns domain size
         return self.n_row, self.n_col
 
 
     def north(self, row, col):
+        # Returns new [row,col]
+        #  if we take the action
         new_row = np.max([row-1, 0])
         new_col = col
         if self.image[new_row, new_col] == 0:
@@ -224,6 +247,8 @@ def north(self, row, col):
 
 
     def northeast(self, row, col):
+        # Returns new [row,col]
+        #  if we take the action
         new_row = np.max([row - 1, 0])
         new_col = np.min([col + 1, self.n_col - 1])
         if self.image[new_row, new_col] == 0:
@@ -233,6 +258,8 @@ def northeast(self, row, col):
 
 
     def northwest(self, row, col):
+        # Returns new [row,col]
+        #  if we take the action
         new_row = np.max([row - 1, 0])
         new_col = np.max([col - 1, 0])
         if self.image[new_row, new_col] == 0:
@@ -242,6 +269,8 @@ def northwest(self, row, col):
 
 
     def south(self, row, col):
+        # Returns new [row,col]
+        #  if we take the action
         new_row = np.min([row + 1, self.n_row - 1])
         new_col = col
         if self.image[new_row, new_col] == 0:
@@ -251,6 +280,8 @@ def south(self, row, col):
 
 
     def southeast(self, row, col):
+        # Returns new [row,col]
+        #  if we take the action
         new_row = np.min([row + 1, self.n_row - 1])
         new_col = np.min([col + 1, self.n_col - 1])
         if self.image[new_row, new_col] == 0:
@@ -260,6 +291,8 @@ def southeast(self, row, col):
 
 
     def southwest(self, row, col):
+        # Returns new [row,col]
+        #  if we take the action
         new_row = np.min([row + 1, self.n_row - 1])
         new_col = np.max([col - 1, 0])
         if self.image[new_row, new_col] == 0:
@@ -269,6 +302,8 @@ def southwest(self, row, col):
 
 
     def east(self, row, col):
+        # Returns new [row,col]
+        #  if we take the action
         new_row = row
         new_col = np.min([col + 1, self.n_col - 1])
         if self.image[new_row, new_col] == 0:
@@ -278,6 +313,8 @@ def east(self, row, col):
 
 
     def west(self, row, col):
+        # Returns new [row,col]
+        #  if we take the action
         new_row = row
         new_col = np.max([col - 1, 0])
         if self.image[new_row, new_col] == 0:
@@ -307,6 +344,9 @@ def neighbors(self, row, col):
 
 
 def trace_path(pred, source, target):
+    # traces back shortest path from
+    #  source to target given pred
+    #  (a predicessor list)
     max_len = 1000
     path = np.zeros((max_len, 1))
     i = max_len - 1
@@ -325,6 +365,8 @@ def trace_path(pred, source, target):
 
 
 def sample_trajectory(M, n_states):
+    # Samples trajectories from random nodes
+    #  in our domain (M)
     G, W = M.get_graph_inv()
     N = G.shape[0]
     if N >= n_states: