Flow matching

docs/FLOWMATCHING.md

@@ -0,0 +1,5 @@
+# Flow Matching
+
+Flow Matching (FM) is a generative modeling paradigm that learns a continuous velocity field that transforms samples drawn from a simple source distribution (e.g., Gaussian noise) into samples from the target distribution (e.g., robot actions). Instead of learning to denoise noisy samples (as in diffusion models), flow matching directly regresses the velocity at each point in time that moves a sample along a probability path.
+
+ref. https://arxiv.org/abs/2210.02747

envs/metaworld_env.py

@@ -17,6 +17,7 @@ def __init__(self, env_name='push-v3', seed=42, render_mode='rgb_array', camera_
             camera_name=camera_name
         )
         self.render_mode = render_mode
+        self.camera_name = camera_name
 
         obs, _ = self.env.reset()
         self.state_dim = self._extract_state(obs).shape[0]
@@ -49,11 +50,45 @@ def _extract_state(self, obs):
             state = obs
         return np.asarray(state, dtype=np.float32)
 
-    def _get_image(self):
-        img = self.env.render()
+    def _get_image(self, camera_name=None):
+        if camera_name is None or camera_name == self.camera_name:
+            img = self.env.render()
+        else:
+            try:
+                img = self.env.render(camera_name=camera_name)
+            except TypeError:
+                # Older wrappers may only support the camera selected at env creation.
+                img = self.env.render()
         img = img.astype(np.uint8)
         return img
 
+    def render(self, camera_name=None):
+        return self._get_image(camera_name=camera_name)
+
+    def _get_stateful_env(self):
+        for candidate in (self.env, getattr(self.env, "unwrapped", None)):
+            if candidate is not None:
+                has_get = hasattr(candidate, "get_env_state")
+                has_set = hasattr(candidate, "set_env_state")
+                if has_get and has_set:
+                    return candidate
+        return None
+
+    def get_env_state(self):
+        stateful_env = self._get_stateful_env()
+        if stateful_env is None:
+            raise AttributeError("Environment does not expose get_env_state/set_env_state")
+        return stateful_env.get_env_state()
+
+    def set_env_state(self, state):
+        stateful_env = self._get_stateful_env()
+        if stateful_env is None:
+            raise AttributeError("Environment does not expose get_env_state/set_env_state")
+        stateful_env.set_env_state(state)
+
+    def sync_from(self, other):
+        self.set_env_state(other.get_env_state())
+
     def reset(self, seed=None):
         obs, info = self.env.reset(seed=seed)
         state = self._extract_state(obs)
@@ -68,4 +103,4 @@ def step(self, action):
         return image, state, reward, done, info
 
     def close(self):
-        self.env.close()
+        self.env.close()

models/flow_matching_head.py

@@ -0,0 +1,96 @@
+"""Flow-matching policy head for action generation."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dataclasses import dataclass
+
+
+@dataclass
+class FlowMatchingConfig:
+    action_dim: int
+    cond_dim: int
+    t_embed_dim: int = 32
+    sample_steps: int = 32
+
+class SinusoidalTime(nn.Module):
+    """This is kept same as the diffusion time embedding."""
+
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, t):
+        half = self.dim // 2
+        freqs = torch.exp(torch.linspace(0, torch.log(torch.tensor(1000.0)), half, device=t.device))
+        args = t.unsqueeze(-1).float() * freqs.unsqueeze(0)
+        emb = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)
+        return emb
+
+class FlowMatchingModel(nn.Module):
+    """
+    Predict velocity field for flow matching: v(x_t, t | cond)
+    I LOVE THIS BLOG from Federico Sarrocco
+    https://federicosarrocco.com/blog/flow-matching
+    (hence, the code follows the same structure)
+    """
+
+    def __init__(self, cfg: FlowMatchingConfig, hidden_dim=128):
+        super().__init__()
+        self.time_emb = SinusoidalTime(cfg.t_embed_dim)
+        in_dim = cfg.action_dim + cfg.t_embed_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):
+        t_emb = self.time_emb(t)
+        x = torch.cat([x_t, t_emb, cond], dim=-1)
+        v_pred = self.net(x)
+        return v_pred
+
+class FlowMatchingPolicyHead(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.model = FlowMatchingModel(cfg)
+
+    def loss(self, actions, cond):
+        """
+        Conditional flow matching with a linear interpolation path.
+
+        We sample a source point x_0 ~ N(0, I), set x_1 to the demonstrated
+        action, and train the model to predict the constant velocity field
+        along x_t = (1 - t) * x_0 + t * x_1.
+        """
+        B = actions.size(0)
+        t = torch.rand(B, device=actions.device)
+
+        x_0 = torch.randn_like(actions)
+        t_expanded = t.unsqueeze(-1)
+        x_t = (1.0 - t_expanded) * x_0 + t_expanded * actions
+        target_v = actions - x_0
+
+        v_pred = self.model(x_t, t, cond)
+        return F.mse_loss(v_pred, target_v)
+
+    @torch.no_grad()
+    def sample(self, cond, n_samples=None):
+        B = cond.size(0) if n_samples is None else n_samples
+        if cond.size(0) != B:
+            cond = cond.expand(B, -1)
+
+        # Start at the source distribution and integrate forward to t = 1.
+        x_t = torch.randn(B, self.cfg.action_dim, device=cond.device)
+        dt = 1.0 / self.cfg.sample_steps
+
+        for step in range(self.cfg.sample_steps):
+            t = torch.full((B,), step * dt, device=cond.device)
+            v = self.model(x_t, t, cond)
+            x_t = x_t + v * dt
+
+        return x_t

models/vision/VISION.md

@@ -0,0 +1 @@
+TODO: write a comprehensive doc to reproduce training, and testing using CLI commands 

models/vision/__init__.py

@@ -0,0 +1,2 @@
+from .tinycnn import TinyCNNEncoder
+from .hf_vit import HFCLIPViT, HFSiglipViT

models/vision/hf_vit.py

@@ -0,0 +1,96 @@
+# models/vision/hf_vit.py
+from __future__ import annotations
+from typing import Literal
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from .registry import VisionEncoder, VisionEncoderCfg, register_vision_encoder
+
+
+class HFViTVisionEncoder(VisionEncoder):
+    """
+    HuggingFace Vision encoder wrapper (CLIP / SigLIP).
+
+    Input:  x (B,3,H,W) float in [0,1] (or uint8 in [0,255])
+    Output: (B,D) pooled CLS token, projected to cfg.d_model if needed
+    """
+    def __init__(
+        self,
+        cfg: VisionEncoderCfg,
+        hf_kind: Literal["clip", "siglip"],
+        default_pretrained: str,
+        mean: tuple[float, float, float],
+        std: tuple[float, float, float],
+    ):
+        super().__init__(cfg)
+
+        try:
+            if hf_kind == "clip":
+                from transformers import CLIPVisionModel
+                model_id = cfg.pretrained or default_pretrained
+                self.backbone = CLIPVisionModel.from_pretrained(model_id)
+            elif hf_kind == "siglip":
+                from transformers import SiglipVisionModel
+                model_id = cfg.pretrained or default_pretrained
+                self.backbone = SiglipVisionModel.from_pretrained(model_id)
+            else:
+                raise ValueError(f"Unsupported hf_kind={hf_kind}")
+        except ImportError as e:
+            raise ImportError("Install transformers: pip install transformers") from e
+
+        # project to cfg.d_model
+        hidden = int(self.backbone.config.hidden_size)
+        out_dim = int(cfg.d_model)
+        self.proj = nn.Identity() if out_dim == hidden else nn.Linear(hidden, out_dim)
+
+        # normalize
+        self.register_buffer("mean", torch.tensor(mean, dtype=torch.float32).view(1, 3, 1, 1), persistent=False)
+        self.register_buffer("std", torch.tensor(std, dtype=torch.float32).view(1, 3, 1, 1), persistent=False)
+
+        self.image_size = int(cfg.image_size or getattr(self.backbone.config, "image_size", 224))
+
+        # cfg has trainable, not freeze
+        if not cfg.trainable:
+            for p in self.backbone.parameters():
+                p.requires_grad = False
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.dtype == torch.uint8: x = x.float() / 255.0
+        else: x = x.float()
+
+        if x.shape[-2:] != (self.image_size, self.image_size):
+            x = F.interpolate(x, size=(self.image_size, self.image_size), mode="bilinear", align_corners=False)
+
+        x = (x - self.mean) / self.std # normalize
+        out = self.backbone(pixel_values=x, return_dict=True).last_hidden_state  # (B, N, hidden)
+        cls = out[:, 0]
+        return self.proj(cls)
+
+
+@register_vision_encoder("hf_clip_vit")
+class HFCLIPViT(HFViTVisionEncoder):
+    def __init__(self, cfg: VisionEncoderCfg):
+        super().__init__(
+            cfg=cfg,
+            hf_kind="clip",
+            default_pretrained="openai/clip-vit-base-patch32",
+            # source for the following mean/std -- default values from HuggingFace documentation
+            # https://hf.co/docs/transformers/main/model_doc/clip#transformers.CLIPImageProcessor.image_mean
+            mean=(0.48145466, 0.4578275, 0.40821073), # CLIP mean
+            std=(0.26862954, 0.26130258, 0.27577711), # CLIP std
+            # i know, they look weird, but trust me they work
+        )
+
+
+@register_vision_encoder("hf_siglip_vit")
+class HFSiglipViT(HFViTVisionEncoder):
+    def __init__(self, cfg: VisionEncoderCfg):
+        super().__init__(
+            cfg=cfg,
+            hf_kind="siglip",
+            default_pretrained="google/siglip-base-patch16-224",
+            mean=(0.5, 0.5, 0.5), # SigLIP mean
+            std=(0.5, 0.5, 0.5), # SigLIP std
+        )

models/vision/registry.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Dict, Type, Optional, List
+
+import torch
+import torch.nn as nn
+
+@dataclass
+class VisionEncoderCfg:
+    name : str = "tinycnn"
+    d_model : int = 128
+    pretrained : Optional[str] = None
+    trainable : bool = False
+    image_size : Optional[int] = None
+
+
+class VisionEncoder(nn.Module):
+    def __init__(self, cfg: VisionEncoderCfg):
+        super().__init__()
+        self.cfg = cfg
+
+    @property
+    def d_model(self) -> int:
+        return self.cfg.d_model
+
+_REGISTRY: Dict[str, Type[VisionEncoder]] = {}
+
+def register_vision_encoder(name: str):
+    name = name.lower()
+
+    def decorator(cls: Type[VisionEncoder]):
+        if name in _REGISTRY:
+            raise ValueError(f"Vision encoder '{name}' already registered by {_REGISTRY[name]}.")
+        _REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+def available_vision_encoders() -> List[str]:
+    return list(_REGISTRY.keys())
+
+def build_vision_encoder(cfg: VisionEncoderCfg) -> VisionEncoder:
+    name = cfg.name.lower()
+    if name not in _REGISTRY:
+        raise ValueError(f"Unknown vision encoder '{name}'. Available: {available_vision_encoders()}")
+    return _REGISTRY[name](cfg)

models/vision/tinycnn.py

@@ -0,0 +1,25 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .registry import VisionEncoder, VisionEncoderCfg, register_vision_encoder
+
+@register_vision_encoder("tinycnn")
+class TinyCNNEncoder(VisionEncoder):
+    def __init__(self, cfg: VisionEncoderCfg):
+        super().__init__(cfg)
+        d_model = cfg.d_model
+
+        self.conv1 = nn.Conv2d(3, 32, kernel_size=5, stride=2, padding=2)
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)
+        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
+        self.proj = nn.Linear(128, d_model)
+        self.ln = nn.LayerNorm(d_model)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = F.relu(self.conv3(x))
+        x = x.mean(dim=[2, 3])
+        x = self.proj(x)
+        x = self.ln(x)
+        return x