Fixed: IQL Training on Acrobot ENV

Author: Vishal-sys-code

configs/iql_acrobot.yaml

@@ -8,5 +8,6 @@ learning_rate: 1e-4
 epochs: 1000
 eval_every: 10
 checkpoint_every: 50
+batches_per_epoch: 100
 
 num_workers: 1

snn-dt/scripts/train.py

@@ -80,7 +80,8 @@ def __init__(self, dataset_path):
 
         # Pre-allocate memory-mapped arrays
         self.states = np.memmap(self.states_mmap_path, dtype=np.float32, mode='w+', shape=(total_transitions, state_dim))
-        self.actions = np.memmap(self.actions_mmap_path, dtype=np.float32, mode='w+', shape=(total_transitions, action_dim))
+        # Use int64 for discrete actions
+        self.actions = np.memmap(self.actions_mmap_path, dtype=np.int64, mode='w+', shape=(total_transitions, action_dim))
         self.rewards = np.memmap(self.rewards_mmap_path, dtype=np.float32, mode='w+', shape=(total_transitions, 1))
         self.next_states = np.memmap(self.next_states_mmap_path, dtype=np.float32, mode='w+', shape=(total_transitions, state_dim))
         self.dones = np.memmap(self.dones_mmap_path, dtype=np.float32, mode='w+', shape=(total_transitions, 1))
@@ -98,9 +99,9 @@ def __init__(self, dataset_path):
             traj_rtg = data['returns_to_go'][i, :clip_len]
 
             # Actions
-            traj_actions = np.zeros((clip_len, action_dim), dtype=np.float32)
+            traj_actions = np.zeros((clip_len, action_dim), dtype=np.int64)
             if clip_len > 1:
-                traj_actions[:clip_len-1] = data['actions'][i, :clip_len-1]
+                traj_actions[:clip_len-1] = data['actions'][i, :clip_len-1].astype(np.int64)
 
             # Rewards and next_states
             rewards = np.zeros((clip_len, 1), dtype=np.float32)
@@ -216,12 +217,15 @@ def train(cfg, logger):
         epoch_losses = []
 
         batch_iter = tqdm(
-            enumerate(train_loader), 
-            total=len(train_loader),
+            enumerate(train_loader),
+            total=cfg.training.batches_per_epoch,
             desc=f"Epoch {epoch+1}/{cfg.training.epochs}"
         )
 
         for i, batch in batch_iter:
+            if i >= cfg.training.batches_per_epoch:
+                break
+            
             model.train()
 
             for k, v in batch.items():

snn-dt/src/models/iql.py

@@ -63,14 +63,23 @@ def get_det_action(self, state):
 
 
 class Critic(nn.Module):
-    def __init__(self, state_size, action_size, hidden_size=256):
+    def __init__(self, state_size, action_size, hidden_size=256, is_discrete=False):
         super(Critic, self).__init__()
-        self.fc1 = nn.Linear(state_size + action_size, hidden_size)
+        self.is_discrete = is_discrete
+        if self.is_discrete:
+            self.action_embedding = nn.Embedding(action_size, hidden_size)
+            self.fc1 = nn.Linear(state_size + hidden_size, hidden_size)
+        else:
+            self.fc1 = nn.Linear(state_size + action_size, hidden_size)
         self.fc2 = nn.Linear(hidden_size, hidden_size)
         self.fc3 = nn.Linear(hidden_size, 1)
 
     def forward(self, state, action):
-        x = torch.cat((state, action), dim=-1)
+        if self.is_discrete:
+            action_emb = self.action_embedding(action.long().squeeze(-1))
+            x = torch.cat((state, action_emb), dim=-1)
+        else:
+            x = torch.cat((state, action), dim=-1)
         x = F.relu(self.fc1(x))
         x = F.relu(self.fc2(x))
         return self.fc3(x)
@@ -106,12 +115,12 @@ def __init__(self, cfg):
         self.is_discrete = cfg.dataset.is_discrete
 
         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.critic1 = Critic(cfg.dataset.state_dim, cfg.dataset.act_dim, cfg.iql.hidden_size, is_discrete=self.is_discrete).to(self.device)
+        self.critic2 = Critic(cfg.dataset.state_dim, cfg.dataset.act_dim, cfg.iql.hidden_size, is_discrete=self.is_discrete).to(self.device)
         self.value_net = Value(cfg.dataset.state_dim, cfg.iql.hidden_size).to(self.device)
 
-        self.critic1_target = Critic(cfg.dataset.state_dim, cfg.dataset.act_dim, cfg.iql.hidden_size).to(self.device)
-        self.critic2_target = Critic(cfg.dataset.state_dim, cfg.dataset.act_dim, cfg.iql.hidden_size).to(self.device)
+        self.critic1_target = Critic(cfg.dataset.state_dim, cfg.dataset.act_dim, cfg.iql.hidden_size, is_discrete=self.is_discrete).to(self.device)
+        self.critic2_target = Critic(cfg.dataset.state_dim, cfg.dataset.act_dim, cfg.iql.hidden_size, is_discrete=self.is_discrete).to(self.device)
         self.critic1_target.load_state_dict(self.critic1.state_dict())
         self.critic2_target.load_state_dict(self.critic2.state_dict())
 
@@ -135,13 +144,8 @@ def learn(self, batch):
 
         # Value loss
         with torch.no_grad():
-            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)
+            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()
@@ -155,7 +159,10 @@ def learn(self, batch):
             exp_a = torch.exp((min_q - v) * self.temperature)
             exp_a = torch.min(exp_a, torch.tensor(100.0, device=self.device))
         _, dist = self.actor.evaluate(states)
-        log_probs = dist.log_prob(actions)
+        if self.is_discrete:
+            log_probs = dist.log_prob(actions.squeeze(-1).long())
+        else:
+            log_probs = dist.log_prob(actions)
         actor_loss = -(exp_a * log_probs).mean()
         self.actor_optimizer.zero_grad()
         actor_loss.backward()
@@ -166,13 +173,8 @@ def learn(self, batch):
             next_v = self.value_net(next_states)
             q_target = rewards + self.gamma * (1 - dones) * next_v
 
-        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)
+        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)