diff --git a/wall_x/trainer/qwen_vl_act_trainer.py b/wall_x/trainer/qwen_vl_act_trainer.py index 05dc648..4a04cd7 100644 --- a/wall_x/trainer/qwen_vl_act_trainer.py +++ b/wall_x/trainer/qwen_vl_act_trainer.py @@ -197,6 +197,9 @@ class QwenVlAct_Trainer: self.train_loop(epoch) self.accelerator.wait_for_everyone() + if (epoch + 1) % self.config.get("epoch_save_interval", 10) == 0: + self.save_checkpoint(epoch) + # Validation after each epoch self.val_loop() self.accelerator.wait_for_everyone() diff --git a/workspace/lerobot_example/config_qact.yml b/workspace/lerobot_example/config_qact.yml index b4d2fea..83e8d50 100644 --- a/workspace/lerobot_example/config_qact.yml +++ b/workspace/lerobot_example/config_qact.yml @@ -5,9 +5,8 @@ log_name: "robotic_training" log_project: "vla_training" model_type: qwen2_5 -pretrained_qwen_vl_path: "/path/to/qwen_vl_model/" -# whether to enable fast tokenizer -use_fast_tokenizer: false +pretrained_qwen_vl_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/" @@ -27,6 +26,7 @@ num_epoch: 100 gradient_accumulation_steps: 32 batch_size_per_gpu: 8 padding_side: left +epoch_save_interval: 10 # Robot configuration - Define degrees of freedom for each component dof_config: