Fine-tuning de NVIDIA Cosmos Predict 2.5 con LoRA/DoRA para generación de videos de robots

Introducción

NVIDIA Cosmos Predict 2.5 es un world model a gran escala capaz de generar videos físicamente plausibles condicionados por texto, imágenes o clips de video. En robótica, adaptar este tipo de modelos a maniobras específicas —por ejemplo, manipulación con pinzas o vistas de cámara concretas— es clave para producir demostraciones sintéticas que alimenten el aprendizaje de políticas. Sin embargo, ajustar un modelo de 2 mil millones de parámetros por completo es costoso y puede provocar olvido catastrófico del conocimiento general.

En este artículo explico cómo realizar un fine-tuning eficiente en parámetros de Cosmos Predict 2.5 utilizando LoRA y DoRA, qué exige el proceso y cómo aprovechar los modelos afinados para generar trayectorias sintéticas útiles en tareas de robot learning. También incluyo consideraciones prácticas relevantes para equipos de investigación y empresas en Latinoamérica.

¿Por qué LoRA/DoRA?

LoRA (Low-Rank Adaptation) y su variante DoRA permiten inyectar adaptadores de tamaño reducido en un modelo base congelado. En lugar de reentrenar todos los pesos, se añaden módulos de bajo rango en puntos críticos (por ejemplo, las proyecciones de atención y capas feedforward). Ventajas principales:

• Drástica reducción de memoria y de parámetros entrenables.
• Posibilidad de entrenar en una sola GPU grande (mínimo 80 GB) o escalar a múltiples GPUs (ej.: 8× H100 para iteración rápida).
• Archivos de adaptador pequeños y portables que se pueden intercambiar para distintos dominios sin tocar el modelo base.

DoRA descompone cada peso en magnitud y dirección antes de aplicar la actualización de bajo rango; conceptualmente es una variante que mantiene la misma lógica de entrenamiento y reemplaza la parametrización de la actualización.

Requisitos y dependencias

Requisitos mínimos mencionados en la guía oficial:

• Python 3.10+
• PyTorch 2.5+ con CUDA
• diffusers (que instala transformers y peft automáticamente)
• accelerate
• Opcional: wandb para monitoreo
• Hardware: al menos una GPU de 80 GB para entrenamiento en una sola GPU; 8× H100s recomendado para iteración rápida

Instalación rápida de dependencias: pip install -U “diffusers[torch]” transformers accelerate peft wandb

Datos y preprocesamiento

La receta utilizada en el ejemplo emplea los mismos datasets del post-training GR00T Dreams:

• Dataset de entrenamiento: 92 videos de manipulación robótica acompañados de prompts que describen tareas de pick-and-place.
• Dataset de prueba: 50 pares (prompt, imagen) donde el modelo debe generar un video condicionando en el prompt y en el frame inicial.

Tras instalar diffusers, el ejemplo se encuentra en examples/cosmos. El script de descarga y preprocesado es download_and_preprocess_datasets.sh; al ejecutarlo se crea una estructura como:

gr1_dataset/train ├── metas/ │ └── *.txt ├── videos/ │ └── *.mp4 └── metadata.csv

Y el set de evaluación queda como un directorio plano con pares .txt y .png:

gr1_dataset/test ├── filename1.txt ├── filename1.png ├── filename2.txt ├── filename2.png └── …

Implementación del entrenamiento

El ejemplo de entrenamiento se encuentra en train_cosmos_predict25_lora.py. Los elementos principales son:

• VideoDataset: carga cada muestra como un par (caption, video) desde el directorio de entrenamiento. Si los videos son más largos que el número de frames a usar (args.num_frames), se muestrea una ventana contigua al azar cada época, lo que aporta aumento temporal. Internamente utiliza VideoProcessor de diffusers para redimensionar y normalizar los frames en tensores con forma (channels, frames, height, width).

• Inicialización y congelado del modelo base: Cosmos Predict 2.5 está compuesto por tres submódulos:

  • Un VAE que codifica videos a latentes.
  • Un text encoder que produce embeddings del prompt.
  • Un DiT que realiza la difusión en el espacio latente.

  Durante el entrenamiento se congelan los pesos del VAE, del encoder de texto y del DiT. Se inyectan adaptadores LoRA en las proyecciones de atención (to_q, to_k, to_v, to_out.0) y en capas feedforward (ff.net.0.proj, ff.net.2). Los parámetros de LoRA se convierten a float32 para estabilidad numérica cuando se usa precisión mixta bf16.

  Pasar use_dora=True activa DoRA en lugar de LoRA; no se requieren cambios adicionales en el bucle de entrenamiento.

• Pérdida: Cosmos Predict 2.5 usa rectified flow. En cada timestep t se construye una interpolación ruidosa xt = σt·noise + (1−σt)·clean para un nivel de ruido σt muestreado desde una distribución logit-normal. El modelo aprende a predecir la velocidad objetivo (noise − clean) mediante MSE. Los dos primeros frames se usan como condicionamiento (sin ruido en sus latentes). En el forward, DiT recibe como entradas las latentes con máscara de condición, embeddings del prompt y el timestep modificado para distinguir frames condicionados.

Generación de trayectorias sintéticas

Una vez entrenado el adaptador, el pipeline afinado puede generar videos condicionados por un frame inicial y un prompt textual. Estas salidas sirven como trayectorias sintéticas para entrenar políticas de control: en lugar de recolectar costosas ejecuciones reales en laboratorio, pueden generarse grandes cantidades de demostraciones en video que luego se conviertan en estados y acciones para el aprendizaje por imitación o métodos offline.

Hay que considerar el dominio de la cámara, las texturas y la física simulada: afinar con datos representativos (por ejemplo, la perspectiva de la cámara y el tipo de manipulación específica) mejora la utilidad de las trayectorias para aprendizaje por transferencia.

Consideraciones prácticas para Latinoamérica

Para grupos de investigación y startups en Latinoamérica, dos puntos son especialmente relevantes:

1. Acceso a hardware: la guía indica que un solo GPU de 80 GB es suficiente para entrenar adaptadores; esto hace viable el trabajo en centros de cómputo académicos, proveedores de nube o nodos regionales que ofrezcan GPU grandes. Alternativamente, el tamaño pequeño de los adaptadores facilita colaborar con instituciones que sí tengan acceso a clusters más potentes.

2. Costos y eficiencia: LoRA/DoRA reduce costos de cómputo y almacenamiento frente a reentrenar modelos completos. Para equipos con presupuestos limitados, esta estrategia permite iterar rápidamente en nuevas tareas robóticas sin requerir enormes inversiones en infraestructura.

Además, generar datos sintéticos acelera la experimentación en robótica aplicada a sectores clave en la región (agroindustria, manufactura ligera, logística), donde replicar entornos reales puede ser caro o peligroso.

Conclusión

El fine-tuning con LoRA/DoRA de Cosmos Predict 2.5 ofrece una vía práctica para adaptar un world model grande a dominios robóticos concretos sin necesidad de reentrenar todo el modelo. La combinación de adaptadores de bajo rango, entrenamiento en precisión mezclada y una receta de rectified flow permite entrenar con requerimientos de hardware y almacenamiento reducidos, generando trayectorias sintéticas útiles para downstream tasks en robot learning.

Para equipos en Latinoamérica, esta aproximación puede facilitar la experimentación aplicada con menor inversión inicial, siempre que se diseñe un dataset de adaptación representativo del dominio objetivo. Los pasos concretos y scripts están disponibles en el repositorio de diffusers (examples/cosmos), incluyendo el preprocesado y la plantilla de entrenamiento con LoRA/DoRA.

Si su equipo planea probar esta receta, comience por preparar un dataset de apoyo (inicialmente pequeño y representativo), verifique disponibilidad de una GPU de 80 GB o acceso a instancias con H100, e iteren afinando solo los adaptadores antes de considerar ajustes más amplios al pipeline.