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
scripts/draw_openloop_plot.py
+87 -59
View File
@@ -1,22 +1,13 @@
import os
import yaml
import torch
import argparse
from tqdm import tqdm
import matplotlib.pyplot as plt
from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction
from wall_x.data.load_lerobot_dataset import load_test_dataset, get_data_configs
model_path = "path/to/model"
action_tokenizer_path = "path/to/action_tokenizer"
save_dir = "path/to/plot"
model = Qwen2_5_VLMoEForAction.from_pretrained(
model_path, action_tokenizer_path=action_tokenizer_path
)
model.eval()
model = model.to("cuda")
model = model.bfloat16()
def load_config(config_path):
"""Load configuration from YAML file."""
with open(config_path, "r") as f:
@@ -27,60 +18,97 @@ def load_config(config_path):
return config
# get test dataloader
path = "path/to/config"
config = load_config(path)
dataload_config = get_data_configs(config["data"])
lerobot_config = dataload_config.get("lerobot_config", {})
dataset = load_test_dataset(config, lerobot_config, seed=42)
dataloader = dataset.get_dataloader()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--pred_horizon", type=int, default=32)
parser.add_argument("--origin_action_dim", type=int, default=7)
args = parser.parse_args()
total_frames = len(dataloader)
origin_action_dim = args.origin_action_dim
pred_horizon = args.pred_horizon
pred_horizon = 32
action_dim = 14
gt_traj = torch.zeros((total_frames, action_dim))
pred_traj = torch.zeros((total_frames, action_dim))
# get train config
model_path = "/path/to/model"
action_tokenizer_path = "/path/to/action/tokenizer"
save_dir = "/path/to/save/dir"
path = "/path/to/train/config"
config = load_config(path)
for idx, batch in enumerate(dataloader):
if idx % pred_horizon == 0 and idx + pred_horizon < total_frames:
batch = batch.to("cuda")
with torch.no_grad():
outputs = model(
**batch,
action_dim=action_dim,
pred_horizon=pred_horizon,
mode="predict",
predict_mode="fast",
# load model with customized robot config
model = Qwen2_5_VLMoEForAction.from_pretrained(
model_path, train_config=config, action_tokenizer_path=action_tokenizer_path
)
model.eval()
model = model.to("cuda")
model = model.bfloat16()
# get test dataloader
dataload_config = get_data_configs(config["data"])
lerobot_config = dataload_config.get("lerobot_config", {})
dataset = load_test_dataset(config, lerobot_config, seed=42)
dataloader = dataset.get_dataloader()
total_frames = len(dataloader)
predict_mode = "fast" if config.get("use_fast_tokenizer", False) else "diffusion"
action_dim = 20 if predict_mode == "diffusion" else origin_action_dim
gt_traj = torch.zeros((total_frames, origin_action_dim))
pred_traj = torch.zeros((total_frames, origin_action_dim))
# use tqdm to show the progress
for idx, batch in tqdm(
enumerate(dataloader), total=total_frames, desc="predicting"
):
if idx % pred_horizon == 0 and idx + pred_horizon < total_frames:
batch = batch.to("cuda")
with torch.no_grad():
outputs = model(
**batch,
action_dim=action_dim,
pred_horizon=pred_horizon,
mode="predict",
predict_mode=predict_mode,
)
pred_traj[idx : idx + pred_horizon] = (
outputs["predict_action"][:, :, :origin_action_dim]
.detach()
.cpu()
.squeeze(0)
)
# Denormalize ground truth actions
gt_action_chunk = batch["action_chunk"][:, :, :origin_action_dim]
dof_mask = batch["dof_mask"].to(gt_action_chunk.dtype)
denormalized_gt = (
model.action_preprocessor.normalizer_action.unnormalize_data(
gt_action_chunk,
[lerobot_config.get("repo_id", "physical-intelligence/libero")],
dof_mask,
).squeeze(0)
)
pred_traj[idx : idx + pred_horizon] = outputs["predict_action"].detach().cpu()
gt_traj[idx : idx + pred_horizon] = denormalized_gt.detach().cpu()
# Denormalize ground truth actions
gt_action_chunk = batch["action_chunk"][:, :, :action_dim]
dof_mask = batch["dof_mask"].to(gt_action_chunk.dtype)
denormalized_gt = model.action_preprocessor.normalizer_action.unnormalize_data(
gt_action_chunk, ["x2_normal"], dof_mask
)
gt_traj[idx : idx + pred_horizon] = denormalized_gt.detach().cpu()
gt_traj_np = gt_traj.numpy()
pred_traj_np = pred_traj.numpy()
timesteps = gt_traj.shape[0]
gt_traj_np = gt_traj.numpy()
pred_traj_np = pred_traj.numpy()
fig, axs = plt.subplots(
origin_action_dim, 1, figsize=(15, 5 * origin_action_dim), sharex=True
)
fig.suptitle("Action Comparison for lerobot", fontsize=16)
timesteps = gt_traj.shape[0]
for i in range(origin_action_dim):
axs[i].plot(range(timesteps), gt_traj_np[:, i], label="Ground Truth")
axs[i].plot(range(timesteps), pred_traj_np[:, i], label="Prediction")
axs[i].set_ylabel(f"Action Dim {i+1}")
axs[i].legend()
axs[i].grid(True)
fig, axs = plt.subplots(action_dim, 1, figsize=(15, 5 * action_dim), sharex=True)
fig.suptitle("Action Comparison for lerobot", fontsize=16)
for i in range(action_dim):
axs[i].plot(range(timesteps), gt_traj_np[:, i], label="Ground Truth")
axs[i].plot(range(timesteps), pred_traj_np[:, i], label="Prediction")
axs[i].set_ylabel(f"Action Dim {i+1}")
axs[i].legend()
axs[i].grid(True)
axs[-1].set_xlabel("Timestep")
plt.tight_layout(rect=[0, 0.03, 1, 0.95])
os.makedirs(save_dir, exist_ok=True)
plt.savefig(os.path.join(save_dir, "lerobot_comparison.png"))
plt.close()
axs[-1].set_xlabel("Timestep")
plt.tight_layout(rect=[0, 0.03, 1, 0.95])
os.makedirs(save_dir, exist_ok=True)
save_path = os.path.join(save_dir, "lerobot_comparison.png")
plt.savefig(save_path)
print(f"Saved plot to {save_path}")
plt.close()