| # 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数据集上训练,包含机器人操作任务。 | |
| ## 许可证 | |
| 请参考原始项目的许可证。 | |