vla / docs /architecture.md
anhtld's picture
Initial commit: DoVLA-CIL codebase (h=16 breakthrough)
adc02fa verified
|
Raw
History Blame
5.18 kB
# Architecture
DoVLA-CIL is organized around one invariant: every intervention in a group starts from the same
serialized simulator state. The codebase keeps task generation, simulation, intervention sampling,
effect extraction, data storage, training, and evaluation separated so real simulator backends can
be added without rewriting the research pipeline.
## Package Boundaries
- `dovla_cil.config`: YAML defaults, typed config objects, environment expansion, CLI overrides,
and resolved-config saving.
- `dovla_cil.vlm`: OpenAI-compatible VLM client, prompt templates, task generation, and optional
semantic failure annotation.
- `dovla_cil.tasks`: task schemas, validators, symbolic predicates, and built-in toy/CausalStress
task libraries.
- `dovla_cil.sim`: simulator protocol, toy backend, registry, and optional ManiSkill/Genesis
skeletons.
- `dovla_cil.interventions`: action schemas, perturbations, language/physics counterfactual
descriptors, and intervention samplers.
- `dovla_cil.effects`: structured effect extraction, reward computation, and deterministic failure
classification.
- `dovla_cil.data`: CIL record/group schemas, JSONL sharding, indices, datasets, group-aware
sampling, and collation support.
- `dovla_cil.models`: DoVLA encoders and heads plus one backbone boundary shared by native state,
native RGB, pinned pretrained CLIP, and future external VLA adapters.
- `dovla_cil.training`: interventional losses, batch collation, trainer, checkpoints, and metrics.
- `dovla_cil.eval`: CausalStress and downstream benchmark placeholders.
- `dovla_cil.experiments`: scaling laws, baselines, reports, and paper artifact helpers.
- `dovla_cil.generation`: local generation pipeline and optional Ray distributed generation.
- `dovla_cil.transfercritic`: optional data-curation critic for set-conditioned marginal utility
selection. It is not used by core training unless explicitly imported by an experiment.
- `dovla_cil.retrieval`: optional critic-gated exemplar retrieval for inference-time policy
conditioning. It is not part of core training unless explicitly wrapped around a policy.
## Data Flow
1. Load or generate validated `TaskSpec` objects.
2. Reset a simulator backend to a task and scene.
3. Serialize the exact simulator state.
4. Render the initial observation and symbolic state.
5. Plan or load an expert action, then sample `K` interventions.
6. For each intervention, restore the exact state, execute the action, and record outcomes.
7. Extract structured effects, reward, failure type, regret, and rank within the group.
8. Write grouped CIL records into shards and indices.
9. Train/evaluate with group-aware datasets and same-state losses.
For ManiSkill, steps 3-6 are vectorized over both distinct states `G` and interventions `K`.
Physical measurement and RGB observation are deliberately separate: GPU PhysX writes a versioned
archive of exact initial and next states, then a CPU renderer observes those fixed states without
changing actions, rewards, or success labels. Images are JPEG-compressed inside one HDF5 archive,
with stable references in each CIL record.
## Core Learning Invariant
Core training uses one `InterventionalFieldHead` to predict an effect embedding and scalar utility
potential for `(state, language, action)`. Same-group edges supervise differences in potential and
effect. A scalar potential makes lattice comparisons integrable and path-independent, while edge
differences cancel state-specific reward offsets. BC on the best action and a small absolute anchor
resolve decoding and field-offset ambiguity. Separate reward/ranking/regret heads are retained only
for the `legacy` ablation.
The pretrained CLIP path changes only observation-language encoding. It uses the same action
encoder, policy decoding, field head, losses, sampler, and evaluator as native DoVLA. Because CLIP
is frozen, image/text features contain no action or reward labels and can be cached once per
`group_id`; group-aware splits and all supervised learning still occur after that cache boundary.
Compact checkpoints omit frozen public weights and record the pinned local model path.
## Simulator Contract
Backends implement `SimulatorBackend`:
```python
seed(seed)
reset_task(task, scene=None)
serialize_state()
restore_state(state_blob)
render_observation()
get_symbolic_state()
execute_action_chunk(action)
close()
```
The toy backend implements this contract today. ManiSkill3 and Genesis wrappers are optional
skeletons that fail cleanly when their packages are not installed.
## Extension Points
- Add new task families in `dovla_cil.tasks.library` and validate with `tasks.validators`.
- Add new simulator adapters through `dovla_cil.sim.registry`.
- Add intervention types by extending `InterventionSampler` metadata conventions.
- Add real visual/language backbones through `models/openvla_adapter.py`.
- Add large-scale runners through `generation/distributed.py` or cluster-specific launchers.
- Add optional data-curation studies through `transfercritic/` without changing core trainers.
- Add optional inference-time retrieval through `retrieval/` without changing model checkpoints.