Rdt 170m

🚀 RDT-170M

RDT-170M is a 170M-parameter imitation learning Diffusion Transformer (RDT(small) in ablation). It offers a hidden size of 1024 and a depth of 14, which are half of those in RDT-1B. Given language instructions and RGB images of up to three views, RDT can predict the next 64 robot actions. It is compatible with almost all modern mobile manipulators, covering single - arm to dual - arm, joint to EEF, position to velocity, and even wheeled locomotion.

All the code, pre - trained model weights, and data are licensed under the MIT license.

For more information, please refer to our project page and paper.

✨ Features

• Versatile Compatibility: Compatible with a wide range of modern mobile manipulators, including single - arm, dual - arm, joint - based, EEF - based, position - controlled, velocity - controlled, and wheeled locomotion robots.
• Multi - Modal Input: Accepts language instructions and RGB images from up to three views to predict robot actions.
• Open - Source and Licensed: All code, pre - trained model weights, and data are available under the MIT license.

📚 Documentation

Model Details

Uses

RDT takes language instructions, RGB images (up to three views), control frequency (if any), and proprioception as input and predicts the next 64 robot actions. It supports the control of almost all robot manipulators through a unified action space, which includes all the main physical quantities of the robot manipulator (e.g., end - effector and joint, position and velocity, and wheeled locomotion). To deploy it on your robot platform, you need to fill the relevant quantities of the raw action vector into the unified space vector. For more information, see [our repository](https://github.com/thu - ml/RoboticsDiffusionTransformer).

⚠️ Important Note

Due to the embodiment gap, RDT cannot yet generalize to new robot platforms (not seen in the pre - training datasets). In this case, we recommend collecting a small dataset of the target robot and then using it to fine - tune RDT. See [our repository](https://github.com/thu - ml/RoboticsDiffusionTransformer) for a tutorial.

💻 Usage Examples

Basic Usage

# Please first clone the repository and install dependencies
# Then switch to the root directory of the repository by "cd RoboticsDiffusionTransformer"

# Import a create function from the code base
from scripts.agilex_model import create_model

# Names of cameras used for visual input
CAMERA_NAMES = ['cam_high', 'cam_right_wrist', 'cam_left_wrist']
config = {
    'episode_len': 1000,  # Max length of one episode
    'state_dim': 14,      # Dimension of the robot's state
    'chunk_size': 64,     # Number of actions to predict in one step
    'camera_names': CAMERA_NAMES,
}
pretrained_vision_encoder_name_or_path = "google/siglip-so400m-patch14-384" 
# Create the model with the specified configuration
model = create_model(
    args=config,
    dtype=torch.bfloat16, 
    pretrained_vision_encoder_name_or_path=pretrained_vision_encoder_name_or_path,
    pretrained='robotics-diffusion-transformer/rdt-1b',
    control_frequency=25,
)

# Start inference process
# Load the pre-computed language embeddings
# Refer to scripts/encode_lang.py for how to encode the language instruction
lang_embeddings_path = 'your/language/embedding/path'
text_embedding = torch.load(lang_embeddings_path)['embeddings']  
images: List(PIL.Image) = ... #  The images from last 2 frames
proprio = ... # The current robot state
# Perform inference to predict the next `chunk_size` actions
actions = policy.step(
    proprio=proprio,
    images=images,
    text_embeds=text_embedding 
)

📄 License

All the code, pre - trained model weights, and data of RDT - 170M are licensed under the MIT license.

📚 Citation

If you find our work helpful, please cite us:

@article{liu2024rdt,
  title={RDT-1B: a Diffusion Foundation Model for Bimanual Manipulation},
  author={Liu, Songming and Wu, Lingxuan and Li, Bangguo and Tan, Hengkai and Chen, Huayu and Wang, Zhengyi and Xu, Ke and Su, Hang and Zhu, Jun},
  journal={arXiv preprint arXiv:2410.07864},
  year={2024}
}

Thank you!