---

name: model-training description: "Configure, smoke-check, train, resume, and troubleshoot Diffusion Planner models with the repository's PyTorch/DDP pipeline." disable-model-invocation: true metadata: disco-role: operating license: NO_LICENSE

Diffusion Planner model training

Use this sub-skill for model-ready data validation, training configuration, CUDA/DDP smoke checks, checkpoint resume, and training failure diagnosis. Follow the staged route; consult the detailed contracts in the reference links when a stage requires them.

Scope and routing

This skill covers the repository's training flag contract, model input/output contracts, normalization, distributed launch, EMA, logging, checkpoints, and resume. The long-running trainer is intentionally reference-only because it writes checkpoints, requires a generated corpus, and is not a safe portable bundled helper; use the bundled checker before a separately managed training entrypoint. This skill stops before nuPlan scenario preprocessing, closed-loop simulation, or custom guidance authoring.

• For raw nuPlan data and model-ready .npz creation, use data-preparation.
• For planner configuration, checkpoint consumption, and simulation, use closed-loop-planning.
• For guidance_fn, collision guidance, or guided sampling, use guidance. Training consumes model-ready .npz records plus a JSON filename list. It produces a run directory containing args.json, checkpoints, and logs; the matching args.json is part of the checkpoint contract for downstream use.

Preconditions and safety boundaries

Before training or resume, confirm:

1. The active Python environment imports the project and required dependencies; use its interpreter, not a private interpreter path.
2. The data directory, relative-filename JSON list, normalization file, and planned writable save directory are valid; choose deliberate resume/overwrite.
3. CUDA/NCCL and intended devices are available for full training. CPU is only for parser, manifest, normalization, and limited API checks.
4. Inputs and small fixtures are local; this skill does not download data or checkpoints.

The checked-in launcher is a template, not a safe command. Do not copy its placeholder values, sudo, or private interpreter invocation. Use the active environment with python -m torch.distributed.run and explicit device selection instead.

Staged decision workflow

0. Classify the request

• Parser/contract smoke: stay local, bounded, and side-effect free.
• One-process API/data smoke: use a tiny fixture, --ddp false, zero workers, and disabled augmentation if the import path requires it.
• Training: proceed only after data, normalization, device, batch, launch, and output checks pass.
• Resume: inspect args.json and latest.pth as a pair first.
• Route preprocessing, simulation, and guidance to the siblings above.

1. Run safe preflight checks

From the project execution root, run the native parser check and the bundled training-contract checker:

python scripts/check_training_contract.py --help
python scripts/check_training_contract.py --check-normalization normalization.json \
  --predicted-neighbor-num 10

The full trainer is a long-running, side-effectful source workflow rather than an included runtime helper. Before invoking a separately managed trainer, apply the flag and tensor contracts in the references and complete these bounded checks.

For a manifest check, add --check-manifest <data-dir> --data-list <json-file> and a small --limit. The checker is no-training, no-download, and no-checkpoint; a parser pass is not evidence that training or CUDA is healthy.

2. Validate the model/data contract

Confirm the manifest resolves every filename under the chosen data directory and that records contain the required ego, neighbor, lane, route, speed-limit, and static-object fields. Check the configured axes for future/time length, agent and neighbor counts, lane/route counts and lengths, and static-object count; the detailed shapes and keys are in the API reference.

Use one normalization layout consistently. normalization.json supplies future-state statistics for ego/neighbors and observation statistics for other inputs. An absent key, wrong vector length, non-positive standard deviation, non-finite value, or mismatched layout is a preflight failure, not a learning-rate problem.

3. Reconcile configuration

Compare CLI flags, data axes, normalization, and checkpoint args.json. Keep hidden_dim divisible by num_heads; dataset/model flags, not descriptive ones alone, determine reshaping. For a one-process smoke, use a matching architecture and stop after an API/forward assertion unless training was asked.

The important defaults are future length 80, history length 21, 32 agent rows, 10 predicted neighbors, 70 lanes, 25 route lanes, 5 static objects, global batch 2048, CUDA, DDP, and EMA enabled. Treat defaults as a starting point; verify them against the actual data and checkpoint rather than relying on them.

4. Choose a launch safely

For one process, set --ddp false, use one visible GPU when available, and make the requested batch the local batch. For DDP, use a real launcher and check that:

• global batch_size % WORLD_SIZE == 0 and each rank gets a non-empty batch;
• --nproc-per-node equals the number of visible devices after remapping;
• all ranks use identical arguments, readable inputs, and a writable output;
• the port is free and no stale rank variables or failed worker group remain.

The script does not validate batch divisibility. A DDP flag without launcher environment is not a reliable fallback; choose a real launcher or explicitly use --ddp false. See workflows for launch shapes.

5. Observe the run and preserve evidence

Record the sanitized command, device mapping, world size/local batch, normalization source, one manifest entry, output directory, and first traceback. Rank 0 should write args.json, checkpoints, TensorBoard output, and optional W&B files; distinguish logger/save-path failure from model or data failure.

6. Resume deliberately

Pass --resume_model_path as the directory containing latest.pth, not the file. Compare its sibling args.json with the requested architecture before loading. The loader accepts a wrapped checkpoint or bare state dict, and optimizer, scheduler, epoch, W&B, and EMA fields may be absent; report a warm start when those states were not restored. Do not use broad strict=False or silently continue through architecture mismatches.

7. Recover and hand off

Triage the first failing contract using the troubleshooting matrix: manifest/data, normalization, one-GPU API, CUDA visibility, DDP/port, save or logging, then checkpoint compatibility. Stop after failed preflight; send data defects to data-preparation and checkpoint/simulation defects to closed-loop-planning.

Verified warnings

• Full training is CUDA/NCCL-oriented; CPU checks cannot validate distributed initialization, GPU memory, or a useful epoch.
• Global batch is integer-divided across ranks without a divisibility check.
• --use_ema false is not flag-only here: the entrypoint still updates EMA; guard or initialize that path in a reviewed adaptation.
• Observed augmentation application is near 1 - augment_prob; verify before reproducibility claims. The “Cosine” scheduler warms up, then uses a fixed multiplicative schedule rather than cosine decay.
• Missing ema_state_dict prevents reproducing EMA evaluation.

Detailed contracts and graph links

• API/config/tensor/checkpoint contract
• Staged commands and launch workflows
• Failure recovery and non-goals
• Bounded contract-check script
• Diffusion Planner root skill
• Data preparation sibling
• Closed-loop planning sibling
• Guidance sibling

Keep source-backed detail in the references; keep this router focused on applicability, order, stop conditions, and handoffs.