Step-by-step process for MMBench dataset collection

[zef1611]

Hi @nicklashansen and team,

First of all, thank you for the great work on "Learning Massively Multitask World Models for Continuous Control" and for releasing the code!

I am reading through the paper and am very interested in the MMBench benchmark introduced in the work. As I am relatively new to this area, I am trying to get a better sense of how the benchmark is constructed programmatically.

Could you provide some high-level guidance and details on the step-by-step process used to collect the demonstrations for the 200 tasks.

I would love to understand how to work with the benchmark directly or potentially extend it in the future, so any instructions would be extremely helpful.

Thanks for your time!

[nicklashansen]

Hi @zef1611, I'm glad to hear that you're interested in our work!

We describe the overall process of collecting demonstrations in Appendix C (Demonstrations) of our paper, copied here for convenience:

To aid in research on pretraining + finetuning of massively multitask agents we provide an ex-
pert demonstration dataset for MMBench. This dataset contains 10-40 demonstrations per task
(depending on task difficulty) for a total of 4020 demonstrations, which are all collected by single-
task TD-MPC2 (Hansen et al., 2024) agents trained from low-dimensional state observations. Each
single-task agent is trained for 5M environment steps (Atari: 10M) to ensure that performance has
converged, although we observe that training converges earlier than that on most tasks. Training a
single-task agent for 5M environment steps takes approx. 6 hours on a single NVIDIA RTX 5090
GPU, totaling an estimated 1,300 GPU hours to train expert policies for every task in MMBench.
To ensure that only successful demonstrations are collected, we reject episodes that do not result
in success (when available) or have episode returns that are substantially (> 25%) lower than the
median for that task. Table 5 provides an overview of our demonstration dataset. For completeness,
we also release all 200 checkpoints that generated the provided demonstrations; these checkpoints
can be used to generate additional checkpoints, or be used for e.g. teacher-student policy distillation.

The exact script used to generate demonstrations can be found here: https://github.com/nicklashansen/newt/blob/main/tdmpc2/generate_demos.py

I hope that helps!