Back to Integrations
/ INTEGRATION · Experiment tracking

Weights & Biases

Tracking robot-learning runs in Weights & Biases — versioning datasets as artifacts, logging rollout videos, and why validation loss alone falls short.

Updated Aug 20267 min read
SHORT ANSWER

Weights & Biases tracks robot-learning runs the same way it tracks any training job, but robot data adds two things generic ML workflows skip: versioning the dataset itself as an artifact so a checkpoint traces back to the exact episodes that trained it, and logging closed-loop rollout success alongside loss, since behavior-cloning loss is a weak proxy for policy performance.

Weights & Biases tracks a training run the way it tracks any deep learning job — config, metrics, checkpoints, system stats — but robot learning stresses two parts of it that most computer vision or NLP workflows barely touch: the dataset itself needs a version identity, because a policy's behavior is a direct function of which episodes trained it, and loss alone does not tell you whether the robot will actually complete the task. Both are addressable with existing W&B primitives; neither happens automatically. The training loop itself is usually a plain PyTorch job — W&B doesn't replace DataLoader or the training step, it instruments the run around it.

How the data gets there

A run starts with wandb.init(), and the dataset enters the picture as a versioned artifact rather than an opaque file path in a config dict:

import wandb
 
run = wandb.init(project="pick-and-place", job_type="train")
 
# Reference the exact dataset version this run consumes.
# If "robot-demos:latest" resolves to a new version between runs,
# W&B still records precisely which version this run used.
dataset_artifact = run.use_artifact("robot-demos:latest", type="dataset")
dataset_dir = dataset_artifact.download()
 
run.config.update({
    "dataset_version": dataset_artifact.version,
    "num_episodes": dataset_artifact.metadata.get("num_episodes"),
    "chunk_len": 16,
    "lr": 1e-4,
})
 
for step, batch in enumerate(train_loader):
    loss = train_step(batch)
    wandb.log({"train/loss": loss}, step=step)
 
# Log a handful of evaluation rollouts as video, not just a scalar
wandb.log({
    "eval/success_rate": success_rate,
    "eval/rollout": wandb.Video(rollout_frames, fps=15, format="mp4"),
})

On the data-producing side — after a collection or curation batch — the dataset itself gets logged as an artifact once, independent of any training run:

import wandb
 
with wandb.init(project="pick-and-place", job_type="dataset-curation") as run:
    artifact = wandb.Artifact(
        name="robot-demos",
        type="dataset",
        metadata={"num_episodes": 4213, "task": "pick-and-place", "embodiment": "franka-fr3"},
    )
    artifact.add_dir("/data/pick_and_place/v7")   # or artifact.add_reference("s3://…")
    run.log_artifact(artifact)

For large, video-heavy datasets, add_reference() against an S3 or GCS URI tracks file metadata and checksums without copying the underlying files into W&B-managed storage — useful when the dataset already lives in object storage and re-uploading terabytes of video is not worth doing per version.

Every subsequent training run that calls use_artifact("robot-demos:latest") — or pins an explicit version, robot-demos:v7 — shows up in the artifact's lineage graph. Given a checkpoint, the question "which episodes produced this" has a direct answer instead of requiring someone to reconstruct it from a changelog.

The workflow in practice

  • Version the dataset at the point it's finalized, not per training run. A curation or collection batch gets one artifact version; every run that trains on it references that version, so ten runs against the same data don't produce ten redundant uploads.
  • Log real-robot evaluation as its own tracked signal, alongside — not instead of — training loss. A wandb.Table with one row per evaluation episode (task, success/fail, failure mode, video link) turns "did it work" into something searchable and comparable across checkpoints, rather than a note in someone's head.
  • Tag runs by dataset lineage, not just hyperparameters, so filtering the run table by "everything trained on robot-demos:v7" is a saved view rather than a manual search through configs.

Sweeps over data mixtures, and team reporting

A wandb sweep config is just a set of parameters your training script reads at startup — nothing restricts those parameters to model hyperparameters. A robot-learning sweep gets more value out of varying the ratio of teleoperated to synthetically augmented demonstrations, or the per-task sampling weight in a multi-task mixture, than out of another pass over learning rate:

# sweep.yaml
method: bayes
metric:
  name: eval/success_rate
  goal: maximize
parameters:
  teleop_fraction:
    min: 0.2
    max: 1.0
  mimic_fraction:
    min: 0.0
    max: 0.8
  lr:
    values: [1e-4, 3e-4, 1e-3]
wandb sweep sweep.yaml          # prints a sweep ID
wandb agent <entity>/<project>/<sweep_id>   # run on each machine you have available

Because the sweep optimizes eval/success_rate rather than validation loss, the winning configuration is whichever data mixture and hyperparameter combination actually improved closed-loop performance — not whichever one best matched demonstrated actions on paper. Pulling the resulting loss curves, success rates, and a grid of rollout videos into a shared W&B Report is how most teams run the recurring "where are we" review without re-deriving the answer from raw run data each time.

Logging from a rig with no reliable network

A training cluster has a stable connection to W&B's servers; a robot on a lab bench, a warehouse floor, or a moving base often doesn't. Rather than losing run history to a dropped connection, wandb.init(mode="offline") (or setting the WANDB_MODE=offline environment variable before the process starts) writes every wandb.log() call and every artifact to a local run directory instead of streaming it out immediately:

# On the robot, no network assumed
WANDB_MODE=offline python evaluate_policy.py --checkpoint outputs/train/pick_place_act/last.ckpt
 
# Later, once the machine is back on a network — or from a laptop after
# copying the run directory off the robot
wandb sync ./wandb/offline-run-20260805_141203-abc123

The run behaves identically from the training script's point of view — wandb.log(), wandb.Video, and artifact logging all work the same way — the only difference is when the data actually reaches W&B's servers. This matters specifically for real-robot evaluation rollouts: a team running eval sessions on a cellular-connected mobile base, or in a facility that blocks outbound traffic from robot-network VLANs by policy, can still log every rollout locally and reconcile the run history in one batch sync afterward, rather than either losing the data or routing evaluation through a live network dependency that has nothing to do with the robot completing its task.

Dataset lineage as an input to the flywheel, not just a paper trail

The lineage graph use_artifact/log_artifact builds isn't only for after-the-fact reproducibility — it's the record a team uses to decide what to collect next. If a robot-demos:v7 → checkpoint → low eval/success_rate chain shows up repeatedly for one task, and a robot-demos:v9 (a version that added targeted demonstrations for that task's failure mode) shows a checkpoint clearing the same eval suite, that's the data flywheel made visible in the run table instead of asserted anecdotally. Tagging runs with the dataset version and the failure modes a wandb.Table surfaced turns "we think more data on X helped" into a comparison you can pull up and show someone.

That same lineage view is where a synthetic-versus-real data mixture sweep earns its keep: logging teleop_fraction and mimic_fraction as run config alongside the artifact version each run consumed means a later reviewer can filter "runs where synthetic augmentation exceeded 50% of the mixture" and read eval/success_rate directly off the table, rather than reconstructing the mixture ratio from a training script that may have since changed. Feeding flagged episodes from a dataset QA pass back into the next artifact version — rather than leaving bad episodes silently mixed into latest — is what keeps that lineage graph trustworthy instead of just detailed.

Gotchas

Reference artifacts still need a retention policy. add_reference() avoids duplicating file bytes into W&B storage, but W&B still tracks metadata for every referenced file — a dataset with millions of small files (individual frame images rather than video) can produce artifact manifests large enough to matter. Prefer video-encoded or chunked storage over one-file-per-frame when artifact-tracking a large dataset.

wandb.log() steps need to stay monotonic per run. Logging out of step order, or from multiple processes writing to the same run without coordinating step numbers, produces metric charts with gaps or overwritten points — a common outcome of naively adding W&B logging to a multi-GPU training loop without gating logging calls to rank 0.

Artifact versions are immutable once logged. log_artifact() on the same name creates a new version rather than mutating the old one — which is the point, for reproducibility — but it means re-running a curation script without changing the name produces v8, v9, and so on, and stale versions accumulate unless someone prunes them deliberately.

Sweeps parallelize by agent, not automatically by data mixture. Running a data-mixture sweep at any real scale means launching multiple wandb agent processes against the sweep ID yourself — on however many machines or jobs you have available — since the sweep controller assigns configurations but does not provision compute.

RUN INIT
wandb.init(project=..., job_type=...)
DATASET ARTIFACT
wandb.Artifact(type="dataset")
EXTERNAL STORAGE
artifact.add_reference("s3://...")
SWEEP METHODS
grid, random, bayes

KEY FACTS

CORE CALLS
wandb.init(), wandb.log(), wandb.Artifact
DATASET VERSIONING
wandb.Artifact(type="dataset"), lineage via use_artifact / log_artifact
MEDIA LOGGING
wandb.Video for rollouts, wandb.Table for per-task breakdowns
SWEEPS
wandb sweep + wandb agent — grid, random, or Bayesian search

/ QUESTIONS

Frequently asked

Put this into practice.

Tell us what your robots need to learn. We will scope the rig, the operators, the protocol, and the first datasets — usually in one call.