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

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)