Update Wall-X to 1.1.0 (#104)

wall_x/model/core/action/head.py

@@ -0,0 +1,27 @@
+"""Action head helpers used by the VLA action processor."""
+
+import math
+
+import torch
+import torch.nn as nn
+
+
+class SinusoidalPosEmb(nn.Module):
+    """Sinusoidal timestep embedding for action flow timesteps."""
+
+    def __init__(self, dim: int, min_period: float = 4e-3, max_period: float = 4.0):
+        super().__init__()
+        if dim % 2 != 0:
+            raise ValueError(f"embedding_dim ({dim}) must be divisible by 2")
+        self.dim = dim
+        self.min_period = min_period
+        self.max_period = max_period
+
+    def forward(self, x):
+        half_dim = self.dim // 2
+        exponent = math.log(10000) / (half_dim - 1)
+        frequencies = torch.exp(
+            torch.arange(half_dim, device=x.device, dtype=torch.float32) * -exponent
+        )
+        emb = x[:, None] * frequencies[None, :]
+        return torch.cat((emb.sin(), emb.cos()), dim=-1)