diff --git a/justfile b/justfile
index 1f26170..5f3c137 100644
--- a/justfile
+++ b/justfile
@@ -1,3 +1,4 @@
+set shell := ["zsh", "-cu"]
 
 breakout:
     python src/main.py env.train.id=BreakoutNoFrameskip-v4 
@@ -8,3 +9,10 @@ crafter:
 minihack:
     python src/main.py env.train.id=MiniHack-River-v0
 
+# watch the latest runs
+watch_latest:
+    . ./.venv/bin/activate
+    cd ./outputs && \
+    cd *([-1]) && \
+    cd *([-1]) && \
+    scripts/play.sh -e -r -h
diff --git a/src/models/world_model.py b/src/models/world_model.py
index b346094..42bbf81 100644
--- a/src/models/world_model.py
+++ b/src/models/world_model.py
@@ -40,6 +40,7 @@ class WorldModel(nn.Module):
         self.pos_emb = nn.Embedding(config.max_tokens, config.embed_dim)
         self.act_emb = nn.Embedding(act_vocab_size, config.embed_dim)
 
+        # FIXME: having slices is unclear. maybe it's better just to have obs and action embeddings?
         self.embedder = Embedder(
             max_blocks=config.max_blocks,
             block_masks=[act_tokens_pattern, obs_tokens_pattern],
@@ -101,7 +102,7 @@ class WorldModel(nn.Module):
 
     def forward(self, tokens: torch.LongTensor, past_keys_values: Optional[KeysValues] = None) -> WorldModelOutput:
 
-        num_steps = tokens.size(1)  # (B, T)
+        num_steps = tokens.size(1)  # (B=8, T=170) where often the last 10 are actons
         assert num_steps <= self.config.max_tokens
         prev_steps = 0 if past_keys_values is None else past_keys_values.size
 
@@ -118,12 +119,13 @@ class WorldModel(nn.Module):
 
     def compute_loss(self, batch: Batch, tokenizer: Tokenizer, **kwargs: Any) -> LossWithIntermediateLosses:
 
-        with torch.no_grad():
-            # [B=8, S=10, Colors=3, H=64, W=64] -> [B=8, S=10, 16]
-            obs_tokens = tokenizer.encode(batch['observations'], should_preprocess=True).tokens  # (BL, K)
+        # with torch.no_grad():
+        # [B=8, S=10, Colors=3, H=64, W=64] -> [B=8, S=10, 16]
+        obs_tokens = tokenizer.encode(batch['observations'], should_preprocess=True).tokens  # (BL, K)
 
         act_tokens = rearrange(batch['actions'], 'b l -> b l 1')
         tokens = rearrange(torch.cat((obs_tokens, act_tokens), dim=2), 'b l k1 -> b (l k1)')  # (B, L(K+1))
+        # So first 10 are observation, the last 10 tokens are actions
 
         outputs = self(tokens)
 
diff --git a/src/trainer.py b/src/trainer.py
index 6c3b8f0..30d0c83 100644
--- a/src/trainer.py
+++ b/src/trainer.py
@@ -97,7 +97,11 @@ class Trainer:
         print(f'{sum(p.numel() for p in self.agent.actor_critic.parameters())} parameters in agent.actor_critic')
 
         self.optimizer_tokenizer = torch.optim.Adam(self.agent.tokenizer.parameters(), lr=cfg.training.learning_rate)
-        self.optimizer_world_model = configure_optimizer(self.agent.world_model, cfg.training.learning_rate, cfg.training.world_model.weight_decay)
+        # self.optimizer_world_model = configure_optimizer([self.agent.tokenizer, self.agent.world_model], cfg.training.learning_rate, cfg.training.world_model.weight_decay)
+        self.optimizer_world_model = torch.optim.Adam(
+            list(self.agent.tokenizer.parameters())+list(self.agent.world_model.parameters()),
+            lr=cfg.training.learning_rate
+        )
         self.optimizer_actor_critic = torch.optim.Adam(self.agent.actor_critic.parameters(), lr=cfg.training.learning_rate)
 
         if cfg.initialization.path_to_checkpoint is not None:
diff --git a/src/utils.py b/src/utils.py
index 07ba101..0f09590 100644
--- a/src/utils.py
+++ b/src/utils.py
@@ -13,41 +13,54 @@ from src.episode import Episode
 from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
 
 
-def configure_optimizer(model, learning_rate, weight_decay, *blacklist_module_names):
+def configure_optimizer(models, learning_rate, weight_decay, *blacklist_module_names):
     """Credits to https://github.com/karpathy/minGPT"""
     # FIXME: check this is still good for LoRA
     # separate out all parameters to those that will and won't experience regularizing weight decay
     decay = set()
     no_decay = set()
+    decay_params = []
+    no_decay_params = []
+    param_dict = {}
     whitelist_weight_modules = (torch.nn.Linear, torch.nn.Conv1d)
     blacklist_weight_modules = tuple(ALL_LAYERNORM_LAYERS+[torch.nn.Embedding])
-    for mn, m in model.named_modules():
-        for pn, p in m.named_parameters():
-            fpn = '%s.%s' % (mn, pn) if mn else pn  # full param name
-            if any([fpn.startswith(module_name) for module_name in blacklist_module_names]):
-                no_decay.add(fpn)
-            elif 'bias' in pn:
-                # all biases will not be decayed
-                no_decay.add(fpn)
-            elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules):
-                # weights of whitelist modules will be weight decayed
-                decay.add(fpn)
-            elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules):
-                # weights of blacklist modules will NOT be weight decayed
-                no_decay.add(fpn)
+    for model in models:
+        for mn, m in model.named_modules():
+            for pn, p in m.named_parameters():
+                fpn = '%s.%s' % (mn, pn) if mn else pn  # full param name
+                if any([fpn.startswith(module_name) for module_name in blacklist_module_names]):
+                    no_decay.add(fpn)
+                    no_decay_params.append(p)
+                elif 'bias' in pn:
+                    # all biases will not be decayed
+                    no_decay.add(fpn)
+                    no_decay_params.append(p)
+                elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules):
+                    # weights of whitelist modules will be weight decayed
+                    decay.add(fpn)
+                    decay_params.append(p)
+                elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules):
+                    # weights of blacklist modules will NOT be weight decayed
+                    no_decay.add(fpn)
+                    no_decay_params.append(p)
+                else:
+                    logger.warning(f"Parameter {fpn} of module {mn} not handled!")
+                    # raise NotImplementedError(f"Parameter {fpn} of module {m} not handled!")
+                    decay.add(fpn)
+                    decay_params.append(p)
 
-    # validate that we considered every parameter
-    param_dict = {pn: p for pn, p in model.named_parameters()}
+        # validate that we considered every parameter
+        param_dict.update({pn: p for pn, p in model.named_parameters()})
     inter_params = decay & no_decay
     union_params = decay | no_decay
-    logger.debug(f"decay {decay} no_decay {no_decay}")
+    # logger.debug(f"decay {decay} no_decay {no_decay}")
     assert len(inter_params) == 0, f"parameters {str(inter_params)} made it into both decay/no_decay sets!"
     assert len(param_dict.keys() - union_params) == 0, f"parameters {str(param_dict.keys() - union_params)} were not separated into either decay/no_decay set!"
 
     # create the pytorch optimizer object
     optim_groups = [
-        {"params": [param_dict[pn] for pn in sorted(list(decay))], "weight_decay": weight_decay},
-        {"params": [param_dict[pn] for pn in sorted(list(no_decay))], "weight_decay": 0.0},
+        {"params": no_decay_params, "weight_decay": weight_decay},
+        {"params": decay_params, "weight_decay": 0.0},
     ]
     optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate)
     return optimizer