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:

  1. Additive Task + Skill Soft Prompts: Separate task_prompt_hub and skill_prompt_hub embeddings replace the single domain prompt. Combined additively: prompt = task_prompt[task_id] + skill_prompt[skill_id].

  2. 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.

  3. Enriched Language Instructions: Training appends skill + object context to task descriptions, e.g., "Turn on the radio. Current: pick up radio from coffee table."

  4. 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

Downloads last month

-

Downloads are not tracked for this model. How to track
Video Preview
loading