参考 环境安装,isaac-sim-仿真平台安装, 按顺序安装所需要的环境。
本文档主要介绍仿真闭环的全流程,包括场景搭建、数据采集、模型训练到部署验证的全流程,整体工作流程如下图所示:
下面以 Stack Cube 任务为例进行介绍,Stack Cube 任务的目标是将一个红色的方块叠到蓝色的方块上。
这里直接使用预先搭建好的仿真环境。下载地址: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 数据采集任务
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/ 下
转化为 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
为了在模型训练前验证采集的数据是否存在问题,可以在仿真环境里 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
先修改 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/ 下
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 目录下
整体仿真架构如下:
这里提供 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 节点用于根据当前观测输出动作。
(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) 执行效果:
(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
