---
license: mit
tags:
  - usam
  - vla
  - robot-learning
  - qwen3-vl
  - lda
  - ablation
base_model:
  - Qwen/Qwen3-VL-4B-Instruct
  - facebook/dinov3-vitb16
---

# USAM — usam-full-loss

Canonical USAM loss recipe: action + aux RGB + aux depth + cosine drift + subtask completion. Baseline.

USAM Stage B1-real pretrain checkpoints, variant **`usam-full-loss`**. Architecture: Qwen3-VL-4B-Instruct backbone (LoRA r=16) + DINOv3 visual encoder + LDA flow-matching action head + USAM auxiliary heads (drift / subtask / depth-RGB geom).

## Provenance

- **SLURM job**: `60949`
- **Run directory**: `runs/usam_real_qwen_b1-60949/`
- **Source config**: `configs/train/stage_b1_real_qwen_pretrain.yaml` (mirrored as `config.yaml` in this repo)

## Loss configuration

```yaml
loss_weights:
  action: 1.0
  rgb: 1.0
  depth: 0.3
  drift: 0.1
  subtask: 0.1
  geom_target: 0.0
  ramp_steps: 50000           # matches max_steps
```

## Saved checkpoints

14 checkpoints, every 2500 steps:

- `checkpoint_step00002500.pt`
- `checkpoint_step00005000.pt`
- `checkpoint_step00007500.pt`
- `checkpoint_step00010000.pt`
- `checkpoint_step00012500.pt`
- `checkpoint_step00015000.pt`
- `checkpoint_step00017500.pt`
- `checkpoint_step00020000.pt`
- `checkpoint_step00022500.pt`
- `checkpoint_step00025000.pt`
- `checkpoint_step00027500.pt`
- `checkpoint_step00030000.pt`
- `checkpoint_step00032500.pt`
- `checkpoint_step00035000.pt`

- **Latest step**: `35000`
- **Checkpoint format**: `trainable+buffers` — state_dict (LoRA + trainable adapters + buffers) plus full AdamW optimizer state, scheduler state, and run metadata. Every saved step is independently resumable for continued training.

## Ablation context

This repo is one variant in the USAM B1-real loss-ablation matrix. The seven variants isolate single-loss contributions:

| Variant | Repo | What it tests |
|---|---|---|
| Full + geom | `usam-full-loss-geom` | Upper bound: does depth-RGB geometric consistency help on top of the baseline? |
| Full (baseline) | `usam-full-loss` | Canonical recipe — `action + rgb + depth + drift + subtask` |
| Action only | `usam-action-only` | Lower bound: pure VLA action loss |
| No aux vision | `usam-no-aux-vision` | Does aux RGB + depth co-training help? |
| No USAM aux | `usam-no-usam-aux` | Do drift + subtask add lift beyond LDA-style co-training? |
| Drift only | `usam-drift-only` | Marginal contribution of drift alone |
| 3-source (DROID) | `usam-full` | Canonical recipe + DROID dataset (3-source full data mix) |

See `docs/ABLATION_STUDY.md` in the source repo for the full design.

## Usage

```python
import torch

ckpt = torch.load(
    "checkpoint_step00035000.pt",
    weights_only=False,
    map_location="cpu",
)
state_dict = ckpt["state_dict"]    # trainable + buffers only
step       = ckpt["step"]          # int
opt_state  = ckpt["optimizer"]     # for resume
sched      = ckpt["scheduler"]     # for resume

# Load into a freshly-constructed USAM model:
missing, unexpected = model.load_state_dict(state_dict, strict=False)
# `missing` will contain the frozen base-model keys (Qwen3-VL + DINOv3),
# which load from their respective HF base checkpoints. See
# usam/_train_helpers.py:2437-2453 for the reference loader.
```