Feature/add multitask diffusion transformer policy implementation

[brysonjones]

What this does

This PR adds support for an implementation of Multitask Diffusion Transformer Policy, which was shown in a demo of Boston Dynamics' Atlas robot performing whole-body manipulation tasks

I wanted to dive into the research of this method, and build an open-source implementation for the community to leverage and build from.

I will simultaneously be releasing a blog post that includes the details of this work as this gets merged in! (Note: the current blog post link is broken until being released)

How it was tested

• I added a simple test script test_multi_task_dit_policy.py to validate any install and import errors
• I've tested launching training runs on Modal with batch sizes of 320 on H200 and B200 GPUs
• I've trained multiple policies for >60k training steps and >20 hours with this policy, and deployed them onto a Trossen WidowAI Arm (for which I have another PR open for to integrate)

How to checkout & try? (for the reviewer)

Run the test script:

uv run pytest tests/policies/test_multi_task_dit_policy.py -v

Train a policy:

lerobot-train \
  --policy.type=multi_task_dit \
  --dataset.repo_id={{your_dataset_name}} \
  --dataset.root={{your/dataset/path}} \
  --output_dir=outputs/train/multi_task_dit \
  --job_name=multi_task_dit_training_test \
  --policy.device=cuda \
  --batch_size=16 \
  --steps=10000 \
  --save_freq=1000 \
  --wandb.enable=true \
  --policy.repo_id=YOUR_HF_USERNAME/multi_task_dit_policy_test

[s1lent4gnt]

Thanks for the PR, @brysonjones — nice work! 🙌

Here are my comments from the first review pass:

I think it would be better to keep everything in a single file, modeling_multi_task_dit.py, and remove the modules/ directory for now. It will be easier to maintain.

To match the original LBM paper, we should remove DINOv3 and stick to the CLIP-based setup. However, I’m fine with keeping the flow-matching objective since it’s used in the Boston Dynamics blog post. Maybe we can add a short comment in the code to clarify this difference from the paper.

We can probably simplify things by using a single config for everything instead of multiple configs.

[s1lent4gnt]

Comments

[brysonjones]

Thanks for the PR, @brysonjones — nice work! 🙌

Here are my comments from the first review pass:

I think it would be better to keep everything in a single file, modeling_multi_task_dit.py, and remove the modules/ directory for now. It will be easier to maintain.

To match the original LBM paper, we should remove DINOv3 and stick to the CLIP-based setup. However, I’m fine with keeping the flow-matching objective since it’s used in the Boston Dynamics blog post. Maybe we can add a short comment in the code to clarify this difference from the paper.

We can probably simplify things by using a single config for everything instead of multiple configs.

Hey @s1lent4gnt!

Thanks for the review and feedback. I think all of this is reasonable and makes sense on my side. Will work on on the updates and push them through soon 👍

[brysonjones]

@s1lent4gnt I think all these points should be addressed now, let me know what you think!

[brysonjones]

@s1lent4gnt Just worked through moving the tokenization of the task to the pre-processor. I think this is cleaner as well!

[s1lent4gnt]

Great work @brysonjones !
LGTM!

[brysonjones]

Thank you for the review and help getting this ready, @s1lent4gnt!

Everything is good from my side, let me know if there's any other adjustments to make before merging in 👍

[brysonjones]

Seems like there were a few missing parts where the tests failed (missing additions to docs toc, transformers lib conditional import)

Have updated and pushed those through