smaller horizon, smaller lstm. 100x faster actor_critic. does it learn though?

src/envs/world_model_env.py

@@ -65,13 +65,12 @@ class WorldModelEnv:
 
         token = action.clone().detach() if isinstance(action, torch.Tensor) else torch.tensor(action, dtype=torch.long)
         token = token.reshape(-1, 1).to(self.device)  # (B, 1)
+        
 
         for k in range(num_passes):  # assumption that there is only one action token.
 
             outputs_wm = self.world_model(token, past_keys_values=self.keys_values_wm)
             output_sequence.append(outputs_wm.output_sequence)
-
-            # if outputs_wm.logits_rewards.shape[1] > 0:
             if k == 0:
                 reward = Categorical(logits=outputs_wm.logits_rewards).sample().float().cpu().numpy().reshape(-1) - 1   # (B,)
                 done = Categorical(logits=outputs_wm.logits_ends).sample().cpu().numpy().astype(bool).reshape(-1)       # (B,)

src/models/actor_critic.py

@@ -36,6 +36,7 @@ class ImagineOutput:
 class ActorCritic(nn.Module):
     def __init__(self, act_vocab_size, use_original_obs: bool = False) -> None:
         super().__init__()
+        shrink = 4
         self.use_original_obs = use_original_obs
         self.conv1 = nn.Conv2d(3, 32, 3, stride=1, padding=1)
         self.maxp1 = nn.MaxPool2d(2, 2)
@@ -43,15 +44,15 @@ class ActorCritic(nn.Module):
         self.maxp2 = nn.MaxPool2d(2, 2)
         self.conv3 = nn.Conv2d(32, 64, 3, stride=1, padding=1)
         self.maxp3 = nn.MaxPool2d(2, 2)
-        self.conv4 = nn.Conv2d(64, 64, 3, stride=1, padding=1)
+        self.conv4 = nn.Conv2d(64, 64//shrink, 3, stride=1, padding=1)
         self.maxp4 = nn.MaxPool2d(2, 2)
 
-        self.lstm_dim = 512
-        self.lstm = nn.LSTMCell(1024, self.lstm_dim)
+        self.lstm_dim = 64
+        self.lstm = nn.LSTMCell(1024//shrink, self.lstm_dim)
         self.hx, self.cx = None, None
 
-        self.critic_linear = nn.Linear(512, 1)
-        self.actor_linear = nn.Linear(512, act_vocab_size)
+        self.critic_linear = nn.Linear(self.lstm_dim, 1)
+        self.actor_linear = nn.Linear(self.lstm_dim, act_vocab_size)
 
     def __repr__(self) -> str:
         return "actor_critic"
@@ -85,7 +86,7 @@ class ActorCritic(nn.Module):
         x = F.relu(self.maxp2(self.conv2(x)))
         x = F.relu(self.maxp3(self.conv3(x)))
         x = F.relu(self.maxp4(self.conv4(x)))
-        x = torch.flatten(x, start_dim=1)
+        x = torch.flatten(x, start_dim=1) # [b=32, 64//shrink, 4, 4]
 
         if mask_padding is None:
             self.hx, self.cx = self.lstm(x, (self.hx, self.cx))

src/trainer.py

@@ -169,6 +169,7 @@ class Trainer:
             if max_grad_norm is not None:
                 torch.nn.utils.clip_grad_norm_(component.parameters(), max_grad_norm)
 
+ 
             optimizer.step()
 
         metrics = {f'{str(component)}/train/total_loss': loss_total_epoch, **intermediate_losses}