One Model for All Tasks: Leveraging Efficient World Models in Multi-Task Planning

El artículo presenta ScaleZero, un modelo único para la planificación multi-tarea que combina una arquitectura de Mezcla de Expertos (MoE) para resolver conflictos de gradientes y una estrategia de Escalamiento Dinámico de Parámetros (DPS) para optimizar la capacidad del modelo, logrando un rendimiento comparable al de agentes especializados con una menor cantidad de interacciones con el entorno.

Lo esencial

¡Claro que sí! Imagina que quieres entrenar a un solo robot para que sea un maestro en todo: que sepa jugar al ajedrez, que aprenda a conducir un coche, que resuelva acertijos de texto y que juegue a videojuegos de arcade al mismo tiempo.

El problema es que, hasta ahora, intentar entrenar a un "cerebro" único para tantas cosas diferentes era como intentar enseñar a un estudiante a ser un cirujano, un piloto de carreras y un poeta al mismo tiempo, sin que se confundiera. El estudiante se agotaba, olvidaba lo que sabía y terminaba haciendo mal todas las tareas.

Este paper, titulado "ScaleZero", presenta una solución brillante con dos grandes ideas para crear ese "robot maestro". Aquí te lo explico con analogías sencillas:

1. El Problema: El "Colapso de la Plasticidad" (El cerebro se congela)

Imagina que tienes un equipo de limpieza (el modelo de IA) que debe limpiar 26 habitaciones diferentes (26 videojuegos distintos).

• La habitación fácil: "Pong" (un juego de tenis simple). Se limpia en 5 minutos.
• La habitación difícil: "Seaquest" (un juego complejo con muchos enemigos y laberintos). Se tarda horas.

En los modelos antiguos (como UniZero), el equipo de limpieza intentaba limpiar todas las habitaciones al mismo tiempo con la misma fuerza. Como la habitación fácil se limpia rápido, el equipo se enfocaba tanto en ella que olvidaba cómo limpiar la difícil. Peor aún, los neurona del cerebro del robot se "dormían" (se volvían inertes) porque no sabían cómo adaptarse a la complejidad. A esto los autores lo llaman "Colapso de la Plasticidad": el robot deja de aprender cosas nuevas y se queda atascado.

2. La Solución 1: ScaleZero (El Edificio con "Expertos" en lugar de un solo gigante)

Para arreglar esto, los autores cambiaron la arquitectura del cerebro del robot. En lugar de tener un solo cerebro gigante que intenta hacer todo, crearon ScaleZero, que funciona como un gran edificio de oficinas con "Expertos".

• La analogía del MoE (Mezcla de Expertos): Imagina que el edificio tiene una recepción inteligente (un "router"). Cuando llega una tarea:
  • Si es un juego de tenis simple, la recepción lo envía al Experto en Deportes.
  • Si es un juego de exploración espacial, lo envía al Experto en Astronomía.
  • Si es un acertijo de texto, lo envía al Experto en Literatura.

Cada "experto" es un pequeño cerebro especializado. Al separar las tareas, el Experto en Astronomía no se distrae con el tenis, y el Experto en Literatura no se confunde con los cohetes. Esto evita que el robot se agote y le permite aprender muchas cosas a la vez sin que una tarea interfiera con la otra.

3. La Solución 2: DPS (El "Crecimiento Inteligente" o "Dinámico")

Incluso con los expertos, hay un problema: ¿Qué pasa si el robot aprende el tenis muy rápido, pero sigue luchando con el ajedrez? ¿Deberíamos seguir gastando energía y tiempo entrenando al experto de tenis? ¡No! Eso es desperdiciar recursos.

Aquí entra la segunda gran idea: DPS (Escalado Dinámico de Parámetros).

• La analogía del "Kit de Herramientas en Crecimiento": Imagina que el robot empieza con una caja de herramientas básica.
  • Al principio, usa todas las herramientas para aprender lo básico de todo.
  • En cuanto el robot domina el tenis (la tarea se "resuelve"), congelamos esas herramientas específicas para el tenis. Ya no las tocamos más.
  • Si el robot sigue luchando con el ajedrez, abrimos una nueva caja de herramientas (añadimos un módulo nuevo, llamado LoRA) específicamente para el ajedrez.

Es como si el robot tuviera un sistema de crecimiento: solo añade más "cerebro" cuando es estrictamente necesario. Esto hace que el entrenamiento sea mucho más rápido y eficiente, ahorrando hasta un 28.5% de tiempo y energía (interacciones con el entorno) porque deja de gastar tiempo en cosas que ya sabe hacer.

¿Qué lograron?

Con estas dos innovaciones (el edificio de expertos + el crecimiento dinámico), ScaleZero logró:

1. Ser un "Todo Terreno": Un solo modelo que juega tan bien como 26 expertos separados en videojuegos, 18 tareas de control de robots y 4 juegos de texto.
2. Ser Eficiente: Aprendió todo esto usando menos datos que los métodos anteriores.

En resumen

Antes, intentar entrenar a un solo agente para muchas tareas era como intentar que un solo estudiante aprendiera todo el contenido de la escuela en un día: se agotaba y fallaba.
ScaleZero es como crear una escuela con aulas especializadas (donde cada experto enseña su materia) y un sistema de becas dinámico (que solo invierte más dinero en las materias donde el estudiante aún tiene dificultades).

El resultado es un agente inteligente, flexible y eficiente, capaz de ser un verdadero "generalista" en el mundo de la inteligencia artificial.

Resumen técnico

Aquí tienes un resumen técnico detallado del artículo "ONE MODEL FOR ALL TASKS: LEVERAGING EFFICIENT WORLD MODELS IN MULTI-TASK PLANNING", presentado en ICLR 2026.

1. El Problema: Colapso de Plasticidad y Conflictos de Gradiente

El trabajo aborda los desafíos fundamentales del Aprendizaje por Refuerzo Multi-Tarea (MTRL) en entornos heterogéneos, donde las tareas varían significativamente en sus espacios de observación, acción y complejidad subyacente.

• Limitaciones de los Modelos Actuales: Los modelos de mundo unificados existentes (como UniZero) suelen utilizar arquitecturas monolíticas (Transformers densos). Aunque funcionan bien en tareas individuales, fallan en conjuntos de tareas diversas.
• Diagnóstico del Fallo: Los autores identifican dos fenómenos críticos que causan el fracaso:
  1. Colapso de la Plasticidad: Las tareas más simples y de convergencia rápida dominan los gradientes, suprimiendo las señales de aprendizaje de tareas complejas. Esto lleva a una pérdida de la capacidad del modelo para adaptarse a nuevos datos.
  2. Interferencia Representacional: Se observa un aumento en la proporción de neuronas inactivas ("dormant neurons") y una inflación descontrolada de la norma del estado latente. Esto indica que el espacio de características se colapsa en una subespacio de baja dimensión, incapaz de representar la diversidad de tareas.
• Asignación Estática de Recursos: Las arquitecturas convencionales aplican los mismos recursos computacionales a todas las tareas, independientemente de su dificultad o estado de aprendizaje, lo que resulta en una ineficiencia computacional.

2. Metodología: ScaleZero y Escalamiento Dinámico de Parámetros (DPS)

Para resolver estos problemas, los autores proponen un enfoque dual que actúa a nivel de arquitectura (iteración única) y a nivel de proceso de aprendizaje (iteraciones globales).

A. Arquitectura: ScaleZero

ScaleZero es un modelo de mundo unificado diseñado para mitigar conflictos de gradiente y mantener la plasticidad. Se basa en una exploración sistemática del espacio de diseño de UniZero, resultando en tres componentes clave:

1. Backbone Mixture-of-Experts (MoE):
   • Reemplaza las capas de redes neuronales densas (FFN) en el Transformer por capas MoE dispersas.
   • Utiliza un mecanismo de enrutamiento (gating) que selecciona dinámicamente un subconjunto de "expertos" especializados para cada entrada.
   • Beneficio: Esto crea vías computacionales condicionales, permitiendo que diferentes tareas utilicen sub-redes distintas, reduciendo drásticamente la interferencia de gradientes y manteniendo la plasticidad del modelo.
2. Codificador ViT (Vision Transformer):
   • Utiliza un ViT entrenado desde cero para la extracción de características visuales, demostrando ser más escalable que las arquitecturas ResNet en este contexto.
3. Normalización Estándar:
   • Reemplaza la normalización SimNorm (usada en UniZero) por LayerNorm estándar, encontrando un mejor equilibrio entre estabilidad y expresividad representacional en entornos heterogéneos.

B. Estrategia de Entrenamiento: Escalamiento Dinámico de Parámetros (DPS)

DPS es una estrategia en línea que asigna dinámicamente la capacidad del modelo según el progreso de aprendizaje, evitando el desperdicio de recursos en tareas ya resueltas.

• Curriculum de Complejidad: El entrenamiento se divide en etapas. Comienza con un "calentamiento" (Stage 0) donde se entrena un modelo base compartido.
• Expansión Progresiva: A medida que las tareas se "resuelven" (alcanzan un umbral de rendimiento), se congelan sus parámetros y se inyectan nuevos módulos LoRA (Low-Rank Adaptation) para las tareas restantes que aún requieren aprendizaje.
• Aislamiento de Parámetros: Al congelar el backbone y los adaptadores anteriores, DPS previene el "olvido catastrófico" y permite una plasticidad dirigida solo hacia las tareas difíciles.
• Eficiencia: Reduce la necesidad de interacciones con el entorno al centrar los recursos computacionales exclusivamente en las tareas no resueltas.

3. Contribuciones Clave

1. Diagnóstico Cuantitativo: Establecen un vínculo empírico y teórico entre el colapso del rendimiento en MTRL y la degradación de la plasticidad interna (aumento de neuronas dormidas e inflación de la norma latente).
2. Arquitectura ScaleZero: Presentan un modelo unificado que supera o iguala a los agentes especializados de tarea única en 48 tareas diversas (Atari, DMControl, Jericho), demostrando que un solo modelo puede manejar la heterogeneidad sin sufrir interferencia destructiva.
3. Estrategia DPS: Introducen un mecanismo de asignación de recursos adaptativo que reduce las interacciones con el entorno en un 28.5% manteniendo un rendimiento competitivo, ofreciendo un mejor equilibrio entre coste computacional y rendimiento final.
4. Análisis Teórico de MoE: Proporcionan una justificación teórica (Teorema 5.1) que demuestra que el límite superior del conflicto de gradientes en una capa MoE es estrictamente menor que en una capa densa, validando la eficacia de la especialización de expertos.

4. Resultados Experimentales

Los experimentos se realizaron en un entorno puramente de RL en línea (sin datos de expertos) en tres benchmarks:

• Atari 100k (26 juegos): ScaleZero (MT) logra una puntuación media normalizada (HNS) superior a la de 26 agentes UniZero entrenados individualmente (ST). Esto demuestra una transferencia de conocimiento positiva neta, superando el problema de la interferencia negativa.
• DeepMind Control Suite (18 tareas continuas): ScaleZero alcanza un rendimiento mediano superior al de los especialistas, validando su versatilidad en control continuo.
• Jericho (4 juegos de texto): El modelo logra un rendimiento comparable a agentes especializados y al estado del arte (CALM+OC), demostrando su agnosticismo a la modalidad (texto, visión, estado).
• Eficiencia (DPS): En el benchmark DMC, la variante ScaleZero-DPS alcanza el mismo rendimiento que el modelo base utilizando solo el 71.5% de las interacciones con el entorno.

5. Significado e Impacto

Este trabajo es significativo porque:

• Rompe el paradigma monolítico: Demuestra que los modelos de mundo unificados no están condenados a fallar en tareas heterogéneas si se diseñan con mecanismos de especialización (MoE) y gestión dinámica de recursos (DPS).
• Hacia Agentes Generalistas: Proporciona una arquitectura robusta y escalable para crear agentes generalistas capaces de aprender múltiples habilidades complejas simultáneamente sin necesidad de datos masivos de expertos o fine-tuning offline.
• Eficiencia de Muestreo: La estrategia DPS ofrece una ruta viable para reducir drásticamente el coste computacional en el entrenamiento de agentes multi-tarea, un factor crítico para la aplicación en el mundo real.

En resumen, ScaleZero representa un avance fundamental en la planificación multi-tarea, combinando la especialización arquitectónica con la adaptabilidad dinámica para superar las limitaciones de plasticidad y eficiencia que han frenado el desarrollo de agentes generalistas.