import os, torch, numpy as np, pickle, argparse from copy import deepcopy from itertools import repeat from tqdm import tqdm from einops import rearrange import matplotlib matplotlib.use('Agg') # Set backend for headless plotting import matplotlib.pyplot as plt import time from torchvision import transforms from ModelTrain.module.utils import load_data from ModelTrain.module.utils import compute_dict_mean, set_seed, detach_dict, calibrate_linear_vel, postprocess_base_action from ModelTrain.module.policy import ACTPolicy, CNNMLPPolicy, DiffusionPolicy from ModelTrain.module.policy_with_hsa import ACTPolicyWithHSA, create_default_hsa_config from ModelTrain.module.policy_jepa_adapter_with_hsa import ACTJEPAHsa, create_default_hsa_config as create_hsa_config_jepa import IPython e = IPython.embed def get_auto_index(dataset_dir): max_idx = 1000 for i in range(max_idx + 1): if not os.path.isfile(os.path.join(dataset_dir, f"qpos_{i}.npy")): return i else: raise Exception(f"Error getting auto index, or more than {max_idx} episodes") def train(args): set_seed(1) ckpt_dir = args["ckpt_dir"] policy_class = args.get("policy_class", "ACT") task_name = args["task_name"] batch_size_train = args["batch_size"] batch_size_val = args["batch_size"] num_steps = args["num_steps"] eval_every = args["eval_every"] validate_every = args["validate_every"] save_every = args["save_every"] resume_ckpt_path = args["resume_ckpt_path"] from ModelTrain.constants import TASK_CONFIGS task_config = TASK_CONFIGS[task_name] dataset_dir = task_config["dataset_dir"] episode_len = task_config["episode_len"] camera_names = task_config["camera_names"] tactile_camera_names = task_config.get("tactile_camera_names", []) stats_dir = task_config.get("stats_dir", None) sample_weights = task_config.get("sample_weights", None) train_ratio = task_config.get("train_ratio", 0.99) name_filter = task_config.get("name_filter", lambda n: True) state_dim = task_config.get('state_dim', 14) action_dim = task_config.get('action_dim', 16) lr_backbone = 1e-05 backbone = "resnet18" if policy_class == "ACT": enc_layers = 4 dec_layers = 7 nheads = 8 policy_config = {'lr':args["lr"], 'num_queries':args["chunk_size"], 'kl_weight':args["kl_weight"], 'hidden_dim':args["hidden_dim"], 'dim_feedforward':args["dim_feedforward"], 'lr_backbone':lr_backbone, 'backbone':backbone, 'enc_layers':enc_layers, 'dec_layers':dec_layers, 'nheads':nheads, 'camera_names':camera_names, 'vq':False, 'vq_class':None, 'vq_dim':None, 'action_dim':action_dim, 'no_encoder':args["no_encoder"]} elif policy_class == "ACTJEPA": enc_layers = 4 dec_layers = 7 nheads = 8 # Handle backward compatibility: vit_ckpt_path or vitg_ckpt_path vit_ckpt = args.get("vit_ckpt_path") or args.get("vitg_ckpt_path") policy_config = {'lr':args["lr"], 'num_queries':args["chunk_size"], 'kl_weight':args["kl_weight"], 'hidden_dim':args["hidden_dim"], 'dim_feedforward':args["dim_feedforward"], 'lr_backbone':lr_backbone, 'backbone':backbone, 'enc_layers':enc_layers, 'dec_layers':dec_layers, 'nheads':nheads, 'camera_names':camera_names, 'tactile_camera_names':tactile_camera_names, 'vq':False, 'vq_class':None, 'vq_dim':None, 'action_dim':action_dim, 'no_encoder':args["no_encoder"], 'use_vitg':True, 'vitg_ckpt_path':vit_ckpt, 'vit_model':args.get("vit_model", "vitg")} elif policy_class == "ACTJEPAAdapter": enc_layers = 4 dec_layers = 7 nheads = 8 # Handle backward compatibility: vit_ckpt_path or vitg_ckpt_path vit_ckpt = args.get("vit_ckpt_path") or args.get("vitg_ckpt_path") policy_config = {'lr':args["lr"], 'num_queries':args["chunk_size"], 'kl_weight':args["kl_weight"], 'hidden_dim':args["hidden_dim"], 'dim_feedforward':args["dim_feedforward"], 'lr_backbone':lr_backbone, 'backbone':backbone, 'enc_layers':enc_layers, 'dec_layers':dec_layers, 'nheads':nheads, 'camera_names':camera_names, 'tactile_camera_names':tactile_camera_names, 'vq':False, 'vq_class':None, 'vq_dim':None, 'action_dim':action_dim, 'no_encoder':args["no_encoder"], 'use_vitg':True, 'vitg_ckpt_path':vit_ckpt, 'vit_model':args.get("vit_model", "vitg"), 'adapter_hidden_dim':args.get("adapter_hidden_dim", 512), 'adapter_depth':args.get("adapter_depth", 3), 'adapter_dropout':args.get("adapter_dropout", 0.1), 'adapter_scale_init':args.get("adapter_scale_init", 0.1), 'adapter_pooling':args.get("adapter_pooling", "attention")} else: if policy_class == "Diffusion": policy_config = {'lr':args["lr"], 'camera_names':camera_names, 'action_dim':action_dim, 'observation_horizon':1, 'action_horizon':8, 'prediction_horizon':args["chunk_size"], 'num_queries':args["chunk_size"], 'num_inference_timesteps':10, 'ema_power':0.75, 'vq':False} else: if policy_class == "CNNMLP": policy_config = {'lr':args["lr"], 'lr_backbone':lr_backbone, 'backbone':backbone, 'num_queries':1, 'camera_names':camera_names} else: raise NotImplementedError config = {'num_steps':num_steps, 'eval_every':eval_every, 'validate_every':validate_every, 'save_every':save_every, 'ckpt_dir':ckpt_dir, 'resume_ckpt_path':resume_ckpt_path, 'episode_len':episode_len, 'state_dim':state_dim, 'lr':args["lr"], 'policy_class':policy_class, 'policy_config':policy_config, 'task_name':task_name, 'seed':args["seed"], 'temporal_agg':args["temporal_agg"], 'camera_names':camera_names, 'load_pretrain':args["load_pretrain"], 'enable_hsa':args.get("enable_hsa", False), 'hsa_weight':args.get("hsa_weight", 1.0), 'hsa_temperature':args.get("hsa_temperature", 0.07), 'hsa_img_size':args.get("hsa_img_size", 224), 'hsa_feature_dim':args.get("hsa_feature_dim", 768), 'hsa_num_heads':args.get("hsa_num_heads", 12), 'robot_type':args.get("robot_type", "Nova 2"), 'wrist_camera':args.get("wrist_camera", "left_wrist"), 'camera_params':args.get("camera_params", None)} if not os.path.isdir(ckpt_dir): os.makedirs(ckpt_dir) config_path = os.path.join(ckpt_dir, "config.pkl") expr_name = ckpt_dir.split("/")[-1] with open(config_path, "wb") as f: pickle.dump(config, f) print(f"Loading data from: {dataset_dir}") use_vitg = args.get("use_vitg", False) # Use all cameras (RGB + tactile) for data loading all_camera_names = camera_names + tactile_camera_names train_dataloader, val_dataloader, stats, _ = load_data(dataset_dir, name_filter, all_camera_names, batch_size_train, batch_size_val, (args["chunk_size"]), (args["skip_mirrored_data"]), (config["load_pretrain"]), policy_class, stats_dir_l=stats_dir, sample_weights=sample_weights, train_ratio=train_ratio, use_vitg=use_vitg, tactile_camera_names=tactile_camera_names) stats_path = os.path.join(ckpt_dir, "dataset_stats.pkl") with open(stats_path, "wb") as f: pickle.dump(stats, f) best_ckpt_info = train_bc(train_dataloader, val_dataloader, config) best_step, min_val_loss, best_state_dict = best_ckpt_info ckpt_path = os.path.join(ckpt_dir, "policy_best.ckpt") torch.save(best_state_dict, ckpt_path) print(f"Best ckpt, val loss {min_val_loss:.6f} @ step{best_step}") def make_policy(policy_class, policy_config, hsa_config=None): """Create policy with optional HSA loss support.""" if policy_class == "ACT": if hsa_config is not None and hsa_config.get('enable_hsa', False): policy = ACTPolicyWithHSA(policy_config, hsa_config) else: policy = ACTPolicy(policy_config) elif policy_class == "ACTJEPA": from ModelTrain.module.policy_jepa import ACTJEPAPolicy policy = ACTJEPAPolicy(policy_config) elif policy_class == "ACTJEPAAdapter": from ModelTrain.module.policy_jepa_adapter import ACTJEPAAdapterPolicy if hsa_config is not None and hsa_config.get('enable_hsa', False): policy = ACTJEPAHsa(policy_config, hsa_config) else: policy = ACTJEPAAdapterPolicy(policy_config) elif policy_class == "CNNMLP": policy = CNNMLPPolicy(policy_config) elif policy_class == "Diffusion": policy = DiffusionPolicy(policy_config) else: raise NotImplementedError return policy def make_optimizer(policy_class, policy): if policy_class == "ACT": optimizer = policy.configure_optimizers() elif policy_class == "ACTJEPA": optimizer = policy.configure_optimizers() elif policy_class == "ACTJEPAAdapter": optimizer = policy.configure_optimizers() elif policy_class == "CNNMLP": optimizer = policy.configure_optimizers() elif policy_class == "Diffusion": optimizer = policy.configure_optimizers() else: raise NotImplementedError return optimizer def get_image(ts, camera_names, rand_crop_resize=False): print("get_image") curr_images = [] for cam_name in camera_names: curr_image = rearrange(ts.observation["images"][cam_name], "h w c -> c h w") curr_images.append(curr_image) else: curr_image = np.stack(curr_images, axis=0) curr_image = torch.from_numpy(curr_image / 255.0).float().cuda().unsqueeze(0) if rand_crop_resize: print("rand crop resize is used!") original_size = curr_image.shape[-2:] ratio = 0.95 curr_image = curr_image[..., int(original_size[0] * (1 - ratio) / 2):int(original_size[0] * (1 + ratio) / 2), int(original_size[1] * (1 - ratio) / 2):int(original_size[1] * (1 + ratio) / 2)] curr_image = curr_image.squeeze(0) resize_transform = transforms.Resize(original_size, antialias=True) curr_image = resize_transform(curr_image) curr_image = curr_image.unsqueeze(0) return curr_image def forward_pass(data, policy, enable_hsa=False): """ Forward pass through policy, handling both standard and HSA modes. Args: data: Batch data (5 items with tactile) or (4 items without tactile) policy: Policy model enable_hsa: Whether to compute HSA loss (requires tactile data) Returns: Forward dictionary with losses """ # Handle both old format (4 items) and new format (5 items with tactile) if len(data) == 5: # New format: RGB images, tactile images, qpos, action, is_pad rgb_data, tactile_data, qpos_data, action_data, is_pad = data # Debug: Print shapes to understand data structure # print(f"DEBUG: rgb_data type={type(rgb_data)}, shape={rgb_data.shape if isinstance(rgb_data, torch.Tensor) else 'N/A'}") # print(f"DEBUG: tactile_data type={type(tactile_data)}, len={len(tactile_data) if isinstance(tactile_data, list) else 'N/A'}") # if isinstance(tactile_data, list) and len(tactile_data) > 0: # print(f"DEBUG: tactile_data[0] type={type(tactile_data[0])}") # if isinstance(tactile_data[0], torch.Tensor): # print(f"DEBUG: tactile_data[0] shape={tactile_data[0].shape}") # elif isinstance(tactile_data[0], list): # print(f"DEBUG: tactile_data[0] is list, len={len(tactile_data[0])}") # if len(tactile_data[0]) > 0: # print(f"DEBUG: tactile_data[0][0] type={type(tactile_data[0][0])}, shape={tactile_data[0][0].shape if isinstance(tactile_data[0][0], torch.Tensor) else 'N/A'}") # Handle tactile_data: could be tensor or list depending on batching if isinstance(tactile_data, torch.Tensor): # Already a tensor (batch, num_tactile, C, H, W) or (batch, C, H, W) if tactile_data.dim() == 4: # Single tactile sensor: (batch, C, H, W) -> add camera dim tactile_data = tactile_data.unsqueeze(1) # (batch, 1, C, H, W) # Concatenate RGB and tactile along camera dimension image_data = torch.cat([rgb_data, tactile_data], dim=1) elif tactile_data and len(tactile_data) > 0: # It's a list - but DataLoader returns list of already-batched tensors # tactile_data is list of (batch, C, H, W) tensors # Check if first element is already batched (has same batch size as rgb_data) if isinstance(tactile_data[0], torch.Tensor) and tactile_data[0].dim() == 4: # Each element is (batch, C, H, W), need to add camera dimension # Stack along camera dimension: list of (B,C,H,W) -> (B, num_tactile, C, H, W) tactile_stacked = torch.stack(tactile_data, dim=1) # Stack along dim=1 (camera dim) # print(f"DEBUG: tactile_stacked shape after stack(dim=1)={tactile_stacked.shape}") elif isinstance(tactile_data[0], list): # List of lists: (batch) of (num_tactile) of (C, H, W) tactile_stacked = torch.stack([torch.stack(batch_tactile) for batch_tactile in tactile_data]) # print(f"DEBUG: tactile_stacked shape after list-of-lists={tactile_stacked.shape}") else: # List of per-sample tensors: (batch) of (C, H, W) tactile_stacked = torch.stack(tactile_data) # (batch, C, H, W) # print(f"DEBUG: tactile_stacked shape after stack={tactile_stacked.shape}") tactile_stacked = tactile_stacked.unsqueeze(1) # (batch, 1, C, H, W) # print(f"DEBUG: tactile_stacked shape after unsqueeze={tactile_stacked.shape}") # Can't concatenate due to different spatial sizes (480x640 vs 224x224) # Pass as list instead - model will handle separately image_data = [rgb_data, tactile_stacked] # print(f"DEBUG: Passing image_data as list: RGB={rgb_data.shape}, Tactile={tactile_stacked.shape}") else: # No tactile data image_data = rgb_data else: # Old format: image_data, qpos, action, is_pad image_data, qpos_data, action_data, is_pad = data # Move to CUDA if isinstance(image_data, list): # List of tensors (hybrid mode with different resolutions) image_data = [img.cuda() for img in image_data] else: image_data = image_data.cuda() qpos_data, action_data, is_pad = (qpos_data.cuda(), action_data.cuda(), is_pad.cuda()) return policy(qpos_data, image_data, action_data, is_pad) def train_bc(train_dataloader, val_dataloader, config): num_steps = config["num_steps"] ckpt_dir = config["ckpt_dir"] seed = config["seed"] policy_class = config["policy_class"] policy_config = config["policy_config"] eval_every = config["eval_every"] validate_every = config["validate_every"] save_every = config["save_every"] # Setup HSA configuration if enabled enable_hsa = config.get("enable_hsa", False) hsa_config = None if enable_hsa: # Compute wrist camera index from camera names wrist_camera_name = config.get("wrist_camera", "left_wrist") camera_names = config['camera_names'] try: wrist_camera_idx = camera_names.index(wrist_camera_name) except ValueError: print(f"Warning: Wrist camera '{wrist_camera_name}' not found in camera_names {camera_names}, using index 1 (left_wrist)") wrist_camera_idx = 1 hsa_config = { 'enable_hsa': True, 'hsa_weight': config.get("hsa_weight", 1.0), 'temperature': config.get("hsa_temperature", 0.07), 'img_size': config.get("hsa_img_size", 224), 'feature_dim': config.get("hsa_feature_dim", 768), 'num_heads': config.get("hsa_num_heads", 12), 'robot_type': config.get("robot_type", "Nova 2"), 'wrist_camera': wrist_camera_name, 'wrist_camera_idx': wrist_camera_idx, 'tactile_camera_idx': 0, # Default to first tactile sensor 'camera_params': config.get("camera_params", None) # Camera calibration for gripper-aware offset } set_seed(seed) policy = make_policy(policy_class, policy_config, hsa_config) # Print HSA configuration if enabled if enable_hsa: print("\n" + "="*60) print("HSA Loss ENABLED") print(f" Policy Class: {policy_class}") print(f" HSA Weight: {hsa_config['hsa_weight']}") print(f" Temperature: {hsa_config['temperature']}") print(f" Robot Type: {hsa_config['robot_type']}") print(f" Wrist Camera: {hsa_config['wrist_camera']}") print("="*60 + "\n") if config["load_pretrain"]: loading_status = policy.deserialize(torch.load(os.path.join("./ckpt/pretrain_all", "policy_step_50000_seed_0.ckpt"))) print(f"loaded! {loading_status}") if config["resume_ckpt_path"] is not None: loading_status = policy.deserialize(torch.load(config["resume_ckpt_path"])) print(f'Resume policy from: {config["resume_ckpt_path"]}, Status: {loading_status}') optimizer = make_optimizer(policy_class, policy) policy.cuda() min_val_loss = np.inf best_ckpt_info = None train_dataloader = repeater(train_dataloader) train_loss = [] val_loss = [] train_hsa = [] # Track HSA loss val_hsa = [] # Track validation HSA loss last_time = time.time() start_time = last_time for step in tqdm(range(num_steps + 1)): if step % validate_every == 0: print("validating") with torch.inference_mode(): policy.eval() validation_dicts = [] for batch_idx, data in enumerate(val_dataloader): forward_dict = forward_pass(data, policy, enable_hsa=enable_hsa) validation_dicts.append(forward_dict) if batch_idx > 50: break validation_summary = compute_dict_mean(validation_dicts) epoch_val_loss = validation_summary["loss"].mean() if epoch_val_loss < min_val_loss: min_val_loss = epoch_val_loss best_ckpt_info = (step, min_val_loss, deepcopy(policy.serialize())) for k in list(validation_summary.keys()): validation_summary[f"val_{k}"] = validation_summary.pop(k) else: print(f"Val loss: {epoch_val_loss:.5f}") val_loss.append(float(epoch_val_loss.item())) # Track HSA validation loss if enabled if enable_hsa and 'val_hsa_total' in validation_summary: val_hsa.append(float(validation_summary['val_hsa_total'].mean().item())) # Print HSA losses prominently if enabled if enable_hsa: hsa_keys = [k for k in validation_summary.keys() if 'hsa' in k.lower()] if hsa_keys: print(" HSA Losses:", end=" ") for k in hsa_keys: print(f"{k}: {validation_summary[k].mean().item():.3f}", end=" ") print() # Print all validation metrics summary_string = "" for k, v in validation_summary.items(): summary_string += f"{k}: {v.mean().item():.3f} " else: print(summary_string) if step > 0: if step % eval_every == 0: ckpt_name = f"policy_step_{step}_seed_{seed}.ckpt" ckpt_path = os.path.join(ckpt_dir, ckpt_name) policy.train() optimizer.zero_grad() data = next(train_dataloader) forward_dict = forward_pass(data, policy, enable_hsa=enable_hsa) loss = forward_dict["loss"] loss.mean().backward() optimizer.step() train_loss.append(float(loss.mean().item())) # Track HSA training loss if enabled if enable_hsa and 'hsa_total' in forward_dict: train_hsa.append(float(forward_dict['hsa_total'].mean().item())) # Print training loss periodically (every 100 steps) if step % 100 == 0 and step > 0: train_summary = f"Step {step} - Train loss: {loss.mean().item():.5f}" if enable_hsa and 'hsa_wrist' in forward_dict: train_summary += f" | L1: {forward_dict['l1'].mean().item():.3f}" train_summary += f" | KL: {forward_dict['kl'].mean().item():.3f}" train_summary += f" | HSA_wrist: {forward_dict['hsa_wrist'].mean().item():.3f}" if 'hsa_total' in forward_dict: train_summary += f" | HSA_total: {forward_dict['hsa_total'].mean().item():.3f}" print(train_summary) if step % save_every == 0: ckpt_path = os.path.join(ckpt_dir, f"policy_step_{step}_seed_{seed}.ckpt") torch.save(policy.serialize(), ckpt_path) cur_time = time.time() last_time = cur_time else: print("train all time:", cur_time - start_time) ckpt_path = os.path.join(ckpt_dir, "policy_last.ckpt") torch.save(policy.serialize(), ckpt_path) best_step, min_val_loss, best_state_dict = best_ckpt_info ckpt_path = os.path.join(ckpt_dir, f"policy_step_{best_step}_seed_{seed}.ckpt") torch.save(best_state_dict, ckpt_path) print(f"Training finished:\nSeed {seed}, val loss {min_val_loss:.6f} at step {best_step}") # Plot training loss with smoothing plt.figure(figsize=(12, 6)) plt.plot(train_loss, label="Training Loss (raw)", color='blue', alpha=0.2, linewidth=0.5) # Add smoothed line if len(train_loss) > 100: window = min(100, len(train_loss) // 10) smoothed = np.convolve(train_loss, np.ones(window)/window, mode='valid') smooth_steps = list(range(window//2, len(train_loss) - window//2)) plt.plot(smooth_steps, smoothed, label=f"Training Loss (smoothed)", color='darkblue', linewidth=2) plt.title("Training Loss Over Steps") plt.xlabel("Step") plt.ylabel("Loss") plt.legend() plt.grid(True, alpha=0.3) plt.savefig(ckpt_dir + "/train_loss.png", dpi=150) plt.close() # Plot validation loss plt.figure(figsize=(10, 6)) plt.plot(val_loss, label="Validation Loss", color='green', linewidth=2, marker='o', markersize=4) plt.title("Validation Loss Over Steps") plt.xlabel("Validation Step") plt.ylabel("Loss") plt.legend() plt.grid(True, alpha=0.3) plt.savefig(ckpt_dir + "/val_loss.png", dpi=150) plt.close() # Plot HSA losses if available if enable_hsa and len(train_hsa) > 0: # Training HSA loss with smoothing plt.figure(figsize=(12, 6)) # Raw data with transparency plt.plot(train_hsa, label="HSA Loss (raw)", color='orange', alpha=0.2, linewidth=0.5) # Smoothed data using moving average window_size = min(100, len(train_hsa) // 10) # Adaptive window if len(train_hsa) > window_size: smoothed = np.convolve(train_hsa, np.ones(window_size)/window_size, mode='valid') smoothed_steps = list(range(window_size//2, len(train_hsa) - window_size//2)) plt.plot(smoothed_steps, smoothed, label=f"HSA Loss (smoothed, window={window_size})", color='darkorange', linewidth=2) plt.title("HSA Training Loss Over Steps") plt.xlabel("Step") plt.ylabel("HSA Loss") plt.legend() plt.grid(True, alpha=0.3) plt.savefig(ckpt_dir + "/train_hsa_loss.png", dpi=150) plt.close() # Validation HSA loss if len(val_hsa) > 0: plt.figure(figsize=(10, 6)) plt.plot(val_hsa, label="Validation HSA Loss", color='red') plt.title("HSA Validation Loss Over Steps") plt.xlabel("Validation Step") plt.ylabel("HSA Loss") plt.legend() plt.grid(True) plt.savefig(ckpt_dir + "/val_hsa_loss.png") plt.close() # Combined plot: Training and Validation HSA if len(val_hsa) > 0: plt.figure(figsize=(12, 6)) # Raw training data (very transparent) plt.plot(train_hsa, color='orange', alpha=0.15, linewidth=0.5, label='Training (raw)') # Smoothed training data window_size = min(100, len(train_hsa) // 10) if len(train_hsa) > window_size: smoothed = np.convolve(train_hsa, np.ones(window_size)/window_size, mode='valid') smoothed_steps = list(range(window_size//2, len(train_hsa) - window_size//2)) plt.plot(smoothed_steps, smoothed, color='darkorange', linewidth=2, label=f'Training (smoothed)') # Validation data val_steps = [i * validate_every for i in range(len(val_hsa))] plt.plot(val_steps, val_hsa, color='red', linewidth=2.5, marker='o', markersize=4, label='Validation') plt.title("HSA Loss: Training vs Validation") plt.xlabel("Step") plt.ylabel("HSA Loss") plt.legend() plt.grid(True, alpha=0.3) plt.savefig(ckpt_dir + "/hsa_loss_combined.png", dpi=150) plt.close() print(f"HSA loss plots saved to {ckpt_dir}/") if len(train_hsa) > 100: print(f" Initial HSA: {train_hsa[0]:.3f} → Final HSA: {train_hsa[-1]:.3f}") return best_ckpt_info def repeater(data_loader): epoch = 0 for loader in repeat(data_loader): for data in loader: yield data else: print(f"Epoch {epoch} done") epoch += 1 # okay decompiling train_module.pyc