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
```bash
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:
```bash
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:
```bash
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:
```bash
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:
```python
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:
```bash
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:
```bash
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
```bash
make train-smoke
```
This creates a small toy dataset and trains for one epoch.