SARM: Mô hình hóa phần thưởng theo giai đoạn

SARM (Stage-Aware Reward Modeling) là một khung mô hình hóa phần thưởng dựa trên video cho các thao tác robot có tầm nhìn xa. Hướng dẫn này bao gồm cách huấn luyện các mô hình phần thưởng SARM và tùy chọn sử dụng chúng với Reward-Aligned Behavior Cloning (RA-BC).

Bài báo: SARM: Stage-Aware Reward Modeling for Long Horizon Robot Manipulation

Tại sao cần Mô hình Phần thưởng?

Behavior cloning tiêu chuẩn coi mọi khung hình trong dữ liệu minh chứng là ngang nhau, nhưng các bộ dữ liệu robot trong thực tế thường lộn xộn. Chúng chứa những khoảnh khắc do dự, các lần sửa lỗi và các quỹ đạo có chất lượng không đồng đều. Các mô hình phần thưởng giải quyết vấn đề này bằng cách học một khái niệm khái quát hóa về tiến độ nhiệm vụ từ các dữ liệu minh chứng: khi được cung cấp các khung hình video và mô tả nhiệm vụ, chúng dự đoán mức độ gần hoàn thành nhiệm vụ của robot (0→1). Tín hiệu “tiến độ” học được này có thể được sử dụng theo nhiều cách, trong đó hai ứng dụng đầy hứa hẹn là: (1) học bắt chước có trọng số (RA-BC), trong đó các khung hình có tiến độ cao nhận trọng số lớn hơn trong quá trình huấn luyện chính sách, và (2) học tăng cường, trong đó mô hình phần thưởng cung cấp phần thưởng dày đặc để cải thiện chính sách trực tuyến hoặc ngoại tuyến.

Tổng quan

SARM có các tính năng sau:

1. Kiến trúc nhận thức giai đoạn: Dự đoán đồng thời giai đoạn tác vụ cấp cao và tiến độ chi tiết trong từng giai đoạn
2. Ghi chú phụ nhiệm vụ: Sử dụng ghi chú phụ nhiệm vụ bằng ngôn ngữ tự nhiên để suy ra các nhãn tiến độ nhất quán
3. Tỷ lệ thời gian: Tính các tiên nghiệm ở cấp độ tập dữ liệu (α̅_k) cho từng nhiệm vụ con để chuẩn hóa tiến độ giữa các bài diễn giải có độ dài khác nhau

SARM huấn luyện trên mục tiêu stage+tau gọn nhẹ cho mỗi khung hình:

• stage: chỉ số stage nguyên k ∈ {0, ..., K-1}
• τ (tau): tiến độ trong giai đoạn τ ∈ [0, 1]
• mã hóa đích: y = k + τ (đây là thứ mà bộ xử lý tập dữ liệu tạo ra)

Vào thời điểm suy luận (và trong RA-BC hạ nguồn), SARM chuyển đổi giá trị thô k + τ thành tiến độ đã chuẩn hóa trong [0, 1] bằng cách sử dụng tỷ lệ thời gian ở cấp độ tập dữ liệu α̅_k (được lưu trong meta/temporal_proportions_*.json).

Điều này khớp với Công thức (2) trong bài báo:

progress_t = P_{k-1} + α̅_k × τ_t

Trong đó:

• τ_t = (t - s_k) / (e_k - s_k) là thời gian được chuẩn hóa trong từng tiểu tác vụ
• P_{k-1} là phần trước tích lũy (tổng các tỷ lệ của các nhiệm vụ con trước đó)
• α̅_k là tỷ lệ thời gian cho nhiệm vụ con k

Điều này đảm bảo rằng các trạng thái tác vụ giống hệt nhau được ánh xạ đến các giá trị tiến độ nhất quán, ngay cả giữa các bài diễn giải có độ dài khác nhau.

Đầu vào và Mục tiêu (Những gì mã mới mong đợi)

SARM được huấn luyện thông qua bộ xử lý của nó (src/lerobot/rewards/sarm/processor_sarm.py), bao gồm:

• Mã hóa hình ảnh và văn bản tác vụ bằng CLIP (ViT-B/32) thành video_features và text_features
• Đệm/cắt ngắn trạng thái robot thành state_features (tối đa max_state_dim)
• Xây dựng các mục tiêu dưới dạng sparse_targets (và dense_targets trong dense_only/dual) bằng cách mã hóa stage+tau y = k + τ
• Masks rewind frames sử dụng tensor lengths cho từng mẫu (rewind là một kỹ thuật tăng cường dữ liệu trong quá trình huấn luyện)

Tối thiểu, mỗi mẫu huấn luyện cần:

• task (chuỗi): mô tả tác vụ
• các hình ảnh reward_model.image_key và trạng thái reward_model.state_key từ bộ dữ liệu

Chế độ chú thích

Bạn có thể chọn từ 3 chế độ chú thích xác định cách tính nhãn tiến độ:

Chi tiết chế độ

pip install -e ".[sarm]"

1. Train SARM → 2. Visualize predictions → 3. (Optional) Train policy with RA-BC

1. Annotate (dense) → 2. Verify → 3. Train SARM → 4. Visualize → 5. (Optional) Train policy with RA-BC

1. Annotate (sparse+dense) → 2. Verify → 3. Train SARM → 4. Visualize → 5. (Optional) Train policy with RA-BC

Bước 1: Gán nhãn cho các nhiệm vụ con

python src/lerobot/data_processing/sarm_annotations/subtask_annotation.py \
  --repo-id your-username/your-dataset \
  --dense-only \
  --dense-subtasks "Bring robot arms up from starting position,Grab near side and do 1st fold,Grab side and do 2nd fold,Grab side and do 3rd fold to finish folding" \
  --video-key observation.images.base \
  --num-workers 4 \
  --push-to-hub

python src/lerobot/data_processing/sarm_annotations/subtask_annotation.py \
  --repo-id your-username/your-dataset \
  --sparse-subtasks "Bring arms up from starting position,Fold the towel (3 folds in total)" \
  --dense-subtasks "Bring robot arms up from starting position,Grab near side and do 1st fold,Grab side and do 2nd fold,Grab side and do 3rd fold to finish folding" \
  --video-key observation.images.base \
  --num-workers 4 \
  --push-to-hub

Đối số chú thích

Lưu ý: Sau khi chú thích hoàn tất, mặc định sẽ tự động trực quan hóa 5 tập. Sử dụng --num-visualizations 0 để bỏ qua bước này.

Bước 2: Xác minh chú thích

python src/lerobot/data_processing/sarm_annotations/subtask_annotation.py \
  --repo-id your-username/your-dataset \
  --visualize-only \
  --visualize-type dense \
  --num-visualizations 5 \
  --video-key observation.images.base \
  --output-dir ./subtask_viz

python src/lerobot/data_processing/sarm_annotations/subtask_annotation.py \
  --repo-id your-username/your-dataset \
  --visualize-only \
  --visualize-type both \
  --num-visualizations 5 \
  --video-key observation.images.base \
  --output-dir ./subtask_viz

Điều này tạo ra các hình ảnh trực quan hóa hiển thị các khung hình video với các ranh giới của các nhiệm vụ con được phủ lên và dòng thời gian của các nhiệm vụ con.

Các đối số trực quan hóa

Mẹo: Nếu chú thích không chính xác, hãy điều chỉnh mô tả các nhiệm vụ con cho cụ thể hơn và chạy lại.

Bước 3: Huấn luyện SARM

lerobot-train \
  --dataset.repo_id=your-username/your-dataset \
  --reward_model.type=sarm \
  --reward_model.annotation_mode=single_stage \
  --reward_model.image_key=observation.images.base \
  --output_dir=outputs/train/sarm_single \
  --batch_size=32 \
  --steps=5000 \
  --trackio.enable=true \
  --trackio.project=sarm \
  --reward_model.repo_id=your-username/your-model-name

lerobot-train \
  --dataset.repo_id=your-username/your-dataset \
  --reward_model.type=sarm \
  --reward_model.annotation_mode=dense_only \
  --reward_model.image_key=observation.images.base \
  --output_dir=outputs/train/sarm_dense \
  --batch_size=32 \
  --steps=5000 \
  --trackio.enable=true \
  --trackio.project=sarm \
  --reward_model.repo_id=your-username/your-model-name

lerobot-train \
  --dataset.repo_id=your-username/your-dataset \
  --reward_model.type=sarm \
  --reward_model.annotation_mode=dual \
  --reward_model.image_key=observation.images.base \
  --output_dir=outputs/train/sarm_dual \
  --batch_size=32 \
  --steps=5000 \
  --trackio.enable=true \
  --trackio.project=sarm \
  --reward_model.repo_id=your-username/your-model-name

Đào tạo đa GPU

Thêm accelerate launch --multi_gpu --num_processes=4 để sử dụng nhiều GPU cho việc huấn luyện.

Tham số huấn luyện

Bước 4: Trực quan hóa các dự đoán

Sử dụng compute_rabc_weights.py với --visualize-only để trực quan hóa các dự đoán của mô hình (và, nếu có, các mục tiêu được suy ra từ chú thích) mà không ghi tệp parquet.

python -m lerobot.rewards.sarm.compute_rabc_weights \
  --dataset-repo-id your-username/your-dataset \
  --reward-model-path your-username/sarm-model \
  --visualize-only \
  --num-visualizations 5 \
  --head-mode sparse \
  --output-dir ./sarm_viz

python -m lerobot.rewards.sarm.compute_rabc_weights \
  --dataset-repo-id your-username/your-dataset \
  --reward-model-path your-username/sarm-model \
  --visualize-only \
  --num-visualizations 5 \
  --head-mode dense \
  --output-dir ./sarm_viz

python -m lerobot.rewards.sarm.compute_rabc_weights \
  --dataset-repo-id your-username/your-dataset \
  --reward-model-path your-username/sarm-model \
  --visualize-only \
  --num-visualizations 5 \
  --head-mode both \
  --output-dir ./sarm_viz

Trực quan hóa cho thấy:

• Biểu đồ tiến độ: Tiến độ dự kiến (và “GT” tùy chọn có nguồn gốc từ chú thích khi có sẵn và --stride 1)
• Xác suất giai đoạn: Biểu đồ diện tích chồng chất của xác suất giai đoạn dự đoán
• Khung mẫu: Các khung chính từ tập với nhãn tiến độ/giai đoạn

Các đối số trực quan hóa

Bước 5 (Tùy chọn): Huấn luyện Policy với RA-BC

Behavior Cloning theo hướng phần thưởng (RA-BC) sử dụng mô hình SARM đã được huấn luyện để gán trọng số cho các mẫu huấn luyện dựa trên mức cải thiện tiến độ dự kiến. Quy trình này yêu cầu hai bước:

1. Tính trước các giá trị tiến độ cho tất cả các khung hình bằng mô hình SARM đã được huấn luyện
2. Đào tạo chính sách với trọng số RA-BC bằng cách sử dụng các giá trị đã tính trước

Cách RA-BC hoạt động

Với mỗi mẫu huấn luyện, RA-BC tính toán độ tiến bộ delta:

r_i = φ(o_{t+Δ}) - φ(o_t)

Trong đó φ là dự đoán tiến độ SARM và Δ là chunk_size của chính sách. Các mẫu có tiến độ dương (các minh chứng tốt) nhận trọng số cao hơn, trong khi các mẫu có tiến độ âm hoặc bằng không bị giảm trọng số.

Việc gán trọng số tuân theo Phương trình 8-9 trong bài báo:

• Trọng số mềm: w̃_i = clip((r_i − (μ − 2σ)) / (4σ + ε), 0, 1)
• Trọng số cuối cùng: w_i = 𝟙{r_i > κ} + 𝟙{0 ≤ r_i ≤ κ} × w̃_i

Bước 5a: Tính các giá trị tiến độ SARM

Đầu tiên, chạy mô hình SARM trên tất cả các khung hình trong bộ dữ liệu của bạn để tính các giá trị tiến độ:

python -m lerobot.rewards.sarm.compute_rabc_weights \
  --dataset-repo-id your-username/your-dataset \
  --reward-model-path your-username/sarm-model \
  --head-mode sparse \
  --num-visualizations 5 \
  --push-to-hub

Script này:

• Xử lý tất cả các khung hình và tính toán các giá trị tiến độ
• Lưu các giá trị tiến độ vào tệp parquet nằm cạnh bộ dữ liệu trên đĩa (mặc định là <dataset_root>/sarm_progress.parquet)
• Tạo trực quan hóa cho N tập đầu tiên (mặc định: 5)

Các đối số:

Định dạng đầu ra (sarm_progress.parquet):

Bước 5b: Huấn luyện Policy với RA-BC

Khi đã có tệp tiến độ, hãy huấn luyện chính sách của bạn với trọng số RA-BC. Tệp tiến độ sẽ được tự động phát hiện từ đường dẫn bộ dữ liệu (sarm_progress.parquet) nếu không được cung cấp rõ ràng. Hiện tại, PI0, PI0.5, SmolVLA và FastWAM được hỗ trợ với RA-BC:

lerobot-train \
  --dataset.repo_id=your-username/your-dataset \
  --policy.type=pi0 \
  --sample_weighting.type=rabc \
  --sample_weighting.head_mode=sparse \
  --sample_weighting.kappa=0.01 \
  --output_dir=outputs/train/policy_rabc \
  --batch_size=32 \
  --steps=40000

Tệp kịch bản huấn luyện tự động:

• Tải các giá trị tiến độ đã tính trước từ tệp parquet
• Sử dụng chunk_size của chính sách để tính toán các mức tăng tiến độ (Δ)
• Tính toán trọng số mẫu dựa trên mức độ cải thiện tiến độ
• Áp dụng tổn thất có trọng số trong quá trình huấn luyện

Lập luận RA-BC:

Điều chỉnh RA-BC Kappa

Tham số kappa là ngưỡng xác định mẫu nào nhận trọng số đầy đủ (w=1). Việc hiểu cách tinh chỉnh tham số này là rất quan trọng để RA-BC hoạt động hiệu quả.

Cách thức hoạt động của trọng số:

Chẩn đoán các vấn đề về kappa:

Theo dõi các chỉ số Trackio này trong quá trình huấn luyện:

Nếu sample_weight_mean_weight ≈ 1.0: Kappa của bạn quá thấp. Hầu hết các mẫu có delta > kappa và bỏ qua hoàn toàn việc gán trọng số mềm. RA-BC trở nên tương đương với BC thông thường.

Đặt kappa dựa trên dữ liệu của bạn:

Giá trị mặc định kappa=0.01 được tinh chỉnh cho tác vụ gấp áo thun trong bài báo (các tập hợp ~90 giây ở 30fps). Đối với bộ dữ liệu của bạn, hãy kiểm tra sample_weighting/delta_mean và sample_weighting/delta_std được ghi log:

# If delta_mean ≈ 0.03 and delta_std ≈ 0.02:
# Most deltas fall in range [0.01, 0.05]

# Option 1: Set kappa = delta_mean (medium selectivity)
--sample_weighting.kappa=0.03

# Option 2: Set kappa = delta_mean + delta_std (high selectivity)
--sample_weighting.kappa=0.05

# Option 3: Set kappa = delta_mean + 2*delta_std (very selective)
--sample_weighting.kappa=0.07

Khi nào RA-BC có thể không hữu ích:

Nếu tập dữ liệu của bạn đã có chất lượng cao (tiến bộ nhất quán trên tất cả các minh chứng), RA-BC sẽ không mang lại nhiều lợi ích vì không có gì để lọc.

Đào tạo đa GPU với RA-BC

accelerate launch \
  --multi_gpu \
  --num_processes=4 \
  src/lerobot/scripts/lerobot_train.py \
  --dataset.repo_id=your-username/your-dataset \
  --policy.type=pi0 \
  --sample_weighting.type=rabc \
  --sample_weighting.kappa=0.01 \
  --output_dir=outputs/train/policy_rabc \
  --batch_size=32 \
  --steps=40000

Mẹo & Thực hành tốt

Chọn một Chế độ

• Bắt đầu với single_stage cho các thí nghiệm nhanh - không cần chi phí gán nhãn
• Sử dụng dense_only khi bạn muốn theo dõi tiến độ chi tiết nhưng các tác vụ không có các giai đoạn cấp cao rõ ràng
• Sử dụng dual cho các tác vụ phức tạp, nơi cả tiến độ tổng quát lẫn chi tiết đều có ý nghĩa

Chất lượng chú thích

1. Cụ thể hóa tên các nhiệm vụ con: Thay vì “fold”, hãy dùng “grab near side and fold toward center”
2. Xác minh bằng trực quan hóa: Luôn kiểm tra một vài tập trước khi huấn luyện
3. Đặt tên nhất quán: Sử dụng cùng tên các nhiệm vụ con trong tất cả các tập

RA-BC

1. Đào tạo SARM trước: Chất lượng của RA-BC phụ thuộc hoàn toàn vào chất lượng của SARM
2. Theo dõi sample_weight_mean_weight: Nếu giá trị ≈ 1.0, hãy tăng kappa (xem Điều chỉnh Kappa của RA-BC)

Trích dẫn

@article{chen2025sarm,
  title={SARM: Stage-Aware Reward Modeling for Long Horizon Robot Manipulation},
  author={Chen, Qianzhong and Yu, Justin and Schwager, Mac and Abbeel, Pieter and Shentu, Yide and Wu, Philipp},
  journal={arXiv preprint arXiv:2509.25358},
  year={2025}
}