Learning Massively Multitask World Models for Continuous Control

AI summary

Overview

• Research area: Reinforcement learning for continuous control — specifically massively multitask, online model-based RL and world models.
• Technical level: Advanced. The paper builds on TD-MPC2, uses latent self-predictive world models, cross-entropy value regression, and planning in latent space. The high-level message is accessible, but the method details assume RL background.
• Scope in one sentence: The authors introduce MMBench, a 200-task, 10-domain benchmark for massively multitask continuous control, and Newt, a language-conditioned multitask world model that is pretrained on demonstrations and then trained with online RL across all tasks at once.

What This Paper Is About

Most continuous-control RL research trains agents on a single task or in strictly offline regimes, which has reinforced the belief that online RL does not scale. The authors ask whether a single policy can instead be trained with online interaction across hundreds of control tasks simultaneously. To answer this, they build a large multitask benchmark (MMBench) and an agent (Newt) that combines demonstration pretraining with joint online RL over all tasks.

Key Contributions

1. MMBench, a benchmark for massively multitask RL. 200 unique continuous control tasks spanning 10 task domains (DMControl, DMControl Extended, Meta-World, ManiSkill3, MuJoCo, MiniArcade, Box2D, RoboDesk, OGBench, Atari). It comprises 159 existing tasks, 22 new tasks and task variants, and 19 entirely new arcade-style tasks (MiniArcade) — 41 new tasks in total. Every task has language instructions, demonstrations, and optionally image observations. All 220 MMBench tasks (including 20 test tasks) are released.

2. Newt, a language-conditioned multitask world model. Based on TD-MPC2, it takes low-dimensional states, language instructions, and optionally RGB observations, and produces actions via planning in the latent space of a decoder-free, self-predictive world model. Components: a CLIP text encoder, a DINOv2 image encoder, a state encoder, latent dynamics, reward prediction, terminal value prediction, and a policy prior.

3. Algorithmic extensions for the massively multitask online setting. Model-based pretraining over all world-model components using demonstrations, a model-based behavior-cloning term added to the policy objective, constrained planning that biases the planner toward the pretrained policy and anneals that bias to zero over the first 12% of training, and demonstration oversampling (50% of each update batch drawn from a separate demonstration buffer).

4. An open release and an accelerated training pipeline. 200+ model checkpoints, 4000+ task demonstrations, training and evaluation code, and asynchronous environment infrastructure. Training is distributed across processes and GPUs and compiled with torch.compile.

Main Findings

• Newt outperforms strong baselines overall. After 100M environment steps in total across all tasks with state observations, Newt is more data-efficient and achieves a higher overall performance than PPO, FastTD3, and multitask TD-MPC2. The authors attribute much of this to better performance in DMControl, DMControl Ext., ManiSkill, and MiniArcade.

• Some domains remain hard for every RL method. Performance is subpar across all RL methods in MuJoCo, Box2D, and Atari, where Newt is often similar to the simpler behavior-cloning baseline. The authors conjecture this may be due to the relative uniqueness of tasks in those domains, noting for example that many Atari games have little in common beyond their action space.

• Language conditioning helps differentiate tasks. Conditioning on language instructions gives a clear gain of 0.371 → 0.438 normalized score, with the largest improvement in domains where tasks cannot be distinguished by observations alone (for example RoboDesk). Language conditioning matches the performance of task indices on training tasks while also providing a mechanism for generalization to unseen tasks.

• Scaling model and batch size pays off in the multitask regime. In contrast to prior single-task RL work showing only marginal gains from model scaling, the authors observe a clear benefit from scaling both model and batch size, up to a point. They conjecture a compute-optimal (model, batch) size exists for a given number of tasks.

• Every way of using demonstrations helps, and combining them helps most. Pretraining, oversampling, and the model-based BC loss are individually helpful; using all of them in conjunction gives the best performance, with the largest benefits in hard-exploration tasks.

• Transfer to unseen tasks and embodiments. On a held-out set of tasks spanning multiple domains, a pretrained Newt achieves a zero-shot score of 0.192 versus 0.013 when trained from scratch, and an average score of 0.868 at 100k environment steps versus 0.480 for the baseline.

• Unseen language instructions can inhibit zero-shot generalization. In 10 unseen manipulation tasks (100 trials per task), replacing the unseen noun in the instruction with the seen but inaccurate word "cube" gives a Push success rate of 21.0 versus 0.3 for the unseen instruction, but a Pick success rate of 0 versus 10.5 for the unseen instruction. Across 6 pushing tasks this swap improves zero-shot success rate by 20.7%. To reflect true capability, all transfer experiments use unseen instructions.

• Open-loop control works over surprisingly long horizons. Newt plans and executes actions with no environment feedback on 8 diverse tasks with planning horizons up to 48 time steps — 16x longer than its training horizon of 3 — closely matching closed-loop control in most tasks. Failure modes include drifting dynamics (Walker Walk, DMControl), failing to decelerate after reaching the target (Lunarlander Takeoff, Box2D), and inability to predict stochastic elements (Assault, Atari).

• Vision gives a small overall gain but large domain-specific swings. Adding 224x224 RGB inputs and finetuning the whole model for 30M environment steps yields only +0.004 overall (0.438 → 0.442), but improves RoboDesk by +0.125 (to 0.500) and Meta-World by +0.069 (to 0.572), while decreasing MiniArcade by 0.004 (to 0.437) and DMControl by 0.029 (to 0.471).

• Training cost is practical but not trivial. On 200 tasks for 100M total environment steps with a 20M-parameter agent: 11.2 days on 1x RTX 3090, 7.3 days on 2x RTX 3090, 4.6 days on 2x RTX 5090, and 4.7 days for state+RGB on 2x RTX PRO 6000. Machines use an AMD EPYC 9354 CPU and at least 128 GB of RAM; the demonstration dataset requires 32 GB of disk space.

• Practical details. Language is encoded with CLIP-ViT/B (512-dimensional embeddings) and images with DINOv2/B (768-dimensional embeddings). State observations are 128-dimensional vectors, actions are 16-dimensional, Newt agents have 20M learnable parameters, and the replay buffer holds 10M transitions. Each task has 10–40 demonstrations, collected by single-task TD-MPC2 agents trained from low-dimensional states; the 200 single-task reference agents were trained for 5M environment steps each (1B total steps).

Methodology in Plain English

The authors start from TD-MPC2, a model-based RL method that learns a world model in a compressed latent space and picks actions by planning inside that latent space rather than decoding raw future observations. They extend it in three directions.

First, they make it multitask and language-conditioned. A frozen CLIP text encoder turns the task instruction into a vector, and a frozen DINOv2 encoder optionally turns camera images into a vector. These are concatenated with the low-dimensional state, and the result drives the encoder, latent dynamics, reward head, value head, and policy prior — all implemented as MLPs. Rewards and values are predicted with a cross-entropy (discrete regression) objective and values are modeled in log-space, because reward scales differ drastically across 200 tasks. Discount factors are set per task, using domain defaults where available.

Second, they make demonstrations do more work than just sit in a replay buffer. Before any online interaction, all learnable components are pretrained on demonstrations (with the value term temporarily disabled so the strong action supervision dominates). When online RL starts, planning is initially biased toward the pretrained policy to avoid an inaccurate value function hurting early behavior, and this bias is annealed to zero over the first 12% of training. Demonstrations are kept in a separate buffer and sampled at 50% of each update. A behavior-cloning term is added to the policy loss, which both uses expert actions directly and distills the actions chosen by the planner into the policy prior.

Third, they make it fast enough to be usable: environment stepping, rendering, frame stacking, image encoding, and language embedding caching are all handled by custom wrappers, and training, inference, and replay are distributed across processes and GPUs with torch.compile. They then compare against behavior cloning, PPO, FastTD3, a multitask TD-MPC2 without language/pretraining/demonstrations/BC, and 200 single-task TD-MPC2 specialists, evaluating in state-based mode unless stated otherwise.

Why This Matters

This paper challenges the assumption that online RL cannot scale in continuous control, and backs that challenge with a released benchmark, an algorithm, and 200+ checkpoints. It provides a multi-domain, multi-embodiment RL counterpart to the multi-embodiment imitation-learning datasets that currently dominate robotics. It also shows that language can serve as the task interface instead of a one-hot task index, which preserves a path to generalization on unseen tasks, and it quantifies where that generalization currently breaks (unseen nouns in instructions).

Real-world applications suggested by the paper's task domains:

• Robotic manipulation: tabletop tasks (Meta-World, ManiSkill3) and desk tasks (RoboDesk) where a single policy handles many object-interaction instructions, including from camera images where vision helped most.
• Robot locomotion and navigation: locomotion tasks (DMControl, MuJoCo) and maze navigation (OGBench) with quadruped and other embodiments.
• Game playing and simulated control: the MiniArcade arcade suite, Box2D, and Atari, relevant to agents that must act across rule sets with shared action spaces.
• Generalist controllers for heterogeneous hardware: the transfer experiments, which finetune one pretrained model to unseen tasks and embodiments with no demonstrations, mirror the practical need to deploy one model across many robots and tasks.

Industry relevance: the wall-time numbers (for example 4.6 days on 2x RTX 5090 for 200 tasks and 100M steps) put massively multitask RL within reach of labs with modest GPU resources, and the released demonstration dataset (32 GB), checkpoints, docker image, and asynchronous wrappers lower the engineering barrier considerably. The combination of a strong behavior-cloning baseline being competitive in some domains and RL winning in others is directly relevant to teams deciding whether to invest in imitation learning or RL pipelines.

Future Directions

• Consistent gains across all domains. Performance in MuJoCo, Box2D, and Atari lagged for every RL method, with Newt often matching simple behavior cloning. The authors call developing methods that improve more consistently across tasks an exciting future direction.
• Better handling of unseen language. Unseen instructions can sharply reduce zero-shot success (for example 0.3 versus 21.0 on pushing tasks), so making agents robust to novel task descriptions — rather than relying on seen nouns — remains open.
• Scaling with the number of tasks. The authors conjecture that a compute-optimal (model, batch) size exists for a given number of tasks, and that scaling to more training tasks will require proportionally larger models and batches than currently used.
• Extending transfer and open-loop capability. Transfer results are described as non-trivial but expected to improve as more training data becomes available, and open-loop planning failures are attributed to specific causes (drifting dynamics, failure to decelerate, stochastic elements) that point to more targeted modeling work.

Target Audience

Researchers and engineers working on reinforcement learning, model-based RL, world models, and multitask or multi-embodiment control. It is most useful for those who want a benchmark and infrastructure to prototype massively multitask RL, for groups studying demonstration-assisted online RL, and for practitioners interested in language-conditioned control policies with vision. Readers without RL background can follow the benchmark and results sections, but the method section assumes familiarity with value-based model-based RL.

Authors’ abstract

General-purpose control demands agents that act across many tasks and embodiments, yet research on reinforcement learning (RL) for continuous control remains dominated by single-task or offline regimes, reinforcing a view that online RL does not scale. Inspired by the foundation model recipe (large-scale pretraining followed by light RL) we ask whether a single agent can be trained on hundreds of tasks with online interaction. To accelerate research in this direction, we introduce a new benchmark with 200 diverse tasks spanning many domains and embodiments, each with language instructions, demonstrations, and optionally image observations. We then present \emph{Newt}, a language-conditioned multitask world model that is first pretrained on demonstrations to acquire task-aware representations and action priors, and then jointly optimized with online interaction across all tasks. Experiments show that Newt yields better multitask performance and data-efficiency than a set of strong baselines, exhibits strong open-loop control, and enables rapid adaptation to unseen tasks. We release our environments, demonstrations, code for training and evaluation, as well as 200+ checkpoints.