Ajuste de Cosmos Predict 2.5 con LoRA/DoRA para video-robots | Keryc
NVIDIA publica una guía técnica para adaptar Cosmos Predict 2.5 a tareas concretas de robótica, mostrando cómo usar LoRA y DoRA para generar trayectorias sintéticas de robots sin tener que reentrenar todo el modelo. ¿El objetivo? Crear videos físicamente plausibles condicionados por texto e imágenes y usarlos como datos escalables para entrenar políticas de robots.
Qué anuncia NVIDIA
Cosmos Predict 2.5 es un world model a gran escala que genera videos físicamente consistentes condicionados por texto, imágenes o clips. NVIDIA muestra un pipeline de fine-tuning parametricamente eficiente usando LoRA y DoRA para adaptar el modelo a dominios específicos (por ejemplo, manipulación robótica o vistas de cámara concretas).
La novedad práctica: en lugar de reentrenar los 2B parámetros del modelo (caro y propenso a olvidar conocimiento general), se inyectan adaptadores pequeños y portables que permiten entrenar en una sola GPU potente y luego cambiar adapters según la tarea.
Por qué esto es útil para robótica
Recoger trayectorias reales es lento y costoso. ¿Y si pudieras generar miles de trayectorias sintéticas que sean físicamente plausibles y específicas a tu setup de cámara o robot? Eso acelera la iteración de políticas.
LoRA/DoRA mantiene el backbone congelado y solo ajusta parámetros de baja dimensión. El adapter resultante ocupa poco espacio, se comparte fácilmente y evita el problema de catastrophic forgetting que aparece al reentrenar todo el modelo.
Cómo implementan LoRA y DoRA en Cosmos Predict 2.5
Cosmos Predict 2.5 tiene tres submódulos: un VAE que codifica videos a latentes, un text encoder y un DiT que realiza la difusión en latente. Durante el ajuste, todos los pesos base se mantienen congelados y se inyectan adaptadores en puntos clave del DiT.
Los módulos objetivo típicos son las proyecciones de atención to_q, to_k, to_v, to_out.0 y algunas capas feedforward ff.net.0.proj, ff.net.2.
Ejemplo mínimo de cómo se carga el pipeline y se añade LoRA (resumido):
Configuras LoraConfig con r (rank), lora_alpha y target_modules, y luego haces dit.add_adapter(lora_config).
Para estabilidad numérica, los parámetros entrenables del adapter se upcastean a float32 cuando se entrena con bf16.
Si activas use_dora=True, el adapter usa la descomposición magnitud-dirección (DoRA) en vez de la actualización low-rank pura. No necesitas cambiar el loop de entrenamiento.
Datos y pipeline de entrenamiento
Dataset de entrenamiento: 92 videos de manipulación robótica con prompts textuales (pick-and-place).
Dataset de evaluación: 50 pares (prompt, imagen inicial) para generar un video por par.
Formato esperado tras preprocesar:
gr1_dataset/train con videos/*.mp4, metas/*.txt y metadata.csv.
gr1_dataset/test con pares filename1.txt, filename1.png, etc.
El VideoDataset carga cada muestra como (caption, video) y si el video es más largo que num_frames extrae una ventana aleatoria cada época para aumentar variantes temporales.
Algoritmo de entrenamiento y pérdida
Cosmos Predict usa rectified flow. En pocas palabras, para un nivel de ruido σ_t se construye x_t = σ_t * noise + (1 - σ_t) * clean y el modelo predice la velocidad que transforma el ruido hacia la muestra limpia. Los primeros dos frames se usan como condición y no se les añade ruido.
La pérdida es una MSE entre la velocidad predicha y la velocidad objetivo (noise - clean_latent) calculada solo sobre los frames no condicionados.
El optimizador es AdamW sobre los parámetros del adapter y el scheduler calienta linealmente la tasa de aprendizaje y luego la decae según una política definida.
Requisitos de hardware: al menos una GPU de 80 GB para entrenamiento en un solo dispositivo; 8× H100 para iteraciones rápidas. NVIDIA muestra que 100 epochs dan resultados sólidos: ~17 h en 1 H100, ~2.5 h en 8 H100s.
Los adapters se guardan en formato pytorch_lora_weights.safetensors cada checkpointing_epochs.
En inferencia cargas el pipeline, aplicas pipe.load_lora_weights("/path/to/lora/checkpoint") y opcionalmente pipe.fuse_lora(lora_scale=1.0) para fusionar los pesos y evitar overheads en tiempo de inferencia.
También hay una opción para hacer la inicialización del ruido reproducible entre arquitecturas: arch_invariant_rand.
Evaluación: métricas y juez LLM
Para medir calidad geométrica usan el error de Sampson:
Temporal Sampson Error: consistencia entre frames consecutivos.
Cross-view Sampson Error: consistencia entre vistas simultáneas.
Además usan Cosmos Reason2 (un LLM) como juez con dos rubricas YAML:
video_physics.yaml: juzga plausibilidad física sin ver el prompt.
video_IF.yaml: evalúa si el video cumple la instrucción dada (prompt + video).
Cada video recibe una puntuación de 1 a 5 en ambas dimensiones.
Resultados y lecciones prácticas
Compararon el modelo base (sin fine-tuning), LoRA y DoRA en r=8 y r=32. Hallazgos principales:
Antes del ajuste el modelo base tiende a fallar con manos de robot (hallucina manos humanas), no sigue siempre la mano correcta y muestra jitter.
Fine-tuning con LoRA y DoRA reduce esos problemas: mejor seguimiento de instrucciones y menos jitter.
Rank mayor (r=32 vs r=8) mejora la capacidad de seguir instrucciones (por ejemplo, usar la mano correcta o agarrar el objeto indicado). No obstante, la consistencia geométrica y la plausibilidad física no mejoran mucho con rank mayor.
Hipótesis interpretativa: las priors geométricas y físicas ya están bien codificadas en el backbone congelado; el adapter solo necesita desplazar la distribución hacia la apariencia y comportamiento del dominio robótico. Por eso r=8 puede ser suficiente para arreglar la apariencia y estructura básica.
Recomendación práctica:
Si tu memoria es limitada o te importa el tamaño del adapter: empieza con LoRA r=8.
Si observas inestabilidad con LoRA a bajo rank o tienes presupuesto de memoria, prueba DoRA r=32. DoRA puede estabilizar el aprendizaje por su descomposición magnitud-dirección.
Para reproducibilidad, usa arch_invariant_rand. Para despliegue rápido, usa pipe.fuse_lora().
Conclusión
La guía de NVIDIA muestra que adaptar un world model grande a un dominio robótico es práctico con estrategias de entrenamiento económico en parámetros. Generar trayectorias sintéticas con Cosmos Predict 2.5 + LoRA/DoRA abre la puerta a ciclos de desarrollo más rápidos para aprendizaje por imitación y simulación de políticas.
Si trabajas con robots y datos reales escasos, este enfoque te permite crear datasets específicos por cámara/robot sin sacrificar el conocimiento general del modelo. ¿La moraleja? No siempre necesitas reentrenar todo: un adapter bien diseñado puede ser la pieza que acelera tu proyecto.
