Update MTVRP

rl4co/envs/routing/mtvrp/env.py

@@ -281,7 +281,7 @@ def get_action_mask(td: TensorDict) -> torch.Tensor:
  & ~exceeds_dist_limit
  & ~td["visited"]
  )
-
+ #print(can_visit)
  # Mask depot: don't visit depot if coming from there and there are still customer nodes I can visit
  can_visit[:, 0] = ~((curr_node == 0) & (can_visit[:, 1:].sum(-1) > 0))
  return can_visit
@@ -349,9 +349,14 @@ def check_solution_validity(td: TensorDict, actions: torch.Tensor):
  curr_time = torch.max(
  curr_time + dist, gather_by_index(td["time_windows"], next_node)[..., 0]
  )
+
+ new_shape = curr_time.size()
+ skip_open_end = td["open_route"].view(*new_shape) & (next_node == 0).view(*new_shape)
+
  assert torch.all(
- curr_time <= gather_by_index(td["time_windows"], next_node)[..., 1]
+ (curr_time <= gather_by_index(td["time_windows"], next_node)[..., 1]) | skip_open_end
  ), "vehicle cannot start service before deadline"
+
  curr_time = curr_time + gather_by_index(td["service_time"], next_node)
  curr_node = next_node
  curr_time[curr_node == 0] = 0.0 # reset time for depot
@@ -450,7 +455,7 @@ def _make_spec(self, td_params: TensorDict):
  def check_variants(td):
  """Check if the problem has the variants"""
  has_open = td["open_route"].squeeze(-1)
- has_tw = (td["time_windows"][:, :, 1] != float("inf")).any(-1)
+ has_tw = (td["time_windows"][:, :, 1] != 4.6).any(-1)
  has_limit = (td["distance_limit"] != float("inf")).squeeze(-1)
  has_backhaul = (td["demand_backhaul"] != 0).any(-1)
  return has_open, has_tw, has_limit, has_backhaul

rl4co/envs/routing/mtvrp/generator.py

@@ -256,7 +256,8 @@ def _default_open(td, remove):
  @staticmethod
  def _default_time_window(td, remove):
  default_tw = torch.zeros_like(td["time_windows"])
- default_tw[..., 1] = float("inf")
+ #default_tw[..., 1] = float("inf")
+ default_tw[..., 1] = 4.6 # max tw
  td["time_windows"][remove] = default_tw[remove]
  td["service_time"][remove] = torch.zeros_like(td["service_time"][remove])
  return td

rl4co/models/nn/env_embeddings/context.py

@@ -32,6 +32,7 @@ def env_context_embedding(env_name: str, config: dict) -> nn.Module:
  "mtsp": MTSPContext,
  "smtwtp": SMTWTPContext,
  "mdcpdp": MDCPDPContext,
+ "mtvrp": MTVRPContext
  }
 
  if env_name not in embedding_registry:
@@ -146,6 +147,50 @@ def _state_embedding(self, embeddings, td):
  state_embedding = td["vehicle_capacity"] - td["used_capacity"]
  return state_embedding
 
+class VRPBContext(EnvContext):
+ """Context embedding for the Capacitated Vehicle Routing Problem (CVRP).
+ Project the following to the embedding space:
+ - current node embedding
+ - remaining capacity (vehicle_capacity - used_capacity)
+ """
+
+ def __init__(self, embed_dim):
+ super(VRPContext, self).__init__(
+ embed_dim=embed_dim, step_context_dim=embed_dim + 1
+ )
+
+ def _state_embedding(self, embeddings, td):
+ mask = (td["used_capacity_backhaul"] == 0)
+ used_capacity = torch.where(mask, td["used_capacity_linehaul"], td["used_capacity_backhaul"])
+ state_embedding = td["vehicle_capacity"] - used_capacity
+ return state_embedding
+
+class MTVRPContext(VRPBContext):
+ """Context embedding for the Capacitated Vehicle Routing Problem (CVRP).
+ Project the following to the embedding space:
+ - current node embedding
+ - remaining capacity (vehicle_capacity - used_capacity)
+ - current time
+ - current route length
+ - if route should be open
+ """
+
+ def __init__(self, embed_dim):
+ super(VRPBContext, self).__init__(
+ embed_dim=embed_dim, step_context_dim=embed_dim + 4
+ )
+
+ def _state_embedding(self, embeddings, td):
+
+ capacity = super()._state_embedding(embeddings, td)
+ current_time = td["current_time"]
+ current_length = td["current_route_length"]
+ is_open = td["open_route"]
+ is_open_tensor = torch.zeros_like(is_open, dtype=torch.float)
+ is_open_tensor[is_open] = 1
+
+ return torch.cat([capacity, current_time, current_length, is_open_tensor], -1)
+
 
 class VRPTWContext(VRPContext):
  """Context embedding for the Capacitated Vehicle Routing Problem (CVRP).

rl4co/models/nn/env_embeddings/init.py

@@ -2,10 +2,8 @@
 import torch.nn as nn
 
 from tensordict.tensordict import TensorDict
-
 from rl4co.models.nn.ops import PositionalEncoding
 
-
 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.
@@ -33,6 +31,7 @@ def env_init_embedding(env_name: str, config: dict) -> nn.Module:
  "smtwtp": SMTWTPInitEmbedding,
  "mdcpdp": MDCPDPInitEmbedding,
  "fjsp": FJSPFeatureEmbedding,
+ "mtvrp":MTVRPInitEmbedding,
  }
 
  if env_name not in embedding_registry:
@@ -146,6 +145,28 @@ def forward(self, td):
  )
  )
  return torch.cat((depot_embedding, node_embeddings), -2)
+
+
+class MTVRPInitEmbedding(VRPInitEmbedding):
+ def __init__(self, embed_dim, linear_bias=True, node_dim: int = 5):
+ # node_dim = 5: x, y, demand, tw start, tw end
+ super(MTVRPInitEmbedding, self).__init__(embed_dim, linear_bias, node_dim)
+
+ def forward(self, td):
+ depot, cities = td["locs"][:, :1, :], td["locs"][:, 1:, :]
+ #durations = td["durations"][..., 1:]
+ time_windows = td["time_windows"][..., 1:, :]
+ # embeddings
+ demands = td["demand_linehaul"][..., None] - td["demand_backhaul"][..., None]
+
+ depot_embedding = self.init_embed_depot(depot)
+ node_embeddings = self.init_embed(
+ torch.cat(
+ (cities, demands[:,1:], time_windows), -1
+ )
+ )
+
+ return torch.cat((depot_embedding, node_embeddings), -2)
 
 
 class SVRPInitEmbedding(nn.Module):
@@ -383,7 +404,6 @@ def forward(self, td):
  # concatenate on graph size dimension
  return torch.cat([depot_embeddings, pick_embeddings, delivery_embeddings], -2)
 
-
 class FJSPFeatureEmbedding(nn.Module):
  def __init__(self, embed_dim, linear_bias=True, norm_coef: int = 100):
  super().__init__()
@@ -443,4 +463,4 @@ def _stepwise_operations_embed(self, td: TensorDict):
  raise NotImplementedError("Stepwise encoding not yet implemented")
 
  def _stepwise_machine_embed(self, td: TensorDict):
- raise NotImplementedError("Stepwise encoding not yet implemented")
+ raise NotImplementedError("Stepwise encoding not yet implemented")