it works better but still flickers when imagining. hmm I do think I need the delta-iris stuff

research_journal.md

@@ -358,3 +358,33 @@ OK it seems slightly better yay! Lets train it overnight and see
 
 next idea is to the delta-IRIS thing where the tokens only have to encode the diff(obs)
 
+# 2023-11-19 16:50:45 
+
+Seems to be working! Now let's plan delta-IRIS
+
+
+So IRIS has
+- Encoder $E(x_0, a_0) = t_0$
+  ```py
+  obs_tokens = self.tokenizer.encode(observations, should_preprocess=True).tokens    # (B, C, H, W) -> (B, K)
+  ```
+- Embed $Emb(t_0) = z_0$
+    ```py
+    embedded_tokens = self.tokenizer.embedding(self.obs_tokens)     # (B, K, E)
+    z = rearrange(embedded_tokens, 'b (h w) e -> b e h w', h=int(np.sqrt(self.num_observations_tokens)))
+    ```
+- Dynamics $D(z_0, a_0) = z_1$
+    ```py
+    outputs_wm = self.world_model(token, past_keys_values=self.keys_values_wm)
+    ```
+- Decoder $D(z_0, a_0) = x_1$
+    ```py
+    rec = self.tokenizer.decode(z, should_postprocess=True)         # (B, C, H, W)
+    ```
+
+
+but we have tokens vs z
+
+Questions:
+- wait why are we just passing in "action_token" to the transformer and not obs? that must have obs in it right... right??? confirm
+- in iris-delta how did they pass everything in? I guess obs_prev was tokenized too? I think the slices are annoying so maybe I should just pass things seperatly

src/agent.py

@@ -34,4 +34,5 @@ class Agent(nn.Module):
         input_ac = obs if self.actor_critic.use_original_obs else torch.clamp(self.tokenizer.encode_decode(obs, should_preprocess=True, should_postprocess=True), 0, 1)
         logits_actions = self.actor_critic(input_ac).logits_actions[:, -1] / temperature
         act_token = Categorical(logits=logits_actions).sample() if should_sample else logits_actions.argmax(dim=-1)
+        # FIXME, is this really just an action and doesn't have an obs in?
         return act_token

src/envs/world_model_env.py

@@ -69,6 +69,7 @@ class WorldModelEnv:
 
         for k in range(num_passes):  # assumption that there is only one action token.
 
+            # FIXME: hold on we are ONLY passing in the action token! should it not be obs too
             outputs_wm = self.world_model(token, past_keys_values=self.keys_values_wm)
             output_sequence.append(outputs_wm.output_sequence)
             if k == 0:

src/models/actor_critic.py

@@ -149,9 +149,7 @@ class ActorCritic(nn.Module):
             outputs_ac = self(obs)
             action_token = Categorical(logits=outputs_ac.logits_actions).sample()
             
-            # TODO this is really slow, I guess we need grad? does it help to put it in eval? no
-            # wm_env.world_model.eval()
-            
+            # FIXME shouldn't we pass in obs too?            
             obs, reward, done, _ = wm_env.step(action_token, should_predict_next_obs=(k < horizon - 1))
 
             all_actions.append(action_token)