address a point of confusion, during inference of actions, cached sta…

…te keys / values stay the same
lucidrains · Dec 31, 2024 · a771a4c · a771a4c
1 parent d7c20e3
commit a771a4c
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/pi_zero_pytorch/pi_zero.py b/pi_zero_pytorch/pi_zero.py
@@ -1218,8 +1218,6 @@ def forward(
 
         cached_state_key_values_iter = iter(default(cached_state_keys_values, []))
 
-        state_cached_keys_values = []
-
         # value residual learning
 
         actions_value_residual = None
@@ -1231,6 +1229,9 @@ def forward(
         # transformer
 
         if not inferencing:
+
+            next_state_cached_keys_values = []
+
             for (
                 (attn, state_ff, actions_ff, memories_ff),
                 (attn_ada_rmsnorm, attn_ada_layerscale, ff_ada_rmsnorm, ff_ada_layerscale),
@@ -1254,7 +1255,7 @@ def forward(
                     memories = memory_tokens
                 )
 
-                state_cached_keys_values.append((state_keys, state_values))
+                next_state_cached_keys_values.append((state_keys, state_values))
 
                 actions_value_residual = default(actions_value_residual, action_values)
 
@@ -1298,6 +1299,10 @@ def forward(
 
         else:
 
+            assert exists(cached_state_keys_values) and len(cached_state_keys_values) > 0
+
+            next_state_cached_keys_values = cached_state_keys_values
+
             for (
                 (attn, state_ff, actions_ff, memories_ff),
                 (attn_ada_rmsnorm, attn_ada_layerscale, ff_ada_rmsnorm, ff_ada_layerscale),
@@ -1318,8 +1323,6 @@ def forward(
                     return_keys_values = True
                 )
 
-                state_cached_keys_values.append((state_keys, state_values))
-
                 actions_value_residual = default(actions_value_residual, action_values)
 
                 actions_attn_out = attn_ada_layerscale(actions_attn_out, time_cond)
@@ -1392,7 +1395,7 @@ def forward(
             if not return_state_keys_values:
                 return pred_actions_flow, written_memory_tokens
 
-            return pred_actions_flow, state_cached_keys_values
+            return pred_actions_flow, next_state_cached_keys_values
 
         flow_loss = self.zero
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "pi-zero-pytorch"
-version = "0.1.4"
+version = "0.1.5"
 description = "π0 in Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }