vla / docs /training.md
anhtld's picture
Initial commit: DoVLA-CIL codebase (h=16 breakthrough)
adc02fa verified
|
Raw
History Blame
8.32 kB

Training

DoVLA-CIL trains from group-aware CIL batches. The first implementation is lightweight and CPU-friendly for toy symbolic observations, while preserving extension points for visual VLA backbones.

Dataset and Sampler

CILDataset loads sharded JSONL datasets and exposes:

  • record indexing
  • get_group(group_id)
  • iter_groups()

GroupAwareBatchSampler supports:

  • full_group: complete groups in each batch
  • pairs: same-group positive/negative pairs for ranking
  • mixed: grouped records with configurable records per group

The collator returns observation tensors where possible and preserves metadata for action chunks, effects, rewards, regrets, ranks, candidate types, group IDs, and failures.

Interventional Action Field Objective

The default lattice_field objective predicts a shared effect vector e_i and utility potential u_i for every action intervention. Within each same-state group, action-lattice edges supervise u_i - u_j with measured utility differences and e_i - e_j with measured effect differences. The potential edge energy combines three same-state terms: magnitude regression on measured utility differences, an order-margin hinge, and a Bradley-Terry preference likelihood on the sign of each measured edge. All three depend only on u_i - u_j, so the objective remains invariant to per-state reward offsets, and the scalar potential makes cycle sums zero by construction. K up to 32 uses the complete same-state graph by default; lower --lattice-neighbors values are an explicit sparse-graph ablation. Best-action BC anchors policy decoding; a small absolute effect/progress/success term anchors the otherwise free group offset.

The preference term is controlled by field_preference in InterventionalLossWeights, for example --loss-weight field_preference=0.5. Setting it to 0 recovers the earlier regression-plus-margin field objective.

Use --objective legacy only for the ablation that combines absolute BC, effect, success, progress, ranking, and regret losses separately.

Train

python scripts/train_dovla.py \
  --dataset data/cil_toy \
  --out runs/dovla_toy \
  --epochs 5 \
  --batch-groups 8 \
  --records-per-group 8 \
  --hidden-dim 256 \
  --lr 0.001 \
  --device auto \
  --objective lattice_field \
  --lattice-neighbors 32

The trainer saves:

  • latest.pt
  • best.pt
  • metrics.json

Validation splits are by group_id, not by individual records, so ranking examples do not leak across train/val.

Model Skeleton

The default DoVLAModel has:

  • toy symbolic observation encoder
  • hashed bag-of-words language encoder
  • action encoder
  • policy head
  • interventional field head for shared effect and utility potential
  • legacy effect/reward/regret heads for controlled ablations

Toy action vectorization/de-vectorization utilities live in the model/action encoder modules.

VLA Backbone Adapters

External VLA models remain optional. dovla_cil/models/openvla_adapter.py defines:

  • VLABackbone protocol
  • ToyVLABackbone
  • PretrainedCLIPBackbone
  • ExternalOpenVLAAdapter placeholder

PretrainedCLIPBackbone is the reproducible pretrained VLM baseline. It replaces the native image/language encoders while retaining the same action encoder, policy head, and Interventional Action Field. CLIP is frozen by default; one normalized image/text feature pair is cached per CIL group_id, while the context projection and all DoVLA components remain trainable. Frozen CLIP weights are omitted from DoVLA checkpoints and reloaded from the pinned model directory.

Download the public model once on a network-enabled node and pin the revision:

hf download openai/clip-vit-base-patch32 \
  --revision 3d74acf9a28c67741b2f4f2ea7635f0aaf6f0268 \
  --local-dir /scratch/$USER/dovla/models/openai-clip-vit-base-patch32-3d74acf

Train from rendered CIL observations without network access:

TRANSFORMERS_OFFLINE=1 HF_HUB_OFFLINE=1 python scripts/train_dovla.py \
  --dataset data/cil_maniskill_rgb \
  --out runs/dovla_clip \
  --observation-mode rgb \
  --backbone clip \
  --backbone-model /scratch/$USER/dovla/models/openai-clip-vit-base-patch32-3d74acf \
  --backbone-feature-cache /scratch/$USER/dovla/caches/cil_clip_features.pt

This is an external pretrained vision-language baseline, not an OpenVLA claim. A future OpenVLA integration still uses VLABackbone; the repository does not vendor or silently download OpenVLA.

Full External VLA Baseline Bridge

The repository also includes a small bridge for real external VLA policy baselines such as SmolVLA or OpenVLA:

python scripts/export_lerobot_dataset.py \\
  --dataset /scratch/$USER/dovla/experiments/maniskill_presuccess_six_task_collection \\
  --out runs/external_vla/lerobot_export \\
  --selection expert \\
  --group-sampling task_balanced \\
  --seed 0

python scripts/run_external_vla_baseline.py \\
  --model-family smolvla \\
  --dataset runs/external_vla/lerobot_export \\
  --checkpoint /scratch/$USER/dovla/models/smolvla_base-c83c316 \\
  --out runs/external_vla/smolvla \\
  --adapter-entrypoint dovla_cil.eval.smolvla_cil_baseline:run_smolvla_cil_baseline \\
  --adapter-config configs/external/smolvla_cil_smoke.json \\
  --dry-run

export_lerobot_dataset.py creates a dependency-light interchange export rather than requiring LeRobot inside DoVLA-CIL. Each JSONL row contains observation.image or the original cil_observation_ref, the flattened numeric action, the full action_chunk, instruction text, success/reward, and CIL provenance (group_id, record_id, state_hash, rank, regret, and candidate type). Expert selection with task-balanced sampling gives the isolated SmolVLA runtime a controlled BC baseline without leaking validation groups through state ordering.

Dry-run mode writes external_vla_baseline_plan.json with a secret-free environment, download, and run plan. The repository includes a SmolVLA adapter; additional external models can provide an --adapter-entrypoint module:function implementing:

def run(spec_dict: dict, plan_dict: dict) -> dict:
    ...

The bundled SmolVLA entrypoint fine-tunes on expert rows and evaluates each prediction by selecting the nearest action actually executed from the same serialized state. Reward, success, and regret come from those measured outcomes. This protocol tests candidate selection and is not presented as online policy rollout. Keeping LeRobot in its isolated runtime prevents dependency conflicts with the stable ManiSkill/DoVLA environment.

Before a measured run, verify and load-test the staged SmolVLA checkpoint:

python scripts/verify_external_checkpoint.py \
  --checkpoint /scratch/$USER/dovla/models/smolvla_base-c83c316 \
  --out outputs/external_vla_smolvla_checkpoint_manifest.json \
  --model-family smolvla

sbatch scripts/slurm/smoke_smolvla_checkpoint.sbatch

The smoke job runs with Hub and Transformers offline flags. Its JSON artifact records package versions, device, policy class, parameter counts, and checkpoint load time, so successful loading is not inferred from file presence alone.

For a matched scientific comparison, export all 3,500 expert groups and run the aligned config:

export DATASET=/scratch/$USER/dovla/experiments/maniskill_presuccess_six_task_collection
export OUT=/scratch/$USER/dovla/experiments/external_vla_export_full_aligned
export SELECTION=expert GROUP_SAMPLING=task_balanced MAX_GROUPS=3500 SEED=0
sbatch scripts/slurm/export_lerobot_dataset.sbatch

export ADAPTER_CONFIG=/workspace/configs/external/smolvla_cil_aligned.json
export OUT=/scratch/$USER/dovla/experiments/smolvla_cil_aligned
sbatch scripts/slurm/run_smolvla_cil_baseline.sbatch

The aligned config uses the same deterministic group shuffle as the core trainer. It trains on 2,800 groups and evaluates on the identical 700 held-out groups. SmolVLA expert-only BC reaches top-1 0.5229, selected success 0.3457, and selected regret 0.1366; DoVLA-IAF seed 0 reaches 0.6171, 0.3786, and 0.0599. Both rows select among measured same-state candidates. These numbers do not constitute an online SmolVLA rollout comparison.

Smoke Training

make train-smoke

This creates a small toy dataset and trains for one epoch.