MPC（Model Predictive Control）代码逻辑：实时追踪物体

import isaacsim
import torch
import os
import numpy as np
from omni.isaac.kit import SimulationApp

simulation_app = SimulationApp(
    {
        "headless": False,
        "width": "1920",
        "height": "1080",
    }
)
## 这里要注意的是，先创建 SimulationApp 后才能导入下面的库，因为下面的库属于它的扩展
import carb
from helper import add_extensions, add_robot_to_scene
from omni.isaac.core import World
from omni.isaac.core.objects import cuboid
from omni.isaac.core.utils.types import ArticulationAction
from omni.isaac.core.utils.numpy.rotations import euler_angles_to_quats
from pxr import PhysxSchema
## CuRobo
from curobo.util.logger import setup_curobo_logger
from curobo.util.usd_helper import UsdHelper
from curobo.geom.sdf.world import CollisionCheckerType
from curobo.geom.types import WorldConfig
from curobo.rollout.rollout_base import Goal
from curobo.types.base import TensorDeviceType
from curobo.types.math import Pose
from curobo.types.robot import JointState
from curobo.types.state import JointState
from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
from curobo.wrap.reacher.mpc import MpcSolver, MpcSolverConfig

my_world = World(stage_units_in_meters=1.0)
stage = my_world.stage
## 下面设置stage的Prim，stage可以理解为集合，而prim就是其中的元素
xform = stage.DefinePrim("/World", "Xform")	# prim包含路径和类型
stage.SetDefaultPrim(xform)
stage.DefinePrim("/curobo", "Xform")
## 最后设定一个默认地面
my_world.scene.add_default_ground_plane()

target = cuboid.VisualCuboid(	# VisualCuboid是无实体有形状的立方体
	"/World/target",	# 该target的路径
    position=np.array([0.5, 0, 0.7]),	# 在空间的位置
    orientation=np.array([0, 1, 0, 0]),	# 物体的方向（单位是角度，可以理解为绕着不同的中心轴旋转）
    color=np.array([1.0, 0, 0]),	# 颜色 RGB
    size=0.05,	# 尺寸（scale*size就是长宽高在 world 中的真实值）
    scale=np.array([1.0, 1.0, 1.0]),	# 长宽高的比例
)

dy_cuboid1 = cuboid.DynamicCuboid(	# DynamicCuboid是有实体且可以移动的立方体
    prim_path="/World/dy_cuboid1",
    position=np.array([0.7, 0.0, 0.5]),
    scale=np.array([0.5, 0.5, 0.5]),
    color=np.array([1.0, 1.0, 1.0]),
    size=1,
    linear_velocity=np.array([0.0, 1.0, 0.0]),	# 这里添加了一个线速度
)
dy_prim = stage.GetPrimAtPath("/World/dy_cuboid1")
## 由于特殊需要，将物体重力设置为无
dy = PhysxSchema.PhysxRigidBodyAPI.Apply(dy_prim)
dy.GetDisableGravityAttr().Set(True)

past_pose = None	# target 的过去位置
cmd_state_full = None	# 存储mpc过程中机器人的关节状态
n_obstacle_cuboids = 30
n_obstacle_mesh = 10
init_curobo = False
init_world = False
step = 0

setup_curobo_logger("warn")
usd_help = UsdHelper()
tensor_args = TensorDeviceType()
usd_help.load_stage(my_world.stage)
add_extensions(simulation_app, None)

## 设定 robot 配置
robot_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"]	# 这里导入 franka 的配置
j_names = robot_cfg["kinematics"]["cspace"]["joint_names"]	# 各关节名称
default_config = robot_cfg["kinematics"]["cspace"]["retract_config"]	# 机器人的初始状态
robot_cfg["kinematics"]["collision_sphere_buffer"] += 0.02 #这里增大一点机器人的碰撞检测范围

## 将 robot 配置导入到 world 中
robot, robot_prim_path = add_robot_to_scene(robot_cfg, my_world)
articulation_controller = robot.get_articulation_controller()	# 获取机器人的关节控制器

## 这里导入一个桌子，world_cfg_table 使用它的 cuboid
world_cfg_table = WorldConfig.from_dict(
    load_yaml(join_path(get_world_configs_path(), "collision_table.yml"))
)
world_cfg_table.cuboid[0].pose[2] -= 0.04
## world_cfg_mesh 使用它的 mesh（mesh是碰撞网格，在碰撞检测时通常计算网格化的环境和物体）
world_cfg_mesh = WorldConfig.from_dict(
    load_yaml(join_path(get_world_configs_path(), "collision_table.yml"))
).get_mesh_world()	# 这里获取环境的网格化
world_cfg_mesh.mesh[0].name += "_mesh"
world_cfg_mesh.mesh[0].pose[2] = -10.0

world_cfg = WorldConfig(cuboid=world_cfg_table.cuboid, mesh=world_cfg_mesh.mesh)

## 设置 mpc 的参数
mpc_config = MpcSolverConfig.load_from_robot_config(
    robot_cfg,
    world_cfg,
    use_cuda_graph=True,
    use_cuda_graph_metrics=True,
    use_cuda_graph_full_step=False,
    self_collision_check=True,
    collision_checker_type=CollisionCheckerType.MESH,
    collision_cache={"obb": n_obstacle_cuboids, "mesh": n_obstacle_mesh},
    use_mppi=True,
    use_lbfgs=False,
    use_es=False,
    store_rollouts=True,
    step_dt=0.02,
)
## 实例化 mpc
mpc = MpcSolver(mpc_config)
retract_cfg = mpc.rollout_fn.dynamics_model.retract_config.clone().unsqueeze(0)	# 初始状态
joint_names = mpc.rollout_fn.joint_names	# 关节名称
current_state = JointState.from_position(retract_cfg, joint_names=joint_names)	# 设置当前状态（初始化）
state = mpc.rollout_fn.compute_kinematics(current_state)	# 根据当前状态，计算末端执行器姿态
retract_pose = Pose(state.ee_pos_seq, quaternion=state.ee_quat_seq)	# 将计算结果记录成初始姿态
## 设置 goal
goal = Goal(
    current_state=current_state,
    goal_state=JointState.from_position(retract_cfg, joint_names=joint_names),
    goal_pose=retract_pose,
)
goal_buffer = mpc.setup_solve_single(goal, 1)
mpc.update_goal(goal_buffer)
mpc_result = mpc.step(current_state, max_attempts=2)

while simulation_app.is_running():
    if not init_world:	# 初始试运行 10 steps
        for _ in range(10):
            my_world.step(render=True)
        init_world = True
    my_world.step(render=True)
    if not my_world.is_playing():	# 手动点击按钮运行
        continue
    step_index = my_world.current_time_step_index
    if step_index <= 2:	# 初始化机器人
        my_world.reset()
        idx_list = [robot.get_dof_index(x) for x in j_names]
        robot.set_joint_positions(default_config, idx_list)
        robot._articulation_view.set_max_efforts(
            values=np.array([5000 for i in range(len(idx_list))]), joint_indices=idx_list
        )
    if not init_curobo:
        init_curobo = True
    step += 1

    step_index = step
    if step_index % 1 == 0:		# 每个step都更新
        print("Updating world")
        obstacles = usd_help.get_obstacles_from_stage(
            ignore_substring=[	# 不被视作障碍物的目录
                robot_prim_path,
                "/World/target",
                "/World/defaultGroundPlane",
                "/curobo",
            ],
            reference_prim_path=robot_prim_path,
        ).get_collision_check_world()
        obstacles.add_obstacle(world_cfg_table.cuboid[0])
        mpc.world_coll_checker.load_collision_model(obstacles)	# mpc加载障碍物
    cube_position, cube_orientation = target.get_world_pose()	# 获取 target 位置
    dy_cube_position, dy_cube_orientation = dy_cuboid1.get_world_pose()
    if past_pose is None:
        past_pose = cube_position + 1.0

    if np.linalg.norm(cube_position - past_pose) > 1e-3:
        ee_translation_goal = cube_position
        ee_orientation_teleop_goal = euler_angles_to_quats([np.pi, -0.5*np.pi, 0])
        ik_goal = Pose(
            position=tensor_args.to_device(ee_translation_goal),
            quaternion=tensor_args.to_device(ee_orientation_teleop_goal),
        )
        goal_buffer.goal_pose.copy_(ik_goal)	# 将目标放入mpc的 goal_buffer 中
        mpc.update_goal(goal_buffer)
        past_pose = cube_position

    sim_js = robot.get_joints_state()
    sim_js_names = robot.dof_names
    cu_js = JointState(
        position=tensor_args.to_device(sim_js.positions),
        velocity=tensor_args.to_device(sim_js.velocities) * 0.0,
        acceleration=tensor_args.to_device(sim_js.velocities) * 0.0,
        jerk=tensor_args.to_device(sim_js.velocities) * 0.0,
        joint_names=sim_js_names,
    )
    cu_js = cu_js.get_ordered_joint_state(mpc.rollout_fn.joint_names)	# 只控制必要的关节
    if cmd_state_full is None:
        current_state.copy_(cu_js)
    else:
        current_state_partial = cmd_state_full.get_ordered_joint_state(mpc.rollout_fn.joint_names)
        current_state.copy_(current_state_partial)
        current_state.joint_names = current_state_partial.joint_names
    current_state.copy_(cu_js)
    mpc_result = mpc.step(current_state, max_attempts=2)	# 计算结果
    ## 更新 cmd_state_full
    cmd_state_full = mpc_result.js_action
    idx_list = []
    common_js_names = []
    for x in sim_js_names:
        if x in cmd_state_full.joint_names:
            idx_list.append(robot.get_dof_index(x))
            common_js_names.append(x)
    cmd_state = cmd_state_full.get_ordered_joint_state(common_js_names)
    cmd_state_full = cmd_state
    ## 末端控制器执行
    art_action = ArticulationAction(
        joint_positions=cmd_state.position.cpu().numpy(),
        joint_indices=idx_list,
    )
    for _ in range(3):
        articulation_controller.apply_action(art_action)
simulation_app.close()	# 最后进行close（这里还是需要手动点击按钮停止，因为一旦开始运行就不会退出循环）