| # Diffusion Policy Checkpoint - Step 100000 |
|
|
| 这是第100000步的diffusion policy训练checkpoint,用于机器人视觉运动控制。 |
|
|
| ## 模型信息 |
|
|
| - **模型类型**: Diffusion Policy |
| - **训练步数**: 100,000 |
| - **视觉骨干网络**: resnet18 |
| - **输入观察步数**: 2 |
| - **动作步数**: 8 |
| - **时间范围**: 16 |
| - **扩散步数**: 100 |
|
|
| ## 输入特征 |
|
|
| - **状态观察**: [8] 维 |
| - **头部相机图像**: [3, 256, 256] (RGB, 256x256) |
|
|
| ## 输出特征 |
|
|
| - **动作**: [7] 维 |
|
|
| ## 训练配置 |
|
|
| - **批次大小**: 16 |
| - **学习率**: 0.0001 |
| - **优化器**: adam |
| - **数据集**: /home/shuo/research/datasets/steer_test_lerobot |
|
|
| ## 文件结构 |
|
|
| ``` |
| 100000/ |
| ├── pretrained_model/ |
| │ ├── model.safetensors # 模型权重 (~1005MB) |
| │ ├── config.json # 模型配置 |
| │ └── train_config.json # 训练配置 |
| └── training_state/ |
| ├── optimizer_state.safetensors # 优化器状态 (~2GB) |
| ├── scheduler_state.json # 学习率调度器状态 |
| ├── optimizer_param_groups.json # 优化器参数组 |
| ├── rng_state.safetensors # 随机数生成器状态 |
| └── training_step.json # 训练步数信息 |
| ``` |
|
|
| ## 使用方法 |
|
|
| ### 加载模型 |
|
|
| ```python |
| import torch |
| from lerobot.common.policies.diffusion import DiffusionPolicy |
| |
| # 加载配置 |
| config_path = "pretrained_model/config.json" |
| with open(config_path, 'r') as f: |
| config = json.load(f) |
| |
| # 创建模型 |
| policy = DiffusionPolicy(config) |
| |
| # 加载权重 |
| checkpoint = torch.load("pretrained_model/model.safetensors", map_location='cpu') |
| policy.load_state_dict(checkpoint) |
| ``` |
|
|
| ### 恢复训练 |
|
|
| ```python |
| # 加载训练状态 |
| training_state = torch.load("training_state/optimizer_state.safetensors") |
| optimizer.load_state_dict(training_state) |
| |
| # 继续训练 |
| # ... 训练代码 |
| ``` |
|
|
| ## 数据集 |
|
|
| 此模型在rainbow_real数据集上训练,包含机器人操作任务。 |
| |
| ## 许可证 |
| |
| 请参考原始项目的许可证。 |
| |