fix(GLOP): add VRPPolar init/edge embeddings (!not usable!)

Author: Furffico

rl4co/models/nn/env_embeddings/edge.py

@@ -99,11 +99,72 @@ def _cost_matrix_to_graph(self, batch_cost_matrix: Tensor, init_embeddings: Tens
                 x=init_embeddings[index],
                 edge_index=edge_index,
                 edge_attr=edge_attr,
+            )  # type: ignore
+            graph_data.append(graph)
+
+        batch = Batch.from_data_list(graph_data)  # type: ignore
+        batch.edge_attr = self.edge_embed(batch.edge_attr)  # type: ignore
+        return batch
+
+
+class VRPPolarEdgeEmbedding(TSPEdgeEmbedding):
+    """TODO"""
+
+    node_dim = 2
+
+    def forward(self, td, init_embeddings: Tensor):
+        with torch.no_grad():
+            if "polar_locs" in td.keys():
+                theta = td["polar_locs"][..., 1]
+            else:
+                shifted_locs = td["locs"] - td["locs"][..., 0:1, :]
+                x, y = shifted_locs[..., 0], shifted_locs[..., 1]
+                theta = torch.atan2(y, x)
+
+            delta_theta_matrix = theta[..., :, None] - theta[..., None, :]
+            edge_attr1 = 1 - torch.cos(delta_theta_matrix)
+            edge_attr2 = get_distance_matrix(td["locs"])
+            cost_matrix = torch.stack((edge_attr1, edge_attr2), dim=-1)
+            del edge_attr1, edge_attr2, delta_theta_matrix
+
+            batch = self._cost_matrix_to_graph(cost_matrix, init_embeddings)
+            del cost_matrix
+
+        batch.edge_attr = self.edge_embed(batch.edge_attr)  # type: ignore
+        return batch
+
+    @torch.no_grad()
+    def _cost_matrix_to_graph(self, batch_cost_matrix: Tensor, init_embeddings: Tensor):
+        """Convert batched cost_matrix to batched PyG graph, and calculate edge embeddings.
+
+        Args:
+            batch_cost_matrix: Tensor of shape [batch_size, n, n, m]
+            init_embedding: init embeddings of shape [batch_size, n, m]
+        """
+        graph_data = []
+        for index, cost_matrix in enumerate(batch_cost_matrix):
+            edge_index, _ = sparsify_graph(
+                cost_matrix[..., 0], self.k_sparse, self_loop=False
+            )
+            edge_index = edge_index.T[torch.all(edge_index != 0, dim=0)].T
+            _, depot_edge_index = torch.topk(
+                cost_matrix[0, :, 1], k=self.k_sparse, largest=False, sorted=False
             )
+            depot_edge_index = depot_edge_index[depot_edge_index != 0]
+            depot_edge_index = torch.stack(
+                (torch.zeros_like(depot_edge_index), depot_edge_index), dim=0
+            )
+            edge_index = torch.concat((depot_edge_index, edge_index), dim=-1).detach()
+            edge_attr = cost_matrix[edge_index[0], edge_index[1]].detach()
+
+            graph = Data(
+                x=init_embeddings[index],
+                edge_index=edge_index,
+                edge_attr=edge_attr,
+            )  # type: ignore
             graph_data.append(graph)
 
-        batch = Batch.from_data_list(graph_data)
-        batch.edge_attr = self.edge_embed(batch.edge_attr)
+        batch = Batch.from_data_list(graph_data)  # type: ignore
         return batch
 
 
@@ -146,8 +207,8 @@ def forward(self, td, init_embeddings: Tensor):
                 x=node_embed,
                 edge_index=edge_index,
                 edge_attr=torch.zeros((m, self.embed_dim), device=device),
-            )
+            )  # type: ignore
             data_list.append(data)
 
-        batch = Batch.from_data_list(data_list)
+        batch = Batch.from_data_list(data_list)  # type: ignore
         return batch

rl4co/models/nn/env_embeddings/init.py

@@ -4,6 +4,7 @@
 from tensordict.tensordict import TensorDict
 
 from rl4co.models.nn.ops import PositionalEncoding
+from rl4co.utils.ops import cartesian_to_polar
 
 
 def env_init_embedding(env_name: str, config: dict) -> nn.Module:
@@ -152,6 +153,55 @@ def forward(self, td):
         return torch.cat((depot_embedding, node_embeddings), -2)
 
 
+class VRPPolarInitEmbedding(nn.Module):
+    """Initial embedding for the Vehicle Routing Problems (VRP).
+    Embed the following node features to the embedding space, based on polar coordinates:
+        - locs: r, theta coordinates of the nodes, with the depot as the origin
+        - demand: demand of the customers
+    """
+
+    def __init__(
+        self,
+        embed_dim,
+        linear_bias=True,
+        node_dim: int = 3,
+        attach_cartesian_coords=False,
+    ):
+        super(VRPPolarInitEmbedding, self).__init__()
+        self.node_dim = node_dim + (
+            2 if attach_cartesian_coords else 0
+        )  # 3: r, theta, demand; 5: r, theta, demand, x, y;
+        self.attach_cartesian_coords = attach_cartesian_coords
+        self.init_embed = nn.Linear(self.node_dim, embed_dim, linear_bias)
+        self.init_embed_depot = nn.Linear(
+            self.node_dim, embed_dim, linear_bias
+        )  # depot embedding
+
+    def forward(self, td):
+        with torch.no_grad():
+            locs = td["locs"]
+            polar_locs = cartesian_to_polar(locs, locs[..., 0:1, :])
+            td["polar_locs"] = polar_locs
+
+            demand = td["demand"]
+            demand_with_depot = torch.concat(
+                (torch.zeros(demand.shape[0], 1, device=demand.device), demand),
+                dim=-1,
+            ).unsqueeze(-1)
+
+            if self.attach_cartesian_coords:
+                x = torch.concat((polar_locs, demand_with_depot, locs), dim=-1)
+            else:
+                x = torch.concat((polar_locs, demand_with_depot), dim=-1)
+
+        depot, cities = x[:, :1, :], x[:, 1:, :]
+        depot_embedding = self.init_embed_depot(depot)
+        node_embeddings = self.init_embed(cities)
+
+        out = torch.cat((depot_embedding, node_embeddings), -2)
+        return out
+
+
 class SVRPInitEmbedding(nn.Module):
     def __init__(self, embed_dim, linear_bias=True, node_dim: int = 3):
         super(SVRPInitEmbedding, self).__init__()