Action normalization

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,10 +17,13 @@ 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]
         self.action_dim = self.env.action_space.shape[0]
+        self.action_low = np.asarray(self.env.action_space.low, dtype=np.float32)
+        self.action_high = np.asarray(self.env.action_space.high, dtype=np.float32)
         self.obs_shape = self._get_image().shape
 
     def _extract_state(self, obs):
@@ -49,11 +52,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 +105,4 @@ def step(self, action):
         return image, state, reward, done, info
 
     def close(self):
-        self.env.close()
+        self.env.close()

models/action_head_utils.py

@@ -0,0 +1,37 @@
+"""Shared action-head building blocks."""
+
+import torch.nn as nn
+
+
+class ResidualMLPBlock(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, dim),
+            nn.SiLU(),
+            nn.Linear(dim, dim),
+        )
+
+    def forward(self, x):
+        return x + self.net(x)
+
+
+class ResidualActionMLP(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, hidden_dim: int, num_blocks: int = 3):
+        super().__init__()
+        self.input_proj = nn.Sequential(
+            nn.Linear(in_dim, hidden_dim),
+            nn.SiLU(),
+        )
+        self.blocks = nn.Sequential(*[ResidualMLPBlock(hidden_dim) for _ in range(num_blocks)])
+        self.out = nn.Sequential(
+            nn.LayerNorm(hidden_dim),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, out_dim),
+        )
+
+    def forward(self, x):
+        x = self.input_proj(x)
+        x = self.blocks(x)
+        return self.out(x)

models/diffusion_head.py

@@ -5,6 +5,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from .action_head_utils import ResidualActionMLP
 
 
 @dataclass
@@ -14,6 +15,7 @@ class DiffusionConfig:
     beta_end: float = 1e-2
     action_dim: int = 4
     cond_dim: int = 128 # conditional input dim
+    hidden_dim: int = 256
 
 
 def make_beta_schedule(cfg: DiffusionConfig):
@@ -62,19 +64,16 @@ class ActionDenoiseModel(nn.Module):
     t:     (B,) integer timestep
     cond:  (B, cond_dim) fused VLA token
     """
-    def __init__(self, cfg: DiffusionConfig, time_emb_dim=32, hidden_dim=128):
+    def __init__(self, cfg: DiffusionConfig, time_emb_dim=32):
         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),
+        self.net = ResidualActionMLP(
+            in_dim=in_dim,
+            out_dim=cfg.action_dim,
+            hidden_dim=cfg.hidden_dim,
         )
 
     def forward(self, x_t, t, cond):

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
+from .action_head_utils import ResidualActionMLP
+
+
+@dataclass
+class FlowMatchingConfig:
+    action_dim: int
+    cond_dim: int
+    t_embed_dim: int = 32
+    sample_steps: int = 32
+    hidden_dim: int = 256
+
+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):
+        super().__init__()
+        self.time_emb = SinusoidalTime(cfg.t_embed_dim)
+        in_dim = cfg.action_dim + cfg.t_embed_dim + cfg.cond_dim
+        self.net = ResidualActionMLP(
+            in_dim=in_dim,
+            out_dim=cfg.action_dim,
+            hidden_dim=cfg.hidden_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
+        )