refactor

Author: legitnull

examples/qwen_gr00t/conf/train/qwen_gr00t.yaml

@@ -16,32 +16,23 @@ system:
   # TODO(yupu): Support resuming from checkpoint
 
 model:
-  # TODO: (yupu) the config layout is still a mess
   model_name: qwen_gr00t
-  # Path to the checkpoint of the pretrained base VLM model, e.g. Qwen3-VL-4B-Instruct
-  checkpoint_dir: /workspace/models/Qwen/Qwen3-VL-4B-Instruct/
   vlm:
     type: qwen3-vl
-  qwenvl:
     base_vlm: /workspace/models/Qwen/Qwen3-VL-4B-Instruct/
     attn_implementation: flash_attention_2
-    vl_hidden_dim: 2048
-  dino:
-    dino_backbone: dinov2_vits14
   action_model:
     # Whether to condition the action model on proprioceptive state (observation.state)
     use_state: false
-    type: flow_matching
+    type: gr00t_action_head
     action_model_type: DiT-B
-    action_hidden_dim: 1024
     hidden_size: 1024
     add_pos_embed: True
     max_seq_len: 1024
     action_dim: 7
     state_dim: 7
     future_action_window_size: 7
     action_horizon: 8
-    past_action_window_size: 0
     repeated_diffusion_steps: 4
     noise_beta_alpha: 1.5
     noise_beta_beta: 1.0
@@ -58,7 +49,13 @@ model:
       num_layers: 16
       output_dim: 1024
       positional_embeddings: None
-  reduce_in_full_precision: True
+
+  prompt_template: "Your task is {instruction}. To identify the key objects for your task. Locate their bounding boxes in [x1,y1,x2,y2] format."
+
+  normalization_mapping:
+    VISUAL: IDENTITY
+    STATE: MIN_MAX
+    ACTION: MIN_MAX
 
   optimizer:
     name: AdamW
@@ -117,24 +114,9 @@ data:
     vision_root: ""
     action_key: eepose
     state_key: eepose
-  # TODO: (yupu) Remove this once we have a proper dataset config
-  vla_data:
-    dataset_py: lerobot_datasets
-    data_root_dir: playground/Datasets/
-    data_mix: libero_goal_old
-    action_type: delta_qpos
-    CoT_prompt: Your task is {instruction}. To identify the key objects for your task. Locate their bounding boxes in [x1,y1,x2,y2] format.
-    CoT_answer: bbox
-    default_image_resolution: [3, 224, 224]
-    load_all_data_for_training: True
-    obs: ["image_0"]
-    video_backend: torchvision_av
   # Path to the training data
   data_path: /workspace/datasets/IPEC-COMMUNITY/libero_goal_no_noops_1.0.0_lerobot/
   tolerance_s: 0.0001
-  # TODO: (yupu) I think these indices should belong to the policy config, maybe put it in the model config?
-  observation_delta_indices: [0]
-  action_delta_indices: [0,1,2,3,4,5,6,7]
   preprocessor:
     name: policy_preprocessor
     steps:
@@ -151,18 +133,12 @@ data:
         config:
           eps: 1e-8
           features: {}
-          norm_map:
-            VISUAL: IDENTITY
-            STATE: MIN_MAX
-            ACTION: MIN_MAX
+          # norm_map is injected at runtime from model.normalization_mapping
   postprocessor:
     name: policy_postprocessor
     steps:
       - registry_name: unnormalizer_processor
         config:
           eps: 1e-8
           features: {}
-          norm_map:
-            VISUAL: IDENTITY
-            STATE: MIN_MAX
-            ACTION: MIN_MAX
+          # norm_map is injected at runtime from model.normalization_mapping

flagscale/inference/inference_qwen_gr00t.py

@@ -1,5 +1,4 @@
 import argparse
-import importlib
 
 import numpy as np
 import torch
@@ -8,7 +7,8 @@
 
 from flagscale.logger import logger
 from flagscale.models.utils.constants import OBS_STATE
-from flagscale.train.utils.train_utils import load_checkpoint
+from flagscale.models.vla import TrainablePolicy
+from flagscale.train.processor import PolicyProcessorPipeline
 
 
 def load_image(image_path: str, size: tuple[int, int] | None = None) -> torch.Tensor:
@@ -33,10 +33,16 @@ def run_inference(config_path: str):
     engine_cfg = cfg.engine
     generate_cfg = cfg.generate
 
-    model_variant = engine_cfg.model_variant
-    policy = getattr(importlib.import_module("flagscale.models.vla"), model_variant)
-    model, preprocessor, postprocessor = load_checkpoint(
-        engine_cfg.model, policy, engine_cfg.device
+    pretrained_dir = engine_cfg.model
+    model = TrainablePolicy.from_pretrained(pretrained_dir, device=engine_cfg.device)
+
+    preprocessor = PolicyProcessorPipeline.from_pretrained(
+        pretrained_dir,
+        config_filename="policy_preprocessor.json",
+    )
+    postprocessor = PolicyProcessorPipeline.from_pretrained(
+        pretrained_dir,
+        config_filename="policy_postprocessor.json",
     )
 
     # TODO: (yupu): model.to(dtype)?

flagscale/models/utils/constants.py

@@ -52,6 +52,19 @@
 POLICY_PREPROCESSOR_DEFAULT_NAME = "policy_preprocessor"
 POLICY_POSTPROCESSOR_DEFAULT_NAME = "policy_postprocessor"
 
+
+def resolve_pretrained_dir(path: Path, sentinel: str) -> Path:
+    """Resolve a checkpoint root dir to the pretrained_model subdir if needed.
+
+    If ``sentinel`` (e.g. ``config.json``) does not exist directly under
+    ``path`` but does exist under ``path / pretrained_model/``, returns the
+    latter.  Otherwise returns ``path`` unchanged.
+    """
+    if not (path / sentinel).exists() and (path / PRETRAINED_MODEL_DIR / sentinel).exists():
+        return path / PRETRAINED_MODEL_DIR
+    return path
+
+
 if "LEROBOT_HOME" in os.environ:
     raise ValueError(
         f"You have a 'LEROBOT_HOME' environment variable set to '{os.getenv('LEROBOT_HOME')}'.\n"

flagscale/models/vla/__init__.py

@@ -1,4 +1,4 @@
-from .action_model import FlowMatchingHead
+from .base_policy import TrainablePolicy
 from .protocols import ActionModel, VLMBackbone
 from .qwen_gr00t import QwenGr00t
 from .registry import (
@@ -11,6 +11,7 @@
 from .vlm import Qwen3VLBackbone, Qwen25VLBackbone, QwenVLBackbone
 
 __all__ = [
+    "TrainablePolicy",
     "VLMBackbone",
     "ActionModel",
     "register_vlm",
@@ -22,5 +23,4 @@
     "QwenVLBackbone",
     "Qwen25VLBackbone",
     "Qwen3VLBackbone",
-    "FlowMatchingHead",
 ]

flagscale/models/vla/action_model/__init__.py

@@ -1,3 +1,3 @@
-from .flow_matching import FlowMatchingHead
+from .gr00t_action_header import GR00TActionHead
 
-__all__ = ["FlowMatchingHead"]
+__all__ = ["GR00TActionHead"]

flagscale/models/vla/action_model/gr00t_action_header.py

@@ -8,7 +8,10 @@
 # Action repeat is inspired by CogACT
 
 
+from __future__ import annotations
+
 from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
 
 import torch
 import torch.nn.functional as F
@@ -17,11 +20,65 @@
 from transformers import PretrainedConfig
 from transformers.feature_extraction_utils import BatchFeature
 
+if TYPE_CHECKING:
+    from omegaconf import DictConfig
+
+from flagscale.models.utils.constants import ACTION
 from flagscale.models.vla.action_model.flow_matching_head.cross_attention_dit import DiT
 from flagscale.models.vla.action_model.flow_matching_head.encoding_utils import (
     SinusoidalPositionalEncoding,
     swish,
 )
+from flagscale.models.vla.registry import register_action_model
+
+
+@dataclass
+class GR00TActionHeadConfig:
+    type: str = "gr00t_action_head"
+    action_model_type: str = "DiT-B"
+    hidden_size: int = 1024
+    action_dim: int = 7
+    state_dim: int = 7
+    future_action_window_size: int = 7
+    action_horizon: int = 8
+    use_state: bool = False
+    repeated_diffusion_steps: int = 4
+    add_pos_embed: bool = True
+    max_seq_len: int = 1024
+    noise_beta_alpha: float = 1.5
+    noise_beta_beta: float = 1.0
+    noise_s: float = 0.999
+    num_timestep_buckets: int = 1000
+    num_inference_timesteps: int = 4
+    num_target_vision_tokens: int = 32
+    diffusion_model_cfg: dict = field(default_factory=dict)
+
+    @classmethod
+    def from_omegaconf(cls, cfg: DictConfig) -> GR00TActionHeadConfig:
+        if cfg.get("diffusion_model_cfg") is None:
+            raise ValueError("diffusion_model_cfg is required in action_model config")
+        diffusion_cfg = dict(cfg.diffusion_model_cfg)
+        return cls(
+            type=cfg.get("type", "gr00t_action_head"),
+            action_model_type=cfg.get("action_model_type", "DiT-B"),
+            hidden_size=cfg.get("hidden_size", 1024),
+            action_dim=cfg.get("action_dim", 7),
+            state_dim=cfg.get("state_dim", 7),
+            future_action_window_size=cfg.get("future_action_window_size", 7),
+            action_horizon=cfg.get("action_horizon", 8),
+            use_state=cfg.get("use_state", False),
+            repeated_diffusion_steps=cfg.get("repeated_diffusion_steps", 4),
+            add_pos_embed=cfg.get("add_pos_embed", True),
+            max_seq_len=cfg.get("max_seq_len", 1024),
+            noise_beta_alpha=cfg.get("noise_beta_alpha", 1.5),
+            noise_beta_beta=cfg.get("noise_beta_beta", 1.0),
+            noise_s=cfg.get("noise_s", 0.999),
+            num_timestep_buckets=cfg.get("num_timestep_buckets", 1000),
+            num_inference_timesteps=cfg.get("num_inference_timesteps", 4),
+            num_target_vision_tokens=cfg.get("num_target_vision_tokens", 32),
+            diffusion_model_cfg=diffusion_cfg,
+        )
+
 
 # TODO try to merge DiT Modules with follow_match_head, they are just the same arch, but diff loss, use diffusers package will be simple
 
@@ -37,7 +94,6 @@ def __init__(self, num_categories, input_dim, hidden_dim):
     def forward(self, x, cat_ids):
         selected_W = self.W[cat_ids]
         selected_b = self.b[cat_ids]
-        # import ipdb; ipdb.set_trace()
         return torch.bmm(x, selected_W) + selected_b.unsqueeze(1)
 
 
@@ -218,12 +274,10 @@ def __init__(self, **kwargs):
 class FlowmatchingActionHead(nn.Module):
     def __init__(
         self,
-        full_config,
+        config: GR00TActionHeadConfig,
     ):
         super().__init__()
-        config = full_config.model.action_model
         self.hidden_size = config.hidden_size  # @JinhuiYE
-        self.full_config = full_config
         action_model_type = config.action_model_type
         action_model_cfg = DiTConfig[action_model_type]
 
@@ -405,3 +459,47 @@ def device(self):
     @property
     def dtype(self):
         return next(iter(self.parameters())).dtype
+
+
+@register_action_model("gr00t_action_head")
+class GR00TActionHead(nn.Module):
+    """
+    GR00T action head wrapper for VLA framework.
+
+    Adapts the FlowmatchingActionHead (tensor-based interface) to the
+    dict-based interface expected by the VLA framework (QwenGr00t).
+    """
+
+    def __init__(self, config: GR00TActionHeadConfig):
+        super().__init__()
+        self._head = FlowmatchingActionHead(config=config)
+
+    def forward(
+        self, vlm_output: dict[str, torch.Tensor], action_input: dict[str, torch.Tensor]
+    ) -> dict[str, torch.Tensor]:
+        vl_embs = vlm_output["hidden_states"]
+        actions = action_input["actions"]
+        state = action_input.get("state")
+        encoder_attention_mask = action_input.get("attention_mask")
+        mask = action_input.get("mask")
+
+        loss = self._head.forward(
+            vl_embs=vl_embs,
+            actions=actions,
+            state=state,
+            encoder_attention_mask=encoder_attention_mask,
+            mask=mask,
+        )
+        return {"loss": loss}
+
+    def predict_action(
+        self, vlm_output: dict[str, torch.Tensor], action_input: dict[str, torch.Tensor]
+    ) -> dict[str, torch.Tensor]:
+        vl_embs = vlm_output["hidden_states"]
+        state = action_input.get("state")
+
+        actions = self._head.predict_action(vl_embs=vl_embs, state=state)
+        return {ACTION: actions}
+
+    def fsdp_units(self) -> list[nn.Module]:
+        return list(self._head.model.transformer_blocks)