Update delta-iris/src/world_model.py

delta-iris/src/world_model.py

@@ -54,6 +54,12 @@ class WorldModel(nn.Module):
     def __repr__(self) -> str:
         return "world_model"
 
+    def blocks_left_in_kv_cache(self):
+        return self.transformer.num_blocks_left_in_kv_cache
+
+    def reset_kv_cache(self):
+        transformer.reset_kv_cache(n=1)
+
     def forward(self, sequence: torch.FloatTensor, use_kv_cache: bool = False) -> dict:      
         prev_steps = self.transformer.keys_values.size if use_kv_cache else 0
         num_steps = sequence.size(1)