Add status-aware now/future decoder

100/config.yaml

@@ -8,7 +8,10 @@ model_params:
    'ff_hidden_dim': 512 
    'ffd': 'siglu' 
    'norm_type': 'rms' 
-   'use_sparse': 'topk'   # relu topk entmax15 sparsemax a-entmax False 
+   'use_sparse': 'topk'   # relu topk entmax15 sparsemax a-entmax False
+   'use_now_future': false
+   'future_dim': 128
+   'future_heads': 8
 
 
 

100/envs/env.py

@@ -142,14 +142,16 @@ def reset(self, td: Optional[TensorDict] = None, batch_size=None, lib_data=False
             td = self.generator(batch_size=batch_size).to('cuda')
         batch_size = [batch_size] if isinstance(batch_size, int) else batch_size
         self.to(td.device)
-        return super().reset(td, batch_size=batch_size,lib_data=lib_data)
+        # The rollout batch is dynamic and ``step`` already bypasses EnvBase.
+        # TorchRL 0.6 otherwise rejects it against the empty EnvBase batch.
+        return self._reset(td, batch_size=batch_size, lib_data=lib_data)
 
     def _reset(self, td: Optional[TensorDict] = None, batch_size: Optional[list] = None, lib_data=False) -> TensorDict:
-        device = td.device
-        if lib_data == False:
-            # Create reset TensorDict
-            td_reset = TensorDict(
-                {
+        device = td.device
+        if lib_data == False:
+            # Create reset TensorDict
+            td_reset = TensorDict(
+                {
                     "locs": td["locs"],
                     "demand_backhaul": td["demand_backhaul"],
                     "demand_linehaul": td["demand_linehaul"],
@@ -166,12 +168,14 @@ def _reset(self, td: Optional[TensorDict] = None, batch_size: Optional[list] = N
                     "used_capacity_backhaul": torch.zeros((*batch_size, 1), device=device),  # for capacity constraints in backhaul
                     "used_capacity_linehaul": torch.zeros((*batch_size, 1), device=device),  # for capacity constraints in linehaul
                     "visited": torch.zeros((*batch_size, td["locs"].shape[-2]), dtype=torch.bool, device=device,),
-                },
-                batch_size=batch_size,
-                device=device,
-            )
-            td_reset.set("action_mask", self.get_action_mask(td_reset))
-        else:
+                },
+                batch_size=batch_size,
+                device=device,
+            )
+            if "p_s_tag" in td.keys():
+                td_reset.set("p_s_tag", td["p_s_tag"])
+            td_reset.set("action_mask", self.get_action_mask(td_reset))
+        else:
             # Demands: linehaul (C) and backhaul (B). Backhaul defaults to 0
             demand_linehaul = torch.cat(
                 [torch.zeros_like(td["demand_linehaul"][..., :1]), td["demand_linehaul"]],
@@ -642,4 +646,4 @@ def __setstate__(self, state, set_seed:bool=True):
         self.__dict__.update(state)
         if set_seed:
             self.rng = torch.manual_seed(0) # = torch.default_generator
-            self.rng.set_state(state["rng"].cpu()) # = set  torch.default_generator
+            self.rng.set_state(state["rng"].cpu()) # = set  torch.default_generator

100/model.py

@@ -21,11 +21,13 @@ def linear_layer(input_dim, output_dim, mean=0., std=1e-2, bias=True, bias_init_
     return linear
 
 @dataclass
-class PrecomputedCache:
+class PrecomputedCache:
     node_embeddings: Tensor
     glimpse_key: Tensor
-    glimpse_val: Tensor
-    logit_key: Tensor
+    glimpse_val: Tensor
+    logit_key: Tensor
+    future_key: Tensor = None
+    future_val: Tensor = None
 
     @property
     def fields(self):
@@ -112,7 +114,15 @@ def forward(self, td, env):
         decoder_k = reshape_by_heads(self.decoder.Wk(node_embed), head_num=args.model_params['head_num'])
         decoder_v = reshape_by_heads(self.decoder.Wv(node_embed), head_num=args.model_params['head_num']) # (batch, head_num, problem+1, qkv_dim)v
         decoder_single_head_k = node_embed.transpose(1, 2) # (batch, embedding, problem+1)
-        cache = PrecomputedCache(node_embed, decoder_k, decoder_v, decoder_single_head_k)
+        future_k, future_v = self.decoder.precompute_future(node_embed)
+        cache = PrecomputedCache(
+            node_embed,
+            decoder_k,
+            decoder_v,
+            decoder_single_head_k,
+            future_k,
+            future_v,
+        )
         # Main decoding: loop until all sequences are done
         step = 0
         while not td["done"].all():
@@ -245,7 +255,51 @@ def forward(self, x):
         h = self.normalization_2(h + self.feed_forward(h))
         return h # shape: (batch, problem, embedding)
 
-class VRP_Decoder(nn.Module):
+class Now_Future_PointerAttention(nn.Module):
+    """Status-aware now/future glimpse from the project decoder."""
+
+    def __init__(self, embed_dim=128, future_dim=128, future_heads=8):
+        super().__init__()
+        if future_dim % future_heads != 0:
+            raise ValueError("future_dim must be divisible by future_heads")
+        self.future_dim = future_dim
+        self.future_heads = future_heads
+        self.project_out = nn.Linear(embed_dim, embed_dim)
+        self.project_future_query = nn.Linear(embed_dim, future_dim, bias=False)
+        self.project_futures = nn.Linear(future_dim, embed_dim, bias=False)
+        self.gate_fc1 = nn.Linear(embed_dim + future_dim, 1, bias=False)
+        self.gate_fc2 = nn.Linear(1, embed_dim, bias=False)
+
+    def forward(
+        self,
+        query,
+        now_key,
+        now_value,
+        future_key,
+        future_value,
+        now_mask,
+        future_mask,
+    ):
+        now_raw = masked_multi_head_attention(
+            query, now_key, now_value, now_mask
+        )
+        now = self.project_out(now_raw)
+        query_flat = query.transpose(1, 2).reshape(
+            query.size(0), query.size(2), -1
+        )
+        future_query = reshape_by_heads(
+            self.project_future_query(query_flat), self.future_heads
+        )
+        future = masked_multi_head_attention(
+            future_query, future_key, future_value, future_mask
+        )
+        gate = torch.sigmoid(
+            self.gate_fc2(F.relu(self.gate_fc1(torch.cat([now, future], -1))))
+        )
+        return now + gate * self.project_futures(future)
+
+
+class VRP_Decoder(nn.Module):
     def __init__(self, **model_params):
         super().__init__()
         self.model_params = model_params
@@ -255,7 +309,34 @@ def __init__(self, **model_params):
         self.Wq_last = nn.Linear(embedding_dim+5, head_num * qkv_dim, bias=False) # 
         self.Wk = nn.Linear(embedding_dim, head_num * qkv_dim, bias=False)
         self.Wv = nn.Linear(embedding_dim, head_num * qkv_dim, bias=False)
-        self.multi_head_combine = nn.Linear(head_num * qkv_dim, embedding_dim)
+        self.multi_head_combine = nn.Linear(head_num * qkv_dim, embedding_dim)
+        self.use_now_future = self.model_params.get("use_now_future", False)
+        future_dim = self.model_params.get("future_dim", embedding_dim)
+        future_heads = self.model_params.get("future_heads", head_num)
+        if self.use_now_future:
+            self.now_future_pointer = Now_Future_PointerAttention(
+                embedding_dim, future_dim, future_heads
+            )
+            self.project_future_node_embeddings = nn.Linear(
+                embedding_dim, 2 * future_dim, bias=False
+            )
+        else:
+            self.now_future_pointer = None
+            self.project_future_node_embeddings = None
+
+    def precompute_future(self, encoded_nodes):
+        if not self.use_now_future:
+            return None, None
+        future_key, future_val = self.project_future_node_embeddings(
+            encoded_nodes
+        ).chunk(2, dim=-1)
+        future_key = reshape_by_heads(
+            future_key, self.model_params.get("future_heads", self.model_params["head_num"])
+        )
+        future_val = reshape_by_heads(
+            future_val, self.model_params.get("future_heads", self.model_params["head_num"])
+        )
+        return future_key, future_val
 
     def forward(self, td, cache, num_starts, prompt=None):
         td = unbatchify(td, num_starts) #  num_starts * bs -> bs , num_starts
@@ -279,8 +360,21 @@ def forward(self, td, cache, num_starts, prompt=None):
         glimpse_q = reshape_by_heads(self.Wq_last(context_embedding), head_num=self.model_params['head_num'])
         mask = td["action_mask"]
         # mha
-        out_concat = multi_head_attention(glimpse_q, cache.glimpse_key, cache.glimpse_val, mask) # (batch, pomo, head_num*qkv_dim)
-        mh_atten_out = self.multi_head_combine(out_concat) # (batch, pomo, embedding)
+        if self.use_now_future:
+            future_mask = (~mask) & (~td["visited"])
+            future_mask[..., 0] = False
+            mh_atten_out = self.now_future_pointer(
+                glimpse_q,
+                cache.glimpse_key,
+                cache.glimpse_val,
+                cache.future_key,
+                cache.future_val,
+                mask,
+                future_mask,
+            )
+        else:
+            out_concat = multi_head_attention(glimpse_q, cache.glimpse_key, cache.glimpse_val, mask)
+            mh_atten_out = self.multi_head_combine(out_concat)
         # sha
         score = torch.matmul(mh_atten_out, cache.logit_key) # (batch, pomo, problem)
         score_scaled = score / self.model_params['sqrt_embedding_dim'] # (batch, pomo, problem)
@@ -297,13 +391,27 @@ def forward(self, td, cache, num_starts, prompt=None):
 
 ########################################
 # NN SUB CLASS / FUNCTIONS
-def reshape_by_heads(qkv, head_num):
+def reshape_by_heads(qkv, head_num):
     # q.(batch, n, head_num*key_dim)   : n can be either 1 or PROBLEM_SIZE
     batch_s = qkv.size(0)
     n = qkv.size(1)
     q_reshaped = qkv.reshape(batch_s, n, head_num, -1) # (batch, n, head_num, key_dim)
     q_transposed = q_reshaped.transpose(1, 2) # (batch, head_num, n, key_dim)
-    return q_transposed
+    return q_transposed
+
+
+def masked_multi_head_attention(q, k, v, visibility):
+    """Boolean-masked MHA with zero output for an empty visibility set."""
+    score = torch.matmul(q, k.transpose(2, 3))
+    score = score / (q.size(-1) ** 0.5)
+    empty = ~visibility.any(dim=-1, keepdim=True)
+    safe_mask = visibility.clone()
+    safe_mask[..., 0] |= empty.squeeze(-1)
+    score = score.masked_fill(~safe_mask[:, None, :, :], float("-inf"))
+    weights = F.softmax(score, dim=-1)
+    out = torch.matmul(weights, v)
+    out = out.transpose(1, 2).reshape(q.size(0), q.size(2), -1)
+    return out.masked_fill(empty, 0)
 
 def multi_head_attention(q, k, v, ninf_mask=None, sparse=False):
     # q (batch, head_num, n, key_dim)   : n can be either 1 or PROBLEM_SIZE

50/config.yaml

@@ -9,7 +9,10 @@ model_params:
    'ffd': 'siglu' 
    'norm_type': 'rms' 
    'use_sparse': 'topk'   # relu topk entmax15 sparsemax a-entmax False 
-   'p_num': 5
+   'p_num': 5
+   'use_now_future': false
+   'future_dim': 128
+   'future_heads': 8
 
 optimizer_params:
     'optimizer':

50/envs/env.py

@@ -146,13 +146,15 @@ def reset(self, td: Optional[TensorDict] = None, batch_size=None) -> TensorDict:
             td = self.generator(batch_size=batch_size).to('cuda')
         batch_size = [batch_size] if isinstance(batch_size, int) else batch_size
         self.to(td.device)
-        return super().reset(td, batch_size=batch_size)
+        # The rollout batch is dynamic and ``step`` already bypasses EnvBase.
+        # TorchRL 0.6 otherwise rejects it against the empty EnvBase batch.
+        return self._reset(td, batch_size=batch_size)
 
     def _reset(self, td: Optional[TensorDict] = None, batch_size: Optional[list] = None,) -> TensorDict:
 
         device = td.device
         # Create reset TensorDict
-        td_reset = TensorDict(
+        td_reset = TensorDict(
             {
                 "locs": td["locs"],
                 "demand_backhaul": td["demand_backhaul"],
@@ -172,9 +174,11 @@ def _reset(self, td: Optional[TensorDict] = None, batch_size: Optional[list] = N
                 "visited": torch.zeros((*batch_size, td["locs"].shape[-2]), dtype=torch.bool, device=device,),
             },
             batch_size=batch_size,
-            device=device,
-        )
-        td_reset.set("action_mask", self.get_action_mask(td_reset))
+            device=device,
+        )
+        if "p_s_tag" in td.keys():
+            td_reset.set("p_s_tag", td["p_s_tag"])
+        td_reset.set("action_mask", self.get_action_mask(td_reset))
         return td_reset
 
     def dataset(self, data_size=None, phase="train"):
@@ -562,4 +566,4 @@ def __setstate__(self, state, set_seed:bool=True):
         self.__dict__.update(state)
         if set_seed:
             self.rng = torch.manual_seed(0) # = torch.default_generator
-            self.rng.set_state(state["rng"].cpu()) # = set  torch.default_generator
+            self.rng.set_state(state["rng"].cpu()) # = set  torch.default_generator