feat(train): 跑通训练脚本
This commit is contained in:
gouhanke
@@ -7,115 +7,27 @@ from tqdm import tqdm
|
||||
from omegaconf import DictConfig, OmegaConf
|
||||
from torch.utils.data import DataLoader
|
||||
from torch.optim import AdamW
|
||||
from pathlib import Path
|
||||
|
||||
# 确保导入路径正确
|
||||
# Ensure correct import path
|
||||
sys.path.append(os.getcwd())
|
||||
|
||||
from roboimi.vla.agent import VLAAgent
|
||||
from hydra.utils import instantiate
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
@hydra.main(version_base=None, config_path="../../vla/conf", config_name="config")
|
||||
def main(cfg: DictConfig):
|
||||
print(OmegaConf.to_yaml(cfg))
|
||||
log.info(f"🚀 Starting VLA Training with Real Data (Device: {cfg.train.device})")
|
||||
|
||||
# --- 1. 实例化 Dataset & DataLoader ---
|
||||
# Hydra 根据 conf/data/custom_hdf5.yaml 实例化类
|
||||
dataset = instantiate(cfg.data)
|
||||
|
||||
dataloader = DataLoader(
|
||||
dataset,
|
||||
batch_size=cfg.train.batch_size,
|
||||
shuffle=True,
|
||||
num_workers=cfg.train.num_workers,
|
||||
pin_memory=(cfg.train.device != "cpu")
|
||||
)
|
||||
log.info(f"✅ Dataset loaded. Size: {len(dataset)}")
|
||||
|
||||
# --- 2. 实例化 Agent ---
|
||||
agent: VLAAgent = instantiate(cfg.agent)
|
||||
agent.to(cfg.train.device)
|
||||
agent.train()
|
||||
|
||||
optimizer = AdamW(agent.parameters(), lr=cfg.train.lr)
|
||||
|
||||
# --- 3. Training Loop ---
|
||||
# 使用一个无限迭代器或者 epoch 循环
|
||||
data_iter = iter(dataloader)
|
||||
pbar = tqdm(range(cfg.train.max_steps), desc="Training")
|
||||
|
||||
for step in pbar:
|
||||
try:
|
||||
batch = next(data_iter)
|
||||
except StopIteration:
|
||||
#而在 epoch 结束时重新开始
|
||||
data_iter = iter(dataloader)
|
||||
batch = next(data_iter)
|
||||
|
||||
# Move to device
|
||||
# 注意：这里需要递归地将字典里的 tensor 移到 GPU
|
||||
batch = recursive_to_device(batch, cfg.train.device)
|
||||
|
||||
# --- 4. Adapter Layer (适配层) ---
|
||||
# Dataset 返回的是具体的相机 key (如 'agentview_image' 或 'top')
|
||||
# Agent 期望的是通用的 'image'
|
||||
# 我们在这里做一个映射，模拟多模态融合前的处理
|
||||
|
||||
# 假设我们只用配置里的第一个 key 作为主视觉
|
||||
# primary_cam_key = cfg.data.obs_keys[0]
|
||||
|
||||
# Dataset 返回 shape: (B, Obs_Horizon, C, H, W)
|
||||
# DebugBackbone 期望: (B, C, H, W) 或者 (B, Seq, Dim)
|
||||
# 这里我们取 Obs_Horizon 的最后一帧 (Current Frame)
|
||||
# input_img = batch['obs'][primary_cam_key][:, -1, :, :, :]
|
||||
|
||||
# agent_input = {
|
||||
# "obs": {
|
||||
# "image": input_img,
|
||||
# "text": batch["language"] # 传递语言指令
|
||||
# },
|
||||
# "actions": batch["actions"] # (B, Chunk, Dim)
|
||||
# }
|
||||
agent_input = {
|
||||
"action": batch["action"],
|
||||
"qpos": batch["qpos"],
|
||||
"images": {}
|
||||
}
|
||||
|
||||
for cam_name in cfg.data.camera_names:
|
||||
key = f"image_{cam_name}"
|
||||
agent_input["images"][cam_name] = batch[key].squeeze(1)
|
||||
|
||||
# --- 5. Forward & Backward ---
|
||||
outputs = agent(agent_input)
|
||||
|
||||
# 处理 Loss 掩码 (如果在真实训练中，需要在这里应用 action_mask)
|
||||
# 目前 DebugHead 内部直接算了 MSE，还没用 mask，我们在下一阶段优化 Policy 时加上
|
||||
loss = outputs['loss']
|
||||
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
if step % cfg.train.log_freq == 0:
|
||||
pbar.set_postfix({"loss": f"{loss.item():.4f}"})
|
||||
|
||||
log.info("✅ Training Loop with Real HDF5 Finished!")
|
||||
|
||||
# --- 6. Save Checkpoint ---
|
||||
save_dir = "checkpoints"
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
save_path = os.path.join(save_dir, "vla_model_final.pt")
|
||||
|
||||
# 保存整个 Agent 的 state_dict
|
||||
torch.save(agent.state_dict(), save_path)
|
||||
log.info(f"💾 Model saved to {save_path}")
|
||||
|
||||
log.info("✅ Training Loop Finished!")
|
||||
|
||||
def recursive_to_device(data, device):
|
||||
"""
|
||||
Recursively move nested dictionaries/lists of tensors to specified device.
|
||||
|
||||
Args:
|
||||
data: Dictionary, list, or tensor
|
||||
device: Target device (e.g., 'cuda', 'cpu')
|
||||
|
||||
Returns:
|
||||
Data structure with all tensors moved to device
|
||||
"""
|
||||
if isinstance(data, torch.Tensor):
|
||||
return data.to(device)
|
||||
elif isinstance(data, dict):
|
||||
@@ -124,5 +36,193 @@ def recursive_to_device(data, device):
|
||||
return [recursive_to_device(v, device) for v in data]
|
||||
return data
|
||||
|
||||
|
||||
@hydra.main(version_base=None, config_path="../../vla/conf", config_name="config")
|
||||
def main(cfg: DictConfig):
|
||||
"""
|
||||
VLA Training Script with ResNet Backbone and Diffusion Policy.
|
||||
|
||||
This script:
|
||||
1. Loads dataset from HDF5 files
|
||||
2. Instantiates VLAAgent with ResNet vision encoder
|
||||
3. Trains diffusion-based action prediction
|
||||
4. Saves checkpoints periodically
|
||||
"""
|
||||
|
||||
# Print configuration
|
||||
print("=" * 80)
|
||||
print("VLA Training Configuration:")
|
||||
print("=" * 80)
|
||||
print(OmegaConf.to_yaml(cfg))
|
||||
print("=" * 80)
|
||||
|
||||
log.info(f"🚀 Starting VLA Training (Device: {cfg.train.device})")
|
||||
|
||||
# Create checkpoint directory
|
||||
checkpoint_dir = Path("checkpoints")
|
||||
checkpoint_dir.mkdir(exist_ok=True)
|
||||
|
||||
# =========================================================================
|
||||
# 1. Instantiate Dataset & DataLoader
|
||||
# =========================================================================
|
||||
log.info("📦 Loading dataset...")
|
||||
try:
|
||||
dataset = instantiate(cfg.data)
|
||||
log.info(f"✅ Dataset loaded successfully. Total samples: {len(dataset)}")
|
||||
except Exception as e:
|
||||
log.error(f"❌ Failed to load dataset: {e}")
|
||||
raise
|
||||
|
||||
dataloader = DataLoader(
|
||||
dataset,
|
||||
batch_size=cfg.train.batch_size,
|
||||
shuffle=True,
|
||||
num_workers=cfg.train.num_workers,
|
||||
pin_memory=(cfg.train.device != "cpu"),
|
||||
drop_last=True # Drop incomplete batches for stable training
|
||||
)
|
||||
log.info(f"✅ DataLoader created. Batches per epoch: {len(dataloader)}")
|
||||
|
||||
# =========================================================================
|
||||
# 2. Instantiate VLA Agent
|
||||
# =========================================================================
|
||||
log.info("🤖 Initializing VLA Agent...")
|
||||
try:
|
||||
agent = instantiate(cfg.agent)
|
||||
agent.to(cfg.train.device)
|
||||
agent.train()
|
||||
log.info(f"✅ Agent initialized and moved to {cfg.train.device}")
|
||||
|
||||
# Count parameters
|
||||
total_params = sum(p.numel() for p in agent.parameters())
|
||||
trainable_params = sum(p.numel() for p in agent.parameters() if p.requires_grad)
|
||||
log.info(f"📊 Total parameters: {total_params:,}")
|
||||
log.info(f"📊 Trainable parameters: {trainable_params:,}")
|
||||
|
||||
except Exception as e:
|
||||
log.error(f"❌ Failed to initialize agent: {e}")
|
||||
raise
|
||||
|
||||
# =========================================================================
|
||||
# 3. Setup Optimizer
|
||||
# =========================================================================
|
||||
optimizer = AdamW(agent.parameters(), lr=cfg.train.lr, weight_decay=1e-5)
|
||||
log.info(f"🔧 Optimizer: AdamW (lr={cfg.train.lr})")
|
||||
|
||||
# =========================================================================
|
||||
# 4. Training Loop
|
||||
# =========================================================================
|
||||
log.info("🏋️ Starting training loop...")
|
||||
|
||||
data_iter = iter(dataloader)
|
||||
pbar = tqdm(range(cfg.train.max_steps), desc="Training", ncols=100)
|
||||
|
||||
best_loss = float('inf')
|
||||
|
||||
for step in pbar:
|
||||
try:
|
||||
batch = next(data_iter)
|
||||
except StopIteration:
|
||||
# Restart iterator when epoch ends
|
||||
data_iter = iter(dataloader)
|
||||
batch = next(data_iter)
|
||||
|
||||
# =====================================================================
|
||||
# Move batch to device
|
||||
# =====================================================================
|
||||
batch = recursive_to_device(batch, cfg.train.device)
|
||||
|
||||
# =====================================================================
|
||||
# Prepare agent input
|
||||
# =====================================================================
|
||||
# Dataset returns: {action, qpos, image_<cam_name>, ...}
|
||||
# Agent expects: {images: dict, qpos: tensor, action: tensor}
|
||||
|
||||
# Extract images into a dictionary
|
||||
images = {}
|
||||
for cam_name in cfg.data.camera_names:
|
||||
key = f"image_{cam_name}"
|
||||
if key in batch:
|
||||
images[cam_name] = batch[key] # (B, obs_horizon, C, H, W)
|
||||
|
||||
# Prepare agent input
|
||||
agent_input = {
|
||||
'images': images, # Dict of camera images
|
||||
'qpos': batch['qpos'], # (B, obs_horizon, obs_dim)
|
||||
'action': batch['action'] # (B, pred_horizon, action_dim)
|
||||
}
|
||||
|
||||
# =====================================================================
|
||||
# Forward pass & compute loss
|
||||
# =====================================================================
|
||||
try:
|
||||
loss = agent.compute_loss(agent_input)
|
||||
except Exception as e:
|
||||
log.error(f"❌ Forward pass failed at step {step}: {e}")
|
||||
raise
|
||||
|
||||
# =====================================================================
|
||||
# Backward pass & optimization
|
||||
# =====================================================================
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
|
||||
# Gradient clipping for stable training
|
||||
torch.nn.utils.clip_grad_norm_(agent.parameters(), max_norm=1.0)
|
||||
|
||||
optimizer.step()
|
||||
|
||||
# =====================================================================
|
||||
# Logging
|
||||
# =====================================================================
|
||||
if step % cfg.train.log_freq == 0:
|
||||
pbar.set_postfix({
|
||||
"loss": f"{loss.item():.4f}",
|
||||
"best_loss": f"{best_loss:.4f}"
|
||||
})
|
||||
log.info(f"Step {step}/{cfg.train.max_steps} | Loss: {loss.item():.4f}")
|
||||
|
||||
# =====================================================================
|
||||
# Checkpoint saving
|
||||
# =====================================================================
|
||||
if step > 0 and step % cfg.train.save_freq == 0:
|
||||
checkpoint_path = checkpoint_dir / f"vla_model_step_{step}.pt"
|
||||
torch.save({
|
||||
'step': step,
|
||||
'model_state_dict': agent.state_dict(),
|
||||
'optimizer_state_dict': optimizer.state_dict(),
|
||||
'loss': loss.item(),
|
||||
}, checkpoint_path)
|
||||
log.info(f"💾 Checkpoint saved: {checkpoint_path}")
|
||||
|
||||
# Save best model
|
||||
if loss.item() < best_loss:
|
||||
best_loss = loss.item()
|
||||
best_model_path = checkpoint_dir / "vla_model_best.pt"
|
||||
torch.save({
|
||||
'step': step,
|
||||
'model_state_dict': agent.state_dict(),
|
||||
'optimizer_state_dict': optimizer.state_dict(),
|
||||
'loss': loss.item(),
|
||||
}, best_model_path)
|
||||
log.info(f"🌟 Best model updated: {best_model_path} (loss: {best_loss:.4f})")
|
||||
|
||||
# =========================================================================
|
||||
# 5. Save Final Model
|
||||
# =========================================================================
|
||||
final_model_path = checkpoint_dir / "vla_model_final.pt"
|
||||
torch.save({
|
||||
'step': cfg.train.max_steps,
|
||||
'model_state_dict': agent.state_dict(),
|
||||
'optimizer_state_dict': optimizer.state_dict(),
|
||||
'loss': loss.item(),
|
||||
}, final_model_path)
|
||||
log.info(f"💾 Final model saved: {final_model_path}")
|
||||
|
||||
log.info("✅ Training completed successfully!")
|
||||
log.info(f"📊 Final Loss: {loss.item():.4f}")
|
||||
log.info(f"📊 Best Loss: {best_loss:.4f}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
main()
|
||||
|
||||
Reference in New Issue
Block a user