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

O artigo apresenta o ScaleZero, um modelo único para planejamento multi-tarefa que combina uma arquitetura Mixture-of-Experts e uma estratégia de Escalonamento Dinâmico de Parâmetros (DPS) online para mitigar conflitos de gradiente e otimizar a alocação de capacidade, alcançando desempenho comparável a agentes especializados com menor interação com o ambiente.

O essencial

Imagine que você está tentando ensinar um único robô a fazer de tudo: jogar xadrez, cozinhar, dirigir um carro e resolver um quebra-cabeça de texto. O problema é que, se você tentar ensinar tudo ao mesmo tempo usando o mesmo "cérebro" padrão, o robô acaba ficando confuso. As regras do xadrez atrapalham as do cozimento, e ele acaba não aprendendo nada direito. Isso é o que os cientistas chamam de conflito de gradientes e colapso da plasticidade.

Este artigo, apresentado na conferência ICLR 2026, apresenta uma solução brilhante chamada ScaleZero. Vamos entender como ele funciona usando analogias do dia a dia.

1. O Problema: O "Cérebro" Quebrado

Antes do ScaleZero, existiam modelos que tentavam fazer tudo de uma vez só (como o UniZero). Imagine que esse modelo é um único funcionário tentando atender 26 clientes diferentes ao mesmo tempo, cada um pedindo algo totalmente diferente.

• O cliente A quer algo rápido e simples (como jogar Pong).
• O cliente B quer algo complexo e difícil (como jogar Seaquest).

O funcionário tenta atender a todos, mas acaba focando no cliente A porque é mais fácil. O cliente B fica esperando, e o funcionário começa a cometer erros, esquecendo o que aprendeu antes. O "cérebro" do robô fica travado, com neurônios "dormindo" e sem capacidade de aprender coisas novas. É como tentar usar uma única chave para abrir 26 portas diferentes: nenhuma abre direito.

2. A Solução de Arquitetura: O "Escritório de Especialistas" (MoE)

Para resolver isso, os autores criaram o ScaleZero. Em vez de ter um único funcionário, eles montaram um escritório cheio de especialistas.

• A Analogia do MoE (Mistura de Especialistas): Imagine que o robô não é uma pessoa, mas sim uma grande sala de reuniões. Dentro dessa sala, existem vários consultores especialistas (os "Experts").
  • Quando chega um pedido de Pong, o "gerente" (o roteador) chama apenas o especialista em jogos rápidos.
  • Quando chega um pedido de Seaquest, ele chama o especialista em exploração complexa.
  • Cada especialista trabalha em sua própria mesa, sem interferir no trabalho do outro.

Isso evita que as tarefas se misturem e confundam o robô. O ScaleZero usa essa estrutura de "Mistura de Especialistas" (MoE) para garantir que, mesmo aprendendo tudo ao mesmo tempo, cada tarefa tenha seu próprio espaço para crescer. É como ter uma equipe de médicos: um cardiologista não tenta tratar uma fratura de perna; ele deixa isso para o ortopedista.

3. A Solução de Processo: O "Orçamento Dinâmico" (DPS)

Agora, imagine que você tem um orçamento limitado de tempo e dinheiro para treinar esse robô. Tradicionalmente, você gastaria o mesmo tempo treinando o robô para tarefas fáceis (que ele já domina) e para tarefas difíceis (que ele ainda não entende). Isso é um desperdício!

Para resolver isso, eles criaram o DPS (Escalonamento Dinâmico de Parâmetros).

• A Analogia do Orçamento Inteligente: Pense no DPS como um gerente de projeto muito esperto.
  • Fase 1: Ele treina o robô em tudo um pouco para criar uma base sólida.
  • Fase 2: Assim que o robô domina uma tarefa fácil (como Pong), o gerente diz: "Parou! Não gaste mais tempo nisso. Vamos focar no que ainda está difícil."
  • Expansão: Quando o robô encontra uma tarefa muito difícil, o gerente "contrata" um novo especialista temporário (usando uma técnica chamada LoRA, que é como adicionar um "anexo" leve ao cérebro do robô) para ajudar naquela tarefa específica, sem precisar reconstruir todo o cérebro.

Isso permite que o robô aprenda de forma mais eficiente, gastando cerca de 28% menos tempo e dados do que os métodos antigos, mas chegando ao mesmo resultado.

4. Os Resultados: O "Polímata" Perfeito

O teste foi feito em três mundos muito diferentes:

1. Jogos de Arcade (Atari): Jogos visuais rápidos e complexos.
2. Controle Robótico (DMC): Tarefas físicas contínuas, como fazer um robô andar ou correr.
3. Jogos de Texto (Jericho): Aventuras baseadas em leitura e lógica.

O ScaleZero conseguiu ser um "generalista" incrível. Ele aprendeu a jogar todos os jogos e controlar todos os robôs usando apenas um único modelo, sem precisar de um robô diferente para cada tarefa. E o melhor: ele fez isso com a mesma eficiência de especialistas que só sabiam fazer uma coisa, mas usando muito menos recursos.

Resumo em uma Frase

O ScaleZero é como transformar um único funcionário sobrecarregado em uma equipe organizada de especialistas, onde cada um cuida do que sabe fazer de melhor, e um gerente inteligente decide exatamente quando e onde alocar recursos para que ninguém perca tempo com o que já foi resolvido.

Isso é um passo gigante para criar Agentes Gerais de IA — robôs que podem realmente aprender qualquer coisa, do mesmo jeito que um humano aprende, sem se confundir no processo.

Resumo técnico

Título: Um Modelo para Todas as Tarefas: Aproveitando Modelos de Mundo Eficientes no Planejamento Multitarefa

1. O Problema

O artigo aborda os desafios críticos enfrentados por Modelos de Mundo Unificados (como o UniZero) quando aplicados a cenários de Aprendizado por Reforço Multitarefa (MTRL) heterogêneo. Embora esses modelos sejam eficazes em tarefas únicas, a aplicação de um único modelo compartilhado em uma ampla variedade de tarefas (com diferentes espaços de observação, ações e complexidades dinâmicas) revela duas falhas fundamentais:

1. Colapso da Plasticidade (Plasticity Collapse): Em ambientes multitarefa, tarefas mais simples e que convergem rapidamente dominam os gradientes, suprimindo os sinais de aprendizado de tarefas complexas. Isso leva a uma perda progressiva da capacidade da rede de se ajustar a novos dados (plasticidade), resultando em um "colapso" onde o modelo falha catastróficamente em tarefas desafiadoras (ex: Seaquest no Atari), enquanto mantém bom desempenho em tarefas simples (ex: Pong).
2. Alocação Estática de Recursos: Arquiteturas convencionais utilizam uma abordagem "tamanho único", aplicando os mesmos recursos de aprendizado (atualizações de modelo e coleta de dados) a todas as tarefas, independentemente de sua dificuldade ou estágio de convergência. Isso gera ineficiência computacional, desperdiçando recursos em tarefas já resolvidas em vez de focar nas que necessitam de mais aprendizado.

O diagnóstico quantitativo realizado pelos autores mostra que o colapso da plasticidade está correlacionado com uma inflação descontrolada da norma do estado latente e um aumento drástico na proporção de neurônios "dormientes" (dormant neurons) no backbone do Transformer.

2. Metodologia

Para resolver esses problemas, os autores propõem duas abordagens complementares: uma solução arquitetônica interna e uma estratégia de processo externa.

A. Arquitetura: ScaleZero

O ScaleZero é um novo modelo de mundo unificado projetado para mitigar conflitos de gradiente e o colapso da plasticidade. Suas principais inovações arquitetônicas incluem:

• Backbone Mixture-of-Experts (MoE): Substituição das camadas densas de feed-forward do Transformer por camadas MoE esparsas. Isso permite o cálculo condicional, onde diferentes tarefas são roteadas para sub-redes de "especialistas" distintas. Isso reduz a interferência representacional e os conflitos de gradiente, mantendo a capacidade do modelo sem aumentar proporcionalmente o custo computacional.
• Codificador ViT: Uso de um Vision Transformer (ViT) treinado do zero para extração de características visuais escaláveis, substituindo arquiteturas baseadas em ResNet.
• Normalização Latente: Uso de LayerNorm padrão em vez de SimNorm, encontrando um melhor equilíbrio entre estabilidade e expressividade representacional em ambientes heterogêneos.

B. Estratégia de Treinamento: Dynamic Parameter Scaling (DPS)

Para otimizar a alocação de recursos ao longo do processo de aprendizado, os autores introduzem o DPS, um mecanismo online que acopla a capacidade do modelo ao progresso de aprendizado:

• Curadoria Adaptativa de Tarefas: O sistema monitora o desempenho de cada tarefa. Uma vez que uma tarefa atinge um limiar de desempenho ("resolvida"), ela é removida do conjunto de tarefas ativas, cessando a coleta de dados e atualizações de gradiente para ela.
• Expansão Progressiva de Capacidade (LoRA): O modelo base é treinado inicialmente. À medida que novas tarefas difíceis são introduzidas ou tarefas existentes não convergem, o DPS injeta adapters LoRA (Low-Rank Adaptation) independentes e leves.
• Isolamento de Parâmetros: Ao avançar para uma nova etapa, os parâmetros do modelo base e dos adapters anteriores são congelados. O treinamento foca apenas no novo adapter e nos fatores de escala, prevenindo o esquecimento catastrófico e permitindo plasticidade direcionada apenas onde é necessária.

3. Principais Contribuições

1. Diagnóstico Quantitativo: Estabelecimento de uma ligação concreta entre a degradação de desempenho em MTRL e dinâmicas internas de aprendizado, especificamente o colapso da plasticidade (medido por neurônios dormientes e inflação de norma latente).
2. ScaleZero: Proposição de uma arquitetura unificada que supera o UniZero e agentes especialistas de tarefa única, demonstrando que a integração de MoE em modelos de mundo é fundamental para o aprendizado multitarefa eficaz.
3. Estratégia DPS: Desenvolvimento de uma estratégia de treinamento adaptativa que reduz o custo computacional e a necessidade de interações com o ambiente, otimizando o uso de recursos dinamicamente.
4. Análise Teórica e Empírica: Fornecimento de uma análise teórica (teoremas sobre limites superiores de conflito de gradiente em MoE) e validação empírica da eficácia do MoE e do DPS.

4. Resultados

Os métodos foram avaliados em três benchmarks distintos cobrindo 48 tarefas em total:

• Atari (26 jogos): O ScaleZero (MT) alcançou uma média de pontuação normalizada superior à média de 26 agentes UniZero treinados individualmente (tarefa única), demonstrando transferência de conhecimento positiva. O modelo superou especialmente tarefas de exploração complexas onde o baseline falhava.
• DeepMind Control Suite (18 tarefas contínuas): O ScaleZero obteve desempenho competitivo, superando a mediana dos agentes de tarefa única, validando sua eficácia em controle contínuo.
• Jericho (4 jogos baseados em texto): O modelo alcançou desempenho comparável a agentes especializados e ao forte baseline CALM+OC, provando sua agnosticidade de modalidade (texto).
• Eficiência (DPS): Na avaliação do DPS no benchmark DMC, a estratégia permitiu que o modelo atingisse desempenho comparável aos agentes de tarefa única utilizando apenas 71,5% das interações com o ambiente (uma redução de 28,5% no custo de amostragem).

5. Significância

Este trabalho representa um avanço significativo rumo a agentes generalistas verdadeiramente eficazes.

• Superação de Limitações Arquitetônicas: Demonstra que modelos monolíticos podem ser superados por arquiteturas esparsas (MoE) que permitem especialização dinâmica sem sacrificar a generalização.
• Eficiência de Amostragem: A estratégia DPS oferece um caminho prático para reduzir drasticamente o custo de treinamento em RL, um gargalo crítico para aplicações do mundo real.
• Fundação para Agentes Gerais: Ao combinar planejamento baseado em modelo de mundo (MCTS) com arquiteturas adaptativas e treinamento eficiente, o ScaleZero estabelece uma base robusta para o desenvolvimento de agentes capazes de operar em ambientes complexos, heterogêneos e multimodais.

O código e os modelos estão disponíveis publicamente no repositório LightZero da OpenDILab.