import torch
import torch.nn as nn
[docs]def env_init_embedding(env_name: str, config: dict) -> nn.Module:
"""Get environment initial embedding. The init embedding is used to initialize the
general embedding of the problem nodes without any solution information.
Consists of a linear layer that projects the node features to the embedding space.
Args:
env: Environment or its name.
config: A dictionary of configuration options for the environment.
"""
embedding_registry = {
"tsp": TSPInitEmbedding,
"atsp": TSPInitEmbedding,
"cvrp": VRPInitEmbedding,
"sdvrp": VRPInitEmbedding,
"pctsp": PCTSPInitEmbedding,
"spctsp": PCTSPInitEmbedding,
"op": OPInitEmbedding,
"dpp": DPPInitEmbedding,
"mdpp": MDPPInitEmbedding,
"pdp": PDPInitEmbedding,
"mtsp": MTSPInitEmbedding,
}
if env_name not in embedding_registry:
raise ValueError(
f"Unknown environment name '{env_name}'. Available init embeddings: {embedding_registry.keys()}"
)
return embedding_registry[env_name](**config)
[docs]class TSPInitEmbedding(nn.Module):
"""Initial embedding for the Traveling Salesman Problems (TSP).
Embed the following node features to the embedding space:
- locs: x, y coordinates of the cities
"""
def __init__(self, embedding_dim, linear_bias=True):
super(TSPInitEmbedding, self).__init__()
node_dim = 2 # x, y
self.init_embed = nn.Linear(node_dim, embedding_dim, linear_bias)
[docs] def forward(self, td):
out = self.init_embed(td["locs"])
return out
[docs]class VRPInitEmbedding(nn.Module):
"""Initial embedding for the Vehicle Routing Problems (VRP).
Embed the following node features to the embedding space:
- locs: x, y coordinates of the nodes (depot and customers separately)
- demand: demand of the customers
"""
def __init__(self, embedding_dim, linear_bias=True):
super(VRPInitEmbedding, self).__init__()
node_dim = 3 # x, y, demand
self.init_embed = nn.Linear(node_dim, embedding_dim, linear_bias)
self.init_embed_depot = nn.Linear(
2, embedding_dim, linear_bias
) # depot embedding
[docs] def forward(self, td):
# [batch, 1, 2]-> [batch, 1, embedding_dim]
depot, cities = td["locs"][:, :1, :], td["locs"][:, 1:, :]
depot_embedding = self.init_embed_depot(depot)
# [batch, n_city, 2, batch, n_city, 1] -> [batch, n_city, embedding_dim]
node_embeddings = self.init_embed(
torch.cat((cities, td["demand"][..., None]), -1)
)
# [batch, n_city+1, embedding_dim]
out = torch.cat((depot_embedding, node_embeddings), -2)
return out
[docs]class PCTSPInitEmbedding(nn.Module):
"""Initial embedding for the Prize Collecting Traveling Salesman Problems (PCTSP).
Embed the following node features to the embedding space:
- locs: x, y coordinates of the nodes (depot and customers separately)
- expected_prize: expected prize for visiting the customers.
In PCTSP, this is the actual prize. In SPCTSP, this is the expected prize.
- penalty: penalty for not visiting the customers
"""
def __init__(self, embedding_dim, linear_bias=True):
super(PCTSPInitEmbedding, self).__init__()
node_dim = 4 # x, y, prize, penalty
self.init_embed = nn.Linear(node_dim, embedding_dim, linear_bias)
self.init_embed_depot = nn.Linear(2, embedding_dim, linear_bias)
[docs] def forward(self, td):
depot, cities = td["locs"][:, :1, :], td["locs"][:, 1:, :]
depot_embedding = self.init_embed_depot(depot)
node_embeddings = self.init_embed(
torch.cat(
(
cities,
td["expected_prize"][..., None],
td["penalty"][..., 1:, None],
),
-1,
)
)
# batch, n_city+1, embedding_dim
out = torch.cat((depot_embedding, node_embeddings), -2)
return out
[docs]class OPInitEmbedding(nn.Module):
"""Initial embedding for the Orienteering Problems (OP).
Embed the following node features to the embedding space:
- locs: x, y coordinates of the nodes (depot and customers separately)
- prize: prize for visiting the customers
"""
def __init__(self, embedding_dim, linear_bias=True):
super(OPInitEmbedding, self).__init__()
node_dim = 3 # x, y, prize
self.init_embed = nn.Linear(node_dim, embedding_dim, linear_bias)
self.init_embed_depot = nn.Linear(
2, embedding_dim, linear_bias
) # depot embedding
[docs] def forward(self, td):
depot, cities = td["locs"][:, :1, :], td["locs"][:, 1:, :]
depot_embedding = self.init_embed_depot(depot)
node_embeddings = self.init_embed(
torch.cat(
(
cities,
td["prize"][..., 1:, None], # exclude depot
),
-1,
)
)
out = torch.cat((depot_embedding, node_embeddings), -2)
return out
[docs]class DPPInitEmbedding(nn.Module):
"""Initial embedding for the Decap Placement Problem (DPP), EDA (electronic design automation).
Embed the following node features to the embedding space:
- locs: x, y coordinates of the nodes (cells)
- probe: index of the (single) probe cell. We embed the euclidean distance from the probe to all cells.
"""
def __init__(self, embedding_dim, linear_bias=True):
super(DPPInitEmbedding, self).__init__()
node_dim = 2 # x, y
self.init_embed = nn.Linear(node_dim, embedding_dim // 2, linear_bias) # locs
self.init_embed_probe = nn.Linear(1, embedding_dim // 2, linear_bias) # probe
[docs] def forward(self, td):
node_embeddings = self.init_embed(td["locs"])
probe_embedding = self.init_embed_probe(
self._distance_probe(td["locs"], td["probe"])
)
return torch.cat([node_embeddings, probe_embedding], -1)
def _distance_probe(self, locs, probe):
# Euclidean distance from probe to all locations
probe_loc = torch.gather(locs, 1, probe.unsqueeze(-1).expand(-1, -1, 2))
return torch.norm(locs - probe_loc, dim=-1).unsqueeze(-1)
[docs]class MDPPInitEmbedding(nn.Module):
"""Initial embedding for the Multi-port Placement Problem (MDPP), EDA (electronic design automation).
Embed the following node features to the embedding space:
- locs: x, y coordinates of the nodes (cells)
- probe: indexes of the probe cells (multiple). We embed the euclidean distance of each cell to the closest probe.
"""
def __init__(self, embedding_dim, linear_bias=True):
super(MDPPInitEmbedding, self).__init__()
node_dim = 2 # x, y
self.init_embed = nn.Linear(node_dim, embedding_dim, linear_bias) # locs
self.init_embed_probe_distance = nn.Linear(
1, embedding_dim, linear_bias
) # probe_distance
self.project_out = nn.Linear(embedding_dim * 2, embedding_dim, linear_bias)
[docs] def forward(self, td):
probes = td["probe"]
locs = td["locs"]
node_embeddings = self.init_embed(locs)
# Get the shortest distance from any probe
dist = torch.cdist(locs, locs, p=2)
dist[~probes] = float("inf")
min_dist, _ = torch.min(dist, dim=1)
min_probe_dist_embedding = self.init_embed_probe_distance(min_dist[..., None])
return self.project_out(
torch.cat([node_embeddings, min_probe_dist_embedding], -1)
)
[docs]class PDPInitEmbedding(nn.Module):
"""Initial embedding for the Pickup and Delivery Problem (PDP).
Embed the following node features to the embedding space:
- locs: x, y coordinates of the nodes (depot, pickups and deliveries separately)
Note that pickups and deliveries are interleaved in the input.
"""
def __init__(self, embedding_dim, linear_bias=True):
super(PDPInitEmbedding, self).__init__()
node_dim = 2 # x, y
self.init_embed_depot = nn.Linear(2, embedding_dim, linear_bias)
self.init_embed_pick = nn.Linear(node_dim * 2, embedding_dim, linear_bias)
self.init_embed_delivery = nn.Linear(node_dim, embedding_dim, linear_bias)
[docs] def forward(self, td):
depot, locs = td["locs"][..., 0:1, :], td["locs"][..., 1:, :]
num_locs = locs.size(-2)
pick_feats = torch.cat(
[locs[:, : num_locs // 2, :], locs[:, num_locs // 2 :, :]], -1
) # [batch_size, graph_size//2, 4]
delivery_feats = locs[:, num_locs // 2 :, :] # [batch_size, graph_size//2, 2]
depot_embeddings = self.init_embed_depot(depot)
pick_embeddings = self.init_embed_pick(pick_feats)
delivery_embeddings = self.init_embed_delivery(delivery_feats)
# concatenate on graph size dimension
return torch.cat([depot_embeddings, pick_embeddings, delivery_embeddings], -2)
[docs]class MTSPInitEmbedding(nn.Module):
"""Initial embedding for the Multiple Traveling Salesman Problem (mTSP).
Embed the following node features to the embedding space:
- locs: x, y coordinates of the nodes (depot, cities)
"""
def __init__(self, embedding_dim, linear_bias=True):
"""NOTE: new made by Fede. May need to be checked"""
super(MTSPInitEmbedding, self).__init__()
node_dim = 2 # x, y
self.init_embed = nn.Linear(node_dim, embedding_dim, linear_bias)
self.init_embed_depot = nn.Linear(
2, embedding_dim, linear_bias
) # depot embedding
[docs] def forward(self, td):
depot_embedding = self.init_embed_depot(td["locs"][..., 0:1, :])
node_embedding = self.init_embed(td["locs"][..., 1:, :])
return torch.cat([depot_embedding, node_embedding], -2)
