finish example for static obstacle

examples/sipp_planner.ipynb

@@ -414,7 +414,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3.8.13 ('gnnmp')",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },

planner/learned/GNN_static_planner.py

@@ -3,6 +3,8 @@
 from utils.utils import seed_everything, create_dot_dict, to_np
 from utils.graphs import knn_graph_from_points
 
+from planner.learned.model.GNN_static import GNNet
+
 from torch_sparse import coalesce
 from torch_geometric.nn import knn_graph
 from torch_geometric.data import Data
@@ -12,16 +14,15 @@
 
 
 class GNNStaticPlanner(LearnedPlanner):
-    def __init__(self, num_batch, model, k_neighbors=50, **kwargs):
+    def __init__(self, num_batch, model_args, k_neighbors=50, **kwargs):
         self.num_batch = num_batch
-        self.model = model
+        self.model = GNNet(**model_args)
         self.k_neigbors = k_neighbors
         self.num_node = 0
 
         super(GNNStaticPlanner, self).__init__(self.model, **kwargs)
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        for model_ in self.model:
-            model_.to(self.device)
+        self.model.to(self.device)
 
     def _num_node(self):
         return self.num_node
@@ -32,18 +33,19 @@ def _plan(self, env, start, goal, timeout, seed=0, **kwargs):
         self.model.eval()
         path = self._explore(env, start, goal, self.model, timeout, k=self.k_neigbors, n_sample=self.num_batch)
 
-        return create_dot_dict(solution=path)
+        return create_dot_dict(solution=path if len(path) else None)
 
     def create_graph(self):
         graph_data = knn_graph_from_points(self.points, self.k_neighbors)
         self.edges = graph_data.edges
         self.edge_index = graph_data.edge_index
         self.edge_cost = graph_data.edge_cost
 
-    def create_data(self, points, edge_index=None, k=50):
+    def create_data(self, points, obstacles, edge_index=None, k=50):
         goal_index = 1
         data = Data(goal=torch.FloatTensor(points[goal_index]))
-        data.v = torch.FloatTensor(points)
+        data.v = torch.FloatTensor(np.array(points))
+        data.obstacles = torch.FloatTensor()
 
         if edge_index is not None:
             data.edge_index = torch.tensor(edge_index.T).to(self.device)
@@ -73,7 +75,7 @@ def _explore(self, env, start, goal, model_gnn, timeout, k, n_sample, loop=10):
 
         while not success:
 
-            data = self.create_data(points, k)
+            data = self.create_data(points, env.get_obstacles(), k=k)
             self.num_node = len(data.v)
             policy = model_gnn(**data.to(self.device).to_dict(), loop=loop)
             policy = policy.cpu()
@@ -89,12 +91,12 @@ def _explore(self, env, start, goal, model_gnn, timeout, k, n_sample, loop=10):
                 end_a, end_b = int(end_a), int(end_b)
                 end_a = explored[end_a]
                 explored_edges.extend([[end_a, end_b], [end_b, end_a]])
-                if env._edge_fp(to_np(data.v[end_a]), to_np(data.v[end_b])):
+                if env.edge_fp(to_np(data.v[end_a]), to_np(data.v[end_b])):
                     explored.append(end_b)
                     prev[end_b] = end_a
 
                     policy[:, end_b] = 0
-                    if env.in_goal_region(to_np(data.v[end_b]), to_np(data.v[1])):
+                    if end_b==1:
                         success = True
                         path = [end_b]
                         node = end_b
@@ -114,7 +116,7 @@ def _explore(self, env, start, goal, model_gnn, timeout, k, n_sample, loop=10):
                 new_points = env.sample_n_points(n_sample, need_negative=True)
                 points = points + list(new_points)
 
-        return create_dot_dict(solution=list(data.v[path].data.cpu().numpy()) if len(path) else None)
+        return list(data.v[path].data.cpu().numpy())
 
 
 

planner/learned/model/GNN_dynamic.py

@@ -2,7 +2,7 @@
 from torch.nn import Sequential as Seq, Linear as Lin, ReLU, Sigmoid
 from torch_geometric.nn import MessagePassing
 
-from base_models import Block
+from planner.learned.model.base_models import Block
 
 
 class MPNN(MessagePassing):

planner/learned/model/GNN_static.py

@@ -2,7 +2,7 @@
 from torch.nn import Sequential as Seq, Linear as Lin, ReLU, Sigmoid
 from torch_geometric.nn import MessagePassing
 
-from base_models import Block
+from planner.learned.model.base_models import Block
 
 
 class MPNN(MessagePassing):
@@ -73,8 +73,8 @@ def forward(self, v, labels, obstacles, edge_index, loop, **kwargs):
             obs_edge_code = self.obs_edge_code(obstacles.view(-1, self.obs_size))
 
             for na, ea in zip(self.node_attentions, self.edge_attentions):
-                x, obs_node_code = na(x, obs_node_code)
-                y, obs_edge_code = ea(y, obs_edge_code)
+                x = na(x, obs_node_code)
+                y = ea(y, obs_edge_code)
 
         for i in range(loop):
             x = self.mpnn(x, edge_index, y)

wrappers/obstacles.py

@@ -0,0 +1,14 @@
+import numpy as np
+
+class ObstaclePositionWrapper():
+    '''
+    return representation of obstacles as concatenated vector of positions
+    '''
+    def __init__(self, baseObject):
+        self.__class__ = type(baseObject.__class__.__name__,
+                              (self.__class__, baseObject.__class__),
+                              {})
+        self.__dict__ = baseObject.__dict__
+
+    def get_obstacles(self):
+        return np.array([list(obj.base_position)+list(obj.half_extents) for obj in self.objects])