Cosmos 3 Edge de NVIDIA corre el control del robot a bordo

NVIDIA presento Cosmos 3 Edge, un modelo de mundo pensado para que la politica de control de un robot corra en el propio robot y no en un servidor remoto. Con 4.000 millones de parametros (incluido un razonador de 2.000 millones basado en Nemotron), el modelo pertenece a la familia Cosmos 3 y comparte el mismo preentrenamiento con datos del mundo fisico que Cosmos 3 Nano y Cosmos 3 Super, pero es lo bastante chico para caber en un NVIDIA Jetson Thor.

Los modelos de mundo aprenden como se mueven e interactuan los objetos, es decir, como caen, se deslizan o responden al contacto. Ese conocimiento previo sirve como punto de partida para entrenar una politica de robot, en lugar de obligarla a aprender toda la fisica desde cero con demostraciones especificas de la tarea. El checkpoint liberado esta disponible en Hugging Face y el codigo, en el repositorio abierto cosmos-framework.

¿Por que importa correr el modelo dentro del robot?

Llevar un modelo de mundo a un robot fisico choca con dos limites practicos. El primero es la memoria del dispositivo: el modelo y su estado de ejecucion tienen que caber en lo que el robot lleva encima. El segundo es la latencia de control: el pipeline completo de inferencia debe ser lo bastante rapido para sostener la frecuencia de control que exige la maquina.

Cosmos 3 Edge responde a los dos. En BF16 los pesos ocupan alrededor de 9 GB, cifra que entra en la memoria del Jetson AGX Thor, de modo que tanto el servidor de politica como el cliente de control viven en el robot. Sobre un Jetson AGX Thor T5000, la politica DROID genera cada bloque de acciones en cerca de 1,53 segundos trabajando a 640 x 540 de resolucion y 15 Hz, mientras que un solo bloque cubre unos 2,13 segundos de movimiento. Como el siguiente bloque queda listo antes de que termine el anterior, el brazo se mueve de forma continua sin ninguna GPU de centro de datos en el circuito.

En las tareas de lazo cerrado de RoboLab, la politica post-entrenada alcanza un 22,9% de exito. NVIDIA lee ese resultado como prueba de que un modelo de mundo de 4B puede funcionar como columna vertebral de una politica en tiempo real ejecutada a bordo.

¿Con que datos se entrena la politica?

La politica del tutorial se entrena sobre el conjunto nvidia/Cosmos3-DROID, que reune 76.000 trayectorias teleoperadas exitosas, cerca de 350 horas repartidas en 86 tareas y 564 escenas, capturadas con un brazo Franka Panda y una pinza Robotiq. El paquete viene en formato LeRobotDataset v3.0 a 640 x 360.

La preparacion tiene tres etapas: filtrar los cuadros inactivos o ajenos a la tarea, seleccionar las demostraciones exitosas para entrenamiento y aplicar recortes aleatorios, reescalado y variacion de color durante el entrenamiento.

hf download nvidia/Cosmos3-DROID --repo-type dataset --local-dir Cosmos3-DROID

Quien quiera usar datos propios debe convertirlos al formato LeRobot Dataset v3, con video por camara, estados de las articulaciones, estado de la pinza, acciones e instruccion de la tarea. Para un montaje tipo DROID con Franka basta cambiar la ruta del conjunto de datos; otra morfologia exige su propia configuracion de experimento, que define el espacio de acciones, la dimensionalidad, la disposicion de camaras y la normalizacion. Cosmos 3 admite varios cuerpos robotticos, entre ellos Franka de doble brazo, UR, WidowX 250 y LeRobot SO101.

¿Cuanto computo hace falta para post-entrenarlo?

Aca conviene bajar las expectativas: esto no es un fine-tune de una sola GPU. La corrida validada por NVIDIA usa 64 nodos de 4 GB200 durante 60.000 iteraciones, unas 68 horas, lo que equivale a cerca de 17.400 horas-GB200. El hardware de entrenamiento validado es una NVIDIA DGX Station con el superchip GB200 Grace Blackwell o el GB300 Grace Blackwell Ultra Desktop, con CUDA 13.0 y contenedor NGC 26.06-py3.

El post-entrenamiento se divide en cuatro pasos: descargar el conjunto de datos, convertir el checkpoint base al formato de checkpoint distribuido (DCP), aplicar el filtro de curacion y lanzar la corrida.

# Descarga del dataset Cosmos3-DROID (split de exitos)
hf download nvidia/Cosmos3-DROID --repo-type dataset --local-dir /path/to/Cosmos3-DROID

# Conversion del checkpoint base a DCP (una sola vez, corre en CPU)
python -m cosmos_framework.scripts.convert_model_to_dcp \ -o /path/to/Cosmos3-Edge-dcp --checkpoint-path Cosmos3-Edge

# Lanzamiento del post-entrenamiento
bash examples/launch_sft_action_policy_droid_nano.sh

El lanzador que trae el repositorio registra la receta de Cosmos 3 Nano. Para entrenar Cosmos 3 Edge hay que hacer tres cambios antes de ejecutarlo: reemplazar la importacion de NANO_MODEL_CONFIG por EDGE_MODEL_CONFIG, apuntar BASE_CHECKPOINT_PATH al checkpoint DCP de Edge y renombrar el script de lanzamiento. Todo lo demas, incluidos el conjunto de datos, el espacio de acciones, el filtro de curacion y el calendario de entrenamiento, queda igual.

¿Como se despliega en el Jetson Thor?

La politica se sirve mediante un servidor WebSocket que habla el protocolo OpenPI, el mismo que usa el ecosistema de politicas DROID. El cliente envia un diccionario de observaciones y el servidor devuelve un bloque de acciones. En el caso de Edge, ese servidor corre nativamente en el Jetson Thor, asi que la peticion nunca sale del robot.

export HF_TOKEN=<tu_token_hf>
# Thor: ejecutar en modo eager; las ruedas Triton estandar no traen kernels sm_110a
export TORCHDYNAMO_DISABLE=1
python -m cosmos_framework.scripts.action_policy_server_robolab \
  --checkpoint_path nvidia/Cosmos3-Edge-Policy-DROID \
  --port 8000 \
  --format-prompt-as-json True
curl localhost:8000/healthz   # -> OK

La politica DROID esta condicionada por el estado, o sea que joint_position y gripper_position son entradas reales del modelo y no relleno. En el robot, la misma llamada se repite dentro de un lazo de replanificacion: cada ciclo lee camaras frescas y las posiciones medidas del brazo y la pinza, ejecuta un prefijo del bloque, y vuelve a preguntar.

while not task_done():
    result = client.infer(get_observation())
    execute_actions(result["action"][:16], hz=15)   # los primeros 16 de 32, luego replanifica

Ejecutar solo un prefijo y replanificar es practica estandar: la politica se autocorrige cada pocos segundos y, al estar condicionada por el estado, cada replanificacion parte desde donde el brazo realmente esta y no desde donde el bloque anterior supuso que estaria.

Validar en simulacion antes de tocar el robot

No hace falta un robot fisico para probar la politica. NVIDIA recomienda evaluarla primero en simulacion para evitar comportamientos inesperados. RoboLab, el banco de pruebas Isaac Lab-Arena que hay detras de la tabla de posiciones del mismo nombre, es abierto y su cliente se conecta al mismo servidor de politica: ejecuta cada bloque de acciones en fisica y devuelve las observaciones renderizadas, con 120 tareas de manipulacion condicionadas por lenguaje.

Que significa para un taller en Chile

La barrera real no esta en la inferencia sino en el entrenamiento. Un Jetson AGX Thor bordea los 3.500 dolares de lista, un monto alto pero al alcance de un laboratorio universitario o una integradora mediana; las 17.400 horas-GB200 del post-entrenamiento, en cambio, quedan fuera de discusion para cualquiera que no arriende computo en la nube. La lectura practica para la region es que el checkpoint ya post-entrenado, publicado bajo la ficha Cosmos3-Edge-Policy-DROID, es el activo aprovechable: se descarga, se sirve en el Thor y se evalua en RoboLab sin gastar un peso en entrenamiento. Quien tenga un brazo Franka, UR o incluso un LeRobot SO101 de bajo costo entra al ecosistema por el lado del despliegue, no por el del entrenamiento.