take first measure for straightening flow issue

README.md

@@ -74,4 +74,15 @@ sampled_actions = model(vision, commands, joint_state, trajectory_length = 32) #
 }
 ```
 
+```bibtex
+@article{Li2024ImmiscibleDA,
+    title   = {Immiscible Diffusion: Accelerating Diffusion Training with Noise Assignment},
+    author  = {Yiheng Li and Heyang Jiang and Akio Kodaira and Masayoshi Tomizuka and Kurt Keutzer and Chenfeng Xu},
+    journal = {ArXiv},
+    year    = {2024},
+    volume  = {abs/2406.12303},
+    url     = {https://api.semanticscholar.org/CorpusID:270562607}
+}
+```
+
 [*dear alice*](https://www.youtube.com/watch?v=z-Ng5ZvrDm4)

pi_zero_pytorch/pi_zero.py

@@ -12,6 +12,8 @@
 
 from torchdiffeq import odeint
 
+from scipy.optimize import linear_sum_assignment
+
 from rotary_embedding_torch import (
     RotaryEmbedding,
     apply_rotary_emb
@@ -117,6 +119,15 @@ def pad_at_dim(
     zeros = ((0, 0) * dims_from_right)
     return F.pad(t, (*zeros, *pad), value = value)
 
+# flow related
+
+def noise_assignment(data, noise):
+    device = data.device
+    data, noise = tuple(rearrange(t, 'b ... -> b (...)') for t in (data, noise))
+    dist = torch.cdist(data, noise)
+    _, assign = linear_sum_assignment(dist.cpu())
+    return torch.from_numpy(assign).to(device)
+
 # losses
 
 def direction_loss(pred, target, dim = -1):
@@ -413,6 +424,7 @@ def __init__(
         lm_loss_weight = 1.,
         flow_loss_weight = 1.,
         direction_loss_weight = 0.,
+        immiscible_flow = False, # https://arxiv.org/abs/2406.12303
         odeint_kwargs: dict = dict(
             atol = 1e-5,
             rtol = 1e-5,
@@ -495,6 +507,10 @@ def __init__(
 
         self.lm_pad_id = lm_pad_id
 
+        # flow related
+
+        self.immiscible_flow = immiscible_flow
+
         # loss related
 
         self.lm_loss_weight = lm_loss_weight
@@ -615,6 +631,10 @@ def forward(
         if not return_actions_flow:
             noise = torch.randn_like(actions)
 
+            if self.immiscible_flow:
+                assignment = noise_assignment(actions, noise)
+                noise = noise[assignment]
+
             flow = actions - noise
             padded_times = rearrange(times, 'b -> b 1 1')
 

pyproject.toml

@@ -30,6 +30,7 @@ dependencies = [
     "einops>=0.8.0",
     "jaxtyping",
     "rotary-embedding-torch>=0.8.5",
+    'scipy',
     "torch>=2.5",
     'torchdiffeq',
     "tqdm"