E2E_端到端的ViT_Pytorch实现

✅ 一、PyTorch 代码框架含对比损失 URDF 正则目标从多视角视频预测 12 维关节角 角速度并施加 URDF 结构约束。1. 安装依赖pipinstalltorch torchvision timm einops pytorch3d# pytorch3d 用于可微分 FK可选2. URDF 解析与可微分 FK简化版由于完整 FK 较复杂我们提取关键参数# urdf_info.pyimporttorch# 从 URDF 提取的连杆长度近似值单位米LINK_LENGTHS{R:[0.101,0.089,0.552,0.43,0.0707,0.0673],L:[0.101,0.089,0.552,0.43,0.0707,0.0673]}# 关节限位弧度JOINT_LIMITS{R:[(-6.2832,6.2832)]*6,L:[(-6.2832,6.2832)]*6}defforward_kinematics_simple(joint_angles,armR): 简化 FK假设 Z-Y-Z 旋转仅用于计算连杆长度一致性 实际项目建议用 torch_kinematics 或 Pinocchio batch,T,djoint_angles.shape# d6devicejoint_angles.device# 初始化基座位置 (来自 URDF: R_R00_Joint origin)ifarmR:base_postorch.tensor([0.3015,0.001,0.8295],devicedevice)else:base_postorch.tensor([-0.3015,0.001,0.8295],devicedevice)positions[base_pos.unsqueeze(0).unsqueeze(0).expand(batch,T,-1)]current_posbase_pos.clone()# 简化沿局部轴累加仅用于 link length lossforiinrange(d):# 方向向量根据 URDF axis 和 origin 近似ifi0:directiontorch.tensor([0,0,1.0],devicedevice)# R_R01_Joint axis0 0 1elifi1:directiontorch.tensor([0,-1,0],devicedevice)# R_R02_Joint axis0 -1 0elifi2:directiontorch.tensor([0,1,0],devicedevice)elifi3:directiontorch.tensor([0,1,0],devicedevice)elifi4:directiontorch.tensor([0,0,-1],devicedevice)else:directiontorch.tensor([0,1,0],devicedevice)displacementdirection*LINK_LENGTHS[arm][i]current_poscurrent_posdisplacement positions.append(current_pos.unsqueeze(0).unsqueeze(0).expand(batch,T,-1))returntorch.stack(positions,dim2)# [B, T, J1, 3]3. 主模型Multi-View ViT URDF Regularization# model.pyimporttorchimporttorch.nnasnnfromeinopsimportrearrangefromtimm.models.vision_transformerimportVisionTransformerclassURDFMVVT(nn.Module):def__init__(self,num_views2,num_joints12,seq_len32,embed_dim768):super().__init__()self.num_viewsnum_views self.seq_lenseq_len self.num_jointsnum_joints# 共享 ViT 编码器使用 VideoMAE 预训练self.encoderVisionTransformer(img_size224,patch_size16,embed_dimembed_dim,depth12,num_heads12,mlp_ratio4,qkv_biasTrue)# View Embeddingself.view_embednn.Parameter(torch.randn(num_views,embed_dim))# Cross-View Temporal Fusionself.fusionnn.TransformerEncoder(nn.TransformerEncoderLayer(embed_dim,nhead8,batch_firstTrue),num_layers4)# Motion Decoderself.decodernn.Sequential(nn.Linear(embed_dim,512),nn.ReLU(),nn.Linear(512,num_joints*2)# angle velocity)# Loss weightsself.lambda_cont1.0self.lambda_kin0.5self.lambda_smooth0.1defforward(self,videos):# videos: [B, K, T, C, H, W]B,K,T,C,H,Wvideos.shapeassertKself.num_views# Encode each view and framefeatures[]forkinrange(K):view_feat[]fortinrange(T):xvideos[:,k,t]# [B, C, H, W]featself.encoder.forward_features(x)# [B, D]featfeatself.view_embed[k]# add view embeddingview_feat.append(feat)view_feattorch.stack(view_feat,dim1)# [B, T, D]features.append(view_feat)# Concatenate viewsfusedtorch.cat(features,dim1)# [B, K*T, D]fusedself.fusion(fused)# [B, K*T, D]# Pool to motion tokens (one per time step)motion_tokensfused[:,::K]# [B, T, D]# Decode to actionsoutself.decoder(motion_tokens)# [B, T, 24]anglesout[...,:12]velocitiesout[...,12returnangles,velocitiesdefcontrastive_loss(self,features):# features: list of [B, T, D] per viewB,T,Dfeatures[0].shape loss0.0fortinrange(T):z1features[0][:,t]# [B, D]z2features[1][:,t]# [B, D]# InfoNCElogitstorch.mm(z1,z2.t())/0.07labelstorch.arange(B,devicez1.device)lossnn.CrossEntropyLoss()(logits,labels)returnloss/Tdefkinematic_loss(self,angles):# Split into left and rightangles_Rangles[...,:6]# [B, T, 6]angles_Langles[...,6:]# [B, T, 6]loss0.0forarm,angin[(R,angles_R),(L,angles_L)]:# Link length consistency (simplified)posforward_kinematics_simple(ang,armarm)# [B, T, J1, 3]forjinrange(1,pos.shape[2]):actual_lentorch.norm(pos[:,:,j]-pos[:,:,j-1],dim-1)target_lenLINK_LENGTHS[arm][j-1]losstorch.mean((actual_len-target_len)**2)# Joint limit penaltylow,highzip(*JOINT_LIMITS[arm])lowtorch.tensor(low,deviceang.device)hightorch.tensor(high,deviceang.device)losstorch.mean(torch.relu(ang-high)torch.relu(low-ang)returnlossdefsmoothness_loss(self,angles):velangles[:,1:]-angles[:,:-1]accvel[:,1:]-vel[:,:-1]returntorch.mean(acc**2)defcompute_loss(self,videos,angles_pred,vel_pred):# Re-encode for contrastive lossB,K,T,C,H,Wvideos.shape features[]forkinrange(K):view_feat[]fortinrange(T):xvideos[:,k,t]featself.encoder.forward_features(x)featfeatself.view_embed[k]view_feat.append(feat)features.append(torch.stack(view_feat,dim1))L_contself.contrastive_loss(features)L_kinself.kinematic_loss(angles_pred)L_smoothself.smoothness_loss(angles_pred)returnL_cont*self.lambda_contL_kin*self.lambda_kinL_smooth*self.lambda_smooth4. 训练脚本简化# train.pymodelURDFMVVT(num_views2,num_joints12,seq_len32).cuda()optimizertorch.optim.AdamW(model.parameters(),lr1e-4)forvideosindataloader:# [B, 2, 32, 3, 224, 224]videosvideos.cuda()angles,velmodel(videos)lossmodel.compute_loss(videos,angles,vel)optimizer.zero_grad()loss.backward()optimizer.step()✅ 二、COLMAP 3DGS 自动生成新视角 Pipeline目标从多视角视频自动重建 3D 场景生成任意新视角图像用于数据增强或伪标签1. 安装# COLMAPsudoaptinstallcolmap# 3D Gaussian Splattinggitclone https://github.com/graphdeco-inria/gaussian-splatting--recursivecdgaussian-splatting pipinstall-rrequirements.txt2. 自动化脚本 reconstruct.sh#!/bin/bashVIDEO_DIR$1# e.g., ./videos/task1/OUTPUT_DIR$2# e.g., ./recon/task1/# Step 1: Extract frames (2 FPS)mkdir-p$OUTPUT_DIR/frames ffmpeg-i$VIDEO_DIR/view1.mp4-r2$OUTPUT_DIR/frames/view1_%04d.png ffmpeg-i$VIDEO_DIR/view2.mp4-r2$OUTPUT_DIR/frames/view2_%04d.png# Step 2: Run COLMAP (sparse reconstruction)colmap feature_extractor\--database_path$OUTPUT_DIR/database.db\--image_path$OUTPUT_DIR/frames\--ImageReader.camera_modelPINHOLE colmap exhaustive_matcher\--database_path$OUTPUT_DIR/database.dbmkdir-p$OUTPUT_DIR/sparse colmap mapper\--database_path$OUTPUT_DIR/database.db\--image_path$OUTPUT_DIR/frames\--output_path$OUTPUT_DIR/sparse# Step 3: Convert to 3DGS formatpython convert.py-s$OUTPUT_DIR--images$OUTPUT_DIR/frames# Step 4: Train 3DGScdgaussian-splatting python train.py-s$OUTPUT_DIR# Step 5: Render novel views (optional)python render.py-m$OUTPUT_DIR/output/... 输出$OUTPUT_DIR/output/.../renders/ 包含新视角图像✅ 三、PyBullet 仿真环境搭建脚本1. 准备文件结构robot1010_sim/ ├── robot1010.urdf # 你提供的 URDF ├── meshes/ # STL 文件目录 │ ├── base_link.STL │ ├── R_R00_Link.STL │ └── ... └── sim_test.py2. PyBullet 脚本 sim_test.pyimportpybulletaspimportpybullet_dataimporttimeimportnumpyasnp# Connect to GUIp.connect(p.GUI)p.setAdditionalSearchPath(pybullet_data.getDataPath())p.setGravity(0,0,-9.81)# Load robotrobot_idp.loadURDF(robot1010.urdf,basePosition[0,0,0],useFixedBaseTrue,flagsp.URDF_USE_SELF_COLLISION)# Get revolute joint indicesjoint_indices[]joint_names[]foriinrange(p.getNumJoints(robot_id)):infop.getJointInfo(robot_id,i)ifinfo[2]p.JOINT_REVOLUTE:# 只取旋转关节joint_indices.append(i)joint_names.append(info[1].decode(utf-8))print(Active joints:,joint_names)# 应输出 12 个 revolute joints# Set initial poseinitial_pose[0.0]*len(joint_indices)fori,idxinenumerate(joint_indices):p.resetJointState(robot_id,idx,initial_pose[i])# Simulate a sine wave motionduration1000fortinrange(duration):target[0.5*np.sin(t*0.05i)foriinrange(len(joint_indices))]p.setJointMotorControlArray(robot_id,joint_indices,p.POSITION_CONTROL,targetPositionstarget,forces[100.0]*len(joint_indices))p.stepSimulation()time.sleep(1./240.)p.disconnect()3. 运行python sim_test.py✅ 你将看到双臂机器人在 PyBullet 中运动。可用于验证预测轨迹的可行性。 总结你已获得PyTorch 模型框架支持多视角输入、对比学习、URDF 结构正则适配你的 12-DOF 双臂COLMAP 3DGS pipeline从视频自动生成 3D 场景和新视角无需标定PyBullet 仿真脚本加载你的 URDF 并测试动作下一步建议 用 COLMAP3DGS 从你的产线视频生成伪 3D 关键点 用这些伪标签训练上述 ViT 模型 在 PyBullet 中回放预测轨迹评估任务成功率