fixing my ablations...

Author: Vishal-sys-code

ablation_studies/experiment_contract_light.yaml

@@ -11,7 +11,7 @@ batch_size: 64
 optimizer: AdamW
 lr: 3e-4
 weight_decay: 1e-2
-epochs: 100
+epochs: 500
 local_lr_eta_local: 0.05
 surrogate_slope_k: 10
 spike_energy_pJ: 5.0
@@ -32,4 +32,4 @@ cql:
   hidden_size: 256
   with_lagrange: false
   cql_weight: 1.0
-  target_action_gap: 10.0
+  target_action_gap: 10.0

ablation_studies/run_experiment.py

@@ -19,6 +19,10 @@
 # --- Add Project Root to sys.path ---
 project_root = Path(__file__).resolve().parent.parent
 sys.path.append(str(project_root))
+# Add snn-dt/src to sys.path to allow for model imports
+snn_dt_src_path = project_root / 'snn-dt' / 'src'
+if snn_dt_src_path.exists():
+    sys.path.insert(0, str(snn_dt_src_path))
 
 # --- Local Imports ---
 from ablation_studies.src.datasets import OfflineSequenceDataset, OfflineTransitionDataset
@@ -52,10 +56,10 @@ def load_config(contract_path, variant_path):
 # --- Model Factory ---
 def get_model(cfg):
     model_name_map = {
-        'dt': ('snn_dt.src.models.dt', 'DecisionTransformer'),
-        'snn_dt': ('snn_dt.src.models.snn_dt', 'SnnDt'),
-        'iql': ('snn_dt.src.models.iql', 'IQL'),
-        'cql': ('snn_dt.src.models.cql', 'CQL'),
+        'dt': ('models.dt', 'DecisionTransformer'),
+        'snn_dt': ('models.snn_dt', 'SnnDt'),
+        'iql': ('models.iql', 'IQL'),
+        'cql': ('models.cql', 'CQL'),
         'ablation_dsformer': ('ablation_studies.src.models.ablation_dsformer', 'AblationDsFormer'),
     }
 
@@ -115,9 +119,9 @@ def evaluate_policy(model, env_name, cfg):
                         action = int(np.argmax(action))
 
                 state, reward, terminated, truncated, _ = env.step(action)
-                done = terminated or truncated or (t + 1 >= cfg.dataset.max_timesteps)
+                done = terminated or truncated or (t >= cfg.sequence_length_N - 1)
 
-                if t < cfg.sequence_length_N - 1:
+                if not done:
                     actions[0, t+1] = torch.tensor(action, device=cfg.device)
                     states[0, t+1] = torch.from_numpy(state).to(cfg.device)
                     rtgs[0, t+1] = rtgs[0, t] - reward
@@ -160,14 +164,18 @@ def train(cfg, logger):
     dataset_args = {'path': str(dataset_path)}
     if not is_transition_model: dataset_args['seq_len'] = cfg.sequence_length_N
     dataset = DatasetClass(**dataset_args)
-    train_loader = DataLoader(dataset, batch_size=cfg.batch_size, shuffle=True, num_workers=min(os.cpu_count(), 4))
+    # NOTE: num_workers is set to 0 to avoid a hanging issue with multiprocessing.
+    train_loader = DataLoader(dataset, batch_size=cfg.batch_size, shuffle=True, num_workers=0)
 
     model = get_model(cfg).to(cfg.device)
     optimizer = torch.optim.AdamW(model.parameters(), lr=float(cfg.lr), weight_decay=float(cfg.weight_decay)) if list(model.parameters()) else None
     loss_fn = torch.nn.MSELoss()
 
     logger.info(json.dumps({"train/param_count": sum(p.numel() for p in model.parameters())}))
 
+    # Create save directory if it doesn't exist
+    Path(cfg.save_dir).mkdir(parents=True, exist_ok=True)
+
     for epoch in range(1, cfg.epochs + 1):
         model.train()
         for batch_idx, batch in enumerate(tqdm(train_loader, desc=f"Epoch {epoch}/{cfg.epochs}", file=sys.stderr)):

ablation_studies/runs/no_plasticity/seed_0/CartPole-v1/ckpt_epoch_10.pt

@@ -7,3 +7,15 @@
 {"train/param_count": 399889}
 {"train/step": 250, "train/loss": 0.1063467487692833}
 {"train/step": 500, "train/loss": 0.08335910737514496}
+{"train/param_count": 399889}
+{"train/step": 250, "train/loss": 0.1063467487692833}
+{"train/step": 500, "train/loss": 0.08335910737514496}
+{"epoch": 10, "val/mean_return": 20.0, "val/std_return": 0.0}
+{"train/step": 750, "train/loss": 0.08903709799051285}
+{"train/step": 1000, "train/loss": 0.06532430648803711}
+{"train/step": 1250, "train/loss": 0.07413662225008011}
+{"epoch": 20, "val/mean_return": 20.0, "val/std_return": 0.0}
+{"train/step": 1500, "train/loss": 0.06820853054523468}
+{"train/step": 1750, "train/loss": 0.06684345006942749}
+{"train/step": 2000, "train/loss": 0.06681139767169952}
+{"epoch": 30, "val/mean_return": 20.0, "val/std_return": 0.0}

ablation_studies/runs/no_plasticity/seed_0/CartPole-v1/ckpt_epoch_20.pt

@@ -1,2 +1,2 @@
 Command: D:\Github\neuromorphic_decision_transformer\ablation_studies\run_experiment.py --variant no_plasticity --env CartPole-v1 --seed 0 --contract experiment_contract_light.yaml
-Git Hash: fd20c7baafb79138e1195b4cb4a1a62d3ace4f60
+Git Hash: b1a848ff72b193feb2a523af890fd78e7a644a82

ablation_studies/runs/no_plasticity/seed_0/CartPole-v1/ckpt_epoch_30.pt

@@ -40,20 +40,19 @@ def main():
                     "--contract", CONTRACT
                 ]
 
-                print(f"[{current_job}/{total_jobs}] Running: Variant={variant}, Env={env}, Seed={seed}")
+                print(f"\n--- [{current_job}/{total_jobs}] Running: Variant={variant}, Env={env}, Seed={seed} ---")
 
                 if args.dry_run:
-                    print(f"Command: {' '.join(cmd)}")
+                    print(f"  Command: {' '.join(cmd)}")
                 else:
                     try:
                         subprocess.run(cmd, check=True)
+                        print(f"--- Finished: Variant={variant}, Env={env}, Seed={seed} (Success) ---")
                     except subprocess.CalledProcessError as e:
-                        print(f"Error running job: {e}")
-                        # Depending on preference, we might want to continue or stop. 
-                        # For now, let's continue to the next one but log the error.
-                        print("Continuing to next job...")
+                        print(f"  Error running job: {e}")
+                        print(f"--- Finished: Variant={variant}, Env={env}, Seed={seed} (Failed) ---")
 
-    print("--- All targeted experimental runs complete! ---")
+    print("\n--- All targeted experimental runs complete! ---")
 
 if __name__ == "__main__":
     main()

ablation_studies/runs/no_plasticity/seed_0/CartPole-v1/metrics.jsonl

@@ -32,8 +32,8 @@ class OfflineTransitionDataset(Dataset):
     Dataset for transition-based models like IQL and CQL.
     Processes trajectories into individual (s, a, r, s', d) transitions.
     """
-    def __init__(self, dataset_path):
-        data = np.load(dataset_path, mmap_mode='r')
+    def __init__(self, path):
+        data = np.load(path, mmap_mode='r')
 
         # Calculate total number of transitions
         total_transitions = int(np.sum(data['mask'])) - data['mask'].shape[0]

ablation_studies/runs/no_plasticity/seed_0/CartPole-v1/run_info.txt

@@ -138,7 +138,7 @@ def forward(self, batch):
 
         state_embed = self.embed_state(batch["states"])
         action_embed = self.embed_action(batch["actions"])
-        rtg_embed = self.embed_return(batch["returns_to_go"])
+        rtg_embed = self.embed_return(batch["returns_to_go"].float())
         time_embed = self.embed_timestep(batch["timesteps"].squeeze(-1))
 
         state_embed, action_embed, rtg_embed = state_embed + time_embed, action_embed + time_embed, rtg_embed + time_embed