huangyanpei/VLA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix normalizer (#57)

Browse Source 
* fix normalizer

* fix val

* update compute stats

* delete norm

* update readme

* minor fix

* fix action normalizer

* fix

* fix

* update

* update

* update

* update

* update

* lint

* lint

* lint

* lint
This commit is contained in:
Lufang Chen
2025-10-24 17:29:12 +08:00
committed by GitHub
parent 35399d187a
commit d821b0cb26
14 changed files with 837 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wall_x/data/config.py
+2
View File
@@ -36,6 +36,8 @@ ACTION_DATASET_NAMES = [
"taco_play",
"utaustin_mutex",
"viola",
"physical-intelligence/libero",
"lerobot/aloha_mobile_cabinet",
]
# Supported multimodal datasets
wall_x/data/load_lerobot_dataset.py
+69 -37
View File
@@ -17,17 +17,10 @@ from wall_x.data.utils import (
)
from transformers import AutoProcessor
from .utils import load_norm_stats, KEY_MAPPINGS
T_co = TypeVar("T_co", covariant=True)
CAMERA_KEY_MAPPINGS = {
"lerobot/aloha_mobile_cabinet": {
"observation.images.cam_high": "face_view",
"observation.images.cam_left_wrist": "left_wrist_view",
"observation.images.cam_right_wrist": "right_wrist_view",
},
}
# Abstract class for dataset
class Dataset(Protocol[T_co]):
@@ -46,6 +39,8 @@ class PreprocessedDataset(Dataset[T_co]):
dataset,
config,
dataload_config,
norm_stats,
lerobot_config,
seed=42,
rank=0,
world_size=1,
@@ -72,6 +67,8 @@ class PreprocessedDataset(Dataset[T_co]):
self.config = config
self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False)
self.dataload_config = dataload_config
self.norm_stats = norm_stats
self.lerobot_config = lerobot_config
self.data_config = X2RDataProcessingConfig().update(
train_test_split=self.dataload_config["train_test_split"],
@@ -82,7 +79,9 @@ class PreprocessedDataset(Dataset[T_co]):
priority_order=self.dataload_config.get("priority_order", None),
)
self._cam_key_mapping = CAMERA_KEY_MAPPINGS[self.hf_dataset.meta.repo_id]
self._cam_key_mapping = KEY_MAPPINGS[self.hf_dataset.meta.repo_id]["camera"]
self._state_key_mapping = KEY_MAPPINGS[self.hf_dataset.meta.repo_id]
self._action_key_mapping = KEY_MAPPINGS[self.hf_dataset.meta.repo_id]
def _vision_preprocess(self, frames):
processed_frames = []
@@ -124,14 +123,14 @@ class PreprocessedDataset(Dataset[T_co]):
def __getitem__(self, index):
data = self._dataset[index]
image_inputs, h, w, resize_h, resize_w = self._vision_preprocess(data)
agent_pos = data["observation.state"]
action = data["action"]
agent_pos = data[self._state_key_mapping["state"]]
action = data[self._action_key_mapping["action"]]
frame_index = data["frame_index"]
instruction_info = {"instruction": data["task"]}
generate_subtask_ratio = self.data_config.generate_subtask_ratio
complete_text, generate_subtask = get_wallx_normal_text(
instruction_info,
33 - 1,
self.dataload_config.get("action_horizon", 33) - 1,
frame_index,
self.data_config.priority_order,
self._cam_key_mapping,
@@ -189,7 +188,7 @@ class PreprocessedDataset(Dataset[T_co]):
sampler=sampler, # Use distributed sampler instead of shuffle=True
num_workers=num_workers,
collate_fn=DataCollator(
self.config, self.dataload_config, self.hf_dataset.meta.stats
self.config, self.dataload_config, self.norm_stats, self.lerobot_config
),
pin_memory=True, # Enable for GPU training
persistent_workers=num_workers > 0, # Only if num_workers > 0
@@ -225,7 +224,7 @@ class PreprocessedDataset(Dataset[T_co]):
sampler=sampler,
num_workers=num_workers,
collate_fn=DataCollator(
self.config, self.dataload_config, self.hf_dataset.meta.stats
self.config, self.dataload_config, self.norm_stats, self.lerobot_config
),
pin_memory=True,
persistent_workers=num_workers > 0,
@@ -241,13 +240,16 @@ class DataCollator:
_processor_cache = {}
_action_tokenizer_cache = {}
def __init__(self, config, dataload_config, stats):
def __init__(self, config, dataload_config, stats, lerobot_config):
self.config = config
self.dataload_config = dataload_config
self.stats = stats
self.min_stat = stats["action"]["min"]
self.max_stat = stats["action"]["max"]
self.delta = self.max_stat - self.min_stat
self.action_min_stat = stats["action"].min
self.action_delta = stats["action"].delta
self.state_min_stat = stats["state"].min
self.state_delta = stats["state"].delta
self.lerobot_config = lerobot_config
self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False)
self.load_processor()
@@ -271,10 +273,11 @@ class DataCollator:
if self.config.get("padding_side", "left") == "left":
processor.tokenizer.padding_side = "left"
new_tokens = ["<|propri|>", "<|action|>"]
processor.tokenizer.add_tokens(new_tokens)
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 += [
new_tokens = [
f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size)
]
processor.tokenizer.add_tokens(new_tokens)
@@ -301,7 +304,6 @@ class DataCollator:
"""
Normalize action data using min-max normalization.
"""
delta = torch.from_numpy(delta)
delta = torch.where(delta == 0, torch.ones_like(delta), delta)
x = (action - min_stat) / delta
x = x * 2 - 1
@@ -318,7 +320,9 @@ class DataCollator:
agent_pos = agent_pos.unsqueeze(1)
agent_pos_mask = (~torch.isnan(agent_pos)).float()
agent_pos.nan_to_num_(nan=0.0)
agent_pos = self._normalize(agent_pos, self.min_stat, self.delta)
agent_pos = self._normalize(
agent_pos, self.state_min_stat, self.state_delta
)
if agent_pos.shape[-1] != 20:
agent_pos = torch.cat(
[
@@ -350,7 +354,9 @@ class DataCollator:
action = action.unsqueeze(1)
dof_mask = (~torch.isnan(action)).float()
action.nan_to_num_(nan=0.0)
action = self._normalize(action, self.min_stat, self.delta)
action = self._normalize(
action, self.action_min_stat, self.action_delta
)
if action.shape[-1] != 20:
action = torch.cat(
[
@@ -393,7 +399,7 @@ class DataCollator:
additional_inputs["text"],
additional_inputs["action_chunk"],
self.train_action_tokenizer if self.use_fast_tokenizer else None,
["x2_normal"] * additional_inputs["text"].__len__(),
[self.lerobot_config["repo_id"]] * additional_inputs["text"].__len__(),
additional_inputs["dof_mask"],
)
@@ -415,7 +421,9 @@ class DataCollator:
inputs.update(additional_inputs)
inputs["dataset_names"] = ["x2_normal"] * inputs["action_chunk"].shape[0]
inputs["dataset_names"] = [self.lerobot_config["repo_id"]] * inputs[
"action_chunk"
].shape[0]
return inputs
@@ -447,18 +455,24 @@ def load_lerobot_data(
dataload_config = get_data_configs(config["data"])
# repo_id = "lerobot/aloha_mobile_cabinet"
repo_id = lerobot_config.get("repo_id", "lerobot/aloha_mobile_cabinet")
repo_id = lerobot_config.get("repo_id", None)
assert repo_id is not None, "repo id is required"
root = lerobot_config.get("root", None)
meta_info = LeRobotDatasetMetadata(repo_id)
meta_info = LeRobotDatasetMetadata(repo_id, root=root)
dataset_fps = meta_info.fps
episodes_num = meta_info.total_episodes
norm_stats_path = config.get("norm_stats_path", None)
assert (
norm_stats_path is not None
), "norm stats is required, please refer to 'wall-x/scripts/compute_norm_stats.py' to compute stats"
norm_stats = load_norm_stats(norm_stats_path, repo_id)
delta_timestamps = {
# action chunk
"action": [
KEY_MAPPINGS[repo_id]["action"]: [
t / dataset_fps
for t in range(dataload_config.get("action_horizon", 32) - 1)
for t in range(dataload_config.get("action_horizon", 33) - 1)
],
}
batch_size = config.get("batch_size_per_gpu", 8)
@@ -486,6 +500,8 @@ def load_lerobot_data(
train_dataset,
config,
dataload_config,
norm_stats,
lerobot_config,
seed=seed,
rank=rank,
world_size=world_size,
@@ -567,11 +583,15 @@ def get_data_configs(config):
class TestDataset(PreprocessedDataset):
def __init__(self, dataset, config, dataload_config, seed=42):
def __init__(
self, dataset, config, dataload_config, norm_stats, lerobot_config, seed=42
):
super().__init__(
dataset,
config,
dataload_config,
norm_stats,
lerobot_config,
seed=seed,
rank=0,
world_size=1,
@@ -587,7 +607,7 @@ class TestDataset(PreprocessedDataset):
self,
batch_size=1,
collate_fn=DataCollator(
self.config, self.dataload_config, self.hf_dataset.meta.stats
self.config, self.dataload_config, self.norm_stats, self.lerobot_config
),
)
@@ -614,29 +634,41 @@ def load_test_dataset(
# Set seed for reproducibility
torch.manual_seed(seed)
dataset_fps = 50
repo_id = lerobot_config.get("repo_id", None)
assert repo_id is not None, "repo id is required"
root = lerobot_config.get("root", None)
meta_info = LeRobotDatasetMetadata(repo_id, root=root)
dataset_fps = meta_info.fps
dataload_config = get_data_configs(config["data"])
norm_stats_path = config.get("norm_stats_path", None)
assert (
norm_stats_path is not None
), "norm stats is required, please refer to 'wall-x/scripts/compute_norm_stats.py' to compute stats"
norm_stats = load_norm_stats(norm_stats_path, repo_id)
delta_timestamps = {
# action chunk
"action": [
KEY_MAPPINGS[repo_id]["action"]: [
t / dataset_fps
for t in range(dataload_config.get("action_horizon", 32) - 1)
for t in range(dataload_config.get("action_horizon", 33) - 1)
],
}
repo_id = lerobot_config.get("repo_id", "lerobot/aloha_mobile_cabinet")
dataset = LeRobotDataset(
repo_id,
episodes=[episode],
delta_timestamps=delta_timestamps,
video_backend="pyav",
root=root,
)
print(f"Selected episodes: {dataset.episodes}")
print(f"Number of episodes selected: {dataset.num_episodes}")
print(f"Number of frames selected: {dataset.num_frames}")
dataset = TestDataset(dataset, config, dataload_config, seed=seed)
dataset = TestDataset(
dataset, config, dataload_config, norm_stats, lerobot_config, seed=seed
)
return dataset
wall_x/data/utils.py
+53
View File
@@ -11,7 +11,28 @@ import random
from collections import OrderedDict
from typing import List, Dict, Any, Optional, Union, Tuple
from transformers import BatchFeature
from dataclasses import dataclass
import json
KEY_MAPPINGS = {
"lerobot/aloha_mobile_cabinet": {
"camera": {
"observation.images.cam_high": "face_view",
"observation.images.cam_left_wrist": "left_wrist_view",
"observation.images.cam_right_wrist": "right_wrist_view",
},
"state": "observation.state",
"action": "action",
},
"physical-intelligence/libero": {
"camera": {
"image": "face_view",
"wrist_image": "left_wrist_view",
},
"state": "state",
"action": "actions",
},
}
CAMERA_NAME_MAPPING = {
"face_view": "front view",
@@ -609,3 +630,35 @@ def replace_action_token(
text = [t.replace("<|action_fast|><|im_end|>\n", "") for t in text]
return text
@dataclass
class NormStats:
min: torch.Tensor
max: torch.Tensor
delta: torch.Tensor
def load_norm_stats(norm_stats_path, dataset_name):
with open(norm_stats_path, "r") as f:
norm_stats = json.load(f)
action_key = KEY_MAPPINGS[dataset_name]["action"]
state_key = KEY_MAPPINGS[dataset_name]["state"]
q01 = torch.tensor(norm_stats["norm_stats"][action_key]["q01"])
q99 = torch.tensor(norm_stats["norm_stats"][action_key]["q99"])
delta = q99 - q01
action_norm_stats = NormStats(
min=q01,
max=q99,
delta=delta,
)
q01 = torch.tensor(norm_stats["norm_stats"][state_key]["q01"])
q99 = torch.tensor(norm_stats["norm_stats"][state_key]["q99"])
delta = q99 - q01
state_norm_stats = NormStats(
min=q01,
max=q99,
delta=delta,
)
return {"action": action_norm_stats, "state": state_norm_stats}
wall_x/model/action_head.py
+47 -7
View File
@@ -3,6 +3,7 @@ import torch
import torch.nn as nn
from torch.distributions import Beta
from wall_x.utils.constant import action_statistic_dof
import logging
class Normalizer(nn.Module):
@@ -14,6 +15,16 @@ class Normalizer(nn.Module):
normalization to map actions to the [-1, 1] range.
"""
def _pad_to_action_dim(self, xs, action_dim):
"""
Pad the action data to the action dimension.
"""
if xs.shape[-1] < action_dim:
padding_shape = list(xs.shape)
padding_shape[-1] = action_dim - padding_shape[-1]
xs = torch.cat([xs, torch.zeros(padding_shape).to(xs.device)], dim=-1)
return xs
def __init__(self, action_statistic_dof, dof_config):
"""
Initialize the normalizer with robot-specific action statistics.
@@ -25,6 +36,8 @@ class Normalizer(nn.Module):
super(Normalizer, self).__init__()
action_statistic = {}
# hard code the action dimension to 20
action_dim = 20
# Process statistics for each robot
for robot_name in action_statistic_dof.keys():
@@ -39,11 +52,17 @@ class Normalizer(nn.Module):
all_dof_delta.extend(action_statistic_dof[robot_name][k]["delta"])
else:
# Use default values if statistics not available
# raise ValueError(f"Statistics not available for {k} of {robot_name}")
logging.warning(
f"Statistics not available for {k} of {robot_name}, using default values"
)
all_dof_min.extend([0.0] * dof_config[k])
all_dof_delta.extend([1.0] * dof_config[k])
all_dof_min = torch.tensor(all_dof_min)
all_dof_delta = torch.tensor(all_dof_delta)
all_dof_min = self._pad_to_action_dim(torch.tensor(all_dof_min), action_dim)
all_dof_delta = self._pad_to_action_dim(
torch.tensor(all_dof_delta), action_dim
)
action_statistic[robot_name]["min"] = all_dof_min
action_statistic[robot_name]["delta"] = all_dof_delta
@@ -61,7 +80,7 @@ class Normalizer(nn.Module):
}
)
def normalize_data(self, xs, dataset_names):
def normalize_data(self, xs, dataset_names, dof_mask=None):
"""
Normalize action data to [-1, 1] range using robot-specific statistics.
@@ -75,10 +94,19 @@ class Normalizer(nn.Module):
new_xs = []
# Filter out multimodal dataset entries
dataset_names = [name for name in dataset_names if name != "x2_multimodal"]
dof_mask = dof_mask if dof_mask is not None else [None] * len(xs)
for x, dataset_name in zip(xs, dataset_names):
for x, dataset_name, mask in zip(xs, dataset_names, dof_mask):
# Apply DOF mask if provided
if mask is not None:
mask = mask[0].bool()
action_space_delta = self.delta[dataset_name][mask]
action_space_min = self.min[dataset_name][mask]
else:
action_space_delta = self.delta[dataset_name]
action_space_min = self.min[dataset_name]
# Apply min-max normalization
x = (x - self.min[dataset_name]) / (self.delta[dataset_name])
x = (x - action_space_min) / (action_space_delta)
# Scale to [-1, 1] range
x = x * 2 - 1
# Clamp to ensure bounds
@@ -210,9 +238,21 @@ class ActionProcessor(nn.Module):
self.hidden_size = config.hidden_size
# Initialize data normalizers for actions and proprioception
self.normalizer_action = Normalizer(action_statistic_dof, config.dof_config)
self.normalizer_action = Normalizer(
action_statistic_dof,
(
config.customized_dof_config
if hasattr(config, "customized_dof_config")
else config.dof_config
),
)
self.normalizer_propri = Normalizer(
action_statistic_dof, config.agent_pos_config
action_statistic_dof,
(
config.customized_agent_pos_config
if hasattr(config, "customized_agent_pos_config")
else config.agent_pos_config
),
)
# Proprioception projection layer (includes history/current state)
wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py
+98 -7
View File
@@ -44,7 +44,9 @@ from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl import (
Qwen2_5_VLSdpaAttention,
)
from wall_x.data.config import ACTION_DATASET_NAMES, MULTIMODAL_DATASET_NAMES
from wall_x.utils.constant import action_statistic_dof
from wall_x.data.utils import load_norm_stats
from pprint import pprint
logger = logging.get_logger(__name__)
@@ -744,10 +746,77 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
config_class = Qwen2_5_VLConfig
_no_split_modules = ["Qwen2_5_VLDecoderLayer_with_MoE", "Qwen2_5_VLVisionBlock"]
@classmethod
def _set_customized_config(cls, config):
"""
Processing norm_stats.json and reconstruct the DoF mapping
"""
dataload_config = config["data"]
if not dataload_config.get("use_lerobot", False):
raise NotImplementedError(
"Not implemented for non-lerobot dataset currently"
)
enable_customized_robot_config = config.get(
"enable_customized_robot_config", False
)
assert (
enable_customized_robot_config
), "enable_customized_robot_config must be true when use lerobot dataset"
customized_dof_config = config["customized_robot_config"][
"customized_dof_config"
]
customized_agent_pos_config = config["customized_robot_config"][
"customized_agent_pos_config"
]
norm_stats_path = config["norm_stats_path"]
norm_stats = load_norm_stats(
norm_stats_path, config["data"]["lerobot_config"]["repo_id"]
)
action_min = norm_stats["action"].min.numpy().tolist()
action_delta = norm_stats["action"].delta.numpy().tolist()
state_min = norm_stats["state"].min.numpy().tolist()
state_delta = norm_stats["state"].delta.numpy().tolist()
name = config["customized_robot_config"]["name"]
dof_key = []
agent_pos_key = []
dof_value = []
agent_pos_value = []
stats_dict = {}
for k, v in customized_dof_config.items():
dof_key.append(k)
dof_value.append(v)
for k, v in customized_agent_pos_config.items():
agent_pos_key.append(k)
agent_pos_value.append(v)
dof_idx = np.array([0] + dof_value).cumsum()
for i in range(len(dof_idx) - 1):
stats_dict[dof_key[i]] = {
"min": action_min[dof_idx[i] : dof_idx[i + 1]],
"delta": action_delta[dof_idx[i] : dof_idx[i + 1]],
}
agent_pos_idx = np.array([0] + agent_pos_value).cumsum()
for i in range(len(agent_pos_idx) - 1):
stats_dict[agent_pos_key[i]] = {
"min": state_min[agent_pos_idx[i] : agent_pos_idx[i + 1]],
"delta": state_delta[agent_pos_idx[i] : agent_pos_idx[i + 1]],
}
action_statistic_dof[name] = stats_dict
print("Customized robot config added")
pprint(action_statistic_dof)
@classmethod
def from_pretrained(
cls,
pretrained_model_path,
train_config,
config_path=None,
processor_path=None,
action_tokenizer_path=None,
@@ -766,7 +835,6 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
Returns:
Qwen2_5_VLMoEForAction: Loaded model instance
"""
# Load model components from pretrained path
config_path = os.path.join(pretrained_model_path, "config.json")
config = cls.config_class.from_pretrained(config_path)
@@ -776,6 +844,18 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
action_tokenizer_path, trust_remote_code=True
)
# Set the customized robot configuration to ensure consistency between cross-embodiment
# representations and the Wall-X action dimensionality.
cls._set_customized_config(train_config)
customized_dof_config = train_config["customized_robot_config"][
"customized_dof_config"
]
customized_agent_pos_config = train_config["customized_robot_config"][
"customized_agent_pos_config"
]
setattr(config, "customized_dof_config", customized_dof_config)
setattr(config, "customized_agent_pos_config", customized_agent_pos_config)
# Initialize model with configuration and processor
model = cls(config, processor=processor, **kwargs)
@@ -789,6 +869,14 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
state_dict = {}
for file in safetensor_files:
sd = load_file(file, device="cpu")
# filter normalizer statistic params
del_keys = []
for key in sd.keys():
if "action_preprocessor.normalizer" in key:
print(f"filter load model weight {key}")
del_keys.append(key)
for key in del_keys:
del sd[key]
state_dict.update(sd)
model.load_state_dict(state_dict, strict=False)
@@ -860,11 +948,14 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
"""
# Create list of fast action token IDs
fast_action_token_list = []
for i in range(self.processor.tokenizer.init_kwargs["action_token_vocab_size"]):
action_token_id = self.processor.tokenizer.convert_tokens_to_ids(
f"<|action_token_{i}|>"
)
fast_action_token_list.append(action_token_id)
if self.use_fast_tokenizer:
for i in range(
self.processor.tokenizer.init_kwargs["action_token_vocab_size"]
):
action_token_id = self.processor.tokenizer.convert_tokens_to_ids(
f"<|action_token_{i}|>"
)
fast_action_token_list.append(action_token_id)
# Get special action token IDs
action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>")
wall_x/trainer/qwen_vl_act_trainer.py
+26 -7
View File
@@ -263,7 +263,6 @@ class QwenVlAct_Trainer:
self.train_dataloader = self.dataset.get_train_dataloader()
self.model.train()
grad_accum_steps = self.config.get("gradient_accumulation_steps", 1)
total = len(self.train_dataloader)
t0 = time.time()
enable_profiling = self.config["profile"]
@@ -341,7 +340,7 @@ class QwenVlAct_Trainer:
self.timers("optimizer").stop()
# Update global step and learning rate after gradient accumulation
if (i + 1) % grad_accum_steps == 0:
if self.accelerator.sync_gradients:
self.lr_scheduler.step()
self.global_step += 1
lr = self.lr_scheduler.get_last_lr()[0]
@@ -522,7 +521,12 @@ class QwenVlAct_Trainer:
if model_type == "wall-oss":
model = Qwen2_5_VLMoEForAction.from_pretrained(
self.config["pretrained_wallx_path"],
**{"use_fast_tokenizer": self.use_fast_tokenizer},
train_config=self.config,
action_tokenizer_path=(
self.config["action_tokenizer_path"]
if self.use_fast_tokenizer
else None
),
)
self.processor = model.processor
model = model.to(torch.bfloat16)
@@ -535,14 +539,15 @@ class QwenVlAct_Trainer:
self.processor = AutoProcessor.from_pretrained(
self.config["pretrained_wallx_path"], use_fast=True
)
new_tokens = ["<|propri|>", "<|action|>"]
self.processor.tokenizer.add_tokens(new_tokens)
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 += [
new_tokens = [
f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size)
]
self.processor.tokenizer.add_tokens(new_tokens)
@@ -557,6 +562,19 @@ class QwenVlAct_Trainer:
action_tokenizer.vocab_size
)
self.processor.action_processor = action_tokenizer
# Set the customized robot configuration to ensure consistency between cross-embodiment
# representations and the Wall-X action dimensionality.
Qwen2_5_VLMoEForAction._set_customized_config(self.config)
customized_dof_config = self.config["customized_robot_config"][
"customized_dof_config"
]
customized_agent_pos_config = self.config["customized_robot_config"][
"customized_agent_pos_config"
]
setattr(config, "customized_dof_config", customized_dof_config)
setattr(config, "customized_agent_pos_config", customized_agent_pos_config)
model = Qwen2_5_VLMoEForAction(
config,
self.use_fast_tokenizer,
@@ -803,7 +821,8 @@ class QwenVlAct_Trainer:
# 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"))
if self.accelerator.is_main_process:
self.processor.save_pretrained(os.path.join(ckpt_path, "processor"))
# Save current iteration steps for dataset resuming
if step != 0:
@@ -845,7 +864,7 @@ class QwenVlAct_Trainer:
# Load full checkpoint including optimizer and scheduler states
self.accelerator.load_state(checkpoint_path)
self.print_rank0(f"Resumed from checkpoint: {checkpoint_path}")
self.print_rank0(f"\033[32mResumed from checkpoint: {checkpoint_path}\033[0m")
def _load_fsdp_state_dict_with_distribute_tensor(self):