Fixed: DSFormer Baseline

Vishal-sys-code · Vishal-sys-code · commit e5a6d76753e6 · 2025-11-09T17:05:02.000+05:30
diff --git a/debug_snn_dt.py b/debug_snn_dt.py
@@ -0,0 +1,85 @@
+import torch
+import sys
+import os
+
+# Add snn-dt to the python path to allow imports from src
+sys.path.insert(0, os.path.abspath("snn-dt"))
+
+from src.models.snn_dt import SnnDt
+
+# A simplified MockConfig, similar to the one in the tests
+class MockConfig:
+    def __init__(self):
+        self.model = self.Model()
+        self.dataset = self.Dataset()
+        self.snn = self.Snn()
+        self.env = "dummy_env"
+
+    class Model:
+        name = "snn_dt"
+        d_model = 128
+        n_heads = 4
+        n_layers = 2
+        
+    class Dataset:
+        state_dim = 4
+        act_dim = 1
+        max_timesteps = 100
+        is_discrete = False
+        
+    class Snn:
+        lif_tau = 20.0
+        surrogate_k = 25.0
+        use_plasticity = False
+
+    class Training:
+        device = "cpu"
+
+def debug_snn_dt():
+    """
+    Instantiates the SnnDt model, runs a forward pass, and checks the spike count.
+    """
+    print("--- Initializing SNN-DT Debug Script ---")
+    
+    # 1. Setup model and config
+    cfg = MockConfig()
+    model = SnnDt(cfg)
+    model.eval()  # Use eval mode to disable training-specific logic like plasticity
+
+    print("SnnDt model instantiated successfully.")
+
+    # 2. Create a batch of dummy data with high magnitude to encourage spiking
+    batch = {
+        "states": torch.randn(16, 20, 4) * 100,
+        "actions": torch.randn(16, 20, 1) * 100,
+        "returns_to_go": torch.randn(16, 20, 1) * 100,
+        "timesteps": torch.randint(0, 100, (16, 20)),
+        "mask": torch.ones(16, 20),
+    }
+
+    print("Batch created. Running a single forward pass...")
+
+    # 3. Run the forward pass and check spikes
+    with torch.no_grad():
+        model(batch)
+    
+    spike_count_1 = model.count_spikes()
+    print(f"Spike count after first pass: {spike_count_1}")
+
+    # 4. Run a second pass to check accumulation
+    print("Running a second forward pass to check accumulation...")
+    with torch.no_grad():
+        model(batch)
+
+    spike_count_2 = model.count_spikes()
+    print(f"Spike count after second pass: {spike_count_2}")
+    
+    # 5. Check the reset mechanism
+    print("Resetting spike counts...")
+    model.reset_spike_counts()
+    print(f"Spike count after reset: {model.count_spikes()}")
+    
+    print("--- Debug Script Finished ---")
+
+if __name__ == "__main__":
+    debug_snn_dt()
diff --git a/results/all_runs/dsformer_CartPole-v1/training.log b/results/all_runs/dsformer_CartPole-v1/training.log
@@ -0,0 +1,24 @@
+2025-11-09 10:11:25,647 [INFO] Checking for dataset...
+2025-11-09 10:11:25,655 [INFO] Dataset found at D:\Github\neuromorphic_decision_transformer\data\CartPole-v1\dataset.npz.
+2025-11-09 10:11:25,656 [INFO] Starting training...
+2025-11-09 10:11:25,818 [INFO] Dataset size: 1000 clips
+2025-11-09 10:11:25,874 [INFO] DataLoader created with num_workers=0 and pin_memory=False.
+2025-11-09 10:11:39,234 [INFO] Starting training loop...
+2025-11-09 10:26:07,900 [INFO] Checking for dataset...
+2025-11-09 10:26:07,902 [INFO] Dataset found at D:\Github\neuromorphic_decision_transformer\data\CartPole-v1\dataset.npz.
+2025-11-09 10:26:07,903 [INFO] Starting training...
+2025-11-09 10:26:08,336 [INFO] Dataset size: 1000 clips
+2025-11-09 10:26:08,369 [INFO] DataLoader created with num_workers=0 and pin_memory=False.
+2025-11-09 10:26:22,564 [INFO] Starting training loop...
+2025-11-09 14:47:13,401 [INFO] Checking for dataset...
+2025-11-09 14:47:13,403 [INFO] Dataset found at D:\Github\neuromorphic_decision_transformer\data\CartPole-v1\dataset.npz.
+2025-11-09 14:47:13,404 [INFO] Starting training...
+2025-11-09 14:47:13,575 [INFO] Dataset size: 1000 clips
+2025-11-09 14:47:13,635 [INFO] DataLoader created with num_workers=0 and pin_memory=False.
+2025-11-09 14:47:29,057 [INFO] Starting training loop...
+2025-11-09 14:50:54,988 [INFO] Checking for dataset...
+2025-11-09 14:50:55,000 [INFO] Dataset found at D:\Github\neuromorphic_decision_transformer\data\CartPole-v1\dataset.npz.
+2025-11-09 14:50:55,001 [INFO] Starting training...
+2025-11-09 14:50:55,110 [INFO] Dataset size: 1000 clips
+2025-11-09 14:50:55,139 [INFO] DataLoader created with num_workers=0 and pin_memory=False.
+2025-11-09 14:51:05,312 [INFO] Starting training loop...
diff --git a/results/all_runs/snn_dt_CartPole-v1/training.log b/results/all_runs/snn_dt_CartPole-v1/training.log
@@ -0,0 +1,6 @@
+2025-11-09 10:06:36,149 [INFO] Checking for dataset...
+2025-11-09 10:06:36,159 [INFO] Dataset found at D:\Github\neuromorphic_decision_transformer\data\CartPole-v1\dataset.npz.
+2025-11-09 10:06:36,164 [INFO] Starting training...
+2025-11-09 10:06:36,398 [INFO] Dataset size: 1000 clips
+2025-11-09 10:06:36,472 [INFO] DataLoader created with num_workers=0 and pin_memory=False.
+2025-11-09 10:06:56,621 [INFO] Starting training loop...
diff --git a/run_tests.py b/run_tests.py
@@ -0,0 +1,9 @@
+import sys
+import os
+import pytest
+
+if __name__ == "__main__":
+    print("Current working directory:", os.getcwd())
+    sys.path.insert(0, os.path.abspath("snn-dt"))
+    print("sys.path:", sys.path)
+    sys.exit(pytest.main(["-x", "snn-dt/tests/test_models.py"]))
diff --git a/snn-dt/scripts/train.py b/snn-dt/scripts/train.py
@@ -23,10 +23,10 @@
 import warnings
 warnings.filterwarnings('ignore')
 
-# from src.models.cql import CQL
-# from src.models.dt import DecisionTransformer
-# from src.models.dsformer import DsFormer
-# from src.models.iql import IQL
+from src.models.cql import CQL
+from src.models.dt import DecisionTransformer
+from src.models.dsformer import DsFormer
+from src.models.iql import IQL
 from src.models.snn_dt import SnnDt
 from src.utils.config import AttrDict
 from src.utils.models import get_model
@@ -195,6 +195,9 @@ def train(cfg, logger):
     for epoch in range(cfg.training.epochs):
         start_time = time.time()
         epoch_losses = []
+
+        if hasattr(model, "reset_spike_counts"):
+            model.reset_spike_counts()
         
         train_iter = iter(train_loader)
         pbar = tqdm(range(cfg.training.batches_per_epoch), desc=f"Epoch {epoch+1}/{cfg.training.epochs}")
@@ -250,10 +253,12 @@ def train(cfg, logger):
             log_str = f"Epoch {epoch+1}/{cfg.training.epochs} | Time: {epoch_time:.2f}s | Loss: {avg_loss:.4f}"
             
             # Spike counting for SNN models
-            if isinstance(model, SnnDt):
+            if hasattr(model, "count_spikes"):
                 spikes = model.count_spikes()
-                log_str += f" | Spikes: {spikes}"
+                log_str += f" | Spikes: {spikes:.2f}"
                 eval_results["spikes"] = spikes
+            else:
+                eval_results["spikes"] = 0.0
 
             metrics.append({"epoch": epoch + 1, "loss": avg_loss, **eval_results, "time_s": epoch_time})
             log_str += f" | Eval Return: {eval_results['return_mean']:.2f}"
@@ -376,6 +381,11 @@ def main():
     
     # Convert to AttrDict for easy access
     cfg = AttrDict(cfg)
+
+    # Adaptive training controls for SNNs
+    if "snn" in cfg.model.name or "dsformer" in cfg.model.name:
+        cfg.training.batches_per_epoch = min(cfg.training.batches_per_epoch, cfg_raw.get("snn_batches_per_epoch", 100))
+        cfg.training.eval_every = max(cfg.training.eval_every, cfg_raw.get("snn_eval_every", 50))
     
     # Construct dataset path from env name, relative to project root
     cfg.dataset.path = str(project_root / f"data/{args.env}/dataset.npz")
diff --git a/snn-dt/src/models/dsformer.py b/snn-dt/src/models/dsformer.py
@@ -1,6 +1,6 @@
 import torch
 import torch.nn as nn
-from norse.torch.module.lif import LIFCell
+from norse.torch.module.lif import LIF, LIFCell, LIFParameters
 
 from src.models.base import BasePolicy
 
@@ -16,8 +16,14 @@ def __init__(self, d_model, n_heads, lif_tau, surrogate_k):
         self.k_proj = nn.Linear(d_model, d_model)
         self.v_proj = nn.Linear(d_model, d_model)
 
-        self.q_lif = LIFCell()
-        self.k_lif = LIFCell()
+        p = LIFParameters(
+            tau_mem_inv=torch.tensor(1.0 / lif_tau),
+            v_th=torch.tensor(0.8),
+            method="super",
+            alpha=surrogate_k,
+        )
+        self.q_lif = LIF(p=p)
+        self.k_lif = LIF(p=p)
 
         self.spike_count = 0
 
@@ -28,15 +34,8 @@ def forward(self, x, state_q, state_k, attn_mask=None):
         k = self.k_proj(x)
         v = self.v_proj(x)
 
-        spikes_q_seq = []
-        spikes_k_seq = []
-        for t in range(seq_len):
-            spikes_q, state_q = self.q_lif(q[:, t], state_q)
-            spikes_k, state_k = self.k_lif(k[:, t], state_k)
-            spikes_q_seq.append(spikes_q)
-            spikes_k_seq.append(spikes_k)
-        spikes_q = torch.stack(spikes_q_seq, dim=1)
-        spikes_k = torch.stack(spikes_k_seq, dim=1)
+        spikes_q, _ = self.q_lif(q)
+        spikes_k, _ = self.k_lif(k)
 
         q_reshaped = spikes_q.view(batch_size, seq_len, self.n_heads, self.head_dim).permute(0, 2, 1, 3)
         k_reshaped = spikes_k.view(batch_size, seq_len, self.n_heads, self.head_dim).permute(0, 2, 1, 3)
@@ -49,8 +48,10 @@ def forward(self, x, state_q, state_k, attn_mask=None):
         
         attn_output = (attn_weights @ v_reshaped).permute(0, 2, 1, 3).reshape(batch_size, seq_len, self.d_model)
 
-        self.spike_count = spikes_q.sum() + spikes_k.sum()
-        return attn_output, state_q, state_k
+        if not hasattr(self, "spike_count"):
+            self.spike_count = 0.0
+        self.spike_count += (spikes_q.sum() + spikes_k.sum()).item()
+        return attn_output
 
 
 class DsFormer(BasePolicy, nn.Module):
@@ -95,6 +96,18 @@ def forward(self, batch):
             action_input = actions
 
         action_embeddings = self.embed_action(action_input)
+
+        # Pad action embeddings to match state/return sequence length
+        if action_embeddings.shape[1] < seq_len:
+            padding_size = seq_len - action_embeddings.shape[1]
+            padding = torch.zeros(
+                action_embeddings.shape[0],
+                padding_size,
+                action_embeddings.shape[2],
+                device=action_embeddings.device,
+            )
+            action_embeddings = torch.cat([action_embeddings, padding], dim=1)
+
         return_embeddings = self.embed_return(batch["returns_to_go"])
         time_embeddings = self.embed_timestep(batch["timesteps"])
 
@@ -110,10 +123,8 @@ def forward(self, batch):
         x = self.embed_ln(stacked_inputs)
 
         attn_mask = nn.Transformer.generate_square_subsequent_mask(x.shape[1], device=x.device)
-        q_states = [None] * len(self.blocks)
-        k_states = [None] * len(self.blocks)
         for i, block in enumerate(self.blocks):
-            x, q_states[i], k_states[i] = block(x, q_states[i], k_states[i], attn_mask=attn_mask)
+            x = block(x, None, None, attn_mask=attn_mask)
 
         action_preds = self.action_predictor(x[:, 1::3])
         return action_preds
@@ -144,7 +155,12 @@ def load(self, path):
         self.load_state_dict(torch.load(path))
 
     def count_spikes(self):
-        return sum(b.spike_count for b in self.blocks)
+        total_spikes = sum(block.spike_count for block in self.blocks)
+        return total_spikes / len(self.blocks) if len(self.blocks) > 0 else 0.0
+
+    def reset_spike_counts(self):
+        for block in self.blocks:
+            block.spike_count = 0.0
 
     def num_params(self):
         return sum(p.numel() for p in self.parameters() if p.requires_grad)
diff --git a/snn-dt/src/models/snn_dt.py b/snn-dt/src/models/snn_dt.py
@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 from norse.torch.module.leaky_integrator import LICell
-from norse.torch.module.lif import LIFCell, LIFParameters
+from norse.torch.module.lif import LIF, LIFCell, LIFParameters
 
 from src.models.base import BasePolicy
 
@@ -27,8 +27,8 @@ def __init__(self, cfg, d_model, n_heads, lif_tau, surrogate_k):
             alpha=surrogate_k,
         )
 
-        self.q_lif = LIFCell(p=p)
-        self.k_lif = LIFCell(p=p)
+        self.q_lif = LIF(p=p)
+        self.k_lif = LIF(p=p)
         self.v_li = LICell()
         
         self.use_plasticity = False # Will be set by SnnDt
@@ -52,16 +52,8 @@ def forward(self, x, state_q, state_k, attn_mask=None):
         k = self.k_proj(x)
         v = self.v_proj(x)
 
-        # Spiking Q and K
-        spikes_q_seq = []
-        spikes_k_seq = []
-        for t in range(seq_len):
-            spikes_q, state_q = self.q_lif(q[:, t], state_q)
-            spikes_k, state_k = self.k_lif(k[:, t], state_k)
-            spikes_q_seq.append(spikes_q)
-            spikes_k_seq.append(spikes_k)
-        spikes_q = torch.stack(spikes_q_seq, dim=1)
-        spikes_k = torch.stack(spikes_k_seq, dim=1)
+        spikes_q, _ = self.q_lif(q)
+        spikes_k, _ = self.k_lif(k)
 
         # Attention
         q_reshaped = spikes_q.view(batch_size, seq_len, self.n_heads, self.head_dim)
@@ -79,13 +71,15 @@ def forward(self, x, state_q, state_k, attn_mask=None):
         routing_gate = self.routing_mlp(attn_output)
         out = attn_output * routing_gate
 
-        self.spike_count = spikes_q.sum() + spikes_k.sum()
+        if not hasattr(self, "spike_count"):
+            self.spike_count = 0.0
+        self.spike_count += (spikes_q.sum() + spikes_k.sum()).item()
         
         # Three-factor plasticity
         if self.training and self.use_plasticity:
             self.update_eligibility_trace(spikes_q, v)
             
-        return out, state_q, state_k
+        return out
 
     def update_eligibility_trace(self, presynaptic_spikes, postsynaptic_potential):
         # Simplified eligibility trace update
@@ -193,10 +187,8 @@ def forward(self, batch):
         
         # Spiking transformer blocks
         attn_mask = nn.Transformer.generate_square_subsequent_mask(x.shape[1], device=x.device)
-        q_states = [None] * len(self.blocks)
-        k_states = [None] * len(self.blocks)
         for i, block in enumerate(self.blocks):
-            x, q_states[i], k_states[i] = block(x, q_states[i], k_states[i], attn_mask=attn_mask)
+            x = block(x, None, None, attn_mask=attn_mask)
         
         action_preds = self.action_predictor(x[:, 1::3])
         return action_preds
@@ -227,7 +219,12 @@ def load(self, path):
         self.load_state_dict(torch.load(path))
 
     def count_spikes(self):
-        return sum(b.spike_count for b in self.blocks)
+        total_spikes = sum(block.spike_count for block in self.blocks)
+        return total_spikes / len(self.blocks) if len(self.blocks) > 0 else 0.0
+
+    def reset_spike_counts(self):
+        for block in self.blocks:
+            block.spike_count = 0.0
 
     def num_params(self):
         return sum(p.numel() for p in self.parameters() if p.requires_grad)
diff --git a/snn-dt/tests/test_models.py b/snn-dt/tests/test_models.py
