⚡ Representation Learning Enables Scalable Multitask Deep Reinforcement Learning

Official code repository for the paper:

📄 Representation Learning Enables Scalable Multitask Deep Reinforcement Learning

Johan Obando-Ceron, Lu Li, Scott Fujimoto, Pierre-Luc Bacon, Aaron Courville, Pablo Samuel Castro

This repository builds on MMBench & Newt — Learning Massively Multitask World Models for Continuous Control.

---

Running agents

Available task sets and their sizes:

Task set	# Tasks	Description
dmcontrol	21	DeepMind Control Suite
dmcontrol-ext	16	Extended DMControl tasks
metaworld	49	Meta-World manipulation tasks
maniskill	36	ManiSkill3 tasks
mujoco	6	MuJoCo locomotion tasks
box2d	8	Box2D tasks
robodesk	6	RoboDesk tasks
ogbench	12	OGBench goal-conditioned tasks
pygame	19	PyGame tasks
atari	27	Atari 2600 games
soup	200	All of the above combined

agent=mrq

python tdmpc2/train.py use_demos=false agent=mrq task=<TASK_SET> seed=<SEED> steps=10000000 horizon=5

agent=tdmpc2

python tdmpc2/train.py use_demos=false agent=tdmpc2 task=<TASK_SET> seed=<SEED> steps=10000000

Model sizes

Size	Parameters	enc_dim	mlp_dim	latent_dim	num_enc_layers	num_q
S	2M	128	256	384	2	3
B	5M	256	512	512	2	5
L	20M	1024	1024	512	3	5
XL	80M	2048	2048	704	4	7

Running on 200 tasks (soup)

To train on the full soup set of 200 tasks spanning all domains:

# Single GPU
python tdmpc2/train.py agent=tdmpc2 task=soup seed=1 steps=100000000 model_size=L

# Multi-GPU (uses all visible GPUs automatically)
python tdmpc2/train.py agent=tdmpc2 task=soup seed=1 steps=100000000 model_size=L multiproc=true

Note: Number of tasks must be divisible by the number of GPUs when using multiproc=true. The soup set contains 200 tasks, so 1, 2, 4, 5, 8, 10, 20, 25, 40, 50, or 200 GPUs are valid configurations.

Few-shot finetuning on held-out tasks

We evaluate generalization by finetuning on 34 held-out tasks that are excluded from the soup training set. These span variants of DMControl, ManiSkill3, OGBench, and PyGame:

34 held-out tasks

cartpole-balance-two-poles-sparse, cartpole-balance-long-sparse, walker-stand-incline, walker-walk-incline, walker-run-incline, walker-arabesque, walker-lie-down, walker-legs-up, walker-headstand, spinner-spin-four, spinner-spin-backward-four, spinner-jump-four, ms-push-apple, ms-push-pear, ms-push-rubiks-cube, ms-push-can, ms-push-sponge, ms-push-banana, ms-push-screwdriver, ms-pick-rubiks-cube, ms-pick-cup, ms-pick-golf-ball, ms-pick-soccer-ball, og-point-var1, og-point-var2, pygame-point-maze-var4, pygame-point-maze-var5, pygame-point-maze-var6, pygame-point-maze-var7, pygame-point-maze-var8, pygame-reacher-easy, pygame-reacher-hard, pygame-reacher-var1, pygame-reacher-var2

To finetune a pretrained soup checkpoint on a held-out task:

python tdmpc2/train.py agent=tdmpc2 task=<HELD_OUT_TASK> seed=1 \
    checkpoint=<path/to/soup/checkpoint.pt> finetune=true steps=1000000

---

Configuration

All hyperparameters are defined in tdmpc2/config.py and can be overridden via command-line using Hydra syntax (key=value). Commonly useful options:

• seed — random seed for reproducibility
• exp_name — experiment name used for logging and checkpoint directories
• enable_wandb / wandb_project / wandb_entity — W&B logging settings
• compile — enable/disable torch.compile (default: true)
• save_video — save evaluation rollout videos (requires env_mode=sync)
• model_size — model size preset: S (2M), B (5M), L (20M), XL (80M)

---

Citation

If you find this work useful, please cite:

@misc{obandoceron2026representationlearningenablesscalable,
      title={Representation Learning Enables Scalable Multitask Deep Reinforcement Learning},
      author={Johan Obando-Ceron and Lu Li and Scott Fujimoto and Pierre-Luc Bacon and Aaron Courville and Pablo Samuel Castro},
      year={2026},
      eprint={2606.05555},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2606.05555},
}

If you use the MMBench environment suite or the Newt baseline, please also cite:

@misc{Hansen2025Newt,
      title={Learning Massively Multitask World Models for Continuous Control},
      author={Nicklas Hansen and Hao Su and Xiaolong Wang},
      year={2025},
      eprint={2511.19584},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2511.19584},
}

---

Contributing

We welcome contributions focused on improving sample efficiency and wall clock time performance in multitask reinforcement learning. If you have a proposal for a more efficient training component or want to add support for a new task, please open an issue or pull request. We are particularly interested in contributions that push the boundaries of scalability across diverse task sets.

---

License

This project is licensed under the MIT License - see the LICENSE file for details. Note that the repository relies on third-party code, which is subject to their respective licenses.