第十四讲、Isaaclab在USD环境中进行策略推断-CSDN博客

本文链接：https://blog.csdn.net/m0_47719040/article/details/147401036

0 前言

官方教程：https://isaac-sim.github.io/IsaacLab/main/source/tutorials/03_envs/policy_inference_in_usd.html
Isaacsim+Isaaclab安装：https://blog.csdn.net/m0_47719040/article/details/146389391?spm=1001.2014.3001.5502

上一讲中我们学习了如何修改环境（主要修改环境中的智能体及其配置参数），本讲将介绍如何将训练好的策略，应用到新的环境（USD）中。

教程对应的脚本为policy_inference_in_usd.py在scripts/tutorials/03_envs目录下。

1 代码执行

STEP1：训练模型

进入安装 isaac lab 时创建的conda虚拟环境
在该环境下进入 isaac sim文件夹中运行source setup_conda_env.sh
终端中输入./isaaclab.sh -p scripts/reinforcement_learning/rsl_rl/train.py --task Isaac-Velocity-Rough-H1-v0 --headless运行你的代码，训练模型。

会在IsaacLab文件夹下生成一个logs文件夹IsaacLab->logs->rsl_rl->h1_rough->对应你的训练时间的文件夹
在这里插入图片描述

STEP2：导出policy
同样在终端输入下述代码运行，会在上面说的文件夹下生成一个exported文件夹,

python scripts/reinforcement_learning/rsl_rl/play.py --task Isaac-Velocity-Rough-H1-v0 --num_envs 64 --checkpoint logs/rsl_rl/h1_rough/EXPERIMENT_NAME(要换成你自己的地址)/model_2999.pt

在这里插入图片描述

STEP3：在仓库场景下使用H1机器人进行推理

./isaaclab.sh -p scripts/tutorials/03_envs/policy_inference_in_usd.py --checkpoint logs/rsl_rl/h1_rough/EXPERIMENT_NAME(要换成你自己的地址)/exported/policy.pt

在这里插入图片描述

2 代码详解

1、scripts/tutorials/03_envs/policy_inference_in_usd.py代码详解。

# 导入命令行参数解析模块
import argparse

# 导入Isaac Lab应用启动器
from isaaclab.app import AppLauncher

# 创建命令行参数解析器，添加描述
parser = argparse.ArgumentParser(description="Tutorial on inferencing a policy on an H1 robot in a warehouse.")
# 添加必需的--checkpoint参数，用于指定模型检查点路径
parser.add_argument("--checkpoint", type=str, help="Path to model checkpoint exported as jit.", required=True)

# 将AppLauncher的命令行参数添加到解析器
AppLauncher.add_app_launcher_args(parser)
# 解析命令行参数
args_cli = parser.parse_args()

# 创建应用启动器实例并启动Omniverse应用
app_launcher = AppLauncher(args_cli)
simulation_app = app_launcher.app

"""以下代码在Omniverse应用启动后运行"""
# 导入后续需要的模块
import io  # 处理二进制流
import os  # 文件路径操作
import torch  # PyTorch深度学习框架

import omni  # Omniverse核心库

# 导入Isaac Lab相关模块
from isaaclab.envs import ManagerBasedRLEnv
from isaaclab.terrains import TerrainImporterCfg
from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR  # 获取Isaac Sim资源目录

# 导入H1机器人在粗糙地形环境的配置类（该配置类会在后续做更详细的解释）
from isaaclab_tasks.manager_based.locomotion.velocity.config.h1.rough_env_cfg import H1RoughEnvCfg_PLAY


def main():
    """主函数"""
    # 加载训练好的JIT策略模型
    policy_path = os.path.abspath(args_cli.checkpoint)  # 获取绝对路径
    # 使用Omniverse客户端读取文件内容（返回元组的第三个元素是文件内容）
    file_content = omni.client.read_file(policy_path)[2]
    # 将字节内容转换为文件流对象
    file = io.BytesIO(memoryview(file_content).tobytes())
    # 加载Torch JIT模型
    policy = torch.jit.load(file)
    
    # 初始化环境配置
    env_cfg = H1RoughEnvCfg_PLAY()
    env_cfg.scene.num_envs = 1  # 设置环境实例数量为1
    env_cfg.curriculum = None    # 禁用课程学习
    
    # 配置地形导入参数
    env_cfg.scene.terrain = TerrainImporterCfg(
        prim_path="/World/ground",  # USD场景中的根路径
        terrain_type="usd",         # 使用预制的USD地形
        usd_path=f"{ISAAC_NUCLEUS_DIR}/Environments/Simple_Warehouse/warehouse.usd"  # 仓库场景路径
    )
    
    # 设置模拟参数
    env_cfg.sim.device = "cpu"   # 使用CPU进行计算
    env_cfg.sim.use_fabric = False  # 禁用Fabric加速
    
    # 创建强化学习环境
    env = ManagerBasedRLEnv(cfg=env_cfg)
    # 重置环境获取初始观测
    obs, _ = env.reset()
    
    # 主循环：当应用运行时持续执行
    while simulation_app.is_running():
        # 使用策略网络生成动作（只使用"policy"观测部分）
        action = policy(obs["policy"])  
        # 执行动作并获取新的观测
        obs, _, _, _, _ = env.step(action)


if __name__ == "__main__":
    main()              # 运行主函数
    simulation_app.close()  # 关闭模拟应用

2、H1RoughEnvCfg_PLAY类详解
该类位于/IsaacLab/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/h1/rough_env_cfg.py文件中。

@configclass  # 使用配置类装饰器，通常用于自动生成配置参数或提供类型检查等特性
class H1RoughEnvCfg_PLAY(H1RoughEnvCfg):  # 继承自H1RoughEnvCfg的Play模式配置类
    def __post_init__(self):  # 后初始化方法，用于覆盖父类配置
        # post init of parent
        super().__post_init__()  # 首先调用父类的后初始化方法

        # make a smaller scene for play
        self.scene.num_envs = 50  # 将并行环境数量从父类默认值减少到50（节省计算资源）
        self.scene.env_spacing = 2.5  # 环境之间的间隔设为2.5米（防止碰撞）
        self.episode_length_s = 40.0  # 每个episode最长持续时间设为40秒
        
        # spawn the robot randomly in the grid (instead of their terrain levels)
        self.scene.terrain.max_init_terrain_level = None  # 取消地形等级限制，允许随机生成
        
        # reduce the number of terrains to save memory
        if self.scene.terrain.terrain_generator is not None:  # 如果存在地形生成器
            self.scene.terrain.terrain_generator.num_rows = 5  # 地形网格行数设为5（减少内存占用）
            self.scene.terrain.terrain_generator.num_cols = 5  # 地形网格列数设为5
            self.scene.terrain.terrain_generator.curriculum = False  # 禁用课程学习模式

        # 设置基础速度指令范围（固定化参数便于控制）
        self.commands.base_velocity.ranges.lin_vel_x = (1.0, 1.0)  # X轴线性速度固定为1.0m/s
        self.commands.base_velocity.ranges.lin_vel_y = (0.0, 0.0)  # Y轴线性速度固定为0（禁止横向移动）
        self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0)  # Z轴角速度范围[-1,1]rad/s（允许旋转）
        self.commands.base_velocity.ranges.heading = (0.0, 0.0)  # 朝向角固定为0（保持初始方向）
        
        # disable randomization for play
        self.observations.policy.enable_corruption = False  # 关闭观测噪声（获得确定性输入）
        
        # remove random pushing
        self.events.base_external_force_torque = None  # 禁用外部力/扭矩干扰事件
        self.events.push_robot = None  # 禁用随机推动机器人事件

所以实际上我们可以通过obs中的速度指令控制机器人运动，比如我们可以将主循环的部分做下述修改后可以发现机器人一直做直线运动。

    # 主循环：当应用运行时持续执行
    while simulation_app.is_running():
    	obs["policy"][0][3 * (4 - 1) + 2] = 0
        # 使用策略网络生成动作（只使用"policy"观测部分）
        action = policy(obs["policy"])  
        # 执行动作并获取新的观测
        obs, _, _, _, _ = env.step(action)