File size: 5,784 Bytes
23e1f91 | 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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 | # Dynamics Model
```
cd RISE/dynamics/dynamics_model
# Navigate to the dynamics model directory before running the following commands
```
## Generated Samples
<p align="center">
<img src="../rise_assets/conveyor.gif" alt="Conveyor demo" width="49%" />
<img src="../rise_assets/backpack.gif" alt="Backpack demo" width="49%" />
</p>
## Data Format
The framework expects data in the LeRobot format. For optimal training performance, we strongly recommend pre-resizing videos to `[256, 192]` resolution.
### Directory Structure
All tasks should be organized in the `dataset` directory with the following structure:
```bash
# copy your dataset under the dataset directory
cp -r path/to/your/dataset dataset/
```
Each dataset is organized as follows:
```
task_A/
├── data/
│ └── chunk-000/
│ ├── episode_000000.parquet
│ ├── episode_000001.parquet
│ ├── episode_000002.parquet
│ └── ...
├── meta/
│ ├── info.json
│ ├── episodes.jsonl
│ ├── episodes_stats.jsonl
│ └── tasks.jsonl
└── videos/
└── chunk-000/
└── [video files]
```
### Video Preprocessing
The `preprocess.sh` script resizes all videos in the dataset to `256x192` resolution using ffmpeg, preserving aspect ratio with center padding. Processed videos are saved in `videos_small` while maintaining the original directory structure.
**Usage:**
```bash
# Process specific datasets
./preprocess.sh dataset1 [dataset2](optional)
```
The output would be as follows with **videos_small**:
```
task_A/
├── data/
├── meta/
└── videos/
└── videos_small/
│ └── chunk-000/
│ └── [video files]
```
## Model Checkpoints
### Base LTX Backbone
Download the LTX-Video backbone components (Text Encoder, Tokenizer, and VAE) using the provided script:
```bash
./download.sh
```
This script automatically downloads all required components from the [LTX-Video HuggingFace repository](https://huggingface.co/Lightricks/LTX-Video) to the `checkpoints` directory.
Alternatively, you can manually download the following components:
1. **Text Encoder**: [text_encoder](https://huggingface.co/Lightricks/LTX-Video/tree/main/text_encoder)
2. **Tokenizer**: [tokenizer](https://huggingface.co/Lightricks/LTX-Video/tree/main/tokenizer)
3. **VAE**: [vae](https://huggingface.co/Lightricks/LTX-Video/tree/main/vae)
4. **Pre-trained dynamics model**: [dynamics_model](https://huggingface.co/OpenDriveLab-org/RISE_Assets/tree/main/dynamics_model/pretrained), pretrained on Galaxea Open World and AgiBot World Alpha jointly.
Place all downloaded weights in the same directory and update the `pretrained_model_name_or_path` field in your configuration file.
## Training
### Pre-training
Pre-training is performed on large-scale robotic datasets to learn general dynamics priors. We utilize the following datasets:
- **Galaxea Open World Dataset**: [Galaxea-Open-World-Dataset](https://huggingface.co/datasets/OpenGalaxea/Galaxea-Open-World-Dataset/tree/main)
- **AgiBot World Alpha**: [AgiBotWorld-Alpha](https://huggingface.co/datasets/agibot-world/AgiBotWorld-Alpha)
#### Steps
1. **Prepare Data**: Convert your datasets to the LeRobot format as described above.
2. **Configure Training**: Edit `configs/ltx_model/pretrain.yaml` according to the comments:
- Set `pretrained_model_name_or_path` to your LTX backbone checkpoint directory
- Set `diffusion_model.model_path` to your pre-trained diffusion checkpoint
- Configure `data.train.data_roots` and `data.val.data_roots` to point to your dataset directories
3. **Launch Training**:
```bash
bash train_task_centric.sh
```
### Fine-tuning
Fine-tuning adapts the pre-trained model to specific task domains using domain-specific datasets.
#### Steps
1. **Prepare Task-Specific Data**: Organize your fine-tuning dataset in the LeRobot format.
2. **Compute Action Normalization Statistics**: Use `norm.py` to compute and save normalization statistics:
```bash
python norm.py --datasets <your_finetune_dataset> --save-config data/utils/action_norm.json
```
This automatically computes min and max values for each dataset and saves them to a JSON configuration file.
3. **Configure Fine-tuning**: Edit `configs/ltx_model/finetune.yaml`:
- Set `pretrained_model_name_or_path` to your LTX backbone checkpoint directory
- Set `diffusion_model.model_path` to your diffusion checkpoint
- Configure `data.train.data_roots` and `data.val.data_roots` for your fine-tuning dataset
- Add `norm_config_path: data/utils/action_norm.json` to both `data.train` and `data.val` sections
- The data loader will automatically use the normalization values from the config file based on dataset names
4. **Launch Fine-tuning**:
```bash
bash task_finetune.sh
```
## Inference
The inference pipeline generates future video sequences conditioned on initial observations and action sequences.
### Steps
1. **Configure Inference**: Edit `configs/ltx_model/infer.yaml`:
- Set `pretrained_model_name_or_path` your LTX backbone checkpoint directory
- Set `diffusion_model.model_path` to your diffusion checkpoint
2. **Update Inference Script**: Edit `infer.sh` with appropriate paths
3. **Run Inference**:
```bash
bash infer.sh
```
### Inference Parameters
- `--config_file`: Path to inference configuration file
- `--image_root`: Directory containing input observation images
- `--output_path`: Directory to save generated videos
- `--act_tokens_path`: Path to action token file (`.pt` format)
- `--norm_constant`: Normalization constant for action tokens (e.g., `FINETUNE_TASK`)
|