From 297dc4ced1e1d07e7c514bb786798acf8374f771 Mon Sep 17 00:00:00 2001 From: keivalya Date: Sun, 21 Dec 2025 20:17:08 -0500 Subject: [PATCH 1/2] Added vision blog --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index f5ab785..da025d0 100644 --- a/README.md +++ b/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. From 0b4fad37716650c593f0fab808b9abadd675cfae Mon Sep 17 00:00:00 2001 From: Nikhil Nakhate Date: Tue, 23 Dec 2025 15:05:54 -0800 Subject: [PATCH 2/2] Adding registry for action expert and adds flow matching head --- .gitignore | 2 +- models/action_expert/__init__.py | 6 ++ models/action_expert/registry.py | 70 ++++++++++++++++ models/diffusion_head.py | 31 ++++--- models/flow_matching_head.py | 139 +++++++++++++++++++++++++++++++ models/vla_diffusion_policy.py | 47 ++++++++--- scripts/test.py | 32 ++++++- scripts/train.py | 23 ++++- 8 files changed, 322 insertions(+), 28 deletions(-) create mode 100644 models/action_expert/__init__.py create mode 100644 models/action_expert/registry.py create mode 100644 models/flow_matching_head.py diff --git a/.gitignore b/.gitignore index c491ee8..b8f4600 100644 --- a/.gitignore +++ b/.gitignore @@ -3,4 +3,4 @@ __pycache__/ MUJOCO_LOG.txt LIBERO/ videos*/ -checkpoints/ \ No newline at end of file +checkpoints*/ \ No newline at end of file diff --git a/models/action_expert/__init__.py b/models/action_expert/__init__.py new file mode 100644 index 0000000..94e603b --- /dev/null +++ b/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" + diff --git a/models/action_expert/registry.py b/models/action_expert/registry.py new file mode 100644 index 0000000..e172826 --- /dev/null +++ b/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) + diff --git a/models/diffusion_head.py b/models/diffusion_head.py index 045f158..d296b49 100644 --- a/models/diffusion_head.py +++ b/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] diff --git a/models/flow_matching_head.py b/models/flow_matching_head.py new file mode 100644 index 0000000..42b7d31 --- /dev/null +++ b/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 diff --git a/models/vla_diffusion_policy.py b/models/vla_diffusion_policy.py index 5f6e5ab..6c118a6 100644 --- a/models/vla_diffusion_policy.py +++ b/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) diff --git a/scripts/test.py b/scripts/test.py index c9f0b07..59af886 100644 --- a/scripts/test.py +++ b/scripts/test.py @@ -8,6 +8,7 @@ from envs.metaworld_env import MetaWorldMT1Wrapper from models.vla_diffusion_policy import VLADiffusionPolicy +from models.action_expert.registry import ActionExpertCfg from utils.tokenizer import SimpleTokenizer @@ -67,11 +68,17 @@ def parse_args(): default="videos", help="Directory to save videos (if --save-video is set)", ) + parser.add_argument( + "--action-expert-name", + type=str, + default=None, + help="Override action expert name from checkpoint (diffusion | flow_matching)", + ) return parser.parse_args() -def load_model_and_tokenizer(checkpoint_path: str, device: torch.device): +def load_model_and_tokenizer(checkpoint_path: str, device: torch.device, action_expert_name_override: str = None): ckpt = torch.load(checkpoint_path, map_location=device) vocab = ckpt["vocab"] @@ -80,6 +87,24 @@ def load_model_and_tokenizer(checkpoint_path: str, device: torch.device): d_model = ckpt["d_model"] diffusion_T = ckpt["diffusion_T"] + # load action expert config + action_expert_cfg = None + if "action_expert_cfg" in ckpt: + action_expert_cfg = ActionExpertCfg(**ckpt["action_expert_cfg"]) + + # allow CLI override + if action_expert_name_override is not None: + if action_expert_cfg is None: + # create default config if checkpoint doesn't have it + action_expert_cfg = ActionExpertCfg( + name=action_expert_name_override, + action_dim=action_dim, + cond_dim=d_model, + T=diffusion_T, + ) + else: + action_expert_cfg.name = action_expert_name_override + vocab_size = max(vocab.values()) + 1 model = VLADiffusionPolicy( @@ -88,6 +113,7 @@ def load_model_and_tokenizer(checkpoint_path: str, device: torch.device): action_dim=action_dim, d_model=d_model, diffusion_T=diffusion_T, + action_expert_cfg=action_expert_cfg, ).to(device) model.load_state_dict(ckpt["model_state_dict"]) @@ -105,7 +131,7 @@ def main(): # load model + tokenizer print(f"[test] Loading checkpoint from {args.checkpoint}") - model, tokenizer = load_model_and_tokenizer(args.checkpoint, device) + model, tokenizer = load_model_and_tokenizer(args.checkpoint, device, action_expert_name_override=args.action_expert_name) # encode instruction instr_tokens = tokenizer.encode(args.instruction) @@ -116,7 +142,7 @@ def main(): env_name=args.env_name, seed=args.seed, render_mode="rgb_array", - camera_name="topview", + camera_name="corner", ) print(f"[test] Meta-World MT1 env: {args.env_name}") diff --git a/scripts/train.py b/scripts/train.py index 5ca41a5..3ec9d11 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -7,6 +7,7 @@ from torch.utils.data import Dataset, DataLoader from models.vla_diffusion_policy import VLADiffusionPolicy +from models.action_expert.registry import ActionExpertCfg class TrainingDataset(Dataset): @@ -54,6 +55,11 @@ def parse_args(): default="checkpoints/model.pt") parser.add_argument("--device", type=str, default="cuda", help="'cuda' or 'cpu'") + + # action expert command-line arguments + parser.add_argument("--action-expert-name", type=str, default="diffusion", + help="diffusion | flow_matching") + return parser.parse_args() @@ -67,12 +73,20 @@ def main(): state_dim = dataset.states.shape[1] action_dim = dataset.actions.shape[1] + action_expert_cfg = ActionExpertCfg( + name=args.action_expert_name, + action_dim=action_dim, + cond_dim=args.d_model, + T=args.diffusion_T, + ) + model = VLADiffusionPolicy( vocab_size=vocab_size, state_dim=state_dim, action_dim=action_dim, d_model=args.d_model, - diffusion_T=args.diffusion_T + diffusion_T=args.diffusion_T, + action_expert_cfg=action_expert_cfg, ).to(device) loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True) @@ -107,6 +121,13 @@ def main(): "action_dim": action_dim, "d_model": args.d_model, "diffusion_T": args.diffusion_T, + # save action expert config + "action_expert_cfg": { + "name": action_expert_cfg.name, + "action_dim": action_expert_cfg.action_dim, + "cond_dim": action_expert_cfg.cond_dim, + "T": action_expert_cfg.T, + }, }, args.save_path, )