training.md

Training

RLDX-1 has a single unified launcher — rldx/experiment/launch_train.py — that covers fine-tuning and mid-training. The difference between the two is just which dataset-selection arguments and which base checkpoint you pass.

This doc covers the CLI surface, the three common training modes, and where to look for canonical example scripts.

Entry point

uv run torchrun --nproc_per_node=<N_GPUS> --master_port=<PORT> \
    rldx/experiment/launch_train.py \
        [ CLI args here ]

launch_train.py parses a single TrainConfig dataclass (rldx/configs/train_config.py) via tyro — every field on that dataclass is exposed as a CLI flag with the same name, with underscores turned into dashes (e.g. global_batch_size → --global-batch-size).

The launcher internally reuses run() from rldx/experiment/experiment.py, which handles dataset construction, model setup, DeepSpeed initialisation, and the HuggingFace Trainer main loop.

Three training modes

(a) Fine-tune from a base checkpoint

The common case: take a pre-trained RLDX-1 checkpoint and specialise it on a single dataset.

uv run torchrun --nproc_per_node=8 rldx/experiment/launch_train.py \
    --base-model-path RLWRLD/RLDX-1-PT \
    --dataset-path /path/to/robocasa_dataset \
    --embodiment-tag GENERAL_EMBODIMENT \
    --modality-config-path rldx/configs/data/robocasa_config.py \
    --video-length 4 \
    --n-cog-tokens 64 \
    --global-batch-size 64 \
    --learning-rate 1e-4 \
    --max-steps 60000 \
    --save-steps 1000 \
    --save-total-limit 5 \
    --output-dir ./ckpt/rldx/finetuned/my_run \
    --experiment-name my_run \
    --use-wandb --wandb-project rldx-finetune

Choosing the base checkpoint. Match --video-length to the base: RLWRLD/RLDX-1-PT (video input) uses --video-length 4, while RLWRLD/RLDX-1-PT-IMG (image input) uses --video-length 1.

Key fine-tune flags:

Flag	Purpose
--base-model-path	HF hub id or local path to the pre-trained checkpoint. Processor + config + weights are all loaded from it.
--dataset-path	LeRobot-format dataset directory (single-dataset fine-tune).
--dataset-paths + --dataset-mix-ratios	Multi-dataset fine-tune with per-dataset mix ratios. Mutually exclusive with --dataset-path.
--embodiment-tag	Which embodiment slot to train. GENERAL_EMBODIMENT routes through the category-specific MLP's "new" head.
--modality-config-path	Python file declaring the modality (video / state / action) config for this embodiment. See rldx/configs/data/ for the built-ins.

Canonical script: run_scripts/train/finetune_rldx_robocasa.sh.

(b) Mid-train with optional add-ons (memory / motion / physics)

Mid-training resumes from a pre-trained checkpoint and additionally turns on one or more add-on components for supervised continuation. All three add-ons are additive flags on top of the base fine-tune command lines:

# Memory
--use-memory \
--memory-length 4 \
--memory-n-cog-tokens 16 \
--concat-memory \
--memory-dropout-prob 0.3

# Motion module
--use-motion \
--motion-insert-layer 9 \
--motion-injection-point vision_encoder \
--motion-pool-type avg

# Physics (tactile + torque)
--use-physics \
--physics-keys tactile torque \
--physics-dims 30 7 \
--physics-loss-weight 0.1

All three can be combined in a single run. Canonical scripts: run_scripts/train/ablations/.

Dataset layout: meta/modality.json

Beyond the LeRobot-format parquet/video files, every dataset RLDX-1 trains on must carry a meta/modality.json describing how to slice the flat state/action vectors into named joint groups and how to remap video columns to modality keys. This file is loaded by rldx/data/dataset/lerobot_episode_loader.py on every dataset open; without it the loader falls back to a single-key default that almost certainly does not match what the modality config you train with expects.

my_dataset/
├── data/
│   └── chunk-000/
│       └── episode_000000.parquet
├── videos/
│   └── chunk-000/
│       └── observation.images.<cam>/episode_000000.mp4
└── meta/
    ├── episodes.jsonl
    ├── episodes_stats.jsonl
    ├── info.json
    ├── modality.json        ← required
    ├── stats.json
    └── tasks.jsonl

Schema ({modality}.{key}.start/end for state and action; video uses original_key to remap a parquet column to a logical name):

{
    "state": {
        "right_arm_joints": { "start": 0,  "end": 7 },
        "left_arm_joints":  { "start": 7,  "end": 14 },
        "neck_joints":      { "start": 44, "end": 46 }
    },
    "action": {
        "right_arm_joints": { "start": 0,  "end": 7 },
        "left_arm_joints":  { "start": 7,  "end": 14 },
        "neck_joints":      { "start": 44, "end": 46 }
    },
    "video": {
        "camera_ego_left":  { "original_key": "observation.images.camera_ego_left" },
        "camera_ego_right": { "original_key": "observation.images.camera_ego_right" }
    }
}

The slice ranges in state and action must match the joint groups declared in the Python MODALITY_CONFIGS entry you pass via --modality-config-path. The dataset's flat observation.state / action parquet columns are sliced according to these ranges and re-keyed under the joint-group name before reaching the model.

For the embodiment slot itself (which selects which Python MODALITY_CONFIGS entry is active), see embodiment_tags.md.

Image pipeline flags

See architecture.md for the geometry. The three knobs are:

--image-max-area 65536       # area budget, default 256*256
--image-resize-m 32          # alignment multiple
--random-crop-fraction 0.9   # None (default) = no-op
--random-rotation-angle 5    # optional train-only rotate
--color-jitter-params brightness 0.3 contrast 0.4 saturation 0.5 hue 0.08  # this is the default

--random-crop-fraction=None is the production default — the refactor confirmed that the current datasets are 256×256 and the crop is skipped. Set it only when you want extra augmentation.

Training knobs cheat-sheet

Flag	Default	Notes
--num-gpus	1	Also needs matching torchrun --nproc_per_node.
--global-batch-size	64	Divided by num_gpus * gradient_accumulation_steps.
--gradient-accumulation-steps	1
--learning-rate	1e-4	Warmup via --warmup-ratio (default 0.05).
--lr-scheduler-type	cosine
--weight-decay	1e-5
--max-steps	30000	Total training steps.
--save-steps	1000	Checkpoint cadence.
--save-total-limit	5	Rolling window of saved checkpoints.
--dataloader-num-workers	2	Workers per rank.
--use-wandb + --wandb-project	off	WANDB_PROJECT env var is also honored.
--experiment-name	"debug"	Run name on wandb and for the output directory.

Tuning which parameters are trained

Defaults reflect the refactor decision: backbone frozen below, action head fully tuned.

Flag	Default	What it controls
--tune-llm	False	Train the entire backbone VLM.
--tune-visual	False	Train the vision tower.
--tune-top-llm-layers	4	When --tune-llm is off, train just the top N LM layers.
--tune-projector	True	Multi-modal projector.
--tune-diffusion-model	True	MSAT action model.
--freeze-cog-tokens	False	Freeze the learnable cognition-token embeddings.
--state-dropout-prob	0.3	Stochastic state dropout for regularisation.
--general-embodiment-train-ratio	0	Mix-in ratio of general-embodiment (cross-embodiment) samples when fine-tuning on a single embodiment. 0 disables the mix.

LoRA fine-tuning

For memory-constrained fine-tunes you can replace full-parameter training of either the action model (MSAT) or the backbone VLM with PEFT LoRA adapters. Both surfaces are independent; you can enable one or both.

# Action-model LoRA — replaces full MSAT tuning
--action-model-use-lora \
--action-model-lora-rank 16 \
--action-model-lora-alpha 32 \
--action-model-lora-dropout 0.0

# Backbone LoRA — replaces top-N backbone VLM layer tuning
--backbone-use-lora \
--backbone-lora-rank 16 \
--backbone-lora-alpha 32 \
--backbone-lora-num-layers -1   # -1 = all layers, N>0 = top-N suffix, 0 = skip

Flag	Default	What it controls
--action-model-use-lora	False	Inject LoRA into MSAT QKV / projection / FFN linears. Overrides --tune-diffusion-model — when on, only the LoRA adapters are trainable.
--action-model-lora-rank	16	Adapter rank r.
--action-model-lora-alpha	32	Scaling factor α.
--action-model-lora-dropout	0.0	LoRA dropout.
--action-model-lora-target-modules	(see rldx/configs/model/rldx.py:136)	MSAT linear submodules to wrap. Targets that don't exist in the current MSAT (e.g. p_qkv when use_physics=False) are filtered before the PEFT call.
--backbone-use-lora	False	Inject LoRA into the backbone VLM's attention + MLP projections. Overrides --tune-top-llm-layers — when on, the backbone is frozen and only LoRA adapters train.
--backbone-lora-rank	16	Adapter rank.
--backbone-lora-alpha	32	Scaling factor.
--backbone-lora-num-layers	-1	-1 = inject into all backbone VLM layers; N > 0 = inject into the top N layers only; 0 = skip backbone LoRA.
--backbone-lora-dropout	0.0	LoRA dropout.

The launcher warns when a LoRA flag overrides one of the full-tuning flags above, so a stale config file does not silently produce a no-op run.

Real-Time Chunking (training-time)

Training-time RTC (the t_tok=1 trick from Black et al. (Training-Time Action Conditioning for Efficient Real-Time Chunking)) bakes clean-prefix conditioning directly into the model. At inference time the prefix is just inpainted at τ=1 instead of running an extra guidance pass, so the chunk-boundary stitching incurs no latency overhead.

--rtc-training-max-delay 4

When --rtc-training-max-delay > 0 each training step samples a per-sample prefix length d ~ U{0, …, max_delay}; those positions are conditioned on ground-truth clean actions and are excluded from the loss. A checkpoint trained this way enables --rtc-inference-mode trained on the server and is compatible with --compile fullgraph. A checkpoint trained with --rtc-training-max-delay 0 (the default) can still serve RTC by using --rtc-inference-mode guided at inference time, but the guided path is slower and incompatible with fullgraph compilation.

See inference_server.md for the corresponding inference-time flags and the mode/compile compatibility matrix.

Checkpoint format

After every save_steps, the CheckpointFormatCallback writes the trainer output plus:

{output_dir}/checkpoint-{step}/
├── config.json                   # model_type="RLDX-1", see rldx/configs/model/rldx.py
├── model-00001-of-00003.safetensors
├── ...
├── experiment_cfg/               # conf.yaml, final_model_config.json, dataset_statistics.json
├── processor/                    # <-- subdir layout required for inference
│   ├── processor_config.json
│   ├── embodiment_id.json
│   └── statistics.json
├── wandb_config.json
└── train.log

Consumers load checkpoints via AutoProcessor.from_pretrained(ckpt/processor) (note the subdir). RLDXPolicy handles this automatically.

Troubleshooting

"RuntimeError: CUDA out of memory"

• Raise --gradient-accumulation-steps to keep global batch constant.
• Switch to ZeRO-3: the stage is read from TrainingConfig.deepspeed_stage (nested; not a top-level CLI flag). Flip it to 3 in your launcher or the defaults in rldx/configs/training/training_config.py, which dispatches to rldx/configs/deepspeed/zero3_config.json.
• Drop --tune-top-llm-layers to 0 to freeze more of the backbone.

"Loaded memory_length=1 from checkpoint" when you expected more

The base checkpoint was trained without memory; --use-memory alone is not enough, you also need --memory-length, --memory-n-cog-tokens, and usually --concat-memory. Double-check that the base model actually supports the memory path (all current RLDX-1 pre-training runs do).

Where to next

• installation.md — environment setup
• evaluation.md — run the trained checkpoint on benchmarks
• inference_server.md — serve checkpoints over ZeroMQ
• architecture.md — what the config flags are actually wiring up