X-VLA BEHAVIOR-1K Checkpoints โ Weighted CE Fix
Fine-tuned X-VLA checkpoints on the BEHAVIOR-1K 2025 challenge demonstrations, with weighted cross-entropy loss for the skill classifier.
What's New: Weighted CE Fix
The original skill classifier suffered from severe class imbalance โ "move to" accounts for 34% of all frames, causing the classifier to always predict that single skill. This release fixes the issue with sqrt-inverse frequency class weighting:
- Weights range from 0.086 (move to) to 3.04 (ignite), a 35x ratio
- Normalized so mean weight = 1.0
- Applied to the auxiliary CE loss on the skill classifier head
This ensures the model learns to distinguish all 34 skill primitives, not just the dominant one.
Checkpoints
| Folder | Task(s) | Learning Rate | Steps | Description |
|---|---|---|---|---|
single_task_0_100k |
Task 0 | 2e-5 | 100k | Single-task (task 0 only) |
single_task_1_100k |
Task 1 | 2e-5 | 100k | Single-task (task 1 only) |
Training is ongoing โ 200k checkpoints and additional LR variants (1e-4) will be uploaded as they complete.
Architecture Innovations
All checkpoints include the following modifications to X-VLA:
Additive Task + Skill Soft Prompts: Separate
task_prompt_hubandskill_prompt_hubembeddings replace the single domain prompt. Combined additively:prompt = task_prompt[task_id] + skill_prompt[skill_id].VLM-based Skill Classifier: A linear head on pooled VLM features predicts the current skill primitive (34 classes). Trained with auxiliary weighted CE loss (ฮป=0.1). At inference, auto-predicts skill when ground truth is unavailable.
Enriched Language Instructions: Training appends skill + object context to task descriptions, e.g., "Turn on the radio. Current: pick up radio from coffee table."
23D Action Space:
base_qvel(3) + trunk_qpos(4) + arm_left(7) + grip_left(1) + arm_right(7) + grip_right(1)with delta-joint transformation on 17 dims.
Usage
from models.modeling_xvla import XVLA
from models.processing_xvla import XVLAProcessor
model = XVLA.from_pretrained("Hoshipu/xvla-behavior1k-weighted-ce/single_task_0_100k")
processor = XVLAProcessor.from_pretrained("Hoshipu/xvla-behavior1k-weighted-ce/single_task_0_100k")
Training Details
- Base model: X-VLA-Pt (pretrained)
- Optimizer: AdamW, betas=(0.9, 0.95)
- Schedule: 1k freeze steps โ 2k warmup โ cosine decay (min_lr_ratio=0.1)
- Batch size: 16 per GPU ร 4 GPUs = 64 effective
- Mixed precision: bf16
- Skill class weights: sqrt-inverse frequency, 34 classes, mean-normalized to 1.0
Related
- Previous checkpoints (without weighted CE): Hoshipu/xvla-behavior1k-checkpoints
- X-VLA: github.com/2toINF/X-VLA
- BEHAVIOR-1K: behavior.stanford.edu/behavior-1k