failo0711/ATG-MoE_TrainingSet

ATG-MoE Pressure-Reducing Valve Assembly Training Set

   

This repository provides the training set used in our paper: "ATG-MoE: Autoregressive trajectory generation with mixture-of-experts for assembly skill learning".

The dataset is specifically designed for the pressure-reducing valve assembly task, featuring multi-skill robotic learning capabilities.

[!IMPORTANT] This release only contains the training set. Evaluation requires a Unity-based simulation environment. The evaluation scripts and instructions will be released in our GitHub repository.

📌 Dataset Overview

The dataset supports multi-skill imitation learning in industrial assembly. It provides high-quality demonstrations mapping multi-view observations to precise skill trajectories.

🛠️ Included Skills

The dataset covers 8 key skills required for the assembly process:

• —Sleeve Placement
• —Large Spring Insertion
• —Rod Placement
• —Rod Seating
• —Nut Seating
• —Spring Insertion
• —Plug Seating
• —Body Seating

📊 Dataset Scale

• —Total Training Episodes: 768
• —Skills: 8 distinct assembly tasks
• —Data per Skill: 96 training episodes

📂 Data Format

The dataset follows the RLBench organization style. Each episode folder contains a complete demonstration trajectory.

Directory Structure

episode0/
├── front_depth/                # Depth maps from front camera
├── front_rgb/                  # RGB images from front camera
├── left_shoulder_depth/        # Depth maps from left shoulder camera
├── left_shoulder_rgb/          # RGB images from left shoulder camera
├── overhead_depth/             # Depth maps from overhead camera
├── overhead_rgb/               # RGB images from overhead camera
├── right_shoulder_depth/       # Depth maps from right shoulder camera
├── right_shoulder_rgb/         # RGB images from right shoulder camera
├── proprioception/             # Robot joint states and end-effector pose
├── camera_matrix.json          # Intrinsic and extrinsic parameters
├── kfs.json                    # Keyframe indices for trajectory
├── lang_emb.pkl                # Language instruction embeddings
├── variation_descriptions.json # Natural language descriptions
└── variation_number.pkl        # Variation ID for the task

📝 Citation

If you find this dataset or our work helpful, please cite:

@misc{huang2026atgmoeautoregressivetrajectorygeneration,
  title={ATG-MoE: Autoregressive trajectory generation with mixture-of-experts for assembly skill learning},
  author={Weihang Huang and Chaoran Zhang and Xiaoxin Deng and Hao Zhou and Zhaobo Xu and Shubo Cui and Long Zeng},
  year={2026},
  eprint={2603.19029},
  archivePrefix={arXiv},
  primaryClass={cs.RO},
  url={https://arxiv.org/abs/2603.19029},
}