Kebijakan DiT Multitask

Kebijakan Multitask Diffusion Transformer (DiT) adalah evolusi dari arsitektur Diffusion Policy asli, yang memanfaatkan DiT besar dengan kondisioning teks dan visi untuk pembelajaran robot multitask. Implementasi ini mendukung tujuan difusi dan pencocokan aliran untuk generasi aksi, memungkinkan robot untuk melakukan berbagai tugas manipulasi yang dikondisikan oleh instruksi bahasa.

Ringkasan Model

Model menggunakan:

• CLIP Vision Encoder: Memproses gambar RGB dari beberapa sudut pandang kamera
• CLIP Text Encoder: Mengkodekan instruksi tugas bahasa (bobot beku dengan proyeksi yang dapat dipelajari)
• Diffusion Transformer: Memprediksi urutan aksi yang dikondisikan oleh observasi dan bahasa
• Dua Tujuan: Mendukung baik difusi (DDPM/DDIM) maupun flow matching untuk generasi aksi

Model ini menarik karena Anda dapat mencapai kehalusan gerak yang sangat tinggi, yang mampu bersaing dengan VLA berparameter miliaran, hanya dengan ~450M parameter dan pelatihan yang jauh lebih sedikit.

Persyaratan Instalasi

Multitask DiT Policy memiliki dependensi tambahan. Instal dengan:

pip install 'lerobot[multi_task_dit]'

Ini akan menginstal semua dependensi yang diperlukan, termasuk pustaka HuggingFace Transformers untuk model CLIP.

Penggunaan

Untuk menggunakan Multitask DiT dalam konfigurasi LeRobot Anda, tentukan tipe kebijakan sebagai:

policy.type=multi_task_dit

Pelatihan

Perintah Pelatihan Dasar

Berikut adalah perintah pelatihan lengkap untuk melatih Multitask DiT pada dataset Anda:

lerobot-train \
  --dataset.repo_id=YOUR_DATASET \
  --output_dir=./outputs/multitask_dit_training \
  --batch_size=32 \
  --steps=5000 \
  --save_freq=500 \
  --log_freq=100 \
  --policy.type=multi_task_dit \
  --policy.device=cuda \
  --policy.repo_id="HF_USER/multitask-dit-your-robot" \
  --trackio.enable=true

Hyperparameter yang Direkomendasikan dan Detail Dataset (Frekuensi Kontrol 30Hz)

Untuk kinerja yang andal, mulailah dengan hiperparameter default yang disarankan berikut:

lerobot-train \
  --dataset.repo_id=YOUR_DATASET \
  --output_dir=./outputs/mutitask_dit_training \
  --batch_size=320 \
  --steps=30000 \
  --policy.type=multi_task_dit \
  --policy.device=cuda \
  --policy.horizon=32 \
  --policy.n_action_steps=24 \
  --policy.objective=diffusion \
  --policy.noise_scheduler_type=DDPM \
  --policy.num_train_timesteps=100 \
  --policy.repo_id="HF_USER/multitask-dit-your-robot" \
  --trackio.enable=true

Parameter Utama:

• Ukuran Batch: 192-320 - Jika Anda memiliki akses ke GPU yang dapat mendukung ini, Anda akan mendapatkan dinamika pelatihan terbaik
• Horizon: 32 - jumlah langkah aksi yang diprediksi, ~1,0 detik pada 30Hz
• n_action_steps: 24 - ~0,8 detik pada 30Hz
• Tujuan: diffusion - mulai dengan diffusion dan bereksperimen dengan flow matching jika kualitas generasi buruk
• Langkah Pelatihan: >30k langkah disarankan untuk satu tugas

Parameter Konfigurasi Pelatihan

Pemilihan Tujuan

Pilih antara difusi dan pencocokan aliran:

# Diffusion objective (default)
--policy.objective=diffusion \
--policy.noise_scheduler_type=DDPM \
--policy.num_train_timesteps=100 \
--policy.num_inference_steps=10 \
--policy.beta_schedule=squaredcos_cap_v2 \
--policy.prediction_type=epsilon \
--policy.clip_sample=true \
--policy.clip_sample_range=1.0

# Flow matching objective
--policy.objective=flow_matching \
--policy.timestep_sampling_strategy=beta \
--policy.num_integration_steps=100 \
--policy.integration_method=euler \
--policy.sigma_min=0.0

Arsitektur Transformer

Sesuaikan kapasitas model berdasarkan ukuran dataset:

# Small datasets (< 100 examples)
--policy.num_layers=4 \
--policy.hidden_dim=512 \
--policy.num_heads=8  # should ideally be hidden_dim // 64

# Medium datasets (100-5k examples) - default
--policy.num_layers=6 \
--policy.hidden_dim=512 \
--policy.num_heads=8  # should ideally be hidden_dim // 64

# Large datasets (> 5k examples)
--policy.num_layers=8 \
--policy.hidden_dim=512 \
--policy.num_heads=8   # should ideally be hidden_dim // 64

Opsi Encoding Posisional:

Model mendukung dua metode pengkodean posisi untuk urutan aksi:

# Rotary Position Embedding (RoPE) - default, recommended
--policy.use_rope=true \
--policy.rope_base=10000.0  # Base frequency for RoPE

# Absolute positional encoding
--policy.use_positional_encoding=true  # Disables RoPE when true

Parameter Transformer Lainnya:

--policy.dropout=0.1  # Dropout rate for DiT blocks (0.0-1.0)
--policy.timestep_embed_dim=256  # Timestep embedding dimension

Konfigurasi Vision Encoder

# Use different CLIP model for more expressivity at the cost of inference time
# experiment with larger or smaller models depending on the complexity of your tasks and size of dataset
--policy.vision_encoder_name=openai/clip-vit-large-patch14

# Use separate vision encoder per camera
# This may be useful when cameras have significantly different characteristics, but
# be wary of increased VRAM footprint.
--policy.use_separate_rgb_encoder_per_camera=true

# Image preprocessing (you may need to resize your images for inference speed ups)
--policy.image_resize_shape=[XXX,YYY] \
--policy.image_crop_shape=[224,224] \
--policy.image_crop_is_random=true  # Random during training, center at inference

Konfigurasi Text Encoder

# Use different CLIP text encoder model
# same as vision: experiment with larger or smaller models depending on the
# complexity of your tasks and size of dataset
--policy.text_encoder_name=openai/clip-vit-large-patch14

Konfigurasi Laju Pembelajaran

Vision encoder menggunakan pengali laju belajar yang terpisah, di mana 1/10 disarankan sebagai titik awal yang ideal:

--policy.optimizer_lr=2e-5 \
--policy.vision_encoder_lr_multiplier=0.1  # Vision encoder LR = 0.1 * optimizer_lr

Panduan Penyesuaian Pelatihan

1. Flow Matching dengan Beta Sampling

Implementasi difusi asli di sini didasarkan pada pekerjaan yang dijelaskan dalam makalah LBM TRI

Selain itu, kami telah mengimplementasikan tujuan flow-matching, yang dijelaskan secara tingkat tinggi dalam artikel blog Boston Dynamics.

Pertimbangkan untuk menguji tujuan flow-matching dan mengevaluasi perbedaan kinerja untuk tugas Anda:

--policy.objective=flow_matching \
--policy.timestep_sampling_strategy=beta \
--policy.timestep_sampling_alpha=1.5 \
--policy.timestep_sampling_beta=1.0 \
--policy.timestep_sampling_s=0.999

Hal ini belum terbukti menjadi solusi ajaib untuk setiap kasus pengguna, tetapi sesekali menghasilkan aksi yang lebih halus dan konsisten.

2. Jumlah Lapisan Transformer

Sesuaikan kapasitas model dengan ukuran dataset Anda:

• Dataset kecil (< 100 contoh): Kurangi menjadi 4 lapisan
• Dataset besar (> 5k contoh): Tingkatkan menjadi 8 layer

3. Penyetelan horizon

Model dapat sensitif terhadap horizon yang Anda pilih. Mulailah dengan horizon sekitar 1 detik berdasarkan frekuensi kontrol Anda:

• frekuensi 30 Hz: horizon=30
• frekuensi 10 Hz: horizon=10

Lalu, lakukan eksperimen dengan meningkatkannya dari sana. Horizon menentukan seberapa jauh ke masa depan model memprediksi aksi.

4. Sensitivitas n_action_steps

Model juga dapat sangat sensitif terhadap n_action_steps. Mulailah dengan nilai sekitar 0,8 detik berdasarkan frekuensi kontrol Anda, lalu sesuaikan dari situ:

• Nilai yang lebih rendah: Lebih reaktif tetapi berpotensi kurang stabil untuk tugas jangka panjang
• Nilai yang lebih tinggi: Lebih baik untuk eksekusi jangka panjang, tetapi kegagalan open-loop memiliki keterbatasan dalam pemulihan

Penyetelan Inferensi

Untuk inferensi yang lebih cepat, gunakan DDIM dengan langkah sampling yang lebih sedikit:

--policy.noise_scheduler_type=DDIM \
--policy.num_inference_steps=10

Melanjutkan Pelatihan

Untuk melanjutkan pelatihan dari checkpoint:

lerobot-train \
  --config_path=./outputs/mutitask_dit_training/checkpoints/last/pretrained_model/train_config.json \
  --resume=true

Direktori checkpoint harus berisi file model.safetensors dan config.json (disimpan secara otomatis selama pelatihan). Saat melanjutkan, konfigurasi dimuat dari checkpoint, sehingga Anda tidak perlu menentukan parameter lain.

Mode Kegagalan Umum dan Pemecahan Masalah

Pelatihan model-model ini bisa menjadi rumit. Berikut adalah mode kegagalan umum dan pendekatan debugging:

Menganggur / Tidak Bergerak

Model dapat “runtuh” selama inferensi, yang menghasilkan gerakan statis atau tidak ada gerakan sama sekali. Hal ini dapat terjadi ketika:

1. Data pelatihan tidak mencukupi: Jika Anda hanya memiliki 20-50 contoh, coba untuk menggandakan ukuran dataset Anda secara kasar. Setelah Anda memiliki lebih dari 300 contoh, jika Anda masih melihat ini, tugasnya mungkin terlalu kompleks.

2. Beberapa tugas serupa: Ketika dataset Anda berisi beberapa tugas serupa (misalnya, mengambil 2 objek yang berbeda), model mungkin terlalu bergantung pada kondisioning bahasa yang mungkin tidak cukup kaya.

Tips debugging:

• Perbesar ukuran dataset (gandakan hingga mencapai lebih dari 300 contoh)
• Latih lebih lama, hingga 100k langkah, bahkan ketika loss mendatar
• Periksa apakah model menerima instruksi bahasa yang tepat atau tingkatkan keragaman instruksi

Menjalankan Tugas yang Salah

Kadang-kadang robot akan mengabaikan instruksi Anda sepenuhnya dan melakukan tugas lain. Hal ini umumnya hanya terjadi jika Anda telah melatih model pada beberapa tugas.

Penyebab potensial:

• Ambiguitas instruksi bahasa
• Data pelatihan spesifik tugas yang tidak memadai
• Kebingungan model antara tugas-tugas serupa dalam dataset multitugas

Tips debugging:

• Verifikasi spesifisitas instruksi bahasa, terutama jika deskripsi serupa di antara beberapa tugas
• Periksa distribusi tugas dalam dataset pelatihan Anda dan tambahkan pembobotan pada tugas yang gagal/diabaikan
• Pertimbangkan penyesuaian halus yang spesifik untuk tugas

Ketidakstabilan Pelatihan

Jika loss pelatihan tidak stabil atau divergen:

• Coba sesuaikan learning rate antara 1e-5 dan 3e-4
• Tingkatkan ukuran batch jika memungkinkan
• Periksa bahwa normalisasi dataset Anda benar
• Verifikasi prapemrosesan gambar berfungsi dengan benar

Pertimbangan Performa

Kebutuhan GPU

• Inferensi: Disarankan menggunakan setidaknya RTX 5070 Ti (atau GPU setara) untuk kinerja kecepatan yang wajar
• Pelatihan: GPU dengan VRAM yang cukup untuk memuat ukuran batch >64 sangat ideal, yang akan bervariasi tergantung pada jumlah pengamatan gambar, dll

Rekomendasi Ukuran Batch

• Minimum: 64 (di bawah nilai ini dapat menyebabkan pelatihan tidak stabil)
• Disarankan: 256-320 (kinerja terbaik, membutuhkan GPU yang lebih besar)

Contoh: Pelatihan pada Dataset Kustom

Berikut adalah contoh lengkap pelatihan pada dataset kustom:

lerobot-train \
  --dataset.repo_id=YOUR_DATASET \
  --output_dir=./outputs/mutitask_dit_training \
  --batch_size=320 \
  --steps=30000 \
  --save_freq=1000 \
  --log_freq=100 \
  --env_eval_freq=1000 \
  --policy.type=multi_task_dit \
  --policy.device=cuda \
  --policy.horizon=32 \
  --policy.n_action_steps=24 \
  --policy.objective=diffusion \
  --policy.noise_scheduler_type=DDPM \
  --policy.num_layers=6 \
  --policy.hidden_dim=512 \
  --policy.vision_encoder_name=openai/clip-vit-base-patch16 \
  --policy.image_resize_shape=[320,240] \
  --policy.image_crop_shape=[224,224] \
  --policy.repo_id="HF_USER/multitask-dit-your-robot" \
  --trackio.enable=true \
  --trackio.project=multitask_dit

Hasil Libero

lerobot-train \
  --dataset.repo_id=HuggingFaceVLA/libero \
  --policy.type=multi_task_dit \
  --policy.push_to_hub=false \
  --output_dir="./outputs/multitask_dit_libero" \
  --job_name="multitask-dit-libero" \
  --trackio.enable=true \
  --trackio.project=multitask_dit_libero \
  --dataset.image_transforms.enable=true \
  --dataset.image_transforms.max_num_transforms=4 \
  --dataset.image_transforms.tfs='{"brightness":{"type":"ColorJitter","kwargs":{"brightness":[0.75,1.25]}},"contrast":{"type":"ColorJitter","kwargs":{"contrast":[0.6,1.4]}},"saturation":{"type":"ColorJitter","kwargs":{"saturation":[0.8,1.2]}},"hue":{"type":"ColorJitter","kwargs":{"hue":[-0.05,0.05]}},"sharpness":{"type":"SharpnessJitter","kwargs":{"sharpness":[0.6,1.4]}},"rotation":{"type":"RandomRotation","kwargs":{"degrees":[-5,5]}},"translation":{"type":"RandomAffine","kwargs":{"degrees":0,"translate":[0.1,0.1]}}}' \
  --dataset.video_backend=torchcodec \
  --policy.use_amp=true \
  --policy.horizon=48 \
  --policy.n_obs_steps=2 \
  --policy.use_rope=true \
  --policy.use_positional_encoding=false \
  --policy.hidden_dim=768 \
  --policy.num_layers=8 \
  --policy.num_heads=12 \
  --policy.dropout=0.1 \
  --policy.timestep_embed_dim=256 \
  --policy.objective=diffusion \
  --policy.optimizer_lr=3e-4 \
  --policy.optimizer_weight_decay=0 \
  --policy.scheduler_warmup_steps=0 \
  --policy.vision_encoder_name=openai/clip-vit-base-patch16 \
  --policy.image_resize_shape=[256,256] \
  --policy.image_crop_is_random=true \
  --policy.text_encoder_name=openai/clip-vit-base-patch16 \
  --policy.vision_encoder_lr_multiplier=0.1 \
  --policy.device=cuda \
  --num_workers=8 \
  --save_freq=4000 \
  --log_freq=100 \
  --steps=100000 \
  --batch_size=320

Hasil:

Referensi

Untuk detail lebih lanjut tentang implementasi teknis dan arsitektur, lihat:

• Pemeriksaan Cermat terhadap Model Perilaku Besar untuk Manipulasi Lincah Multi-Tugas
• Model Perilaku Besar dan Atlas Menemukan Posisi Baru
• Membedah dan Membuka Sumber Kebijakan Transformer Difusi Multitugas