Skip to content

Latest commit

 

History

History
146 lines (119 loc) · 6 KB

sim_loop_tutorial.md

File metadata and controls

146 lines (119 loc) · 6 KB

环境安装

参考 环境安装,isaac-sim-仿真平台安装, 按顺序安装所需要的环境。

仿真闭环教程

本文档主要介绍仿真闭环的全流程,包括场景搭建、数据采集、模型训练到部署验证的全流程,整体工作流程如下图所示: 仿真闭环工作流程 下面以 Stack Cube 任务为例进行介绍,Stack Cube 任务的目标是将一个红色的方块叠到蓝色的方块上。

1 仿真场景搭建

这里直接使用预先搭建好的仿真环境。下载地址:Download Link , 然后放置于项目根目录下的 assets/sim_env。

下面介绍如何在 Isaac Sim 中打开场景。首先启动 isaac sim 4.5:

cd ~/isaac_sim # 进入 isaac sim 的根目录
bash isaac-sim.sh

出现 isaac sim 的 GUI 界面后,点击菜单栏的 File -> Open, 选择 FlexUMI 根目录下的 assets/sim_env/stack_cube.usd,点击确定,即可加载场景到 Isaac Sim: Stack Cube 场景

2 数据采集

在项目根目录下运行以下脚本执行数据采集操作:

执行单次 Stack Cube 数据采集任务

  bash scripts/sim/run_sim.sh --config-name stack_cube_data_collect

重复执行60次,得到 60 个 episode

  bash scripts/sim/data_collect.sh --nums 60

采集的数据会保存到项目根目录的 data_collect/ 下

3 数据后处理及质量验证

(1) 数据后处理

转化为 hdf5 格式用于 ACT 模型的训练:

  conda activate FlexUMI
  python scripts/sim/to_hdf5.py --base_folder data_collect/ --output_folder ./dataset/stack_cube/

假定原始数据位于根目录的 data_collect/ 下,转换后的数据将保存到 dataset/stack_cube/ 下

转为为 zarr 格式用于 DP 模型的训练:

  conda activate FlexUMI
  python scripts/sim/to_umi_format.py --base_folder data_collect/ --output_folder ./dataset/stack_cube_zarr.zip --num_episodes <epoch_nums>

将<num_episodes> 替换为收集的 epoch 的数量,例如 60;转换后的数据将保存为 stack_cube_zarr.zip

(2) 数据质量验证

为了在模型训练前验证采集的数据是否存在问题,可以在仿真环境里 replay 采集的轨迹,观察能否正常完成任务。

以 hdf5 格式的数据回放为例,修改 stack_cube_replay.yaml

      src_cube_cfg: 修改红色方块的位置和方向
      target_cube_cfg: 修改蓝色方块的位置和方向
      file_path: 修改需要回放的 hdf5 数据文件的路径

假设要回放 episode 0 的数据,方块的位姿在文件 data_collect/epo_0/data/cube_red 和 data_collect/epo_0/data/cube_blue 下面,里面有 7个数,前三个表示位置,后四个旋转 ,直接复制过来即可。

在 仿真里回放 episode 0 进行验证:

    bash scripts/sim/run_sim.sh --config-name stack_cube_replay.yaml

4 模型训练

(1) ACT 模型训练

先修改 src/policies/act/constants.py 中和 stack_cube 有关的参数,替换dataset_dir、num_episodes 为你的数据集路径和采集的epoch数,然后运行:

python3 src/policies/act/train.py \
--task_name stack_cube \
--ckpt_dir assets/ckpt/stack_cube \
--policy_class ACT --kl_weight 10 --chunk_size 20 \
--hidden_dim 512 --batch_size 8 --dim_feedforward 3200 \
--num_epochs 20000 --lr 1e-5 --seed 0

训练得到的模型权重会保存到根目录的 assets/ckpt/stack_cube/ 下

(2) DP 模型训练,

  python3 src/policies/dp/train.py \
    --config-name=train_diffusion_unet_timm_umi_workspace \
    task.dataset_path=./dataset/stack_cube_zarr.zip

task.dataset_path 后面替换为你数据集的路径,数据集是单个zarr.zip文件。训练得到的 DP 模型的权重会保存到 data/outputs/xx.xx.xx_train_diffusion_unet_timm_umi/checkpoints 目录下

5 部署执行

整体仿真架构如下: 仿真闭环的流程图 这里提供 Stack Cube 任务预训练的模型权重: 注意:运行 ACT 算法需要同时下载 policy_best.ckpt 和 dataset_stats.pkl

文件名 下载地址 存放路径 说明
policy_best.ckpt 下载链接 assets/ckpt/stack_cube/ 用于 stack cube 的 ACT 模型的权重
dataset_stats.pkl 下载链接 assets/ckpt/stack_cube/ 数据集相关参数
stack_cube_simv3.ckpt/ 下载链接 assets/ckpt/stack_cube/ 用于 stack cube 的 DP 模型的权重

整个闭环仿真包含三个节点,需要启动三个终端来运行,其中 ROS 节点用于通信,Isaac Sim 节点用于仿真,Policy 节点用于根据当前观测输出动作。

测试 ACT 算法

(1) 启动 ROS 节点

  roscore

(2) 启动 Policy 节点

运行 ACT 算法:

    python3 src/policies/act/infer_sim.py \
    --task_name stack_cube \
    --ckpt_dir assets/ckpt/stack_cube \
    --policy_class ACT --kl_weight 10 --chunk_size 20 --hidden_dim 512 --batch_size 8 --dim_feedforward 3200 \
    --num_epochs 2000  --lr 1e-5 \
    --seed 0 --eval --temporal_agg

(3) 启动仿真节点

启动 isaac sim,加载仿真环境

  bash scripts/sim/run_sim.sh --config-name stack_cube_act.yaml

(4) 执行效果:

stack cube act video

测试 DP 算法

(1) 启动 ROS 节点

  roscore

(2) 启动 Policy 节点

  python3 src/policies/dp/infer_sim.py --ckpt_path ./assets/ckpt/stack_cube/stack_cube_simv3.ckpt --n_action_steps 4

ckpt_path 后面替换为你的 DP 模型权重的路径

(3) 启动仿真节点

  bash scripts/sim/run_sim.sh --config-name stack_cube_dp.yaml