Jocasta-diff-VLA

Code and loader: alanrbtx/Diffusion-VLA-tuning

Jocasta-diff-VLA is a Dream-VLA-7B adaptation to eight simulated Panda manipulation task families in robosuite. It uses the Dream-VLX authors' continuous Flow Matching action head and a diffusion language-model backbone. One jointly trained policy achieved 332/512 sustained successes (64.8%) on the fixed development evaluation.

This repository contains an inference checkpoint, not just a LoRA language adapter: the LoRA weights, trained action head, noisy-action projector, proprioception projector, dataset normalization, processor files and training configuration are included. The pinned original Dream-VLA base weights are required separately. The public loader downloads both when needed. No optimizer or training-resume state is distributed.

Load

Use the accompanying repository with its pinned submodules and patch:

git clone --recurse-submodules https://github.com/alanrbtx/Diffusion-VLA-tuning.git
cd Diffusion-VLA-tuning
git -C third_party/Dream-VLX apply ../../patches/dream-vlx-panda.patch
python -m pip install -r requirements.lock -r requirements-multitask.txt

Run inference on one visible CUDA GPU:

from load_policy import load_policy, predict_action

policy, processor, cfg = load_policy('RobotLearningLab/Jocasta-diff-VLA')
# primary and wrist: uint8 RGB arrays, each (224,224,3)
# proprio: float32 (8,), xyz + axis-angle orientation + two gripper joints
actions = predict_action(policy, processor, cfg, primary, wrist, proprio,
                         'pick up the red cube', seed=84907)
# actions.shape == (8,7)

For reproducibility, pass the desired Hub commit with revision=.... The base is fixed to Dream-org/Dream-VLA-7B@0cd0e25c1771c5b468c74585b5cf35c4a4aa315c; base_model_path=... reuses a local copy. This is a custom VLA package; AutoModel.from_pretrained or loading only the adapter will not reconstruct the trained action policy.

Action columns are world-frame translation deltas in metres, extrinsic XYZ Euler deltas in radians, and gripper openness. They require the documented conversion to normalized Panda OSC commands; see panda_adapter.py. Proprioception uses the recorded quaternion-to-axis-angle convention in robot_state.py. In simulation, keep CUDA inference and EGL rendering on separate physical GPUs. Model inference itself creates no EGL context.

Training

Setting	Value
Base	Dream-VLA-7B, pinned revision above
Upstream code	DreamLM/Dream-VLX a72fa4106eae97a923b531342fb692742809782f
Objective	Author Flow Matching velocity MSE
Optimizer	AdamW, learning rate1e-4
Updates	100,000; fixed final checkpoint
Global batch	64 (4 GPUs × microbatch8 × accumulation2)
LoRA	rank32, alpha16, dropout0, all linear layers
Training/sampling seeds	84907 /190019
Inputs	two224RGB cameras,8D proprioception, instruction
Inference	4 Flow Matching steps; action chunk8
Training data	8,140 successful trajectories;4,841,083 timesteps
Task/language diversity	8 families,30 goals,510 train instruction strings

The model was initialized from the original Dream-VLA base with fresh action modules, not from the earlier single-task Lift model. Training used family/goal balancing and the author's image augmentation. A separate120-string held-out-language development bank was excluded from training; the results below use canonical instructions and do not constitute a held-out-language evaluation. Raw demonstration data are not included; the code documents the exact cache schema.

Evaluation

64 fixed development scenes per family, seeds30000–30063, original physical scene configuration and horizons,10 settling steps, hard_reset=False. Success requires the original task predicate to hold for10 consecutive physical steps. All512 episodes were independently audited, including45,254 policy calls and365,450 physical steps. One training seed; no reserved test scenes, checkpoint selection or metric-driven reruns.

Task	Successes	Rate
Lift	64/64	100.0%
PickPlace	40/64	62.5%
Stack	43/64	67.2%
Door	34/64	53.1%
NutAssembly	32/64	50.0%
ToolHang	9/64	14.1%
Push	61/64	95.3%
Drawer	49/64	76.6%
Total	332/512	64.8%

Original pretrained Dream-VLA achieved 0/64 on the same PickPlace scenes; Jocasta-diff-VLA achieved 40/64. The original policy used one image, one action,4 discrete-diffusion steps and native Bridge statistics. The adapted policy uses two images/proprio,8 actions,4 Flow Matching steps and Panda statistics. This is a comparison of complete transferred policies, not an isolated backbone or diffusion-mechanism effect. No baseline result is claimed for the other seven tasks. See evaluation-results.json for counts and per-family Wilson intervals.

Scope and limitations

These are development results on our robosuite/Panda task definitions, not official LIBERO scores. ToolHang is weak (9/64). Other robots, unseen task families, real hardware and the reserved test split have not been validated. This release implements manipulation skills, not autonomous replication or learned mission-memory inheritance.

The release loader is checked against the original checkpoint on saved observations from all eight task families. validation.json records the outcome. release.json records file hashes; the original adapter and action-component weight bytes are retained unchanged. Reference inference environment: Python3.10.18, PyTorch2.5.1+cu124, Transformers4.51.3, BF16/SDPA on H200. No task evaluation was rerun for packaging.

Attribution

Architecture, Flow Matching loss/sampler, action heads and projectors originate from Dream-VLX, with upstream OpenVLA/OpenVLA-OFT attribution retained in the code. Panda adaptations and release: RobotLearningLab. Upstream code and model terms remain applicable; this adaptation is released under Apache-2.0.

@article{ye2025dreamvla,
  title={Dream-VL & Dream-VLA: Open Vision-Language and Vision-Language-Action Models with Diffusion Language Model Backbone},
  author={Ye, Jiacheng and Gong, Shansan and Gao, Jiahui and Fan, Junming and Wu, Shuang and Bi, Wei and Bai, Haoli and Shang, Lifeng and Kong, Lingpeng},
  journal={arXiv preprint arXiv:2512.22615},
  year={2025}
}

Validated code revision: c21ef87ef2d749b1eb2d06f7eb91e4c2b843bcd6.