import argparse import sys import os BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) print(BASE_DIR) sys.path.append(BASE_DIR) sys.path.append(BASE_DIR+'/ModelTrain') sys.path.append(BASE_DIR+'/ModelTrain/detr') sys.path.append(BASE_DIR+'/robomimic-r2d2') from module.train_module import train def arg_config(): parser = argparse.ArgumentParser() parser.add_argument('--ckpt_dir', action='store', type=str, help='ckpt_dir', default='./ckpt/dobot_pick_random_1013',required=False) parser.add_argument('--task_name', action='store', type=str, default='dobot_pick_random_1013',help='task_name', required=False) parser.add_argument('--batch_size', action='store', type=int, help='batch_size', default=16, required=False) parser.add_argument('--seed', action='store', type=int, help='seed', default=0,required=False) parser.add_argument('--num_steps', action='store', type=int, help='num_steps', default=30000, required=False) parser.add_argument('--lr', action='store', type=float, help='lr', default=2e-5,required=False) parser.add_argument('--load_pretrain', action='store_true', default=False) # Ignore this parameter and leave the default setting parser.add_argument('--eval_every', action='store', type=int, default=100, help='eval_every', required=False) # Ignore this parameter and leave the default setting parser.add_argument('--validate_every', action='store', type=int, default=100, help='validate_every',required=False) parser.add_argument('--save_every', action='store', type=int, default=10000, help='save_every', required=False) parser.add_argument('--resume_ckpt_path', action='store', type=str, help='resume_ckpt_path', default=None, required=False) parser.add_argument('--skip_mirrored_data', action='store_true') parser.add_argument('--kl_weight', action='store', type=int, help='KL divergence weight,recommended set 10 or 100', default=10,required=False) parser.add_argument('--chunk_size', action='store', type=int, help='The model predicts the length of the output action sequence at a time', default=45,required=False) parser.add_argument('--hidden_dim', action='store', type=int, help='hidden_dim', default=512, required=False) parser.add_argument('--dim_feedforward', action='store', type=int, help='dim_feedforward', default=3200,required=False) parser.add_argument('--temporal_agg', action='store_true', default=True) parser.add_argument('--no_encoder', action='store_true', default=False) # Policy class selection parser.add_argument('--policy_class', action='store', type=str, default='ACT', choices=['ACT', 'ACTJEPA', 'ACTJEPAAdapter', 'CNNMLP', 'Diffusion'], help='Policy class to use (ACT for standard, ACTJEPA for hybrid RGB+tactile, ACTJEPAAdapter for learnable adapters)') # ViT-related arguments (used by ACTJEPA and ACTJEPAAdapter) parser.add_argument('--use_vitg', action='store_true', default=False, help='Use ViTG encoder for tactile images (deprecated, use --policy_class ACTJEPA instead)') parser.add_argument('--vit_model', action='store', type=str, default='vitg', choices=['vitg', 'vitl'], help='ViT model type for tactile processing: vitg (1408-dim) or vitl (1024-dim)') parser.add_argument('--vit_ckpt_path', action='store', type=str, help='Path to ViT checkpoint file (.pt)', required=False) parser.add_argument('--vitg_ckpt_path', action='store', type=str, help='Path to ViTG checkpoint file (deprecated, use --vit_ckpt_path)', required=False) # Adapter-related arguments (used by ACTJEPAAdapter) parser.add_argument('--adapter_hidden_dim', action='store', type=int, default=512, help='Hidden dimension for residual adapter MLP') parser.add_argument('--adapter_depth', action='store', type=int, default=3, help='Number of layers in adapter MLP') parser.add_argument('--adapter_dropout', action='store', type=float, default=0.1, help='Dropout rate for adapter') parser.add_argument('--adapter_scale_init', action='store', type=float, default=0.1, help='Initial value for residual scaling factor') parser.add_argument('--adapter_pooling', action='store', type=str, default='attention', choices=['attention', 'mean'], help='Pooling type for aggregating patch tokens') # CLIP Encoder arguments parser.add_argument('--clip_model', action='store', type=str, default='ViT-B-16', choices=['ViT-B-32', 'ViT-B-16', 'ViT-L-14', 'ViT-H-14'], help='CLIP model variant for RGB cameras (default: ViT-B-16)') parser.add_argument('--clip_pretrained', action='store', type=str, default='openai', help='CLIP pretrained weights (default: openai, can also use laion2b_s34b_b88k, etc.)') parser.add_argument('--freeze_clip', action='store_true', default=False, help='Freeze CLIP encoder weights (default: False, i.e., trainable)') # Text conditioning arguments parser.add_argument('--enable_text', action='store_true', default=True, help='Enable text conditioning using CLIP text encoder') parser.add_argument('--text_prompt', action='store', type=str, default="Insert the peg into the hole", help='Text prompt for task conditioning (e.g., "insert USB cable", "pick red block")') # HSA Loss arguments parser.add_argument('--enable_hsa', action='store_true', default=False, help='Enable HSA (Hard Sample Aware) loss for tactile-visual alignment') parser.add_argument('--hsa_weight', action='store', type=float, default=1.0, help='Weight for HSA loss term (default: 1.0)') parser.add_argument('--hsa_temperature', action='store', type=float, default=0.07, help='Temperature parameter for HSA contrastive loss (default: 0.07)') parser.add_argument('--hsa_img_size', action='store', type=int, default=224, help='Image size for HSA feature extraction (default: 224)') parser.add_argument('--hsa_feature_dim', action='store', type=int, default=768, help='Feature dimension for HSA backbone (default: 768)') parser.add_argument('--hsa_num_heads', action='store', type=int, default=12, help='Number of attention heads for HSA backbone (default: 12 for ViT-L, use 16 for ViT-G)') parser.add_argument('--robot_type', action='store', type=str, default='Nova 2', choices=['Nova 2', 'Nova 5'], help='Robot type for forward kinematics (default: Nova 2)') parser.add_argument('--wrist_camera', action='store', type=str, default='left_wrist', help='Name of the wrist camera for HSA (default: left_wrist)') # Use parse_known_args to ignore unknown arguments (e.g., from inference scripts) args, unknown = parser.parse_known_args() return vars(args) if __name__ == '__main__': args = arg_config() train(args)