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

config/trainer.yaml

@@ -54,7 +54,7 @@ training:
   should: True
   learning_rate: 0.0001
   tokenizer:
-    batch_num_samples: 32
+    batch_num_samples: 128
     grad_acc_steps: 1
     max_grad_norm: 10.0
     start_after_epochs: 5
@@ -71,7 +71,7 @@ training:
     grad_acc_steps: 1
     max_grad_norm: 10.0
     start_after_epochs: 50
-    steps_per_epoch: 200
+    steps_per_epoch: 20
     imagine_horizon: ${common.sequence_length}
     burn_in: 20
     gamma: 0.995

config/world_model/default.yaml

@@ -1,6 +1,6 @@
 _target_: models.TransformerConfig
 tokens_per_block: 17
-max_blocks: 20
+max_blocks: 10
 attention: 'causal'
 num_layers: 10
 num_heads: 4

research_journal.md

@@ -163,3 +163,39 @@ hm maybe it's just the face it has to backprop throguh the whole LLM :( damn...
 well running eval on the transformer brought it down from 100sec to 60, but it's still huge. 
 
 But then why is the model training fast? It makes not sense
+ 
+# 2023-11-16 12:54:48
+
+Why is agent so slow? Lets find out
+- look at diagram
+- look at train_agent
+  - to tokenizer.compute_loss is just tokenizer
+  - world_model.compute_loss user tokenizer with no grad
+  - actor_critic? takes an hour!!
+    - imagine (with grad?)
+      - x20 = horizon
+        - self(obs)
+        - WorldModelEnv.step this has no grad!
+          - transformer
+          - tokenizer with no grad
+    - compute_lambda_returns with no grad
+
+
+So changes:
+- the world model step always had no grad!
+- I just made the lstm smaller and the horizon smaller
+- from 1h to 3m. Reasonable.
+
+
+Experiment: 
+- try no grad on the model? ok it now takes 20 minutes to train... still slow
+with a smaller lstm and only 10 steos ut tajes 8 mins,
+
+![](img/2023-11-16-13-01-11.png)
+
+Experience collection (train_dataset): 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:05<00:00, 36.11it/s]
+Training tokenizer: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:55<00:00,  3.61it/s]
+Training world_model: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [01:12<00:00,  2.76it/s]
+Training tokenizer: 55sec
+Training world_model 72 sec
+train actor_critic 3min. It looks like it scales with lstm size!

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}