Adding registry for action expert and adds flow matching head

.gitignore

@@ -3,4 +3,4 @@ __pycache__/
 MUJOCO_LOG.txt
 LIBERO/
 videos*/
-checkpoints/
+checkpoints*/

README.md

@@ -22,6 +22,7 @@ This project intentionally keeps the codebase small (~150 LOC for the core model
 > I recommend reading the following blogs to get started with mini-VLA implementation.
 > - [Building Vision-Language-Action Model from scratch (Basics)](https://open.substack.com/pub/keivalya/p/building-vision-language-action-from?utm_campaign=post-expanded-share&utm_medium=web)
 > - [Building VLA models from scratch — II (Math, Code, and Intuition)](https://medium.com/@keivalyap/building-vla-models-from-scratch-ii-0180020dbc85)
+> - [Upgrading mini-VLA with CLIP/SigLIP vision encoders](https://medium.com/@keivalyap/mini-vla-with-vision-encoders-f9ba8d8d2988)
 
 This project is not meant to be state-of-the-art instead, it provides a clear, hackable template for understanding VLA design.
 

models/action_expert/__init__.py

@@ -0,0 +1,6 @@
+from .registry import ActionExpertCfg, ActionExpert, build_action_expert, register_action_expert, available_action_experts
+
+# Import registered action experts to trigger registration
+import models.diffusion_head  # registers "diffusion"
+import models.flow_matching_head  # registers "flow_matching"
+

models/action_expert/registry.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Dict, Type, Optional, List
+
+import torch
+import torch.nn as nn
+
+@dataclass
+class ActionExpertCfg:
+    name: str = "diffusion"
+    action_dim: int = 4
+    cond_dim: int = 128
+    T: int = 16
+    # Diffusion-specific params
+    beta_start: float = 1e-4
+    beta_end: float = 1e-2
+    # Flow matching doesn't need additional params beyond T, action_dim, cond_dim
+
+
+class ActionExpert(nn.Module):
+    """
+    Base class for action generation heads (diffusion, flow matching, etc.)
+    All action experts must implement:
+    - loss(actions, cond) -> loss tensor
+    - sample(cond, n_samples=None) -> actions tensor
+    """
+    def __init__(self, cfg: ActionExpertCfg):
+        super().__init__()
+        self.cfg = cfg
+
+    def loss(self, actions: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        """
+        actions: (B, action_dim)  ground-truth actions
+        cond:    (B, cond_dim)    fused VLA token
+        returns: loss tensor
+        """
+        raise NotImplementedError
+
+    @torch.no_grad()
+    def sample(self, cond: torch.Tensor, n_samples: Optional[int] = None) -> torch.Tensor:
+        """
+        cond: (B, cond_dim) or (1, cond_dim)
+        n_samples: optional number of samples to generate
+        returns: (B, action_dim) sampled actions
+        """
+        raise NotImplementedError
+
+
+_REGISTRY: Dict[str, Type[ActionExpert]] = {}
+
+def register_action_expert(name: str):
+    name = name.lower()
+
+    def decorator(cls: Type[ActionExpert]):
+        if name in _REGISTRY:
+            raise ValueError(f"Action expert '{name}' already registered by {_REGISTRY[name]}.")
+        _REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+def available_action_experts() -> List[str]:
+    return list(_REGISTRY.keys())
+
+def build_action_expert(cfg: ActionExpertCfg) -> ActionExpert:
+    name = cfg.name.lower()
+    if name not in _REGISTRY:
+        raise ValueError(f"Unknown action expert '{name}'. Available: {available_action_experts()}")
+    return _REGISTRY[name](cfg)
+

models/diffusion_head.py

@@ -6,6 +6,8 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
+from .action_expert.registry import ActionExpert, ActionExpertCfg, register_action_expert
+
 
 @dataclass
 class DiffusionConfig:
@@ -88,13 +90,22 @@ def forward(self, x_t, t, cond):
         eps_pred = self.net(x)
         return eps_pred
 
-class DiffusionPolicyHead(nn.Module):
-    def __init__(self, cfg: DiffusionConfig):
-        super().__init__()
-        self.cfg = cfg
-        self.denoise_model = ActionDenoiseModel(cfg)
-        betas, alphas, alpha_bar = make_beta_schedule(cfg)
-        # register as buffers so they move with the module’s device
+@register_action_expert("diffusion")
+class DiffusionPolicyHead(ActionExpert):
+    def __init__(self, cfg: ActionExpertCfg):
+        super().__init__(cfg)
+        # Convert ActionExpertCfg to DiffusionConfig for internal use
+        diffusion_cfg = DiffusionConfig(
+            T=cfg.T,
+            beta_start=cfg.beta_start,
+            beta_end=cfg.beta_end,
+            action_dim=cfg.action_dim,
+            cond_dim=cfg.cond_dim,
+        )
+        self.diffusion_cfg = diffusion_cfg
+        self.denoise_model = ActionDenoiseModel(diffusion_cfg)
+        betas, alphas, alpha_bar = make_beta_schedule(diffusion_cfg)
+        # register as buffers so they move with the module's device
         self.register_buffer("betas", betas)
         self.register_buffer("alphas", alphas)
         self.register_buffer("alpha_bar", alpha_bar)
@@ -116,7 +127,7 @@ def loss(self, actions, cond):
         """
         B = actions.size(0)
         device = actions.device
-        t = torch.randint(0, self.cfg.T, (B,), device=device) # uniform sampling t
+        t = torch.randint(0, self.diffusion_cfg.T, (B,), device=device) # uniform sampling t
         noise = torch.randn_like(actions)
         x_t = self.q_sample(actions, t, noise)  # noisy actions
         eps_pred = self.denoise_model(x_t, t, cond)
@@ -135,8 +146,8 @@ def sample(self, cond, n_samples=None):
             B = n_samples
             cond = cond.expand(B, -1)
 
-        x_t = torch.randn(B, self.cfg.action_dim, device=cond.device)
-        for t_step in reversed(range(self.cfg.T)):
+        x_t = torch.randn(B, self.diffusion_cfg.action_dim, device=cond.device)
+        for t_step in reversed(range(self.diffusion_cfg.T)):
             t = torch.full((B,), t_step, device=cond.device, dtype=torch.long)
             eps_pred = self.denoise_model(x_t, t, cond)
             beta_t = self.betas[t_step]

models/flow_matching_head.py

@@ -0,0 +1,139 @@
+"""Flow matching policy head for action generation."""
+
+import math
+from dataclasses import dataclass
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .action_expert.registry import ActionExpert, ActionExpertCfg, register_action_expert
+
+
+@dataclass
+class FlowMatchingConfig:
+    T: int = 50  # number of time steps
+    action_dim: int = 4
+    cond_dim: int = 128 # conditional input dim
+
+
+class SinusoidalTimeEmbedding(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, t: torch.Tensor):
+        """
+        t: (B,) continuous time values in [0, 1]
+        returns: (B, dim)
+        """
+        half_dim = self.dim // 2
+        device = t.device
+        freqs = torch.exp(
+            torch.linspace(
+                math.log(1.0),
+                math.log(1000.0),
+                half_dim,
+                device=device
+            )
+        )
+        # Scale continuous time [0, 1] to [0, 1000] to match frequency range
+        t_scaled = t.float() * 1000.0
+        # (B, half_dim)
+        args = t_scaled.unsqueeze(-1) * freqs.unsqueeze(0)
+        emb = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)
+        if self.dim % 2 == 1:
+            emb = torch.cat([emb, torch.zeros_like(emb[..., :1])], dim=-1)
+        return emb
+
+class FlowMatchingModel(nn.Module):
+    """
+    epsilon_theta(x_t, t, cond)
+    x_t:   (B, action_dim)
+    t:     (B,) continuous time in [0, 1]
+    cond:  (B, cond_dim) fused VLA token
+    """
+    def __init__(self, cfg: FlowMatchingConfig, time_emb_dim=32, hidden_dim=128):
+        super().__init__()
+        self.cfg = cfg
+        self.time_emb = SinusoidalTimeEmbedding(time_emb_dim)
+
+        in_dim = cfg.action_dim + time_emb_dim + cfg.cond_dim
+
+        self.net = nn.Sequential(
+            nn.Linear(in_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, cfg.action_dim),
+        )
+
+    def forward(self, x_t, t, cond):
+        """
+        x_t: (B, action_dim)
+        t:   (B,) continuous time in [0, 1]
+        cond: (B, cond_dim)
+        """
+        t_emb = self.time_emb(t)  # (B, time_emb_dim)
+        x = torch.cat([x_t, t_emb, cond], dim=-1)
+        v_pred = self.net(x)
+        return v_pred
+
+@register_action_expert("flow_matching")
+class FlowMatchingPolicyHead(ActionExpert):
+    def __init__(self, cfg: ActionExpertCfg):
+        super().__init__(cfg)
+        # Convert ActionExpertCfg to FlowMatchingConfig for internal use
+        flow_cfg = FlowMatchingConfig(
+            T=cfg.T,
+            action_dim=cfg.action_dim,
+            cond_dim=cfg.cond_dim,
+        )
+        self.flow_cfg = flow_cfg
+        self.flow_matching_model = FlowMatchingModel(flow_cfg)
+
+    def q_sample(self, x_1, t, x_0):
+        """
+        Flow matching interpolation: x_t = (1 - t) * x_0 + t * x_1
+        x_0: (B, action_dim)  source (noise)
+        x_1: (B, action_dim)  target (actions)
+        t:   (B,) continuous in [0, 1]
+        """
+        # unsqueeze t to (B, 1) for broadcasting with (B, action_dim)
+        t = t.unsqueeze(-1)  # (B, 1)
+        return (1 - t) * x_0 + t * x_1
+
+    def loss(self, actions, cond):
+        """
+        actions: (B, action_dim)  ground-truth actions
+        cond:    (B, cond_dim)    fused VLA token
+        """
+        B = actions.size(0)
+        device = actions.device
+        t = torch.rand((B,), device=device) # uniform sampling t
+        x_0 = torch.randn_like(actions)
+        v_target = actions - x_0
+        x_t = self.q_sample(actions, t, x_0)  # noisy actions
+        v_pred = self.flow_matching_model(x_t, t, cond)
+        return F.mse_loss(v_pred, v_target)
+
+    @torch.no_grad()
+    def sample(self, cond, n_samples=None):
+        """
+        cond: (B, cond_dim) or (1, cond_dim)
+        returns: (B, action_dim) sampled actions x_0
+        """
+        self.eval()
+        if n_samples is None:
+            B = cond.size(0)
+        else:
+            B = n_samples
+            cond = cond.expand(B, -1)
+
+        x_t = torch.randn(B, self.flow_cfg.action_dim, device=cond.device)
+        dt = 1.0 / self.flow_cfg.T
+        for t_step in range(self.flow_cfg.T):
+            t = torch.full((B,), t_step * dt, device=cond.device, dtype=torch.float32)
+            v_pred = self.flow_matching_model(x_t, t, cond)
+
+            x_t = x_t + v_pred * dt
+        return x_t

models/vla_diffusion_policy.py

@@ -1,26 +1,48 @@
 """VLA Diffusion Policy Model."""
 
+from __future__ import annotations
+
 import torch.nn as nn
+from typing import Optional
 from .encoders import ImageEncoderTinyCNN, TextEncoderTinyGRU, StateEncoderMLP
 from .fusion import FusionMLP
-from .diffusion_head import DiffusionConfig, DiffusionPolicyHead
+
+from .action_expert.registry import ActionExpertCfg, build_action_expert
+import models.action_expert  # Import to trigger registration of action experts
 
 
 class VLADiffusionPolicy(nn.Module):
-    def __init__(self, vocab_size, state_dim, action_dim,
-                 d_model=128, diffusion_T=16):
+    def __init__(
+        self,
+        vocab_size,
+        state_dim,
+        action_dim,
+        d_model=128,
+        diffusion_T=16,
+        action_expert_cfg: Optional[ActionExpertCfg] = None,
+    ):
         super().__init__()
         self.img_encoder = ImageEncoderTinyCNN(d_model=d_model)
         self.txt_encoder = TextEncoderTinyGRU(vocab_size=vocab_size, d_word=64, d_model=d_model)
         self.state_encoder = StateEncoderMLP(state_dim=state_dim, d_model=d_model)
         self.fusion = FusionMLP(d_model=d_model)
 
-        cfg = DiffusionConfig(
-            T=diffusion_T,
-            action_dim=action_dim,
-            cond_dim=d_model,
-        )
-        self.diffusion_head = DiffusionPolicyHead(cfg)
+
+        if action_expert_cfg is None:
+            action_expert_cfg = ActionExpertCfg(
+                name="diffusion",
+                action_dim=action_dim,
+                cond_dim=d_model,
+                T=diffusion_T,
+            )
+        else:
+            # Ensure action expert config matches provided parameters
+            action_expert_cfg.cond_dim = d_model
+            action_expert_cfg.action_dim = action_dim
+            action_expert_cfg.T = diffusion_T
+
+        self.action_expert_cfg = action_expert_cfg
+        self.action_expert = build_action_expert(action_expert_cfg)
 
     def encode_obs(self, image, text_tokens, state):
         img_token = self.img_encoder(image)  # (B, d_model)
@@ -31,10 +53,10 @@ def encode_obs(self, image, text_tokens, state):
 
     def loss(self, image, text_tokens, state, actions):
         """
-        Compute the loss of the diffusion policy head given the image, text tokens, state, and actions.
+        Compute the loss of the action expert given the image, text tokens, state, and actions.
         """
         cond = self.encode_obs(image, text_tokens, state)
-        return self.diffusion_head.loss(actions, cond)
+        return self.action_expert.loss(actions, cond)
 
     def act(self, image, text_tokens, state):
         """
@@ -44,5 +66,4 @@ def act(self, image, text_tokens, state):
         returns: (B, action_dim)
         """
         cond = self.encode_obs(image, text_tokens, state)
-        actions = self.diffusion_head.sample(cond)
-        return actions
+        return self.action_expert.sample(cond)