# Dita full fine-tuning and LoRA

The training preset follows the upstream **LIBERO-Spatial** recipe at
[RoboDita/Dita `db6907d`](https://github.com/RoboDita/Dita/tree/db6907d8c0b9878570570e10d3a0599420072044).
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:

```bash
--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](https://github.com/RoboDita/Dita/blob/db6907d8c0b9878570570e10d3a0599420072044/scripts/train_diffusion_oxe.py#L69),
[LIBERO commands](https://github.com/RoboDita/Dita/blob/db6907d8c0b9878570570e10d3a0599420072044/README.md#L98),
[LoRA targets and initialization](https://github.com/RoboDita/Dita/blob/db6907d8c0b9878570570e10d3a0599420072044/scripts/finetune_realdata.py#L242).

## Full fine-tuning

Install the current source:

```bash
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.

```bash
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

```bash
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`:

```bash
--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:

```bash
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:

```python
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

```bash
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](../../scripts/README.md#native-lerobot-train-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.