How kiyas-mahmud/webagent works
1. What this repo is
- This is an offline ML research/training repo for a failure-aware web agent.
- The central design is backbone-agnostic: one model template stays fixed, while only the front-end encoder/backbone changes across many variants.
- The core prediction stack always ends in one fused 768-d embedding feeding four task heads for the four “pillars” of the agent: detect, diagnose, recover, remember.
2. Top-level system boundary
- The repository is organized around a single pipeline:
dataset preparation → verification → dataloading → model forward pass → multi-task loss → training/checkpointing → evaluation/reporting.
- There is no web serving layer here; the architecture is for training and evaluation, not deployment.
3. Data subsystem
- scripts/setup_data.py and src/web_agent/data/verify_dataset.py are the entry points for turning the raw dataset into a usable local training corpus.
- The dataset is built around Mind2Web-style step data with failure augmentation, plus a shared images/ folder that should not be split.
- src/web_agent/data/dataset.py, gold_dataset.py, and dataloader.py define the main sample format and batching logic, including the five mode filters mentioned in the README.
- gold_dataset.py / gold_dataloader.py appear to support the “gold” subset used for the highest-quality training path.
4. Model boundary: fixed core, swappable encoder
- The src/web_agent/models/ package is the main architectural boundary.
- It is split into:
- encoders/ for backbone-specific input processing,
- fusion.py for combining modalities,
- adapter.py for projecting backbone outputs into the fixed hidden size,
- heads.py for the four task heads,
- loss.py for the shared weighted objective,
- model.py for assembling the full end-to-end network.
- This is the key diagram separation: backbone varies, downstream task stack does not.
5. Two encoder families
- The README describes two major patterns:
- Dual-encoder models: vision_enc + text_enc -> cross-attention -> 768-d
- Unified VLM models: VLM(image+text) -> linear adapter -> 768-d
- In code, these correspond to different implementations under src/web_agent/models/encoders/vision.py, text.py, and vlm.py.
- This is a meaningful architectural split because it changes where multimodal fusion happens: either before the shared 768-d representation or inside the backbone itself.
6. Task heads and loss
- The model always produces one shared latent embedding, then branches into four fixed heads.
- heads.py likely maps that shared representation into the four outputs required by the project’s pillars.
- loss.py combines them into a single weighted multi-task loss, which the README says is fixed across all backbones.
- labels.py is the likely source of the canonical target definitions used by both training and evaluation.
7. Memory subsystem
- src/web_agent/memory/index.py is a distinct subsystem for the remember pillar.
- Architecturally, this suggests an external retrieval/index component rather than a pure parametric prediction head.
- It should be drawn as a sidecar capability that consumes past interaction traces or embeddings and can influence the memory-related output path.
8. Training pipeline
- src/web_agent/train/ is staged:
- smoke_test.py for tiny correctness runs,
- mini_train.py for a small but real training slice,
- train.py and trainer.py for the full training loop.
- The README emphasizes a strict progression: SMOKE → MINI → FULL → EVAL → CHECKPOINT.
- This makes the training pipeline a gated system, not a single monolithic “train” command.
9. Evaluation subsystem
- src/web_agent/eval/metrics.py and evaluate.py implement model assessment.
- The README indicates evaluation spans three test splits, so the evaluator is not a single-score endpoint.
- results/results_table.csv and the results/ folder act as the persistent reporting layer for comparison across backbones and runs.
10. Baseline models
- src/web_agent/baselines/ contains standalone reference implementations:
b1_random.py, b2_mindact.py, b3_clip_mlp.py, b4_layoutlmv3.py.
- These are important in the architecture diagram as comparative branches, not part of the main training stack.
- They provide performance anchors against the shared four-head architecture.
11. Configuration and orchestration
- configs/base.yaml and configs/backbones/*.yaml define the model variants.
- The config layer is the main mechanism for selecting which encoder/backbone is plugged into the shared model core.
- scripts/run_model.py is the main orchestration entry point, with stage selection for smoke/mini/full/eval.
- scripts/run_gold.py appears to be a specialized path for the gold subset.
12. Runtime and environment assumptions
- The project is designed to run locally or on notebook/GPU environments, with the README calling out Kaggle T4 and Vast.ai A100 as compute targets.
- notebooks/kaggle_train.ipynb is intentionally thin: it installs the package and calls the same script entry point, so notebooks are wrappers rather than a separate code path.
13. Cross-cutting utilities
- src/web_agent/utils/ provides shared infrastructure: checkpointing, class weighting, logging, and seeding.
- These are supporting services used across training/eval rather than user-facing features.
- src/web_agent/config.py likely normalizes config loading and run parameters across scripts.
14. Testing and verification
- tests/test_shapes.py shows the repo uses at least structural tests for tensor shapes and likely finite-loss sanity.
- Combined with the smoke stage and dataset verification, the repository has a strong “verify early” architecture: data sanity, shape sanity, then scaled training.
15. Docs as architecture source of truth
- The docs/ and learn/ directories are unusually important here: they contain the project spec, architecture design, implementation plan, and phased training notes.
- For diagramming, treat these as the authoritative design layer that explains why the code is organized into fixed heads + swappable backbones + staged training.
Main entry points: run_model.py Run model · evaluate.py Evaluate · trainer.py Trainer · dataloader.py Batching
Components 28
Data
Dataset preparation, verification, loading, and gold subset handling.
- Setup data · Builds the local training corpus from raw data and images.
scripts/setup_data.py - Verify data · Checks dataset integrity and readiness before training.
src/web_agent/data/verify_dataset.py - Dataset · Main dataset format for step-level web-agent examples.
src/web_agent/data/dataset.py - Gold set · Higher-quality subset used for the gold training path.
src/web_agent/data/gold_dataset.py - Batching · Standard batching and mode filtering for training samples.
src/web_agent/data/dataloader.py - Gold loader · Loader tuned for the gold subset training path.
src/web_agent/data/gold_dataloader.py
Model
Backbone-agnostic multimodal model stack with fixed heads and shared loss.
- Text enc · Text-side encoder used in dual-encoder backbones.
src/web_agent/models/encoders/text.py - Vision enc · Image-side encoder used in dual-encoder backbones.
src/web_agent/models/encoders/vision.py - VLM enc · Unified multimodal backbone for image-plus-text inputs.
src/web_agent/models/encoders/vlm.py - Fusion · Combines modalities or backbone outputs before the fixed latent.
src/web_agent/models/fusion.py - Adapter · Projects backbone outputs into the shared 768-d representation.
src/web_agent/models/adapter.py - Task heads · Four fixed outputs for detect, diagnose, recover, and remember.
src/web_agent/models/heads.py - Multi-task loss · Shared weighted objective across all backbones and heads.
src/web_agent/models/loss.py - Model core · End-to-end network wiring encoders, fusion, adapter, heads, and loss.
src/web_agent/models/model.py - Memory index · External memory subsystem for the remember pillar.
src/web_agent/memory/index.py
Training
Staged training entry points and trainer loop.
- Smoke test · Tiny correctness run for early validation.
src/web_agent/train/smoke_test.py - Mini train · Small but real training slice before full runs.
src/web_agent/train/mini_train.py - Trainer · Shared training loop and checkpoint coordination.
src/web_agent/train/trainer.py - Full train · Main full-scale training entry point.
src/web_agent/train/train.py
Evaluation
Metrics and evaluation reporting.
Ops
Configuration, utilities, and orchestration.
- Run model · Main stage selector for smoke, mini, full, and eval flows.
scripts/run_model.py - Run gold · Specialized orchestration path for the gold subset.
scripts/run_gold.py - Config · Loads base and backbone-specific run configuration.
src/web_agent/config.py - Utilities · Checkpointing, class weights, logging, and seeding support.
src/web_agent/utils
Connections 35
- Setup data → Verify data: prepare corpus · Raw data is turned into a local dataset and then verified.
- Verify data → Dataset: validate · Verification produces a ready dataset schema.
- Dataset → Batching: batch · Standard samples are batched for training.
- Dataset → Gold set: subset · A higher-quality gold subset is derived from the same corpus.
- Gold set → Gold loader: batch · Gold samples use a dedicated loader path.
- Batching → Smoke test: feed · Standard batches drive smoke runs.
- Batching → Mini train: feed · Standard batches drive mini training.
- Batching → Trainer: feed · Standard batches feed the full trainer loop.
- Gold loader → Full train: feed · Gold runs can route through the full training entry.
- Text enc → Fusion: encode · Text features enter the multimodal fusion stage.
- Vision enc → Fusion: encode · Vision features enter the multimodal fusion stage.
- VLM enc → Adapter: encode · Unified VLM backbones produce outputs for projection.
- Fusion → Adapter: combine · Fused features are projected into the fixed latent size.
- Adapter → Task heads: 768-d latent · Shared embedding fans out to the four task heads.
- Task heads → Multi-task loss: predict · Head outputs are combined under a weighted objective.
- Memory index → Task heads: remember signal · External memory supports the remember-related output path.
- Model core → Multi-task loss: train step · The assembled model computes the shared multitask loss.
- Smoke test → Trainer: reuse loop · Smoke runs exercise the same trainer machinery.
- Mini train → Trainer: reuse loop · Mini training reuses the shared trainer machinery.
- Full train → Trainer: run loop · Full training is executed through the trainer.
- Run model → Smoke test: stage · The main runner can dispatch smoke runs.
- Run model → Mini train: stage · The main runner can dispatch mini runs.
- Run model → Full train: stage · The main runner can dispatch full training.
- Run model → Evaluate: stage · The main runner can dispatch evaluation.
- Run gold → Gold loader: gold path · Specialized orchestration targets the gold data path.
- Config → Run model: select backbone · Config chooses which backbone variant is assembled.
- Config → Model core: build model · Configuration determines the instantiated model stack.
- Utilities → Trainer: support · Shared utilities support training and checkpointing.
- Utilities → Evaluate: support · Shared utilities also support evaluation runs.
- Model core → Evaluate: predict · Trained models are evaluated on test splits.
- Evaluate → Metrics: score · Evaluation delegates scoring to the metrics layer.
- Evaluate → Results: report · Evaluation writes comparison artifacts to results.
- Baselines → Results: compare · Baseline runs provide performance anchors in reporting.
- Docs → Config: specify · Documentation describes the configuration and staged workflow.
- Docs → Model core: design · Docs define the fixed-head, swappable-backbone architecture.