From ba968e84f13f5edf79d6a038f997d050b7c14829 Mon Sep 17 00:00:00 2001
From: Bryson Jones <bkjones97@gmail.com>
Date: Wed, 10 Dec 2025 16:26:57 -0800
Subject: [PATCH] fix bugs when testing on hardware

---
 docs/source/multitask_dit.mdx                 | 21 +++++++++----------
 .../configuration_multi_task_dit.py           | 11 +++++++---
 .../multi_task_dit/modeling_multi_task_dit.py |  2 +-
 3 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/docs/source/multitask_dit.mdx b/docs/source/multitask_dit.mdx
index 960b1c740..824298002 100644
--- a/docs/source/multitask_dit.mdx
+++ b/docs/source/multitask_dit.mdx
@@ -40,17 +40,16 @@ Here's a complete training command for training Multi-Task DiT on your dataset:
 
 ```bash
 lerobot-train \
-  --dataset.repo_id=$DATASET_ID \
-  --output_dir=$OUTPUT_DIR \
-  --job_name=$JOB_NAME \
+  --dataset.repo_id={{MY_DATASET_ID}} \
+  --output_dir={{MY_OUTPUT_DIR}} \
   --policy.type=multi_task_dit \
   --policy.device=cuda \
+  --policy.repo_id={{MY_REPO_ID}}
   --batch_size=32 \
   --steps=5000 \
   --save_freq=500 \
   --log_freq=100 \
-  --wandb.enable=true \
-  --policy.repo_id=$REPO_ID
+  --wandb.enable=true
 ```
 
 ### Recommended Hyperparameters and Dataset Details (30Hz Control Frequency)
@@ -59,15 +58,15 @@ For reliable performance, start with these suggested default hyperparameters:
 
 ```bash
 lerobot-train \
-  --dataset.repo_id=$DATASET_ID \
-  --output_dir=$OUTPUT_DIR \
-  --job_name=$JOB_NAME \
+  --dataset.repo_id={{MY_DATASET_ID}} \
+  --output_dir={{MY_OUTPUT_DIR}} \
   --policy.type=multi_task_dit \
   --policy.device=cuda \
   --batch_size=320 \
   --steps=30000 \
   --policy.horizon=32 \
   --policy.n_action_steps=24 \
+  --policy.repo_id={{MY_REPO_ID}} \
   --policy.objective=diffusion \
   --policy.noise_scheduler_type=DDPM \
   --policy.num_train_timesteps=100 \
@@ -263,8 +262,8 @@ Here's a complete example training on a custom dataset:
 
 ```bash
 lerobot-train \
-  --dataset.repo_id=your_username/your_dataset \
-  --output_dir=outputs/multitask_dit_training \
+  --dataset.repo_id={{MY_DATASET_ID}} \
+  --output_dir={{MY_OUTPUT_DIR}} \
   --policy.type=multi_task_dit \
   --policy.device=cuda \
   --batch_size=320 \
@@ -283,7 +282,7 @@ lerobot-train \
   --policy.image_crop_shape=[224,224] \
   --wandb.enable=true \
   --wandb.project=multitask_dit \
-  --policy.repo_id=your_username/multitask_dit_policy
+  --policy.repo_id={{MY_REPO_ID}}
 ```
 
 ## References
diff --git a/src/lerobot/policies/multi_task_dit/configuration_multi_task_dit.py b/src/lerobot/policies/multi_task_dit/configuration_multi_task_dit.py
index dc7708551..7a6605692 100644
--- a/src/lerobot/policies/multi_task_dit/configuration_multi_task_dit.py
+++ b/src/lerobot/policies/multi_task_dit/configuration_multi_task_dit.py
@@ -15,7 +15,6 @@
 # limitations under the License.
 
 from dataclasses import dataclass, field
-from typing import Literal
 
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import NormalizationMode
@@ -37,7 +36,7 @@ class MultiTaskDiTConfig(PreTrainedConfig):
     n_action_steps: int = 24  # Actions executed per policy call (~0.8s at 30Hz)
 
     # Objective Selection
-    objective: Literal["diffusion", "flow_matching"] = "diffusion"
+    objective: str = "diffusion"  # "diffusion" or "flow_matching"
 
     # --- Diffusion-specific (used when objective="diffusion") ---
     noise_scheduler_type: str = "DDPM"  # "DDPM" or "DDIM"
@@ -54,7 +53,7 @@ class MultiTaskDiTConfig(PreTrainedConfig):
     sigma_min: float = 0.0  # Minimum noise in flow interpolation path
     num_integration_steps: int = 100  # ODE integration steps at inference
     integration_method: str = "euler"  # ODE solver: "euler" or "rk4"
-    timestep_sampling_strategy: Literal["uniform", "beta"] = "beta"
+    timestep_sampling_strategy: str = "beta"  # "uniform" or "beta"
 
     timestep_sampling_s: float = 0.999  # (beta only) Max timestep threshold
     timestep_sampling_alpha: float = 1.5  # (beta only) Beta distribution alpha
@@ -112,6 +111,12 @@ class MultiTaskDiTConfig(PreTrainedConfig):
 
     def _validate(self):
         """Validate configuration parameters."""
+        # Objective validation
+        if self.objective not in ["diffusion", "flow_matching"]:
+            raise ValueError(
+                f"objective must be 'diffusion' or 'flow_matching', got '{self.objective}'"
+            )
+
         # Transformer validation
         if self.hidden_dim <= 0:
             raise ValueError("hidden_dim must be positive")
diff --git a/src/lerobot/policies/multi_task_dit/modeling_multi_task_dit.py b/src/lerobot/policies/multi_task_dit/modeling_multi_task_dit.py
index 82c8cd750..c43782b36 100644
--- a/src/lerobot/policies/multi_task_dit/modeling_multi_task_dit.py
+++ b/src/lerobot/policies/multi_task_dit/modeling_multi_task_dit.py
@@ -21,8 +21,8 @@ Supports both diffusion and flow matching objectives for action generation.
 
 References:
 - https://arxiv.org/abs/2507.05331
+- https://bostondynamics.com/blog/large-behavior-models-atlas-find-new-footing/
 - https://brysonkjones.substack.com/p/dissecting-and-open-sourcing-multitask-diffusion-transformer-policy
-- https://brysonkjones.substack.com/p/dissecting-multitask-diffusion-transformer-policy
 """
 
 import math