| # 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`) |
| |
