Fixed: IQL Training for Baseline

Author: Vishal-sys-code

results/all_runs/iql_CartPole-v1/training.log

@@ -18,3 +18,9 @@
 2025-11-06 11:37:53,843 [INFO] Checking for dataset...
 2025-11-06 11:37:54,177 [INFO] Dataset found at D:\Github\neuromorphic_decision_transformer\data\CartPole-v1\dataset.npz.
 2025-11-06 11:37:54,178 [INFO] Starting training...
+2025-11-06 12:36:23,691 [INFO] Checking for dataset...
+2025-11-06 12:36:23,705 [INFO] Dataset found at D:\Github\neuromorphic_decision_transformer\data\CartPole-v1\dataset.npz.
+2025-11-06 12:36:23,706 [INFO] Starting training...
+2025-11-06 12:37:13,769 [INFO] Dataset size: 22826 clips
+2025-11-06 12:37:14,136 [INFO] DataLoader created with num_workers=0 and pin_memory=False.
+2025-11-06 12:37:30,656 [INFO] Starting training loop...

snn-dt/scripts/train.py

@@ -57,62 +57,73 @@ class OfflineTransitionDataset(Dataset):
     def __init__(self, dataset_path):
         data = np.load(dataset_path)
 
-        # First, calculate the total number of transitions
-        total_transitions = 0
-        for i in range(data["states"].shape[0]):
-            mask = data["mask"][i]
-            clip_len = int(mask.sum())
-            clip_len = min(clip_len, data["states"].shape[1]) # Cap clip_len
-            total_transitions += clip_len
-
-        # Pre-allocate numpy arrays
-        state_dim = data["states"].shape[2]
-        action_dim = data["actions"].shape[2]
+        states_list, actions_list, rewards_list, next_states_list, dones_list = [], [], [], [], []
 
-        self.states = np.zeros((total_transitions, state_dim), dtype=np.float32)
-        self.actions = np.zeros((total_transitions, action_dim), dtype=np.float32)
-        self.rewards = np.zeros((total_transitions, 1), dtype=np.float32)
-        self.next_states = np.zeros((total_transitions, state_dim), dtype=np.float32)
-        self.dones = np.zeros((total_transitions, 1), dtype=np.float32)
-        
-        current_idx = 0
+        # Check if actions exist and get action_dim, otherwise handle gracefully
+        if "actions" not in data or data["actions"].shape[0] == 0:
+            action_dim = 0 # Placeholder, this dataset will be empty
+        else:
+            action_dim = data["actions"].shape[2]
+
         for i in range(data["states"].shape[0]):
             mask = data["mask"][i]
             clip_len = int(mask.sum())
-            clip_len = min(clip_len, data["states"].shape[1]) # Cap clip_len
-
+            
             if clip_len == 0:
                 continue
 
-            # Get the trajectory data
+            # Trajectory data
             traj_states = data["states"][i, :clip_len]
-            traj_actions = data["actions"][i, :clip_len]
             traj_rtg = data["returns_to_go"][i, :clip_len]
 
-            # Populate states and actions
-            self.states[current_idx : current_idx + clip_len] = traj_states
-            self.actions[current_idx : current_idx + clip_len] = traj_actions
+            # Add all states for this trajectory to the list
+            states_list.append(traj_states)
 
-            # Populate rewards, next_states, and dones
+            # Actions: N-1 real actions, plus one dummy action for the terminal state
+            traj_actions_list = []
             if clip_len > 1:
-                self.rewards[current_idx : current_idx + clip_len - 1] = (traj_rtg[:-1] - traj_rtg[1:]).reshape(-1, 1)
-                self.next_states[current_idx : current_idx + clip_len - 1] = traj_states[1:]
-                self.dones[current_idx : current_idx + clip_len - 1] = 0.0
-
-            # Final transition
-            if clip_len > 0:
-                self.rewards[current_idx + clip_len - 1] = traj_rtg[-1]
-                self.next_states[current_idx + clip_len - 1] = np.zeros_like(traj_states[-1])
-                self.dones[current_idx + clip_len - 1] = 1.0
+                traj_actions_list.append(data["actions"][i, :clip_len - 1])
+            
+            # Add a dummy action for the terminal state
+            dummy_action = np.zeros((1, action_dim), dtype=np.float32)
+            traj_actions_list.append(dummy_action)
+            actions_list.append(np.concatenate(traj_actions_list, axis=0))
+
+            # Rewards and next_states
+            rewards = np.zeros((clip_len, 1), dtype=np.float32)
+            next_states = np.zeros_like(traj_states)
+            dones = np.zeros((clip_len, 1), dtype=np.float32)
+
+            if clip_len > 1:
+                # Rewards for non-terminal states
+                rewards[:-1] = (traj_rtg[:-1] - traj_rtg[1:]).reshape(-1, 1)
+                # Next_states for non-terminal states
+                next_states[:-1] = traj_states[1:]
+            
+            # Handle terminal transition
+            rewards[-1] = traj_rtg[-1]  # This is RTG, not a true reward, but matches original logic
+            # next_states[-1] is already zeros
+            dones[-1] = 1.0
 
-            current_idx += clip_len
-
-        # Convert to torch tensors
-        self.states = torch.from_numpy(self.states).float()
-        self.actions = torch.from_numpy(self.actions).float()
-        self.rewards = torch.from_numpy(self.rewards).float()
-        self.next_states = torch.from_numpy(self.next_states).float()
-        self.dones = torch.from_numpy(self.dones).float()
+            rewards_list.append(rewards)
+            next_states_list.append(next_states)
+            dones_list.append(dones)
+
+        # Handle case where no valid trajectories were found
+        if not states_list:
+            self.states = torch.empty(0, data["states"].shape[2], dtype=torch.float32)
+            self.actions = torch.empty(0, action_dim, dtype=torch.float32)
+            self.rewards = torch.empty(0, 1, dtype=torch.float32)
+            self.next_states = torch.empty(0, data["states"].shape[2], dtype=torch.float32)
+            self.dones = torch.empty(0, 1, dtype=torch.float32)
+            return
+
+        # Concatenate and convert to torch tensors
+        self.states = torch.from_numpy(np.concatenate(states_list, axis=0)).float()
+        self.actions = torch.from_numpy(np.concatenate(actions_list, axis=0)).float()
+        self.rewards = torch.from_numpy(np.concatenate(rewards_list, axis=0)).float()
+        self.next_states = torch.from_numpy(np.concatenate(next_states_list, axis=0)).float()
+        self.dones = torch.from_numpy(np.concatenate(dones_list, axis=0)).float()
 
     def __len__(self):
         return len(self.states)

snn-dt/src/models/iql.py

@@ -103,8 +103,9 @@ def __init__(self, cfg):
         self.tau = cfg.iql.tau
         self.temperature = cfg.iql.temperature
         self.expectile = cfg.iql.expectile
+        self.is_discrete = cfg.dataset.is_discrete
 
-        self.actor = Actor(cfg.dataset.state_dim, cfg.dataset.act_dim, cfg.iql.hidden_size).to(self.device)
+        self.actor = Actor(cfg.dataset.state_dim, cfg.dataset.act_dim, cfg.iql.hidden_size, is_discrete=self.is_discrete).to(self.device)
         self.critic1 = Critic(cfg.dataset.state_dim, cfg.dataset.act_dim, cfg.iql.hidden_size).to(self.device)
         self.critic2 = Critic(cfg.dataset.state_dim, cfg.dataset.act_dim, cfg.iql.hidden_size).to(self.device)
         self.value_net = Value(cfg.dataset.state_dim, cfg.iql.hidden_size).to(self.device)
@@ -114,10 +115,10 @@ def __init__(self, cfg):
         self.critic1_target.load_state_dict(self.critic1.state_dict())
         self.critic2_target.load_state_dict(self.critic2.state_dict())
 
-        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=cfg.training.lr)
-        self.critic1_optimizer = optim.Adam(self.critic1.parameters(), lr=cfg.training.lr)
-        self.critic2_optimizer = optim.Adam(self.critic2.parameters(), lr=cfg.training.lr)
-        self.value_optimizer = optim.Adam(self.value_net.parameters(), lr=cfg.training.lr)
+        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=float(cfg.training.lr))
+        self.critic1_optimizer = optim.Adam(self.critic1.parameters(), lr=float(cfg.training.lr))
+        self.critic2_optimizer = optim.Adam(self.critic2.parameters(), lr=float(cfg.training.lr))
+        self.value_optimizer = optim.Adam(self.value_net.parameters(), lr=float(cfg.training.lr))
 
     def forward(self, batch):
         # IQL has a custom learn method, so forward is a no-op for now
@@ -134,8 +135,13 @@ def learn(self, batch):
 
         # Value loss
         with torch.no_grad():
-            q1 = self.critic1_target(states, actions)
-            q2 = self.critic2_target(states, actions)
+            if self.is_discrete:
+                actions_one_hot = F.one_hot(actions.squeeze().long(), num_classes=self.actor.logits.out_features).float()
+                q1 = self.critic1_target(states, actions_one_hot)
+                q2 = self.critic2_target(states, actions_one_hot)
+            else:
+                q1 = self.critic1_target(states, actions)
+                q2 = self.critic2_target(states, actions)
             min_q = torch.min(q1, q2)
         value = self.value_net(states)
         value_loss = loss_fn(min_q - value, self.expectile).mean()
@@ -159,8 +165,15 @@ def learn(self, batch):
         with torch.no_grad():
             next_v = self.value_net(next_states)
             q_target = rewards + self.gamma * (1 - dones) * next_v
-        q1 = self.critic1(states, actions)
-        q2 = self.critic2(states, actions)
+
+        if self.is_discrete:
+            actions_one_hot = F.one_hot(actions.squeeze().long(), num_classes=self.actor.logits.out_features).float()
+            q1 = self.critic1(states, actions_one_hot)
+            q2 = self.critic2(states, actions_one_hot)
+        else:
+            q1 = self.critic1(states, actions)
+            q2 = self.critic2(states, actions)
+        
         critic1_loss = F.mse_loss(q1, q_target)
         critic2_loss = F.mse_loss(q2, q_target)