File size: 1,649 Bytes
b2d799b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | ---
license: apache-2.0
tags:
- robotics
- rdt
- libero
- diffusion
- transformers
---
# RDT-1B LIBERO Checkpoint
RDT-1B fine-tuned on LIBERO Spatial benchmark. Best performing checkpoint.
## Model Information
- Base Model: RDT-1B (Residual Diffusion Transformer)
- Training Framework: DeepSpeed ZeRO Stage 2
- Precision: BF16
## Checkpoint Contents
This checkpoint includes:
### For Inference
- `ema/model.safetensors` - EMA model weights (recommended for inference)
- `config.json` - Model configuration
### For Training
- `pytorch_model/` - DeepSpeed distributed training checkpoint
- `bf16_zero_pp_rank_*_optim_states.pt` - Optimizer states (ZeRO Stage 2)
- `mp_rank_00_model_states.pt` - Model states
- `scheduler.bin` - Learning rate scheduler state
- `random_states_*.pkl` - Random number generator states
- `zero_to_fp32.py` - Utility to convert DeepSpeed checkpoint to FP32
## Usage
### For Inference
```python
from transformers import AutoModel
import torch
# Load the EMA model for inference
model = AutoModel.from_pretrained(
"TJ-chen/RDT-1B-LIBERO-Spatial",
subfolder="ema",
trust_remote_code=True
)
model.eval()
```
### For Continued Training
Download the complete checkpoint and use DeepSpeed to resume training:
```bash
# The checkpoint can be loaded with DeepSpeed ZeRO Stage 2
# Make sure your training script is configured with the same DeepSpeed settings
```
## Citation
If you use this model, please cite:
```bibtex
@article{rdt2024,
title={Residual Diffusion Transformer for Robotic Manipulation},
author={Your Name},
journal={arXiv preprint},
year={2024}
}
```
## License
Apache 2.0
|