exponential moving averaging model is pretty indispensible these days

pi_zero_pytorch/pi_zero.py

@@ -17,6 +17,8 @@
 
 from scipy.optimize import linear_sum_assignment
 
+from ema_pytorch import EMA
+
 from rotary_embedding_torch import (
     RotaryEmbedding,
     apply_rotary_emb
@@ -587,6 +589,24 @@ def load_pretrained_vlm_weights_(
     ):
         raise NotImplementedError
 
+    def create_ema(
+        self,
+        beta = 0.99,
+        **ema_kwargs
+    ) -> EMA:
+
+        ema_pi_zero = EMA(
+            self,
+            beta = beta,
+            include_online_model = False,
+            forward_method_names = (
+                'sample_actions',
+            ),
+            **ema_kwargs
+        )
+
+        return ema_pi_zero
+
     @torch.inference_mode()
     def sample_actions(
         self,

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "pi-zero-pytorch"
-version = "0.0.24"
+version = "0.0.25"
 description = "π0 in Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }
@@ -28,6 +28,7 @@ dependencies = [
     "beartype",
     "einx>=0.3.0",
     "einops>=0.8.0",
+    "ema-pytorch>=0.7.3",
     "jaxtyping",
     "rotary-embedding-torch>=0.8.5",
     'scipy',