From ecee11fc33834c8785d911eaf5b7a572bf57dd1a Mon Sep 17 00:00:00 2001 From: wassname Date: Sun, 19 Nov 2023 10:58:45 +0800 Subject: [PATCH] notes --- src/models/tokenizer/tokenizer.py | 2 ++ src/models/world_model.py | 1 + 2 files changed, 3 insertions(+) diff --git a/src/models/tokenizer/tokenizer.py b/src/models/tokenizer/tokenizer.py index 2504b42..a96ec82 100644 --- a/src/models/tokenizer/tokenizer.py +++ b/src/models/tokenizer/tokenizer.py @@ -46,6 +46,8 @@ class Tokenizer(nn.Module): def compute_loss(self, batch: Batch, **kwargs: Any) -> LossWithIntermediateLosses: assert self.lpips is not None observations = self.preprocess_input(rearrange(batch['observations'], 'b t c h w -> (b t) c h w')) + # TODO: in the delta-IRIS paper (https://openreview.net/forum?id=o8IDoZggqO) they encode(x0, a0, x1) -> z1 and decode(x0, a0, z1). In esense the tokens only need to encode the change + # note they also do dynamics(x0, a0, z1) -> z2. decode(x1, a1, z2) -> x2 z, z_quantized, reconstructions = self(observations, should_preprocess=False, should_postprocess=False) # Codebook loss. Notes: diff --git a/src/models/world_model.py b/src/models/world_model.py index 9999ff4..b346094 100644 --- a/src/models/world_model.py +++ b/src/models/world_model.py @@ -47,6 +47,7 @@ class WorldModel(nn.Module): ) # why have this? Well I worry that the transformer can't adapt, since so much is frozen + # TODO: If I get the dynamics model working, maybe try without it self.post_embed = nn.Sequential( nn.Linear(config.embed_dim, config.embed_dim), nn.ReLU(),