diff --git a/research_journal.md b/research_journal.md index eac579f..492e009 100644 --- a/research_journal.md +++ b/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 diff --git a/src/agent.py b/src/agent.py index 5534232..ee884fd 100644 --- a/src/agent.py +++ b/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 diff --git a/src/envs/world_model_env.py b/src/envs/world_model_env.py index 73e3852..4e447d4 100644 --- a/src/envs/world_model_env.py +++ b/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: diff --git a/src/models/actor_critic.py b/src/models/actor_critic.py index ec1256d..4d68c78 100644 --- a/src/models/actor_critic.py +++ b/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)