From f2194a2a2be992c0df7650b5b14bb16351174302 Mon Sep 17 00:00:00 2001 From: zhenai Date: Wed, 16 Sep 2026 11:21:40 +0800 Subject: [PATCH] Lv0 is good --- scripts/rsl_rl/play.py | 4 +- .../tasks/manager_based/twofootbot/Lv0.py | 124 +++++++++++++++++- .../tasks/manager_based/twofootbot/Lv1.py | 114 ++++++++++++++++ .../tasks/manager_based/twofootbot/Lv2.py | 101 ++++++++++++++ .../manager_based/twofootbot/__init__.py | 34 ++++- .../twofootbot/sequence_command.py | 78 +++++++++++ 6 files changed, 445 insertions(+), 10 deletions(-) create mode 100644 source/TwoFootBot/TwoFootBot/tasks/manager_based/twofootbot/Lv1.py create mode 100644 source/TwoFootBot/TwoFootBot/tasks/manager_based/twofootbot/Lv2.py create mode 100644 source/TwoFootBot/TwoFootBot/tasks/manager_based/twofootbot/sequence_command.py diff --git a/scripts/rsl_rl/play.py b/scripts/rsl_rl/play.py index 7f6e66e..23502c6 100644 --- a/scripts/rsl_rl/play.py +++ b/scripts/rsl_rl/play.py @@ -60,8 +60,8 @@ with contextlib.suppress(ImportError): # -- argparse ---------------------------------------------------------------- parser = argparse.ArgumentParser(description="Train an RL agent with RSL-RL.") -parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.") -parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).") +parser.add_argument("--video", action="store_true", default=True, help="Record videos during training.") +parser.add_argument("--video_length", type=int, default=800, help="Length of the recorded video (in steps).") parser.add_argument( "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations." ) diff --git a/source/TwoFootBot/TwoFootBot/tasks/manager_based/twofootbot/Lv0.py b/source/TwoFootBot/TwoFootBot/tasks/manager_based/twofootbot/Lv0.py index d6cdce6..fa27f88 100644 --- a/source/TwoFootBot/TwoFootBot/tasks/manager_based/twofootbot/Lv0.py +++ b/source/TwoFootBot/TwoFootBot/tasks/manager_based/twofootbot/Lv0.py @@ -6,6 +6,7 @@ import math from pathlib import Path +import torch import isaaclab.sim as sim_utils from isaaclab.actuators import ImplicitActuatorCfg from isaaclab.assets import ArticulationCfg, AssetBaseCfg @@ -29,6 +30,86 @@ from . import mdp # Pre-defined configs ## +JOINT_NAMES = [ # 定义机器人需要控制的八个关节名称。 + "jlf0_link", + "jlf1_link", + "jlf2_link", + "jlf3_link", + "jrf0_link", + "jrf1_link", + "jrf2_link", + "jrf3_link", +] +BOT_PACKAGE_PATH = Path(__file__).resolve().parents[6] / "bot/2foot" # 获取机器人 ROS 包目录。 +URDF_PATH = BOT_PACKAGE_PATH / "urdf/2foot.urdf" # 指定正式使用的 2foot URDF 文件。 + +def spawn_twofootbot_from_urdf(prim_path, cfg, translation=None, orientation=None, **kwargs): + """Spawn the URDF and enable contact reporting on every rigid link.""" + from pxr import Usd, UsdPhysics + from isaaclab.sim.spawners.from_files.from_files import spawn_from_urdf + from isaaclab.sim.utils import safe_set_attribute_on_usd_prim + + spawn_cfg = cfg.replace(activate_contact_sensors=False) + prim = spawn_from_urdf(prim_path, spawn_cfg, translation, orientation, **kwargs) + for child_prim in Usd.PrimRange(prim): + if not child_prim.HasAPI(UsdPhysics.RigidBodyAPI): + continue + applied_schemas = child_prim.GetAppliedSchemas() + if "PhysxRigidBodyAPI" not in applied_schemas: + child_prim.AddAppliedSchema("PhysxRigidBodyAPI") + if "PhysxContactReportAPI" not in applied_schemas: + child_prim.AddAppliedSchema("PhysxContactReportAPI") + safe_set_attribute_on_usd_prim( + child_prim, "physxRigidBody:sleepThreshold", 0.0, camel_case=False + ) + safe_set_attribute_on_usd_prim( + child_prim, "physxContactReport:threshold", 0.0, camel_case=False + ) + return prim + +TWOFOOTBOT_URDF_CFG = ArticulationCfg( + spawn=sim_utils.UrdfFileCfg( + func=spawn_twofootbot_from_urdf, + asset_path=str(URDF_PATH), + fix_base=False, + activate_contact_sensors=False, + ros_package_paths=[{"name": "2foot", "path": str(BOT_PACKAGE_PATH)}], + joint_drive=sim_utils.UrdfFileCfg.JointDriveCfg( + drive_type="force", + target_type="position", + gains=sim_utils.UrdfFileCfg.JointDriveCfg.PDGainsCfg(stiffness=17.0, damping=0.3), + ), + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, # 不启用重力。 + max_linear_velocity=100.0, + max_angular_velocity=100.0, + max_depenetration_velocity=5.0, + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=True, # # 启用自碰撞检测。 + solver_position_iteration_count=8, + solver_velocity_iteration_count=1, + ), + ), + init_state=ArticulationCfg.InitialStateCfg( # 初始状态配置。单位： 弧度 + pos=(0.0, 0.0, 0.18), + joint_pos={"jlf0_link": 0.0, "jlf1_link": 0.0, "jlf2_link": 0.0, "jlf3_link": 0.0, + "jrf0_link": 0.0, "jrf1_link": 0.0, "jrf2_link": 0.0, "jrf3_link": 0.0}, + joint_vel={".*": 0.0}, + ), + actuators={ # 动力配置。 + "legs": ImplicitActuatorCfg( + joint_names_expr=JOINT_NAMES, # 控制所有腿的关节。 + effort_limit_sim=0.5, # 力矩限制。单位：牛顿米。 + velocity_limit_sim=5, # 速度限制 (rad/s) + stiffness=5.0, # 刚度。单位：牛顿/米。 + damping=0.3, # 阻尼。单位：牛顿秒/米。 + ) + }, +) + + + DEMO_ROOT = Path(__file__).resolve().parents[6] URDF_PATH = DEMO_ROOT / "bot" / "2foot" / "urdf" / "2foot.urdf" BOT_PACKAGE_PATH = DEMO_ROOT / "bot" / "2foot" @@ -42,6 +123,14 @@ JOINT_NAMES = [ "jrf2_link", "jrf3_link", ] +IMU_OFFSET_B = (0.0, 0.0, 0.091) + + +def imu_lin_vel(env, asset_cfg: SceneEntityCfg, offset_b: tuple[float, float, float]) -> torch.Tensor: + """Return the linear velocity at the IMU position, expressed in the body frame.""" + asset = env.scene[asset_cfg.name] + offset = asset.data.root_ang_vel_b.new_tensor(offset_b).expand_as(asset.data.root_ang_vel_b) + return asset.data.root_lin_vel_b + torch.linalg.cross(asset.data.root_ang_vel_b, offset, dim=1) TWOFOOTBOT_CFG = ArticulationCfg( spawn=sim_utils.UrdfFileCfg( @@ -61,7 +150,7 @@ TWOFOOTBOT_CFG = ArticulationCfg( max_depenetration_velocity=5.0, ), articulation_props=sim_utils.ArticulationRootPropertiesCfg( - enabled_self_collisions=False, + enabled_self_collisions=True, solver_position_iteration_count=8, solver_velocity_iteration_count=1, ), @@ -99,7 +188,7 @@ class TwofootbotdemoSceneCfg(InteractiveSceneCfg): ) # robot - robot: ArticulationCfg = TWOFOOTBOT_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + robot: ArticulationCfg = TWOFOOTBOT_URDF_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") contact_forces = None # lights @@ -130,12 +219,16 @@ class CommandsCfg: base_velocity = velocity_mdp.UniformVelocityCommandCfg( asset_name="robot", resampling_time_range=(10.0, 10.0), + # 相对航向环境参数，用于控制或影响相对航向相关的环境配置 + # 该参数的默认值为0.0，表示在不设置特殊航向环境时的基准值 rel_standing_envs=0.0, + # 相对航向环境参数，默认值为0.0 rel_heading_envs=0.0, + # 航向命令标志，默认为False heading_command=False, debug_vis=True, ranges=velocity_mdp.UniformVelocityCommandCfg.Ranges( - lin_vel_x=(0.35, 0.35), + lin_vel_x=(0.2, 0.2), lin_vel_y=(0.0, 0.0), ang_vel_z=(0.0, 0.0), heading=(-math.pi, math.pi), @@ -152,14 +245,27 @@ class ObservationsCfg: """Observations for policy group.""" # observation terms (order preserved) - base_lin_vel = ObsTerm(func=velocity_mdp.base_lin_vel, noise=UniformNoiseCfg(n_min=-0.1, n_max=0.1)) +# 创建观测项，用于定义环境中的各种观测信息 +# 每个观测项都包含一个函数和一个可选的噪声配置 +# 基线线性速度观测项，添加均匀噪声，噪声范围在-0.1到0.1之间 + base_lin_vel = ObsTerm( + func=imu_lin_vel, + params={"asset_cfg": SceneEntityCfg("robot"), "offset_b": IMU_OFFSET_B}, + noise=UniformNoiseCfg(n_min=-0.1, n_max=0.1), + ) +# 基线角速度观测项，添加均匀噪声，噪声范围在-0.2到0.2之间 base_ang_vel = ObsTerm(func=velocity_mdp.base_ang_vel, noise=UniformNoiseCfg(n_min=-0.2, n_max=0.2)) +# 投影重力观测项，不添加噪声 projected_gravity = ObsTerm(func=velocity_mdp.projected_gravity) +# 速度命令观测项，使用特定的命令名称"base_velocity" velocity_commands = ObsTerm( func=velocity_mdp.generated_commands, params={"command_name": "base_velocity"} ) +# 相对关节位置观测项，使用场景中的机器人实体 joint_pos_rel = ObsTerm(func=velocity_mdp.joint_pos_rel, params={"asset_cfg": SceneEntityCfg("robot")}) +# 相对关节速度观测项，使用场景中的机器人实体 joint_vel_rel = ObsTerm(func=velocity_mdp.joint_vel_rel, params={"asset_cfg": SceneEntityCfg("robot")}) +# 上一个动作观测项 actions = ObsTerm(func=velocity_mdp.last_action) def __post_init__(self) -> None: @@ -175,7 +281,9 @@ class EventCfg: """Configuration for events.""" # reset +# 定义一个基础重置事件项，用于重置机器人的初始状态 reset_base = EventTerm( + # 指定重置函数为velocity_mdp模块中的reset_root_state_uniform函数 func=velocity_mdp.reset_root_state_uniform, mode="reset", params={ @@ -183,7 +291,9 @@ class EventCfg: "velocity_range": {"x": (-0.05, 0.05), "y": (-0.05, 0.05), "z": (-0.05, 0.05)}, }, ) +# 定义一个重置关节的事件项 reset_joints = EventTerm( + # 指定重置关节的函数，使用velocity_mdp模块中的reset_joints_by_scale方法 func=velocity_mdp.reset_joints_by_scale, mode="reset", params={ @@ -227,7 +337,7 @@ class TerminationsCfg: # (3) The base has collapsed onto the ground. fallen_height = DoneTerm( func=mdp.root_height_below_minimum, - params={"minimum_height": 0.08}, + params={"minimum_height": 0.10}, ) @@ -238,7 +348,7 @@ class TerminationsCfg: @configclass class TwofootbotdemoEnvCfg(ManagerBasedRLEnvCfg): - scene: TwofootbotdemoSceneCfg = TwofootbotdemoSceneCfg(num_envs=1024, env_spacing=1.5) + scene: TwofootbotdemoSceneCfg = TwofootbotdemoSceneCfg(num_envs=8192, env_spacing=1.5) # Basic settings observations: ObservationsCfg = ObservationsCfg() actions: ActionsCfg = ActionsCfg() @@ -253,7 +363,7 @@ class TwofootbotdemoEnvCfg(ManagerBasedRLEnvCfg): """Post initialization.""" # general settings self.decimation = 4 - self.episode_length_s = 20.0 + self.episode_length_s = 40.0 #time_out # viewer settings self.viewer.eye = (3.0, 2.5, 1.8) # simulation settings diff --git a/source/TwoFootBot/TwoFootBot/tasks/manager_based/twofootbot/Lv1.py b/source/TwoFootBot/TwoFootBot/tasks/manager_based/twofootbot/Lv1.py new file mode 100644 index 0000000..6815ec7 --- /dev/null +++ b/source/TwoFootBot/TwoFootBot/tasks/manager_based/twofootbot/Lv1.py @@ -0,0 +1,114 @@ +from isaaclab.managers import RewardTermCfg as RewTerm +from isaaclab.utils.configclass import configclass +from isaaclab.managers import ( + RewardTermCfg, +) +import math +import isaaclab_tasks.manager_based.locomotion.velocity.mdp as mdp +from isaaclab.managers import SceneEntityCfg +from isaaclab.sensors import ContactSensorCfg + + + +from TwoFootBot.tasks.manager_based.twofootbot.Lv0 import ( + CommandsCfg, + TwofootbotdemoEnvCfg, + RewardsCfg, + TwofootbotdemoSceneCfg, + +) + + +@configclass +class Lv1SceneCfg(TwofootbotdemoSceneCfg): + # 新增：接触力传感器 + contact_forces = ContactSensorCfg( + prim_path="{ENV_REGEX_NS}/Robot/.*", # 或 "/World/envs/env_.*/Robot/.*" + history_length=3, + track_air_time=True, + debug_vis=False, + ) + + + + + +@configclass +class Lv1RewardsCfg(RewardsCfg): + """Configuration for the first level of the TwoFootBot task.""" +# 定义跟踪角速度z方向的指数奖励项配置 + # track_ang_vel_z_exp = RewardTermCfg( + # func=mdp.track_ang_vel_z_exp, # 奖励函数，用于跟踪角速度z方向的指数 + # weight=0.5, # 奖励项权重，控制该奖励项的重要性 + # params={"command_name": "base_velocity", "std": 0.25}, # 命令名称和标准差参数 + # ) + + #奖励抬脚 + feet_air_time = RewTerm( + func=mdp.feet_air_time_positive_biped, # 调用的函数，用于计算双足机器人脚部空中时间 + weight=0.5, # 该奖励项的权重，表示在总奖励中的重要性 + params={ # 函数所需的参数配置 + "command_name": "base_velocity", # 命令名称，可能与基础速度控制相关 + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=["lf3_link", "rf3_link"]), # 传感器配置，指定监测接触力的脚部刚体 + "threshold": 0.4, # 阈值，用于判断脚部是否接触地面 + }, + ) + #打滑扣分 + # feet_slide = RewTerm( + # # 指定使用的函数为mdp.feet_slide + # func=mdp.feet_slide, + # # 设置权重为-0.1，表示脚部滑动会带来负奖励 + # weight=-5e-3, + # # 配置参数，包括传感器和资产的设置 + # params={ + # # 配置接触力传感器，监测所有脚踝滚动链接的接触力 + # "sensor_cfg": SceneEntityCfg("contact_forces", body_names=["lf3_link", "rf3_link"]), + # # 配置机器人实体，指定监测所有脚踝滚动链接 + # "asset_cfg": SceneEntityCfg("robot", body_names=["lf3_link", "rf3_link"]), + # }, + # ) + + # 定义奖励项配置：动作速率的L2范数惩罚 + action_rate_l2 = RewardTermCfg(func=mdp.action_rate_l2, weight=-5e-3) + # 定义奖励项配置：关节力矩最小化 + joint_torque_minimization = RewardTermCfg(func=mdp.joint_torques_l2, weight=-1e-4) + # 定义奖励项配置：关节速度最小化 + joint_velocity_minimization = RewardTermCfg(func=mdp.joint_vel_l2, weight=-1e-2) + + + + +@configclass +class Lv1CommandsCfg(CommandsCfg): + base_velocity = mdp.UniformVelocityCommandCfg( + # 资产名称设置为"robot" + asset_name="robot", + # 重采样时间范围设置为(10.0, 10.0) + resampling_time_range=(4.0, 6.0), + # 相对站立环境参数设置为0.02 + rel_standing_envs=0.02, + # 相对朝向环境参数设置为1.0 + rel_heading_envs=1.0, + # 启用朝向命令 + heading_command=False, + # 设置朝向控制刚度为0.5 + #heading_control_stiffness=0.5, + debug_vis=True, + ranges=mdp.UniformVelocityCommandCfg.Ranges( + lin_vel_x=(0.2, 0.5), lin_vel_y=(-0.0, 0.0), ang_vel_z=(-0.0, 0.0), heading=(-math.pi, math.pi) + ), + ) + + +@configclass +class Lv1Cfg(TwofootbotdemoEnvCfg): + """Configuration for the first level of the TwoFootBot task.""" + commands: Lv1CommandsCfg = Lv1CommandsCfg() + rewards: Lv1RewardsCfg = Lv1RewardsCfg() + scene: Lv1SceneCfg = Lv1SceneCfg(num_envs=8192*2, env_spacing=1.5) + decimation = 4 + + + def __post_init__(self) -> None: + """Post initialization.""" + self.episode_length_s = 80.0 # time_out \ No newline at end of file diff --git a/source/TwoFootBot/TwoFootBot/tasks/manager_based/twofootbot/Lv2.py b/source/TwoFootBot/TwoFootBot/tasks/manager_based/twofootbot/Lv2.py new file mode 100644 index 0000000..671caca --- /dev/null +++ b/source/TwoFootBot/TwoFootBot/tasks/manager_based/twofootbot/Lv2.py @@ -0,0 +1,101 @@ +from isaaclab.managers import SceneEntityCfg +from isaaclab.utils.configclass import configclass +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab_tasks.utils import preset + + +import isaaclab_tasks.manager_based.locomotion.velocity.mdp as mdp +from TwoFootBot.tasks.manager_based.twofootbot.Lv1 import Lv1Cfg +from TwoFootBot.tasks.manager_based.twofootbot.Lv0 import EventCfg + + + + +@configclass +class Lv2EventCfg(EventCfg): + # 创建一个物理材质事件项 + physics_material = EventTerm( + # 指定函数为随机化刚体材质的函数 + func=mdp.randomize_rigid_body_material, + # 设置事件模式为启动时执行 + mode="startup", + # 设置事件参数 + params={ + # 配置场景实体，指定为机器人，并匹配所有身体部位 + "asset_cfg": SceneEntityCfg("robot", body_names=".*"), + # 设置静态摩擦力的范围为固定值0.8 + "static_friction_range": (0.8, 0.8), + # 设置动态摩擦力的范围为固定值0.6 + "dynamic_friction_range": (0.6, 0.6), + # 设置恢复系数的范围为固定值0.0（完全非弹性碰撞） + "restitution_range": (0.0, 0.0), + # 设置材质分桶数量为64 + "num_buckets": 64, + }, + ) + + # 定义一个基础质量添加事件项 + add_base_mass = EventTerm( + # 设置事件项的执行函数为随机化刚体质量 + func=mdp.randomize_rigid_body_mass, + # 设置事件项的执行模式为"startup"，表示在启动时执行 + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names="head_link"), + "mass_distribution_params": (1 / 1.25, 1.25), + "operation": "scale", + "distribution": "log_uniform", + }, + ) + + # 预设基座质心随机化配置 + base_com = preset( # 使用preset函数创建预设配置 + default=EventTerm( # 默认事件项配置 + func=mdp.randomize_rigid_body_com, # 执行的函数：随机化刚体质心 + mode="startup", # 事件模式：启动时执行 + params={ # 函数参数配置 + "asset_cfg": SceneEntityCfg("robot", body_names="head_link"), # 场景实体配置：机器人基座 + "com_range": {"x": (-0.05, 0.05), "y": (-0.05, 0.05), "z": (-0.01, 0.01)}, + }, + ), + newton_mjwarp=None, + ) + + + + # 定义一个重置机器人关节的事件项 + reset_robot_joints = EventTerm( + # 指定事件处理函数为mdp.reset_joints_by_scale + func=mdp.reset_joints_by_scale, + # 设置事件模式为"reset" + mode="reset", + # 设置事件参数 + params={ + # 设置位置范围为(0.5, 1.5) + "position_range": (0.5, 1.5), + # 设置速度范围为(0.0, 0.0) + "velocity_range": (0.0, 0.0), + }, + ) + + # interval + # 创建一个名为push_robot的EventTerm对象，用于执行机器人推动操作 + push_robot = EventTerm( + func=mdp.push_by_setting_velocity, # 指定执行函数为mdp中的push_by_setting_velocity + mode="interval", # 设置执行模式为间隔模式 + interval_range_s=(0.0, 3.0), # 设置执行间隔时间范围为10.0到15.0秒 + params={"velocity_range": {"x": (-0.5, 0.5), "y": (-0.05, 0.05)}}, # 设置推动速度范围，x和y方向的速度范围均为-0.5到0.5 + ) + + +@configclass +class Lv2Cfg(Lv1Cfg): + """Configuration for the first level of the TwoFootBot task.""" + events: Lv2EventCfg = Lv2EventCfg() + + decimation = 4 + + + def __post_init__(self) -> None: + """Post initialization.""" + self.episode_length_s = 120.0 #time_out diff --git a/source/TwoFootBot/TwoFootBot/tasks/manager_based/twofootbot/__init__.py b/source/TwoFootBot/TwoFootBot/tasks/manager_based/twofootbot/__init__.py index 2df1fdc..a611f63 100644 --- a/source/TwoFootBot/TwoFootBot/tasks/manager_based/twofootbot/__init__.py +++ b/source/TwoFootBot/TwoFootBot/tasks/manager_based/twofootbot/__init__.py @@ -51,12 +51,44 @@ gym.register( }, )''' + +#$ python scripts/rsl_rl/train.py --task Template-Twofootbot-mlp-v0 gym.register( - id="Template-Twofootbotv0-Demo", + id="Template-Twofootbot-mlp-v0", entry_point="isaaclab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ "env_cfg_entry_point": f"{__name__}.Lv0:TwofootbotdemoEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:PPORunnerCfg", }, +) + + +# $ python scripts/rsl_rl/train.py \ +# --task Template-Twofootbot-mlp-v1 \ +# --headless \ +# --resume \ +# --load_run 2026-09-15_13-16-41-V0 \ +# --checkpoint model_4999.pt --device cuda:0 + +gym.register( + id="Template-Twofootbot-mlp-v1", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.Lv1:Lv1Cfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:PPORunnerCfg", + }, +) + +# Lv2Cfg + +gym.register( + id="Template-Twofootbot-mlp-v2", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.Lv2:Lv2Cfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:PPORunnerCfg", + }, ) \ No newline at end of file diff --git a/source/TwoFootBot/TwoFootBot/tasks/manager_based/twofootbot/sequence_command.py b/source/TwoFootBot/TwoFootBot/tasks/manager_based/twofootbot/sequence_command.py new file mode 100644 index 0000000..fd4a61e --- /dev/null +++ b/source/TwoFootBot/TwoFootBot/tasks/manager_based/twofootbot/sequence_command.py @@ -0,0 +1,78 @@ +"""自定义速度命令：直走 -> 原地转向 -> 直走 循环。 + +使用 heading_command=True，转向由 heading 目标平滑驱动。 +""" + +from __future__ import annotations + +import math +import torch +from collections.abc import Sequence + +import isaaclab_tasks.manager_based.locomotion.velocity.mdp as mdp +from isaaclab.utils.configclass import configclass + + +class SequenceHeadingWalkTurnCommand(mdp.UniformVelocityCommand): + """直走 -> 原地转向 -> 直走 循环，转向通过 heading 目标实现。 + + 阶段循环：0(直走) -> 1(转向) -> 2(直走) -> 0(直走) -> ... + 每个阶段持续 resampling_time_range 指定的时间。 + + 与直接写 vel_command_b[2] 的区别： + 这里只设置 heading_target，由基类的 _update_command 根据 heading 误差 + 实时计算角速度，转向更平滑。 + """ + + cfg: "SequenceHeadingWalkTurnCommandCfg" + + def __init__(self, cfg: "SequenceHeadingWalkTurnCommandCfg", env): + super().__init__(cfg, env) + # 当前阶段：0=直走，1=转向，2=直走 + self._phase = torch.zeros(self.num_envs, dtype=torch.long, device=self.device) + + def _resample_command(self, env_ids: Sequence[int]): + # 先让基类采样默认速度 / heading，避免遗漏默认行为 + super()._resample_command(env_ids) + + env_ids_t = torch.as_tensor(env_ids, dtype=torch.long, device=self.device) + phase = self._phase[env_ids_t] + num = len(env_ids_t) + + # 前进速度：在配置范围内随机 + vx_min, vx_max = self.cfg.ranges.lin_vel_x + vx = torch.empty(num, device=self.device).uniform_(vx_min, vx_max) + + # 当前世界朝向 + current_heading = self.robot.data.heading_w[env_ids_t] + + # 阶段 0 和 2：直走，heading 锁定当前朝向 + mask_walk = (phase == 0) | (phase == 2) + if mask_walk.any(): + ids = env_ids_t[mask_walk] + self.vel_command_b[ids, 0] = vx[mask_walk] + self.vel_command_b[ids, 1] = 0.0 + # vel_command_b[2] 由基类 _update_command 根据 heading 误差计算 + self.heading_target[ids] = current_heading[mask_walk] + + # 阶段 1：原地转向，heading 目标 = 当前朝向 + 90°（左转） + mask_turn = phase == 1 + if mask_turn.any(): + ids = env_ids_t[mask_turn] + self.vel_command_b[ids, 0] = 0.0 + self.vel_command_b[ids, 1] = 0.0 + self.heading_target[ids] = current_heading[mask_turn] + math.pi / 2.0 + + # 推进 phase + self._phase[env_ids_t] = (self._phase[env_ids_t] + 1) % 3 + + def _update_command(self): + # 基类会处理 heading_command=True 时的角速度计算 + super()._update_command() + + +@configclass +class SequenceHeadingWalkTurnCommandCfg(mdp.UniformVelocityCommandCfg): + """序列 heading 命令的配置类。""" + + class_type: type = SequenceHeadingWalkTurnCommand \ No newline at end of file