TRAINING.md

Dita full fine-tuning and LoRA

The training preset follows the upstream LIBERO-Spatial recipe at RoboDita/Dita db6907d. LIBERO-10 uses a higher peak LR (5e-4). The upstream real-robot LoRA script uses a different scheduler; the LoRA branch here uses the LIBERO preset below. This aligns the requested optimizer settings, not the complete upstream data pipeline or benchmark results.

The five Hub models' runtime 2.2.0 configurations include the training fields and the saved image preprocessing device. LIBERO-Long saves peak LR 5e-4; the other models save 1e-4. Install the matching wheel from the model repository before loading these configurations.

Image preprocessing

policy.image_preprocessing_device selects the FP32 image resize and ImageNet normalization path for both training and inference:

Value	Execution device
cpu (default and published model setting)	CPU, preserving the reference preprocessing results
policy	The model's actual device, including the local CUDA device in distributed training

Pixel validation has no configuration switch: the policy performs one aminmax reduction on its model device and rejects non-finite or out-of-range pixels. Shape/channel checks use tensor metadata. Preprocessing remains FP32 under AMP.

To fine-tune the DROID initialization using CUDA preprocessing, add:

--policy.image_preprocessing_device=policy

The CLI overrides the initialization's saved cpu value. The new checkpoint saves policy; loading it for inference restores the same path automatically. Both full policies and LoRA checkpoints save this field through the Dita config. Keep that saved value when resuming or deploying. CUDA interpolation may differ slightly from the CPU reference, so changing this setting on an already trained model requires checking loss, gradients, and inference actions. The published models use cpu; this runtime update does not assert GPU-preprocessed rollout equivalence.

Implementation

1. freeze_backbone=false enables full fine-tuning of DINOv2, QFormer/FiLM, and the policy Transformer/action head. CLIP remains a frozen processor. Existing checkpoints retain an explicitly saved freeze setting.
2. gradient_checkpointing=true independently enables activation checkpointing in the ViT and QFormer. Both require training and enabled autograd before entering a checkpoint wrapper. Evaluation and no_grad/inference_mode skip the wrappers. Training uses use_reentrant=False, allowing gradients to reach LoRA adapters even when the base weights and input tensors are frozen.
3. LeRobot's native --peft.* branch wraps all Dita nn.Linear modules with LoRA, matching the upstream target selection. This includes the visual backbone, QFormer/FiLM, Transformer and action head. Only adapter weights are trained; freeze_backbone=true does not exclude backbone adapters.
4. The optimizer includes only trainable parameters, in the upstream group order: policy no-decay, policy decay, backbone no-decay, backbone decay. LoRA weights use the same grouping rules after adapter injection.
5. The registered dita scheduler is serialized by LeRobot so optimizer and scheduler state can be restored on resume.

Setting	Source default
freeze_backbone	false
gradient_checkpointing	true
optimizer_lr	0.0001
optimizer_backbone_lr_scale	0.1
optimizer_betas	[0.9, 0.95]
optimizer_eps	1e-8
optimizer_weight_decay	0.05
Bias and parameters with at most one dimension	Weight decay 0
optimizer_grad_clip_norm	0 (disabled by LeRobot)
scheduler_warmup_steps	1000
scheduler_decay_steps	100000
scheduler_min_lr_scale	0.01
lora_dropout	0.0

The schedule starts at zero LR, linearly warms up to the peak at step 1000, then follows the upstream cosine expression to 1% of the peak at step 100000. The backbone/policy LR ratio stays at 0.1 throughout. Short --steps runs do not compress this schedule. The upstream expression continues the cosine beyond the configured period; finish at scheduler_decay_steps or explicitly choose a longer period for longer runs.

Sources: parameter groups and scheduler, LIBERO commands, LoRA targets and initialization.

Full fine-tuning

Install the current source:

pip install -e ".[training]"

Use explicit overrides when starting from an older converted checkpoint. Its saved optimizer_lr, weight decay, clipping and freeze fields take precedence over new dataclass defaults. This starts a new training run; resuming an old run preserves that run's saved optimizer/scheduler state instead.

lerobot-train \
  --policy.path=albus2024/dita_droid_pretrain \
  --policy.device=cuda \
  --policy.dtype=bfloat16 \
  --policy.input_features=null \
  --policy.gripper_mode=binary \
  --policy.freeze_backbone=false \
  --policy.gradient_checkpointing=true \
  --policy.optimizer_lr=0.0001 \
  --policy.optimizer_backbone_lr_scale=0.1 \
  --policy.optimizer_weight_decay=0.05 \
  --policy.optimizer_grad_clip_norm=0 \
  --policy.scheduler_warmup_steps=1000 \
  --policy.scheduler_decay_steps=100000 \
  --policy.scheduler_min_lr_scale=0.01 \
  --policy.push_to_hub=false \
  --dataset.repo_id=YOUR_USERNAME/YOUR_LIBERO_SPATIAL_DATASET \
  --output_dir=outputs/train/dita_full \
  --batch_size=64 --steps=100000

The upstream LIBERO example uses eight GPUs, each with batch size 64. The command above is a single-process example; matching effective batch size is a separate training choice. Set --policy.optimizer_lr=0.0005 for the upstream LIBERO-10 LR. For continuous position grippers use --policy.gripper_mode=continuous and the appropriate target dataset.

LoRA fine-tuning

pip install -e ".[lora]"

Add the following to the full fine-tuning command and choose a fresh output directory, for example outputs/train/dita_lora:

--peft.method_type=LORA \
--peft.r=32

The default alpha is automatically 2 * r (64 for rank 32), initialization is Gaussian for A and zero for B, and dropout is zero. Override with --peft.lora_alpha, --peft.init_type, --peft.target_modules or --policy.lora_dropout. LeRobot's native CLI rank default is 16, so specify --peft.r=32 explicitly for the example above. Do not set policy.use_peft=true when starting from a base checkpoint: that flag identifies an already saved adapter and is set by LeRobot when it wraps the policy.

Adapter checkpoints contain adapter_model.safetensors, adapter_config.json, the Dita config.json, processors, and training state. The adapter config records the base model location; that base model must remain available. Adapters do not contain a standalone copy of the full Dita weights.

Resume and inference

Resume either training branch from its saved training configuration:

lerobot-train \
  --resume=true \
  --config_path=outputs/train/dita_lora/checkpoints/last/pretrained_model/train_config.json

LeRobot's native policy factory loads full policies and saved adapters. For LIBERO evaluation, use the saved pretrained_model directory as --policy.path in the normal lerobot-eval command; the saved use_peft field selects adapter loading automatically. The same applies to LeRobot's recording and deployment entry points that use make_policy.

For Python, use make_policy to load adapters, followed by the saved processors:

from pathlib import Path
import torch
from lerobot.configs.policies import PreTrainedConfig
from lerobot.policies.factory import make_policy, make_pre_post_processors
import lerobot_policy_dita  # Register the Dita config, scheduler and processors.

path = Path("outputs/train/dita_lora/checkpoints/last/pretrained_model")
config = PreTrainedConfig.from_pretrained(path)
config.pretrained_path = path
config.device = "cuda"
policy = make_policy(config, ds_meta=dataset.meta).eval()
pre, post = make_pre_post_processors(config, pretrained_path=path)
with torch.inference_mode():
    action = post(policy.select_action(pre(observation_batch)))

Here dataset.meta is the target LeRobot dataset metadata and observation_batch contains the image(s) and task text. Dita's direct DitaPolicy.from_pretrained loads full weights; use the factory above for adapter checkpoints. In both branches the inference entry point automatically skips activation checkpoint wrappers, even if checkpointing was enabled during training and the saved backbone freeze setting is false.

Validation

OMP_NUM_THREADS=1 pytest -q tests/test_dita_training.py
pytest -q

The training tests use the actual Dita modules at reduced width/depth. They check full/frozen training with checkpointing on/off, gradient parity against direct forward execution, inference without checkpoint calls, optimizer groups, schedule/resume behavior, and LoRA updates plus adapter reload through LeRobot.

For a short full-size check, use the real three-episode LIBERO subset described in the training smoke test, apply the optimizer overrides above, and run each branch with --steps=5 --batch_size=2 --num_workers=0 --save_freq=5 --env_eval_freq=0. Resume to step 6, then load the saved policy/adapter and run inference. Keeping the original warmup means the first optimizer call has zero LR; later calls must update the intended weights.

Results on 2026-09-14

Environment: LeRobot 0.6.1, PyTorch 2.9.1+cu128, PEFT 0.19.1, Transformers 5.5.4, timm 1.0.24; RTX 4090 D. The repository suite passed all 321 tests. Native training used the DROID initialization, the same three real LIBERO episodes (843 frames), BF16, batch size 2 and the unmodified 1000-step warmup.

Check	Full fine-tuning	LoRA, rank 32
Trainable parameters	213,225,984	13,000,928
Native train/save/resume	5 steps + resumed step 6	5 steps + resumed step 6
Intended updates	ViT, QFormer, Transformer	Adapters in all three components
Base parameter tensors changed	372	0
Optimizer state	Finite; step 6	Finite; step 6
Restored policy/backbone LR	6e-7 / 6e-8	6e-7 / 6e-8
Saved model inference	BF16, FP32, FP16; finite [1, 7] actions	Same
Checkpoint calls in inference	0	0
Checkpoint flags enabled vs disabled, same RNG	Exactly equal actions	Exactly equal actions

The local validation artifacts are in workspace/validation/dita-training-20260914/: native training logs, both checkpoints, verify.py, pytest.log, and validation_report.json. They are ignored by Git. These short checks do not establish convergence or robot task success.

Runtime 2.1.0 release validation

The bundled wheel separately passed all 321 tests and strict loading of all five Hub models. BF16 inference on the checked observation matched runtime 2.0.1 exactly, with zero activation-checkpoint calls. LIBERO processors restored finite actions; the DROID initialization retained its unbound processor templates.

Using the candidate DROID configuration without optimizer overrides, native BF16 full fine-tuning and rank-32 LoRA each completed three steps and resumed for a fourth on the same three real LIBERO episodes. Both restored the scheduler and finite optimizer state, with policy/backbone LR 4e-7 / 4e-8 at step 4. The intended parameters updated in the ViT, QFormer and Transformer in both branches.

Artifacts, hash-locked publication manifests, upload revisions and downloaded Hub verification are under workspace/validation/dita-runtime-2.1.0/. Each model repository also contains runtime_validation_report.json and archives the previous configuration and runtime metadata under history/runtime-2.0.1/.

Runtime 2.2.0 image preprocessing validation

All five full-size checkpoints were checked on real observations, with batch size 2, FP32/BF16/FP16, fixed RNG, and deterministic kernels. LIBERO checks used two samples from the cached three-episode subset; DROID used two SO101 samples with a 100-action horizon. Comparing the new cpu mode against the previous CPU image path produced exactly equal loss, all 372 parameter gradients, and inference actions in every case. CPU and policy modes both produced finite results.

Switching existing weights to policy changes interpolation numerics. Across the five checkpoints, the largest measured differences were:

Compute dtype	Loss absolute difference	Gradient relative L2 difference	Normalized action absolute difference
FP32	2.87e-6	9.35e-5	6.62e-6
BF16	1.98e-3	1.59e-2	1.70e-2
FP16	3.65e-4	5.43e-3	8.22e-4

These are maxima over a small paired numerical check, not accuracy bounds or a LIBERO success-rate evaluation. No binary gripper decisions changed in those samples. The published configs retain cpu; new fine-tuning can explicitly select policy and carry that choice into deployment. A DROID initialization overridden from cpu to policy completed two real-data updates, saved the override, and reproduced its inference actions exactly after reloading.

Validation artifacts are under workspace/validation/dita-preprocessing-20260916/. The numerical comparison uses deterministic kernels to distinguish preprocessing effects from mixed-precision backward-kernel nondeterminism. Standard training does not need to enable deterministic algorithms to use this configuration.