Merge pull request #8 from StarrickLiu/main

Fix Readme, fix Timer in training

wall_x/data/load_lerobot_dataset.py

@@ -207,7 +207,7 @@ class DataCollator:
         self.load_processor()
 
     def load_processor(self):
-        processor_path = self.config["pretrained_qwen_vl_path"]
+        processor_path = self.config["pretrained_wallx_path"]
         action_tokenizer_path = self.config["action_tokenizer_path"]
 
         # Use cached processors if available

wall_x/trainer/qwen_vl_act_trainer.py

@@ -97,7 +97,7 @@ class QwenVlAct_Trainer:
                 - qwen_vl_act_config_path (str): Path to model configuration file
                 - learning_rate (float): Base learning rate for training
                 - num_epoch (int): Number of training epochs
-                - pretrained_qwen_vl_path (str): Path to pretrained model
+                - pretrained_wallx_path (str): Path to pretrained model
                 - And other training hyperparameters
             logger: Logger instance for tracking metrics
             accelerator (Accelerator, optional): Hugging Face Accelerate instance for distributed training
@@ -253,6 +253,11 @@ class QwenVlAct_Trainer:
             profiler.__enter__()
 
         try:
+            
+            # Setup timers for First iteration
+            self.timers("interval-time", log_level=0).start(barrier=False)
+            self.timers("data-load", log_level=0).start(barrier=False)
+
             for i, batch in enumerate(self.train_dataloader, self.initial_step):
                 # Move batch to device
                 if isinstance(self.dataset, PreprocessedDataset):
@@ -406,7 +411,7 @@ class QwenVlAct_Trainer:
         """
         # Load pretrained model
         model = Qwen2_5_VLMoEForAction.from_pretrained(
-            self.config["pretrained_qwen_vl_path"], 
+            self.config["pretrained_wallx_path"], 
             **{"use_fast_tokenizer": self.use_fast_tokenizer}
         )
         self.processor = model.processor

workspace/README.md

@@ -2,26 +2,25 @@
 
 This document explains the key configuration parameters that can be modified for Wall-X training.
 
+## Quick Start Checklist
+1. **Update run.sh**: Set `code_dir` and `config_path` to your actual paths
+2. **Configure GPUs**: Set `CUDA_VISIBLE_DEVICES` for your available GPUs  
+3. **Update config paths**: Replace all `/path/to/` placeholders in `config_qact.yml` with actual paths
+4. **Configure robot**: Set `dof_config` and `agent_pos_config` for your robot
+5. **Set dataset**: Choose appropriate `repo_id` for your dataset
+6. **Adjust batch size**: Set `batch_size_per_gpu` based on GPU memory
+7. **Run training**: Execute `bash ./workspace/lerobot_example/run.sh`
+
 ## Enable FAST tokenizer
 To fine-tune using the FAST tokenizer, please download the repository and update the `action_tokenizer_path`. Make sure to set `use_fast_tokenizer` to `true`:
 ```bash
 git clone https://huggingface.co/physical-intelligence/fast
 ```
 
-## Quick Start Checklist
-1. **Update run.sh**: Set `code_dir` and `config_path` to your actual paths
-2. **Configure GPUs**: Set `CUDA_VISIBLE_DEVICES` for your available GPUs  
-3. **Update config paths**: Replace all `/path/to/` placeholders in config.yml with actual paths
-4. **Configure robot**: Set `dof_config` and `agent_pos_config` for your robot
-5. **Set dataset**: Choose appropriate `repo_id` for your dataset
-6. **Adjust batch size**: Set `batch_size_per_gpu` based on GPU memory
-7. **Run training**: Execute `bash ./workspace/lerobot_example/run.sh`
-
 ## Required Paths (Must Modify)
 ```yaml
-processor_path: "/path/to/model/"                    # Path to model processor
-pretrained_qwen_vl_path: "/path/to/qwen_vl_model/"  # Path to pretrained Qwen VL model
-qwen_vl_act_config_path: "/path/to/config.json"     # Path to model config file
+pretrained_wallx_path: "/path/to/wallx_model/"      # Path to pretrained Qwen VL model
+use_fast_tokenizer: false                           # True: train FAST, False: train Flow
 action_tokenizer_path: "/path/to/fast/"             # Path to action tokenizer
 save_path: "/path/to/workspace/"                    # Path to save training outputs
 ```

workspace/lerobot_example/config_qact.yml

@@ -5,7 +5,7 @@
 log_name: "robotic_training"
 log_project: "vla_training"
 model_type: qwen2_5
-pretrained_qwen_vl_path: "/path/to/wallx_model/"
+pretrained_wallx_path: "/path/to/wallx_model/"
 use_fast_tokenizer: false # True: train FAST, False: train Flow
 action_tokenizer_path: "/path/to/fast/"
 save_path: "/path/to/workspace/"