Update train from QwenVL (#50)

* update from vlm

* update

* update

* update

.gitattributes

@@ -0,0 +1 @@
+workspace/lerobot_example/evaluation/lerobot_openloop.png filter=lfs diff=lfs merge=lfs -text

wall_x/data/load_lerobot_dataset.py

@@ -255,14 +255,6 @@ class DataCollator:
         processor_path = self.config["pretrained_wallx_path"]
         action_tokenizer_path = self.config["action_tokenizer_path"]
 
-        # Use cached processors if available
-        if processor_path not in self._processor_cache:
-            self._processor_cache[processor_path] = AutoProcessor.from_pretrained(
-                processor_path, use_fast=True
-            )
-            if self.config.get("padding_side", "left") == "left":
-                self._processor_cache[processor_path].tokenizer.padding_side = "left"
-
         if (
             self.use_fast_tokenizer
             and action_tokenizer_path not in self._action_tokenizer_cache
@@ -273,6 +265,28 @@ class DataCollator:
                 )
             )
 
+        # Use cached processors if available
+        if processor_path not in self._processor_cache:
+            processor = AutoProcessor.from_pretrained(processor_path, use_fast=True)
+            if self.config.get("padding_side", "left") == "left":
+                processor.tokenizer.padding_side = "left"
+
+            if self.use_fast_tokenizer and self.config.get("model_type") == "qwen2_5":
+                action_tokenizer = self._action_tokenizer_cache[action_tokenizer_path]
+                new_tokens = ["<|propri|>", "<|action|>"]
+                new_tokens += [
+                    f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size)
+                ]
+                processor.tokenizer.add_tokens(new_tokens)
+                begin_idx_token = "<|action_token_0|>"
+                token_id = processor.tokenizer.convert_tokens_to_ids(begin_idx_token)
+                processor.tokenizer.init_kwargs["action_token_start_index"] = token_id
+                processor.tokenizer.init_kwargs["action_token_vocab_size"] = (
+                    action_tokenizer.vocab_size
+                )
+
+            self._processor_cache[processor_path] = processor
+
         self.processor = self._processor_cache[processor_path]
 
         if not self.use_fast_tokenizer:
@@ -282,15 +296,6 @@ class DataCollator:
                 action_tokenizer_path
             ]
 
-        if self.use_fast_tokenizer:
-            self.action_mapper = {}
-            for i in range(self.train_action_tokenizer.vocab_size):
-                token = f"<|action_token_{i}|>"
-                token_id = self.processor.tokenizer.convert_tokens_to_ids(token)
-                self.action_mapper[token_id] = i
-        else:
-            self.action_mapper = None
-
     @classmethod
     def _normalize(cls, action, min_stat, delta):
         """

wall_x/trainer/qwen_vl_act_trainer.py

@@ -15,9 +15,9 @@ from torch.distributed.tensor import distribute_tensor
 from accelerate import Accelerator
 from safetensors.torch import load_file
 from transformers.optimization import get_cosine_with_min_lr_schedule_with_warmup
-
+from transformers import AutoProcessor
 from wall_x.utils.timers import Timers
-from wall_x.model.qwen2_5_based import Qwen2_5_VLMoEForAction
+from wall_x.model.qwen2_5_based import Qwen2_5_VLMoEForAction, Qwen2_5_VLConfig
 from wall_x.data.config import ACTION_DATASET_NAMES, MULTIMODAL_DATASET_NAMES
 from wall_x.data.load_lerobot_dataset import (
     PreprocessedDataset,
@@ -212,7 +212,6 @@ class QwenVlAct_Trainer:
         - Memory cleanup
         """
         self.accelerator.wait_for_everyone()
-
         # Optional validation before training starts
         if self.config.get("resume", None) is not None and self.config["resume"].get(
             "validate_first", False
@@ -225,7 +224,7 @@ class QwenVlAct_Trainer:
             self.train_loop(epoch)
             self.accelerator.wait_for_everyone()
 
-            if (epoch + 1) % self.config.get("epoch_save_interval", 10) == 0:
+            if (epoch + 1) % self.config.get("epoch_save_interval", 1) == 0:
                 self.save_checkpoint(epoch)
 
             # Validation after each epoch
@@ -317,7 +316,6 @@ class QwenVlAct_Trainer:
                     self.timers("forward-compute").stop()
 
                     loss = outputs.loss
-
                     # Check for NaN loss
                     if torch.isnan(loss):
                         print(
@@ -519,12 +517,61 @@ class QwenVlAct_Trainer:
         - Model preparation for distributed training
         """
         # Load pretrained model
+        model_type = self.config.get("model_type", "qwen2_5")
+        assert model_type in ["wall-oss", "qwen2_5"]
+        if model_type == "wall-oss":
             model = Qwen2_5_VLMoEForAction.from_pretrained(
                 self.config["pretrained_wallx_path"],
                 **{"use_fast_tokenizer": self.use_fast_tokenizer},
             )
             self.processor = model.processor
             model = model.to(torch.bfloat16)
+        elif model_type == "qwen2_5":
+
+            config = Qwen2_5_VLConfig.from_pretrained(
+                self.config["qwen_vl_act_config_path"]
+            )
+            flow_loss_weight = self.config.get("flow_loss_weight", 1.0)
+            self.processor = AutoProcessor.from_pretrained(
+                self.config["pretrained_wallx_path"], use_fast=True
+            )
+            if self.config.get("use_fast_tokenizer", False):
+                action_tokenizer_path = self.config["action_tokenizer_path"]
+                action_tokenizer = AutoProcessor.from_pretrained(
+                    action_tokenizer_path, trust_remote_code=True
+                )
+                # process for use fast
+                new_tokens = ["<|propri|>", "<|action|>"]
+                new_tokens += [
+                    f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size)
+                ]
+                self.processor.tokenizer.add_tokens(new_tokens)
+                begin_idx_token = "<|action_token_0|>"
+                token_id = self.processor.tokenizer.convert_tokens_to_ids(
+                    begin_idx_token
+                )
+                self.processor.tokenizer.init_kwargs["action_token_start_index"] = (
+                    token_id
+                )
+                self.processor.tokenizer.init_kwargs["action_token_vocab_size"] = (
+                    action_tokenizer.vocab_size
+                )
+                self.processor.action_processor = action_tokenizer
+            model = Qwen2_5_VLMoEForAction(
+                config,
+                self.use_fast_tokenizer,
+                self.processor,
+                flow_loss_weight=flow_loss_weight,
+            )
+
+            model = model.to(torch.bfloat16)
+            model = self.load_qwen_pretrain_weight(
+                model, self.config["pretrained_wallx_path"]
+            )
+            model.resize_token_embeddings(len(self.processor.tokenizer))
+            model = model.to(torch.bfloat16)
+        else:
+            raise NotImplementedError(f"Invalid model type: {model_type}")
 
         # Configure optimizer based on training strategy
         if "freeze_vlm" in self.config and self.config["freeze_vlm"]:
@@ -752,8 +799,12 @@ class QwenVlAct_Trainer:
         else:
             ckpt_path = f"{save_path}/{epoch}_{step}"
 
+        # If FSDP SHARDED_STATE_DICT is used, please refer to the wall-x/workspace/README.md
+        # merge checkpoint section to merge the weights into a single safetensors if needed.
         self.accelerator.save_state(ckpt_path)
 
+        self.processor.save_pretrained(os.path.join(ckpt_path, "processor"))
+
         # Save current iteration steps for dataset resuming
         if step != 0:
             _rank = self.accelerator.process_index
@@ -773,9 +824,11 @@ class QwenVlAct_Trainer:
         """
         checkpoint_path = self.config["resume"]["ckpt"]
 
+        if self.config.get("resume", {}).get("load_ckpt_only", False):
             if self.config.get("FSDP2", False):
                 self._load_fsdp_state_dict_with_distribute_tensor()
-        elif self.config.get("resume", {}).get("load_ckpt_only", False):
+
+            else:
                 # Load only model weights
                 ckpt_path = self.config["resume"]["ckpt"] + "/model.safetensors"
                 state_dict = load_file(ckpt_path, device="cpu")

workspace/README.md

@@ -4,10 +4,12 @@ This document explains the key configuration parameters and memory requirements
 
 ## Quick Start Checklist
 
-### 🚀 **Step 1: Download Pre-trained Model**
+### 🚀 **Step 1: Prepare Model**
-Choose one of the available models:
+Choose one of our pretrained models:
 - **WALL-OSS-FLOW**: https://huggingface.co/x-square-robot/wall-oss-flow
 - **WALL-OSS-FAST**: https://huggingface.co/x-square-robot/wall-oss-fast
+Or from Qwen-2.5-VL
+- Download https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct, settings refer to `config_qact_from_vlm.yml`
 
 ### ⚙️ **Step 2: Configure Environment**
 - Update `run.sh`: Set `code_dir` and `config_path` to your actual paths
@@ -89,6 +91,13 @@ Keep `agent_pos_config` consistent with `dof_config`.
 - `resume.ckpt`: Path to checkpoint for resuming training
 - `resume.load_ckpt_only`: Only load model weights, not optimizer state
 
+## Merge checkpoint
+- If FSDP SHARDED_STATE_DICT is used, please run command below to merge checkpoint into a single safetensors
+```bash
+    # refer to accelerate/commands/merge.py
+    accelerate merge-weights /path/to/sharded_tensors /path/to/model.safetensors
+```
+
 ## Memory Usage
 
 Below are the memory consumption benchmarks for different training configurations using the `lerobot/aloha_mobile_cabinet` dataset:
@@ -106,3 +115,9 @@ Below are the memory consumption benchmarks for different training configuration
 
 - For single GPU training: Ensure at least 48GB VRAM (e.g., RTX 6000 Ada, A6000)
 - For multi-GPU training: Enable FSDP2 for optimal memory distribution
+
+## Reproduce
+
+Openloop plot `wall-x/workspace/lerobot_example/evaluation/lerobot_openloop.png`
+
+To reproduce the results, use the config file wall-x/workspace/lerobot_example/config_qact_from_vlm.yml with a global batch size of 128, adjusted via `gradient_accumulation_steps` and numbers of gpu.

workspace/lerobot_example/config_qact.yml

@@ -4,7 +4,7 @@
 # Model and paths configuration
 log_name: "robotic_training"
 log_project: "vla_training"
-model_type: qwen2_5
+model_type: wall-oss
 pretrained_wallx_path: "/path/to/wallx_model/"  # Must set
 save_path: "/path/to/workspace/"                # Must set
 use_fast_tokenizer: False                       # True: train FAST, False: train Flow

workspace/lerobot_example/config_qact_from_vlm.yml

@@ -0,0 +1,117 @@
+# Train from Qwen-2.5-VL
+
+# Model and paths configuration
+log_name: "robotic_training"
+log_project: "vla_training"
+model_type: qwen2_5
+pretrained_wallx_path: "/path/to/wallx_model/"  # Must set
+save_path: "/path/to/workspace/"                # Must set
+use_fast_tokenizer: True                       # True: train FAST, False: train Flow
+action_tokenizer_path: "/path/to/fast/"         # Must set if use_fast_tokenizer is true
+qwen_vl_act_config_path: "wall-x/workspace/lerobot_example/qwen25_config.json"
+
+
+# Torch Profile
+profile: False
+profile_save_path: /path/to/profile/
+profile_wait_iters: 10
+profile_warmup_iters: 5
+profile_active_iters: 2
+
+# Training hyperparameters
+num_warmup_steps: 100
+num_training_steps: 64000000
+learning_rate: 0.00009
+min_lr: 0.00005
+num_epoch: 100
+gradient_accumulation_steps: 1
+batch_size_per_gpu: 8
+padding_side: left
+epoch_save_interval: 10
+
+# Training optimization settings
+FSDP2: True
+torch_compile: False
+
+# Robot configuration - Define degrees of freedom for each component
+dof_config:
+  follow_left_ee_cartesian_pos: 3      # Left end-effector Cartesian position
+  follow_left_ee_rotation: 3           # Left end-effector rotation
+  follow_left_gripper: 1               # Left gripper control
+  follow_right_ee_cartesian_pos: 3     # Right end-effector Cartesian position
+  follow_right_ee_rotation: 3          # Right end-effector rotation
+  follow_right_gripper: 1              # Right gripper control
+  head_actions: 2                      # Head/camera movement
+  height: 1                            # Mobile base height control
+  car_pose: 3                          # Mobile base pose (x, y, theta)
+
+# Agent proprioception configuration (typically matches DOF config)
+agent_pos_config:
+  follow_left_ee_cartesian_pos: 3
+  follow_left_ee_rotation: 3
+  follow_left_gripper: 1
+  follow_right_ee_cartesian_pos: 3
+  follow_right_ee_rotation: 3
+  follow_right_gripper: 1
+  head_actions: 2
+  height: 1
+  car_pose: 3
+
+# # Checkpoint resuming configuration
+# resume:
+#   ckpt: "/path/to/resume_model/"
+#   load_ckpt_only: true
+
+# Data configuration
+data:
+  use_lerobot: true
+
+  # LeRobot dataset configuration
+  lerobot_config:
+    repo_id: "lerobot/aloha_mobile_cabinet"
+    root: null
+    episodes: null
+    image_transforms: null
+    delta_timestamps: null
+    tolerance_s: 1e-4
+    revision: null
+    force_cache_sync: false
+    download_videos: true
+    video_backend: null
+
+  action_horizon: 32
+  train_test_split: 0.95
+
+  # Action keys for observation and prediction
+  obs_action_keys:
+    - follow_left_ee_cartesian_pos
+    - follow_left_ee_rotation
+    - follow_left_gripper
+    - follow_right_ee_cartesian_pos
+    - follow_right_ee_rotation
+    - follow_right_gripper
+    - head_actions
+    - height
+    - car_pose
+
+  predict_action_keys:
+    - follow_left_ee_cartesian_pos
+    - follow_left_ee_rotation
+    - follow_left_gripper
+    - follow_right_ee_cartesian_pos
+    - follow_right_ee_rotation
+    - follow_right_gripper
+    - head_actions
+    - height
+    - car_pose
+
+  # Image resolution configuration for different camera views
+  resolution:
+    face_view: 256
+    left_wrist_view: 256
+    right_wrist_view: 256
+    move1_view: 256
+    move2_view: 256
+    top_view: 256
+    wall_view: 256
+    multi_modal: 256

workspace/lerobot_example/evaluation/lerobot_openloop.png

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aae8646e566b79b64a669956d6d2c779d72809010c802f75b2623ee371444b47
+size 985979

workspace/lerobot_example/qwen25_config.json

@@ -0,0 +1,106 @@
+{
+    "architectures": [
+      "Qwen2_5_VLForConditionalGeneration"
+    ],
+    "attention_dropout": 0.0,
+    "bos_token_id": 151643,
+    "eos_token_id": 151645,
+    "vision_start_token_id": 151652,
+    "vision_end_token_id": 151653,
+    "vision_token_id": 151654,
+    "image_token_id": 151655,
+    "video_token_id": 151656,
+    "hidden_act": "silu",
+    "hidden_size": 2048,
+    "initializer_range": 0.02,
+    "intermediate_size": 11008,
+    "max_position_embeddings": 128000,
+    "max_window_layers": 70,
+    "model_type": "qwen2_5_vl",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 36,
+    "num_key_value_heads": 2,
+    "rms_norm_eps": 1e-06,
+    "rope_theta": 1000000.0,
+    "sliding_window": 32768,
+    "tie_word_embeddings": true,
+    "torch_dtype": "bfloat16",
+    "transformers_version": "4.41.2",
+    "_attn_implementation": "flash_attention_2",
+    "use_cache": true,
+    "use_sliding_window": false,
+    "vision_config": {
+      "depth": 32,
+      "hidden_act": "silu",
+      "hidden_size": 1280,
+      "intermediate_size": 3420,
+      "num_heads": 16,
+      "in_chans": 3,
+      "out_hidden_size": 2048,
+      "patch_size": 14,
+      "spatial_merge_size": 2,
+      "spatial_patch_size": 14,
+      "window_size": 112,
+      "fullatt_block_indexes": [
+        7,
+        15,
+        23,
+        31
+      ],
+      "tokens_per_second": 2,
+      "temporal_patch_size": 2
+    },
+    "rope_scaling": {
+      "type": "mrope",
+      "mrope_section": [
+        16,
+        24,
+        24
+      ]
+    },
+    "vocab_size": 151936,
+    "num_experts": 2,
+    "experts":[
+      {
+        "hidden_size": 2048,
+        "intermediate_size": 11008,
+        "hidden_act": "silu"
+      },
+      {
+        "hidden_size": 2048,
+        "intermediate_size": 2048,
+        "hidden_act": "silu"
+      }
+    ],
+    "dof_config": {
+        "follow_left_ee_cartesian_pos": 3,
+        "follow_left_ee_rotation": 3,
+        "follow_left_gripper": 1,
+        "follow_right_ee_cartesian_pos": 3,
+        "follow_right_ee_rotation": 3,
+        "follow_right_gripper": 1,
+        "head_actions": 2,
+        "height": 1,
+        "car_pose": 3
+    },
+    "agent_pos_config": {
+      "follow_left_ee_cartesian_pos": 3,
+      "follow_left_ee_rotation": 3,
+      "follow_left_gripper": 1,
+      "follow_right_ee_cartesian_pos": 3,
+      "follow_right_ee_rotation": 3,
+      "follow_right_gripper": 1,
+      "head_actions": 2,
+      "height": 1,
+      "car_pose": 3
+    },
+    "noise_scheduler": {
+      "beta_alpha": 1.5,
+      "beta_beta": 1.0,
+      "s": 0.999,
+      "num_inference_timesteps": 5
+    },
+    "dim_inputs": [2048,2048],
+    "attention_moe": false,
+    "mlp_moe": true
+  }