Jocasta-diff-VLA

Jocasta-diff-VLA — inference-чекпоинт от RobotLearningLab на базе Dream-VLA-7B для восьми семейств симулированных задач манипуляции Panda в robosuite. Модель использует диффузионный языковой бэкбон и непрерывную Flow Matching action head от авторов Dream-VLX. Одна совместно обученная политика набрала 332/512 (64.8%) успехов на фиксированной development-выборке.

⚙ Ключевые особенности

• 🧩 Один policy для 8 семейств задач: Lift, PickPlace, Stack, Door, NutAssembly, ToolHang, Push и Drawer.
• ⚙ В репозитории полный inference-чекпоинт: LoRA, обученная action head, проекторы шумных действий и проприоцепции, нормализация датасета, processor и конфиг обучения.
• 🎯 На входе две RGB-камеры 224×224, 8D проприоцепция и текстовая инструкция; на выходе — chunk из 8 действий по 7 компонентов.
• ⚡ Инференс: 4 шага Flow Matching, BF16/SDPA, одна видимая CUDA GPU; в симуляции CUDA-инференс и EGL-рендеринг лучше держать на разных физических GPU.
• 🧬 База зафиксирована на Dream-org/Dream-VLA-7B@0cd0e25c1771c5b468c74585b5cf35c4a4aa315c; исходные веса нужны отдельно, а AutoModel.from_pretrained не восстановит обученную политику.

🧩 Пример

from load_policy import load_policy, predict_action
policy, processor, cfg = load_policy('RobotLearningLab/Jocasta-diff-VLA')
# primary, wrist: uint8 RGB (224,224,3); proprio: float32 (8,)
actions = predict_action(policy, processor, cfg, primary, wrist, proprio,
                         'pick up the red cube', seed=84907)
# actions.shape == (8, 7)

Подробнее

Результат: 332/512 (64.8%) на фиксированной dev-выборке; на PickPlace — 40/64 против 0/64 у исходной Dream-VLA.
Слабое место — ToolHang: 9/64 (14.1%); другие роботы, реальное железо и reserved test split не проверялись.
Обучение: 8,140 успешных траекторий, 4,841,083 шага, 100,000 обновлений, LoRA rank 32.
Референсный инференс: Python 3.10.18, PyTorch 2.5.1+cu124, Transformers 4.51.3, BF16/SDPA на H200; одна CUDA GPU.
Это development-результаты на задачах robosuite/Panda, не официальные LIBERO scores; для остальных семи задач baseline не заявлен.