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/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|>")