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

[lint] Update lint (#16)

Browse Source 
* update lint

* update readme

* update ruff lint
This commit is contained in:
Lufang Chen
2025-09-11 13:18:33 +08:00
committed by GitHub
parent a89dce95aa
commit e9332a283d
28 changed files with 2406 additions and 1074 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wall_x/trainer/qwen_vl_act_trainer.py
+271 -118
View File
@@ -12,25 +12,29 @@ from datetime import datetime
from torch.optim import AdamW
from accelerate import Accelerator
from safetensors.torch import load_file
from accelerate.utils import DistributedType
from transformers.optimization import get_cosine_with_min_lr_schedule_with_warmup
from wall_x.utils.timers import Timers
from wall_x.model.qwen2_5_based import Qwen2_5_VLMoEForAction
from wall_x.data.config import ACTION_DATASET_NAMES, MULTIMODAL_DATASET_NAMES
from wall_x.data.load_lerobot_dataset import PreprocessedDataset, get_data_configs, load_lerobot_data
from wall_x.data.load_lerobot_dataset import (
PreprocessedDataset,
get_data_configs,
load_lerobot_data,
)
def timer(func):
"""
Decorator to measure function execution time.
Args:
func: Function to be timed
Returns:
Wrapped function with timing functionality
"""
@wraps(func)
def wrapper(*args, **kwargs):
start_time = time.time()
@@ -40,13 +44,14 @@ def timer(func):
f"\033[92m[current time: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}] Function {func.__name__} took {end_time - start_time:.2f} seconds to execute\033[0m"
)
return result
return wrapper
def print_rank_last(message):
"""
Print message only on the last rank in distributed training.
Args:
message (str): Message to print
"""
@@ -60,7 +65,7 @@ def print_rank_last(message):
def seed_all(seed):
"""
Set random seeds for reproducible training.
Args:
seed (int): Random seed value
"""
@@ -72,11 +77,11 @@ def seed_all(seed):
class QwenVlAct_Trainer:
"""
Vision-Language-Action trainer for Qwen-VL models with robotic action prediction.
This trainer handles multi-modal learning combining vision, language, and action data
for robotic control applications. It supports distributed training, mixed precision,
gradient accumulation, and various optimization strategies including MoE (Mixture of Experts).
Features:
- Multi-modal data processing (vision + language + actions)
- Distributed training with Accelerate
@@ -87,10 +92,17 @@ class QwenVlAct_Trainer:
"""
@timer
def __init__(self, config, logger, accelerator: Accelerator = None, seed=42, data_config_path=None):
def __init__(
self,
config,
logger,
accelerator: Accelerator = None,
seed=42,
data_config_path=None,
):
"""
Initialize the Vision-Language-Action trainer.
Args:
config (dict): Training configuration dictionary containing:
- processor_path (str): Path to data preprocessing processor
@@ -103,7 +115,7 @@ class QwenVlAct_Trainer:
accelerator (Accelerator, optional): Hugging Face Accelerate instance for distributed training
seed (int, optional): Random seed for reproducibility. Defaults to 42.
data_config_path (str, optional): Path to data configuration file
Raises:
ValueError: If required configuration keys are missing
"""
@@ -117,41 +129,51 @@ class QwenVlAct_Trainer:
self.logger = logger
self.accelerator = accelerator
self.seed = seed
# Initialize random seeds for reproducibility
seed_all(self.seed)
# Training state variables
self.start_epoch = 0
self.global_step = 0
self.num_epoch = self.config["num_epoch"]
self.initial_step = 0
# Data and model configuration
self.dataload_config = get_data_configs(self.config["data"])
self.data_config_path = data_config_path
self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False)
# Load model and initialize training components
self.load_model()
self.action_dim = sum(self.config["dof_config"].values())
# Distributed training setup
self.rank = self.accelerator.process_index
self.world_size = self.accelerator.num_processes
print(f"rank {self.accelerator.process_index} after load model memory usage: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB", flush=True)
print(
f"rank {self.accelerator.process_index} after load model memory usage: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB",
flush=True,
)
# Load training data
self.load_qact_data()
print(f"rank {self.accelerator.process_index} after load qact data usage: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB", flush=True)
print(
f"rank {self.accelerator.process_index} after load qact data usage: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB",
flush=True,
)
# Resume from checkpoint if specified
if "resume" in self.config:
self.resume_from_checkpoint()
# Initialize special token IDs
self.propri_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|propri|>")
self.action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>")
self.propri_token_id = self.processor.tokenizer.convert_tokens_to_ids(
"<|propri|>"
)
self.action_token_id = self.processor.tokenizer.convert_tokens_to_ids(
"<|action|>"
)
# Initialize evaluation metrics
self.base_l1_loss = None
@@ -162,12 +184,14 @@ class QwenVlAct_Trainer:
# Adjust global step if resuming from checkpoint
if self.initial_step != 0:
self.global_step = self.initial_step // self.config.get("gradient_accumulation_steps", 1)
self.global_step = self.initial_step // self.config.get(
"gradient_accumulation_steps", 1
)
def print_rank0(self, msg, flush=True):
"""
Print message only on rank 0 (main process).
Args:
msg: Message to print
flush (bool): Whether to flush output buffer
@@ -178,7 +202,7 @@ class QwenVlAct_Trainer:
def fit(self):
"""
Main training loop executing multiple epochs with validation.
Handles the complete training process including:
- Training loop execution
- Validation after each epoch
@@ -186,9 +210,11 @@ class QwenVlAct_Trainer:
- Memory cleanup
"""
self.accelerator.wait_for_everyone()
# Optional validation before training starts
if self.config.get("resume", None) is not None and self.config["resume"].get("validate_first", False):
if self.config.get("resume", None) is not None and self.config["resume"].get(
"validate_first", False
):
self.val_loop()
self.accelerator.wait_for_everyone()
@@ -196,24 +222,24 @@ class QwenVlAct_Trainer:
for epoch in range(self.start_epoch, self.num_epoch):
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()
# Memory cleanup
gc.collect()
def train_loop(self, epoch):
"""
Execute training for a single epoch.
Args:
epoch (int): Current epoch number
Handles:
- Data loading and batching
- Forward/backward passes
@@ -227,7 +253,9 @@ class QwenVlAct_Trainer:
if getattr(self, "train_dataloader", None) is not None:
self.train_sampler.set_epoch(epoch)
else:
self.train_dataloader, self.train_sampler = self.dataset.get_train_dataloader()
self.train_dataloader, self.train_sampler = (
self.dataset.get_train_dataloader()
)
self.train_sampler.set_epoch(epoch)
else:
self.train_dataloader = self.dataset.get_train_dataloader()
@@ -236,16 +264,23 @@ class QwenVlAct_Trainer:
grad_accum_steps = self.config.get("gradient_accumulation_steps", 1)
total = len(self.train_dataloader)
t0 = time.time()
enable_profiling = self.config['profile']
enable_profiling = self.config["profile"]
# Optional PyTorch profiler for performance analysis
if enable_profiling:
profiler = torch.profiler.profile(
activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
schedule=torch.profiler.schedule(wait=self.config['profile_wait_iters'],
warmup=self.config['profile_warmup_iters'],
active=self.config['profile_active_iters']),
on_trace_ready=torch.profiler.tensorboard_trace_handler(self.config['profile_save_path'], worker_name="worker0"),
activities=[
torch.profiler.ProfilerActivity.CPU,
torch.profiler.ProfilerActivity.CUDA,
],
schedule=torch.profiler.schedule(
wait=self.config["profile_wait_iters"],
warmup=self.config["profile_warmup_iters"],
active=self.config["profile_active_iters"],
),
on_trace_ready=torch.profiler.tensorboard_trace_handler(
self.config["profile_save_path"], worker_name="worker0"
),
record_shapes=True,
profile_memory=True,
with_stack=True,
@@ -253,7 +288,7 @@ 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)
@@ -261,28 +296,38 @@ class QwenVlAct_Trainer:
for i, batch in enumerate(self.train_dataloader, self.initial_step):
# Move batch to device
if isinstance(self.dataset, PreprocessedDataset):
batch = {k: v.to(self.accelerator.device, non_blocking=True) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
batch = {
k: (
v.to(self.accelerator.device, non_blocking=True)
if isinstance(v, torch.Tensor)
else v
)
for k, v in batch.items()
}
self.timers("data-load").stop()
with self.accelerator.accumulate(self.model):
# Forward pass
self.timers("forward-compute", log_level=0).start(barrier=False)
outputs = self.model(**batch, mode="train")
self.timers("forward-compute").stop()
loss = outputs.loss
# Check for NaN loss
if torch.isnan(loss):
print(f"Warning: NaN loss detected in epoch: {epoch}, step: {i}", flush=True)
print(
f"Warning: NaN loss detected in epoch: {epoch}, step: {i}",
flush=True,
)
continue
# Backward pass
self.timers("backward-compute", log_level=0).start(barrier=False)
self.accelerator.backward(loss)
self.timers("backward-compute").stop()
# Gradient clipping
total_norm = self.accelerator.clip_grad_norm_(
self.model.parameters(), self.config.get("max_grad_norm", 1.0)
@@ -299,33 +344,79 @@ class QwenVlAct_Trainer:
self.lr_scheduler.step()
self.global_step += 1
lr = self.lr_scheduler.get_last_lr()[0]
# Gather loss across all processes for logging
train_loss = self.accelerator.gather(loss.detach()).mean().item()
train_loss = (
self.accelerator.gather(loss.detach()).mean().item()
)
_log_dict = {
"lr": lr,
"train_loss": train_loss,
}
# Log component losses
if "cross_entropy_loss" in outputs and outputs.cross_entropy_loss is not None:
_log_dict["cross_entropy_loss"] = self.accelerator.gather(outputs.cross_entropy_loss.detach()).mean().item()
if (
"cross_entropy_loss" in outputs
and outputs.cross_entropy_loss is not None
):
_log_dict["cross_entropy_loss"] = (
self.accelerator.gather(
outputs.cross_entropy_loss.detach()
)
.mean()
.item()
)
if "flow_loss" in outputs and outputs.flow_loss is not None:
_log_dict["flow_loss"] = self.accelerator.gather(outputs.flow_loss.detach()).mean().item()
_log_dict["flow_loss"] = (
self.accelerator.gather(outputs.flow_loss.detach())
.mean()
.item()
)
# Log per-dataset channel losses
if "channel_loss_dict" in outputs and outputs.channel_loss_dict is not None:
for dataset_name_i in ACTION_DATASET_NAMES + MULTIMODAL_DATASET_NAMES:
count_sum = self.accelerator.gather(outputs.channel_loss_count_dict[dataset_name_i]).sum().item()
if (
"channel_loss_dict" in outputs
and outputs.channel_loss_dict is not None
):
for dataset_name_i in (
ACTION_DATASET_NAMES + MULTIMODAL_DATASET_NAMES
):
count_sum = (
self.accelerator.gather(
outputs.channel_loss_count_dict[dataset_name_i]
)
.sum()
.item()
)
if count_sum > 0:
channel_loss = self.accelerator.gather(outputs.channel_loss_dict[dataset_name_i].detach()).sum().item() / count_sum
_log_dict[f"channel_loss_{dataset_name_i}"] = channel_loss
channel_loss = (
self.accelerator.gather(
outputs.channel_loss_dict[
dataset_name_i
].detach()
)
.sum()
.item()
/ count_sum
)
_log_dict[f"channel_loss_{dataset_name_i}"] = (
channel_loss
)
# Log action accuracy for fast tokenizer
if "action_accuracy" in outputs.channel_loss_dict and self.use_fast_tokenizer:
if (
"action_accuracy" in outputs.channel_loss_dict
and self.use_fast_tokenizer
):
_log_dict["action_accuracy"] = (
self.accelerator.gather(outputs.channel_loss_dict["action_accuracy"].detach()).mean().item()
self.accelerator.gather(
outputs.channel_loss_dict[
"action_accuracy"
].detach()
)
.mean()
.item()
)
# Log metrics
@@ -334,23 +425,26 @@ class QwenVlAct_Trainer:
# Log gradient norm
if self.logger is not None and self.accelerator.sync_gradients:
self.logger.log({"total_norm": total_norm}, step=self.global_step)
self.logger.log(
{"total_norm": total_norm}, step=self.global_step
)
self.timers("interval-time").stop()
# Setup timers for next iteration
if i < len(self.train_dataloader) - 1:
self.timers("interval-time", log_level=0).start(barrier=False)
self.timers("data-load", log_level=0).start(barrier=False)
# Periodic logging
t1 = time.time()
if i % 1 == 0:
lr = self.lr_scheduler.get_last_lr()[0]
self.training_log(epoch, self.num_epoch, i, total, loss, lr, t1 - t0)
self.training_log(
epoch, self.num_epoch, i, total, loss, lr, t1 - t0
)
t0 = time.time()
if enable_profiling:
profiler.step()
@@ -362,7 +456,7 @@ class QwenVlAct_Trainer:
def val_loop(self):
"""
Execute validation loop with gradient computation disabled.
Evaluates model performance on validation set and logs validation loss.
"""
# Initialize validation dataloader
@@ -374,34 +468,45 @@ class QwenVlAct_Trainer:
self.model.eval()
self.val_loss = 0
# Validation loop
for i, batch in enumerate(
tqdm(self.val_dataloader, desc="Validating", total=len(self.val_dataloader),
disable=not self.accelerator.is_main_process)
tqdm(
self.val_dataloader,
desc="Validating",
total=len(self.val_dataloader),
disable=not self.accelerator.is_main_process,
)
):
if isinstance(self.dataset, PreprocessedDataset):
batch = {k: v.to(self.accelerator.device, non_blocking=True) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
batch = {
k: (
v.to(self.accelerator.device, non_blocking=True)
if isinstance(v, torch.Tensor)
else v
)
for k, v in batch.items()
}
with torch.no_grad():
outputs = self.model(**batch, mode="train")
loss = outputs.loss
self.val_loss += self.accelerator.gather(loss.detach()).mean().item()
# Calculate average validation loss
self.val_loss /= len(self.val_dataloader)
# Log validation metrics
if self.logger is not None:
self.logger.log({"val_loss": self.val_loss}, step=self.global_step)
self.model.train()
@timer
def load_model(self):
"""
Load and configure the Vision-Language-Action model.
Handles:
- Model loading from pretrained weights
- Processor initialization
@@ -411,8 +516,8 @@ class QwenVlAct_Trainer:
"""
# Load pretrained model
model = Qwen2_5_VLMoEForAction.from_pretrained(
self.config["pretrained_wallx_path"],
**{"use_fast_tokenizer": self.use_fast_tokenizer}
self.config["pretrained_wallx_path"],
**{"use_fast_tokenizer": self.use_fast_tokenizer},
)
self.processor = model.processor
model = model.to(torch.bfloat16)
@@ -428,7 +533,7 @@ class QwenVlAct_Trainer:
moe_params.append(param)
param_groups = [{"params": moe_params, "lr": self.config["learning_rate"]}]
self.optimizer = AdamW(param_groups, weight_decay=0.1)
elif "action_expert_learning_rate" in self.config:
# Separate learning rates for VLM and action expert parameters
moe_params = []
@@ -442,17 +547,26 @@ class QwenVlAct_Trainer:
# Configure parameter groups
if self.config.get("train_action_expert_only", False):
self.print_rank0("Training action expert only", flush=True)
param_groups = [{"params": moe_params, "lr": self.config["action_expert_learning_rate"]}]
param_groups = [
{
"params": moe_params,
"lr": self.config["action_expert_learning_rate"],
}
]
else:
param_groups = [
{"params": vlm_params, "lr": self.config["learning_rate"]},
{"params": moe_params, "lr": self.config["action_expert_learning_rate"]},
{
"params": moe_params,
"lr": self.config["action_expert_learning_rate"],
},
]
self.optimizer = AdamW(param_groups, weight_decay=0.1)
self.print_rank0(
f"Setting MoE learning rate to {self.config['action_expert_learning_rate']}, "
f"VLM learning rate to {self.config['learning_rate']}", flush=True
f"VLM learning rate to {self.config['learning_rate']}",
flush=True,
)
else:
# Standard optimizer configuration
@@ -479,9 +593,13 @@ class QwenVlAct_Trainer:
if hasattr(model, "enable_input_require_grads"):
self.model.enable_input_require_grads()
else:
def make_inputs_require_grad(module, input, output):
output.requires_grad_(True)
self.model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
self.model.get_input_embeddings().register_forward_hook(
make_inputs_require_grad
)
# Prepare model, optimizer, and scheduler for distributed training
self.model, self.optimizer, self.lr_scheduler = self.accelerator.prepare(
@@ -492,7 +610,7 @@ class QwenVlAct_Trainer:
def load_qact_data(self):
"""
Load and configure training data for Vision-Language-Action learning.
Supports LeRobot dataset format and handles distributed data loading
across multiple processes.
"""
@@ -511,18 +629,20 @@ class QwenVlAct_Trainer:
def load_qwen_pretrain_weight(self, model, pretrain_weight_path):
"""
Load pretrained Qwen weights with MoE adaptation.
Args:
model: Model instance to load weights into
pretrain_weight_path (str): Path to pretrained weight files
Returns:
Model with loaded pretrained weights
Handles weight key renaming for MoE architecture compatibility.
"""
# Load all safetensors files
weight_files = sorted([f for f in os.listdir(pretrain_weight_path) if f.endswith(".safetensors")])
weight_files = sorted(
[f for f in os.listdir(pretrain_weight_path) if f.endswith(".safetensors")]
)
merged_weights = {}
# Merge weights from all files
@@ -534,19 +654,31 @@ class QwenVlAct_Trainer:
# Rename weights for MoE compatibility
renamed_weights = {}
for key, value in merged_weights.items():
if key.startswith("model.layers") and "mlp." in key and model.config.mlp_moe:
if (
key.startswith("model.layers")
and "mlp." in key
and model.config.mlp_moe
):
# Rename MLP weights for MoE structure
layer_num = key.split(".layers.")[1].split(".mlp")[0]
new_key = key.replace(f"layers.{layer_num}.mlp.", f"layers.{layer_num}.moe.experts.0.")
new_key = key.replace(
f"layers.{layer_num}.mlp.", f"layers.{layer_num}.moe.experts.0."
)
renamed_weights[new_key] = value
elif key.startswith("model.layers") and "self_attn." in key and model.config.attention_moe:
elif (
key.startswith("model.layers")
and "self_attn." in key
and model.config.attention_moe
):
# Rename attention weights for MoE structure
layer_num = key.split(".layers.")[1].split(".self_attn")[0]
proj_types = ["q_proj", "k_proj", "v_proj", "o_proj"]
for proj in proj_types:
if proj in key:
new_key = key.replace(f"layers.{layer_num}.self_attn.{proj}",
f"layers.{layer_num}.self_attn.{proj}_experts.0")
new_key = key.replace(
f"layers.{layer_num}.self_attn.{proj}",
f"layers.{layer_num}.self_attn.{proj}_experts.0",
)
renamed_weights[new_key] = value
break
else:
@@ -560,10 +692,19 @@ class QwenVlAct_Trainer:
return model
def training_log(self, current_epoch, total_epoch, current_train_iter, total_train_iter, loss, lr, time_per_step):
def training_log(
self,
current_epoch,
total_epoch,
current_train_iter,
total_train_iter,
loss,
lr,
time_per_step,
):
"""
Log training progress and performance metrics.
Args:
current_epoch (int): Current epoch number
total_epoch (int): Total number of epochs
@@ -573,26 +714,32 @@ class QwenVlAct_Trainer:
lr (float): Current learning rate
time_per_step (float): Time taken for current step
"""
timers_to_log = ["interval-time", "data-load", "forward-compute", "backward-compute", "optimizer"]
timers_to_log = [
"interval-time",
"data-load",
"forward-compute",
"backward-compute",
"optimizer",
]
log_string = f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]"
log_string += " epoch {:3d}/{:3d} |".format(current_epoch, total_epoch)
log_string += " iter {:6d}/{:6d} |".format(current_train_iter, total_train_iter)
log_string += " loss {:.6f} |".format(loss)
log_string += " lr {:.6f} |".format(lr)
log_string += " time_per_step_avg {:.6f}s |".format(time_per_step)
print_rank_last(log_string)
self.timers.log(timers_to_log, normalizer=1)
def save_checkpoint(self, epoch, step=0):
"""
Save training checkpoint.
Args:
epoch (int): Current epoch number
step (int, optional): Current step number. Defaults to 0.
Saves model state, optimizer state, and training progress information.
"""
save_path = self.config["save_path"]
@@ -600,7 +747,7 @@ class QwenVlAct_Trainer:
ckpt_path = f"{save_path}/{epoch}"
else:
ckpt_path = f"{save_path}/{epoch}_{step}"
self.accelerator.save_state(ckpt_path)
# Save current iteration steps for dataset resuming
@@ -608,14 +755,16 @@ class QwenVlAct_Trainer:
_rank = self.accelerator.process_index
if isinstance(self.dataset, PreprocessedDataset):
torch.save(
{"epoch": epoch, "step": step},
os.path.join(ckpt_path, f"epoch_{epoch}_step_{step}_rank_{_rank}.pth")
{"epoch": epoch, "step": step},
os.path.join(
ckpt_path, f"epoch_{epoch}_step_{step}_rank_{_rank}.pth"
),
)
def resume_from_checkpoint(self):
"""
Resume training from a saved checkpoint.
Handles both full checkpoint loading and model-only loading based on configuration.
"""
checkpoint_path = self.config["resume"]["ckpt"]
@@ -624,45 +773,47 @@ class QwenVlAct_Trainer:
# Load only model weights
ckpt_path = self.config["resume"]["ckpt"] + "/model.safetensors"
state_dict = load_file(ckpt_path, device="cpu")
# Add module prefix if needed for distributed training
new_state_dict = {}
for key in state_dict:
if not key.startswith("module."):
new_key = "module." + key
new_state_dict[new_key] = state_dict[key]
err = self.model.load_state_dict(new_state_dict, strict=False)
self.model.load_state_dict(new_state_dict, strict=False)
else:
# Load full checkpoint including optimizer and scheduler states
self.accelerator.load_state(checkpoint_path)
self.print_rank0(f"Resumed from checkpoint: {checkpoint_path}")
def log_l1_details(self, all_label, all_pred, all_task, all_dof_mask):
"""
Log detailed L1 loss metrics by degrees of freedom.
Args:
all_label (torch.Tensor): Ground truth action labels
all_pred (torch.Tensor): Predicted actions
all_task (list): Task identifiers
all_dof_mask (torch.Tensor): Degrees of freedom mask
Computes and logs L1 loss for each DOF component separately for detailed analysis.
"""
all_task = all_task[:len(all_label)]
all_task = all_task[: len(all_label)]
# Apply DOF mask
all_label = all_label * all_dof_mask
all_pred = all_pred * all_dof_mask
# Compute baseline L1 loss (predict mean action)
if self.base_l1_loss is None:
mean_action = all_label.mean(dim=0)
self.base_l1_loss = nn.functional.l1_loss(all_label, mean_action)
self.logger.log({"base_l1_loss": self.base_l1_loss.item()}, step=self.global_step)
self.logger.log(
{"base_l1_loss": self.base_l1_loss.item()}, step=self.global_step
)
# Log L1 loss for each DOF component
start_idx = 0
@@ -672,8 +823,10 @@ class QwenVlAct_Trainer:
dof_label = all_label[:, :, start_idx:end_idx]
dof_pred = all_pred[:, :, start_idx:end_idx]
dof_l1 = nn.functional.l1_loss(dof_pred, dof_label)
self.print_rank0(f"DOF {dof}, L1 loss: {dof_l1.item()}", flush=True)
self.logger.log({f"detail/l1_loss_{dof}": dof_l1.item()}, step=self.global_step)
start_idx = end_idx
self.logger.log(
{f"detail/l1_loss_{dof}": dof_l1.item()}, step=self.global_step
)
start_idx = end_idx